diff --git a/.agents/skills/database-access/SKILL.md b/.agents/skills/database-access/SKILL.md index 13fb7b7e286..c439d0fe352 100644 --- a/.agents/skills/database-access/SKILL.md +++ b/.agents/skills/database-access/SKILL.md @@ -1,54 +1,73 @@ --- name: database-access -description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope +description: Organize database repositories and proxy dependency wiring, manage connection ownership, or port Prisma operations to Rust and SQLx without expanding migration scope metadata: internal: true --- # Database access +## Situations + +- Adding a query or an operation to a repository, read [adding a query](references/python-repository.md#adding-a-query) +- Moving existing Prisma calls behind a repository, read [moving existing access](references/python-repository.md#moving-existing-access-into-a-repository) +- Porting a repository from Prisma to Rust and SQLx, read [porting to SQLx](references/python-repository.md#porting-a-repository-to-sqlx) +- Injecting a repository into a route, read [Python and FastAPI](references/python-fastapi.md) or [Rust and Axum](references/rust-axum.md) + ## Adding new database access - Put Python database access behind an operation in `litellm/repositories` -- Give callers typed inputs and outputs, not Prisma clients, models, or query dictionaries -- Keep query execution and backend selection inside the repository +- Give each repository its own FastAPI provider, so a route builds only the repositories it uses +- A provider borrows existing handles and performs no database I/O during construction +- Define the operations a route calls as a Protocol in `litellm/models/.py`, next to its domain types + - The Protocol takes and returns domain types, never Prisma models, raw rows or serialized payloads + - Parsing and serialization of stored values live in the repository + - Add a Protocol when a repository gets a provider, not up front for every repository +- Keep query execution inside the repository, with one implementation per repository and no backend switch or fallback +- A repository is either all Prisma or all native, never a mix of both - Prisma is fine for the implementation, adding a query does not require a Rust port - Native Rust callers use `litellm-rust/crates/db` directly -## Asked to move existing access behind a repository +## Framework wiring -- Identify the calls being moved and their transaction boundary -- Move the existing Prisma implementation into the relevant repository -- Replace those direct database calls with repository operations -- Preserve return values, errors, and transaction behavior, with focused regression tests -- Keep this an extraction, not a backend migration or unrelated cleanup +- Python routes receive a per-repository dependency such as `UserBannerRepositoryDep` and call its operations, such as `banners.get_banner()` +- `DatabaseDep` lives in `litellm/proxy/deps/database.py`, repository providers in `litellm/proxy/deps/repositories.py` + - A provider returns the feature's Protocol, so routes never see which backend implements it + - Dependencies used by one feature only, like the banner audit binding, stay next to that feature's routes +- Repositories never import FastAPI, so jobs and helpers can use them without a request +- Use named FastAPI providers with `Annotated[..., Depends(...)]`, ordinary helpers and jobs receive explicit arguments +- Preserve each route's no-database behavior, authorization precedence and background-task scheduling +- Read [Python and FastAPI](references/python-fastapi.md) for provider examples, lifespan state and route verification +- Read [Rust and Axum](references/rust-axum.md) when adding database access to Rust handlers -## An equivalent Rust or SQLx query already exists +## Connection ownership -- Leave it and the current backend selection unchanged unless integration is explicitly requested -- Its existence alone does not justify switching a Prisma caller to Rust -- Repository extraction can keep using Prisma even when the Rust query is available +- Keep connection setup and shutdown separate from repositories, which own no pools +- Reuse existing handles, never connect per request or disconnect from a repository operation +- Give each pool one owner for setup and shutdown, borrowed handles never close it +- Finish requests, jobs and queued database work before closing connections +- Clean up resources opened before a startup failure, without a second disconnect fallback +- Add a backend to the existing lifecycle owner and pass its handle only to repositories that need it +- Preserve cancellation, timeout, error and transaction semantics across the native bridge -## Asked to port a Prisma query to Rust or SQLx +## Changing ownership or reloading resources -- Identify the existing inputs, outputs, errors, and transaction boundary -- Check `litellm-rust/crates/db` for an equivalent query and reuse it if its behavior matches -- Otherwise follow `litellm-rust/crates/db/AGENTS.md` - - Put static SQL in `queries/` and compile-check it with SQLx - - Keep typed inputs and reusable row projections in `src/` - - Refresh SQLx query metadata after SQL changes -- Keep Python conversion in the Python bridge, not in `litellm-rust/crates/db` -- When caller integration is requested, expose the Rust operation through the bridge and call it from the repository - - A query-only port can land without switching callers -- Verify behavior against real PostgreSQL before switching callers - - Check relevant null handling, conflicts, ordering, errors, and transaction semantics -- Remove the old implementation only when no callers need it +- Inventory initialization, reload, tests, background jobs and shutdown together +- Move every writer when changing authoritative storage, copying a replaceable global into state is insufficient +- A frozen repository does not make its clients immutable +- Decide whether each request captures one resource generation or reads it again during processing +- Build and validate replacements before publishing them, preserving the old generation on failure +- Serialize reloads and coordinate publication with shutdown +- Drain streams and background users before retiring old resources, retaining every generation grows resource use +- Define multi-worker coordination separately from an in-process swap +- Choose mechanisms such as `ArcSwap` only after defining publication, capture and retirement ## Migration context and boundaries -- Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations incrementally with Rust APIs -- Mixed Prisma and Rust implementations are expected during the transition, unrelated existing calls can stay as they are -- Both backends may have separate pools through PgBouncer, size them together and verify pooling-mode compatibility +- Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations with Rust APIs one repository at a time +- Prisma and native repositories coexist during the transition, unrelated existing calls can stay as they are +- Prisma and SQLx keep separate client pools to PgBouncer, which pools connections to PostgreSQL + - Size both layers together and preserve reader/writer routing and pooling-mode compatibility - PgBouncer does not share transactions between Prisma and SQLx, keep each atomic unit of work on one backend and transaction - Never retry a failed Rust write through Prisma when it may already have committed - Prisma owns schema migrations today, eventual SQLx migration ownership and Prisma removal are separate planned steps, not part of a query port diff --git a/.agents/skills/database-access/references/python-fastapi.md b/.agents/skills/database-access/references/python-fastapi.md new file mode 100644 index 00000000000..2f6d1b001fc --- /dev/null +++ b/.agents/skills/database-access/references/python-fastapi.md @@ -0,0 +1,101 @@ +# Python and FastAPI + +## Existing pattern + +- `UserBannerStore` in `litellm/models/user_banner.py` is the Protocol the banner routes call +- `get_database` in `litellm/proxy/deps/database.py` reads the lifespan-owned handle from request state +- `get_user_banner_repository` in `litellm/proxy/deps/repositories.py` builds the banner repository from that handle +- Choose each repository's replica routing once in its provider, never at call sites +- A handler receives `UserBannerStore | None`, never a Prisma client or a concrete repository +- `UserBannerRepository` implements `UserBannerStore` structurally, keeping the existing Prisma table, serialization and configuration-sync behavior +- The banner audit provider binds the legacy audit callable at request time, without performing audit I/O +- A native repository can later implement the same Protocol without changing handler signatures, see [porting a repository to SQLx](python-repository.md#porting-a-repository-to-sqlx) + +## Dependency injection + +- Add one provider and one `Annotated` alias per repository in `litellm/proxy/deps/repositories.py` +- Annotate the provider and alias with the feature's Protocol, not the concrete repository +- A route depends only on the repositories it calls, so a request never builds repositories it does not use +- Do not build repositories once at lifespan, a reconnect can replace the Prisma client they wrap + +```python +class UserBannerStore(Protocol): + async def get_banner(self) -> UserBanner: ... + + async def save_banner(self, banner: UserBanner) -> None: ... +``` + +```python +async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None: + return None if database is None else UserBannerRepository(database.prisma) + + +UserBannerRepositoryDep: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)] +``` + +- FastAPI caches each provider within a request, so routes and sub-dependencies share one instance +- Ordinary helpers and jobs receive explicit arguments, never unresolved `Depends` defaults +- A route still reading the `prisma_client` global or building a repository inline switches to the repository's `...RepositoryDep`, keeping its no-database behavior and error precedence + +```python +async def get_user_banner(banners: UserBannerRepositoryDep) -> UserBanner: + if banners is None: + return UserBanner() + return await banners.get_banner() +``` + +## Lifespan state + +- Prisma startup owns connection setup, authentication refresh, health checks and migrations +- `Database` holds the initialized Prisma client and the `NativeDatabase` handle, without holding fields per feature + - The native handle opens no pool yet, its `connect` and `close` are noops +- The proxy lifespan builds the handle with `Database.open`, registers `Database.aclose` on its exit stack and publishes it in request state + - Callers only open and close `Database`, which closes Prisma and the native handle together, and `open` closes Prisma itself when the native connect fails +- Starlette shallow-copies lifespan state into request state, which is distinct from `app.state` +- The database module never imports a repository, legacy globals borrow its handle without owning shutdown +- Drain requests, jobs, logging, spend and audit work before database teardown +- The exit stack closes the handle on exit, including startup failures after the client is created +- When adding native resources, initialize them in the serving process, respecting runtime fork restrictions +- Await native shutdown instead of depending on Python garbage collection +- Apply the [connection ownership rules](../SKILL.md#connection-ownership) when changing the owner + +## Timing and errors + +- FastAPI resolves dependencies before entering the handler +- Handler exception handling cannot catch failures raised during dependency resolution +- A universal required-database dependency can override authorization and validation behavior +- Banner GET without a database returns the default disabled banner +- Banner PATCH checks the admin role before returning the existing missing-database error +- Invalid stored banner data falls back to the default, but storage failures still propagate +- Banner writes finish before the success audit is queued as a FastAPI background task +- Queue post-response work with `BackgroundTasks`, not bare `asyncio.create_task`, so the in-flight drain waits for it before the database closes +- Do not move background work into an awaited provider merely to remove an import + +## Verification + +- Exercise real FastAPI handlers using `dependency_overrides` for external boundaries +- Keep production parsing and persistence code real +- Restore overrides after each test +- Verify complete response bodies, status codes and OpenAPI against the base +- Include combined failure conditions to pin error precedence +- Check saved values, fresh revisions, audit contents and failed-write behavior +- Check lifespan state with a running lifespan, not only a TestClient that skips startup +- Verify providers do not connect, query, disconnect or audit eagerly +- Use real PostgreSQL and live proxy requests when changing persistence behavior + +## Scope + +- Preserve direct Python call sites or migrate them explicitly +- Inventory imports, state writers, side effects and error ordering before migrating a feature +- Preserve configuration sync and serialization, and state which legacy reads remain +- Do not migrate every route in a large module just to establish the pattern +- Do not invent providers for resources without consumers + +## Framework references + +- [Lifespan events](https://fastapi.tiangolo.com/advanced/events/) and [Starlette lifespan state](https://starlette.dev/lifespan/) +- [`SessionDep`-style `Annotated` aliases](https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-session-dependency), which `DatabaseDep` and `UserBannerRepositoryDep` follow +- [Dependency composition and caching](https://fastapi.tiangolo.com/tutorial/dependencies/sub-dependencies/) +- [Background tasks and dependencies with `yield`](https://fastapi.tiangolo.com/advanced/advanced-dependencies/#background-tasks-and-dependencies-with-yield-technical-details) +- [Dependency overrides](https://fastapi.tiangolo.com/advanced/testing-dependencies/) +- [Testing lifespan](https://fastapi.tiangolo.com/advanced/testing-events/) diff --git a/.agents/skills/database-access/references/python-repository.md b/.agents/skills/database-access/references/python-repository.md new file mode 100644 index 00000000000..3b16ef24de6 --- /dev/null +++ b/.agents/skills/database-access/references/python-repository.md @@ -0,0 +1,119 @@ +# Python repositories + +## Adding a query + +Use this when a feature needs a read or write that no repository offers yet, including one more operation on an existing repository + +### Before starting + +- Search `litellm/repositories` for a repository that already owns the table, and extend it instead of adding a second one +- Check `litellm/models/.py` for an existing Protocol and domain types +- Decide the route's behavior when no database is configured, and where that check sits relative to authorization and validation + +### Steps + +1. Domain types and Protocol in `litellm/models/.py` + - The Protocol takes and returns domain types, never Prisma models, raw rows or serialized payloads + - Add the Protocol together with the provider, not for repositories nothing injects +2. Repository in `litellm/repositories/_repository.py` + - Implement the operation with Prisma, a Rust port is not required + - Compose the Prisma client and type the tables it touches with `TableActions` from `prisma_protocols.py`, never inherit `PrismaTableRepository` + - Parse and serialize stored values here, so callers only see domain types + - Never import FastAPI, connect, disconnect or choose replicas per call +3. Provider and `Annotated` alias in `litellm/proxy/deps/repositories.py`, see [dependency injection](python-fastapi.md#dependency-injection) +4. The route takes `RepositoryDep`, handles `None` with the behavior decided above, and calls the operation +5. Jobs and helpers that need the operation receive the repository as an explicit argument, built by their caller from the handle it already holds + +### Done when + +- Handler tests drive the real route with `dependency_overrides` and the real repository code, see [verification](python-fastapi.md#verification) +- Tests cover the no-database path, error precedence, and what was actually stored +- Persistence changes are checked against real PostgreSQL through a live proxy + +### Out of scope + +- Porting the query to Rust, see [porting a repository to SQLx](#porting-a-repository-to-sqlx) +- Moving unrelated direct Prisma calls in the same module + +## Moving existing access into a repository + +Use this when code calls Prisma directly, through `prisma_client.db.` or a helper, and the calls should live behind a repository. The backend stays Prisma + +### Before starting + +- List every call being moved, its callers, and the transaction or batch it runs in +- List the module's imports, global state writers, side effects such as audit and config sync, and error ordering +- Record current return values, raised errors and response bodies, these are what the change must preserve + +### Steps + +1. Move the existing Prisma code into the relevant repository in `litellm/repositories` without rewriting the queries +2. Add domain types and a Protocol in `litellm/models/.py` when a route will inject the repository +3. Replace each direct call with the repository operation, through a provider for routes and explicit arguments for jobs and helpers, see [Python and FastAPI](python-fastapi.md) +4. Calls sharing a transaction move into the same repository or unit of work in one change +5. State which legacy reads remain and why, when the change does not move all of them + +### Done when + +- Focused regression tests pin the preserved return values, errors and transaction behavior +- Response bodies, status codes and OpenAPI match the base +- Configuration sync and serialization behave the same, including invalid stored data + +### Out of scope + +- Switching the backend, even when an equivalent Rust query exists +- Converting other routes in the module just to establish the pattern +- Cleanups unrelated to the moved calls + +## Porting a repository to SQLx + +Status: TBD. No repository has been ported yet. `database.native` is a `NativeDatabase` from `crates/python-bridge/src/db.rs` whose `connect` and `close` are noops that open no pool, and the PyO3 repository surface does not exist. A real pool takes typed writer/reader settings and a token provider, never the Prisma `DATABASE_URL`, which carries Prisma-only params and rotating IAM tokens. The constraints below are settled, the steps are a sketch to revise after the first port lands + +Use this only when a port is explicitly requested. Repository extraction does not require one + +### When not to port + +- An equivalent Rust or SQLx query existing is not a reason to switch a Prisma caller +- Leave existing backend selection unchanged unless integration is explicitly requested +- Never port part of a repository, it is either all Prisma or all native + +### Before starting + +- Identify the repository's inputs, outputs, errors and transaction boundary +- List repositories that share a transaction with it, they move together +- Check `litellm-rust/crates/db` for an equivalent query and reuse it when its behavior matches + +### Steps + +1. Rust queries, following `litellm-rust/crates/db/AGENTS.md` + - Static SQL in `queries/`, compile-checked with SQLx, typed inputs and row projections in `src/` + - Dynamic filters and sort keys become closed Rust enums that select static SQL, every value is bound + - Refresh SQLx query metadata after SQL changes +2. PyO3 surface exposing typed operations, never SQL text, table names or column names from Python. TBD: where it lives and how the handle reaches it +3. Native repository implementing the same Protocol in `litellm/repositories` + - Converts typed PyO3 results to domain types, preserving defaults, invalid-data handling and storage errors + - Python conversion stays in the bridge, not in `crates/db` +4. Switch the provider in `litellm/proxy/deps/repositories.py`, the Protocol and routes stay unchanged + +```diff + async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None: +- return None if database is None else UserBannerRepository(database.prisma) ++ return None if database is None else NativeUserBannerRepository(database.native) +``` + +5. Delete the Prisma repository in the same change and switch every caller, including jobs and helpers + +### Done when + +- Behavior is verified against real PostgreSQL, covering null handling, conflicts, ordering, errors and transaction semantics +- `DATABASE_URL=... cargo test -p litellm-db --locked -- --include-ignored` passes against a Prisma-migrated database +- Live proxy requests show unchanged responses. TBD: a Python-side real-PostgreSQL test harness + +### Rollback + +- Roll back the release, never retry an uncertain native write through Prisma + +### Out of scope + +- SQLx owning schema migrations, and removing Prisma +- Converting unrelated routes to dependency injection diff --git a/.agents/skills/database-access/references/rust-axum.md b/.agents/skills/database-access/references/rust-axum.md new file mode 100644 index 00000000000..83368d27233 --- /dev/null +++ b/.agents/skills/database-access/references/rust-axum.md @@ -0,0 +1,4 @@ +# Rust and Axum + +- This should use same `crates/db` what is exposed to Python world. +- Details are TBD. diff --git a/.github/workflows/test-rust-db.yml b/.github/workflows/test-rust-db.yml index bb88b593d5b..80ec256fa6f 100644 --- a/.github/workflows/test-rust-db.yml +++ b/.github/workflows/test-rust-db.yml @@ -42,6 +42,11 @@ concurrency: env: CARGO_TERM_COLOR: always + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" jobs: db-queries: @@ -87,6 +92,11 @@ jobs: workspaces: litellm-rust cache-on-failure: true + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Install Python dependencies working-directory: . run: uv sync --frozen --extra extra-proxy diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 04af770b24f..cf9fed192f6 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -78,6 +78,11 @@ concurrency: env: CARGO_TERM_COLOR: always + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" jobs: rust-lint: @@ -101,6 +106,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - run: cargo clippy --workspace --all-targets --locked -- -D warnings rust-test: @@ -143,6 +153,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Check generated trace contracts working-directory: . run: uv run scripts/generate_trace_types.py --check @@ -189,6 +204,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - run: uv build --wheel --out-dir dist - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 68823d99519..2c05ee827a3 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -73,9 +73,21 @@ jobs: timeout-minutes: 5 uses: ./.github/actions/cache-cargo-build + - name: Set up sccache + if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Build editable Rust bridge if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' timeout-minutes: 15 + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" run: | mkdir -p rust-bridge-dist uv run --no-project --with maturin==1.15.0 python -c 'import maturin; maturin.build_editable("rust-bridge-dist")' @@ -565,6 +577,7 @@ jobs: tests/unit/proxy/rerank_endpoints tests/unit/proxy/realtime_endpoints tests/unit/proxy/ui_crud_endpoints + tests/unit/proxy/deps tests/unit/proxy/config_resolvers tests/unit/proxy/utils python-version: "3.12" diff --git a/AGENTS.md b/AGENTS.md index ce4b63173d2..c3176351b8b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -90,6 +90,8 @@ Do not put names of customers or customer company names in code, PR descriptions CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`). Download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI +For database operations, connection ownership and proxy dependency wiring, follow [database-access](.agents/skills/database-access/SKILL.md) + Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: ` Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors): diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 531a78d8982..e956b27d5da 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -276,6 +276,8 @@ When the provider reports its own usage or spend numbers, compare the proxy's nu Provider work (a new provider, a new parameter, a new usage or cost field) is the most common kind of contribution, and it is where PRs get reworked most often. Follow these rules +**Preserve zero prices.** In [`model_prices_and_context_window.json`](model_prices_and_context_window.json), an explicit `0.0` means that the priced operation is free. A missing field or `null` means no rate is specified, so the relevant calculator may use a fallback. For example, when no batch input rate is configured, batch cache writes use `cache_creation_input_token_cost` if set, otherwise the regular input rate, then apply the batch discount. A zero ordinary input rate does not make separately priced cache reads or writes free. Test `rate is None` before substituting a fallback, since `row.get("cache_creation_input_token_cost") or fallback` would discard a valid zero. Speech keeps character pricing ahead of token pricing even when the character rate is zero. If neither rate is configured, speech cost calculation raises an error. Its model lookup preserves unset rates instead of treating compatibility defaults as explicit zero prices. Perplexity keeps a configured legacy search rate ahead of the current search rate, including zero. When registering custom prices, `None` leaves an existing rate unchanged and `0.0` overrides it + **Find the existing machinery before you add plumbing.** LiteLLM already has generic paths that move data between the endpoints, the provider transformations, and cost tracking. A second, hand-rolled path for the same data is a common reason a provider PR gets rewritten before merge. [`ARCHITECTURE.md`](ARCHITECTURE.md) shows where the layers live, and the decision records in [`docs/adrs/`](docs/adrs/AGENTS.md) explain why they are shaped that way. For example, [the ADR on provider usage fields](docs/adrs/adr_001_provider_usage_fields_ride_as_usage_extras.md) covers how provider-specific usage fields reach cost tracking. If you cannot find a mechanism for what you need, say so in the PR or ask in [#pr-review on Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA) before you build one **Bill a new provider usage field from the usage object.** Leave the field on the usage object as an extra and read it in that provider's `cost_calculator.py`. A provider without one gets a new `cost_calculator.py` and an arm for it in `cost_per_token` (`litellm/cost_calculator.py`), which picks the calculator by provider. If the provider's transform builds `Usage` from named fields, as the Anthropic, Gemini, and Bedrock Converse ones do, pass the new field there too. Streamed requests rebuild usage in `ChunkProcessor.calculate_usage` (`litellm/litellm_core_utils/streaming_chunk_builder_utils.py`), which does not copy extras yet the way the Responses-to-chat bridge does (minus a short exclusion list), so until it does, bill the field on streams by naming it there, as Anthropic's `inference_geo` is. A provider that reports its own charge for the request, as OpenRouter does in `usage.cost`, is billed at that charge from `_hidden_params` before any calculator runs, unless the deployment sets its own prices. When the field counts a built-in tool such as web search, add the provider to `get_cost_for_web_search_request` in `litellm/llms/__init__.py`, and add a gate to `litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py` only when no existing gate fires. Then show the field's cost in your proof on every endpoint it reaches, streaming included diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index a1b4804ae29..a31e926256a 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -10,6 +10,7 @@ from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast +from pydantic import TypeAdapter, ValidationError from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, @@ -50,6 +51,7 @@ def _error_names_resource(error: Exception, resource_id: str) -> bool: CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48) GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) +_FILE_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( @@ -550,6 +552,10 @@ class CheckBatchCost: llm_provider: str, bare_model_name: str, ) -> Optional[Tuple[str, str]]: + stored_model_id: Final = self._get_stored_model_id(job) + if stored_model_id is not None and self._batch_deployment_exists(stored_model_id): + return stored_model_id, job.unified_object_id + deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider) if deployment_id is None: verbose_proxy_logger.info( @@ -662,23 +668,36 @@ class CheckBatchCost: ) @staticmethod - def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]: - from litellm.types.utils import LiteLLMBatch - + def _parse_file_object(job: "ManagedObjectRow") -> Optional[dict[str, object]]: file_object: Final = job.file_object try: - parsed_file_object: Final[object] = ( + return _FILE_OBJECT_ADAPTER.validate_python( json.loads(file_object) if isinstance(file_object, str) else file_object ) - except (json.JSONDecodeError, ValueError): + except (json.JSONDecodeError, ValueError, ValidationError): return None - if not isinstance(parsed_file_object, dict): + + @classmethod + def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]: + from litellm.types.utils import LiteLLMBatch + + parsed_file_object: Final = cls._parse_file_object(job) + if parsed_file_object is None: return None try: return LiteLLMBatch.model_validate(parsed_file_object).input_file_id except Exception: return None + @classmethod + def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]: + """The deployment that created the batch, written into file_object by the managed + files hook at create time. Older rows have no such key and fall back to matching + the bare model name against the router.""" + parsed_file_object: Final = cls._parse_file_object(job) + model_id: Final = parsed_file_object.get("model_id") if parsed_file_object else None + return model_id if isinstance(model_id, str) and model_id else None + async def _track_completed_batch_cost( self, job: "ManagedObjectRow", diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5b54e7d72a0..0d46efbbb08 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -21,6 +21,7 @@ from typing import ( TypedDict, Union, cast, + runtime_checkable, ) from uuid import NAMESPACE_URL, uuid5 @@ -374,6 +375,12 @@ class _ManagedResourcePrismaClient(Protocol): def db(self) -> _ManagedResourceDatabase: ... +@runtime_checkable +class _WriterRoutedDatabase(Protocol): + @property + def writer(self) -> _ManagedResourceDatabase: ... + + class _SchedulerWithJobLookup(Protocol): def get_job(self, job_id: str) -> object: ... @@ -479,6 +486,11 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag return prisma_client.db.litellm_managedobjecttable +def _writer_database(db: _ManagedResourceDatabase) -> _ManagedResourceDatabase: + """Reads that feed a write must see the create that just happened; a lagging read replica may not""" + return db.writer if isinstance(db, _WriterRoutedDatabase) else db + + _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00" @@ -512,6 +524,25 @@ def _without_model_output( ) +def _with_model_id( + file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse, + model_id: str | None, +) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse: + return file_object if model_id is None else file_object.model_copy(update={"model_id": model_id}) + + +async def _stored_model_id(managed_object_table: _ManagedObjectTableActions, unified_object_id: str) -> str | None: + row: Final = await managed_object_table.find_first(where={"unified_object_id": unified_object_id}) + if row is None: + return None + try: + parsed: Final = _JSON_VALUE_ADAPTER.validate_python(_decode_json_blob(row.file_object)) + except (json.JSONDecodeError, ValidationError): + return None + model_id: Final = parsed.get("model_id") if isinstance(parsed, dict) else None + return model_id if isinstance(model_id, str) and model_id else None + + def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]: hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {} @@ -697,9 +728,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): persist_attribution: bool = False, create_if_missing: bool = True, batch_processed: bool = False, + creating_model_id: str | None = None, ) -> None: """Persist a managed object row, caching it and upserting it in the DB. + creating_model_id is the router deployment that created the object. It is kept + inside the stored file_object JSON as `model_id`, because the raw provider id an + unmanaged batch stores as unified_object_id says nothing about which deployment + made it, and the cost poller needs that deployment to read the result back. A + later observe of the row (create_if_missing=False) keeps the stored value rather + than overwriting it with whichever deployment served the retrieve. + persist_attribution is set only by the batch and background response creates, the callers that can speak for the creator; it gates the api_key, request_tags, and, for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or @@ -720,7 +759,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): It is written only in the upsert create branch. """ verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") - stored_object: Final = _without_model_output(file_object) + managed_database: Final = _managed_database(self.prisma_client) + managed_object_table: Final = managed_database.litellm_managedobjecttable + preserved_model_id: Final = ( + creating_model_id + if creating_model_id is not None or create_if_missing + else await _stored_model_id(_writer_database(managed_database).litellm_managedobjecttable, unified_object_id) + ) + stored_object: Final = _with_model_id(_without_model_output(file_object), preserved_model_id) litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, @@ -752,7 +798,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "status": file_object.status, "updated_by": user_api_key_dict.user_id, } - managed_object_table: Final = _managed_object_table(self.prisma_client) if not create_if_missing: await managed_object_table.update_many( where={"unified_object_id": unified_object_id}, @@ -2170,6 +2215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}), persist_attribution=is_batch_create, create_if_missing=is_batch_create, + creating_model_id=model_id if is_batch_create else None, ) if not is_batch_create: await self.store_batch_output_file_ownership( diff --git a/litellm-rust/.agents/skills/rust-redaction/SKILL.md b/litellm-rust/.agents/skills/rust-redaction/SKILL.md new file mode 100644 index 00000000000..2e906e10607 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-redaction/SKILL.md @@ -0,0 +1,43 @@ +--- +name: rust-redaction +description: Keep secrets out of Debug output in litellm-rust with SecretValue and veil. Use when a struct or enum holds a password, token, key, credential URL or credential provider, when writing or reviewing a Debug impl, or when a type with secrets gets logged or traced +metadata: + internal: true +--- + +# Rust redaction + +Every secret in a Debug output is hidden the same way, so a reviewer can tell at a glance that a field is covered and a new field cannot slip through a hand-written impl + +## Where SecretValue lives + +`SecretValue` lives in `crates/secret-value` and is imported as `litellm_secret_value::SecretValue` everywhere. No crate re-exports it, so there is one import path. Depend on `litellm-secret-value` directly. Do not pull in `litellm-auth-types` or `litellm-auth` just to get it + +`crates/secret-value` holds only the type. Credential acquisition (`TokenProvider`, `CredentialPlan`, provider param tables) stays in `litellm-auth-types`, and generic clients such as `litellm-redis-client` must not depend on it + +## Rules + +- Never hand-write `impl Debug` to hide a value, and never invent markers like `""`, `""` or `.map(|_| "...")`. A hand-written impl silently prints any field added later +- A secret that crosses a crate boundary or lives in config is typed as `SecretValue`. It redacts itself, compares in constant time and only gives up the plaintext through `expose()`. A type whose only secrets are `SecretValue` fields derives plain `Debug` + ```rust + #[derive(Clone, Debug, PartialEq)] + pub struct HostEndpoint { + pub host: String, + pub port: u16, + pub username: Option, + pub password: Option, + pub tls: bool, + pub db: i64, + } + ``` +- Keep a secret as `SecretValue` until the last conversion into the third-party type (redis-rs, reqwest, an SDK) and call `expose()` there. Calling `expose()` early spreads plaintext `String`s through our own types +- A secret that has to stay a raw type, because a serde wire shape or a third-party struct owns the field, is covered by deriving `veil::Redact` on the holder and marking the field `#[redact(with = "[REDACTED]")]`. `Redact` replaces `Debug`, so do not derive both +- A URL that can carry userinfo is a secret as a whole. Do not parse it to show the host. Keep it as a `SecretValue` +- A bag of unparsed settings that may hold credentials (`additional_fields`, raw kwargs) is redacted as a whole with `#[redact(with = "[REDACTED]")]` +- An opaque handle over `Arc` whose trait has no `Debug` gets a newtype whose `Debug` prints only the type name with `finish_non_exhaustive()`. That impl shows nothing, so it is not a redaction impl and the first rule does not apply. A type holding such a handle can derive `Debug` +- Never put a secret in a tracing field or an error message. Log the type with `?value`, or log only non-secret fields. See the rust-tracing and rust-error skills + +## Tests + +- One test per type that holds a secret, built the way production builds it, asserting the plaintext is absent from `format!("{value:?}")` and `[REDACTED]` is present +- Use a distinctive plaintext (`"credential-value"`) so the absence check cannot pass by accident diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 9e7a55eb232..5eaf876e3b2 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,5 +1,7 @@ # Rust workspace rules +For keeping secrets out of Debug output, follow [.agents/skills/rust-redaction/SKILL.md](.agents/skills/rust-redaction/SKILL.md) + For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md) For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e959f77917b..f91deaaa401 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3380,6 +3380,7 @@ dependencies = [ "litellm-auth-azure", "litellm-auth-gcp", "litellm-auth-types", + "litellm-secret-value", ] [[package]] @@ -3416,6 +3417,7 @@ dependencies = [ "azure_identity", "litellm-auth-types", "litellm-error", + "litellm-secret-value", "moka", "rstest", "serde_json", @@ -3435,6 +3437,7 @@ dependencies = [ "http 1.4.2", "litellm-auth-types", "litellm-error", + "litellm-secret-value", "moka", "rstest", "serde", @@ -3450,10 +3453,10 @@ name = "litellm-auth-types" version = "0.1.0" dependencies = [ "litellm-error", + "litellm-secret-value", "rstest", "serde", "serde_json", - "subtle", "thiserror 2.0.19", "tokio", "veil", @@ -3523,6 +3526,7 @@ dependencies = [ "litellm-cache-testing", "litellm-error", "litellm-http", + "litellm-secret-value", "percent-encoding", "reqwest", "rstest", @@ -3668,10 +3672,10 @@ dependencies = [ name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ - "litellm-auth", "litellm-error", "litellm-host", "litellm-host-python", + "litellm-secret-value", "proptest", "pyo3", "rstest", @@ -3692,9 +3696,9 @@ dependencies = [ name = "litellm-config" version = "0.1.0" dependencies = [ - "litellm-auth-types", "litellm-error", "litellm-router-types", + "litellm-secret-value", "rstest", "serde", "serde_json", @@ -3804,7 +3808,6 @@ dependencies = [ "envy", "futures-util", "http-body-util", - "litellm-auth-types", "litellm-circuit-breaker", "litellm-config", "litellm-error", @@ -3817,6 +3820,7 @@ dependencies = [ "litellm-llms", "litellm-redis-client", "litellm-redis-testing", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", @@ -3839,9 +3843,9 @@ dependencies = [ "axum", "axum-login", "futures-util", - "litellm-auth-types", "litellm-config", "litellm-error", + "litellm-secret-value", "litellm-secrets", "rstest", "serde", @@ -3879,6 +3883,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-router", + "litellm-secret-value", "litellm-secrets", "rstest", "serde", @@ -3895,9 +3900,9 @@ name = "litellm-gateway-management" version = "0.1.0" dependencies = [ "base64 0.22.1", - "litellm-auth-types", "litellm-error", "litellm-gateway-auth", + "litellm-secret-value", "rand 0.8.7", "rstest", "thiserror 2.0.19", @@ -3912,9 +3917,9 @@ dependencies = [ "base64 0.22.1", "futures-util", "http 1.4.2", - "litellm-auth-types", "litellm-config", "litellm-error", + "litellm-secret-value", "litellm-secrets", "moka", "rmcp", @@ -3942,9 +3947,9 @@ dependencies = [ "futures-util", "governor", "jsonwebtoken", - "litellm-auth-types", "litellm-error", "litellm-gateway-auth", + "litellm-secret-value", "rand 0.8.7", "rstest", "serde", @@ -3964,8 +3969,8 @@ name = "litellm-host" version = "0.1.0" dependencies = [ "futures-util", - "litellm-auth", "litellm-coroutine", + "litellm-secret-value", "rstest", "serde_json", "thiserror 2.0.19", @@ -4094,6 +4099,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", @@ -4124,6 +4130,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "reqwest", @@ -4154,6 +4161,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "mime_guess", @@ -4188,6 +4196,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "reqwest", @@ -4207,6 +4216,7 @@ dependencies = [ "litellm-auth", "litellm-http", "litellm-inference", + "litellm-secret-value", "litellm-secrets", "tokio", ] @@ -4254,6 +4264,7 @@ dependencies = [ "litellm-llms-types", "litellm-python-compat", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "reqwest", "rstest", @@ -4347,6 +4358,7 @@ dependencies = [ "litellm-llms-types", "litellm-model-catalog", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "litellm-secrets-aws", "litellm-secrets-types", @@ -4396,6 +4408,7 @@ dependencies = [ "futures-util", "litellm-circuit-breaker", "litellm-redis-testing", + "litellm-secret-value", "redis", "rstest", "thiserror 2.0.19", @@ -4432,12 +4445,23 @@ name = "litellm-router-types" version = "0.1.0" dependencies = [ "litellm-auth-types", + "litellm-secret-value", "rstest", "serde", "serde_json", "serde_with", ] +[[package]] +name = "litellm-secret-value" +version = "0.1.0" +dependencies = [ + "rstest", + "serde", + "subtle", + "veil", +] + [[package]] name = "litellm-secrets" version = "0.1.0" @@ -4451,6 +4475,7 @@ dependencies = [ "litellm-error", "litellm-http", "litellm-python-compat", + "litellm-secret-value", "litellm-secrets-aws", "litellm-secrets-azure", "litellm-secrets-cyberark", @@ -4479,6 +4504,7 @@ dependencies = [ "base64 0.22.1", "litellm-auth-aws", "litellm-core-utils", + "litellm-secret-value", "litellm-secrets-types", "litellm-tracing", "rstest", @@ -4498,6 +4524,7 @@ dependencies = [ "litellm-auth-types", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "percent-encoding", "reqwest", @@ -4519,6 +4546,7 @@ dependencies = [ "base64 0.22.1", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "litellm-tracing", "moka", @@ -4549,6 +4577,7 @@ dependencies = [ "litellm-auth-types", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "moka", "percent-encoding", @@ -4570,6 +4599,7 @@ dependencies = [ "litellm-core-utils", "litellm-error", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "reqwest", "rstest", @@ -4589,7 +4619,7 @@ dependencies = [ name = "litellm-secrets-types" version = "0.1.0" dependencies = [ - "litellm-auth-types", + "litellm-secret-value", "moka", "rstest", "serde", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 49a33e08456..75cac2a5d63 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -37,6 +37,7 @@ litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framer = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } litellm-auth-types = { path = "crates/auth-types" } +litellm-secret-value = { path = "crates/secret-value" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index f99972f813b..7dcb0e513a7 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true moka.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 9fda67746d1..f40745ab7a2 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -13,7 +13,8 @@ use azure_identity::{ use sha2::{Digest, Sha256}; use litellm_auth_types::{Error, ErrorDetail}; -use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; +use litellm_auth_types::{InputSource, ResolvedCredential, Sourced}; +use litellm_secret_value::SecretValue; use super::credential_provider_cache::{ AzureCredentialProviderCache, AzureCredentialProviderCacheKey, @@ -498,7 +499,8 @@ mod tests { use azure_core::{Bytes, Result}; use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; - use litellm_auth_types::{InputSource, SecretValue, Sourced}; + use litellm_auth_types::{InputSource, Sourced}; + use litellm_secret_value::SecretValue; fn deployment(value: T) -> Sourced { Sourced::new(value, InputSource::Deployment) diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 25b71741342..a92630af26a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,7 +1,9 @@ use litellm_auth_types::{ - CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, - SecretValue, Sourced, TokenProviderHandle, + CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced, + TokenProviderHandle, }; +use litellm_secret_value::SecretValue; + use litellm_auth_types::{Error, ErrorDetail}; use std::future::Future; use std::pin::Pin; @@ -500,8 +502,9 @@ mod tests { use litellm_auth_types::ResolvedCredential; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, - CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, + CredentialResolver, CredentialResolverHandle, InputSource, Sourced, }; + use litellm_secret_value::SecretValue; #[derive(Debug)] struct FileResolver; diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index a042937a047..6f8290645cf 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -1,8 +1,10 @@ use std::collections::BTreeMap; use litellm_auth_types::{ - CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle, + CredentialResolverHandle, Error, InputSource, Sourced, TokenProviderHandle, }; +use litellm_secret_value::SecretValue; + use serde_json::{Map, Value}; use strum::EnumString; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index 7844b35ecc2..ad60f0db817 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true google-sdk = ["dep:google-cloud-auth", "dep:http"] [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true base64.workspace = true diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index 8ac1e0f54ba..650d5831796 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -258,7 +258,7 @@ fn auth_acquisition_error(error: gcp_auth::Error) -> Error { #[cfg(test)] mod tests { - use litellm_auth_types::SecretValue; + use litellm_secret_value::SecretValue; use super::*; diff --git a/litellm-rust/crates/auth-gcp/src/config.rs b/litellm-rust/crates/auth-gcp/src/config.rs index 5444897aaaf..133e1c26903 100644 --- a/litellm-rust/crates/auth-gcp/src/config.rs +++ b/litellm-rust/crates/auth-gcp/src/config.rs @@ -4,7 +4,9 @@ use base64::{ Engine, alphabet, engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig}, }; -use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced}; +use litellm_auth_types::{Error, ErrorDetail, InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use serde_json::{Map, Value}; use litellm_auth_types::VertexParams; diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml index 400f2640221..f607f0d1da0 100644 --- a/litellm-rust/crates/auth-types/Cargo.toml +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -7,9 +7,9 @@ repository.workspace = true [dependencies] litellm-error.workspace = true +litellm-secret-value.workspace = true serde.workspace = true serde_json.workspace = true -subtle.workspace = true thiserror.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/auth-types/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs index 4f69c6fffc8..36e21f7794a 100644 --- a/litellm-rust/crates/auth-types/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -5,7 +5,8 @@ use std::sync::Arc; use veil::Redact; -use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle}; +use crate::{Error, ResolvedCredential, TokenProviderHandle}; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { @@ -104,8 +105,8 @@ mod tests { CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; - use crate::SecretValue; use crate::{Error, ErrorDetail}; + use litellm_secret_value::SecretValue; #[derive(Debug)] struct HostResolver; diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index 99651977867..ab3bd3febbf 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -5,7 +5,6 @@ mod error; pub mod http; mod params; mod policy; -mod secret; mod token; use serde::{Deserialize, Serialize}; @@ -55,5 +54,4 @@ pub use error::{Error, ErrorDetail}; pub use http::CredentialPlacement; pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; -pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth-types/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs index 8fff59e284a..2232350965a 100644 --- a/litellm-rust/crates/auth-types/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -71,7 +71,8 @@ impl ProviderAuthPolicy { #[cfg(test)] mod tests { use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; - use crate::{CredentialPlacement, ResolvedCredential, SecretValue}; + use crate::{CredentialPlacement, ResolvedCredential}; + use litellm_secret_value::SecretValue; const RULES: &[CredentialRule] = &[CredentialRule { kind: CredentialPlanKind::Static, diff --git a/litellm-rust/crates/auth-types/src/token.rs b/litellm-rust/crates/auth-types/src/token.rs index a3126052065..b948dda44fc 100644 --- a/litellm-rust/crates/auth-types/src/token.rs +++ b/litellm-rust/crates/auth-types/src/token.rs @@ -5,7 +5,8 @@ use std::time::SystemTime; use veil::Redact; -use crate::{Error, SecretValue}; +use crate::Error; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, PartialEq, Eq)] pub enum ResolvedCredential { diff --git a/litellm-rust/crates/auth-types/tests/token.rs b/litellm-rust/crates/auth-types/tests/token.rs index 7f7027942b0..39908ae7b8a 100644 --- a/litellm-rust/crates/auth-types/tests/token.rs +++ b/litellm-rust/crates/auth-types/tests/token.rs @@ -8,9 +8,9 @@ use std::{ time::{Duration, SystemTime}, }; -use litellm_auth_types::{ - Error, ErrorDetail, ResolvedCredential, SecretValue, TokenProviderHandle, -}; +use litellm_auth_types::{Error, ErrorDetail, ResolvedCredential, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use rstest::rstest; fn credential(index: usize, access_token: bool) -> ResolvedCredential { diff --git a/litellm-rust/crates/auth/Cargo.toml b/litellm-rust/crates/auth/Cargo.toml index ee4900ebcc2..674ae65a59f 100644 --- a/litellm-rust/crates/auth/Cargo.toml +++ b/litellm-rust/crates/auth/Cargo.toml @@ -16,3 +16,6 @@ litellm-auth-types.workspace = true litellm-auth-aws = { workspace = true, optional = true } litellm-auth-azure = { workspace = true, optional = true } litellm-auth-gcp = { workspace = true, optional = true } + +[dev-dependencies] +litellm-secret-value.workspace = true diff --git a/litellm-rust/crates/auth/tests/facade.rs b/litellm-rust/crates/auth/tests/facade.rs index f1092b15def..5d29bac1072 100644 --- a/litellm-rust/crates/auth/tests/facade.rs +++ b/litellm-rust/crates/auth/tests/facade.rs @@ -1,7 +1,8 @@ use litellm_auth::{ CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, - ProviderAuthPolicy, ResolvedCredential, SecretValue, + ProviderAuthPolicy, ResolvedCredential, }; +use litellm_secret_value::SecretValue; const RULES: &[CredentialRule] = &[CredentialRule { kind: CredentialPlanKind::Static, diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index dbd8c9ac908..878db4b8010 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true futures-util.workspace = true litellm-auth-gcp.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index 079ec44dbce..852d8bbf035 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -1,7 +1,9 @@ use std::{future::Future, pin::Pin}; use litellm_auth_gcp::{VertexAuth, VertexConfig}; -use litellm_auth_types::{InputSource, SecretValue, Sourced}; +use litellm_auth_types::{InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use litellm_cache::Error; pub trait TokenSource: Send + Sync + 'static { diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index 1cd05710366..0133af16b80 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -16,7 +16,7 @@ serde_json.workspace = true thiserror.workspace = true [dev-dependencies] -litellm-auth.workspace = true +litellm-secret-value.workspace = true proptest.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 493e91bb5ea..9431b9d1326 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -925,11 +925,11 @@ assert all(value is failure for name, value in logger.calls if name.endswith('_h mod payload_tests { use std::ffi::CStr; - use litellm_auth::SecretValue; use litellm_host::hooks::CallHooks; use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest}; use litellm_host::lifecycle::ExecutionEvent; use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py}; + use litellm_secret_value::SecretValue; use proptest::prelude::*; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 5b0cd22332f..278aea6200e 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth-types.workspace = true +litellm-secret-value.workspace = true litellm-router-types.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/config/src/mcp.rs b/litellm-rust/crates/config/src/mcp.rs index cbe992f9956..00480fa5be1 100644 --- a/litellm-rust/crates/config/src/mcp.rs +++ b/litellm-rust/crates/config/src/mcp.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, fmt}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use crate::Object; diff --git a/litellm-rust/crates/config/src/redis.rs b/litellm-rust/crates/config/src/redis.rs index 5674f993025..7a91691c54b 100644 --- a/litellm-rust/crates/config/src/redis.rs +++ b/litellm-rust/crates/config/src/redis.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, fmt, time::Duration}; -use litellm_auth_types::SecretValue; use litellm_router_types::Spelled; +use litellm_secret_value::SecretValue; use serde::{Deserialize, Deserializer, de::Error as _}; use crate::{AdditionalFields, Error, Object, Value}; diff --git a/litellm-rust/crates/config/src/settings.rs b/litellm-rust/crates/config/src/settings.rs index 9571c4fa5de..1bfc19983f3 100644 --- a/litellm-rust/crates/config/src/settings.rs +++ b/litellm-rust/crates/config/src/settings.rs @@ -1,6 +1,6 @@ use std::fmt; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use crate::{ diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml index 27f20203843..ba9097e85d8 100644 --- a/litellm-rust/crates/gateway-auth/Cargo.toml +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -6,9 +6,9 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum = { workspace = true, features = ["matched-path"] } axum-login.workspace = true -litellm-auth-types.workspace = true litellm-config.workspace = true litellm-error.workspace = true litellm-secrets.workspace = true diff --git a/litellm-rust/crates/gateway-auth/src/authentication.rs b/litellm-rust/crates/gateway-auth/src/authentication.rs index 4ecfaa375b4..349913619f5 100644 --- a/litellm-rust/crates/gateway-auth/src/authentication.rs +++ b/litellm-rust/crates/gateway-auth/src/authentication.rs @@ -1,8 +1,8 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use axum::http::{Request, request::Parts}; -use litellm_auth_types::SecretValue; use litellm_config::Config; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/gateway-auth/src/http.rs b/litellm-rust/crates/gateway-auth/src/http.rs index ddbf240509a..ad1718161c0 100644 --- a/litellm-rust/crates/gateway-auth/src/http.rs +++ b/litellm-rust/crates/gateway-auth/src/http.rs @@ -4,7 +4,7 @@ use axum::{ middleware::Next, response::Response, }; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{AccessRequest, Auth, AuthenticatedRequest, Credential, Error}; diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 1fe8fa60e0e..0deaf3f7a82 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -1,6 +1,6 @@ use std::{future::Future, pin::Pin, time::SystemTime}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{Error, hash_token}; diff --git a/litellm-rust/crates/gateway-auth/src/ui.rs b/litellm-rust/crates/gateway-auth/src/ui.rs index 3d655bde187..ee4dd670aa7 100644 --- a/litellm-rust/crates/gateway-auth/src/ui.rs +++ b/litellm-rust/crates/gateway-auth/src/ui.rs @@ -3,7 +3,7 @@ use std::{convert::Infallible, sync::Arc}; use axum::extract::FromRequestParts; use axum::http::{header::AUTHORIZATION, request::Parts}; use axum_login::{AuthUser, AuthnBackend, UserId}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/gateway-auth/tests/auth.rs b/litellm-rust/crates/gateway-auth/tests/auth.rs index 781ba385381..54180e30504 100644 --- a/litellm-rust/crates/gateway-auth/tests/auth.rs +++ b/litellm-rust/crates/gateway-auth/tests/auth.rs @@ -8,9 +8,9 @@ use axum::{ routing::get, }; use futures_util::future::BoxFuture; -use litellm_auth_types::SecretValue; use litellm_config::Config; use litellm_gateway_auth::{Auth, RequireMasterKey, hash_token}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use rstest::{fixture, rstest}; use tower::ServiceExt; diff --git a/litellm-rust/crates/gateway-auth/tests/keys.rs b/litellm-rust/crates/gateway-auth/tests/keys.rs index 4d16d3573bc..1f6fc461a5b 100644 --- a/litellm-rust/crates/gateway-auth/tests/keys.rs +++ b/litellm-rust/crates/gateway-auth/tests/keys.rs @@ -1,10 +1,10 @@ use std::time::{Duration, SystemTime}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ Error, keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key}, }; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; struct Lookup { diff --git a/litellm-rust/crates/gateway-auth/tests/ui.rs b/litellm-rust/crates/gateway-auth/tests/ui.rs index 2e5f72d605d..8d465e1198d 100644 --- a/litellm-rust/crates/gateway-auth/tests/ui.rs +++ b/litellm-rust/crates/gateway-auth/tests/ui.rs @@ -1,6 +1,6 @@ use axum_login::{AuthUser, AuthnBackend}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{UiBackend, UiCredentials}; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; #[fixture] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 034c3a9826e..69ca561e4b5 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-cache-response.workspace = true axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs index f200fd7b141..b8e6ef4ff4f 100644 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -2,10 +2,10 @@ use litellm_gateway_auth::AuthenticatedRequest; use std::sync::Arc; use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse}; -use litellm_auth::SecretValue; use litellm_inference_ocr::types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput}; use litellm_llms::base_llm::ocr::transformation::decode_request_value; use litellm_llms_types::formats::ocr::OcrDocument; +use litellm_secret_value::SecretValue; use serde_json::Value; use crate::{ diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 737dc0c830c..185cbc8594d 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -13,7 +13,9 @@ use futures_util::future::BoxFuture; use litellm_gateway_inference::{Deployment, Gateway, router}; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; use litellm_inference::resources::CoreResources; -use litellm_secrets::{SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::source::SecretSource; + use serde_json::Value; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; diff --git a/litellm-rust/crates/gateway-management/Cargo.toml b/litellm-rust/crates/gateway-management/Cargo.toml index 502760f8fbe..1ff5dd5319d 100644 --- a/litellm-rust/crates/gateway-management/Cargo.toml +++ b/litellm-rust/crates/gateway-management/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true base64.workspace = true -litellm-auth-types.workspace = true litellm-error.workspace = true litellm-gateway-auth.workspace = true rand.workspace = true diff --git a/litellm-rust/crates/gateway-management/src/keys.rs b/litellm-rust/crates/gateway-management/src/keys.rs index 45d915b7b3d..feeb40b2403 100644 --- a/litellm-rust/crates/gateway-management/src/keys.rs +++ b/litellm-rust/crates/gateway-management/src/keys.rs @@ -1,8 +1,8 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::keys::{KeyHash, KeyStatus}; +use litellm_secret_value::SecretValue; use rand::{RngCore, rngs::OsRng}; use crate::Error; diff --git a/litellm-rust/crates/gateway-management/tests/keys.rs b/litellm-rust/crates/gateway-management/tests/keys.rs index bf0221cf1be..8a970bb407c 100644 --- a/litellm-rust/crates/gateway-management/tests/keys.rs +++ b/litellm-rust/crates/gateway-management/tests/keys.rs @@ -7,7 +7,6 @@ use std::{ }; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ hash_token, keys::{KeyHash, KeyLookup, KeyStatus, verify_key}, @@ -16,6 +15,7 @@ use litellm_gateway_management::{ Error, keys::{Keys, Revocation}, }; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; use support::{MemoryStore, Unavailable}; diff --git a/litellm-rust/crates/gateway-mcp/Cargo.toml b/litellm-rust/crates/gateway-mcp/Cargo.toml index bed48bad58d..65077eec0ea 100644 --- a/litellm-rust/crates/gateway-mcp/Cargo.toml +++ b/litellm-rust/crates/gateway-mcp/Cargo.toml @@ -7,11 +7,11 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum = { workspace = true, features = ["json", "query"] } futures-util.workspace = true litellm-config.workspace = true litellm-secrets.workspace = true -litellm-auth-types.workspace = true litellm-error.workspace = true base64.workspace = true http.workspace = true diff --git a/litellm-rust/crates/gateway-mcp/src/configured.rs b/litellm-rust/crates/gateway-mcp/src/configured.rs index cc7a52fc029..33cbd243530 100644 --- a/litellm-rust/crates/gateway-mcp/src/configured.rs +++ b/litellm-rust/crates/gateway-mcp/src/configured.rs @@ -8,8 +8,8 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use futures_util::future::try_join_all; use http::{HeaderName, HeaderValue}; -use litellm_auth_types::SecretValue; use litellm_config::{McpAuth, McpServer, McpTransport}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use rmcp::{ RoleClient, ServiceExt, diff --git a/litellm-rust/crates/gateway-ui/Cargo.toml b/litellm-rust/crates/gateway-ui/Cargo.toml index 1c65665be72..a6104455e96 100644 --- a/litellm-rust/crates/gateway-ui/Cargo.toml +++ b/litellm-rust/crates/gateway-ui/Cargo.toml @@ -23,10 +23,10 @@ tower-http = { version = "0.6.11", features = ["fs", "set-header"] } tower-sessions.workspace = true [dev-dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true serde_json.workspace = true tower = { version = "0.5", features = ["util"] } -litellm-auth-types.workspace = true rstest.workspace = true tempfile.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/gateway-ui/tests/routes.rs b/litellm-rust/crates/gateway-ui/tests/routes.rs index 95246f33a35..86b7db63ed9 100644 --- a/litellm-rust/crates/gateway-ui/tests/routes.rs +++ b/litellm-rust/crates/gateway-ui/tests/routes.rs @@ -16,8 +16,8 @@ use tower_sessions::{ session::{Id, Record}, }; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::UiBackend; +use litellm_secret_value::SecretValue; struct App { router: Router, diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 4d831aa2314..6eeaa43e013 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum.workspace = true envy = "0.4.2" http-body-util = "0.1" @@ -17,7 +18,6 @@ thiserror.workspace = true futures-util.workspace = true litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true -litellm-auth-types.workspace = true litellm-config.workspace = true litellm-error.workspace = true litellm-circuit-breaker.workspace = true diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis.rs index 197f40e7988..c40b289029b 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis.rs @@ -152,12 +152,12 @@ async fn reachable(client: RedisClient) -> Option { fn client(params: Params, breaker: Option) -> Result { tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured"); let endpoint = match params.endpoint { - RedisEndpoint::Url(url) => Endpoint::Url(url.expose().to_owned()), + RedisEndpoint::Url(url) => Endpoint::Url(url), RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint { host: host.host, port: host.port, - username: host.username.map(|value| value.expose().to_owned()), - password: host.password.map(|value| value.expose().to_owned()), + username: host.username, + password: host.password, tls: host.tls, db: host.db, }), @@ -170,8 +170,8 @@ fn client(params: Params, breaker: Option) -> Result futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async { Ok(None) }) } diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index c876c1cbcbb..7de7a97c870 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -159,7 +159,7 @@ async fn mounts_ui_without_exposing_credentials_or_authorizing_inference(inferen std::fs::write(assets.path().join("index.html"), "dashboard").unwrap(); let backend = litellm_gateway_auth::UiBackend::new( "admin".into(), - litellm_auth_types::SecretValue::new("ui-password"), + litellm_secret_value::SecretValue::new("ui-password"), ) .unwrap(); let ui = litellm_gateway_ui::router( diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index d3354fa8622..e0ccab4e0ad 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true -litellm-auth.workspace = true litellm-coroutine.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index 3cff781792f..39af7c410e0 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -22,7 +22,7 @@ pub struct RequestContext { /// Optional-param names that carry credentials and must be redacted when logged. pub secret_fields: Vec, /// The credential the route resolved for the provider call. - pub api_key: Option, + pub api_key: Option, } #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index 8a779d44976..bcde3c8d30f 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-cache-response.workspace = true litellm-core-utils.workspace = true diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index 1c724a47956..f948e8fffaf 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -58,7 +58,7 @@ pub(super) async fn execute( custom_llm_provider: provider_name.to_owned(), optional_params: Value::Object(call.optional_params), secret_fields: Vec::new(), - api_key: call.api_key.map(litellm_auth::SecretValue::new), + api_key: call.api_key.map(litellm_secret_value::SecretValue::new), }; let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?; let identity = litellm_host::interceptors::ProviderIdentity { diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 4b3e32f04bc..dd229d3d657 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -140,7 +140,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -154,7 +154,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index 694d9b4de31..ac52cb65a0d 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true bytes.workspace = true futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 9d0bf09197e..2e647688a1a 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -2,7 +2,9 @@ use std::time::Duration; use bytes::Bytes; use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream}; -use litellm_auth::{AuthServices, SecretValue}; +use litellm_auth::AuthServices; +use litellm_secret_value::SecretValue; + use litellm_core_utils::{ dot_notation_indexing::delete_nested_value, get_provider_specific_headers::get_provider_specific_headers, settings::Lookup, diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 480b3c4f3b7..663c87453a9 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -107,7 +107,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -121,7 +121,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-ocr/Cargo.toml b/litellm-rust/crates/inference-ocr/Cargo.toml index bbbdaddb050..ad53126b9fe 100644 --- a/litellm-rust/crates/inference-ocr/Cargo.toml +++ b/litellm-rust/crates/inference-ocr/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true base64.workspace = true bytes.workspace = true futures-util.workspace = true diff --git a/litellm-rust/crates/inference-ocr/src/handler.rs b/litellm-rust/crates/inference-ocr/src/handler.rs index 426d4297f28..c126a192153 100644 --- a/litellm-rust/crates/inference-ocr/src/handler.rs +++ b/litellm-rust/crates/inference-ocr/src/handler.rs @@ -128,7 +128,7 @@ mod tests { decode_request(OcrWireRequest { model: model.into(), document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_key: Some(litellm_secret_value::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/inference-ocr/src/provider_config.rs b/litellm-rust/crates/inference-ocr/src/provider_config.rs index 19e50af0dd1..e76d291254b 100644 --- a/litellm-rust/crates/inference-ocr/src/provider_config.rs +++ b/litellm-rust/crates/inference-ocr/src/provider_config.rs @@ -1,4 +1,6 @@ -use litellm_auth::{InputSource, SecretValue, Sourced}; +use litellm_auth::{InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::get_llm_provider_logic::{ CustomLlmProvider, LlmProviders, get_custom_llm_provider, }; @@ -407,7 +409,7 @@ mod tests { fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("explicit-key"), + litellm_secret_value::SecretValue::new("explicit-key"), InputSource::Deployment, )), api_base: Some(Sourced::new( @@ -415,7 +417,7 @@ mod tests { InputSource::Deployment, )), dynamic_api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("dynamic-key"), + litellm_secret_value::SecretValue::new("dynamic-key"), InputSource::Environment, )), dynamic_api_base: Some(Sourced::new( @@ -455,7 +457,7 @@ mod tests { ) { let dynamic_key = dynamic_value.map(|value| { Sourced::new( - litellm_auth::SecretValue::new(value), + litellm_secret_value::SecretValue::new(value), InputSource::Environment, ) }); @@ -463,7 +465,7 @@ mod tests { dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("explicit-key"), + litellm_secret_value::SecretValue::new("explicit-key"), InputSource::Deployment, )), api_base: Some(Sourced::new( @@ -502,14 +504,14 @@ mod tests { OcrCredentialInputs { api_key: explicit_key.map(|value| { Sourced::new( - litellm_auth::SecretValue::new(value), + litellm_secret_value::SecretValue::new(value), InputSource::Deployment, ) }), api_base: explicit_base .map(|value| Sourced::new(value.into(), InputSource::Deployment)), dynamic_api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("dynamic-key"), + litellm_secret_value::SecretValue::new("dynamic-key"), InputSource::Environment, )), dynamic_api_base: Some(Sourced::new( diff --git a/litellm-rust/crates/inference-ocr/src/types.rs b/litellm-rust/crates/inference-ocr/src/types.rs index 98c609a158f..1b2dbb52d77 100644 --- a/litellm-rust/crates/inference-ocr/src/types.rs +++ b/litellm-rust/crates/inference-ocr/src/types.rs @@ -1,7 +1,9 @@ use std::{collections::BTreeMap, path::PathBuf, time::Duration}; use bytes::Bytes; -use litellm_auth::{InputSource, SecretValue, TokenProviderHandle}; +use litellm_auth::{InputSource, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::call_arguments::CallArguments; use litellm_llms::base_llm::ocr::{ error::Error, diff --git a/litellm-rust/crates/inference-ocr/src/wire.rs b/litellm-rust/crates/inference-ocr/src/wire.rs index 56d22be6523..45a2865823f 100644 --- a/litellm-rust/crates/inference-ocr/src/wire.rs +++ b/litellm-rust/crates/inference-ocr/src/wire.rs @@ -1,6 +1,8 @@ use std::{collections::BTreeMap, time::Duration}; -use litellm_auth::{InputSource, SecretValue}; +use litellm_auth::InputSource; +use litellm_secret_value::SecretValue; + use litellm_llms::base_llm::ocr::{error::Error, transformation::decode_request_value}; use litellm_llms_types::formats::ocr::OcrDocument; use serde::Deserialize; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs index 6f591af73e0..65c7f309b82 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs @@ -3,9 +3,9 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, }; -use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, -}; +use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use rstest::rstest; use super::*; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs index 75a91092daf..9f59bb9a2a2 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs @@ -239,7 +239,7 @@ impl litellm_host_native::services::HostCallHandler for CallerTokenHost { OcrOp::AcquireAzureAdToken(reply) => { self.trace.lock().unwrap().push("token".into()); reply.send(litellm_auth::ResolvedCredential::Static( - litellm_auth::SecretValue::new("caller-token"), + litellm_secret_value::SecretValue::new("caller-token"), )); Ok(()) } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs index 5a48bba9e1a..6871d1ca1f6 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs @@ -80,7 +80,7 @@ fn wire(model: &str, base: &str, document: Value, options: Value) -> OcrWireRequ OcrWireRequest { model: model.into(), document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_key: Some(litellm_secret_value::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/inference-responses/Cargo.toml b/litellm-rust/crates/inference-responses/Cargo.toml index 083179c50c8..f90651634c5 100644 --- a/litellm-rust/crates/inference-responses/Cargo.toml +++ b/litellm-rust/crates/inference-responses/Cargo.toml @@ -25,6 +25,7 @@ tokio-tungstenite.workspace = true tracing.workspace = true [dev-dependencies] +litellm-secret-value.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache.workspace = true litellm-cache-memory.workspace = true diff --git a/litellm-rust/crates/inference-responses/tests/caching.rs b/litellm-rust/crates/inference-responses/tests/caching.rs index 462461695c3..56e25fbcb98 100644 --- a/litellm-rust/crates/inference-responses/tests/caching.rs +++ b/litellm-rust/crates/inference-responses/tests/caching.rs @@ -98,7 +98,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -112,7 +112,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-testing/Cargo.toml b/litellm-rust/crates/inference-testing/Cargo.toml index 748b1df0d3b..5f690532d9c 100644 --- a/litellm-rust/crates/inference-testing/Cargo.toml +++ b/litellm-rust/crates/inference-testing/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true publish = false [dependencies] +litellm-secret-value.workspace = true litellm-auth.workspace = true futures-util.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/inference-testing/src/lib.rs b/litellm-rust/crates/inference-testing/src/lib.rs index 2e20d5a9e91..f330e02e99a 100644 --- a/litellm-rust/crates/inference-testing/src/lib.rs +++ b/litellm-rust/crates/inference-testing/src/lib.rs @@ -5,7 +5,8 @@ use litellm_http::{ ClientVariant, HttpClientConfig, HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver, }; -use litellm_secrets::{SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::source::SecretSource; use litellm_inference::resources::CoreResources; diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index 80095cac36c..01ed023999d 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -199,7 +199,7 @@ impl AnthropicBeta { &[BetaProvider::BedrockConverse, BetaProvider::Databricks] } Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse], - Self::AdvisorTool20260301 | Self::Compact20260904 => &[ + Self::AdvisorTool20260301 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, BetaProvider::Bedrock, @@ -207,6 +207,12 @@ impl AnthropicBeta { BetaProvider::VertexAi, BetaProvider::Databricks, ], + Self::Compact20260904 => &[ + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::BedrockMantle, + BetaProvider::Databricks, + ], Self::Bash20241022 | Self::Bash20250124 | Self::McpServers20251204 @@ -256,14 +262,18 @@ impl AnthropicBeta { Self::InterleavedThinking20250514 | Self::WebSearch20250305 => { &[BetaProvider::BedrockConverse, BetaProvider::Bedrock] } - Self::MidConversationOutputConfig20260701 - | Self::MidConversationToolChanges20260701 => &[ + Self::MidConversationOutputConfig20260701 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, BetaProvider::BedrockMantle, BetaProvider::VertexAi, BetaProvider::Databricks, ], + Self::MidConversationToolChanges20260701 => &[ + BetaProvider::BedrockConverse, + BetaProvider::BedrockMantle, + BetaProvider::Databricks, + ], Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, @@ -483,7 +493,7 @@ mod tests { let config = beta_headers_config(); let config_columns: Vec = config .keys() - .filter(|column| column.as_str() != "description") + .filter(|column| column.as_str() != "description" && !column.ends_with("_header_only")) .cloned() .collect(); let beta_provider_columns: Vec = BetaProvider::VARIANTS @@ -502,12 +512,16 @@ mod tests { let config = beta_headers_config(); for provider in BetaProvider::VARIANTS { for beta in &AnthropicBeta::KNOWN { - let expected = config - .get(provider.as_ref()) - .and_then(serde_json::Value::as_object) - .and_then(|column| column.get(beta.as_str())) - .and_then(serde_json::Value::as_str) - .map(str::to_string); + let cell = |column: &str| { + config + .get(column) + .and_then(serde_json::Value::as_object) + .and_then(|values| values.get(beta.as_str())) + .and_then(serde_json::Value::as_str) + .map(str::to_string) + }; + let expected = + cell(&format!("{provider}_header_only")).or_else(|| cell(provider.as_ref())); let actual = beta.on(*provider).map(|name| name.to_string()); assert_eq!(actual, expected, "provider {provider}, beta {beta}"); if matches!(beta, AnthropicBeta::AdvancedToolUse20251120) diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 962f61ef4ec..79ee17e8b97 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true test-support = ["litellm-http/test-support"] [dependencies] +litellm-secret-value.workspace = true litellm-llms-types.workspace = true litellm-core-utils.workspace = true litellm-error.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index eb796a900f5..164dbce1264 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,4 +1,3 @@ -use litellm_auth::SecretValue; use litellm_core_utils::decline::Decline; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, @@ -10,6 +9,7 @@ use litellm_llms_types::{ }, providers::anthropic::DEFAULT_HEADERS, }; +use litellm_secret_value::SecretValue; use serde::Deserialize; use serde_json::{Map, Value, json}; diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 7991a57353e..aec5d9ebc27 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,5 +1,7 @@ use crate::base_llm::messages::context::MessagesModelCapabilities; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{ has_header, header_value, header_values, with_header, without_headers, diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 6657cbf262f..fb6d7259bf9 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -1,7 +1,8 @@ use litellm_auth::CredentialPlacement; use litellm_llms_types::{ formats::messages::{ - ContextEdit, ContextManagement, Message, MessagesOptionalParams, MessagesRequest, Speed, + ContentBlock, ContentBlockType, ContextEdit, ContextManagement, Message, + MessagesOptionalParams, MessagesRequest, Speed, ThinkingConfig, ThinkingDisplay, }, providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS}, recognized::Recognized, @@ -203,6 +204,10 @@ fn feature_betas(request: &MessagesRequest) -> BetaSet { .then_some(AnthropicBeta::FastMode20260201), messages_carry_output_config(&request.messages) .then_some(AnthropicBeta::PerTurnControl20260701), + uses_thinking_display_updates(params) + .then_some(AnthropicBeta::ThinkingDisplayUpdates20260818), + messages_carry_tool_changes(&request.messages) + .then_some(AnthropicBeta::MidConversationToolChanges20260701), has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301), is_tool_search_used(tools).then_some(AnthropicBeta::AdvancedToolUse20251120), ] @@ -253,6 +258,36 @@ fn messages_carry_output_config(messages: &[Message]) -> bool { .any(|message| message.extra.contains_key("output_config")) } +fn uses_thinking_display_updates(params: &MessagesOptionalParams) -> bool { + let display = match params.thinking.as_ref().and_then(Recognized::known) { + Some(ThinkingConfig::Enabled(thinking)) => thinking.display.as_ref(), + Some(ThinkingConfig::Adaptive(thinking)) => thinking.display.as_ref(), + _ => None, + }; + display.and_then(Recognized::known) == Some(&ThinkingDisplay::Updates) +} + +fn is_tool_reference_change(block: &ContentBlock) -> bool { + matches!( + &block.block_type, + Some(ContentBlockType::Other(block_type)) + if block_type == "tool_addition" || block_type == "tool_removal" + ) && block + .extra + .get("tool") + .and_then(|tool| tool.get("type")) + .and_then(Value::as_str) + == Some("tool_reference") +} + +fn messages_carry_tool_changes(messages: &[Message]) -> bool { + messages + .iter() + .filter(|message| message.role == "system") + .flat_map(Message::blocks) + .any(is_tool_reference_change) +} + fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error { Error::InvalidRequest(crate::ErrorDetail::UnsupportedParameter { model: model.into(), @@ -1071,6 +1106,41 @@ mod tests { json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}), &[] )] + #[case::adaptive_display_updates( + json!({"thinking": {"type": "adaptive", "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::enabled_display_updates( + json!({"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::summarized_display(json!({"thinking": {"type": "adaptive", "display": "summarized"}}), &[])] + #[case::disabled_thinking(json!({"thinking": {"type": "disabled"}}), &[])] + #[case::system_tool_addition( + json!({"messages": [ + {"role": "user", "content": "Use the ping tool"}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] + #[case::system_tool_removal( + json!({"messages": [ + {"role": "system", "content": [{"type": "tool_removal", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] + #[case::user_tool_addition( + json!({"messages": [ + {"role": "user", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &[] + )] + #[case::tool_addition_without_a_tool_reference( + json!({"messages": [ + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_definition"}}]}, + ]}), + &[] + )] fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { assert_eq!(feature_betas(&request(fields)), betas(expected)); } @@ -1093,6 +1163,33 @@ mod tests { ); } + #[rstest] + #[case::vertex_ai( + BetaProvider::VertexAi, + &["mid-conversation-tool-changes-2026-07-01", "thinking-display-updates-2026-08-18"] + )] + #[case::azure_ai( + BetaProvider::AzureAi, + &["mid-conversation-tool-changes-2026-07-01", "thinking-display-updates-2026-08-18"] + )] + #[case::bedrock_converse(BetaProvider::BedrockConverse, &["thinking-display-updates-2026-08-18"])] + #[case::databricks(BetaProvider::Databricks, &["thinking-display-updates-2026-08-18"])] + fn tool_change_and_display_betas_reach_the_providers_that_accept_them( + #[case] provider: BetaProvider, + #[case] expected: &[&str], + ) { + let fields = json!({ + "thinking": {"type": "adaptive", "display": "updates"}, + "messages": [ + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ], + }); + assert_eq!( + provider_feature_betas(&request(fields), provider), + betas(expected) + ); + } + #[rstest] fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() { let client = headers(&[("anthropic-beta", "claude-code-20250219")]); diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index bf934fb4d4c..a2ba2cee0d3 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -1,10 +1,10 @@ -use litellm_auth::SecretValue; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; +use litellm_secret_value::SecretValue; use crate::{ Error, @@ -545,6 +545,11 @@ mod tests { &[("anthropic-beta", "web-search-2025-03-05")], &[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")] )] + #[case::compaction_needs_the_compact_2026_09_04_beta( + json!({"compaction": {"type": "summarize"}}), + &[], + &[("x-api-key", "k"), ("anthropic-beta", "compact-2026-09-04")] + )] #[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])] #[case::safeguards_need_the_dangerous_tool_use_beta( json!({"safeguards": [{"type": "dangerous_tool_use"}]}), diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 36dfe2c890e..870306b1a83 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -1250,7 +1250,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 2f7957fe56e..25f93846967 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -211,7 +211,7 @@ mod tests { #[fixture] fn connection() -> OcrConnection { OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_base: Some("https://example.com".into()), ..Default::default() } @@ -312,7 +312,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/base_llm/auth.rs b/litellm-rust/crates/llms/src/base_llm/auth.rs index de271a269d3..ab2cf00614e 100644 --- a/litellm-rust/crates/llms/src/base_llm/auth.rs +++ b/litellm-rust/crates/llms/src/base_llm/auth.rs @@ -5,7 +5,9 @@ //! that split: `validate_environment` shapes the forwarded headers and names the credential //! as an [`AuthScheme`], and [`resolve_auth`] turns the scheme into headers and a signer. -use litellm_auth::{AuthServices, CredentialPlacement, SecretValue, TokenProviderHandle}; +use litellm_auth::{AuthServices, CredentialPlacement, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{AwsCredentialSource, SigV4Signer}; use litellm_auth_gcp::VertexConfig; use litellm_http::request::with_header; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index c90fd99809e..9574e99d357 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -5,7 +5,9 @@ use serde::{ use serde_json::{Map, Value}; use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; -use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; +use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; use litellm_http::outbound::RequestSigner; use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index 976c9161dd2..6cf0167a30f 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -1,5 +1,7 @@ use litellm_auth::AwsParams; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 79c30df9b19..5b89345a1c5 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -1,6 +1,8 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region}; use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{header_value, with_header, without_headers}; diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 3ec60f12021..1d380c917f6 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -7,7 +7,9 @@ use crate::{ }; use futures_util::StreamExt; use litellm_auth::AwsParams; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{ diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs index ae13704f4a0..f139b93e259 100644 --- a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -1,6 +1,8 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::AwsCredentialSource; use litellm_auth_aws::bedrock_model_id_and_region; use litellm_core_utils::settings::resolve_non_empty; diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index 76ff27b56fc..5bad06adda2 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -735,7 +735,7 @@ mod tests { assert!(matches!( CohereParseConfig.resolve_headers( &OcrConnection { - api_key: Some(litellm_auth::SecretValue::new(" ")), + api_key: Some(litellm_secret_value::SecretValue::new(" ")), ..Default::default() }, &|_| None, diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 168c3e63920..0dbcea01439 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool}, diff --git a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs index 2c3a7eaef1d..d5944484522 100644 --- a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::formats::messages::{ MessagesOptionalParams, MessagesRequest, MessagesResponse, diff --git a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs index 012a657fceb..dafe5fbd16b 100644 --- a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index fa4e2254e86..679aa65f7fd 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -219,7 +219,7 @@ mod tests { #[default(vec![])] extra_headers: Vec<(String, String)>, ) -> OcrConnection { OcrConnection { - api_key: api_key.map(litellm_auth::SecretValue::new), + api_key: api_key.map(litellm_secret_value::SecretValue::new), extra_headers, ..OcrConnection::default() } diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index cf74871894b..deb163daa07 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -3,7 +3,8 @@ use litellm_llms_types::formats::responses::{ }; use serde_json::{Map, Value}; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; use crate::{ Error, ErrorDetail, diff --git a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs index e998f8785cb..83fb8ac3fa2 100644 --- a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs @@ -4,7 +4,9 @@ //! the `max_completion_tokens` -> `max_tokens` rename and the usage //! `*_tokens` null-to-zero sanitize. -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::core_helpers::unix_now; use litellm_llms_types::formats::chat_completions::{ ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, diff --git a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs index 5b4e4f77f2a..eeb3d1cb7f8 100644 --- a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{has_bearer_auth, without_headers}; use litellm_llms_types::{ diff --git a/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs b/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs index 900fba9e916..76155ba55e5 100644 --- a/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::formats::responses::ResponsesApiResponse; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index bf0ddc1c7a8..beda6899d04 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -690,7 +690,7 @@ mod tests { #[test] fn explicit_key_precedes_environment_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("passed-key")), + api_key: Some(litellm_secret_value::SecretValue::new("passed-key")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); @@ -700,7 +700,7 @@ mod tests { #[test] fn blank_explicit_key_uses_environment_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new(" ")), + api_key: Some(litellm_secret_value::SecretValue::new(" ")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index a70336748ad..621a09f4fb0 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index 46e5a68dcb8..e5e3a056fab 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -798,6 +798,18 @@ mod tests { ]}}), &["compact-2026-01-12", "context-management-2025-06-27"] )] + #[case::compaction(json!({"compaction": {"type": "summarize"}}), &["compact-2026-09-04"])] + #[case::thinking_display_updates( + json!({"thinking": {"type": "adaptive", "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::tool_addition( + json!({"messages": [ + {"role": "user", "content": "Use the ping tool"}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]} + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] #[case::advisor_tool_is_not_a_vertex_beta( json!({"tools": [{"type": "advisor_20260301", "name": "advisor"}]}), &[] diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index 77881cd2e78..b902967d9dc 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -137,7 +137,7 @@ impl VertexAiOcrConfig { connection .api_key .as_ref() - .map(litellm_auth::SecretValue::expose), + .map(litellm_secret_value::SecretValue::expose), config, &|name: &str| connection.secret(name), ) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index a4e1c57ab1e..e3ef45f7876 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true litellm-spend-clickhouse.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/credentials.rs b/litellm-rust/crates/python-bridge/src/credentials.rs index a4bb064aae0..adc52a7e77a 100644 --- a/litellm-rust/crates/python-bridge/src/credentials.rs +++ b/litellm-rust/crates/python-bridge/src/credentials.rs @@ -2,7 +2,9 @@ //! keyword arguments and acquired on the host's own thread when the call asks for one. use crate::callable::wrap_failure; -use litellm_auth::{ResolvedCredential, SecretValue}; +use litellm_auth::ResolvedCredential; +use litellm_secret_value::SecretValue; + use pyo3::{ exceptions::PyTypeError, gc::{PyTraverseError, PyVisit}, diff --git a/litellm-rust/crates/python-bridge/src/db.rs b/litellm-rust/crates/python-bridge/src/db.rs new file mode 100644 index 00000000000..03fb30b986f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/db.rs @@ -0,0 +1,19 @@ +use litellm_host_python::ready_future; +use pyo3::{prelude::*, types::PyNone}; + +#[pyclass(frozen)] +pub(crate) struct NativeDatabase {} + +// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first +// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens +#[pymethods] +impl NativeDatabase { + #[staticmethod] + fn connect(py: Python<'_>) -> PyResult> { + ready_future(py, Bound::new(py, NativeDatabase {})?.as_any()) + } + + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + ready_future(py, PyNone::get(py).as_any()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 683e90d02a2..fdcb969a423 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -2,6 +2,7 @@ mod cache; mod callable; mod coercion; mod credentials; +mod db; mod diagnostics; mod errors; mod execution; @@ -21,6 +22,8 @@ mod _native { use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; use pyo3::{prelude::*, types::PyModule}; + #[pymodule_export] + use crate::db::NativeDatabase; #[cfg(feature = "panic-test")] #[pymodule_export] use crate::diagnostics::_panic_for_test; @@ -112,6 +115,7 @@ mod tests { "NativeDiagnosticProcessor", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", + "NativeDatabase", "trace_encode_error", "bundled_model_cost_map", "TokenCounter", diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index c2c3fe51600..73a2922804a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,4 +1,3 @@ -use litellm_auth::SecretValue; use litellm_auth::VertexParams; use litellm_host_python::{from_py, present}; use litellm_inference_ocr::{ @@ -6,6 +5,7 @@ use litellm_inference_ocr::{ wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, }; use litellm_llms::base_llm::ocr::{error::Error, transformation::DROP_PARAMS_PARAM}; +use litellm_secret_value::SecretValue; use pyo3::{prelude::*, types::PyDict}; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 669c10e602a..af6485a56e1 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -2,9 +2,11 @@ use std::{future::Future, pin::Pin, sync::Arc}; use litellm_core_utils::settings::Lookup; use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_secret_value::SecretValue; use litellm_secrets::{ - Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, + Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, }; + use pyo3::{ exceptions::PyException, prelude::*, diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs index 7afd06e2373..c1203dcbd89 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -1,7 +1,9 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets::{SecretManager, SecretManagerState}; -use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, SecretValue}; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem}; + use pyo3::prelude::*; use serde_json::Value; diff --git a/litellm-rust/crates/python-bridge/src/secrets/operations.rs b/litellm-rust/crates/python-bridge/src/secrets/operations.rs index b458e3f34e8..027da1e4718 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/operations.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/operations.rs @@ -81,7 +81,7 @@ impl From for PythonMutationError { pub(super) async fn write_python_provider( manager: &SecretManager, name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, ) -> Result { match manager { SecretManager::Cyberark(client) => { @@ -128,7 +128,7 @@ pub(super) async fn rotate_python_provider( manager: &SecretManager, current_name: &str, new_name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, ) -> Result { match manager { SecretManager::Cyberark(client) => { @@ -186,7 +186,7 @@ pub(super) enum PythonMutationResponse { pub(super) async fn write_python_provider_with_context( manager: &SecretManager, name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, context: &litellm_secrets_types::SecretWriteContext, ) -> Result { if let (SecretManager::HashicorpVault(client), SecretOperationContext::Hashicorp(operation)) = @@ -235,7 +235,7 @@ pub(super) async fn rotate_python_provider_with_context( manager: &SecretManager, current_name: &str, new_name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, context: &SecretOperationContext, ) -> Result { if let (SecretManager::HashicorpVault(client), SecretOperationContext::Hashicorp(context)) = diff --git a/litellm-rust/crates/python-bridge/src/secrets/provider.rs b/litellm-rust/crates/python-bridge/src/secrets/provider.rs index ff10f596311..55343fb63bb 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/provider.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/provider.rs @@ -1,7 +1,9 @@ use std::time::Duration; use super::operations::PythonReadRequest; -use litellm_secrets::{KeyManagementSystem, SecretValue}; +use litellm_secret_value::SecretValue; +use litellm_secrets::KeyManagementSystem; + use litellm_secrets_types::{ AwsOperationContext, CyberarkOperationContext, GoogleOperationContext, HashicorpOperationContext, SecretOperationContext, diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index 2f040bff773..c5d9e3565b3 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -2,7 +2,9 @@ use std::sync::Arc; use futures_util::future::BoxFuture; use litellm_host_python::{PythonContext, attach_blocking}; -use litellm_secrets::{Error, SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::{Error, source::SecretSource}; + use pyo3::prelude::*; use super::error::external_error; diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index deb07175e7a..752b533886e 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -4,10 +4,9 @@ use futures_util::future::BoxFuture; use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::PythonContext; use litellm_http::Client; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; -use litellm_secrets::{ - Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver, SecretValue, -}; +use litellm_secrets::{Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver}; use super::config::SecretManagerSnapshot; diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index b956ec297b6..0d52b0aa629 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -307,7 +307,7 @@ impl NativeSecretManager { super::operations::write_python_provider_with_context( &backend, &secret_name, - &litellm_secrets::SecretValue::new(secret_value), + &litellm_secret_value::SecretValue::new(secret_value), &context, ) .await, @@ -365,7 +365,7 @@ impl NativeSecretManager { &backend, ¤t_secret_name, &new_secret_name, - &litellm_secrets::SecretValue::new(new_secret_value), + &litellm_secret_value::SecretValue::new(new_secret_value), &context, ) .await, diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs index 41658f0baf8..b99dff7aeac 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs @@ -1,9 +1,8 @@ use std::collections::HashMap; -use litellm_secrets::{ - SecretValue, - hashicorp::{HashicorpVault, RawOperationError}, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets::hashicorp::{HashicorpVault, RawOperationError}; + use litellm_secrets_types::{HashicorpOperationContext, SecretWriteContext}; use serde_json::value::RawValue; diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index eab57b36423..3e9447a6bfa 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-circuit-breaker.workspace = true +litellm-secret-value.workspace = true futures-util.workspace = true redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } thiserror.workspace = true diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index a02b0847663..48dbc2b2dfb 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -29,7 +29,7 @@ impl RedisClient { pub fn new(config: RedisConfig) -> Result { let connection = match &config.endpoint { Endpoint::Url(url) => { - Connection::standalone(url.as_str().into_connection_info()?, config.timeout)? + Connection::standalone(url.expose().into_connection_info()?, config.timeout)? } Endpoint::Host(host) => { Connection::standalone(host.connection_info()?, config.timeout)? @@ -204,6 +204,7 @@ mod tests { use std::time::{Duration, Instant}; use litellm_redis_testing::RedisServer; + use litellm_secret_value::SecretValue; use rstest::rstest; use litellm_circuit_breaker::{Breaker, Config, State}; @@ -630,8 +631,12 @@ mod tests { db: 3, })) .unwrap(); - let in_db_0 = connect(Endpoint::Url(server.url().to_owned())).unwrap(); - let via_url = connect(Endpoint::Url(format!("{}/3", server.url()))).unwrap(); + let in_db_0 = connect(Endpoint::Url(SecretValue::new(server.url().to_owned()))).unwrap(); + let via_url = connect(Endpoint::Url(SecretValue::new(format!( + "{}/3", + server.url() + )))) + .unwrap(); in_db_3 .set("k", "3", Duration::from_secs(60)) diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index 64934c68891..c40b9eb6b36 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -1,8 +1,9 @@ -use std::{fmt, time::Duration}; +use std::time::Duration; use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; use litellm_circuit_breaker::Breaker; +use litellm_secret_value::SecretValue; use crate::Namespace; @@ -15,28 +16,28 @@ pub struct RedisConfig { pub breaker: Option, } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub enum Endpoint { - Url(String), + Url(SecretValue), Host(HostEndpoint), Cluster(ClusterEndpoint), } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub struct HostEndpoint { pub host: String, pub port: u16, - pub username: Option, - pub password: Option, + pub username: Option, + pub password: Option, pub tls: bool, pub db: i64, } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub struct ClusterEndpoint { pub nodes: Vec, - pub username: Option, - pub password: Option, + pub username: Option, + pub password: Option, pub tls: bool, } @@ -58,11 +59,11 @@ impl HostEndpoint { pub(crate) fn connection_info(&self) -> RedisResult { let settings = RedisConnectionInfo::default().set_db(self.db); let settings = match &self.username { - Some(username) => settings.set_username(username), + Some(username) => settings.set_username(username.expose()), None => settings, }; let settings = match &self.password { - Some(password) => settings.set_password(password), + Some(password) => settings.set_password(password.expose()), None => settings, }; Ok(addr(&self.host, self.port, self.tls) @@ -83,70 +84,35 @@ fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr { } } -impl fmt::Debug for Endpoint { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Url(url) => formatter - .debug_tuple("Url") - .field(&redacted_url(url)) - .finish(), - Self::Host(host) => formatter - .debug_struct("Host") - .field("host", &host.host) - .field("port", &host.port) - .field("username", &host.username) - .field("password", &host.password.as_ref().map(|_| "")) - .field("tls", &host.tls) - .field("db", &host.db) - .finish(), - Self::Cluster(cluster) => formatter - .debug_struct("Cluster") - .field("nodes", &cluster.nodes) - .field("username", &cluster.username) - .field("password", &cluster.password.as_ref().map(|_| "")) - .field("tls", &cluster.tls) - .finish(), - } - } -} - -fn redacted_url(url: &str) -> String { - match url - .split_once("://") - .and_then(|(scheme, rest)| Some((scheme, rest.rsplit_once('@')?.1))) - { - Some((scheme, host)) => format!("{scheme}://@{host}"), - None => url.to_owned(), - } -} - #[cfg(test)] mod tests { + use litellm_secret_value::SecretValue; use rstest::rstest; - use super::{Endpoint, HostEndpoint}; + use super::{ClusterEndpoint, Endpoint, HostEndpoint}; #[rstest] - #[case::url_with_password(Endpoint::Url("redis://user:secret@cache:6379/0".to_owned()), "redis://@cache:6379/0")] - #[case::url_with_at_in_password(Endpoint::Url("rediss://:p@ss@cache:6380".to_owned()), "rediss://@cache:6380")] - #[case::host_password( + #[case::url(Endpoint::Url(SecretValue::new("redis://credential-value@cache:6379/0")))] + #[case::host( Endpoint::Host(HostEndpoint { host: "cache".to_owned(), port: 6379, - username: None, - password: Some("secret".to_owned()), + username: Some(SecretValue::new("credential-value")), + password: Some(SecretValue::new("credential-value")), tls: false, db: 0, - }), - "" + }) )] - fn debug_output_never_shows_the_password(#[case] endpoint: Endpoint, #[case] shown: &str) { + #[case::cluster(Endpoint::Cluster(ClusterEndpoint { + nodes: Vec::new(), + username: Some(SecretValue::new("credential-value")), + password: Some(SecretValue::new("credential-value")), + tls: false, + }))] + fn debug_output_never_shows_the_secret(#[case] endpoint: Endpoint) { let debug = format!("{endpoint:?}"); - assert!( - !debug.contains("secret") && !debug.contains("p@ss"), - "{debug}" - ); - assert!(debug.contains(shown), "{debug}"); + assert!(!debug.contains("credential-value"), "{debug}"); + assert!(debug.contains("[REDACTED]"), "{debug}"); } } diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index aaf26f105da..41729bfd107 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -33,11 +33,11 @@ impl Connection { .response_timeout(timeout) .retries(RETRIES); let builder = match &endpoint.username { - Some(username) => builder.username(username), + Some(username) => builder.username(username.expose()), None => builder, }; let builder = match &endpoint.password { - Some(password) => builder.password(password), + Some(password) => builder.password(password.expose()), None => builder, }; let config = ClusterConfig::new() diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs index 64a174b1543..e503739a942 100644 --- a/litellm-rust/crates/redis-client/src/test_support.rs +++ b/litellm-rust/crates/redis-client/src/test_support.rs @@ -2,6 +2,7 @@ use std::time::Duration; use litellm_circuit_breaker::Breaker; use litellm_redis_testing::{RedisCluster, RedisServer}; +use litellm_secret_value::SecretValue; use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; use crate::{ @@ -30,7 +31,7 @@ impl Server { pub(crate) fn endpoint(&self) -> Endpoint { match self { - Self::Standalone(server) => Endpoint::Url(server.url().to_owned()), + Self::Standalone(server) => Endpoint::Url(SecretValue::new(server.url().to_owned())), Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { nodes: cluster .nodes() diff --git a/litellm-rust/crates/router-types/Cargo.toml b/litellm-rust/crates/router-types/Cargo.toml index 2e9fec28685..f0470ce8a8e 100644 --- a/litellm-rust/crates/router-types/Cargo.toml +++ b/litellm-rust/crates/router-types/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true description = "The types of litellm/types/router.py: a deployment's litellm_params as a config or a caller spells them" [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/router-types/src/litellm_params.rs b/litellm-rust/crates/router-types/src/litellm_params.rs index 186a1d98036..aa2c47f6e3c 100644 --- a/litellm-rust/crates/router-types/src/litellm_params.rs +++ b/litellm-rust/crates/router-types/src/litellm_params.rs @@ -1,6 +1,8 @@ use std::{collections::BTreeMap, fmt, ops::Deref}; -use litellm_auth_types::{AwsParams, ParamSpec, SecretValue, VertexParams}; +use litellm_auth_types::{AwsParams, ParamSpec, VertexParams}; +use litellm_secret_value::SecretValue; + use serde::Deserialize; use serde_json::Value; use serde_with::{OneOrMany, formats::PreferMany, serde_as}; diff --git a/litellm-rust/crates/router-types/tests/litellm_params.rs b/litellm-rust/crates/router-types/tests/litellm_params.rs index 7d021fd857a..c5c0d16ee8f 100644 --- a/litellm-rust/crates/router-types/tests/litellm_params.rs +++ b/litellm-rust/crates/router-types/tests/litellm_params.rs @@ -1,4 +1,6 @@ -use litellm_auth_types::{AwsParams, SecretValue, VertexParams}; +use litellm_auth_types::{AwsParams, VertexParams}; +use litellm_secret_value::SecretValue; + use litellm_router_types::{LitellmParams, Spelled}; use rstest::rstest; use serde_json::{Map, Value, json}; diff --git a/litellm-rust/crates/secret-value/Cargo.toml b/litellm-rust/crates/secret-value/Cargo.toml new file mode 100644 index 00000000000..467f852a278 --- /dev/null +++ b/litellm-rust/crates/secret-value/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "litellm-secret-value" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +serde.workspace = true +subtle.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/secret-value/src/lib.rs b/litellm-rust/crates/secret-value/src/lib.rs new file mode 100644 index 00000000000..f11d22d3701 --- /dev/null +++ b/litellm-rust/crates/secret-value/src/lib.rs @@ -0,0 +1,5 @@ +#![forbid(unsafe_code)] + +mod secret_value; + +pub use secret_value::SecretValue; diff --git a/litellm-rust/crates/auth-types/src/secret.rs b/litellm-rust/crates/secret-value/src/secret_value.rs similarity index 95% rename from litellm-rust/crates/auth-types/src/secret.rs rename to litellm-rust/crates/secret-value/src/secret_value.rs index 7e6789deef7..27e4227080d 100644 --- a/litellm-rust/crates/auth-types/src/secret.rs +++ b/litellm-rust/crates/secret-value/src/secret_value.rs @@ -32,9 +32,11 @@ impl Hash for SecretValue { #[cfg(test)] mod tests { + use rstest::rstest; + use super::SecretValue; - #[test] + #[rstest] fn debug_redacts_plaintext() { let debug = format!("{:?}", SecretValue::new("credential-value")); @@ -42,7 +44,7 @@ mod tests { assert!(debug.contains("REDACTED")); } - #[test] + #[rstest] fn equality_compares_plaintext_values() { assert_eq!(SecretValue::new("same"), SecretValue::new("same")); assert_ne!(SecretValue::new("same"), SecretValue::new("different")); diff --git a/litellm-rust/crates/secrets-aws/Cargo.toml b/litellm-rust/crates/secrets-aws/Cargo.toml index 5d3bd413484..d17dfcfeec2 100644 --- a/litellm-rust/crates/secrets-aws/Cargo.toml +++ b/litellm-rust/crates/secrets-aws/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth-aws.workspace = true litellm-secrets-types.workspace = true litellm-core-utils.workspace = true diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs index 2220e387838..b25872ab3f6 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -21,11 +21,12 @@ use litellm_auth_aws::constants::{ AWS_ACCESS_KEY_ID, AWS_REGION, AWS_REGION_NAME, AWS_SECRET_ACCESS_KEY, }; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ AwsOperationContext, BaseSecretManager, KeyManagementSettings, RotationError, Secret, - SecretDeleter, SecretRotator, SecretValue, SecretWriteContext, SecretWriter, - async_rotate_secret, + SecretDeleter, SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, }; + use serde_json::Value; use crate::{Error, auth}; diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs index c482a168090..56999633a78 100644 --- a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs @@ -10,11 +10,13 @@ use aws_sdk_secretsmanager::{ Client, config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, }; +use litellm_secret_value::SecretValue; use litellm_secrets_aws::{AwsSecretsManagerV2, Error, RotationResponse}; use litellm_secrets_types::{ AwsOperationContext, BaseSecretManager, KeyManagementSettings, Secret, SecretDeleter, - SecretValue, SecretWriteContext, SecretWriter, + SecretWriteContext, SecretWriter, }; + use rstest::{fixture, rstest}; use serde_json::json; use wiremock::{ diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml index 2aff6129c37..c4bfb24898c 100644 --- a/litellm-rust/crates/secrets-azure/Cargo.toml +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true tokio.workspace = true litellm-auth-azure.workspace = true diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs index d97b7a22cfc..40028828298 100644 --- a/litellm-rust/crates/secrets-azure/src/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -3,7 +3,9 @@ use std::sync::Arc; use litellm_auth_azure::{AzureAuthInputs, AzureAuthService, ConfigValue}; use litellm_auth_types::{InputSource, Sourced}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::{AzureOperationContext, BaseSecretManager, Secret, SecretValue}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{AzureOperationContext, BaseSecretManager, Secret}; + use percent_encoding::{AsciiSet, NON_ALPHANUMERIC}; use serde::Deserialize; diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs index a0ba9e6ddf8..c3ba01672a8 100644 --- a/litellm-rust/crates/secrets-azure/tests/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -1,7 +1,9 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets_azure::{AzureKeyVault, Error}; -use litellm_secrets_types::{Secret, SecretValue}; +use litellm_secrets_types::Secret; + use rstest::{fixture, rstest}; use serde::Deserialize; use wiremock::{ diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml index 2ae2c1ea5c3..643f322f7fa 100644 --- a/litellm-rust/crates/secrets-cyberark/Cargo.toml +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-secrets-types.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index c5cb0308b13..cf077a2bbc8 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -9,11 +9,12 @@ use litellm_core_utils::settings::Lookup; use litellm_http::{ Client, ClientIdentity, ClientVariant, HttpClientConfig, HttpClientPool, TlsSource, Verify, }; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, CyberarkOperationContext, RotationError, SecretCache, SecretDeleter, - SecretRotator, SecretValue, SecretWriteContext, SecretWriter, async_rotate_secret, - validate_secret_name, + SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, validate_secret_name, }; + use moka::future::Cache; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode}; diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs index 783ba6ff67a..e9b480e026b 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs @@ -9,8 +9,10 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; +use litellm_secret_value::SecretValue; use litellm_secrets_cyberark::{CyberArkSecretManager, DeleteOutcome, Error}; -use litellm_secrets_types::{BaseSecretManager, CyberarkOperationContext, SecretValue}; +use litellm_secrets_types::{BaseSecretManager, CyberarkOperationContext}; + use rstest::{fixture, rstest}; use serde::Deserialize; use wiremock::{ diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml index 06b3747def5..3c5c2130788 100644 --- a/litellm-rust/crates/secrets-google/Cargo.toml +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true moka.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/secrets-google/src/auth.rs b/litellm-rust/crates/secrets-google/src/auth.rs index 45fc99d8d5f..7022ecc5ae4 100644 --- a/litellm-rust/crates/secrets-google/src/auth.rs +++ b/litellm-rust/crates/secrets-google/src/auth.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use litellm_auth_gcp::{GoogleCredentials, VertexConfig}; use litellm_auth_types::{InputSource, Sourced}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; pub(crate) fn credentials( project: Option, diff --git a/litellm-rust/crates/secrets-google/src/kms.rs b/litellm-rust/crates/secrets-google/src/kms.rs index 78ebae008f7..5622a3a4d56 100644 --- a/litellm-rust/crates/secrets-google/src/kms.rs +++ b/litellm-rust/crates/secrets-google/src/kms.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use google_cloud_kms_v1::client::KeyManagementService; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{Error, auth}; diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs index 463d1f6fd73..577a2ccbd4d 100644 --- a/litellm-rust/crates/secrets-google/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -2,9 +2,9 @@ use std::{sync::Arc, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::{ - BaseSecretManager, GoogleOperationContext, Secret, SecretCache, SecretValue, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{BaseSecretManager, GoogleOperationContext, Secret, SecretCache}; + use serde::Deserialize; use litellm_auth_gcp::GoogleCredentials; diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml index 8731a8deb7a..4fcd496931d 100644 --- a/litellm-rust/crates/secrets-hashicorp/Cargo.toml +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-error.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs index 4232667ae16..a63161b17b2 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/config.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -1,7 +1,7 @@ use std::{path::PathBuf, time::Duration}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::Error; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index 447beacda85..dae14439630 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -12,11 +12,12 @@ use std::{ }; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, HashicorpOperationContext, RotationError, SecretCache, SecretDeleter, - SecretRotator, SecretValue, SecretWriteContext, SecretWriter, async_rotate_secret, - validate_secret_name, + SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, validate_secret_name, }; + use rustify::errors::ClientError as RustifyClientError; use serde_json::Value; use tokio::sync::Mutex; diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs index 668871d3fc7..d431e472400 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -1,11 +1,13 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_hashicorp::{Error, HashicorpVault, HashicorpVaultConfig}; use litellm_secrets_types::{ - BaseSecretManager, HashicorpOperationContext, RotationError, SecretDeleter, SecretValue, - SecretWriteContext, SecretWriter, + BaseSecretManager, HashicorpOperationContext, RotationError, SecretDeleter, SecretWriteContext, + SecretWriter, }; + use rstest::{fixture, rstest}; use serde::Deserialize; use serde_json::json; diff --git a/litellm-rust/crates/secrets-types/Cargo.toml b/litellm-rust/crates/secrets-types/Cargo.toml index 6dd847ec989..b45c41bcf1c 100644 --- a/litellm-rust/crates/secrets-types/Cargo.toml +++ b/litellm-rust/crates/secrets-types/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth-types.workspace = true +litellm-secret-value.workspace = true moka.workspace = true tokio = { workspace = true, features = ["sync"] } serde.workspace = true diff --git a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs index 737fad28662..41dadd1baa6 100644 --- a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs +++ b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs @@ -1,4 +1,5 @@ -use crate::{Error, SecretValue, SecretWriteContext}; +use crate::{Error, SecretWriteContext}; +use litellm_secret_value::SecretValue; pub fn validate_secret_name(name: &str) -> Result<(), Error> { if name.split('/').any(|segment| segment == "..") diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs index 9b6e82c3a6b..a45acb4df11 100644 --- a/litellm-rust/crates/secrets-types/src/config.rs +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use strum::IntoStaticStr; -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[derive( Clone, diff --git a/litellm-rust/crates/secrets-types/src/context.rs b/litellm-rust/crates/secrets-types/src/context.rs index 1ec493edf42..b71195e2d21 100644 --- a/litellm-rust/crates/secrets-types/src/context.rs +++ b/litellm-rust/crates/secrets-types/src/context.rs @@ -1,6 +1,7 @@ use std::{collections::BTreeMap, time::Duration}; -use crate::{Error, KeyManagementSystem, SecretValue}; +use crate::{Error, KeyManagementSystem}; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, Default, Eq, PartialEq)] pub enum SecretOperationContext { diff --git a/litellm-rust/crates/secrets-types/src/lib.rs b/litellm-rust/crates/secrets-types/src/lib.rs index 4d7eaf8313e..2563ccbe39d 100644 --- a/litellm-rust/crates/secrets-types/src/lib.rs +++ b/litellm-rust/crates/secrets-types/src/lib.rs @@ -18,5 +18,4 @@ pub use context::{ HashicorpOperationContext, SecretOperationContext, SecretWriteContext, }; pub use error::Error; -pub use litellm_auth_types::SecretValue; pub use value::{PythonSecretRead, Secret}; diff --git a/litellm-rust/crates/secrets-types/src/value.rs b/litellm-rust/crates/secrets-types/src/value.rs index 11f57ef4448..8c92902edad 100644 --- a/litellm-rust/crates/secrets-types/src/value.rs +++ b/litellm-rust/crates/secrets-types/src/value.rs @@ -1,4 +1,4 @@ -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[derive(Clone, PartialEq, Eq, veil::Redact)] pub enum Secret { diff --git a/litellm-rust/crates/secrets-types/tests/config.rs b/litellm-rust/crates/secrets-types/tests/config.rs index 45fd9d0ee03..e7ba63343da 100644 --- a/litellm-rust/crates/secrets-types/tests/config.rs +++ b/litellm-rust/crates/secrets-types/tests/config.rs @@ -1,6 +1,6 @@ -use litellm_secrets_types::{ - AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, Secret}; + use rstest::{fixture, rstest}; use serde_json::json; diff --git a/litellm-rust/crates/secrets-types/tests/context.rs b/litellm-rust/crates/secrets-types/tests/context.rs index 8f8eeabdb62..f8d754564dc 100644 --- a/litellm-rust/crates/secrets-types/tests/context.rs +++ b/litellm-rust/crates/secrets-types/tests/context.rs @@ -1,10 +1,11 @@ use std::{collections::BTreeMap, time::Duration}; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ AwsOperationContext, AzureOperationContext, CyberarkOperationContext, GoogleOperationContext, - HashicorpOperationContext, KeyManagementSystem, SecretOperationContext, SecretValue, - SecretWriteContext, + HashicorpOperationContext, KeyManagementSystem, SecretOperationContext, SecretWriteContext, }; + use rstest::{fixture, rstest}; #[fixture] diff --git a/litellm-rust/crates/secrets-types/tests/rotation.rs b/litellm-rust/crates/secrets-types/tests/rotation.rs index 36c1a065ffa..089890c60bd 100644 --- a/litellm-rust/crates/secrets-types/tests/rotation.rs +++ b/litellm-rust/crates/secrets-types/tests/rotation.rs @@ -1,10 +1,12 @@ use std::sync::atomic::{AtomicUsize, Ordering}; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, Error, HashicorpOperationContext, RotationError, SecretDeleter, - SecretOperationContext, SecretRotator, SecretValue, SecretWriteContext, SecretWriter, - async_rotate_secret, validate_secret_name, + SecretOperationContext, SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, + validate_secret_name, }; + use rstest::{fixture, rstest}; struct Manager { diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 023d61136e4..21ee1bdc760 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -14,6 +14,7 @@ azure = ["dep:litellm-secrets-azure"] cyberark = ["dep:litellm-secrets-cyberark"] [dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true litellm-python-compat.workspace = true litellm-secrets-types.workspace = true diff --git a/litellm-rust/crates/secrets/src/compatibility.rs b/litellm-rust/crates/secrets/src/compatibility.rs index 368c9931b51..3d6ba854325 100644 --- a/litellm-rust/crates/secrets/src/compatibility.rs +++ b/litellm-rust/crates/secrets/src/compatibility.rs @@ -3,9 +3,10 @@ use litellm_python_compat::{Value, literal::literal_eval}; use litellm_secrets_types::PythonSecretRead; use crate::{ - Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretManager, SecretValue, + Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretManager, get_secret_from_manager, }; +use litellm_secret_value::SecretValue; pub async fn get_secret_from_python_manager( manager: &SecretManager, diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index f42e09c6bdb..2d2b81d197f 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -5,7 +5,7 @@ use litellm_core_utils::settings::Lookup; use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret}; #[cfg(any(feature = "aws", feature = "google"))] -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[cfg(any( feature = "google", diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index 8d0f0561608..e4459576e94 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -12,9 +12,7 @@ mod state; pub use compatibility::{get_secret_from_python_manager, read_secret_from_python_manager}; pub use error::Error; pub use handler::{ExternalSecretManager, SecretManager, get_secret_from_manager}; -pub use litellm_secrets_types::{ - AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, -}; +pub use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, Secret}; pub use native::load_native_manager; pub use oidc::{OidcProvider, OidcReference, OidcResolver}; pub use resolver::{FailurePolicy, SecretResolver, normalize_nonempty_secret_str}; diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index 24e6912c797..18fc9479ef1 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -9,7 +9,8 @@ use litellm_http::Client; use moka::future::Cache; use serde::Deserialize; -use crate::{Error, SecretValue}; +use crate::Error; +use litellm_secret_value::SecretValue; const GOOGLE_TOKEN_MAX_TTL: Duration = Duration::from_secs(3540); const GITHUB_TOKEN_TTL: Duration = Duration::from_secs(295); diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs index 830e7c22cd5..631b60cba10 100644 --- a/litellm-rust/crates/secrets/src/resolver.rs +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -4,7 +4,8 @@ use crate::compatibility::python_manager_string; use litellm_core_utils::{serde_compat::parse_str_bool, settings::Lookup}; use crate::state::{LookupTarget, normalize_secret_name}; -use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue}; +use crate::{Error, OidcResolver, Secret, SecretManagerState}; +use litellm_secret_value::SecretValue; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum FailurePolicy { diff --git a/litellm-rust/crates/secrets/src/source.rs b/litellm-rust/crates/secrets/src/source.rs index 3c86fd25e9a..1ee224b763a 100644 --- a/litellm-rust/crates/secrets/src/source.rs +++ b/litellm-rust/crates/secrets/src/source.rs @@ -3,7 +3,8 @@ use std::{collections::HashMap, sync::Arc}; use futures_util::future::{BoxFuture, try_join_all}; use litellm_core_utils::settings::Lookup; -use crate::{Error, SecretResolver, SecretValue}; +use crate::{Error, SecretResolver}; +use litellm_secret_value::SecretValue; pub type Secrets = Arc; diff --git a/litellm-rust/crates/secrets/tests/aws.rs b/litellm-rust/crates/secrets/tests/aws.rs index abad04f31c5..6f30d1b5161 100644 --- a/litellm-rust/crates/secrets/tests/aws.rs +++ b/litellm-rust/crates/secrets/tests/aws.rs @@ -2,11 +2,13 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets::{ AccessMode, Error, FailurePolicy, KeyManagementSettings, OidcResolver, Secret, SecretManager, - SecretManagerState, SecretResolver, SecretValue, aws::AwsSecretsManagerV2, + SecretManagerState, SecretResolver, aws::AwsSecretsManagerV2, secret_manager_would_be_consulted, }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; fn state(server: &MockServer, settings: KeyManagementSettings) -> SecretManagerState { diff --git a/litellm-rust/crates/secrets/tests/azure.rs b/litellm-rust/crates/secrets/tests/azure.rs index 6e81017b0f6..201e31f44b8 100644 --- a/litellm-rust/crates/secrets/tests/azure.rs +++ b/litellm-rust/crates/secrets/tests/azure.rs @@ -68,10 +68,11 @@ async fn azure_handler_reads_missing_and_failed_secrets() { async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_null( #[case] body: serde_json::Value, ) { + use litellm_secret_value::SecretValue; use litellm_secrets::{ - OidcResolver, SecretManager, SecretManagerState, SecretResolver, SecretValue, - azure::AzureKeyVault, + OidcResolver, SecretManager, SecretManagerState, SecretResolver, azure::AzureKeyVault, }; + use std::sync::Arc; use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; let server = MockServer::start().await; diff --git a/litellm-rust/crates/secrets/tests/common_read_contract.rs b/litellm-rust/crates/secrets/tests/common_read_contract.rs index e6928038799..4cf3513fadf 100644 --- a/litellm-rust/crates/secrets/tests/common_read_contract.rs +++ b/litellm-rust/crates/secrets/tests/common_read_contract.rs @@ -9,8 +9,9 @@ use std::sync::Arc; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets::{ - KeyManagementSettings, SecretManager, SecretValue, + KeyManagementSettings, SecretManager, aws::AwsSecretsManagerV2, azure::AzureKeyVault, cyberark::CyberArkSecretManager, @@ -18,6 +19,7 @@ use litellm_secrets::{ google::GoogleSecretManager, hashicorp::{HashicorpVault, HashicorpVaultConfig}, }; + use rstest::rstest; use serde_json::json; use wiremock::{ diff --git a/litellm-rust/crates/secrets/tests/cyberark.rs b/litellm-rust/crates/secrets/tests/cyberark.rs index 0a6c394937e..09342234118 100644 --- a/litellm-rust/crates/secrets/tests/cyberark.rs +++ b/litellm-rust/crates/secrets/tests/cyberark.rs @@ -4,10 +4,12 @@ async fn cyberark_handler_reads_values_and_surfaces_errors() { use std::time::Duration; + use litellm_secret_value::SecretValue; use litellm_secrets::{ - Error, KeyManagementSettings, SecretManager, SecretValue, cyberark::CyberArkSecretManager, + Error, KeyManagementSettings, SecretManager, cyberark::CyberArkSecretManager, get_secret_from_manager, }; + use wiremock::{ Mock, MockServer, ResponseTemplate, matchers::{body_string, path}, diff --git a/litellm-rust/crates/secrets/tests/google.rs b/litellm-rust/crates/secrets/tests/google.rs index 5542c864dfd..a75e34335de 100644 --- a/litellm-rust/crates/secrets/tests/google.rs +++ b/litellm-rust/crates/secrets/tests/google.rs @@ -7,10 +7,12 @@ use std::sync::Arc; #[case::failure(503)] #[tokio::test] async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) { + use litellm_secret_value::SecretValue; use litellm_secrets::{ Error, FailurePolicy, KeyManagementSettings, OidcResolver, SecretManager, - SecretManagerState, SecretResolver, SecretValue, google::GoogleSecretManager, + SecretManagerState, SecretResolver, google::GoogleSecretManager, }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; let server = MockServer::start().await; Mock::given(method("GET")) diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs index 867a662cbbb..2d250e1de9d 100644 --- a/litellm-rust/crates/secrets/tests/oidc.rs +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -315,7 +315,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] Box< dyn std::future::Future< Output = Result< - litellm_secrets::SecretValue, + litellm_secret_value::SecretValue, litellm_secrets::azure::Error, >, > + Send @@ -331,7 +331,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] if self.0 { Err(litellm_secrets::azure::Error::MissingCredentials) } else { - Ok(litellm_secrets::SecretValue::new("azure-token")) + Ok(litellm_secret_value::SecretValue::new("azure-token")) } }) } @@ -346,7 +346,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] let result = resolver .get_secret_str( "oidc/azure/api://audience/path", - Some(litellm_secrets::SecretValue::new("fallback")), + Some(litellm_secret_value::SecretValue::new("fallback")), ) .await; if failed { diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index 8435e051315..185f2256780 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -1,9 +1,10 @@ use std::{future::Future, pin::Pin, sync::Arc}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets::{ Error, ExternalSecretManager, FailurePolicy, KeyManagementSettings, KeyManagementSystem, - OidcResolver, Secret, SecretManager, SecretManagerState, SecretResolver, SecretValue, + OidcResolver, Secret, SecretManager, SecretManagerState, SecretResolver, secret_manager_would_be_consulted, }; diff --git a/litellm/__init__.py b/litellm/__init__.py index 5f14c605cd0..f64531f86fb 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -7,7 +7,6 @@ _apply_process_defaults() ### INIT VARIABLES ######################### import threading import os -import datetime from collections.abc import Mapping, Sequence from typing import ( @@ -550,6 +549,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None output_parse_pii: bool = False ############################################# from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete +from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired model_cost = get_model_cost_map(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount @@ -756,16 +756,6 @@ def is_openai_finetune_model(key: str) -> bool: return key.startswith("ft:") and not key.count(":") > 1 -def _is_past_deprecation_date(value: Mapping[str, object]) -> bool: - deprecation_date: Final = value.get("deprecation_date") - if not isinstance(deprecation_date, str): - return False - try: - return datetime.date.fromisoformat(deprecation_date) <= datetime.date.today() - except ValueError: - return False - - def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None: for key, value in tuple(model_cost_map.items()): if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): @@ -873,12 +863,12 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje if "-to-" not in key and "fireworks-ai-default" not in key: fireworks_ai_models.add(key) elif value.get("litellm_provider") == "together_ai": - if _is_past_deprecation_date(value): + if is_hidden_as_retired(value): together_ai_models.discard(key) elif not key.startswith("together-ai-"): together_ai_models.add(key) elif value.get("litellm_provider") == "baseten": - if _is_past_deprecation_date(value): + if is_hidden_as_retired(value): baseten_models.discard(key) else: baseten_models.add(key) @@ -1051,7 +1041,10 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje elif value.get("litellm_provider") == "coralbricks": coralbricks_models.add(key) elif value.get("litellm_provider") == "aiand": - aiand_models.add(key) + if is_hidden_as_retired(value): + aiand_models.discard(key) + else: + aiand_models.add(key) elif value.get("litellm_provider") == "reducto": reducto_models.add(key) elif value.get("litellm_provider") == "bedrock_mantle": diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index c1f484f8d64..de290070ed3 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -1469,11 +1469,11 @@ _UTILS_MODULE_IMPORT_MAP: Final = { "AsyncHTTPHandler": ("litellm.llms.custom_httpx.http_handler", "AsyncHTTPHandler"), "HTTPHandler": ("litellm.llms.custom_httpx.http_handler", "HTTPHandler"), "get_num_retries_from_retry_policy": ( - "litellm.router_utils.get_retry_from_policy", + "litellm.litellm_core_utils.get_retry_from_policy", "get_num_retries_from_retry_policy", ), "reset_retry_policy": ( - "litellm.router_utils.get_retry_from_policy", + "litellm.litellm_core_utils.get_retry_from_policy", "reset_retry_policy", ), "get_coroutine_checker": ( diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 6d3b09aed68..1e003a6e4bc 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -22,7 +22,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( A2AStreamingContext, ) from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager -from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS +from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS, A2A_COMPLETION_BRIDGE_WORKING_MESSAGE from litellm.interactions.agents.utils import merge_agent_headers from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import ModelResponse @@ -268,7 +268,7 @@ class A2ACompletionBridgeHandler: ctx=ctx, state="working", final=False, - message_text="Processing request...", + message_text=A2A_COMPLETION_BRIDGE_WORKING_MESSAGE, ) yield working_event diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index f44e3f6a81b..dcbe7a6506e 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -191,6 +191,23 @@ def _set_litellm_params_on_logging_obj( } +async def _record_send_message_billing( + kwargs: Mapping[str, object], + request: "SendMessageRequest", + response_dict: Mapping[str, object], + litellm_params: Mapping[str, object], + agent_id: str | None, +) -> None: + prompt_tokens, completion_tokens, _ = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)( + request=request, + response_dict=response_dict, + ) + _set_usage_on_logging_obj(kwargs=kwargs, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens) + if "error" not in response_dict: + _set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params) + _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id) + + def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: Mapping[str, object]) -> str: """ Extract agent info and set model/custom_llm_provider for cost tracking. @@ -483,13 +500,21 @@ async def asend_message( if custom_llm_provider: if request is None: raise ValueError("request is required for completion bridge") - return await _send_message_via_completion_bridge( + bridge_response: Final = await _send_message_via_completion_bridge( request=request, custom_llm_provider=custom_llm_provider, api_base=api_base, litellm_params=litellm_params, agent_extra_headers=agent_extra_headers, ) + await _record_send_message_billing( + kwargs=kwargs, + request=request, + response_dict=bridge_response.model_dump(mode="json", exclude_none=True), + litellm_params=litellm_params, + agent_id=agent_id, + ) + return bridge_response # Standard A2A client flow if request is None: @@ -538,30 +563,14 @@ async def asend_message( # Wrap in LiteLLM response type for _hidden_params support response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id)) - # Calculate token usage from request and response - response_dict: Final[dict[str, object]] = a2a_response.root.model_dump(mode="json", exclude_none=True) - ( - prompt_tokens, - completion_tokens, - _, - ) = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)( - request=request, - response_dict=response_dict, - ) - - # Set usage on logging obj for standard logging payload - _set_usage_on_logging_obj( + await _record_send_message_billing( kwargs=kwargs, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, + request=request, + response_dict=a2a_response.root.model_dump(mode="json", exclude_none=True), + litellm_params=litellm_params, + agent_id=agent_id, ) - # Merge agent pricing params into the logging obj so cost is calculated - _set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params) - - # Set agent_id on logging obj for SpendLogs tracking - _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id) - return response @@ -646,6 +655,7 @@ async def asend_message_streaming( metadata: dict[str, object] | None = None, proxy_server_request: dict[str, object] | None = None, agent_extra_headers: dict[str, str] | None = None, + agent_name: str | None = None, **kwargs: object, ) -> AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"]: """ @@ -661,6 +671,7 @@ async def asend_message_streaming( agent_id: Optional agent ID for tracking in SpendLogs metadata: Optional metadata dict (contains user_api_key, user_id, team_id, etc.) proxy_server_request: Optional proxy server request data + agent_name: Agent name logged as a2a_agent/ for completion bridge agents Yields: SendStreamingMessageResponse chunks from the agent @@ -685,6 +696,7 @@ async def asend_message_streaming( print(chunk) ``` """ + called_at: Final = datetime.datetime.now() litellm_params = litellm_params or {} custom_llm_provider: Final = litellm_params.get("custom_llm_provider") @@ -705,20 +717,47 @@ async def asend_message_streaming( request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params) ) - async for chunk in A2ACompletionBridgeHandler.handle_streaming( + bridge_stream: Final = A2ACompletionBridgeHandler.handle_streaming( request_id=str(request.id), params=params, litellm_params=litellm_params, api_base=api_base, agent_extra_headers=agent_extra_headers, - ): - yield chunk + ) + caller_logging_obj: Final = kwargs.get("litellm_logging_obj") + bridge_logging_obj: Final = ( + caller_logging_obj + if isinstance(caller_logging_obj, Logging) + else _build_streaming_logging_obj( + request=request, + agent_name=agent_name or "unknown", + agent_id=agent_id, + litellm_params=litellm_params, + metadata=metadata, + proxy_server_request=proxy_server_request, + start_time=called_at, + ) + ) + _set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params) + _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id) + + bridge_iterator: Final = A2AStreamingIterator( + stream=bridge_stream, + request=request, + logging_obj=bridge_logging_obj, + agent_name=agent_name or "unknown", + ) + try: + async for chunk in bridge_iterator: + yield chunk + except (GeneratorExit, asyncio.CancelledError): + await bridge_iterator.abill_on_disconnect() + raise return if request is None: raise ValueError("request is required") - called_at: Final = datetime.datetime.now() _raw_logging_obj: Final = kwargs.get("litellm_logging_obj") logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None @@ -742,12 +781,12 @@ async def asend_message_streaming( assert a2a_client is not None - agent_name: Final = _get_a2a_model_info(a2a_client, kwargs) + card_agent_name: Final = _get_a2a_model_info(a2a_client, kwargs) if logging_obj is None: logging_obj = _build_streaming_logging_obj( request=request, - agent_name=agent_name, + agent_name=card_agent_name, agent_id=agent_id, litellm_params=litellm_params, metadata=metadata, @@ -755,7 +794,7 @@ async def asend_message_streaming( start_time=called_at, ) - verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name) + verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, card_agent_name) agent_card: Final = _get_a2a_client_agent_card(a2a_client) card_url: Final = get_agent_card_url(agent_card) if agent_card else None @@ -766,18 +805,23 @@ async def asend_message_streaming( agent_card=agent_card, card_url=card_url, api_base=api_base, - agent_name=agent_name, + agent_name=card_agent_name, ) _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id) - async for chunk in A2AStreamingIterator( + url_agent_iterator: Final = A2AStreamingIterator( stream=stream, request=request, logging_obj=logging_obj, - agent_name=agent_name, - ): - yield chunk + agent_name=card_agent_name, + ) + try: + async for chunk in url_agent_iterator: + yield chunk + except (GeneratorExit, asyncio.CancelledError): + await url_agent_iterator.abill_on_disconnect() + raise async def create_a2a_client( diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 35fe6e76f4d..bc65bac1a7c 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -515,6 +515,8 @@ class PydanticAITransformation: "contextId": context_id, "kind": "artifact-update", "taskId": task_id, + "append": i > 0, + "lastChunk": is_last_chunk, "artifact": { "artifactId": artifact_id, "parts": [ diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py index c2d581321cd..75262bde392 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py @@ -206,6 +206,8 @@ class WatsonxOrchestrateTransformation: "contextId": context_id, "kind": "artifact-update", "taskId": task_id, + "append": i > 0, + "lastChunk": is_last, "artifact": { "artifactId": artifact_id, "parts": [{"kind": "text", "text": chunk_text}], diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index 328681b8dbe..7b2b6884cae 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -3,16 +3,18 @@ A2A Streaming Iterator with token tracking and logging support. """ import asyncio -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from datetime import datetime from typing import TYPE_CHECKING, Final +import anyio from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.a2a_protocol.cost_calculator import A2ACostCalculator from litellm.a2a_protocol.utils import A2ARequestUtils +from litellm.constants import A2A_COMPLETION_BRIDGE_WORKING_MESSAGE from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -23,6 +25,23 @@ if TYPE_CHECKING: _OPAQUE_PAYLOAD: Final = TypeAdapter(object) +def _chunk_payload(chunk: "SendStreamingMessageResponse | Mapping[str, object]") -> Mapping[str, object]: + if isinstance(chunk, Mapping): + return chunk + return chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} + + +def _artifact_update(payload: Mapping[str, object]) -> tuple[str, bool] | None: + result: Final = as_str_mapping(payload.get("result")) + if result is None or result.get("kind") != "artifact-update": + return None + artifact: Final = as_str_mapping(result.get("artifact")) + artifact_id: Final = artifact.get("artifactId") if artifact is not None else None + if not isinstance(artifact_id, str): + return None + return artifact_id, result.get("append") is True + + class A2AStreamingIterator: """ Async iterator for A2A streaming responses with token tracking. @@ -32,7 +51,7 @@ class A2AStreamingIterator: def __init__( self, - stream: AsyncIterator["SendStreamingMessageResponse"], + stream: AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"], request: "SendStreamingMessageRequest", logging_obj: LiteLLMLoggingObj, agent_name: str = "unknown", @@ -44,14 +63,16 @@ class A2AStreamingIterator: self.start_time = logging_obj.start_time # Collect chunks for token counting - self.chunks: list[SendStreamingMessageResponse] = [] - self.collected_text_parts: list[str] = [] - self.final_chunk: SendStreamingMessageResponse | None = None + self.chunks: list[SendStreamingMessageResponse | Mapping[str, object]] = [] + self.latest_message_text = "" + self.artifact_texts: Mapping[str, str] = {} + self.saw_error = False + self.final_chunk: SendStreamingMessageResponse | Mapping[str, object] | None = None def __aiter__(self): return self - async def __anext__(self) -> "SendStreamingMessageResponse": + async def __anext__(self) -> "SendStreamingMessageResponse | Mapping[str, object]": try: chunk: Final = await self.stream.__anext__() if self.logging_obj.completion_start_time is None: @@ -59,38 +80,54 @@ class A2AStreamingIterator: # Store chunk self.chunks.append(chunk) + payload: Final = _chunk_payload(chunk) + self.saw_error = self.saw_error or payload.get("error") is not None # Extract text from chunk for token counting - self._collect_text_from_chunk(chunk) + self._collect_text_from_chunk(payload) # Check if this is the final chunk (completed status) - if self._is_completed_chunk(chunk): + if self._is_completed_chunk(payload): self.final_chunk = chunk return chunk except StopAsyncIteration: - # Stream ended - handle logging - if self.final_chunk is None and self.chunks: - self.final_chunk = self.chunks[-1] await self._handle_stream_complete() raise - def _collect_text_from_chunk(self, chunk: "SendStreamingMessageResponse") -> None: - """Extract text from a streaming chunk and add to collected parts.""" + async def abill_on_disconnect(self) -> None: + if not self.chunks: + return + with anyio.CancelScope(shield=True): + await self._handle_stream_complete() + + def _collect_text_from_chunk(self, payload: Mapping[str, object]) -> None: + """Record the reply text a streaming chunk carries.""" try: - chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} - text: Final = A2ARequestUtils.extract_text_from_response(chunk_dict) - if text: - self.collected_text_parts.append(text) + text: Final = A2ARequestUtils.extract_text_from_stream_event(payload) + if not text: + return + artifact_update: Final = _artifact_update(payload) + if artifact_update is None and text == A2A_COMPLETION_BRIDGE_WORKING_MESSAGE: + return + if artifact_update is None: + self.latest_message_text = text + return + artifact_id, append = artifact_update + artifact_text: Final = self.artifact_texts.get(artifact_id, "") + text if append else text + self.artifact_texts = {**self.artifact_texts, artifact_id: artifact_text} except Exception: verbose_logger.debug("Failed to extract text from A2A streaming chunk") - def _is_completed_chunk(self, chunk: "SendStreamingMessageResponse") -> bool: + def _reply_text(self) -> str: + artifact_text: Final = " ".join(text for text in self.artifact_texts.values() if text) + return artifact_text or self.latest_message_text + + def _is_completed_chunk(self, payload: Mapping[str, object]) -> bool: """Check if chunk indicates stream completion.""" try: - chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} - result: Final = chunk_dict.get("result", {}) + result: Final = payload.get("result", {}) if isinstance(result, dict): status: Final = result.get("status", {}) if isinstance(status, dict): @@ -101,6 +138,8 @@ class A2AStreamingIterator: async def _handle_stream_complete(self) -> None: """Handle logging and token counting when stream completes.""" + if self.final_chunk is None and self.chunks: + self.final_chunk = self.chunks[-1] try: end_time: Final = datetime.now() @@ -109,9 +148,7 @@ class A2AStreamingIterator: input_text: Final = A2ARequestUtils.extract_text_from_message(input_message) prompt_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(input_text) - # Use the last (most complete) text from chunks - output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else "" - completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(output_text) + completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(self._reply_text()) total_tokens: Final = prompt_tokens + completion_tokens @@ -128,7 +165,7 @@ class A2AStreamingIterator: self.logging_obj.model_call_details["stream"] = False # Calculate cost using A2ACostCalculator - response_cost: Final = A2ACostCalculator.calculate_a2a_cost(self.logging_obj) + response_cost: Final = 0.0 if self.saw_error else A2ACostCalculator.calculate_a2a_cost(self.logging_obj) self.logging_obj.model_call_details["response_cost"] = response_cost # Build result for logging @@ -167,9 +204,7 @@ class A2AStreamingIterator: # Add final chunk result if available if self.final_chunk: try: - chunk_dict: Final = as_str_mapping( - _OPAQUE_PAYLOAD.validate_python(self.final_chunk.model_dump(mode="json", exclude_none=True)) - ) + chunk_dict: Final = as_str_mapping(_OPAQUE_PAYLOAD.validate_python(_chunk_payload(self.final_chunk))) if chunk_dict is not None: result["result"] = chunk_dict.get("result", {}) except Exception: diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index 17d1b86a609..a9ec0de116e 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -6,13 +6,17 @@ import hashlib from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Protocol +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import has_attributes +from litellm.litellm_core_utils.core_helpers import as_str_mapping, has_attributes if TYPE_CHECKING: from a2a.types import SendMessageRequest, SendStreamingMessageRequest +_OBJECT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + class _SupportsModelDump(Protocol): def model_dump(self, *, mode: str) -> object: ... @@ -34,6 +38,9 @@ class A2ARequestUtils: """ if message is None: return "" + # a2a-sdk Part RootModels keep kind/text on part.root, so read pydantic messages as dicts + if has_attributes(message, _SupportsModelDump, "model_dump"): + return A2ARequestUtils.extract_text_from_message(message.model_dump(mode="json")) # Handle both dict and object access if isinstance(message, dict): @@ -63,8 +70,8 @@ class A2ARequestUtils: Returns: Text from response message parts """ - result: Final = response_dict.get("result", {}) - if not isinstance(result, dict): + result: Final = as_str_mapping(response_dict.get("result")) + if result is None: return "" # Direct message format (A2A spec): detect by explicit kind tag only. @@ -72,10 +79,36 @@ class A2ARequestUtils: # type that happens to include a "parts" field. if result.get("kind") == "message": return A2ARequestUtils.extract_text_from_message(result) + if result.get("kind") == "task": + return A2ARequestUtils.extract_text_from_task(result) message: Final = result.get("message", {}) return A2ARequestUtils.extract_text_from_message(message) + @staticmethod + def extract_text_from_task(task: Mapping[str, object]) -> str: + artifacts: Final = task.get("artifacts") + artifact_items: Final = _OBJECT_ITEMS.validate_python(artifacts) if isinstance(artifacts, list) else () + artifact_texts: Final = tuple( + A2ARequestUtils.extract_text_from_message(artifact) for artifact in artifact_items + ) + if any(artifact_texts): + return " ".join(text for text in artifact_texts if text) + status: Final = as_str_mapping(task.get("status")) + return A2ARequestUtils.extract_text_from_message(status.get("message")) if status is not None else "" + + @staticmethod + def extract_text_from_stream_event(event: Mapping[str, object]) -> str: + result: Final = as_str_mapping(event.get("result")) + if result is None: + return "" + if result.get("kind") == "artifact-update": + return A2ARequestUtils.extract_text_from_message(result.get("artifact")) + if result.get("kind") == "status-update": + status: Final = as_str_mapping(result.get("status")) + return A2ARequestUtils.extract_text_from_message(status.get("message")) if status is not None else "" + return A2ARequestUtils.extract_text_from_response(event) + @staticmethod def get_input_message_from_request( request: "SendMessageRequest | SendStreamingMessageRequest", @@ -128,17 +161,8 @@ class A2ARequestUtils: Returns: Tuple of (prompt_tokens, completion_tokens, total_tokens) """ - # Count input tokens. Dump the message to a dict first so extraction hits - # the dict branch — request-side parts are a2a-sdk Part RootModels whose - # kind/text live on part.root, which the object branch cannot read. This - # mirrors how the response side already works (it operates on model_dump). input_message: Final = A2ARequestUtils.get_input_message_from_request(request) - input_text: Final = A2ARequestUtils.extract_text_from_message( - input_message.model_dump(mode="json") - if has_attributes(input_message, _SupportsModelDump, "model_dump") - else input_message - ) - prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text) + prompt_tokens: Final = A2ARequestUtils.count_tokens(A2ARequestUtils.extract_text_from_message(input_message)) # Count output tokens output_text: Final = A2ARequestUtils.extract_text_from_response(response_dict) diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index cee5ce15469..253074b0e1f 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -1,5 +1,5 @@ { - "description": "Mapping of Anthropic beta headers for each provider. Keys are input header names, values are provider-specific header names (or null if unsupported). Only headers present in mapping keys with non-null values can be forwarded.", + "description": "Mapping of Anthropic beta headers for each provider. Keys are input header names, values are provider-specific header names (or null if unsupported). Only headers present in mapping keys with non-null values can be forwarded. A _header_only column lists betas that provider accepts only as the anthropic-beta HTTP header: they are forwarded there and dropped from the anthropic_beta body field. Proxies that predate these columns ignore them and read the provider column alone.", "anthropic": { "advisor-tool-2026-03-01": "advisor-tool-2026-03-01", "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", @@ -46,6 +46,7 @@ "bash_20250124": null, "code-execution-2025-08-25": "code-execution-2025-08-25", "compact-2026-01-12": null, + "compact-2026-09-04": "compact-2026-09-04", "computer-use-2025-01-24": "computer-use-2025-01-24", "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", @@ -59,6 +60,7 @@ "mcp-client-2025-11-20": "mcp-client-2025-11-20", "mcp-client-2025-04-04": "mcp-client-2025-04-04", "mcp-servers-2025-12-04": null, + "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", "output-128k-2025-02-19": null, "structured-output-2024-03-01": null, "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", @@ -202,6 +204,7 @@ "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, "mcp-servers-2025-12-04": null, + "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", "output-128k-2025-02-19": null, "structured-output-2024-03-01": null, "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", @@ -218,6 +221,10 @@ "web-search-2025-03-05": "web-search-2025-03-05", "inline-tools-2026-09-15": "inline-tools-2026-09-15" }, + "vertex_ai_header_only": { + "compact-2026-09-04": "compact-2026-09-04", + "inline-tools-2026-09-15": "inline-tools-2026-09-15" + }, "databricks": { "advisor-tool-2026-03-01": null, "advanced-tool-use-2025-11-20": null, diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 07730a0fbf8..be2ff07e088 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -30,7 +30,7 @@ from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from functools import cache from importlib.resources import files -from typing import Final, TypeAlias +from typing import Final, Literal, TypeAlias import httpx from pydantic import TypeAdapter @@ -47,6 +47,9 @@ BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]] _BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) _DECODED_JSON: Final = TypeAdapter(object) +BetaPlacement: TypeAlias = Literal["header", "body"] +_HEADER_ONLY_COLUMN_SUFFIX: Final = "_header_only" + class GetAnthropicBetaHeadersConfig: """ @@ -241,6 +244,31 @@ def get_provider_name(provider: str) -> str: return aliases.get(provider, provider) +def _beta_column(name: str) -> Mapping[str, str | None]: + column: Final[Mapping[str, str | None]] = _load_beta_headers_config().get(name, {}) + return column + + +def _header_only_column(canonical_provider: str) -> Mapping[str, str | None]: + return _beta_column(f"{canonical_provider}{_HEADER_ONLY_COLUMN_SUFFIX}") + + +def _provider_beta_mapping(provider: str, placement: BetaPlacement) -> Mapping[str, str | None]: + canonical_provider: Final = get_provider_name(provider) + column: Final = _beta_column(canonical_provider) + header_only: Final = _header_only_column(canonical_provider) + match placement: + case "header": + return {**column, **header_only} + case "body": + return {beta: name for beta, name in column.items() if beta not in header_only} + + +def get_header_only_betas(provider: str) -> frozenset[str]: + header_only: Final = _header_only_column(get_provider_name(provider)) + return frozenset(beta for beta, name in header_only.items() if name is not None) + + def has_beta_headers_config(provider: str) -> bool: config: Final = _bundled_beta_headers_config() aliases: Final = config.get("provider_aliases", {}) @@ -251,6 +279,7 @@ def has_beta_headers_config(provider: str) -> bool: def filter_and_transform_beta_headers( beta_headers: list[str], provider: str, + placement: BetaPlacement = "header", ) -> list[str]: """ Filter and transform beta headers based on provider's mapping configuration. @@ -263,6 +292,7 @@ def filter_and_transform_beta_headers( Args: beta_headers: List of Anthropic beta header values provider: Provider name (e.g., "anthropic", "bedrock", "vertex_ai") + placement: Whether the betas travel as the anthropic-beta header or the anthropic_beta body field Returns: List of filtered and transformed beta headers for the provider @@ -270,11 +300,7 @@ def filter_and_transform_beta_headers( if not beta_headers: return [] - config: Final = _load_beta_headers_config() - provider = get_provider_name(provider) - - # Get the header mapping for this provider - provider_mapping: Final = config.get(provider, {}) + provider_mapping: Final = _provider_beta_mapping(provider, placement) filtered_headers: Final[set[str]] = set() @@ -426,6 +452,7 @@ def update_request_with_filtered_beta( filtered_body_betas: Final = filter_and_transform_beta_headers( beta_headers=existing_body_betas, provider=provider, + placement="body", ) if filtered_body_betas: diff --git a/litellm/constants.py b/litellm/constants.py index c586a5f767a..c17e776f54b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -902,6 +902,7 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset( SUBSCRIPTION_BACKED_PROVIDERS: Final = frozenset({"chatgpt"}) A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"}) +A2A_COMPLETION_BRIDGE_WORKING_MESSAGE: Final = "Processing request..." LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [ "openai", @@ -2046,6 +2047,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = get_env_int("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60) AUTH_OBJECTS_TARGET: Final = "auth_objects" +HEALTH_CHECKS_TARGET: Final = "health_checks" DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = get_env_int("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600) SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ee102e82ca4..f5b9014d87c 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -623,7 +623,11 @@ def cost_per_token( ) if lyria_generation_cost is not None: return 0.0, lyria_generation_cost - speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) + speech_model_info: Final = cached_get_model_info_helper( + model=model_without_prefix, + custom_llm_provider=custom_llm_provider, + preserve_unset_costs=True, + ) cost_metric: Final = select_cost_metric_for_model(speech_model_info) prompt_cost: float = 0.0 completion_cost: float = 0.0 @@ -3084,7 +3088,7 @@ def batch_cost_calculator( batch_rates.cache_creation, batch_rates.cache_creation_rates_by_ttl, ) - elif input_cost_per_token: + elif input_cost_per_token is not None: details: Final = parse_prompt_tokens_details(usage) cache_read_tokens: Final = details["cache_hit_tokens"] cache_creation_tokens: Final = details["cache_creation_tokens"] @@ -3100,7 +3104,10 @@ def batch_cost_calculator( cache_read_cost_key: Final = get_service_tier_cost_key("cache_read_input_token_cost", None) total_prompt_cost += calculate_cost_component(model_info, cache_read_cost_key, cache_read_tokens) / 2 - cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") or input_cost_per_token + configured_cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") + cache_creation_cost: Final = ( + input_cost_per_token if configured_cache_creation_cost is None else configured_cache_creation_cost + ) total_prompt_cost += cache_creation_tokens * cache_creation_cost / 2 text_rate: Final = _batch_or_half(batch_rates.output, output_cost_per_token) image_rate: Final = _batch_or_half( diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 3b73f65c917..ae01aa8d0f7 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -218,7 +218,7 @@ def supported_decisions_providers() -> tuple[str, ...]: ) -def _provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig: +def get_decisions_provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig: provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None) provider_config: Final = ( None if provider is None else ProviderConfigManager.get_provider_decisions_config(model, provider) @@ -408,7 +408,7 @@ def _provider_ir_request( return supported -def _outbound_headers( +def decisions_outbound_headers( extra_headers: Mapping[str, str] | None, custom_llm_provider: str, kwargs: Mapping[str, object] ) -> Mapping[str, str]: forwarded: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( @@ -466,7 +466,7 @@ def prepare_call( api_base=call_api_base, api_key=api_key, ) - provider_config: Final = _provider_config(upstream_model, provider) + provider_config: Final = get_decisions_provider_config(upstream_model, provider) if not upstream_model: raise litellm.BadRequestError( message="A model name is required for the Decisions API", @@ -478,7 +478,7 @@ def prepare_call( except ValueError as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error try: - headers: Final = _outbound_headers(extra_headers, provider, kwargs) + headers: Final = decisions_outbound_headers(extra_headers, provider, kwargs) request_safety_identifier: Final = _request_safety_identifier(safety_identifier, provider_config, kwargs) request: Final = _fields_without_unsupported_media( fields, provider_config, canonical_model, provider, kwargs diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index a9fb5c29bad..dfe508f3630 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -3,8 +3,9 @@ import copy import hashlib import os import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args @@ -70,6 +71,8 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions" # field to suppress a guardrail on the direct-SDK path that never reaches the # proxy's metadata sanitizer. _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) +_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16) +_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16) _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) @@ -208,6 +211,46 @@ def guardrail_request_data_with_streaming( } +_GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object]) + + +def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: + """Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over + any plain name in either order, else its first plain name stays""" + entries: Final = tuple(chain.from_iterable(entry_lists)) + settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries) + named_with_settings: Final = frozenset[str]().union(*settings_names) + return [ + entry + for index, entry in enumerate(entries) + if not _superseded_guardrail_entry(entries, settings_names, named_with_settings, index) + ] + + +def guardrail_entry_names(entry: object) -> tuple[str, ...]: + if isinstance(entry, str): + return (entry,) + if isinstance(entry, dict): + return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return () + + +def _settings_entry_names(entry: object) -> frozenset[str]: + return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset() + + +def _superseded_guardrail_entry( + entries: tuple[object, ...], + settings_names: tuple[frozenset[str], ...], + named_with_settings: frozenset[str], + index: int, +) -> bool: + entry: Final = entries[index] + if isinstance(entry, str): + return entry in named_with_settings or entry in entries[:index] + return bool(settings_names[index]) and settings_names[index] <= frozenset[str]().union(*settings_names[:index]) + + def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool: if event_type in _REALTIME_STREAMING_HOOKS: return True @@ -219,6 +262,25 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = ) +def _str_collection(value: object) -> frozenset[str]: + if not isinstance(value, (list, tuple, set, frozenset)): + return frozenset() + items: Final = cast("Iterable[object]", value) # cast-ok: a metadata collection of guardrail names + return frozenset(item for item in items if isinstance(item, str)) + + +def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> bool: + if marker is None: + return False + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list) and marker in executed: + return True + return False + + _REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) @@ -828,13 +890,16 @@ class CustomGuardrail(CustomLogger): return False - def _pre_call_marker(self) -> str | None: + def _marker(self, token: str) -> str | None: name: Final = self.guardrail_name if not name: return None - return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" + return f"{token}:{name}" - def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None: + def _pre_call_marker(self, *, raised: bool = False) -> str | None: + return self._marker(_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN) + + def mark_pre_call_hook_ran(self, data: dict[str, object], *, raised: bool = False) -> None: """ Record that this guardrail's ``async_pre_call_hook`` already ran for this request, so the deployment-level hook does not run it a second time. @@ -843,8 +908,11 @@ class CustomGuardrail(CustomLogger): router later spreads a deployment's model-level ``guardrails`` into the top-level request kwargs, which would otherwise re-trigger the same hook from ``async_pre_call_deployment_hook``. + + With ``raised``, records instead that the hook raised as a policy pipeline + step the pipeline then moved past. The deployment-level hook still runs it. """ - marker: Final = self._pre_call_marker() + marker: Final = self._pre_call_marker(raised=raised) if marker is None: return for meta_key in ("metadata", "litellm_metadata"): @@ -859,17 +927,57 @@ class CustomGuardrail(CustomLogger): return data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} - def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool: - marker: Final = self._pre_call_marker() - if marker is None: - return False - for meta_key in ("metadata", "litellm_metadata"): - meta = data.get(meta_key) - if isinstance(meta, dict): - executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) - if isinstance(executed, list) and marker in executed: - return True - return False + def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool: + """Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided.""" + return self._pre_call_hook_already_ran(data) or self._pre_call_hook_already_ran(data, raised=True) + + def _pre_call_hook_already_ran(self, data: Mapping[str, object], *, raised: bool = False) -> bool: + return _request_carries_marker(data, self._pre_call_marker(raised=raised)) + + def deployment_hook_ran(self, data: Mapping[str, object]) -> bool: + """Whether a deployment-level hook ran this guardrail as a policy pipeline step of this request.""" + return _request_carries_marker(data, self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN)) + + def _record_deployment_hook_run( + self, + data: dict[str, object], # mutable-ok: the request payload the run is recorded on + ) -> None: + """ + Record that a deployment-level hook ran this guardrail when it is a policy pipeline step of + this request, and name it in the applied guardrails when the request's metadata lists it. + """ + name: Final = self.guardrail_name + marker: Final = self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN) + buckets: Final = tuple( + cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed + for bucket in (data.get("metadata"), data.get("litellm_metadata")) + if isinstance(bucket, dict) + ) + policy_state: Final = next( + (bucket for bucket in buckets if name in _str_collection(bucket.get("_pipeline_managed_guardrails"))), + None, + ) + if marker is None or policy_state is None: + return + recorded: Final = policy_state.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if not isinstance(recorded, list): + policy_state[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] + elif marker not in recorded: + cast("list[object]", recorded).append(marker) # cast-ok: markers are only compared as strings + if not any(name in _str_collection(bucket.get("guardrails")) for bucket in buckets): + return + named_lists: Final = tuple( + cast("list[object]", bucket["applied_guardrails"]) # cast-ok: entries are only compared with names + for bucket in buckets + if isinstance(bucket.get("applied_guardrails"), list) + ) + for applied in named_lists: + if name not in applied: + applied.append(name) + if named_lists: + return + metadata_bucket: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str keys + metadata_bucket["applied_guardrails"] = [name] @classmethod def uses_apply_guardrail_interface(cls) -> bool: @@ -918,12 +1026,15 @@ class CustomGuardrail(CustomLogger): target: Final = self._deployment_hook_target() if target is not self: kwargs["guardrail_to_apply"] = self - result: Final = await target.async_pre_call_hook( - user_api_key_dict=_user_api_key_auth_from_request(kwargs), - cache=dc, - data=kwargs, - call_type="completion" if call_type == CallTypes.completion else "acompletion", - ) + try: + result: Final = await target.async_pre_call_hook( + user_api_key_dict=_user_api_key_auth_from_request(kwargs), + cache=dc, + data=kwargs, + call_type="completion" if call_type == CallTypes.completion else "acompletion", + ) + finally: + self._record_deployment_hook_run(kwargs) if result is not None and isinstance(result, dict): result_messages: Final = result.get("messages") @@ -1007,6 +1118,9 @@ class CustomGuardrail(CustomLogger): finally: if target is not self: request_data.pop("guardrail_to_apply", None) + self._record_deployment_hook_run( + cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed + ) if not self._is_valid_response_type(result): return None diff --git a/litellm/integrations/otel/routing.py b/litellm/integrations/otel/routing.py index 82dd6237edc..339775855be 100644 --- a/litellm/integrations/otel/routing.py +++ b/litellm/integrations/otel/routing.py @@ -21,6 +21,11 @@ _ROUTING_SCALAR_FIELDS: Final = frozenset( "classifier_cost", "classifier_failure_reason", "classifier_error_type", + "classifier_input_budget_tokens", + "classifier_input_tokens", + "classifier_input_count_method", + "classifier_input_truncated", + "classifier_input_budget_status", "classifier_confidence", "classifier_primary_rule", "classifier_capability_boundary", diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index e6318525b88..42fe0837c29 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -66,7 +66,7 @@ from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import EndUserRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository -from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata +from litellm.router_api import deployment_metadata from litellm.types.guardrails import GuardrailEventHooks from litellm.types.integrations.prometheus import * from litellm.types.integrations.prometheus import ( @@ -82,6 +82,7 @@ from litellm.types.proxy.carried_budget_state import ( UserBudgetSnapshot, ) from litellm.types.utils import ( + CallTypes, StandardLoggingGuardrailInformation, StandardLoggingPayload, StandardLoggingZeroCostDiagnostic, @@ -279,6 +280,11 @@ def _logged_fields(value: object) -> Mapping[str, object]: return _NO_LOGGED_FIELDS +_GOOGLE_PASS_THROUGH_PROVIDERS: Final = frozenset( + (litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value) +) + + def _is_llm_provider(name: str) -> bool: routable: Final = ( name in litellm.provider_list or JSONProviderRegistry.exists(name) or is_registered_custom_provider(name) @@ -2942,6 +2948,22 @@ class PrometheusLogger(CustomLogger): return False + @staticmethod + def _google_pass_through_provider(request_data: Mapping[str, object]) -> str | None: + declared_provider: Final = request_data.get("custom_llm_provider") + if request_data.get("call_type") != CallTypes.pass_through.value or not isinstance(declared_provider, str): + return None + return declared_provider if declared_provider in _GOOGLE_PASS_THROUGH_PROVIDERS else None + + @staticmethod + def _failed_request_requested_model(request_data: Mapping[str, object]) -> object: + requested_model: Final = request_data.get("model", "") + if PrometheusLogger._google_pass_through_provider(request_data) is None: + return requested_model + request_metadata: Final = _logged_fields(_logged_fields(request_data.get("litellm_params")).get("metadata")) + model_group: Final = request_metadata.get("model_group") + return model_group if isinstance(model_group, str) and model_group else requested_model + @staticmethod def _extract_api_provider_from_request_data(request_data: dict) -> str | None: """ @@ -2950,22 +2972,27 @@ class PrometheusLogger(CustomLogger): A request can fail before a deployment is resolved, so the provider is not always known. Prefer the resolved ``custom_llm_provider`` on ``litellm_params``, then any provider recovered onto a partial - ``standard_logging_object`` (e.g. a stream that broke mid-flight), and - finally infer it from the requested model name (e.g. ``gpt-4o-mini`` -> - ``openai``) since the proxy's failure ``request_data`` usually carries - only the client-supplied model. Return ``None`` when it cannot be - determined so the label emits empty rather than a guess. + ``standard_logging_object`` (e.g. a stream that broke mid-flight). On a + Gemini or Vertex pass-through failure only, next take the provider the + failure payload names (``gemini`` or ``vertex_ai``), so every other + pass-through route keeps its earlier label. Finally infer it from the + requested model name (e.g. ``gpt-4o-mini`` -> ``openai``) since the + proxy's failure ``request_data`` usually carries only the client-supplied + model. Return ``None`` when it cannot be determined so the label emits + empty rather than a guess. """ litellm_params: Final = request_data.get("litellm_params") or {} - provider = litellm_params.get("custom_llm_provider") - if provider: - return provider standard_logging_object: Final = request_data.get("standard_logging_object") or {} - provider = standard_logging_object.get("custom_llm_provider") - if provider: - return provider + resolved_provider: Final = litellm_params.get("custom_llm_provider") or standard_logging_object.get( + "custom_llm_provider" + ) + if resolved_provider: + return resolved_provider + google_provider: Final = PrometheusLogger._google_pass_through_provider(request_data) + if google_provider is not None: + return google_provider model: Final = litellm_params.get("model") or request_data.get("model") - if not model: + if not isinstance(model, str) or not model: return None try: return litellm.get_llm_provider(model=model)[1] or None @@ -2995,8 +3022,8 @@ class PrometheusLogger(CustomLogger): ).get("custom_llm_provider") if isinstance(provider, str) and provider: return provider - deployment_metadata: Final = find_deployment_metadata(litellm_params) - return _deployment_model_provider(deployment_metadata.get("deployment") if deployment_metadata else None) + picked_deployment: Final = deployment_metadata(litellm_params) + return _deployment_model_provider(picked_deployment.get("deployment") if picked_deployment else None) @staticmethod def _extract_api_provider_from_exception(exception: Exception) -> str | None: @@ -3052,7 +3079,7 @@ class PrometheusLogger(CustomLogger): team_alias=user_api_key_dict.team_alias, org_id=user_api_key_dict.org_id, org_alias=user_api_key_dict.organization_alias, - requested_model=_bounded_requested_model_label(request_data.get("model", "")), + requested_model=_bounded_requested_model_label(self._failed_request_requested_model(request_data)), status_code=str(status_code), exception_status=str(status_code), exception_class=self._get_exception_class_name(original_exception), @@ -5182,7 +5209,7 @@ def _tag_matches_wildcard_configured_pattern(tags: Sequence[str], configured_tag """ import re - from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter pattern_router: Final = PatternMatchRouter() regex_pattern: Final = pattern_router.pattern_to_regex(configured_tag) diff --git a/litellm/router_utils/access_windows.py b/litellm/litellm_core_utils/access_windows.py similarity index 100% rename from litellm/router_utils/access_windows.py rename to litellm/litellm_core_utils/access_windows.py diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/litellm_core_utils/add_retry_fallback_headers.py similarity index 100% rename from litellm/router_utils/add_retry_fallback_headers.py rename to litellm/litellm_core_utils/add_retry_fallback_headers.py diff --git a/litellm/router_utils/clientside_credential_handler.py b/litellm/litellm_core_utils/clientside_credential_handler.py similarity index 100% rename from litellm/router_utils/clientside_credential_handler.py rename to litellm/litellm_core_utils/clientside_credential_handler.py diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 28f45d2bb2c..20e10f63f83 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -9,12 +9,13 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import TypeIs # noqa: TID251 # narrows a duck-typed value to the Protocol its caller reads import litellm from litellm._logging import verbose_logger from litellm.constants import KUBERNETES_POD_ROUTING_KEY +from litellm.litellm_core_utils.hidden_params import set_extra_attribute from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason @@ -710,9 +711,9 @@ def preserve_upstream_non_openai_attributes( _assign_attributes(model_response, original_chunk.model_dump(include=frozenset(candidate_keys)).items()) -def _assign_attributes(target: object, values: Iterable[tuple[str, object]]) -> None: +def _assign_attributes(target: BaseModel, values: Iterable[tuple[str, object]]) -> None: for key, value in values: - setattr(target, key, value) + set_extra_attribute(target, key, value) def safe_deep_copy(data): @@ -982,25 +983,36 @@ def _usage_field(usage: object, name: str) -> object: return getattr(usage, name, None) -def _cost_number(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, int | float): +def finite_cost(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(value): return None return float(value) +def single_log_line(value: object) -> str: + return str(value).replace("\r", "").replace("\n", "") + + +def finite_reported_cost(value: object) -> float | None: + cost: Final = finite_cost(value) + if cost is None and isinstance(value, float): + verbose_logger.warning("Ignoring a non-finite provider-reported cost %s", value) + return cost + + def _reported_cost_value(value: object) -> float | None: """A reported cost is a number, or a breakdown whose total lives under ``total_cost``.""" raw: Final[object] = value if isinstance(value, Mapping): - return _cost_number(_usage_field(raw, "total_cost")) - return _cost_number(value) + return finite_reported_cost(_usage_field(raw, "total_cost")) + return finite_reported_cost(value) def response_cost_in_hidden_params(response: _CarriesHiddenParams) -> float | None: additional_headers: Final[object] = response.hidden_params.get("additional_headers") if not isinstance(additional_headers, Mapping): return None - return _cost_number(_USAGE_BLOCK_ADAPTER.validate_python(additional_headers).get(RESPONSE_COST_HEADER)) + return finite_reported_cost(_USAGE_BLOCK_ADAPTER.validate_python(additional_headers).get(RESPONSE_COST_HEADER)) def override_stamped_response_cost(response: _CarriesHiddenParams, cost: float) -> None: diff --git a/litellm/router_utils/deployment_pricing.py b/litellm/litellm_core_utils/deployment_pricing.py similarity index 100% rename from litellm/router_utils/deployment_pricing.py rename to litellm/litellm_core_utils/deployment_pricing.py diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 5a27f5446ad..85c232ad3f4 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.bug_report import ( ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.secret_redaction import redact_string as redact_secret_string +from litellm.llms.bedrock.common_utils import bedrock_model_reaches_converse_by_default from litellm.llms.bigmodel.common_utils import raise_bigmodel_business_error from litellm.types.utils import LlmProviders @@ -97,6 +98,7 @@ class ExceptionCheckers: known_exception_substrings: Final = [ "exceed context limit", "this model's maximum context length is", + "requested token count exceeds the model's maximum context length", "string too long. expected a string with maximum length", "model's maximum context limit", "is longer than the model's context length", @@ -938,13 +940,22 @@ _BEDROCK_MANTLE_CONTEXT_WINDOW_GENERIC_MESSAGE: Final = ( ) -def _bedrock_mantle_context_window_response(original_exception: _ProviderHTTPException) -> httpx.Response: +def _bedrock_response_with_status(original_exception: _ProviderHTTPException, status_code: int) -> httpx.Response: original_response: Final = getattr(original_exception, "response", None) if not isinstance(original_response, httpx.Response) or getattr(original_response, "_request", None) is None: - return httpx.Response(status_code=400, request=httpx.Request(method="POST", url="https://api.openai.com/v1/")) - if original_response.status_code == 400: + return httpx.Response( + status_code=status_code, request=httpx.Request(method="POST", url="https://api.openai.com/v1/") + ) + if original_response.status_code == status_code: return original_response - return httpx.Response(status_code=400, headers=original_response.headers, request=original_response.request) + return httpx.Response(status_code=status_code, headers=original_response.headers, request=original_response.request) + + +def _came_from_bedrock_converse(original_exception: _ProviderHTTPException) -> bool: + original_response: Final = getattr(original_exception, "response", None) + if not isinstance(original_response, httpx.Response) or getattr(original_response, "_request", None) is None: + return False + return original_response.request.url.path.endswith(("/converse", "/converse-stream")) def _get_bedrock_mantle_context_window_message(error_str: str) -> str | None: @@ -980,7 +991,7 @@ def _map_bedrock_exception( message=mantle_context_window_message, model=model, llm_provider=custom_llm_provider, - response=_bedrock_mantle_context_window_response(original_exception), + response=_bedrock_response_with_status(original_exception, 400), ) if ( "too many tokens" in error_str @@ -1003,6 +1014,18 @@ def _map_bedrock_exception( llm_provider="bedrock", response=getattr(original_exception, "response", None), ) + elif ( + "The provided model identifier is invalid" in error_str + and bool(model) + and _came_from_bedrock_converse(original_exception) + and bedrock_model_reaches_converse_by_default(model) + ): + raise NotFoundError( + message=f"BedrockException - {error_str}", + model=model, + llm_provider="bedrock", + response=_bedrock_response_with_status(original_exception, 404), + ) elif "Malformed input request" in error_str: raise BadRequestError( message=f"BedrockException - {error_str}", diff --git a/litellm/litellm_core_utils/fallback_utils.py b/litellm/litellm_core_utils/fallback_utils.py index cc36f47b516..87888311508 100644 --- a/litellm/litellm_core_utils/fallback_utils.py +++ b/litellm/litellm_core_utils/fallback_utils.py @@ -3,13 +3,13 @@ from typing import Final import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, +) from litellm.litellm_core_utils.core_helpers import ( filter_internal_params, safe_deep_copy, ) -from litellm.router_utils.add_retry_fallback_headers import ( - add_fallback_headers_to_response, -) from .asyncify import run_async_function diff --git a/litellm/router_utils/get_retry_from_policy.py b/litellm/litellm_core_utils/get_retry_from_policy.py similarity index 100% rename from litellm/router_utils/get_retry_from_policy.py rename to litellm/litellm_core_utils/get_retry_from_policy.py diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index e40299a218f..2c7d00552b4 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -119,7 +119,7 @@ def _is_route_segment(custom_llm_provider: str, segment: str, model_after_segmen def wildcard_health_check_models(wildcard: str, custom_llm_provider: str) -> tuple[str, ...]: from litellm.litellm_core_utils.llm_request_utils import rank_chat_models_from_llm_provider - from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter provider_prefix: Final = f"{custom_llm_provider}/" segments: Final = wildcard.removeprefix(provider_prefix).split("/") diff --git a/litellm/litellm_core_utils/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index a74337d2549..c04d7614b90 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -1,8 +1,9 @@ from collections.abc import Iterator, Mapping, MutableMapping +from functools import cache, lru_cache from types import MappingProxyType from typing import Final, Protocol, cast -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, create_model from litellm.types.llms.base import HiddenParams @@ -60,6 +61,38 @@ def stored_attribute(obj: object, name: str) -> object: return extras.get(name) +@cache +def _assigned_extras_join_fields_set() -> bool: + probe: Final = create_model("_ExtraAssignmentProbe", __config__=ConfigDict(extra="allow"))() + setattr(probe, "assigned", None) + return "assigned" in probe.model_fields_set + + +@lru_cache(maxsize=1024) +def _writes_plain_extra(model_class: type[BaseModel], name: str) -> bool: + config: Final = model_class.model_config + return ( + model_class.__setattr__ is BaseModel.__setattr__ + and config.get("extra") == "allow" + and not config.get("frozen") + and not config.get("validate_assignment") + and not name.startswith("_") + and name not in model_class.__pydantic_fields__ + and name not in model_class.__class_vars__ + and not hasattr(model_class, name) + and _assigned_extras_join_fields_set() + ) + + +def set_extra_attribute(model: BaseModel, name: str, value: object) -> None: + extras: Final = model.__pydantic_extra__ + if extras is None or not _writes_plain_extra(type(model), name): + setattr(model, name, value) + return + extras[name] = value + model.__pydantic_fields_set__.add(name) + + def _get_hidden_params_storage(obj: object) -> object | None: if isinstance(obj, dict): return obj.get(_HIDDEN_PARAMS_ATTR) diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index f3dbd612bee..453d4358b35 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -35,6 +35,12 @@ from litellm.types.utils import ( BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" +OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" +ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" +OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key" +IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY}) + MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups" """Where auth records the model access groups that authorized the request, for the spend writer. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4dd33fa5565..924603dc7e8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -76,6 +76,7 @@ from litellm.litellm_core_utils.classifier_logging import ( is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + finite_cost, get_provider_response_headers_from_hidden_params, is_expected_client_error, native_messages_reported_cost, @@ -84,6 +85,7 @@ from litellm.litellm_core_utils.core_helpers import ( proxy_stamped_used_client_oauth_token, reconstruct_model_name, set_response_cost_in_hidden_params, + single_log_line, ) from litellm.litellm_core_utils.error_normalization import normalize_error from litellm.litellm_core_utils.get_litellm_params import get_litellm_params @@ -163,7 +165,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.mcp import MCPPostCallResponseObject from litellm.types.prompts.init_prompts import PromptSpec -from litellm.types.rerank import RerankResponse +from litellm.types.rerank import RerankResponse, get_rerank_usage from litellm.types.router import RouterRateLimitError, SilentExperimentFlag from litellm.types.utils import ( DEPLOYMENT_SCOPED_PRICING_FIELDS, @@ -353,6 +355,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_RESPONSE_MAPPING: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( ("user_api_key_auth", "user_api_key_budget_reservation") @@ -567,7 +570,13 @@ def _resolve_mantle_region_for_cost( _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) +_SERVED_ANTHROPIC_BODY: Final = TypeAdapter(dict[str, object]) _REDACTED_RESULT: Final = TypeAdapter(object) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _without_guardrail_information(metadata: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in metadata.items() if key != _GUARDRAIL_INFORMATION_KEY} def _provider_response_id(source: object) -> str | None: @@ -2088,7 +2097,7 @@ class Logging(LiteLLMLoggingBaseClass): litellm_model_name=litellm_model_name, router_model_id=router_model_id or hidden_params.get("model_id"), ) - return hidden_params["response_cost"] + return self._recorded_response_cost(hidden_params["response_cost"]) elif router_model_id is None and "model_id" in hidden_params: # use model_id if not already set router_model_id = hidden_params["model_id"] @@ -2188,7 +2197,7 @@ class Logging(LiteLLMLoggingBaseClass): litellm_model_name=litellm_model_name, router_model_id=router_model_id, ) - return total_response_cost + return self._recorded_response_cost(total_response_cost) except Exception as e: # error calculating cost debug_info = StandardLoggingModelCostFailureDebugInformation( error_str=str(e), @@ -2214,6 +2223,15 @@ class Logging(LiteLLMLoggingBaseClass): _response_cost_calculator = response_cost_calculator + def _recorded_response_cost(self, cost: float | None) -> float | None: + if cost is None or finite_cost(cost) is not None: + return cost + verbose_logger.warning( + "Recording a non-finite response_cost as 0 for litellm_call_id=%s", single_log_line(self.litellm_call_id) + ) + self.cost_breakdown = None + return 0.0 + def uncached_price(self, result: object) -> UncachedPrice: previous_cost_breakdown: Final = self.cost_breakdown previous_billed_token_rates: Final = self.billed_token_rates @@ -2566,6 +2584,23 @@ class Logging(LiteLLMLoggingBaseClass): attempt.model_call_details = dict(self.model_call_details) return attempt + def guardrail_scoped_copy(self) -> "Logging": + request_params: Final = _LITELLM_PARAMS.validate_python( + self.litellm_params or {} # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped + ) + litellm_params: Final = { + **request_params, + **{ + key: _without_guardrail_information(_METADATA_MAPPING.validate_python(value)) + for key, value in request_params.items() + if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) + }, + } + scoped: Final = copy.copy(self) + scoped.litellm_params = litellm_params + scoped.model_call_details = {**self.model_call_details, "litellm_params": litellm_params} + return scoped + def dispatch_attempt_failure_handlers( self, exception: Exception, traceback_exception: str ) -> Coroutine[None, None, None]: @@ -5034,13 +5069,15 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.warning("LiteLLM: the anthropic_messages stream assembled no response, logging an empty one") return litellm.ModelResponse(model=self.model) else: - from litellm.types.llms.anthropic import AnthropicResponse - - pydantic_result: Final = AnthropicResponse.model_validate({"stop_sequence": None, **(served_body or {})}) - import httpx - + anthropic_body: Final = ( + served_body + if served_body is not None + else _SERVED_ANTHROPIC_BODY.validate_python( + result.model_dump() if isinstance(result, BaseModel) else result + ) + ) result = litellm.AnthropicConfig().transform_parsed_response( - completion_response=pydantic_result.model_dump(), + completion_response={"stop_sequence": None, **anthropic_body}, raw_response=httpx.Response( status_code=200, headers={}, @@ -6781,7 +6818,11 @@ class StandardLoggingPayloadSetup: total_tokens=0, ) - usage: Final = response_obj.get("usage", None) or {} + usage: Final = response_obj.get("usage", None) + if usage is None: + rerank_usage: Final = get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) + if rerank_usage is not None: + return Usage(**rerank_usage) if usage is None or (not isinstance(usage, dict) and not isinstance(usage, Usage)): return Usage( prompt_tokens=0, @@ -6817,7 +6858,7 @@ class StandardLoggingPayloadSetup: return _empty _raw: Final = response_obj.get("usage", None) if _raw is None: - return _empty + return dict(get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) or _empty) if isinstance(_raw, ResponseAPIUsage): return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump() if isinstance(_raw, dict): @@ -7560,7 +7601,15 @@ def get_standard_logging_object_payload( ) custom_pricing: Final = False if spilled_over else logging_obj.custom_pricing_for(init_response_obj) raw_response_cost: Final = kwargs.get("response_cost") - llm_response_cost: Final[float] = raw_response_cost or 0.0 + numeric_response_cost: Final = raw_response_cost if isinstance(raw_response_cost, int | float) else None + priced_response_cost: Final = finite_cost(numeric_response_cost) + if numeric_response_cost is not None and priced_response_cost is None: + verbose_logger.warning( + "Logging a non-finite response_cost=%s as 0 for litellm_call_id=%s", + numeric_response_cost, + single_log_line(logging_obj.litellm_call_id), + ) + llm_response_cost: Final[float] = priced_response_cost or 0.0 guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information")) finalized_cost: Final = _finalized_cost_of(logging_obj) response_cost: Final[float] = ( diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 045869e4547..54213f40ac0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -208,11 +208,14 @@ class StandardBuiltInToolCostTracking: usage=usage, response_object=response_object ) + response_model: Final = _output_item_field(response_object, "model") if model_info is not None and resolved_usage is not None and resolved_provider is not None: result: Final = get_cost_for_web_search_request( custom_llm_provider=resolved_provider, usage=resolved_usage, model_info=model_info, + response_model=response_model if isinstance(response_model, str) else None, + lookup_model_info=StandardBuiltInToolCostTracking._safe_get_model_info, ) if result is not None: return result diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 404dfab0a6d..50c26f4d4ca 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -28,6 +28,7 @@ from litellm.types.utils import ( DataResidency, ImageResponse, ModelInfo, + ModelInfoBase, PassthroughCallTypes, PromptTokensDetailsWrapper, ServiceTier, @@ -197,15 +198,15 @@ def get_billable_input_tokens(usage: Usage) -> int: def select_cost_metric_for_model( - model_info: ModelInfo, + model_info: ModelInfoBase, ) -> Literal["cost_per_character", "cost_per_token"]: """ Select 'cost_per_character' if model_info has 'input_cost_per_character' Select 'cost_per_token' if model_info has 'input_cost_per_token' """ - if model_info.get("input_cost_per_character"): + if model_info.get("input_cost_per_character") is not None: return "cost_per_character" - elif model_info.get("input_cost_per_token"): + elif model_info.get("input_cost_per_token") is not None: return "cost_per_token" else: raise ValueError( diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index 9c2ec0b8984..84f161e0c62 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -1,4 +1,3 @@ -import datetime from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass from typing import Final, cast @@ -9,6 +8,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.model_retirement import is_past_deprecation_date def _form_field_value(value: object) -> str: @@ -168,7 +168,9 @@ class RankedChatModel: def rank_chat_models_from_llm_provider(custom_llm_provider: str) -> tuple[RankedChatModel, ...]: - known_models: Final[Iterable[str]] = litellm.models_by_provider.get(custom_llm_provider, ()) + from litellm.utils import bundled_provider_models + + known_models: Final = bundled_provider_models(custom_llm_provider) ranked_models: Final = ( RankedChatModel( model=model, route_rank=route_rank, cost=cost, is_past_deprecation_date=is_past_deprecation_date @@ -187,18 +189,11 @@ _COST_MAP_DEPRECATION: Final = TypeAdapter(_CostMapDeprecation) def _is_past_deprecation_date(cost_map_key: str) -> bool: - deprecation_date: Final = _COST_MAP_DEPRECATION.validate_python( - litellm.model_cost.get(cost_map_key) or {} # pyright: ignore[reportUnknownMemberType] # untyped cost map JSON - ).get("deprecation_date") - if not isinstance(deprecation_date, str): - return False - try: - return ( - datetime.date.fromisoformat(deprecation_date) - <= datetime.datetime.now(datetime.timezone.utc).astimezone().date() + return is_past_deprecation_date( + _COST_MAP_DEPRECATION.validate_python( + litellm.model_cost.get(cost_map_key) or {} # pyright: ignore[reportUnknownMemberType] # untyped cost map JSON ) - except ValueError: - return False + ) def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float, bool]]: diff --git a/litellm/litellm_core_utils/model_retirement.py b/litellm/litellm_core_utils/model_retirement.py new file mode 100644 index 00000000000..3bfb834733d --- /dev/null +++ b/litellm/litellm_core_utils/model_retirement.py @@ -0,0 +1,19 @@ +import datetime +from collections.abc import Mapping +from typing import Final + +PROVIDERS_HIDING_RETIRED_MODELS: Final = frozenset({"aiand", "baseten", "together_ai"}) + + +def is_past_deprecation_date(value: Mapping[str, object]) -> bool: + deprecation_date: Final = value.get("deprecation_date") + if not isinstance(deprecation_date, str): + return False + try: + return datetime.date.fromisoformat(deprecation_date) <= datetime.datetime.now(datetime.timezone.utc).date() + except ValueError: + return False + + +def is_hidden_as_retired(value: Mapping[str, object]) -> bool: + return value.get("litellm_provider") in PROVIDERS_HIDING_RETIRED_MODELS and is_past_deprecation_date(value) diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/litellm_core_utils/pattern_match_deployments.py similarity index 98% rename from litellm/router_utils/pattern_match_deployments.py rename to litellm/litellm_core_utils/pattern_match_deployments.py index 4dd217cf59b..2b8d1e9c11b 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/litellm_core_utils/pattern_match_deployments.py @@ -279,7 +279,9 @@ class PatternMatchRouter: PatternMatchRouter.set_deployment_model_name( matched_pattern=pattern_match, litellm_deployment_litellm_model=deployment["litellm_params"]["model"] ) - for deployment in self.patterns[pattern] + for deployment in self.patterns.get( # pyright: ignore[reportUnknownMemberType] # untyped deployment dicts + pattern, () + ) ) diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index c831e92c09d..767eab962c8 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -2,7 +2,12 @@ import re from collections.abc import Mapping from typing import Final -from litellm.constants import SESSION_ID_FROM_TRACE_ID_PARAM_KEY, SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + SESSION_ID_FROM_TRACE_ID_PARAM_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping _HTTP_HEADER_NAME_PATTERN: Final = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") _FORBIDDEN_AFFINITY_HEADERS: Final = frozenset( @@ -98,6 +103,41 @@ def get_stable_session_id(litellm_params: object | None) -> str | None: return _non_generated_litellm_session_id(litellm_params) or _metadata_session_id(litellm_params) +def get_caller_session_id(litellm_params: object | None) -> str | None: + """Resolve a caller conversation without accepting a generated session or a promoted trace id.""" + if litellm_params is None: + return None + metadata_values: Final = _metadata_values(litellm_params) + preserved_session: Final = next( + ( + value + for metadata in metadata_values + if isinstance(value := _get_value(metadata, CALLER_SESSION_ID_METADATA_KEY), str) and value + ), + None, + ) + if preserved_session: + return preserved_session + direct_session: Final = _get_value(litellm_params, "session_id") + if isinstance(direct_session, str) and direct_session: + return direct_session + + proxy_request: Final = _get_value(litellm_params, "proxy_server_request") + headers: Final = as_str_mapping(_get_value(proxy_request, "headers")) + trace_header: Final = next( + (value for key, value in (headers or {}).items() if key.lower() == "x-litellm-trace-id"), None + ) + candidates: Final = ( + get_caller_litellm_session_id(litellm_params), + *( + _get_value(metadata, "session_id") + for metadata in metadata_values + if not _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) + ), + ) + return next((value for value in candidates if isinstance(value, str) and value and value != trace_header), None) + + def add_provider_affinity_header( headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers diff --git a/litellm/router_utils/ptu_shares.py b/litellm/litellm_core_utils/ptu_shares.py similarity index 99% rename from litellm/router_utils/ptu_shares.py rename to litellm/litellm_core_utils/ptu_shares.py index 56c11b41cbc..d56c9f2ea14 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/litellm_core_utils/ptu_shares.py @@ -12,13 +12,13 @@ from datetime import date from typing import Final, Generic, TypeAlias, TypeVar from litellm.litellm_core_utils.ptu_pricing import parsed_ptu_shares, ptu_terms +from litellm.litellm_core_utils.routing_common_utils import get_request_team_id, team_may_use_deployment from litellm.llms.azure.ptu_capacity import ( PTUCapacity, deployment_ptu_capacity, is_azure_deployment, is_str_keyed_mapping, ) -from litellm.router_utils.common_utils import get_request_team_id, team_may_use_deployment _DeploymentT: Final = TypeVar("_DeploymentT", bound=Mapping[str, object]) WildcardRoute: TypeAlias = Callable[[str], Sequence[Mapping[str, object]]] diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/litellm_core_utils/reasoning_effort_capability.py similarity index 99% rename from litellm/router_utils/reasoning_effort_capability.py rename to litellm/litellm_core_utils/reasoning_effort_capability.py index 3563fcd96ae..c315a31b65d 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/litellm_core_utils/reasoning_effort_capability.py @@ -38,7 +38,7 @@ from litellm.litellm_core_utils.reasoning_effort_strength import ( ReasoningEffortRounding, nearest_declared_reasoning_effort, ) -from litellm.litellm_core_utils.reasoning_effort_utils import reasoning_effort_from_thinking +from litellm.litellm_core_utils.reasoning_effort_utils import reasoning_effort_from_thinking_silently from litellm.llms.anthropic.common_utils import AnthropicModelInfo, validated_object_mapping from litellm.types.llms.anthropic import AnthropicThinkingParam from litellm.types.llms.openai import REASONING_EFFORT @@ -143,7 +143,7 @@ def with_adaptive_thinking_as_effort( params: Mapping[str, object], *, model: str, custom_llm_provider: str, drop_params: bool ) -> Mapping[str, object]: translated: Final = dict(params) - AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently( model=model, optional_params=translated, custom_llm_provider=custom_llm_provider, drop_params=drop_params ) return translated @@ -151,7 +151,7 @@ def with_adaptive_thinking_as_effort( def thinking_budget_effort(thinking: object, *, model: str, custom_llm_provider: str, drop_params: bool) -> str | None: if custom_llm_provider in _PROVIDERS_SENDING_THINKING_AS_EFFORT: - return reasoning_effort_from_thinking(thinking, drop_params=drop_params) + return reasoning_effort_from_thinking_silently(thinking, drop_params=drop_params) adaptive_config: Final = as_str_mapping( with_adaptive_thinking_as_effort( {"thinking": thinking}, model=model, custom_llm_provider=custom_llm_provider, drop_params=drop_params @@ -162,7 +162,7 @@ def thinking_budget_effort(thinking: object, *, model: str, custom_llm_provider: return adaptive_effort if isinstance(adaptive_effort, str) else None if "thinking" in (get_supported_openai_params(model=model, custom_llm_provider=custom_llm_provider) or ()): return None - return reasoning_effort_from_thinking(thinking, drop_params=drop_params) + return reasoning_effort_from_thinking_silently(thinking, drop_params=drop_params) def provider_reads_thinking(thinking: object, *, model: str, custom_llm_provider: str, drop_params: bool) -> bool: diff --git a/litellm/litellm_core_utils/reasoning_effort_utils.py b/litellm/litellm_core_utils/reasoning_effort_utils.py index acd20aeb37d..e9c94e32f8f 100644 --- a/litellm/litellm_core_utils/reasoning_effort_utils.py +++ b/litellm/litellm_core_utils/reasoning_effort_utils.py @@ -1,6 +1,12 @@ +import reprlib +from collections.abc import Collection, Mapping +from dataclasses import dataclass from typing import Final, Literal +from pydantic import TypeAdapter + import litellm +from litellm._logging import verbose_logger from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, @@ -11,6 +17,54 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping OpenAIStyleReasoningEffort = Literal["minimal", "low", "medium", "high"] +_DROP_WARNING: Final = "Dropping %s=%s for model=%r because %s" +_DROP_PARAMS_DROP_WARNING: Final = "Dropping %s=%s for model=%r (drop_params=True) because %s" +_NON_NUMERIC_BUDGET_REASON: Final = "it is not a number. Send budget_tokens as an integer to set a reasoning level." +_NON_STRING_EFFORT_REASON: Final = ( + "it is not a string. Send reasoning_effort as a level name such as 'high' to set a reasoning level." +) +_REASONING_PARAMS: Final = ("reasoning", "reasoning_effort", "thinking") +OPAQUE_VALUE: Final = TypeAdapter(object) + + +def warn_dropped_reasoning( + param: str, value: object, model: str | None, reason: str, *, drop_params_gated: bool +) -> None: + verbose_logger.warning( + _DROP_PARAMS_DROP_WARNING if drop_params_gated else _DROP_WARNING, param, reprlib.repr(value), model, reason + ) + + +def warn_non_string_reasoning_effort(value: object, model: str | None) -> None: + effort: Final = as_str_mapping(value) + requested_level: Final = effort.get("effort") if effort is not None else value + if requested_level is None or isinstance(value, str): + return + warn_dropped_reasoning("reasoning_effort", value, model, _NON_STRING_EFFORT_REASON, drop_params_gated=False) + + +def warn_dropped_reasoning_params( + dropped: Collection[object], + request: Mapping[str, object], + model: str | None, + reason: str, + *, + drop_params_gated: bool, +) -> None: + for param in _REASONING_PARAMS: + if param in dropped and param in request: + warn_dropped_reasoning(param, request[param], model, reason, drop_params_gated=drop_params_gated) + + +@dataclass(frozen=True, slots=True) +class DroppedThinkingBudget: + budget_tokens: object + + def warn(self, model: str) -> None: + warn_dropped_reasoning( + "thinking.budget_tokens", self.budget_tokens, model, _NON_NUMERIC_BUDGET_REASON, drop_params_gated=True + ) + def reasoning_effort_from_thinking_budget( budget_tokens: float, @@ -29,7 +83,9 @@ def reasoning_effort_from_thinking_budget( return "minimal" -def reasoning_effort_from_thinking(thinking: object, *, drop_params: bool) -> OpenAIStyleReasoningEffort | None: +def _effort_or_dropped_budget( + thinking: object, *, drop_params: bool +) -> OpenAIStyleReasoningEffort | DroppedThinkingBudget | None: fields: Final = as_str_mapping(thinking) if fields is None or fields.get("type") != "enabled": return None @@ -37,10 +93,27 @@ def reasoning_effort_from_thinking(thinking: object, *, drop_params: bool) -> Op if isinstance(budget, (int, float)) and not isinstance(budget, bool): return reasoning_effort_from_thinking_budget(budget) if drop_params or litellm.drop_params is True: - return None + return DroppedThinkingBudget(budget) raise non_numeric_thinking_budget_error(budget) +def reasoning_effort_from_thinking( + thinking: object, *, model: str, drop_params: bool +) -> OpenAIStyleReasoningEffort | None: + effort: Final = _effort_or_dropped_budget(thinking, drop_params=drop_params) + if not isinstance(effort, DroppedThinkingBudget): + return effort + effort.warn(model) + return None + + +def reasoning_effort_from_thinking_silently( + thinking: object, *, drop_params: bool +) -> OpenAIStyleReasoningEffort | None: + effort: Final = _effort_or_dropped_budget(thinking, drop_params=drop_params) + return None if isinstance(effort, DroppedThinkingBudget) else effort + + def non_numeric_thinking_budget_error(budget: object) -> UnsupportedParamsError: return UnsupportedParamsError( message=( diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 548d088d833..981859c620b 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -249,6 +249,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ "reasoningContentBlocks", "citations", "compaction_blocks", + "compaction_start", "compaction_delta", ) if key in fields diff --git a/litellm/router_utils/common_utils.py b/litellm/litellm_core_utils/routing_common_utils.py similarity index 100% rename from litellm/router_utils/common_utils.py rename to litellm/litellm_core_utils/routing_common_utils.py diff --git a/litellm/router_utils/routing_groups.py b/litellm/litellm_core_utils/routing_groups.py similarity index 100% rename from litellm/router_utils/routing_groups.py rename to litellm/litellm_core_utils/routing_groups.py diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 24c401281ef..909e1222ae6 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -1,4 +1,5 @@ import json +import math from collections.abc import Callable from typing import Any, Final @@ -14,16 +15,23 @@ def strip_null_bytes(value: str) -> str: return value.replace("\x00", "") +def _finite_or_none(value: float | None) -> float | None: + return None if isinstance(value, float) and not math.isfinite(value) else value + + def safe_json_structure( data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, value_transform: Callable[[str | None, str], str] | None = None, key: str | None = None, + *, + non_finite_as_null: bool = False, ) -> object: """ Rebuild data out of JSON-native pieces while detecting circular references. If a circular reference is detected then a marker string is returned. - NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. + NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. With + non_finite_as_null, NaN and infinite floats become null, since jsonb rejects them. value_transform, when given, is applied to every string leaf (and to the str() fallback for non-serializable objects) with the mapping key the leaf @@ -43,7 +51,7 @@ def safe_json_structure( cleaned = obj.replace("\x00", "") if "\x00" in obj else obj return _transform(key, cleaned) if isinstance(obj, (int, float, bool, type(None))): - return obj + return _finite_or_none(obj) if non_finite_as_null else obj # Check for circular reference. if id(obj) in seen: return "CircularReference Detected" @@ -65,8 +73,9 @@ def safe_json_structure( seen.remove(id(obj)) return tuple_result elif isinstance(obj, set): - sortable_values: Final[list[Any]] = [_serialize(item, seen, depth + 1, key) for item in obj] - set_result: Final = sorted(sortable_values) + members: Final[list[Any]] = [_serialize(item, seen, depth + 1, key) for item in obj] + present: Final = sorted(member for member in members if member is not None) + set_result: Final = present + [None] * members.count(None) seen.remove(id(obj)) return set_result elif isinstance(obj, BaseModel): @@ -88,6 +97,10 @@ def safe_dumps( data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, value_transform: Callable[[str | None, str], str] | None = None, + *, + non_finite_as_null: bool = False, ) -> str: """Serialize data to JSON text through safe_json_structure.""" - return json.dumps(safe_json_structure(data, max_depth, value_transform), default=str) + return json.dumps( + safe_json_structure(data, max_depth, value_transform, non_finite_as_null=non_finite_as_null), default=str + ) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 94351e72ab1..4857973d51d 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -14,11 +14,18 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, + UPSTREAM_USAGE_COST, chunk_usage, + finite_cost, provider_reported_usage_cost, stamp_litellm_bill_on_usage, ) -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.hidden_params import ( + HIDDEN_PARAMS_ATTR, + get_hidden_params, + set_hidden_param, + set_hidden_params, +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import partition_annotations from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -322,6 +329,20 @@ class ChunkProcessor: models[0] if models else None, ) + @staticmethod + def _get_upstream_provider(chunks: Sequence[object]) -> str | None: + def _provider_from_chunk(chunk: object) -> str | None: + hidden_params: Final = get_hidden_params(chunk) + if hidden_params is None: + return None + provider: Final[object] = hidden_params.get("upstream_provider") + return provider if isinstance(provider, str) and provider else None + + return next( + (provider for chunk in chunks if (provider := _provider_from_chunk(chunk)) is not None), + None, + ) + @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, @@ -499,10 +520,22 @@ class ChunkProcessor: chunks, first_chunk_model, ) - if provider_response_model is not None: - response._hidden_params = dict( # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter - response._hidden_params, # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params getter - provider_response_model=provider_response_model, + upstream_provider: Final = self._get_upstream_provider(chunks) + provider_hidden_params: Final[Mapping[str, str]] = MappingProxyType( + { + key: value + for key, value in ( + ("provider_response_model", provider_response_model), + ("upstream_provider", upstream_provider), + ) + if value is not None + } + ) + if provider_hidden_params: + response_hidden_params: Final[Mapping[str, builtins.object]] = get_hidden_params(response) or {} + set_hidden_params( + response, + dict(response_hidden_params, **provider_hidden_params), ) return response @@ -1014,7 +1047,9 @@ class ChunkProcessor: prompt_tokens_details, cache_creation_token_details ) - if usage_chunk_dict["cost"] is not None: + if Usage.dropped_non_finite_cost(usage_chunk): + cost = None + elif usage_chunk_dict["cost"] is not None: cost = usage_chunk_dict["cost"] prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details) @@ -1441,11 +1476,19 @@ def _set_stream_builder_response_cost( _set_hidden_response_cost(response, response_cost) +def _carry_upstream_usage_cost(response: ModelResponse, chunks: Sequence[object]) -> None: + chunk_hidden_params: Final = (get_hidden_params(chunk) for chunk in reversed(chunks)) + stamped: Final = (hidden for hidden in chunk_hidden_params if hidden is not None and UPSTREAM_USAGE_COST in hidden) + for hidden in stamped: + set_hidden_param(response, UPSTREAM_USAGE_COST, hidden[UPSTREAM_USAGE_COST]) + return + + def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: if logging_obj is None: return - if _logged_provider(logging_obj) in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS and isinstance( - getattr(usage, "cost", None), (int, float) + if _logged_provider(logging_obj) in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS and ( + finite_cost(getattr(usage, "cost", None)) is not None ): return computed_cost: Final = logging_obj.response_cost_calculator(result=response) @@ -1791,6 +1834,7 @@ def stream_chunk_builder( response, logging_obj, ) + _carry_upstream_usage_cost(response, chunks) # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost( response, @@ -1833,6 +1877,7 @@ def stream_chunk_builder( ) break + _carry_upstream_usage_cost(response, chunks) # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost( response, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 682b2a2f538..484edd869cb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -24,10 +24,16 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost +from litellm.litellm_core_utils.core_helpers import ( + UPSTREAM_USAGE_COST, + chunk_usage, + provider_reported_usage_cost, + stamp_litellm_bill_on_usage, +) from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, get_hidden_params, + set_extra_attribute, set_hidden_param, stored_attribute, ) @@ -97,12 +103,17 @@ def _stamp_usage_cost( usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream, complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream, ) -> None: - bill: Final = (get_hidden_params(complete_response) or {}).get("response_cost") + complete_hidden: Final = get_hidden_params(complete_response) or {} + bill: Final = complete_hidden.get("response_cost") usage: Final[object] = getattr(usage_chunk, "usage", None) if usage is None or isinstance(bill, bool) or not isinstance(bill, (int, float)): return - if bill > 0 or isinstance(getattr(usage, "cost", None), (int, float)): - setattr(usage, "cost", float(bill)) # noqa: B010 # cost is an optional extra on the usage model + if bill <= 0 and not isinstance(getattr(usage, "cost", None), (int, float)): + return + chunk_hidden: Final = get_hidden_params(usage_chunk) + if chunk_hidden is not None and UPSTREAM_USAGE_COST in complete_hidden: + chunk_hidden[UPSTREAM_USAGE_COST] = complete_hidden[UPSTREAM_USAGE_COST] + stamp_litellm_bill_on_usage(usage_chunk, usage, float(bill)) def _stamp_served_service_tier(response: ModelResponseStream, complete_streaming_response: ModelResponse) -> None: @@ -266,6 +277,7 @@ class _VertexChunkLike(Protocol): class _ParsedChunkHiddenParams(LiteLLMBaseModel): provider_specific_fields: Mapping[str, object] | None = None + upstream_provider: str | None = None def _provider_response_model(chunk: object) -> str | None: @@ -282,9 +294,17 @@ def _parsed_provider_hidden_params(hidden: object) -> _ParsedChunkHiddenParams | return None +def _provider_upstream_provider(chunk: object) -> str | None: + parsed: Final = _parsed_provider_hidden_params(get_hidden_params(chunk)) + if parsed is None or not parsed.upstream_provider: + return None + return parsed.upstream_provider + + def _provider_hidden_params( chunk: object, provider_response_model: str | None, + upstream_provider: str | None, ) -> Mapping[str, object] | None: hidden: Final = get_hidden_params(chunk) parsed: Final = _parsed_provider_hidden_params(hidden) @@ -297,6 +317,10 @@ def _provider_hidden_params( for key, value in ( ("provider_response_model", provider_response_model), ("provider_specific_fields", provider_specific_fields), + ( + "upstream_provider", + parsed.upstream_provider if parsed is not None and parsed.upstream_provider else upstream_provider, + ), ) if value is not None } @@ -349,6 +373,7 @@ class CustomStreamWrapper: self.system_fingerprint: str | None = None self._provider_response_model: str | None = None + self._upstream_provider: str | None = None self.received_finish_reason: str | None = None self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ @@ -1035,7 +1060,7 @@ class CustomStreamWrapper: _initial_delta: Final = model_response.choices[0].delta.model_dump() _initial_delta.pop("role", None) - model_response.choices[0].delta = Delta(**_initial_delta) + set_extra_attribute(model_response.choices[0], "delta", Delta(**_initial_delta)) return model_response def _prebuilt_plain_delta(self, chunk: object) -> Delta | None: @@ -1061,13 +1086,13 @@ class CustomStreamWrapper: def _set_role_on_fresh_delta(self, model_response: ModelResponseStream) -> None: delta: Final = model_response.choices[0].delta if self.sent_first_chunk is False: - delta.role = "assistant" + set_extra_attribute(delta, "role", "assistant") self.sent_first_chunk = True return if not _holds_only_scalars(delta): self.strip_role_from_delta(model_response) return - delta.role = None + set_extra_attribute(delta, "role", None) if "provider_specific_fields" not in delta.model_fields_set: delta.provider_specific_fields = None @@ -1186,8 +1211,8 @@ class CustomStreamWrapper: and _holds_only_scalars(upstream_delta) ) if can_reuse_choice and upstream_choice is not None: - model_response.choices[0].finish_reason = None # rebind-ok: keep reused choice - model_response.choices[0].index = upstream_choice.index # rebind-ok: keep reused choice + set_extra_attribute(model_response.choices[0], "finish_reason", None) + set_extra_attribute(model_response.choices[0], "index", upstream_choice.index) _prepend_upstream_choice_extras( model_response.choices[0], upstream_choice, @@ -1212,11 +1237,7 @@ class CustomStreamWrapper: else: return model_response.system_fingerprint = upstream_chunk.system_fingerprint - setattr( - model_response, - "citations", - stored_attribute(upstream_chunk, "citations"), - ) + set_extra_attribute(model_response, "citations", stored_attribute(upstream_chunk, "citations")) preserve_upstream_non_openai_attributes( model_response=model_response, original_chunk=original_chunk, @@ -1646,10 +1667,17 @@ class CustomStreamWrapper: provider_response_model: Final = _provider_response_model(chunk) if provider_response_model is not None: self._provider_response_model = provider_response_model + upstream_provider: Final = _provider_upstream_provider(chunk) + if self._upstream_provider is None and upstream_provider is not None: + self._upstream_provider = upstream_provider prebuilt_delta: Final = self._prebuilt_plain_delta(chunk) plain_delta_created: bool = prebuilt_delta is not None # rebind-ok: the first chunk builds its delta below model_response = self.model_response_creator( - hidden_params=_provider_hidden_params(chunk, self._provider_response_model), + hidden_params=_provider_hidden_params( + chunk, + self._provider_response_model, + self._upstream_provider, + ), delta=prebuilt_delta, ) response_obj: dict[str, Any] = {} @@ -2486,6 +2514,25 @@ class CustomStreamWrapper: def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: self._earlier_rounds_spend = earlier + def partial_response(self) -> "ModelResponse | TextCompletionResponse | None": + partial: Final = ( + litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages if isinstance(self.messages, list) else None, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + if self.chunks + else None + ) + usage: Final[object] = getattr(partial, "usage", None) + if partial is None or not isinstance(usage, Usage): + return None + if self.model: + partial.model = self.model + backfill_missing_cache_usage_fields(usage) + return partial + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2495,23 +2542,15 @@ class CustomStreamWrapper: moves it into terminal failure usage before opening the next attempt, so a recovered attempt logs its own usage. """ - if self.logging_obj is None or not self.chunks: + if self.logging_obj is None: return try: - partial_response: Final = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages if isinstance(self.messages, list) else None, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) + partial_response: Final = self.partial_response() if partial_response is None: return - usage: Final = cast(Usage | None, getattr(partial_response, "usage", None)) - if usage is None: + usage: Final[object] = getattr(partial_response, "usage", None) + if not isinstance(usage, Usage): return - if self.model: - partial_response.model = self.model - backfill_missing_cache_usage_fields(usage) partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 if self._earlier_rounds_spend is None: self.logging_obj.model_call_details["combined_usage_object"] = usage diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 800a0939cec..323af72940a 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -56,9 +56,9 @@ from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEnc if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace +from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import get_httpx_client, header_value -from litellm.router_utils.deployment_pricing import deployment_pricing_models from litellm.rust_bridge.tokenizer import get_encoding from litellm.types.llms.anthropic import ( AnthropicContentParamSource, diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index c18ce8cea68..84c91be9386 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -9,6 +9,8 @@ from litellm.types.utils import CallTypes from . import * if TYPE_CHECKING: + from collections.abc import Callable + from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, ) @@ -30,7 +32,29 @@ def get_cost_for_google_maps_grounding_request( return cost_per_google_maps_grounding_request(usage=usage, model_info=model_info) -def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", model_info: "ModelInfo") -> float | None: +def _with_served_model_search_pricing( + model_info: "ModelInfo", + response_model: str | None, + served_model_provider: str, + lookup_model_info: "Callable[[str, str | None], ModelInfo | None] | None", +) -> "ModelInfo": + if model_info.get("search_context_cost_per_query") is not None or response_model is None: + return model_info + served_model_info: Final = ( + lookup_model_info(response_model, served_model_provider) if lookup_model_info is not None else None + ) + if served_model_info is None: + return model_info + return {**model_info, "search_context_cost_per_query": served_model_info.get("search_context_cost_per_query")} + + +def get_cost_for_web_search_request( + custom_llm_provider: str, + usage: "Usage", + model_info: "ModelInfo", + response_model: str | None = None, + lookup_model_info: "Callable[[str, str | None], ModelInfo | None] | None" = None, +) -> float | None: """ Get the cost for a web search request for a given model. @@ -38,6 +62,8 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo custom_llm_provider: The custom LLM provider. usage: The usage object. model_info: The model info. + response_model: The model the provider response names. + lookup_model_info: Looks up a model's cost map entry, returning None when it is not mapped. """ if custom_llm_provider == "gemini": from .gemini.cost_calculator import cost_per_web_search_request @@ -46,17 +72,25 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo elif custom_llm_provider == "anthropic": from .anthropic.cost_calculation import get_cost_for_anthropic_web_search - return get_cost_for_anthropic_web_search(model_info=model_info, usage=usage) + return get_cost_for_anthropic_web_search( + model_info=_with_served_model_search_pricing(model_info, response_model, "anthropic", lookup_model_info), + usage=usage, + ) elif custom_llm_provider.startswith("vertex_ai"): # Anthropic Claude models on Vertex AI populate server_tool_use.web_search_requests # (same as the direct Anthropic API), not prompt_tokens_details.web_search_requests # (which is the Gemini field). Route claude-* models to the Anthropic calculator. model_key: Final[str] = model_info.get("key", "") if model_info else "" - if "claude" in model_key.lower(): + if "claude" in (response_model or model_key).lower(): from .anthropic.cost_calculation import get_cost_for_anthropic_web_search verbose_logger.debug("vertex_ai/claude model detected — routing web search cost to Anthropic calculator") - return get_cost_for_anthropic_web_search(model_info=model_info, usage=usage) + return get_cost_for_anthropic_web_search( + model_info=_with_served_model_search_pricing( + model_info, response_model, "vertex_ai", lookup_model_info + ), + usage=usage, + ) from .vertex_ai.gemini.cost_calculator import ( cost_per_web_search_request as cost_per_web_search_request_vertex_ai, @@ -71,7 +105,10 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo return None from .anthropic.cost_calculation import get_cost_for_anthropic_web_search - return get_cost_for_anthropic_web_search(model_info=model_info, usage=usage) + return get_cost_for_anthropic_web_search( + model_info=_with_served_model_search_pricing(model_info, response_model, "azure_ai", lookup_model_info), + usage=usage, + ) elif custom_llm_provider == "perplexity": # Perplexity handles search costs internally in its own cost calculator # Return 0.0 to indicate costs are already accounted for diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 865247afc63..d773e31c733 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -48,6 +48,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import place_mid_conversation_system, split_leading_system_run, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_non_string_reasoning_effort from litellm.llms.base_llm.base_utils import ( ResponseFormatSchema, as_json_object, @@ -1951,6 +1952,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if isinstance(effort_value, dict): effort_value = effort_value.get("effort") if not isinstance(effort_value, str): + warn_non_string_reasoning_effort(OPAQUE_VALUE.validate_python(value), model) continue mapped_thinking = AnthropicConfig._map_reasoning_effort( reasoning_effort=effort_value, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 82977f97854..e1f9a7d540f 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -37,7 +37,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of, ) -from litellm.litellm_core_utils.reasoning_effort_utils import non_numeric_thinking_budget_error +from litellm.litellm_core_utils.reasoning_effort_utils import DroppedThinkingBudget, non_numeric_thinking_budget_error from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.wif import ( aget_anthropic_wif_token, @@ -58,6 +58,7 @@ from litellm.types.llms.anthropic import ( ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER, ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES, + LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS, AllAnthropicToolsValues, AnthropicMcpServerTool, AnthropicMessagesToolChoice, @@ -126,6 +127,8 @@ _JSON_VALUE: Final = TypeAdapter(JsonValue) _CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") +NATIVE_COMPACTION_PROVIDERS: Final = frozenset({"anthropic", "azure_ai", "vertex_ai"}) + def is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: if custom_llm_provider == "anthropic": @@ -204,7 +207,7 @@ def requires_native_compaction_beta( optional_params: Mapping[str, object], messages: Sequence[object], ) -> bool: - return custom_llm_provider == "anthropic" and ( + return custom_llm_provider in NATIVE_COMPACTION_PROVIDERS and ( optional_params.get("compaction") is not None or carries_signed_compaction_block(messages) ) @@ -246,6 +249,27 @@ def context_management_betas(context_management: JsonValue) -> tuple[str, ...]: ) +def count_tokens_feature_betas( + custom_llm_provider: str, + optional_params: Mapping[str, object], + messages: Sequence[object], +) -> tuple[str, ...]: + return tuple( + beta + for beta, used in ( + ( + ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value, + requires_native_compaction_beta(custom_llm_provider, optional_params, messages), + ), + ( + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, + AnthropicModelInfo().is_mid_conversation_tool_change_used(messages), + ), + ) + if used + ) + + def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None) -> bool: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.utils import supports_prompt_caching @@ -610,11 +634,11 @@ def _numeric_budget_tokens(budget_tokens: object) -> int | None: return None -def _legacy_budget_tokens(budget_tokens: object, *, drop_params: bool) -> int | None: +def _legacy_budget_tokens(budget_tokens: object, *, drop_params: bool) -> int | DroppedThinkingBudget: if (numeric := _numeric_budget_tokens(budget_tokens)) is not None: return numeric if drop_params or litellm.drop_params is True: - return None + return DroppedThinkingBudget(budget_tokens) raise non_numeric_thinking_budget_error(budget_tokens) @@ -645,10 +669,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): file_ids: Final = get_file_ids_from_messages(messages) return len(file_ids) > 0 - def is_thinking_display_updates_used(self, thinking: AnthropicThinkingParam | None) -> bool: - if not isinstance(thinking, dict): - return False - return thinking.get("type") in ("adaptive", "enabled") and thinking.get("display") == "updates" + def is_thinking_display_updates_used(self, thinking: object) -> bool: + match thinking: + case {"type": "adaptive" | "enabled", "display": "updates"}: + return True + case _: + return False def thinking_betas(self, thinking: AnthropicThinkingParam | None) -> tuple[str, ...]: return ( @@ -1097,6 +1123,22 @@ class AnthropicModelInfo(BaseLLMModelInfo): custom_llm_provider: str, drop_params: bool, ) -> None: + dropped: Final = AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently( + model=model, + optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + drop_params=drop_params, + ) + if dropped is not None: + dropped.warn(model) + + @staticmethod + def translate_legacy_thinking_for_adaptive_model_silently( + model: str, + optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers + custom_llm_provider: str, + drop_params: bool, + ) -> DroppedThinkingBudget | None: """Translate legacy ``thinking.type=enabled`` to adaptive for the adaptive-thinking models that reject it (4.7+ and the 5 families). Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the @@ -1105,22 +1147,25 @@ class AnthropicModelInfo(BaseLLMModelInfo): ``output_config.effort`` is never overridden. """ if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): - return + return None if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): - return + return None thinking: Final = validated_object_mapping(optional_params.get("thinking")) if thinking is None or thinking.get("type") != "enabled": - return + return None budget_tokens: Final = _legacy_budget_tokens(thinking.get("budget_tokens", 0), drop_params=drop_params) display: Final = thinking.get("display") - supported_display: Final = {"display": display} if display in ("summarized", "omitted") else {} + keeps_display: Final = display in ("summarized", "omitted") or ( + display == "updates" and custom_llm_provider in LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS + ) + supported_display: Final = {"display": display} if keeps_display else {} carried_keys: Final = carried_thinking_keys( thinking, rebuilt_keys=frozenset(("type", "budget_tokens", "display")), drop_params=drop_params ) optional_params["thinking"] = {**carried_keys, "type": "adaptive", **supported_display} - if budget_tokens is None: - return + if isinstance(budget_tokens, DroppedThinkingBudget): + return budget_tokens effort: Final = AnthropicModelInfo._legacy_budget_to_effort( model=model, budget_tokens=budget_tokens, custom_llm_provider=custom_llm_provider ) @@ -1129,6 +1174,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): "effort": effort, **(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})), } + return None @staticmethod def _legacy_budget_to_effort(model: str, budget_tokens: int, custom_llm_provider: str) -> str: diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 8ea6293bbe6..bc9649ed867 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -134,6 +134,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking, + reasoning_effort_from_thinking_silently, ) from litellm.llms.anthropic.common_utils import ( eager_input_streaming_flag, @@ -297,6 +298,7 @@ class AnthropicAdapter: custom_llm_provider=custom_llm_provider, api_base=api_base, drop_params=drop_params, + report_reasoning_drops=True, ) return translated_body, tool_name_mapping @@ -666,7 +668,9 @@ class LiteLLMAnthropicMessagesAdapter: def translate_anthropic_thinking_to_reasoning_effort( thinking: AnthropicThinkingParam, *, + model: str, drop_params: bool = True, + report_reasoning_drops: bool = False, ) -> str | None: """ Translate Anthropic's thinking parameter to OpenAI's reasoning_effort. @@ -685,7 +689,9 @@ class LiteLLMAnthropicMessagesAdapter: if thinking_type == "disabled": return "none" elif thinking_type == "enabled": - return reasoning_effort_from_thinking(thinking, drop_params=drop_params) + if report_reasoning_drops: + return reasoning_effort_from_thinking(thinking, model=model, drop_params=drop_params) + return reasoning_effort_from_thinking_silently(thinking, drop_params=drop_params) elif thinking_type == "adaptive": # Adaptive thinking: effort is controlled by output_config.effort, # not budget_tokens. Return a default; caller should override with @@ -752,7 +758,7 @@ class LiteLLMAnthropicMessagesAdapter: return {"thinking": thinking} else: reasoning_effort: Final = LiteLLMAnthropicMessagesAdapter.translate_anthropic_thinking_to_reasoning_effort( - thinking + thinking, model=model ) if reasoning_effort: return { @@ -1147,6 +1153,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, drop_params: bool = True, + report_reasoning_drops: bool = False, ) -> None: """Translate Anthropic thinking to either thinking or reasoning_effort. @@ -1207,7 +1214,10 @@ class LiteLLMAnthropicMessagesAdapter: return reasoning_effort: Final = declared_effort or self.translate_anthropic_thinking_to_reasoning_effort( - cast(AnthropicThinkingParam, thinking), drop_params=drop_params + cast(AnthropicThinkingParam, thinking), + model=model, + drop_params=drop_params, + report_reasoning_drops=report_reasoning_drops, ) if not reasoning_effort: return @@ -1267,6 +1277,7 @@ class LiteLLMAnthropicMessagesAdapter: api_base: str | None = None, preserve_midturn_system: bool = False, drop_params: bool = True, + report_reasoning_drops: bool = False, ) -> tuple[ChatCompletionRequest, dict[str, str]]: """ This is used by the beta Anthropic Adapter, for translating anthropic `/v1/messages` requests to the openai format. @@ -1321,6 +1332,7 @@ class LiteLLMAnthropicMessagesAdapter: new_kwargs=new_kwargs, custom_llm_provider=custom_llm_provider, drop_params=drop_params, + report_reasoning_drops=report_reasoning_drops, ) ## CONVERT STOP_SEQUENCES self._translate_stop_sequences_to_openai( diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 6f97a019a96..1bf00cbdb2f 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -19,7 +19,15 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.reasoning_effort_capability import ( + RequestedReasoningEffort, + deployment_reasoning_effort_levels, + keeps_effort, + model_group_from_metadata, + resolve_deployment_reasoning_effort, +) from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, strip_provider_specific_fields_from_anthropic_messages, tools_marked_for_a_tool_config_point, ) @@ -30,14 +38,6 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway -from litellm.router_utils.reasoning_effort_capability import ( - RequestedReasoningEffort, - deployment_reasoning_effort_levels, - keeps_effort, - model_group_from_metadata, - resolve_deployment_reasoning_effort, - with_adaptive_thinking_as_effort, -) from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -185,9 +185,11 @@ def _with_deployment_output_effort( ).value def translate(source: Mapping[str, object]) -> Mapping[str, object]: - return with_adaptive_thinking_as_effort( - source, model=model, custom_llm_provider=custom_llm_provider, drop_params=drop_params + translated: Final = dict(source) + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=model, optional_params=translated, custom_llm_provider=custom_llm_provider, drop_params=drop_params ) + return translated requested_output: Final = as_str_mapping(params.get("output_config")) caller_output: Final = None if requested_output is None else resolve(requested_output) diff --git a/litellm/llms/anthropic/messages/native/interceptors/advisor.py b/litellm/llms/anthropic/messages/native/interceptors/advisor.py index 4c1d2451939..641ac6fca88 100644 --- a/litellm/llms/anthropic/messages/native/interceptors/advisor.py +++ b/litellm/llms/anthropic/messages/native/interceptors/advisor.py @@ -22,7 +22,7 @@ import litellm import litellm.constants as _c from litellm.litellm_core_utils.url_utils import validate_url from litellm.llms.anthropic.common_utils import strip_advisor_blocks_from_messages -from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure +from litellm.router_api import mark_advisor_orchestration_failure from litellm.types.llms.anthropic import ANTHROPIC_ADVISOR_TOOL_TYPE from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index ef855eff180..5281ad99620 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -21,6 +21,7 @@ from litellm.constants import ANTHROPIC_BILLING_METADATA_PREFIX from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_non_string_reasoning_effort from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.messages.convertible_blocks import drop_non_string_tool_result_texts from litellm.llms.anthropic.messages.safeguards import proxy_serves_safeguards @@ -40,6 +41,7 @@ from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.types.router import GenericLiteLLMParams from ...common_utils import ( + NATIVE_COMPACTION_PROVIDERS, OBJECT_LIST_ADAPTER, AnthropicError, AnthropicModelInfo, @@ -103,6 +105,7 @@ _MESSAGES_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): _workload_identity_eligible: ClassVar[bool] = True + _report_reasoning_drops: ClassVar[bool] = True @property def custom_llm_provider(self) -> str | None: @@ -126,7 +129,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): "tool_choice", "thinking", "context_management", - *(("compaction",) if self._resolved_provider == "anthropic" else ()), + *(("compaction",) if self._resolved_provider in NATIVE_COMPACTION_PROVIDERS else ()), "output_format", "inference_geo", "speed", @@ -695,12 +698,17 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ): anthropic_messages_optional_request_params.pop("output_format") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict + requested_reasoning_effort: Final = OPAQUE_VALUE.validate_python( + anthropic_messages_optional_request_params.get("reasoning_effort") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict + ) self._translate_reasoning_effort_to_anthropic( model=model, optional_params=anthropic_messages_optional_request_params, max_tokens=max_tokens, custom_llm_provider=self._resolved_provider, ) + if self._report_reasoning_drops: + warn_non_string_reasoning_effort(requested_reasoning_effort, model) AnthropicModelInfo.maybe_drop_disabled_thinking( model=model, @@ -708,12 +716,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): custom_llm_provider=self._resolved_provider, ) - AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + dropped_budget: Final = AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently( model=model, optional_params=anthropic_messages_optional_request_params, custom_llm_provider=self._resolved_provider, drop_params=getattr(litellm_params, "drop_params", None) is True, ) + if dropped_budget is not None and self._report_reasoning_drops: + dropped_budget.warn(model) self._translate_adaptive_effort_for_non_adaptive_model( model=model, diff --git a/litellm/llms/anthropic/messages/responses_adapters/transformation.py b/litellm/llms/anthropic/messages/responses_adapters/transformation.py index 94c608f90e6..0630bef0311 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/messages/responses_adapters/transformation.py @@ -534,6 +534,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: output_config: dict[str, object] | None = None, auto_summary_level: str = "detailed", *, + model: str, drop_params: bool = True, ) -> dict[str, object] | None: """ @@ -555,7 +556,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(output_config, dict) and output_config.get("effort"): effort = output_config["effort"] elif thinking_type == "enabled": - budget_effort: Final = reasoning_effort_from_thinking(thinking, drop_params=drop_params) + budget_effort: Final = reasoning_effort_from_thinking(thinking, model=model, drop_params=drop_params) if budget_effort is None: return None effort = budget_effort @@ -667,6 +668,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: thinking, output_config=cast(dict[str, object] | None, output_config), auto_summary_level=auto_summary_level, + model=model, drop_params=drop_params, ) if reasoning: diff --git a/litellm/llms/anthropic/messages/utils.py b/litellm/llms/anthropic/messages/utils.py index 31619a04b74..95c7914befa 100644 --- a/litellm/llms/anthropic/messages/utils.py +++ b/litellm/llms/anthropic/messages/utils.py @@ -214,7 +214,7 @@ def normalize_reasoning_effort_value( if chain is None: return effort - from litellm.router_utils.reasoning_effort_capability import ( + from litellm.litellm_core_utils.reasoning_effort_capability import ( deployment_rounds_reasoning_effort, resolve_supported_reasoning_efforts, ) diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 27e003ce8b6..2da071ebdd2 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -9,7 +9,7 @@ from dataclasses import dataclass, field from functools import reduce from itertools import accumulate, groupby from types import MappingProxyType -from typing import Annotated, Final, Literal, Protocol, TypeAlias +from typing import Annotated, ClassVar, Final, Literal, Protocol, TypeAlias import httpx from pydantic import ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError @@ -81,6 +81,10 @@ _DEPLOYMENT_OPTIONS: Final = frozenset( ) +class _BaselineAnthropicMessagesConfig(AnthropicMessagesConfig): + _report_reasoning_drops: ClassVar[bool] = False + + class _StrictModel(LiteLLMBaseModel): model_config = ConfigDict(extra="forbid", frozen=True, strict=True) @@ -661,7 +665,7 @@ def prepare_native_baseline_body(request: Mapping[str, object], model: str) -> M TypeAdapter(tuple[str, ...]).validate_python(owned.get("additional_drop_params") or ()), dict(options), ) - config: Final = AnthropicMessagesConfig() + config: Final = _BaselineAnthropicMessagesConfig() body: Final = config.transform_anthropic_messages_request( model=resolved_model, messages=strip_provider_specific_fields_from_anthropic_messages(messages), diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 4b3f263201a..c324f0f0980 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -84,6 +84,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): headers: Final = self.get_required_headers( api_key=api_key, litellm_params=litellm_params, + messages=messages, ) if not {"x-api-key", "Authorization"} & headers.keys(): raise AnthropicError( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index 6e3cfce5eb4..c02c92a95d6 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -4,11 +4,13 @@ Azure AI Anthropic CountTokens API transformation logic. Extends the base Anthropic CountTokens transformation with Azure authentication. """ +from collections.abc import Sequence from typing import Final from urllib.parse import urlsplit from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.anthropic.common_utils import count_tokens_feature_betas, merge_anthropic_beta_headers from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -28,6 +30,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): self, api_key: str | None, litellm_params: dict[str, object] | None = None, + messages: Sequence[object] = (), ) -> dict[str, str]: """ Get the required headers for the Azure AI Anthropic CountTokens API. @@ -38,6 +41,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): Args: api_key: The Azure AI API key litellm_params: Optional LiteLLM parameters for additional auth config + messages: The request messages, read for betas the counted features need Returns: Dictionary of required headers with both x-api-key and Azure authentication @@ -46,7 +50,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): headers: Final = { "Content-Type": "application/json", "anthropic-version": "2023-06-01", - "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + "anthropic-beta": merge_anthropic_beta_headers( + ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, count_tokens_feature_betas("azure_ai", {}, messages) + ), **({"x-api-key": api_key} if api_key else {}), } diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 5381447f28f..5214d31fb2e 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( prompt_cache_breakpoint_for_wire, with_prompt_cache_breakpoint, ) +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning_params from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.azure_ai.common_utils import ( api_key_header_for_base, @@ -40,6 +41,14 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +_REQUEST_BODY: Final = TypeAdapter(dict[str, object]) + + +def _request_model(request: Mapping[str, object]) -> str | None: + model: Final = request.get("model") + return model if isinstance(model, str) else None + + class AzureFoundryErrorStrings(str, enum.Enum): SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'" @@ -510,7 +519,13 @@ class AzureAIStudioConfig(OpenAIConfig): request_data = self._drop_extra_params_from_request_data(request_data, error_text) if "Extra inputs are not permitted" in error_text and self._error_has_tool_level_extra_fields(error_text): request_data = self._drop_tool_level_extra_fields(request_data, error_text) - data: Final = drop_params_from_unprocessable_entity_error(e=e, data=request_data) + data: Final = drop_params_from_unprocessable_entity_error( + e=e, + data=request_data, + model=_request_model(_REQUEST_BODY.validate_python(request_data)), + drop_params_gated="Extra inputs are not permitted" in error_text + and not self._error_has_tool_level_extra_fields(error_text), + ) return data def _drop_tool_level_extra_fields(self, request_data: dict, error_text: str) -> dict: @@ -526,6 +541,14 @@ class AzureAIStudioConfig(OpenAIConfig): def _drop_extra_params_from_request_data(self, request_data: dict, error_text: str) -> dict: params_to_drop: Final = self._extract_params_to_drop_from_error_text(error_text) if params_to_drop: + request: Final = _REQUEST_BODY.validate_python(request_data) + warn_dropped_reasoning_params( + params_to_drop, + request, + _request_model(request), + "Azure AI rejected it as an extra parameter its deployment does not pass through.", + drop_params_gated=False, + ) for param in params_to_drop: if param in request_data: request_data.pop(param, None) diff --git a/litellm/llms/azure_ai/decisions/transformation.py b/litellm/llms/azure_ai/decisions/transformation.py index fb5330ed715..22c4b59f5e8 100644 --- a/litellm/llms/azure_ai/decisions/transformation.py +++ b/litellm/llms/azure_ai/decisions/transformation.py @@ -12,7 +12,7 @@ from litellm.llms.azure_ai.common_utils import ( AZURE_ENTRA_LITELLM_PARAM_KEYS, get_azure_ai_auth_headers, ) -from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits _FOUNDRY_ROUTE_SUFFIX: Final = re.compile(r"(/api/projects/[^/]+)?(/openai/v1|/openai|/models|/v1)?$") _AZURE_AI_CREDENTIAL_PARAM_KEYS: Final = AZURE_ENTRA_LITELLM_PARAM_KEYS | frozenset({"azure_ad_token_provider"}) @@ -26,6 +26,16 @@ class AzureAIDecisionsConfig(BaseDecisionsConfig): api_base_env = ("AZURE_AI_API_BASE",) api_key_required = False + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + limits: Final = super().get_input_limits(model, custom_llm_provider, headers) + return ( + limits + if limits.max_request_tokens is not None + else super().get_input_limits("Microsoft-Decision-1", custom_llm_provider, headers) + ) + def missing_api_base_message(self, custom_llm_provider: str) -> str: return "Missing AZURE_AI_API_BASE - set AZURE_AI_API_BASE or pass api_base, e.g. https://.services.ai.azure.com" diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index e8e4b77e757..efd3e3fe9b5 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -1,7 +1,7 @@ import json from abc import ABC from collections.abc import Callable, Mapping, Sequence -from dataclasses import replace +from dataclasses import dataclass, replace from itertools import chain from types import MappingProxyType from typing import Final @@ -511,6 +511,12 @@ def _is_decisions_model(model: str, custom_llm_provider: str) -> bool: return is_decisions_model_mode(model_mode) +@dataclass(frozen=True, slots=True) +class DecisionsInputLimits: + max_request_tokens: int | None = None + max_question_tokens: int | None = None + + class BaseDecisionsConfig(ABC): path: str = "/v1/systemone" api_key_env: tuple[str, ...] = () @@ -533,6 +539,19 @@ class BaseDecisionsConfig(ABC): def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None: return None + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + try: + limit: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get( + "max_input_tokens" + ) + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model + return DecisionsInputLimits() + return DecisionsInputLimits( + max_request_tokens=limit if isinstance(limit, int) and not isinstance(limit, bool) and limit > 0 else None + ) + def missing_api_base_message(self, custom_llm_provider: str) -> str: return f"api_base is required for Decisions provider '{custom_llm_provider}'" diff --git a/litellm/llms/bedrock/chat/chat_completions/transformation.py b/litellm/llms/bedrock/chat/chat_completions/transformation.py index 123e4a9e6b4..dd2277a6c35 100644 --- a/litellm/llms/bedrock/chat/chat_completions/transformation.py +++ b/litellm/llms/bedrock/chat/chat_completions/transformation.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( inline_remote_image_urls, inline_remote_media, ) +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token @@ -42,6 +43,7 @@ from litellm.llms.bedrock.common_utils import ( bedrock_rejects_stop_sequences, bedrock_runtime_chat_completions_serves_reasoning_inline, split_bedrock_region_path, + strip_bedrock_routing_prefix, ) from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig @@ -121,6 +123,13 @@ def without_refused_reasoning_effort(model: str, params: Mapping[str, object]) - effort: Final = params.get("reasoning_effort") if not isinstance(effort, str) or effort not in chat_completions_reasoning_efforts_refused_for(model): return params + warn_dropped_reasoning( + "reasoning_effort", + effort, + strip_bedrock_routing_prefix(model), + "AWS's Chat Completions endpoint rejects that level for this model family, so AWS applies its default effort.", + drop_params_gated=False, + ) return _without_params(params, frozenset(("reasoning_effort",))) @@ -445,6 +454,14 @@ class AmazonBedrockRuntimeChatCompletionsConfig(OpenAILikeChatConfig): ), status_code=400, ) + if malformed_effort: + warn_dropped_reasoning( + "reasoning_effort", + raw_params["reasoning_effort"], + strip_bedrock_routing_prefix(model), + "Bedrock's Chat Completions endpoint takes it only as a string, so AWS applies its default effort.", + drop_params_gated=True, + ) rejected_stop: Final = frozenset(("stop",)) if bedrock_rejects_stop_sequences(model) else frozenset[str]() return dict( without_refused_reasoning_effort( diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 443e722168d..3800489ec05 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -50,6 +50,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, @@ -63,6 +64,7 @@ from litellm.llms.bedrock.chat.tool_result_images import place_tool_result_image from litellm.llms.bedrock.common_utils import ( bedrock_chat_uses_native_structured_output, bedrock_cleaned_output_schema, + bedrock_model_is_unknown, bedrock_model_supports_regex_lookaround, bedrock_native_output_format, bedrock_native_output_is_opt_in, @@ -186,6 +188,8 @@ def _tools_the_model_accepts( _OUTPUT_CONFIG_BLOCK: Final = TypeAdapter(OutputConfigBlock) +_DEEPSEEK_THINKING_REASON: Final = "DeepSeek on Bedrock reasons natively and takes no thinking field." +_DEEPSEEK_R1_EFFORT_REASON: Final = "DeepSeek R1 on Bedrock always reasons and rejects reasoning_effort." def _has_native_text_format(output_config: object) -> bool: @@ -529,13 +533,14 @@ class AmazonConverseConfig(BaseConfig): """ Handle the reasoning_effort parameter based on the model type. - - GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields. + - GPT-OSS, DeepSeek V3 and unknown models (``bedrock_model_is_unknown``): passed through unchanged + via additionalModelRequestFields. - OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields. - Nova 2 models: transformed to reasoningConfig. - Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on adaptive Claude 4.6 / 4.7). """ - if "gpt-oss" in model or "deepseek" in model: + if "gpt-oss" in model or "deepseek" in model or bedrock_model_is_unknown(model): optional_params["reasoning_effort"] = reasoning_effort elif self._is_openai_gpt_reasoning_model(model): reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort} @@ -638,6 +643,10 @@ class AmazonConverseConfig(BaseConfig): or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) ) + def _warn_native_reasoning_drop(self, param: str, value: object, model: str, reason: str) -> None: + if param in self.get_supported_openai_params(model): + warn_dropped_reasoning(param, value, model, reason, drop_params_gated=False) + def _model_rejects_reasoning_effort_param(self, model: str, base_model: str) -> bool: """Whether the model returns a 400 for every ``reasoning_effort`` shape on Converse. @@ -681,7 +690,7 @@ class AmazonConverseConfig(BaseConfig): "parallel_tool_calls", ] - if "arn" in model: # we can't infer the model from the arn, so just add all params + if "arn" in model or bedrock_model_is_unknown(model): supported_params.append("tools") supported_params.append("tool_choice") supported_params.append("thinking") @@ -1141,9 +1150,8 @@ class AmazonConverseConfig(BaseConfig): "tool_choice": {"type": "auto", "disable_parallel_tool_use": not value} } if param == "thinking" and drop_thinking_param: - verbose_logger.debug( - "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.", - model, + self._warn_native_reasoning_drop( + param, OPAQUE_VALUE.validate_python(value), model, _DEEPSEEK_THINKING_REASON ) elif param == "thinking" and not self._is_openai_gpt_reasoning_model(model): if ( @@ -1177,10 +1185,7 @@ class AmazonConverseConfig(BaseConfig): drop_params=drop_params, ) elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_effort_param: - verbose_logger.debug( - "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it always reasons and rejects it.", - model, - ) + self._warn_native_reasoning_drop(param, value, model, _DEEPSEEK_R1_EFFORT_REASON) elif ( param == "reasoning_effort" and isinstance(value, str) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 852179e8caa..877fb4ca692 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -2,7 +2,7 @@ import copy import json import time from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, cast, get_args +from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config @@ -74,10 +74,6 @@ _GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = ( "{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}" ) -_INVOKE_PROVIDERS: Final = TypeAdapter(tuple[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, ...]).validate_python( - get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL) -) - def _response_header(headers: Mapping[str, str], name: str) -> str | None: return headers.get(name, None) @@ -562,62 +558,6 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): """ return False - @staticmethod - def get_bedrock_invoke_provider( - model: str, - ) -> litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None: - """ - Helper function to get the bedrock provider from the model - - handles 4 scenarios: - 1. model=invoke/anthropic.claude-3-5-sonnet-20240620-v1:0 -> Returns `anthropic` - 2. model=anthropic.claude-3-5-sonnet-20240620-v1:0 -> Returns `anthropic` - 3. model=llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n -> Returns `llama` - 4. model=us.amazon.nova-pro-v1:0 -> Returns `nova` - """ - if model.startswith("invoke/"): - model = model.replace("invoke/", "", 1) - - # Special case: Check for "nova" in model name first (before "amazon") - # This handles amazon.nova-* models which would otherwise match "amazon" (Titan) - if "nova" in model.lower(): - if "nova" in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova") - - _split_model: Final = model.split(".")[0] - if _split_model in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model) - - # If not a known provider, check for pattern with two slashes - provider = AmazonInvokeConfig._get_provider_from_model_path(model) - if provider is not None: - return provider - - for provider in _INVOKE_PROVIDERS: - if provider in model: - return provider - return None - - @staticmethod - def _get_provider_from_model_path( - model_path: str, - ) -> litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None: - """ - Helper function to get the provider from a model path with format: provider/model-name - - Args: - model_path (str): The model path (e.g., 'llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n' or 'anthropic/model-name') - - Returns: - Optional[str]: The provider name, or None if no valid provider found - """ - parts: Final = model_path.split("/") - if len(parts) >= 1: - provider: Final = parts[0] - if provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider) - return None - def convert_messages_to_prompt(self, model: str, messages, provider, custom_prompt_dict) -> tuple[str, list | None]: # handle anthropic prompts and amazon titan prompts prompt = "" diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index d88d7e6280c..d6f5d5b6bd7 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1098,6 +1098,11 @@ def _model_cost_entry(key: str) -> Mapping[str, object] | None: return None if raw is None else _MODEL_COST_ENTRY_ADAPTER.validate_python(raw) +def _cost_map_row_names_a_provider(key: str) -> bool: + entry: Final = _model_cost_entry(key) + return entry is not None and entry.get("litellm_provider") is not None + + def _bedrock_price_map_entries(model: str) -> tuple[Mapping[str, object] | None, ...]: return tuple( _model_cost_entry(key) @@ -1567,7 +1572,8 @@ def bedrock_model_supports_regex_lookaround(model: str, litellm_params: Mapping[ The deployment's ``model_info.supports_regex_lookaround`` wins, then the ``model_prices_and_context_window.json`` entry of its ``base_model``, then the - entry of ``model`` itself. A model nobody flagged keeps its schema as sent. + entry of ``model`` itself. A model nobody flagged keeps its schema as sent, unless + no candidate names a model LiteLLM knows, whatever route prefix it carries. """ params: Final = litellm_params or {} deployment_flag: Final = deployment_model_info(params).get("supports_regex_lookaround") @@ -1576,7 +1582,41 @@ def bedrock_model_supports_regex_lookaround(model: str, litellm_params: Mapping[ base_model: Final = params.get("base_model") candidates: Final = (*((base_model,) if isinstance(base_model, str) else ()), model) flags: Final = (_bedrock_converse_model_flag(candidate, "supports_regex_lookaround") for candidate in candidates) - return next((flag for flag in flags if flag is not None), True) + flag: Final = next((flag for flag in flags if flag is not None), None) + if flag is not None: + return flag + return not all(bedrock_model_is_unknown(split_bedrock_region_path(candidate)[1]) for candidate in candidates) + + +def _bedrock_model_id_without_region(model: str) -> str: + routed_id: Final = model.removeprefix("bedrock/") + region, separator, regional_id = routed_id.partition("/") + return regional_id if separator and region in get_all_bedrock_regions() else routed_id + + +def bedrock_model_is_unknown(model: str) -> bool: + model_id: Final = _bedrock_model_id_without_region(model) + if "/" in model_id: + return False + base: Final = get_bedrock_base_model(model_id) + regional_ids: Final = (f"{region}.{base}" for region in get_bedrock_cross_region_inference_regions()) + if any(_cost_map_row_names_a_provider(key) for key in (*_bedrock_model_info_candidates(model_id), *regional_ids)): + return False + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + return BaseAWSLLM.get_bedrock_invoke_provider(model_id) is None + + +def _names_no_invoke_provider(model: str) -> bool: + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + model_id: Final = _bedrock_model_id_without_region(model) + names_one_model: Final = bool(model_id) and ("/" not in model_id or model_id.startswith("arn:")) + return names_one_model and BaseAWSLLM.get_bedrock_invoke_provider(model) is None + + +def bedrock_model_reaches_converse_by_default(model: str) -> bool: + return BedrockModelInfo.named_bedrock_route(model, None) is None and _names_no_invoke_provider(model) _BedrockConverseModelFlag: TypeAlias = Literal[ @@ -1585,10 +1625,13 @@ _BedrockConverseModelFlag: TypeAlias = Literal[ ] -def _bedrock_converse_model_flag(model: str, key: _BedrockConverseModelFlag) -> bool | None: +def _bedrock_model_info_candidates(model: str) -> tuple[str, ...]: base: Final = get_bedrock_base_model(model) - candidates: Final = dict.fromkeys((model, base, _BEDROCK_MODEL_VERSION_SUFFIX_RE.sub("", base))) - for candidate in candidates: + return tuple(dict.fromkeys((model, base, _BEDROCK_MODEL_VERSION_SUFFIX_RE.sub("", base)))) + + +def _bedrock_converse_model_flag(model: str, key: _BedrockConverseModelFlag) -> bool | None: + for candidate in _bedrock_model_info_candidates(model): with contextlib.suppress(Exception): model_info = get_cached_model_info()( model=candidate, @@ -1766,8 +1809,19 @@ class BedrockModelInfo(BaseLLMModelInfo): default (``bedrock_runtime_chat_completions_is_default``) and ``chat_completions/`` opts any other model in; ``request_params`` (the caller's chat params) sends such a request to Converse when it needs a feature only Converse serves, and ``converse/`` pins a model to Converse. Every other OpenAI-family - model stays on Converse without the prefix. + model stays on Converse without the prefix. A model the cost map does not list as Converse goes to + Invoke only when its id names an Invoke provider, and to Converse otherwise. """ + named_route: Final = BedrockModelInfo.named_bedrock_route(model, request_params) + if named_route is not None: + return named_route + return "converse" if _names_no_invoke_provider(model) else "invoke" + + @staticmethod + def named_bedrock_route( + model: str, + request_params: Mapping[str, object] | None, + ) -> BedrockRoute | None: route_mappings: dict[ str, Literal[ @@ -1822,7 +1876,7 @@ class BedrockModelInfo(BaseLLMModelInfo): return "converse" if _OPENAI_FAMILY_MODEL_RE.search(base_model): return "converse" - return "invoke" + return None @staticmethod def _explicit_converse_route(model: str) -> bool: diff --git a/litellm/llms/bedrock/rerank/transformation.py b/litellm/llms/bedrock/rerank/transformation.py index 442903fc05e..c2f5035528d 100644 --- a/litellm/llms/bedrock/rerank/transformation.py +++ b/litellm/llms/bedrock/rerank/transformation.py @@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works from typing import Final +from pydantic import TypeAdapter + from litellm._uuid import uuid from litellm.types.llms.bedrock import ( BedrockRerankBedrockRerankingConfiguration, @@ -90,7 +92,9 @@ class BedrockRerankConfig: example input: {"results":[{"index":0,"relevanceScore":0.6847912669181824},{"index":1,"relevanceScore":0.5980774760246277}]} """ - _billed_units = RerankBilledUnits(**response.get("usage", {"search_units": 1})) # by default 1 search unit + _billed_units: Final = TypeAdapter(RerankBilledUnits).validate_python( + response.get("usage", {"search_units": 1}) + ) _tokens: Final = RerankTokens(**response.get("usage", {})) rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) diff --git a/litellm/llms/clf_ai_gateway/chat/transformation.py b/litellm/llms/clf_ai_gateway/chat/transformation.py index 25e5d0f51f7..04b689c291e 100644 --- a/litellm/llms/clf_ai_gateway/chat/transformation.py +++ b/litellm/llms/clf_ai_gateway/chat/transformation.py @@ -4,8 +4,8 @@ import litellm from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig -from litellm.router_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model from litellm.secret_managers.main import get_secret_str SUPPORTED_OPENAI_PARAMS: Final = ( diff --git a/litellm/llms/compaction.py b/litellm/llms/compaction.py index 16afa2e2c25..034958c63ca 100644 --- a/litellm/llms/compaction.py +++ b/litellm/llms/compaction.py @@ -9,7 +9,7 @@ from pydantic import TypeAdapter from litellm.types.llms.openai import AllMessageValues if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext CompactionProtocol: TypeAlias = Literal["chat", "messages"] _MAPPING: Final = TypeAdapter(Mapping[str, object]) @@ -32,7 +32,9 @@ def get_native_compaction_provider(params: Mapping[str, object]) -> NativeCompac return compaction if compaction.supports_native_compaction(params) else None -async def dispatch(router: Router, protocol: CompactionProtocol, payload: Mapping[str, object]) -> Mapping[str, object]: +async def dispatch( + router: RoutingContext, protocol: CompactionProtocol, payload: Mapping[str, object] +) -> Mapping[str, object]: if protocol == "messages": return _MAPPING.validate_python( await router.aanthropic_messages(custom_llm_provider=None, client=None, **payload) @@ -45,4 +47,4 @@ async def dispatch(router: Router, protocol: CompactionProtocol, payload: Mappin {key: value for key, value in payload.items() if key not in ("model", "messages", "stream")} ), ) - return _MAPPING.validate_python(response.model_dump()) + return _MAPPING.validate_python(response.model_dump()) # pyright: ignore[reportAttributeAccessIssue] # stream=False: acompletion returns a ModelResponse diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 9297e76624d..c32e75bafa8 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -110,6 +110,15 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]: raise mapped_exc(message) from exc +def _unread_payload(content: object) -> bytes: + """Bytes the upstream sent before the connection dropped that aiohttp buffered but never handed out, since its + reads raise the stored payload error ahead of the buffer.""" + if not isinstance(content, aiohttp.StreamReader): + return b"" + buffered: Final = b"".join(content._buffer) # pyright: ignore[reportPrivateUsage] # no public read drains it + return buffered[content._buffer_offset :] # pyright: ignore[reportPrivateUsage] # see above + + def _release_transport_anchor(_transport: object) -> None: """Release a transport anchored to a response stream. See ``_anchor_transport_to``.""" @@ -144,6 +153,12 @@ class AiohttpResponseStream(httpx.AsyncByteStream): raise httpx.ReadError(str(e)) from e except aiohttp.http_exceptions.TransferEncodingError as e: raise httpx.ReadError(str(e)) from e + except aiohttp.ClientPayloadError: + unread: Final = _unread_payload(self._aiohttp_response.content) + if unread: + yield unread + with map_aiohttp_exceptions(): + raise except Exception: # For other exceptions, use the normal mapping with map_aiohttp_exceptions(): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 61fc23589ac..9a2b352dd86 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -249,7 +249,7 @@ if TYPE_CHECKING: FakeAnthropicMessagesStreamIterator, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig - from litellm.responses.streaming_iterator import ResponsesWsFrameGuard + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard, ResponsesWsTurnAdmission from litellm.types.google_genai.main import GenerateContentResponse from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -7269,7 +7269,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, - turn_admission: Callable[[], Awaitable[None]] | None = None, + turn_admission: "ResponsesWsTurnAdmission | None" = None, frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs: object, ) -> Exception | None: diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index ecaa0a2ef16..5b7c2a2c410 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -6,10 +6,12 @@ from collections.abc import Coroutine, Mapping, Sequence from typing import Final, Literal, cast, overload import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, extract_search_results_text, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -18,6 +20,16 @@ from litellm.utils import supports_reasoning, supports_vision from ...openai.chat.gpt_transformation import OpenAIGPTConfig SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "xhigh", "max", "ultra") +_NO_BUDGET_REASON: Final = "DeepSeek takes thinking only as enabled or disabled, with no budget." + + +def _warn_on_dropped_budget(thinking: object, model: str) -> None: + fields: Final = as_str_mapping(thinking) + if fields is None or fields.get("type") != "enabled" or "budget_tokens" not in fields: + return + warn_dropped_reasoning( + "thinking.budget_tokens", fields["budget_tokens"], model, _NO_BUDGET_REASON, drop_params_gated=False + ) class DeepSeekChatConfig(OpenAIGPTConfig): @@ -53,6 +65,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): # Then re-add only if valid for DeepSeek thinking_value: Final = optional_params.pop("thinking", None) reasoning_effort: Final = optional_params.pop("reasoning_effort", None) + _warn_on_dropped_budget(OPAQUE_VALUE.validate_python(thinking_value), model) # Handle thinking parameter - accept both enabled and disabled, ignore budget_tokens if isinstance(thinking_value, dict) and thinking_value.get("type") in ("enabled", "disabled"): diff --git a/litellm/llms/fal_ai/chat/transformation.py b/litellm/llms/fal_ai/chat/transformation.py index 115b1f840de..848ae498cdd 100644 --- a/litellm/llms/fal_ai/chat/transformation.py +++ b/litellm/llms/fal_ai/chat/transformation.py @@ -12,7 +12,9 @@ from typing import TYPE_CHECKING, Final import httpx from pydantic import ConfigDict, TypeAdapter +import litellm from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.base import LiteLLMBaseModel @@ -29,6 +31,10 @@ PROVIDER_PREFIX: Final[str] = "fal_ai/" PASSTHROUGH_PARAMS: Final[frozenset[str]] = frozenset(("reasoning", "temperature", "top_p")) REASONING_DISABLED_EFFORTS: Final[frozenset[str]] = frozenset(("none", "minimal")) REASONING_ENABLED_EFFORTS: Final[frozenset[str]] = frozenset(("low", "medium", "high")) +_UNSUPPORTED_EFFORT_REASON: Final = ( + f"fal only turns reasoning off for {sorted(REASONING_DISABLED_EFFORTS)} " + f"and on for {sorted(REASONING_ENABLED_EFFORTS)}." +) class _FalUsage(LiteLLMBaseModel): @@ -120,9 +126,10 @@ class FalAIChatConfig(BaseConfig): return False if isinstance(value, str) and value in REASONING_ENABLED_EFFORTS: return True - if drop_params: - return None - raise FalAIError(status_code=400, message=f"Unsupported reasoning_effort {value!r} for {model}") + if not (drop_params or litellm.drop_params is True): + raise FalAIError(status_code=400, message=f"Unsupported reasoning_effort {value!r} for {model}") + warn_dropped_reasoning("reasoning_effort", value, model, _UNSUPPORTED_EFFORT_REASON, drop_params_gated=True) + return None def _translate_param(self, param: str, value: object, model: str, drop_params: bool) -> tuple[str, object] | None: if param in ("temperature", "top_p"): diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index fa5794d71b3..92c6c8439b3 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_response_only_fields, unpack_legacy_defs, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -97,6 +98,7 @@ def _json_schema_response_format(schema: object, name: str) -> Mapping[str, obje EFFORT_KWARG_KEYS: Final = frozenset({"enable_thinking", "thinking", "reasoning_budget", "low_effort"}) +_NOT_A_REASONING_MODEL_REASON: Final = "litellm's model map does not list it as a reasoning model." def _bool_from_kwargs(kwargs: Mapping[str, object], keys: tuple[str, ...]) -> bool | None: @@ -107,6 +109,13 @@ def _bool_from_kwargs(kwargs: Mapping[str, object], keys: tuple[str, ...]) -> bo return None +def warn_dropped_effort_kwargs(chat_template_kwargs: object, effort: object, model: str) -> None: + if effort != "none": + warn_dropped_reasoning( + "chat_template_kwargs", chat_template_kwargs, model, _NOT_A_REASONING_MODEL_REASON, drop_params_gated=False + ) + + def effort_from_chat_template_kwargs(kwargs: Mapping[str, object]) -> object: enable_thinking: Final = _bool_from_kwargs(kwargs, ("enable_thinking", "thinking")) if enable_thinking is False: @@ -438,10 +447,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): if effort is None: return () if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): - verbose_logger.debug( - "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs effort keys.", - model, - ) + warn_dropped_effort_kwargs(OPAQUE_VALUE.validate_python(chat_template_kwargs), effort, model) return () return (("reasoning_effort", effort),) diff --git a/litellm/llms/fireworks_ai/completion/transformation.py b/litellm/llms/fireworks_ai/completion/transformation.py index 1f5d4f306a9..611e60cedd4 100644 --- a/litellm/llms/fireworks_ai/completion/transformation.py +++ b/litellm/llms/fireworks_ai/completion/transformation.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from typing import Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUserMessage from litellm.utils import supports_reasoning @@ -12,6 +13,7 @@ from ..chat.transformation import ( NIM_VLLM_STRIP_PARAMS, FireworksAIConfig, effort_from_chat_template_kwargs, + warn_dropped_effort_kwargs, ) from ..common_utils import FireworksAIMixin, resolve_fireworks_resource_name @@ -129,10 +131,7 @@ class FireworksAITextCompletionConfig(FireworksAIMixin, BaseTextCompletionConfig ) return result if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): - verbose_logger.debug( - "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs effort keys.", - model, - ) + warn_dropped_effort_kwargs(OPAQUE_VALUE.validate_python(chat_template_kwargs), effort, model) return result return {**result, "reasoning_effort": effort} diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 63e6cb2b5da..9f9503f93b6 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -99,7 +99,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): thinking_effort: Final = ( None if "reasoning_effort" in non_default_params - else reasoning_effort_from_thinking(thinking, drop_params=drop_params) + else reasoning_effort_from_thinking(thinking, model=model, drop_params=drop_params) ) if thinking_effort is not None: non_default_params["reasoning_effort"] = thinking_effort diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 0b77a04c69a..1b7da6f6635 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -244,5 +244,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): return RerankResponse( id=response.get("id") or str(uuid.uuid4()), results=rerank_results, - meta=rerank_meta, + meta=TypeAdapter(RerankResponseMeta).validate_python(response["meta"]) + if "meta" in response and response["meta"] is not None + else rerank_meta, ) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 15e348453e0..db8dd3fa7f0 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -25,15 +25,16 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_internal_message_fields, strip_none_values_from_message, ) +from litellm.litellm_core_utils.reasoning_effort_capability import ( + declared_reasoning_efforts_for_model, + nearest_declared_reasoning_effort, +) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) -from litellm.router_utils.reasoning_effort_capability import ( - declared_reasoning_efforts_for_model, - nearest_declared_reasoning_effort, -) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.mistral import ( MistralCacheKeyParams, @@ -74,6 +75,12 @@ _MISTRAL_REQUEST_ROUTING: Final = TypeAdapter(MistralRequestRouting) _MISTRAL_CACHE_KEY_PARAMS: Final = TypeAdapter(MistralCacheKeyParams) _MISTRAL_REASONING_EFFORT_PARAMS: Final = TypeAdapter(MistralReasoningEffortParams) _MISTRAL_FILE_PART: Final = TypeAdapter(MistralFilePart) +_REASONING_PROMPT_REASON: Final = "litellm sends Magistral a reasoning system prompt instead of the requested level." +_NON_LEVEL_EFFORT_REASON: Final = ( + "the Mistral API takes it only as a level such as 'high', so Mistral applies its default effort." +) +_DOCUMENTED_MAGISTRAL_EFFORTS: Final = frozenset({"low", "medium", "high"}) +_REASONING_ON_THINKING_TYPES: Final = frozenset({"enabled", "adaptive"}) def _reasoning_level_and_summary(reasoning_effort: object) -> tuple[object, object]: @@ -83,6 +90,16 @@ def _reasoning_level_and_summary(reasoning_effort: object) -> tuple[object, obje return requested.get("effort"), requested.get("summary") +def _warn_unless_reasoning_prompt_honors(param: str, value: object, model: str) -> None: + if param == "reasoning_effort" and isinstance(value, str) and value in _DOCUMENTED_MAGISTRAL_EFFORTS: + return + thinking: Final = as_str_mapping(value) + thinking_type: Final = thinking.get("type") if thinking is not None else None + if param == "thinking" and isinstance(thinking_type, str) and thinking_type in _REASONING_ON_THINKING_TYPES: + return + warn_dropped_reasoning(param, value, model, _REASONING_PROMPT_REASON, drop_params_gated=False) + + def _is_json_array(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a decoded JSON array is a list return isinstance(value, list) @@ -253,6 +270,7 @@ class MistralConfig(OpenAIGPTConfig): if param in ("reasoning_effort", "thinking") and "magistral" in model.lower(): # Flag that we need to add reasoning system prompt optional_params["_add_reasoning_prompt"] = True + _warn_unless_reasoning_prompt_honors(param, OPAQUE_VALUE.validate_python(value), model) if param in ("parallel_tool_calls", "prompt_cache_key"): optional_params[param] = value return optional_params @@ -260,13 +278,18 @@ class MistralConfig(OpenAIGPTConfig): def _reasoning_level_to_send( self, model: str, requested: MistralReasoningEffortParams, drop_params: bool ) -> str | None: - level, summary = _reasoning_level_and_summary(requested.get("reasoning_effort")) + requested_effort: Final = requested.get("reasoning_effort") + level, summary = _reasoning_level_and_summary(requested_effort) sendable: Final = ( _accepted_reasoning_effort(model, level, self.custom_llm_provider) if isinstance(level, str) else None ) if sendable is not None and summary is None: return sendable if drop_params or litellm.drop_params is True: + if sendable is None and level is not None: + warn_dropped_reasoning( + "reasoning_effort", requested_effort, model, _NON_LEVEL_EFFORT_REASON, drop_params_gated=True + ) return sendable raise UnsupportedParamsError( message=( diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 8257a75b4a5..224674739f5 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -21,6 +21,8 @@ from litellm._logging import verbose_logger from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.prompt_templates.common_utils import messages_with_file_parts_as_image_urls +from litellm.litellm_core_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -53,7 +55,6 @@ from litellm.llms.oci.common_utils import ( sign_oci_request, validate_oci_environment, ) -from litellm.router_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model from litellm.types.llms.oci import ( CohereChatRequest, OCIChatRequestPayload, @@ -428,6 +429,13 @@ class OCIChatConfig(BaseConfig): alias = param_map.get(key) if key == "reasoning_effort" and alias and not _model_supports_reasoning_effort(model): if drop_params or litellm.drop_params: + warn_dropped_reasoning( + key, + value, + model, + "the model map lists no reasoning_effort levels for it, so OCI would reject reasoningEffort.", + drop_params_gated=True, + ) continue raise OCIError( status_code=400, @@ -451,6 +459,14 @@ class OCIChatConfig(BaseConfig): if key == "n" and (value is None or value == 1): continue if drop_params or litellm.drop_params: + if key == "reasoning_effort": + warn_dropped_reasoning( + key, + value, + model, + "OCI's Cohere chat API has no reasoning effort field.", + drop_params_gated=True, + ) continue raise OCIError( status_code=400, diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index f9e7071db41..168b562b358 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_images_from_message, extract_reasoning_content, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -29,7 +30,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ModelResponse, ModelResponseStream -from ..common_utils import OllamaError +from ..common_utils import OllamaError, ollama_think if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -172,10 +173,7 @@ class OllamaChatConfig(BaseConfig): if value.get("json_schema") and value["json_schema"].get("schema"): optional_params["format"] = value["json_schema"]["schema"] if param == "reasoning_effort" and value is not None: - if model.startswith("gpt-oss"): - optional_params["think"] = value - else: - optional_params["think"] = value in {"low", "medium", "high"} + optional_params["think"] = ollama_think(model, OPAQUE_VALUE.validate_python(value)) ### FUNCTION CALLING LOGIC ### # Ollama 0.4+ supports native tool calling - pass tools directly # and let Ollama handle model capability detection @@ -237,6 +235,19 @@ class OllamaChatConfig(BaseConfig): litellm_params["function_name"] = function_name tools: Final = optional_params.pop("tools", None) + def transform_tool_call(typed_tool: ChatCompletionAssistantToolCall) -> OllamaToolCall: + function: Final = OllamaToolCallFunction( + name=typed_tool["function"].get("name") or "", + arguments=json.loads(typed_tool["function"]["arguments"]) + if "arguments" in typed_tool["function"] + else {}, + ) + tool_id: Final = typed_tool.get("id") + ollama_tool_call: Final = ( + OllamaToolCall(function=function) if tool_id is None else OllamaToolCall(id=tool_id, function=function) + ) + return ollama_tool_call + new_messages: Final = [] for m in messages: if isinstance( @@ -250,16 +261,7 @@ class OllamaChatConfig(BaseConfig): for tool in tool_calls: typed_tool = ChatCompletionAssistantToolCall(**tool) if typed_tool["type"] == "function": - arguments = {} - if "arguments" in typed_tool["function"]: - arguments = json.loads(typed_tool["function"]["arguments"]) - ollama_tool_call = OllamaToolCall( - function=OllamaToolCallFunction( - name=typed_tool["function"].get("name") or "", - arguments=arguments, - ) - ) - new_tools.append(ollama_tool_call) + new_tools.append(transform_tool_call(typed_tool)) reasoning_content, parsed_content = extract_reasoning_content(cast(dict, m)) content_str = convert_content_list_to_str(cast(AllMessageValues, m)) images = extract_images_from_message(cast(AllMessageValues, m)) diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index f313c62f026..fcf6a396a13 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -8,11 +8,14 @@ from pydantic import BaseModel, TypeAdapter from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ModelResponse, ModelResponseStream _PROVIDER_FIELDS: Final = TypeAdapter(dict[str, object]) +_THINK_ON_EFFORTS: Final = frozenset({"low", "medium", "high"}) +_THINK_OFF_REASON: Final = "Ollama takes think only as on or off, and litellm maps this level to off." _DECODED_JSON: Final = TypeAdapter(object) @@ -296,3 +299,12 @@ class OllamaModelInfo(BaseLLMModelInfo): Return the base model name for Ollama (no-op). """ return model + + +def ollama_think(model: str, reasoning_effort: object) -> object: + if model.startswith("gpt-oss"): + return reasoning_effort + think: Final = reasoning_effort in _THINK_ON_EFFORTS + if isinstance(reasoning_effort, str) and not think and reasoning_effort != "none": + warn_dropped_reasoning("reasoning_effort", reasoning_effort, model, _THINK_OFF_REASON, drop_params_gated=False) + return think diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 239b4e32645..6dbcb050e11 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, inline_remote_image_urls, ) +from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -43,7 +44,7 @@ from litellm.types.utils import ( generate_id, ) -from ..common_utils import OllamaError, OllamaModelInfo, convert_image +from ..common_utils import OllamaError, OllamaModelInfo, convert_image, ollama_think if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -249,10 +250,7 @@ class OllamaConfig(BaseConfig): elif param == "stop": optional_params["stop"] = value elif param == "reasoning_effort" and value is not None: - if model.startswith("gpt-oss"): - optional_params["think"] = value - else: - optional_params["think"] = value in {"low", "medium", "high"} + optional_params["think"] = ollama_think(model, OPAQUE_VALUE.validate_python(value)) elif param == "response_format" and isinstance(value, dict): if value["type"] == "json_object": optional_params["format"] = "json" diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index aa14a985afe..c113cb351d8 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -32,6 +32,7 @@ import litellm from litellm.exceptions import error_response_without_credentials from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning_params from litellm.litellm_core_utils.token_counter import token_counter, usage_tokenizer_model from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -165,17 +166,11 @@ _DECODED_JSON: Final = TypeAdapter(object) def drop_params_from_unprocessable_entity_error( e: openai.UnprocessableEntityError | httpx.HTTPStatusError, data: Mapping[str, object], + *, + model: str | None, + drop_params_gated: bool, ) -> dict[str, object]: - """ - Helper function to read OpenAI UnprocessableEntityError and drop the params that raised an error from the error message. - - Args: - e (UnprocessableEntityError): The UnprocessableEntityError exception - data (Dict[str, Any]): The original data dictionary containing all parameters - - Returns: - Dict[str, Any]: A new dictionary with invalid parameters removed - """ + """Return a copy of ``data`` without the params a 422 response names as invalid.""" invalid_params: Final[list[str]] = [] if isinstance(e, httpx.HTTPStatusError): error_json: Final = decoded_json_object(_DECODED_JSON.validate_python(e.response.json())) @@ -203,6 +198,13 @@ def drop_params_from_unprocessable_entity_error( invalid_params.append(error_dict["loc"][1]) new_data: Final = {k: v for k, v in data.items() if k not in invalid_params} + warn_dropped_reasoning_params( + invalid_params, + data, + model, + "the provider answered 422 naming it an invalid param.", + drop_params_gated=drop_params_gated, + ) return new_data diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index bef99fd0719..6e22ef9112e 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -16,6 +16,7 @@ from openai._base_client import make_request_options from openai._constants import RAW_RESPONSE_HEADER from openai._legacy_response import LegacyAPIResponse from openai._types import RequestOptions +from openai.pagination import SyncCursorPage from openai.types import CreateEmbeddingResponse from openai.types.beta.assistant_deleted import AssistantDeleted from openai.types.chat import ChatCompletion, ChatCompletionChunk @@ -40,6 +41,7 @@ from litellm.types.utils import ( LiteLLMBatch, ModelResponse, ModelResponseStream, + openai_batch_status, ) from litellm.utils import ( CustomStreamWrapper, @@ -940,7 +942,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.UnprocessableEntityError as e: ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800 if litellm.drop_params is True or drop_params is True: - inference_params = drop_params_from_unprocessable_entity_error(e, inference_params) + inference_params = drop_params_from_unprocessable_entity_error( + e, inference_params, model=model, drop_params_gated=True + ) else: raise e # e.message @@ -1097,7 +1101,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.UnprocessableEntityError as e: ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800 if litellm.drop_params is True or drop_params is True: - data = drop_params_from_unprocessable_entity_error(e, data) + data = drop_params_from_unprocessable_entity_error(e, data, model=model, drop_params_gated=True) else: raise e # e.message @@ -1252,7 +1256,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.UnprocessableEntityError as e: ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800 if litellm.drop_params is True or drop_params is True: - data = drop_params_from_unprocessable_entity_error(e, data) + data = drop_params_from_unprocessable_entity_error(e, data, model=model, drop_params_gated=True) else: raise e except ( @@ -2180,6 +2184,18 @@ class OpenAIFilesAPI(BaseLLM): return response +def _with_openai_batch_fields(batches: Sequence[Batch]) -> Sequence[Batch]: + return [ + batch.model_copy( + update={ + "object": batch.object or "batch", + "status": openai_batch_status(batch.status), + } + ) + for batch in batches + ] + + class OpenAIBatchesAPI(BaseLLM): """ OpenAI methods to support for batches @@ -2370,7 +2386,7 @@ class OpenAIBatchesAPI(BaseLLM): ): verbose_logger.debug("listing batches, after= %s, limit= %s", after, limit) response: Final = await openai_client.batches.list(after=after, limit=limit) - return response + return response.model_copy(update={"data": _with_openai_batch_fields(response.data)}) def list_batches( self, @@ -2382,7 +2398,7 @@ class OpenAIBatchesAPI(BaseLLM): organization: str | None, after: str | None = None, limit: int | None = None, - client: OpenAI | None = None, + client: OpenAI | AsyncOpenAI | None = None, litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( @@ -2407,7 +2423,9 @@ class OpenAIBatchesAPI(BaseLLM): ) return self.alist_batches(openai_client=openai_client, after=after, limit=limit) response: Final = openai_client.batches.list(after=after, limit=limit) - return response + if not isinstance(response, SyncCursorPage): + return response + return response.model_copy(update={"data": _with_openai_batch_fields(response.data)}) class OpenAIAssistantsAPI(BaseLLM): diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index 9aaf552e79a..69f3d76c034 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -13,7 +13,8 @@ from typing import TYPE_CHECKING, Final, cast import httpx from pydantic import ConfigDict, TypeAdapter -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.core_helpers import finite_cost +from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, set_hidden_params from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -211,22 +212,33 @@ class OpenrouterConfig(OpenAIGPTConfig): # OpenRouter returns cost information in the usage object when usage.include=true try: response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) + response_provider: Final = response_json.get("provider") + if isinstance(response_provider, str) and response_provider: + set_hidden_params( + model_response, + {**(get_hidden_params(model_response) or {}), "upstream_provider": response_provider}, + ) if "usage" in response_json and response_json["usage"]: response_cost: Final = _JSON_OBJECT.validate_python(response_json["usage"]).get("cost") if response_cost is not None: # Store cost in hidden params for the cost calculator to use if not hasattr(model_response, HIDDEN_PARAMS_ATTR): set_hidden_params(model_response, {}) - hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic response metadata - dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) + cost_hidden_params: Final = ( + cast( # cast-ok: preserve mapping operations on dynamic response metadata + dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) + ) ) - if "additional_headers" not in hidden_params: - hidden_params["additional_headers"] = {} + if "additional_headers" not in cost_hidden_params: + cost_hidden_params["additional_headers"] = {} additional_headers: Final = cast( # cast-ok: preserve mapping operations on response metadata - dict[str, object], hidden_params["additional_headers"] + dict[str, object], cost_hidden_params["additional_headers"] ) - if isinstance(response_cost, str | int | float): - additional_headers["llm_provider-x-litellm-response-cost"] = float(response_cost) + reported_cost: Final = ( + finite_cost(float(response_cost)) if isinstance(response_cost, str | int | float) else None + ) + if reported_cost is not None: + additional_headers["llm_provider-x-litellm-response-cost"] = reported_cost except Exception: # If we can't extract cost, continue without it - don't fail the response pass @@ -278,7 +290,7 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): for choice in chunk["choices"]: choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning") new_choices.append(choice) - return ModelResponseStream( + model_response: Final = ModelResponseStream( id=chunk["id"], object="chat.completion.chunk", created=chunk["created"], @@ -286,6 +298,12 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): model=chunk["model"], choices=new_choices, ) + typed_chunk: Final[Mapping[str, object]] = cast( # cast-ok: stream chunks are decoded JSON objects + Mapping[str, object], chunk + ) + upstream_provider: Final[object] = typed_chunk.get("provider") + if isinstance(upstream_provider, str) and upstream_provider: + set_hidden_params(model_response, {"upstream_provider": upstream_provider}) except KeyError as e: raise OpenRouterException( message=f"KeyError: {e}, Got unexpected response from OpenRouter: {chunk}", @@ -294,3 +312,5 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): ) except Exception as e: raise e + else: + return model_response diff --git a/litellm/llms/openrouter/common_utils.py b/litellm/llms/openrouter/common_utils.py index 96e53a5aae7..64caf83d745 100644 --- a/litellm/llms/openrouter/common_utils.py +++ b/litellm/llms/openrouter/common_utils.py @@ -1,5 +1,92 @@ +from collections.abc import Mapping, Sequence +from typing import Final + +import httpx +from pydantic import BaseModel, ConfigDict + +import litellm +from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues + +_DEFAULT_API_BASE: Final = "https://openrouter.ai/api/v1" +_CHAT_CATALOG_ROUTE: Final = "models/user" +_CHAT_COMPLETIONS_PATH: Final = "/chat/completions" +_BATCH_ONLY_SUFFIX: Final = ":batch" +_CATALOG_TIMEOUT: Final = httpx.Timeout(10.0) class OpenRouterException(BaseLLMException): pass + + +class _CatalogModel(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + + +class _Catalog(BaseModel): + model_config = ConfigDict(frozen=True) + + data: tuple[_CatalogModel, ...] + + +def _chat_catalog_url(api_base: str) -> str: + return join_api_base(api_base.rstrip("/").removesuffix(_CHAT_COMPLETIONS_PATH), _CHAT_CATALOG_ROUTE) + + +class OpenRouterModelInfo(BaseLLMModelInfo): + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return ( + api_key + or litellm.api_key + or litellm.openrouter_key + or get_secret_str("OPENROUTER_API_KEY") + or get_secret_str("OR_API_KEY") + ) + + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + return api_base or litellm.api_base or get_secret_str("OPENROUTER_API_BASE") or _DEFAULT_API_BASE + + @staticmethod + def get_base_model(model: str) -> str: + return model.removeprefix("openrouter/") + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseLLMModelInfo.validate_environment returns a dict + resolved_key: Final = self.get_api_key(api_key) + if resolved_key is None: + return {**headers} + return {**headers, "Authorization": f"Bearer {resolved_key}"} + + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + resolved_key: Final = self.get_api_key(api_key) + if resolved_key is None: + raise ValueError("Set OPENROUTER_API_KEY or the deployment's api_key to list OpenRouter models") + headers: Final = self.validate_environment( + headers={}, model="", messages=(), optional_params={}, litellm_params={}, api_key=resolved_key + ) + response: Final = litellm.module_level_client.get( # pyright: ignore[reportUnknownMemberType] # legacy HTTP wrapper has untyped query parameters + url=_chat_catalog_url(self.get_api_base(api_base)), headers=headers, timeout=_CATALOG_TIMEOUT + ) + response.raise_for_status() + return [ + f"openrouter/{model.id}" + for model in _Catalog.model_validate_json(response.content).data + if not model.id.endswith(_BATCH_ONLY_SUFFIX) + ] diff --git a/litellm/llms/pass_through/typesafe_logging_handler.py b/litellm/llms/pass_through/typesafe_logging_handler.py index 9eae57774c2..7a41d2a5d14 100644 --- a/litellm/llms/pass_through/typesafe_logging_handler.py +++ b/litellm/llms/pass_through/typesafe_logging_handler.py @@ -6,6 +6,7 @@ import httpx from pydantic import TypeAdapter, ValidationError import litellm +from litellm.litellm_core_utils.core_helpers import finite_reported_cost from litellm.litellm_core_utils.decisions_pricing import apply_decisions_pricing from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( @@ -100,9 +101,10 @@ class TypeSafePassthroughLoggingHandler: f"{custom_llm_provider}/{model}" for model in candidate_models if model is not None ) pricing: Final = _pricing_for(candidate_model_keys) + reported_cost: Final = finite_reported_cost(usage.cost) if custom_llm_provider == "openrouter" else None response_cost: Final = ( - usage.cost - if custom_llm_provider == "openrouter" and usage.cost is not None + reported_cost + if reported_cost is not None else input_tokens * pricing.input_cost_per_token + output_tokens * pricing.output_cost_per_token ) usage_object: Final = Usage( diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index 2d762adf958..c40a86b221b 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -6,6 +6,7 @@ Helper util for handling perplexity-specific cost calculation from datetime import datetime from typing import Final +from litellm.litellm_core_utils.core_helpers import finite_reported_cost from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing from litellm.types.utils import Usage from litellm.utils import get_model_info @@ -34,10 +35,9 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non total_cost = cost_info.get("total_cost") elif isinstance(cost_info, (int, float)) and not isinstance(cost_info, bool): total_cost = float(cost_info) - if total_cost is not None: - # Return total cost as completion_cost (prompt_cost=0) since Perplexity - # doesn't break down by input/output in their cost object - return (0.0, float(total_cost)) + reported_cost: Final = None if total_cost is None else finite_reported_cost(float(total_cost)) + if reported_cost is not None: + return (0.0, reported_cost) ## FALLBACK: Calculate cost manually if Perplexity doesn't provide it ## GET MODEL INFO @@ -100,8 +100,9 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non num_search_queries = getattr(usage.prompt_tokens_details, "web_search_requests", 0) or 0 # Check both possible keys for search cost (legacy and current) - search_cost_value: Final = model_info.get("search_queries_cost_per_query") or model_info.get( - "search_context_cost_per_query" + legacy_search_cost: Final = model_info.get("search_queries_cost_per_query") + search_cost_value: Final = ( + model_info.get("search_context_cost_per_query") if legacy_search_cost is None else legacy_search_cost ) if num_search_queries > 0 and search_cost_value is not None: # Handle both dict and float formats diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index 1f3b0fe729d..0921a892849 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -84,9 +84,9 @@ class PerplexitySearchConfig(BaseSearchConfig): def transform_search_request( self, query: str | list[str], - optional_params: dict, - **kwargs, - ) -> dict: + optional_params: dict[str, object], + **kwargs: object, + ) -> dict[str, object]: """ Transform Search request to Perplexity API format. @@ -98,11 +98,7 @@ class PerplexitySearchConfig(BaseSearchConfig): Args: query: Search query (string or list of strings) - optional_params: Optional parameters for the request - - max_results: Maximum number of search results (1-20) - - search_domain_filter: List of domains to filter (max 20) - - max_tokens_per_page: Max tokens per page (default 1024) - - country: Country code filter (e.g., 'US', 'GB', 'DE') + optional_params: Provider parameters, forwarded unchanged when not None. Returns: Dict with typed request data following PerplexitySearchRequest spec @@ -111,24 +107,10 @@ class PerplexitySearchConfig(BaseSearchConfig): "query": query, } - # Add optional parameters following Perplexity API spec (only if not None) - max_results: Final = optional_params.get("max_results") - if max_results is not None: - request_data["max_results"] = max_results - - search_domain_filter: Final = optional_params.get("search_domain_filter") - if search_domain_filter is not None: - request_data["search_domain_filter"] = search_domain_filter - - max_tokens_per_page: Final = optional_params.get("max_tokens_per_page") - if max_tokens_per_page is not None: - request_data["max_tokens_per_page"] = max_tokens_per_page - - country: Final = optional_params.get("country") - if country is not None: - request_data["country"] = country - - return dict(request_data) + return { + **{key: value for key, value in optional_params.items() if value is not None}, + **request_data, + } def transform_search_response( self, diff --git a/litellm/llms/tencent/chat/transformation.py b/litellm/llms/tencent/chat/transformation.py index 1f0bcd218d6..23a6ccd52c6 100644 --- a/litellm/llms/tencent/chat/transformation.py +++ b/litellm/llms/tencent/chat/transformation.py @@ -9,12 +9,15 @@ from typing import Final, TypedDict from typing_extensions import ReadOnly import litellm +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.secret_managers.main import get_secret_str from litellm.utils import supports_reasoning from ...openai.chat.gpt_transformation import OpenAIGPTConfig +_ADAPTIVE_BUDGET_REASON: Final = "it is not an integer, so litellm sends adaptive thinking with no budget." + class ThinkingPayload(TypedDict, total=False): """Tencent TokenHub `thinking` object. @@ -102,6 +105,10 @@ class TencentChatConfig(OpenAIGPTConfig): if isinstance(budget, int): coerced_with_budget: Final[ThinkingPayload] = {"type": "adaptive", "budget_tokens": budget} return coerced_with_budget + if budget is not None: + warn_dropped_reasoning( + "thinking.budget_tokens", budget, model, _ADAPTIVE_BUDGET_REASON, drop_params_gated=False + ) coerced: Final[ThinkingPayload] = {"type": "adaptive"} return coerced diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py index a0eb0c75d35..02436109a4a 100644 --- a/litellm/llms/together_ai/chat/transformation.py +++ b/litellm/llms/together_ai/chat/transformation.py @@ -18,7 +18,7 @@ import litellm from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields -from litellm.router_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model +from litellm.litellm_core_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_function_calling, supports_reasoning, supports_response_schema diff --git a/litellm/llms/together_ai/rerank/transformation.py b/litellm/llms/together_ai/rerank/transformation.py index f1c3ca9e87d..eb7ec53416d 100644 --- a/litellm/llms/together_ai/rerank/transformation.py +++ b/litellm/llms/together_ai/rerank/transformation.py @@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works from typing import Final +from pydantic import TypeAdapter + from litellm._uuid import uuid from litellm.types.rerank import ( RerankBilledUnits, @@ -19,7 +21,7 @@ from litellm.types.rerank import ( class TogetherAIRerankConfig: def _transform_response(self, response: dict) -> RerankResponse: - _billed_units: Final = RerankBilledUnits(**response.get("usage", {})) + _billed_units: Final = TypeAdapter(RerankBilledUnits).validate_python(response.get("usage", {})) _tokens: Final = RerankTokens(**response.get("usage", {})) rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) diff --git a/litellm/llms/typesafe/decisions/transformation.py b/litellm/llms/typesafe/decisions/transformation.py index 952c874bc8e..871596b7f00 100644 --- a/litellm/llms/typesafe/decisions/transformation.py +++ b/litellm/llms/typesafe/decisions/transformation.py @@ -6,7 +6,7 @@ from types import MappingProxyType from typing import Final, Literal from litellm.constants import TYPESAFE_DANGER_ZONE_HEADER_PREFIX -from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, image_media_type +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits, image_media_type from litellm.types.decisions import ( DecisionsAsset, DecisionsIRMessages, @@ -19,6 +19,7 @@ from litellm.types.decisions import ( ) _IMAGES_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}images" +_LONG_CONTEXT_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}long-context" _ASSET_FORMATS: Final[Mapping[str, Literal["png", "jpg"]]] = MappingProxyType( {"image/png": "png", "image/jpeg": "jpg", "image/jpg": "jpg"} ) @@ -84,6 +85,16 @@ class TypeSafeDecisionsConfig(BaseDecisionsConfig): def get_default_api_base(self) -> str | None: return "https://api.typesafe.ai" + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + long_context_values: Final = tuple( + value.strip().lower() for name, value in headers.items() if name.lower() == _LONG_CONTEXT_HEADER + ) + if long_context_values and all(value == "true" for value in long_context_values): + return DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144) + return DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768) + def supports_images_in( self, decision_input: DecisionsIRState | DecisionsIRMessages, headers: Mapping[str, str] ) -> bool: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 0fa80cb2cad..ad64d852ef9 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1257,7 +1257,7 @@ def get_vertex_model_id_from_url(url: str) -> str | None: `https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/publishers/google/models/${MODEL_ID}:streamGenerateContent` """ - match: Final = re.search(r"/models/([^:]+)", url) + match: Final = re.search(r"/models/([^:?#]+)", url) return match.group(1) if match else None diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index c506192f253..1ef5267df14 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc from litellm.litellm_core_utils.prompt_templates.factory import ( encode_tool_call_id_with_signature, ) +from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -201,6 +202,13 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr _SAMPLING_DEPRECATION_UNMAPPED_MODELS: Final = "models outside litellm's model map" +_NON_INTEGER_BUDGET_REASON: Final = ( + "Gemini takes budget_tokens only as an integer, so it uses its default thinking budget." +) +_LEVEL_ONLY_BUDGET_REASON: Final = ( + "Gemini 3 and later set reasoning by thinking level, not budget, so litellm does not send the budget. " + "Send reasoning_effort to pick a level." +) _sampling_deprecation_lock: Final = threading.Lock() @@ -1466,6 +1474,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): params["includeThoughts"] = False else: params["includeThoughts"] = sends_thoughts + if thinking_budget is not None: + warn_dropped_reasoning( + "thinking.budget_tokens", + thinking_budget, + model, + _LEVEL_ONLY_BUDGET_REASON, + drop_params_gated=False, + ) # Follow provider defaults unless explicitly opted into legacy behavior. if litellm.enable_gemini_default_thinking_level_low is True: params["thinkingLevel"] = ( @@ -1478,8 +1494,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # For older Gemini models, use thinkingBudget if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero(thinking_budget): params["includeThoughts"] = sends_thoughts - if thinking_budget is not None and isinstance(thinking_budget, int): + if isinstance(thinking_budget, int): params["thinkingBudget"] = thinking_budget + elif thinking_budget is not None: + warn_dropped_reasoning( + "thinking.budget_tokens", + thinking_budget, + model, + _NON_INTEGER_BUDGET_REASON, + drop_params_gated=False, + ) return params diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py index da82766c2b9..98fa8ed2b77 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py @@ -1,9 +1,15 @@ from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter + from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params -from litellm.llms.anthropic.common_utils import AnthropicModelInfo, context_management_betas +from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, + context_management_betas, + requires_native_compaction_beta, +) from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, messages_carry_output_config, @@ -11,6 +17,7 @@ from litellm.llms.anthropic.messages.native.transformation import ( from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS, + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, ) from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.types.llms.vertex_ai import VertexPartnerProvider @@ -19,6 +26,8 @@ from litellm.types.router import GenericLiteLLMParams from ....vertex_llm_base import VertexBase from ..output_params_utils import sanitize_vertex_anthropic_output_params +_OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object]) + _VERSION_HEADER_VERTEX_READS_FROM_BODY: Final = "anthropic-version" _VERTEX_DEFAULT_VERSION: Final = "vertex-2023-10-16" _VERSIONS_SERVED_AS_VERTEX_DEFAULT: Final = frozenset({"", "2023-01-01", "2023-06-01"}) @@ -190,14 +199,23 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert if messages_carry_output_config(messages): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value) - beta_values.update( + typed_params: Final = _OPTIONAL_PARAMS.validate_python(optional_params) + if requires_native_compaction_beta("vertex_ai", typed_params, messages): + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value) + + tool_change_betas: Final = ( + (ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,) + if anthropic_model_info.is_mid_conversation_tool_change_used(messages) + else () + ) + all_beta_values: Final = beta_values.union( anthropic_model_info.thinking_betas( optional_params.get("thinking"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # optional_params is the base config's bare dict - ) + ), + tool_change_betas, ) - - if beta_values: - request_headers["anthropic-beta"] = ",".join(beta_values) + if all_beta_values: + request_headers["anthropic-beta"] = ",".join(sorted(all_beta_values)) return request_headers, api_base diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 79fd9f9e8b1..3d1690fca33 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -2,19 +2,23 @@ ## Handler file for calling claude-3 on vertex ai import base64 import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final import httpx -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm +from litellm.anthropic_beta_headers_manager import get_header_only_betas +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, inline_remote_image_urls from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse from ....anthropic.chat.transformation import AnthropicConfig +from ....anthropic.common_utils import AnthropicError, requires_native_compaction_beta from .output_params_utils import sanitize_vertex_anthropic_output_params if TYPE_CHECKING: @@ -36,6 +40,32 @@ def _vertex_labels_header(labels: object, header_names: Mapping[str, object]) -> return base64.b64encode(json.dumps(string_labels).encode()).decode() if string_labels else None +_BETA_VALUE: Final[TypeAdapter[str | tuple[str, ...] | None]] = TypeAdapter(str | tuple[str, ...] | None) + +_REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object]) + + +def _beta_names(value: object) -> frozenset[str] | None: + try: + parsed: Final = _BETA_VALUE.validate_python(value) + except ValidationError: + return None + names: Final = parsed.split(",") if isinstance(parsed, str) else parsed or () + return frozenset(name.strip() for name in names if name.strip()) + + +def _client_betas(value: object, field: str, drop_params: bool) -> frozenset[str]: + names: Final = _beta_names(value) + if names is not None: + return names + if drop_params: + return frozenset() + raise AnthropicError( + status_code=400, + message=f"{field} must be a string or a list of strings, got {type(value).__name__}", + ) + + class VertexAIError(Exception): def __init__(self, status_code, message): self.status_code = status_code @@ -157,21 +187,43 @@ class VertexAIAnthropicConfig(AnthropicConfig): if context_management: self._add_context_management_beta_headers(beta_set, context_management) - extra_headers: Final = optional_params.get("extra_headers") or {} - anthropic_beta_value: Final = extra_headers.get("anthropic-beta", "") - if isinstance(anthropic_beta_value, str) and anthropic_beta_value: - for beta in anthropic_beta_value.split(","): - beta = beta.strip() - if beta: - beta_set.add(beta) - elif isinstance(anthropic_beta_value, list): - beta_set.update(anthropic_beta_value) - + request_params: Final = _NON_DEFAULT_PARAMS.validate_python(optional_params) + compaction_betas: Final = ( + frozenset({ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value}) + if requires_native_compaction_beta( + "vertex_ai", request_params, _REQUEST_MESSAGES.validate_python(data["messages"]) + ) + else frozenset[str]() + ) + drop_params: Final = drops_unsupported_params(litellm_params.get("drop_params")) + extra_headers: Final = _NON_DEFAULT_PARAMS.validate_python(request_params.get("extra_headers") or {}) data.pop("extra_headers", None) + injected_betas: Final = frozenset(beta_set) | _client_betas( + extra_headers.get("anthropic-beta"), "The anthropic-beta header", drop_params + ) + client_body_betas: Final = _client_betas( + _NON_DEFAULT_PARAMS.validate_python(data).get("anthropic_beta"), "anthropic_beta", drop_params + ) - if beta_set: - data["anthropic_beta"] = list(beta_set) - headers["anthropic-beta"] = ",".join(beta_set) + requested_betas: Final = injected_betas | client_body_betas + header_only_betas: Final = get_header_only_betas("vertex_ai") if requested_betas else frozenset[str]() + body_betas: Final = requested_betas - header_only_betas + if body_betas: + data["anthropic_beta"] = sorted(body_betas) + elif "anthropic_beta" in data: + del data["anthropic_beta"] + header_betas: Final = ( + injected_betas + | compaction_betas + | (client_body_betas & header_only_betas) + | _client_betas( + _NON_DEFAULT_PARAMS.validate_python(headers).get("anthropic-beta"), + "The anthropic-beta header", + drop_params, + ) + ) + if header_betas: + headers["anthropic-beta"] = ",".join(sorted(header_betas)) return data diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 180761f15d9..e341c8f5f48 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -6,12 +6,19 @@ Unlike Gemini models which use Google's token counting API, partner models use their respective publisher-specific count-tokens endpoints. """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Any, Final +from pydantic import TypeAdapter + from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.anthropic.common_utils import OBJECT_LIST_ADAPTER, is_claude_model_target, tools_without_custom_field +from litellm.llms.anthropic.common_utils import ( + OBJECT_LIST_ADAPTER, + count_tokens_feature_betas, + is_claude_model_target, + tools_without_custom_field, +) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -23,6 +30,9 @@ _COUNT_TOKENS_PUBLISHERS: Final = ( ) +_REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object]) + + def count_tokens_request_body( request_data: Mapping[str, object], model: str, drop_params: object ) -> Mapping[str, object]: @@ -43,6 +53,14 @@ def count_tokens_publisher(model: str) -> str | None: return next((publisher for marker, publisher in _COUNT_TOKENS_PUBLISHERS if marker in model), None) +def count_tokens_headers(access_token: str, request_data: Mapping[str, object]) -> Mapping[str, str]: + authorization: Final = {"Authorization": f"Bearer {access_token}"} + betas: Final = count_tokens_feature_betas( + "vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"]) + ) + return {**authorization, "anthropic-beta": ",".join(betas)} if betas else authorization + + class VertexAIPartnerModelsTokenCounter(VertexBase): """ Token counter for Vertex AI Partner Models. @@ -175,8 +193,7 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): api_base=litellm_params.get("api_base"), ) - # Prepare headers - headers: Final = {"Authorization": f"Bearer {access_token}"} + headers: Final = dict(count_tokens_headers(access_token, request_data)) # Get async HTTP client from litellm import LlmProviders diff --git a/litellm/llms/vllm/common_utils.py b/litellm/llms/vllm/common_utils.py index 6484cdc88e3..014fabd5187 100644 --- a/litellm/llms/vllm/common_utils.py +++ b/litellm/llms/vllm/common_utils.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import httpx @@ -10,6 +11,14 @@ from litellm.types.llms.openai import AllMessageValues from litellm.utils import add_path_to_api_base +def endpoint_under_api_base(endpoint: str, api_base: str | None) -> str: + return ( + endpoint.lstrip("/").removeprefix("v1/") + if api_base is not None and api_base.rstrip("/").endswith("/v1") + else endpoint + ) + + class VLLMError(BaseLLMException): def __init__( self, @@ -39,9 +48,14 @@ class VLLMModelInfo(BaseLLMModelInfo): api_key: str | None = None, api_base: str | None = None, ) -> dict: - if api_key is not None: - headers["x-api-key"] = api_key - return headers + upstream_key: Final = self.upstream_api_key(api_key, api_base) + return {**headers, **self.upstream_auth_headers(upstream_key)} if upstream_key else headers + + def upstream_api_key(self, api_key: str | None, api_base: str | None) -> str | None: + return api_key + + def upstream_auth_headers(self, api_key: str) -> Mapping[str, str]: + return {"Authorization": f"Bearer {api_key}", "x-api-key": api_key} @staticmethod def get_api_base(api_base: str | None = None) -> str | None: diff --git a/litellm/llms/vllm/passthrough/transformation.py b/litellm/llms/vllm/passthrough/transformation.py index e7da6e54410..2b253e7fa38 100644 --- a/litellm/llms/vllm/passthrough/transformation.py +++ b/litellm/llms/vllm/passthrough/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Final from httpx import Response +from litellm.constants import PLACEHOLDER_API_KEY from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.openai_like.passthrough.transformation import ( @@ -12,8 +13,14 @@ from litellm.llms.openai_like.passthrough.transformation import ( logged_text_completion_stream, relays_text_completion, ) +from litellm.secret_managers.main import ( + GOOGLE_OIDC_SECRET_PREFIX, + UnresolvableSecretReference, + get_secret_str, + resolve_local_oidc_reference, +) -from ..common_utils import VLLMModelInfo +from ..common_utils import VLLMModelInfo, endpoint_under_api_base if TYPE_CHECKING: from httpx import URL @@ -22,6 +29,30 @@ if TYPE_CHECKING: class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig): + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return None if api_key == PLACEHOLDER_API_KEY else api_key + + def upstream_api_key(self, api_key: str | None, api_base: str | None) -> str | None: + if api_key is not None or not self._is_vllm_api_base(api_base): + return api_key + vllm_api_key: Final = get_secret_str("VLLM_API_KEY") + if vllm_api_key is None or vllm_api_key.startswith(GOOGLE_OIDC_SECRET_PREFIX): + return None + match resolve_local_oidc_reference(vllm_api_key): + case UnresolvableSecretReference(): + return None + case str() as resolved: + return resolved + + @staticmethod + def _is_vllm_api_base(api_base: str | None) -> bool: + vllm_api_base: Final = get_secret_str("VLLM_API_BASE") + return api_base is not None and vllm_api_base is not None and api_base.rstrip("/") == vllm_api_base.rstrip("/") + + def upstream_auth_headers(self, api_key: str) -> Mapping[str, str]: + return {"Authorization": f"Bearer {api_key}"} + def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in request_data @@ -40,7 +71,7 @@ class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig): raise Exception("VLLM api base not found") return ( - self.format_url(endpoint, base_target_url, request_query_params), + self.format_url(endpoint_under_api_base(endpoint, base_target_url), base_target_url, request_query_params), base_target_url, ) diff --git a/litellm/main.py b/litellm/main.py index 5346741e4d1..2320d1a3e01 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -113,6 +113,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( messages_without_pdf_file_parts, ) from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header +from litellm.litellm_core_utils.reasoning_effort_capability import ( + model_group_from_metadata, + reasoning_effort_needs_allowing, + resolve_chat_reasoning, +) from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) @@ -144,11 +149,6 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.realtime_api.main import realtime_health_check from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway -from litellm.router_utils.reasoning_effort_capability import ( - model_group_from_metadata, - reasoning_effort_needs_allowing, - resolve_chat_reasoning, -) from litellm.secret_managers.main import get_secret_bool, get_secret_str from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper from litellm.types.completion import ( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0799a298942..cbe502e5cae 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -8715,7 +8715,7 @@ "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -8742,7 +8742,11 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, @@ -8763,7 +8767,7 @@ "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -8790,7 +8794,11 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 }, "azure/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, @@ -23070,6 +23078,24 @@ "output_cost_per_token_batches": 1e-06, "source": "https://developers.openai.com/api/docs/pricing" }, + "deepgram/aura": { + "input_cost_per_character": 1.5e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0150/1000 characters = $0.000015 per character", + "note": "Aura-1 text-to-speech voices (aura--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, + "deepgram/aura-2": { + "input_cost_per_character": 3e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.030/1000 characters = $0.00003 per character", + "note": "Aura-2 text-to-speech voices (aura-2--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, "deepgram/base": { "input_cost_per_second": 0.00020833, "litellm_provider": "deepgram", @@ -25276,6 +25302,38 @@ "cache_creation_input_token_cost": 1e-06, "prompt_cache_min_tokens": 2048 }, + "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_1hr": 2.2e-06, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "deprecation_date": "2026-10-15", + "litellm_provider": "bedrock_converse", + "supports_tool_search": true, + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06, + "token_count_model": { + "model": "anthropic.claude-haiku-4-5-20251001-v1:0", + "litellm_provider": "bedrock" + } + }, "eu.anthropic.claude-3-5-sonnet-20240620-v1:0": { "input_cost_per_token": 3e-06, "litellm_provider": "bedrock", @@ -29553,6 +29611,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30159,6 +30218,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_pdf_input": true, @@ -30202,6 +30262,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30266,6 +30327,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30326,6 +30388,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30393,6 +30456,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -30856,8 +30920,9 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/embeddings#model-versions", - "tpm": 10000000 + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "tpm": 10000000, + "input_cost_per_token_batches": 7.5e-08 }, "gemini/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, @@ -31209,6 +31274,9 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, "rpm": 100000, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ @@ -31389,6 +31457,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "rpm": 1000, @@ -31713,7 +31782,10 @@ "input_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, - "supports_url_context": true + "supports_url_context": true, + "cache_read_input_token_cost_above_200k_tokens_batches": 2.5e-07, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "deprecation_date": "2026-07-28", @@ -31871,7 +31943,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "supports_computer_use": true }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -31937,7 +32010,8 @@ "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, - "supports_audio_input": true + "supports_audio_input": true, + "supports_computer_use": true }, "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -32184,7 +32258,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "supports_computer_use": true }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -32280,7 +32355,10 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "output_cost_per_token_flex": 6e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06 }, "gemini/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -32343,7 +32421,10 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "output_cost_per_token_flex": 6e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06 }, "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -32384,6 +32465,7 @@ "text" ], "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -32458,6 +32540,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32522,6 +32605,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32582,6 +32666,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32649,6 +32734,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -47603,6 +47689,24 @@ "output_cost_per_token": 0.0, "source": "https://cohere.com/pricing" }, + "rerank-v5.0-fast": { + "input_cost_per_token": 2e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, + "rerank-v5.0-pro": { + "input_cost_per_token": 5e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -51283,6 +51387,11 @@ "max_input_tokens": 200000, "max_output_tokens": 64000, "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -51311,6 +51420,11 @@ "max_input_tokens": 200000, "max_output_tokens": 64000, "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -52057,6 +52171,11 @@ "max_input_tokens": 200000, "max_output_tokens": 64000, "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -52162,6 +52281,11 @@ "max_input_tokens": 200000, "max_output_tokens": 64000, "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -52692,6 +52816,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -59693,12 +59818,13 @@ "cache_read_input_token_cost": 1.375e-06, "output_cost_per_token": 8.25e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -72228,6 +72354,7 @@ }, "aiand/deepseek-ai/deepseek-v4-flash": { "cache_read_input_token_cost": 8e-08, + "deprecation_date": "2026-10-27", "input_cost_per_token": 1.5e-07, "litellm_provider": "aiand", "max_input_tokens": 1048576, @@ -72257,6 +72384,22 @@ "supports_vision": false }, "aiand/deepseek-ai/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2026-11-05", + "input_cost_per_token": 3e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/deepseek-ai/deepseek-v4.1-flash-1017": { "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 3e-07, "litellm_provider": "aiand", @@ -72286,8 +72429,24 @@ "supports_tool_choice": true, "supports_vision": true }, + "aiand/meta/muse-spark-open": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 943718, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "aiand/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2026-10-27", "input_cost_per_token": 7.5e-07, "litellm_provider": "aiand", "max_input_tokens": 262144, @@ -72378,6 +72537,7 @@ }, "aiand/zai-org/glm-5.2": { "cache_read_input_token_cost": 3e-07, + "deprecation_date": "2026-10-27", "input_cost_per_token": 1e-06, "litellm_provider": "aiand", "max_input_tokens": 1048576, @@ -75343,6 +75503,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "rpm": 1000.0, @@ -76443,6 +76604,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76488,6 +76650,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76540,6 +76703,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76598,6 +76762,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -81100,5 +81265,61 @@ "output_cost_per_token": 0.0, "source": "https://inference.baseten.co/v1/models", "supports_vision": true + }, + "gpt-6-rosalind": { + "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_12hr": 3.75e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 7.5e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.75e-05, + "cache_read_input_token_cost": 1.5e-06, + "cache_read_input_token_cost_above_272k_tokens": 3e-06, + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_above_272k_tokens": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6-rosalind": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } diff --git a/litellm/models/deployment_warning.py b/litellm/models/deployment_warning.py new file mode 100644 index 00000000000..1fcec9163bb --- /dev/null +++ b/litellm/models/deployment_warning.py @@ -0,0 +1,24 @@ +from collections.abc import Mapping +from typing import Literal, Protocol + +from pydantic import AwareDatetime, ConfigDict, Field + +from litellm.types.llms.base import LiteLLMBaseModel + + +class AdminIgnore(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + source: Literal["admin"] = "admin" + ignored_by: str + ignored_at: AwareDatetime + ignored_until: AwareDatetime | None = Field( + default=None, + description="When the snooze ends. Null means the warning is ignored until an admin restores it.", + ) + + +class DeploymentWarningIgnoreStore(Protocol): + async def get_ignores(self) -> Mapping[str, AdminIgnore]: ... + + async def save_ignores(self, ignores: Mapping[str, AdminIgnore]) -> None: ... diff --git a/litellm/models/gateway_activity.py b/litellm/models/gateway_activity.py new file mode 100644 index 00000000000..a37e4e66817 --- /dev/null +++ b/litellm/models/gateway_activity.py @@ -0,0 +1,31 @@ +from typing import Literal, Protocol, TypeAlias + +from pydantic import ConfigDict + +from litellm.types.llms.base import LiteLLMBaseModel + + +class GatewayRouteActivity(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["route"] + date: str + category: str + route: str + successful_requests: int + failed_requests: int + + +class GatewayStatusCodeActivity(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["status"] + status_code: int + failed_requests: int + + +GatewayActivity: TypeAlias = GatewayRouteActivity | GatewayStatusCodeActivity + + +class GatewayActivityStore(Protocol): + async def get_activity(self, start_date: str, end_date: str) -> tuple[GatewayActivity, ...]: ... diff --git a/litellm/models/user_banner.py b/litellm/models/user_banner.py new file mode 100644 index 00000000000..c6ca3906abb --- /dev/null +++ b/litellm/models/user_banner.py @@ -0,0 +1,48 @@ +from typing import Final, Literal, Protocol + +from pydantic import ConfigDict, Field, model_validator + +from litellm.types.llms.base import LiteLLMBaseModel + +USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000 +UserBannerSeverity = Literal["info", "warning", "error"] + + +class UserBannerUpdate(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + enabled: bool = Field( + default=False, + description="If true, the banner is shown to all authenticated dashboard users.", + ) + message: str = Field( + default="", + max_length=USER_BANNER_MAX_MESSAGE_LENGTH, + description="Banner text shown to dashboard users. Markdown is supported.", + ) + severity: UserBannerSeverity = Field( + default="info", + description="Visual style of the banner.", + ) + + @model_validator(mode="after") + def _require_message_when_enabled(self) -> "UserBannerUpdate": + if self.enabled and not self.message.strip(): + raise ValueError("message must be non-empty when the banner is enabled") + return self + + +class UserBanner(UserBannerUpdate): + revision: str = Field( + default="", + description=( + "Server-stamped opaque publish identity; a fresh value is generated on every " + "update so clients re-surface dismissed banners on republish." + ), + ) + + +class UserBannerStore(Protocol): + async def get_banner(self) -> UserBanner: ... + + async def save_banner(self, banner: UserBanner) -> None: ... diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index c76b9be7c33..41ab7fcaf88 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3491,6 +3491,7 @@ "additionalProperties": { "$ref": "#/components/schemas/GuardrailMetric" }, + "description": "Spend by guardrail ID, already included in key and model spend. Populated on user, team and organization aggregated activity only, without a model filter", "title": "Guardrails", "type": "object" }, @@ -3812,10 +3813,12 @@ "GuardrailKeyMetric": { "properties": { "metadata": { - "$ref": "#/components/schemas/KeyMetadata" + "$ref": "#/components/schemas/KeyMetadata", + "description": "API key alias and existence information" }, "metrics": { - "$ref": "#/components/schemas/GuardrailSpendMetrics" + "$ref": "#/components/schemas/GuardrailSpendMetrics", + "description": "Guardrail spend and unpriced units for this API key" } }, "required": [ @@ -3828,6 +3831,7 @@ "GuardrailMetadata": { "properties": { "guardrail_name": { + "description": "Guardrail display name, falling back to the guardrail ID", "title": "Guardrail Name", "type": "string" } @@ -3844,14 +3848,17 @@ "additionalProperties": { "$ref": "#/components/schemas/GuardrailKeyMetric" }, + "description": "Per-key guardrail spend and unpriced units, limited to keys in the day's API key breakdown", "title": "Api Key Breakdown", "type": "object" }, "metadata": { - "$ref": "#/components/schemas/GuardrailMetadata" + "$ref": "#/components/schemas/GuardrailMetadata", + "description": "Guardrail name information" }, "metrics": { - "$ref": "#/components/schemas/GuardrailSpendMetrics" + "$ref": "#/components/schemas/GuardrailSpendMetrics", + "description": "Guardrail spend and unpriced units across all matching API keys" } }, "required": [ @@ -3865,11 +3872,13 @@ "properties": { "spend": { "default": 0.0, + "description": "Guardrail spend in USD, already included in key and model spend", "title": "Spend", "type": "number" }, "untracked_units": { "default": 0, + "description": "Guardrail units with no known price, excluded from spend", "title": "Untracked Units", "type": "integer" } @@ -58006,7 +58015,7 @@ "paths": { "/toolset/{toolset_name}/mcp": { "delete": { - "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", + "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.", "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { @@ -58045,7 +58054,7 @@ ] }, "get": { - "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", + "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.", "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_get", "parameters": [ { @@ -58084,7 +58093,7 @@ ] }, "head": { - "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", + "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.", "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_head", "parameters": [ { @@ -58123,7 +58132,7 @@ ] }, "options": { - "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", + "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.", "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_options", "parameters": [ { @@ -58162,7 +58171,7 @@ ] }, "patch": { - "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", + "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.", "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_patch", "parameters": [ { @@ -58201,7 +58210,7 @@ ] }, "post": { - "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", + "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.", "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_post", "parameters": [ { @@ -58240,7 +58249,7 @@ ] }, "put": { - "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", + "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: API keys other than an admin's key with no\nteam must hold the toolset through one of their grant sources, or the\nrequest will be rejected with a 403.", "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", "parameters": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f165355da5a..07588053959 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -124,6 +124,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.repositories.database import Database from litellm.tracing import TraceReceiver Span = _Span | Any @@ -133,6 +134,7 @@ else: class ProxyLifespanState(TypedDict): tracing_receiver: ReadOnly["TraceReceiver | None"] + database: ReadOnly["Database | None"] class ReconcileOutcome(NamedTuple): @@ -3758,6 +3760,7 @@ class SpendLogsMetadata(TypedDict): kubernetes_pod_routing: ReadOnly[StandardLoggingKubernetesPodRouting | None] internal_call_origin: InternalCallOrigin | None billed_retry_attempt: ReadOnly[NotRequired[bool | None]] + upstream_provider: ReadOnly[NotRequired[str | None]] litellm_roi_estimator: ReadOnly[NotRequired[bool | None]] is_silent_experiment: ReadOnly[NotRequired[bool | None]] guardrail_information: list[StandardLoggingGuardrailInformation] | None diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 219ce906340..5e1d861e7e9 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -431,6 +431,7 @@ async def _handle_stream_message( proxy_server_request: dict[str, object] | None = None, *, agent_extra_headers: dict[str, str] | None = None, + agent_name: str | None = None, user_api_key_dict: UserAPIKeyAuth | None = None, request_data: dict[str, object] | None = None, proxy_logging_obj: ProxyLogging | None = None, @@ -504,6 +505,7 @@ async def _handle_stream_message( metadata=metadata, proxy_server_request=proxy_server_request, agent_extra_headers=agent_extra_headers, + agent_name=agent_name, ) if ( @@ -932,6 +934,7 @@ async def invoke_agent_a2a( metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), agent_extra_headers=agent_extra_headers, + agent_name=str(agent_name), user_api_key_dict=user_api_key_auth_for_call, request_data=data, proxy_logging_obj=proxy_logging_obj, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b8d9530d2eb..9c0bb93ea16 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -84,6 +84,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, CeilingResolver, resolve_agent_access_group_ceiling, ) @@ -166,6 +167,7 @@ from litellm.repositories.table_repositories import ( from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.router import Router +from litellm.router_api import RegistryReader from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.router import DeploymentTypedDict @@ -524,7 +526,7 @@ def _log_budget_lookup_failure(entity: str, error: Exception) -> None: ) -def _get_router_zero_cost_cache(llm_router: Router) -> dict[str, bool] | None: +def _get_router_zero_cost_cache(llm_router: RegistryReader) -> dict[str, bool] | None: """ Return the router's per-instance zero-cost cache, or ``None`` for objects that don't expose one (e.g. ``MagicMock`` stand-ins in unit tests). @@ -540,7 +542,7 @@ def _get_router_zero_cost_cache(llm_router: Router) -> dict[str, bool] | None: return cache if isinstance(cache, dict) else None -def is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None) -> bool: +def is_model_cost_zero(model: str | list[str] | None, llm_router: RegistryReader | None) -> bool: """ Check if a model has zero cost (no configured pricing). @@ -666,7 +668,7 @@ _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) _TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) -def _deployments_served_for(model: str, llm_router: "Router") -> Sequence[DeploymentTypedDict]: +def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]: """ The deployments a request to ``model`` routes to, for the gates that waive budget checks. @@ -936,53 +938,58 @@ _GUARDRAIL_MODIFICATION_KEYS: Final[tuple] = ( ) -def _guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: +def requests_guardrail_modification(request_body: Mapping[str, object]) -> bool: + """Whether a request body sets the plural ``guardrails`` list or a per-request toggle that influences whether + default-on guardrails run (``disable_global_guardrails``, ``disable_global_guardrail`` singular, and + ``opted_out_global_guardrails``), at the top level or in either metadata key""" + return ( + _requests_guardrail_modification_in(request_body.get("metadata")) + or _requests_guardrail_modification_in(request_body.get("litellm_metadata")) + or _requests_guardrail_modification_in(request_body) + ) + + +def _requests_guardrail_modification_in(container: object) -> bool: + coerced: Final = _guardrail_container_as_dict(container) + if coerced is None: + return False + return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) + + +def _guardrail_container_as_dict(container: object) -> Mapping[str, object] | None: + """Accept dict or JSON-string (from multipart/form-data or extra_body). + + Without this, an attacker can smuggle guardrail keys past the check by + sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — + ``isinstance(dict)`` on the string returns False, the check returns + no-modification, and ``add_litellm_data_to_request`` parses the string + to a dict downstream. + """ + if isinstance(container, Mapping): + return container + if isinstance(container, str): + parsed: Final = _safe_json_loads_obj(container) + return parsed if isinstance(parsed, dict) else None + return None + + +def guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: """ Reject user-supplied metadata flags that would modify guardrail behavior - unless the team has explicit permission. Checked keys include the plural - ``guardrails`` list plus the per-request toggles that influence whether - default-on guardrails run (``disable_global_guardrails``, - ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``). + unless the team has explicit permission. User-supplied values for the bypass toggles are also silently ignored by ``_get_admin_metadata`` at read time; this check adds defense in depth by failing loudly at the auth layer so operators see an explicit 403 instead of a confusing silent-ignore. """ + if requests_guardrail_modification(request_body): + require_guardrail_modification_permission(team_object) + + +def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails - def _coerce_to_dict(container: object) -> dict | None: - """Accept dict or JSON-string (from multipart/form-data or extra_body). - - Without this, an attacker can smuggle guardrail keys past the check by - sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — - ``isinstance(dict)`` on the string returns False, the check returns - no-modification, and ``add_litellm_data_to_request`` parses the string - to a dict downstream. - """ - if isinstance(container, dict): - return container - if isinstance(container, str): - parsed: Final = _safe_json_loads_obj(container) - return parsed if isinstance(parsed, dict) else None - return None - - def _user_requested_modification(container: object) -> bool: - coerced: Final = _coerce_to_dict(container) - if coerced is None: - return False - return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) - - # Check both metadata keys — callers can populate either depending on the - # endpoint. Cover the top-level too so root-level injection is rejected. - modifies: Final = ( - _user_requested_modification(request_body.get("metadata")) - or _user_requested_modification(request_body.get("litellm_metadata")) - or _user_requested_modification(request_body) - ) - if not modifies: - return - if not can_modify_guardrails(team_object): raise HTTPException( status_code=403, @@ -1484,7 +1491,7 @@ async def common_checks( _enforce_user_param_check(general_settings, request, request_body, route) _global_proxy_budget_check(global_proxy_spend, skip_all_budget_checks, route) - _guardrail_modification_check(_typed_request_body(request_body), team_object) + guardrail_modification_check(_typed_request_body(request_body), team_object) # 10 [OPTIONAL] Organization RBAC checks organization_role_based_access_check(user_object=user_object, route=route, request_body=request_body) @@ -1684,6 +1691,7 @@ def get_actual_routes(allowed_routes: list) -> list: KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id" _KEY_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) +_OBJECTS_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: @@ -3455,28 +3463,35 @@ async def delete_cache_key_objects( proxy_logging_obj: ProxyLogging | None, ) -> None: """ - Evict a batch of key objects, for callers that delete keys in bulk rather than through - `/key/delete`. Auth resolves a cached key object without re-reading its team, so a key left - cached after its row is gone keeps buying access until its TTL expires. - - Evicting locally only reaches this worker; `_delete_cache_key_object` itself broadcasts each - token, so a deleted key left in a peer worker's in-memory cache still authenticates there until - its TTL expires. + Evict a batch of deleted key objects here and on every other worker. Auth resolves a cached key + object without re-reading its team, so a key left cached after its row is gone keeps buying + access until its TTL expires. Best-effort per key: the rows are already deleted by the time this runs, so an unreachable cache backend must not abort the caller partway through its own cascade. - Every token is dropped from this worker's memory first, then from the auth cache and the - internal usage cache in one batch each, so a deleted key stops authenticating here the instant - its rows are gone instead of refilling from Redis on the next miss. Every other worker then - gets the eviction in one pipelined broadcast: a publish per key would pass the publisher's - pending cap on a bulk delete, and the dropped keys would keep authenticating there until their - TTL. + Every token is dropped from this worker's memory first, then from the auth cache's Redis layer, + so a deleted key stops authenticating here the instant its rows are gone instead of refilling + from Redis on the next miss. Every other worker then gets the eviction in one pipelined + broadcast: a publish per key would pass the publisher's pending cap on a bulk delete, and the + dropped keys would keep authenticating there until their TTL. The internal usage cache is not + an auth refill source, so its batch delete runs beside that chain and a hung Redis costs one + bounded wait there instead of another in line. """ evict_local(cache_keys=hashed_tokens, user_api_key_cache=user_api_key_cache) + usage_eviction: Final = ( + () + if proxy_logging_obj is None + else (evict_shared(cache_keys=hashed_tokens, cache=proxy_logging_obj.internal_usage_cache.dual_cache),) + ) + _ = await asyncio.gather( + _evict_auth_cache_then_broadcast(hashed_tokens=hashed_tokens, user_api_key_cache=user_api_key_cache), + *usage_eviction, + ) + + +async def _evict_auth_cache_then_broadcast(hashed_tokens: Sequence[str], user_api_key_cache: UserApiKeyCache) -> None: await evict_shared(cache_keys=hashed_tokens, cache=user_api_key_cache) - if proxy_logging_obj is not None: - await evict_shared(cache_keys=hashed_tokens, cache=proxy_logging_obj.internal_usage_cache.dual_cache) await broadcast_evictions(cache_keys=hashed_tokens) @@ -4828,7 +4843,7 @@ async def get_skill_names_from_access_groups( def _resolve_all_team_model_sentinel_for_auth_check( models: list[str], - llm_router: Router | None, + llm_router: RegistryReader | None, team_id: str | None, ) -> list[str]: if SpecialModelNames.all_team_models.value not in models or team_id is None or llm_router is None: @@ -4842,7 +4857,7 @@ def _resolve_all_team_model_sentinel_for_auth_check( def check_model_access_helper( model: str, - llm_router: Router | None, + llm_router: RegistryReader | None, models: list[str], team_model_aliases: dict[str, str] | None = None, team_id: str | None = None, @@ -4893,7 +4908,7 @@ _check_model_access_helper: Final = check_model_access_helper def _owning_team_names_for_team_deployment( model: str, - llm_router: Router | None, + llm_router: RegistryReader | None, team_id: str | None, team_model_aliases: Mapping[str, str] | None, ) -> tuple[str, ...]: @@ -4911,7 +4926,7 @@ def _owning_team_names_for_team_deployment( def _can_object_call_model( model: str | list[str], - llm_router: Router | None, + llm_router: RegistryReader | None, models: list[str], team_model_aliases: dict[str, str] | None = None, team_id: str | None = None, @@ -5038,7 +5053,7 @@ def _resolve_team_alias( def _live_team_alias_target( - model: str, team_model_aliases: Mapping[str, str], team_id: str | None, llm_router: Router | None + model: str, team_model_aliases: Mapping[str, str], team_id: str | None, llm_router: RegistryReader | None ) -> str: target: Final = team_model_aliases.get(model) if target is None: @@ -5085,6 +5100,36 @@ async def can_agent_call_model( ) +async def _agent_access_group_ceilings( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver +) -> tuple[AgentAccessGroupCeiling, ...]: + if not valid_token.agent_id: + return () + + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings + + managed: Final = managed_agent_policy(valid_token) + if managed is not None: + return await resolve_managed_agent_ceilings(managed) + unmanaged: Final = await resolve_ceiling(valid_token.agent_id) + return (unmanaged,) if unmanaged is not None else () + + +def _managed_agent_model_grants(valid_token: UserAPIKeyAuth) -> frozenset[str]: + managed: Final = managed_agent_policy(valid_token) + granted: Final = (managed.object_permission or {}).get("models") if managed is not None else None + if not isinstance(granted, (list, tuple)): + return frozenset() + return frozenset(name for name in _OBJECTS_ADAPTER.validate_python(granted) if isinstance(name, str)) + + +async def agent_model_grants( + valid_token: UserAPIKeyAuth, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling +) -> frozenset[str]: + ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) + return _managed_agent_model_grants(valid_token).union(*(ceiling.models for ceiling in ceilings)) + + async def _check_agent_access_group_model_access( model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] valid_token: UserAPIKeyAuth | None, @@ -5092,20 +5137,9 @@ async def _check_agent_access_group_model_access( resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> Literal[True]: """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows.""" - if not model or valid_token is None or not valid_token.agent_id: + if not model or valid_token is None: return True - - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings - - managed: Final = managed_agent_policy(valid_token) - unmanaged: Final = await resolve_ceiling(valid_token.agent_id) if managed is None else None - ceilings: Final = ( - await resolve_managed_agent_ceilings(managed) - if managed is not None - else (unmanaged,) - if unmanaged is not None - else () - ) + ceilings: Final = await _agent_access_group_ceilings(valid_token, resolve_ceiling) dispatched: Final = _resolve_team_alias( model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router ) @@ -5506,7 +5540,7 @@ async def can_key_call_model( model: str | list[str], llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, - llm_router: litellm.Router | None, + llm_router: RegistryReader | None, prisma_client: DatabaseClient | None = None, ) -> Literal[True]: """ @@ -5585,7 +5619,7 @@ async def can_key_call_resolved_model( model: str, llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, - llm_router: litellm.Router | None, + llm_router: RegistryReader | None, ) -> None: from litellm.proxy.proxy_server import ( prisma_client, @@ -5733,7 +5767,7 @@ def can_org_access_model( async def can_team_access_model( model: str | list[str], team_object: LiteLLM_TeamTable | None, - llm_router: Router | None, + llm_router: RegistryReader | None, team_model_aliases: dict[str, str] | None = None, key_model_aliases: Mapping[str, str] | None = None, prisma_client: DatabaseClient | None = None, @@ -5836,7 +5870,7 @@ async def key_access_group_grants_model( model: str | list[str], valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, - llm_router: Router | None, + llm_router: RegistryReader | None, key_model_aliases: Mapping[str, str] | None = None, ) -> bool: """ @@ -5872,7 +5906,7 @@ _key_access_group_grants_model: Final = key_access_group_grants_model def can_project_access_model( model: str | list[str], project_object: LiteLLM_ProjectTable, - llm_router: Router | None, + llm_router: RegistryReader | None, key_model_aliases: Mapping[str, str] | None = None, ) -> Literal[True]: """ @@ -5893,7 +5927,7 @@ def can_project_access_model( def can_customer_access_model( model: str | list[str], end_user_object: LiteLLM_EndUserTable, - llm_router: Router | None, + llm_router: RegistryReader | None, valid_token: UserAPIKeyAuth | None, ) -> Literal[True]: team_model_aliases: Final = team_model_aliases_for_auth_check(valid_token) if valid_token is not None else None @@ -6056,7 +6090,7 @@ async def _team_grants_personal_key_model( model: str, team_object: LiteLLM_TeamTableCachedObj, valid_token: UserAPIKeyAuth, - llm_router: Router | None, + llm_router: RegistryReader | None, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, @@ -6094,7 +6128,7 @@ async def _loaded_teams_allow_model( model: str, teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], valid_token: UserAPIKeyAuth, - llm_router: Router | None, + llm_router: RegistryReader | None, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, @@ -6360,7 +6394,7 @@ async def can_personal_key_call_model_via_teams( model: str | Sequence[str], user_object: LiteLLM_UserTable | None, valid_token: UserAPIKeyAuth, - llm_router: Router | None, + llm_router: RegistryReader | None, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, @@ -7152,7 +7186,7 @@ async def check_team_member_model_access( model: str | list[str], team_object: LiteLLM_TeamTable, valid_token: UserAPIKeyAuth, - llm_router: Router | None, + llm_router: RegistryReader | None, prisma_client: Optional["PrismaClient"], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py index 95ff11261d9..5065c69e512 100644 --- a/litellm/proxy/auth/fallback_budget.py +++ b/litellm/proxy/auth/fallback_budget.py @@ -44,7 +44,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo _is_model_cost_zero, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export is_model_cost_zero, # pyright: ignore[reportPrivateUsage] # the zero-cost predicate the auth-time budget checks use; no public equivalent ) -from litellm.router import Router +from litellm.router_api import RegistryReader, RoutingContext from litellm.types.llms.base import LiteLLMBaseModel @@ -118,7 +118,7 @@ async def is_token_within_budget_for_model( *, model: str, valid_token: UserAPIKeyAuth, - llm_router: Router, + llm_router: RegistryReader, team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None, ) -> bool: """ @@ -187,19 +187,19 @@ class RouterFallbackBudgetCheck: is_enforced: Callable[[], bool] team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None - async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool: + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: RoutingContext) -> bool: if not self.is_enforced(): return True valid_token: Final = _user_api_key_auth_from_request(request_kwargs) if valid_token is None: return True - budget_models: Final = tuple(dict.fromkeys((model, llm_router.registry().resolve(model)))) + budget_models: Final = tuple(dict.fromkeys((model, ctx.registry().resolve(model)))) for budget_model in budget_models: try: if not await is_token_within_budget_for_model( model=budget_model, valid_token=valid_token, - llm_router=llm_router, + llm_router=ctx, team_member_model_budget_limiter=self.team_member_model_budget_limiter, ): return False diff --git a/litellm/proxy/auth/fallback_model_access.py b/litellm/proxy/auth/fallback_model_access.py index fe80dc7a90b..b550d4d6866 100644 --- a/litellm/proxy/auth/fallback_model_access.py +++ b/litellm/proxy/auth/fallback_model_access.py @@ -18,7 +18,7 @@ from pydantic import ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, can_token_call_search_tool -from litellm.router import Router +from litellm.router_api import RegistryReader, RoutingContext from litellm.search import asearch from litellm.types.llms.base import LiteLLMBaseModel @@ -31,7 +31,7 @@ class _FallbackAccessSettings(LiteLLMBaseModel): enforce_fallback_model_access: bool = True -async def is_model_authorized_for_token(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool: +async def is_model_authorized_for_token(*, model: str, valid_token: UserAPIKeyAuth, llm_router: RegistryReader) -> bool: try: await can_key_call_resolved_model( model=model, @@ -98,7 +98,7 @@ class RouterFallbackAccessCheck: is_enforced: Callable[[], bool] - async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool: + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: RoutingContext) -> bool: valid_token: Final = _user_api_key_auth_from_request(request_kwargs) if valid_token is None: return True @@ -106,7 +106,7 @@ class RouterFallbackAccessCheck: return await is_search_tool_authorized_for_token(search_tool_name=model, valid_token=valid_token) if not self.is_enforced(): return True - return await is_model_authorized_for_token(model=model, valid_token=valid_token, llm_router=llm_router) + return await is_model_authorized_for_token(model=model, valid_token=valid_token, llm_router=ctx) router_fallback_access_check: Final = RouterFallbackAccessCheck(is_enforced=_enforced_by_general_settings) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index b21315529eb..4bcf020a9c4 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -1,23 +1,26 @@ # What is this? ## Common checks for /v1/models and `/model/info` +import asyncio +import contextvars import copy import functools import itertools import re from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.responses.main import strip_responses_routing_prefix from litellm.router import Router -from litellm.router_utils.fallback_event_handlers import get_fallback_model_group -from litellm.router_utils.pattern_match_deployments import PatternMatchRouter +from litellm.router_api import fallback_chain from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager, get_valid_models, potential_model_names_in_lookup_order @@ -30,6 +33,17 @@ _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) +_MODEL_LISTING_POOL: Final = ThreadPoolExecutor(thread_name_prefix="litellm-model-listing") + +_PROVIDERS_LISTED_ONLY_THROUGH_A_DEPLOYMENT: Final = frozenset({LlmProviders.OPENROUTER.value}) + +_Listed: Final = TypeVar("_Listed") + + +async def run_model_listing(listing: Callable[[], _Listed]) -> _Listed: + context: Final = contextvars.copy_context() + return await asyncio.get_running_loop().run_in_executor(_MODEL_LISTING_POOL, context.run, listing) + def _check_wildcard_routing(model: str) -> bool: """ @@ -209,6 +223,7 @@ def get_complete_model_list( include_model_access_groups: bool | None = False, only_model_access_groups: bool | None = False, team_id: str | None = None, + expand_wildcard_routes: bool = True, ) -> list[str]: """Logic for returning complete model list for a given key + team pair""" @@ -254,11 +269,10 @@ def get_complete_model_list( return_wildcard_routes=return_wildcard_routes, llm_router=llm_router, team_id=team_id, + expand_wildcard_routes=expand_wildcard_routes, ) - complete_model_list: Final = unique_models + all_wildcard_models - - return complete_model_list + return list(dict.fromkeys(unique_models + all_wildcard_models)) def _hydrate_litellm_credential_name( @@ -425,6 +439,15 @@ def _cached_wildcard_routing(llm_router: Router | None) -> Callable[[str, str | return functools.cache(wildcard_routed_models) +def _deployed_by_name(llm_router: Router | None) -> Callable[[str, str | None], bool]: + if llm_router is None: + return lambda _model, _team_id: False + registry: Final = llm_router.registry() + deployed_names: Final = frozenset(registry.names(kind="group_set")) + team_deployed_names: Final = frozenset(registry.index("team_model")) + return lambda model, team_id: model in deployed_names or (team_id, model) in team_deployed_names + + def _routes_to_model( routed_models: Callable[[str, str | None], frozenset[str]] | None, team_id: str | None ) -> Callable[[str, str], bool]: @@ -455,7 +478,11 @@ def get_known_models_from_wildcard( wildcard_provider_prefix, wildcard_suffix = wildcard_model_to_expand.split("/", 1) except ValueError: # safely fail return [] - deployment_pattern_naming_listed_ids: Final = _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params) + deployment_pattern_naming_listed_ids: Final = ( + wildcard_model_to_expand + if wildcard_model_to_expand != wildcard_model + else _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params) + ) # Use provider from litellm_params when available, otherwise from wildcard prefix # (e.g., "openai" from "openai/*" - needed for BYOK where wildcard isn't in router) @@ -514,6 +541,27 @@ def _under_wildcard_prefix(model: str, wildcard_provider_prefix: str, listed_pro return f"{wildcard_provider_prefix}/{model_suffix if leading_is_listed_providers_own_prefix else model}" +@functools.lru_cache(maxsize=256) +def _wildcard_name_regex(name_pattern: str) -> re.Pattern[str]: + return re.compile(PatternMatchRouter().pattern_to_regex(name_pattern)) + + +def _routed_deployment_model(name_pattern: str, listed_name: str, deployment_model: str) -> str: + name_match: Final = _wildcard_name_regex(name_pattern).match(listed_name) + if name_match is None: + return listed_name + return PatternMatchRouter.set_deployment_model_name( + matched_pattern=name_match, litellm_deployment_litellm_model=deployment_model + ) + + +def _names_no_deployment_serves_by_name( + deployment: Mapping[str, object], names: Sequence[str], deployed_by_name: Callable[[str, str | None], bool] +) -> tuple[str, ...]: + team_id: Final = _deployment_team_id(deployment) + return tuple(name for name in names if not deployed_by_name(name, team_id)) + + def expand_wildcard_deployments_for_model_info( deployments: list[dict[str, object]], llm_router: Router | None = None, @@ -527,6 +575,7 @@ def expand_wildcard_deployments_for_model_info( """ expanded: Final[list[dict[str, object]]] = [] routed_models: Final = _cached_wildcard_routing(llm_router) + deployed_by_name: Final = _deployed_by_name(llm_router) for deployment in deployments: model_name = str(deployment.get("model_name") or "") raw_params = deployment.get("litellm_params") @@ -560,22 +609,29 @@ def expand_wildcard_deployments_for_model_info( expanded.append(deployment) continue - for name in expanded_names: + for name in _names_no_deployment_serves_by_name(deployment, expanded_names, deployed_by_name): row = copy.deepcopy(deployment) row["model_name"] = name params = row.get("litellm_params") if isinstance(params, dict): - params["model"] = name + params["model"] = _routed_deployment_model(model_name, name, litellm_model) expanded.append(row) return expanded +def _known_models_without_a_deployment(wildcard_model: str) -> Sequence[str]: + if wildcard_model.partition("/")[0] in _PROVIDERS_LISTED_ONLY_THROUGH_A_DEPLOYMENT: + return () + return get_known_models_from_wildcard(wildcard_model=wildcard_model, litellm_params=None) + + def _get_wildcard_models( unique_models: list[str], return_wildcard_routes: bool | None = False, llm_router: Router | None = None, team_id: str | None = None, + expand_wildcard_routes: bool = True, ) -> list[str]: models_to_remove: Final = set() routes_to_model: Final = _routes_to_model(_cached_wildcard_routing(llm_router), team_id) @@ -586,6 +642,8 @@ def _get_wildcard_models( all_wildcard_models.append(model) models_to_remove.add(model) + if not expand_wildcard_routes: + continue model_list = llm_router.registry().deployments(model, team_id=team_id) if llm_router else None if model_list: @@ -598,7 +656,7 @@ def _get_wildcard_models( ) ) else: - all_wildcard_models.extend(get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)) + all_wildcard_models.extend(_known_models_without_a_deployment(model)) for model in models_to_remove: unique_models.remove(model) @@ -650,7 +708,7 @@ def get_all_fallbacks( try: # Use existing function to get fallback model group - fallback_model_group, _ = get_fallback_model_group(fallbacks=fallbacks_config, model_group=model) + fallback_model_group, _ = fallback_chain(fallbacks=fallbacks_config, model_group=model) if fallback_model_group is None: return [] diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 67410e867d6..2b0522c09b5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -40,6 +40,7 @@ from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_event, phase_span, seed_request_identity from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias, resolve_served_model from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers @@ -182,13 +183,7 @@ from litellm.proxy.utils import ( from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository -from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model -from litellm.router_utils.fallback_event_handlers import ( - PROVIDER_SCOPED_RESOURCE_KEYS, - fallback_lookup_groups, - fallbacks_disabled_for_request, - get_fallback_model_group_for_lookup_groups, -) +from litellm.router_api import fallbacks_disabled, references_provider_scoped_resource, request_fallback_chain from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes @@ -811,11 +806,7 @@ async def _check_team_member_model_budget_for_model( model=model_name, ) except litellm.BudgetExceededError: - if ( - request_data.get("model") != model_name - or llm_router is None - or fallbacks_disabled_for_request(kwargs=request_data) - ): + if request_data.get("model") != model_name or llm_router is None or fallbacks_disabled(request_data): raise fallback_succeeded: Final = await _try_team_member_model_budget_fallback( user_id=user_id, @@ -857,9 +848,9 @@ async def _try_team_member_model_budget_fallback( user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> bool: - if request.headers.get("content-type", "").lower().startswith("multipart/form-data") or any( - request_data.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS - ): + if request.headers.get("content-type", "").lower().startswith( + "multipart/form-data" + ) or references_provider_scoped_resource(request_data): return False from litellm.proxy import proxy_server @@ -896,10 +887,7 @@ async def _try_team_member_model_budget_fallback( if router_fallbacks is None: return False - fallback_models, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=list(router_fallbacks), - lookup_groups=fallback_lookup_groups(kwargs=request_data, model_group=model_name), - ) + fallback_models, _ = request_fallback_chain(list(router_fallbacks), request_data, model_name) fallback_targets: Final[Sequence[object]] = fallback_models or () fallback_candidates: Final[tuple[str, ...]] = tuple( model for target in fallback_targets if (model := _budget_fallback_target_model(target)) is not None diff --git a/litellm/proxy/bug_report_config.py b/litellm/proxy/bug_report_config.py index d7b920de4d5..6c0209644b4 100644 --- a/litellm/proxy/bug_report_config.py +++ b/litellm/proxy/bug_report_config.py @@ -19,8 +19,8 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, build_environment_report, ) +from litellm.litellm_core_utils.routing_groups import VALID_ROUTING_STRATEGIES from litellm.proxy._types import ConfigGeneralSettings -from litellm.router_utils.routing_groups import VALID_ROUTING_STRATEGIES from litellm.types.caching import LiteLLMCacheType from litellm.types.guardrails import GuardrailEventHooks, LitellmParams, SupportedGuardrailIntegrations from litellm.types.secret_managers.main import KeyManagementSystem diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a79d8fe29db..5599086d5c1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -53,6 +53,7 @@ from litellm.constants import ( ) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.otel.runtime import phase_event +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.litellm_core_utils.bug_report import ( allowlisted, bug_report_notice, @@ -60,6 +61,7 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.core_helpers import ( + finite_cost, get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, @@ -80,6 +82,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias, resolve_served_model from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.served_output_texts import ( record_served_output_texts, @@ -140,16 +143,15 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) +from litellm.responses.mcp.parked_stream_dispatch import in_flight_round_biller, parked_success from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router from litellm.router_api import RoutingContext from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX -from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict -from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model from litellm.types.extraction import ExtractionResponse from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError -from litellm.types.router_weights import validate_router_weights +from litellm.types.router_weights import RouterWeights, validate_router_weights _LateResponseT = TypeVar("_LateResponseT", bound=Response) _LlmCallT = TypeVar("_LlmCallT") @@ -535,6 +537,16 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return await _dispatch_disconnect_billing_event(logging_obj, partial_response) +async def _bill_in_flight_round_on_disconnect(request_data: Mapping[str, object]) -> None: + bill_in_flight_round: Final = in_flight_round_biller(request_data.get("litellm_logging_obj")) + if bill_in_flight_round is None: + return + try: + await bill_in_flight_round() + except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown + verbose_proxy_logger.debug("Failed to bill the in-flight round after client disconnect: %s", e) + + async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool: try: await logging_obj.dispatch_success_handlers( @@ -882,6 +894,61 @@ async def _resolve_per_request_model_group_alias( return target +async def hierarchical_router_settings_updates( + requested_model: object, + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router, + proxy_config: ProxyConfig, + proxy_logging_obj: ProxyLogging, + apply_model_group_alias: bool, +) -> Mapping[str, object]: + # Apply hierarchical router_settings (Key > Team) + # Global router_settings are already on the Router object itself. + # This sits with the other alias rewrites, and ahead of the guardrail + # merge and the pre-call hooks, so everything that keys off the model + # group -- model-level guardrails, per-model budgets and rate limits, + # the logging object -- sees the group that will actually serve. + from litellm.proxy.proxy_server import prisma_client + + router_settings: Final[Mapping[str, object] | None] = _ROUTER_SETTINGS_ADAPTER.validate_python( + await proxy_config._get_hierarchical_router_settings( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ) + if router_settings is None: + return MappingProxyType({}) + + # If router_settings found (from key or team), apply them + # Pass settings as per-request overrides instead of creating a new Router + # This avoids expensive Router instantiation on each request + def validated_router_weights() -> RouterWeights | None: + try: + return validate_router_weights(router_settings.get("weights")) + except ValidationError: + verbose_proxy_logger.warning("Ignoring invalid saved router weights; update team/key router_settings") + return None + + router_weights: Final = validated_router_weights() + alias_target: Final = ( + await _resolve_per_request_model_group_alias( + requested_model=requested_model, + router_settings=router_settings, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) + if apply_model_group_alias + else None + ) + settings_updates: Final[Mapping[str, object]] = { + "router_settings_override": router_settings, + "_router_weights": router_weights, + **({"model": alias_target} if alias_target is not None else {}), + } + return MappingProxyType(settings_updates) + + _REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None) _DECODED_JSON: Final = TypeAdapter(object) _ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -1801,14 +1868,16 @@ def _get_cost_breakdown_from_logging_obj( def _header_value(value: float | None) -> str | None: - return str(value) if value is not None else None + return str(value) if finite_cost(value) is not None else None def _response_cost_header_value(hidden_params: Mapping[str, object]) -> float | str: if hidden_params.get("cache_hit") is True: return 0.0 response_cost: Final = hidden_params.get("response_cost") - return response_cost if isinstance(response_cost, (int, float, str)) and response_cost else "" + if isinstance(response_cost, str): + return response_cost + return response_cost if isinstance(response_cost, int | float) and finite_cost(response_cost) else "" def cost_headers( @@ -2083,6 +2152,13 @@ class _ClientStreamFailureReport: ) +def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]: + logged_params: Final = object_mapping(logged_on.model_call_details.get("litellm_params")) or MappingProxyType({}) + return MappingProxyType( + {key: logged_params[key] for key in ("metadata", "litellm_metadata") if key in logged_params} + ) + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict, served_model: str | None = None): self.data = data @@ -2474,46 +2550,22 @@ class ProxyBaseLLMRequestProcessing: ): self.data["model"] = user_api_key_dict.aliases[self.data["model"]] - # Apply hierarchical router_settings (Key > Team) - # Global router_settings are already on the Router object itself. - # This sits with the other alias rewrites, and ahead of the guardrail - # merge and the pre-call hooks, so everything that keys off the model - # group -- model-level guardrails, per-model budgets and rate limits, - # the logging object -- sees the group that will actually serve. if llm_router is not None and proxy_config is not None: - from litellm.proxy.proxy_server import prisma_client - - router_settings: Final = await proxy_config._get_hierarchical_router_settings( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - ) - - # If router_settings found (from key or team), apply them - # Pass settings as per-request overrides instead of creating a new Router - # This avoids expensive Router instantiation on each request - if router_settings is not None: - self.data["router_settings_override"] = router_settings - try: - self.data["_router_weights"] = validate_router_weights(router_settings.get("weights")) - except ValidationError: - self.data["_router_weights"] = None - verbose_proxy_logger.warning( - "Ignoring invalid saved router weights; update team/key router_settings" - ) - alias_target: Final = await _resolve_per_request_model_group_alias( + self.data.update( + await hierarchical_router_settings_updates( requested_model=self.data.get("model"), - router_settings=router_settings, user_api_key_dict=user_api_key_dict, llm_router=llm_router, + proxy_config=proxy_config, + proxy_logging_obj=proxy_logging_obj, + apply_model_group_alias=True, ) - if alias_target is not None: - self.data["model"] = alias_target + ) aliased_model: Final = self.data.get("model") if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str): self.data.update( - llm_router.priority_header_update( + llm_router.health().priority_header_fields( model=aliased_model, kwargs=self.data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER) ) ) @@ -2779,12 +2831,9 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _resolve_fallback_models(model: str, fallbacks: list) -> list[str] | None: - from litellm.router_utils.fallback_event_handlers import get_fallback_model_group + from litellm.router_api import fallback_chain - fallback_model_group, generic_fallback_idx = get_fallback_model_group( - fallbacks=fallbacks, - model_group=model, - ) + fallback_model_group, generic_fallback_idx = fallback_chain(fallbacks, model) if fallback_model_group is not None: return fallback_model_group return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None @@ -4138,8 +4187,8 @@ class ProxyBaseLLMRequestProcessing: HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the unwrapped inner iterator. """ + from litellm.litellm_core_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.router_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper unwrapped: Final = response.inner if isinstance(response, HiddenParamsAsyncIteratorWrapper) else response @@ -4151,11 +4200,21 @@ class ProxyBaseLLMRequestProcessing: _captured_user_api_key_dict: Final = user_api_key_dict _captured_logging_obj: Final = logging_obj - async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: + async def _on_deferred_stream_complete( + assembled_response: object, cache_hit: object, logged_on: LiteLLMLoggingObj | None = None + ) -> None: await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( - captured_data=_captured_data, + captured_data=_captured_data + if logged_on is None + else { + **_captured_data, + **{ + key: value for key, value in _metadata_logged_on(logged_on).items() if key in _captured_data + }, + "litellm_logging_obj": logged_on, + }, captured_user_api_key_dict=_captured_user_api_key_dict, - captured_logging_obj=_captured_logging_obj, + captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on, assembled_response=assembled_response, cache_hit=cache_hit, ) @@ -4180,14 +4239,8 @@ class ProxyBaseLLMRequestProcessing: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) - case (assembled_response, cache_hit): - await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) + case _ if (success := parked_success(args)) is not None: + await success.dispatch(_captured_native_logging_obj) case _: verbose_proxy_logger.error( "Deferred stream logging dropped: unexpected stored args shape %s", @@ -4487,7 +4540,7 @@ class ProxyBaseLLMRequestProcessing: """ if isinstance(chunk, dict): # Use safe_dumps for proper JSON serialization with circular reference detection - chunk_str: Final = safe_dumps(chunk) + chunk_str: Final = safe_dumps(chunk, non_finite_as_null=True) return f"{STREAM_SSE_DATA_PREFIX}{chunk_str}\n\n" else: return chunk @@ -4524,6 +4577,7 @@ class ProxyBaseLLMRequestProcessing: client_disconnected, ) if recorded_client_disconnect: + await _bill_in_flight_round_on_disconnect(_REQUEST_FIELDS.validate_python(request_data)) deferred_stream_logging_armed: Final = _deferred_stream_logging_is_armed(request_data) ProxyLogging.fire_deferred_stream_logging(request_data) # A disconnect-time success event (the deferred-guardrail flush @@ -4862,7 +4916,11 @@ class ProxyBaseLLMRequestProcessing: obj, model_name, litellm_logging_obj ) if maybe_modified is not None: - lines[idx] = "data: " + safe_dumps(maybe_modified) + ("\r" if ln.endswith("\r") else "") + lines[idx] = ( + "data: " + + safe_dumps(maybe_modified, non_finite_as_null=True) + + ("\r" if ln.endswith("\r") else "") + ) return "\n".join(lines) return None except Exception: @@ -4973,8 +5031,9 @@ class ProxyBaseLLMRequestProcessing: if stream_usage is None: return None service_tier: Final = obj.get("service_tier") + served_model: Final = litellm_logging_obj.model if litellm_logging_obj is not None else None cost_val: Final = ProxyBaseLLMRequestProcessing._streamed_usage_cost( - ModelResponse(usage=stream_usage), + ModelResponse(model=served_model, usage=stream_usage), model_name, service_tier if isinstance(service_tier, str) else None, litellm_logging_obj, diff --git a/litellm/proxy/common_utils/cache_aware_routing.py b/litellm/proxy/common_utils/cache_aware_routing.py index 08728b97b54..dfab1847ab8 100644 --- a/litellm/proxy/common_utils/cache_aware_routing.py +++ b/litellm/proxy/common_utils/cache_aware_routing.py @@ -16,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_f from litellm.llms.anthropic.cache_aware_routing import AnthropicCacheRouting, TokenCounter from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model +from litellm.router_deprecated import renamed_router_param from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import Deployment, PreRoutingHookResponse @@ -23,6 +24,7 @@ from litellm.types.utils import StandardLoggingRoutingDecision if TYPE_CHECKING: from litellm.router import Router + from litellm.router_api import RoutingContext _MESSAGES: Final = TypeAdapter(list[Mapping[str, object]]) _MAPPING: Final = TypeAdapter(Mapping[str, object]) @@ -231,7 +233,7 @@ async def select_cached_model( async def _choose_cached_model( *, - router: Router, + ctx: RoutingContext, config: ComplexityRouterConfig, params_for_model: Callable[[str, str], Mapping[str, object]], response: PreRoutingHookResponse | None, @@ -249,8 +251,10 @@ async def _choose_cached_model( ) from litellm.router_strategy.complexity_router.context_compaction import compaction_pending + router: Final = proxy_server.llm_router if ( - proxy_server.llm_router is not router + router is None + or not ctx.belongs_to(router) or router.registry().plugins("routing") or has_request_transforms() or compaction_pending(request_kwargs) @@ -318,9 +322,10 @@ async def _choose_cached_model( ) +@renamed_router_param("router") async def choose_cached_model( *, - router: Router, + ctx: RoutingContext, config: ComplexityRouterConfig, params_for_model: Callable[[str, str], Mapping[str, object]], response: PreRoutingHookResponse | None, @@ -332,7 +337,7 @@ async def choose_cached_model( try: return await asyncio.wait_for( _choose_cached_model( - router=router, + ctx=ctx, config=config, params_for_model=params_for_model, response=response, diff --git a/litellm/proxy/common_utils/credential_hydration.py b/litellm/proxy/common_utils/credential_hydration.py index f6a6c682cf0..664290fb3fc 100644 --- a/litellm/proxy/common_utils/credential_hydration.py +++ b/litellm/proxy/common_utils/credential_hydration.py @@ -12,10 +12,10 @@ from types import MappingProxyType from typing import Final, cast # noqa: TID251 # narrows untyped credential values at the decrypt boundary import litellm +from litellm.litellm_core_utils.clientside_credential_handler import clientside_credential_keys from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.utils import PrismaClient from litellm.repositories.credentials_repository import CredentialsRepository -from litellm.router_utils.clientside_credential_handler import clientside_credential_keys from litellm.types.router import ( GenericLiteLLMParams, server_owned_wif_fields_present, diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index ad4c3bd8770..9db67aa3808 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -21,7 +21,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.constants import CLAUDE_CODE_CLIENT, CLAUDE_CODE_PICKER_PATTERN, GATEWAY_CLIENT_HEADER -from litellm.router_utils.common_utils import resolve_model_group_alias +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias if TYPE_CHECKING: from litellm.router import Router @@ -219,7 +219,7 @@ def caller_alias_maps( def alias_map(aliases: object) -> Mapping[str, str]: try: - entries: Final = _ALIAS_ENTRIES.validate_python(aliases, strict=True) + entries: Final = _ALIAS_ENTRIES.validate_python(aliases) except ValidationError: return _NO_ALIASES return MappingProxyType( @@ -287,6 +287,15 @@ def _dispatch(model_id: str, layers: Sequence[_AliasLayer]) -> tuple[str, bool]: return reduce(_through_layer, layers, (model_id, False)) +def dispatched_model(model_id: str, aliases: CallerAliases) -> str: + target, _ = _dispatch(model_id, _alias_layers(aliases)) + return target + + +def requestable_alias_names(aliases: CallerAliases) -> frozenset[str]: + return frozenset(chain.from_iterable(layer.aliases for layer in _alias_layers(aliases))) + + def _lookup_id(name: str, lookup_by_response: Mapping[str, str], lookup_ids: Container[str]) -> str | None: if name in lookup_by_response: return lookup_by_response[name] diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 53b2a05d40f..a948c390375 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -28,6 +28,8 @@ _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( _AUTH_FAILURE_STATUS_CODES: Final = frozenset({status.HTTP_401_UNAUTHORIZED, status.HTTP_403_FORBIDDEN}) +_HTTP_ERROR_STATUSES: Final = range(status.HTTP_400_BAD_REQUEST, 600) + _ERROR_BODY: Final = TypeAdapter(dict[str, object]) _FAILOVER_SQLSTATES: Final = frozenset({"25006", "57P01", "57P02", "57P03"}) @@ -47,13 +49,26 @@ def attribute_of(value: object, name: str, default: object = None) -> object: def error_status_code(exc: object, default: int) -> int: - """The HTTP status an exception carries as ``status_code`` or, the way ``ProxyException`` - stores it, as a stringified ``code``; ``default`` when it carries neither.""" - carried: Final = attribute_of(exc, "status_code") - if isinstance(carried, int) and not isinstance(carried, bool): + """The 4xx or 5xx HTTP status an exception carries as ``status_code`` (an int or a decimal string) or, + the way ``ProxyException`` stores it, as a stringified ``code``; ``default`` when it carries neither.""" + carried: Final = _http_error_status(attribute_of(exc, "status_code")) + if carried is not None: return carried stringified: Final = attribute_of(exc, "code") - return int(stringified) if isinstance(stringified, str) and stringified.isdecimal() else default + from_code: Final = _http_error_status(stringified) if isinstance(stringified, str) else None + return default if from_code is None else from_code + + +def _http_error_status(value: object) -> int | None: + match value: + case bool(): + return None + case int() if value in _HTTP_ERROR_STATUSES: + return value + case str() if len(value) == 3 and value.isascii() and value.isdigit() and int(value) in _HTTP_ERROR_STATUSES: + return int(value) + case _: + return None def openai_error_type(exc: object, status_code: int) -> str: @@ -105,9 +120,9 @@ def proxy_exception_for(exc: Exception, default_status_code: int) -> ProxyExcept code=status.HTTP_503_SERVICE_UNAVAILABLE, ) case HTTPException(status_code=status_code, detail=detail): - return _openai_shaped_proxy_exception(exc, str(detail), status_code) + return openai_shaped_proxy_exception(exc, str(detail), status_code) case _: - return _openai_shaped_proxy_exception( + return openai_shaped_proxy_exception( exc, str(exc) or type(exc).__name__, error_status_code(exc, default_status_code) ) @@ -176,14 +191,19 @@ def _route_failure_error_type(exc: Exception, raised_status: int, answered_statu return openai_error_type(exc, answered_status) -def _openai_shaped_proxy_exception(exc: Exception, message: str, status_code: int) -> ProxyException: +def openai_shaped_proxy_exception(exc: Exception, message: str, status_code: int) -> ProxyException: answered_status: Final = answered_status_code(message, status_code) - return ProxyException( - message=message, - type=_catch_all_error_type(exc, answered_status), - param=openai_error_param(exc), - code=answered_status, - ) + return _proxy_exception(exc, message, answered_status, _catch_all_error_type(exc, answered_status)) + + +def status_typed_proxy_exception(exc: HTTPException) -> ProxyException: + message: Final = str(exc.detail) + answered_status: Final = answered_status_code(message, exc.status_code) + return _proxy_exception(exc, message, answered_status, openai_error_type(exc, answered_status)) + + +def _proxy_exception(exc: Exception, message: str, answered_status: int, error_type: str) -> ProxyException: + return ProxyException(message=message, type=error_type, param=openai_error_param(exc), code=answered_status) def _catch_all_error_type(exc: Exception, status_code: int) -> str: diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py index 91f8307283e..e37ce8d14d4 100644 --- a/litellm/proxy/db/accounting_ledger.py +++ b/litellm/proxy/db/accounting_ledger.py @@ -31,6 +31,13 @@ class AccountingConflictError(ValueError): pass +@dataclass(frozen=True, slots=True) +class UsageWrite: + event: InboundUsageEvent | OutboundUsageEvent + envelope: str + targets: tuple[DeliveryTarget, ...] = () + + _Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"] @@ -107,6 +114,35 @@ SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_ev _LOCK_EVENT: Final = """ SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE """ +_INSERT_USAGE_BATCH: Final = """ +INSERT INTO "LiteLLM_AccountingEvent" + (event_id, kind, schema_version, payload_hash, occurred_at, payload) +SELECT event_id, kind::"AccountingEventKind", 1, payload_hash, occurred_at::timestamp, payload +FROM jsonb_to_recordset($1::jsonb) AS records( + event_id text, kind text, payload_hash text, occurred_at text, payload jsonb) +ORDER BY event_id +ON CONFLICT (event_id) DO NOTHING RETURNING event_id +""" +_LOCK_USAGE_BATCH: Final = """ +SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent" +WHERE event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb)) +ORDER BY event_id FOR UPDATE +""" +_LOAD_USAGE_BATCH: Final = """ +SELECT event.payload, ARRAY( + SELECT target::text FROM "LiteLLM_AccountingDelivery" delivery + WHERE delivery.event_id = event.event_id ORDER BY target::text +) AS targets +FROM "LiteLLM_AccountingEvent" event +WHERE event.event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb)) + AND event.kind = 'outbound_usage' AND event.schema_version = 1 +""" +_INSERT_USAGE_DELIVERIES: Final = """ +INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) +SELECT records.event_id, destination.value::"AccountingDeliveryTarget", 1 +FROM jsonb_to_recordset($1::jsonb) AS records(event_id text, targets jsonb) +CROSS JOIN LATERAL jsonb_array_elements_text(records.targets) AS destination +""" _COMPLETED_FAMILY: Final = """ root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') AND root.recorded_at < $1::timestamp @@ -199,6 +235,53 @@ class AccountingLedger: return cls(transaction) + async def insert_usage_batch(self, writes: tuple[UsageWrite, ...]) -> None: + records: Final = tuple(_prepare(write.event, write.targets, write.envelope, 1, 1) for write in writes) + if not records: + return + unique: Final = {record.event_id: record for record in records} + if any(record.payload_hash != unique[record.event_id].payload_hash for record in records): + raise AccountingConflictError("Batch contains conflicting copies of a usage identity") + batch: Final = _json( + tuple( + { + "event_id": record.event_id, + "kind": record.kind.value, + "payload_hash": record.payload_hash, + "occurred_at": record.occurred_at.astimezone(timezone.utc).isoformat(), + "payload": _JSON.validate_json(record.payload), + } + for record in unique.values() + ) + ) + async with self._transaction() as db: + inserted: Final = _STRINGS.validate_python(await db.query_raw(_INSERT_USAGE_BATCH, batch)) + accepted: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_USAGE_BATCH, _json(tuple(unique)))) + if len(accepted) != len(unique) or any( + row["payload_hash"] != unique[row["event_id"]].payload_hash for row in accepted + ): + raise AccountingConflictError("Usage identity conflicts with accepted content or delivery set") + deliveries: Final = _json( + tuple({"event_id": row["event_id"], "targets": unique[row["event_id"]].targets} for row in inserted) + ) + await db.execute_raw(_INSERT_USAGE_DELIVERIES, deliveries) + + async def load_usage_batch(self, event_ids: tuple[str, ...]) -> tuple[UsageWrite, ...]: + if not event_ids: + return () + async with self._transaction() as db: + rows: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python( + await db.query_raw(_LOAD_USAGE_BATCH, _json(event_ids)) + ) + return tuple( + UsageWrite( + OutboundUsageEvent.model_validate(_JSON.validate_python(row["payload"])["event"]), + _json(row["payload"]), + TypeAdapter(tuple[DeliveryTarget, ...]).validate_python(row["targets"]), + ) + for row in rows + ) + async def insert_event( self, event: AccountingEvent, diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index fe804a82e96..d9fb236632b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -34,6 +34,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.core_helpers import finite_cost, single_log_line from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( @@ -652,7 +653,14 @@ class DBSpendUpdateWriter: end_time=end_time, llm_router=llm_router, ) - payload["spend"] = response_cost or 0.0 + billable_cost: Final = finite_cost(response_cost) + if response_cost is not None and billable_cost is None: + verbose_proxy_logger.warning( + "Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s", + response_cost, + single_log_line(payload["request_id"]), + ) + payload["spend"] = billable_cost or 0.0 if isinstance(payload["startTime"], datetime): payload["startTime"] = payload["startTime"].isoformat() if isinstance(payload["endTime"], datetime): @@ -702,7 +710,7 @@ class DBSpendUpdateWriter: spawn_logging_task( asyncio.get_running_loop(), self._batch_database_updates( - response_cost=response_cost, + response_cost=billable_cost, user_id=user_id, hashed_token=hashed_token, team_id=team_id, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 3a4375ac54a..459f95972ce 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -8,7 +8,7 @@ import asyncio import json from collections.abc import Mapping, Sequence from datetime import datetime -from functools import reduce +from functools import partial, reduce from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast @@ -109,6 +109,16 @@ _ValueT = TypeVar("_ValueT") _DECODED_JSON: Final = TypeAdapter(object) +def _zero_for_non_finite_constant(constant: str) -> float: + verbose_proxy_logger.warning( + "Redis spend buffer - reading a parked %s spend as 0, the spend it held is lost", constant + ) + return 0.0 + + +_load_buffered_transaction: Final = partial(json.loads, parse_constant=_zero_for_non_finite_constant) + + def _spend_log_json_default(value: object) -> str: return value.isoformat() if isinstance(value, datetime) else str(value) @@ -118,7 +128,7 @@ def _encode_spend_log_row(row: SpendLogRow) -> str: def _decode_spend_log_row(encoded: str) -> dict[str, object] | None: - decoded: Final = _DECODED_JSON.validate_python(json.loads(encoded)) + decoded: Final = _DECODED_JSON.validate_python(_load_buffered_transaction(encoded)) return decoded if isinstance(decoded, dict) else None @@ -825,13 +835,13 @@ class RedisUpdateBuffer: if slot_result is None: daily_results.append(None) else: - list_of_daily = [json.loads(t) for t in slot_result] + list_of_daily = [_load_buffered_transaction(t) for t in slot_result] aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily) daily_results.append(aggregated) window_spend: Final = ( WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - tuple(json.loads(transaction) for transaction in raw_results[6]) + tuple(_load_buffered_transaction(transaction) for transaction in raw_results[6]) ) if raw_results[6] is not None else None @@ -880,7 +890,7 @@ class RedisUpdateBuffer: ) if list_of_transactions is None: return None - return [json.loads(transaction) for transaction in list_of_transactions] + return [_load_buffered_transaction(transaction) for transaction in list_of_transactions] async def get_all_daily_spend_update_transactions_from_redis_buffer( self, @@ -1010,9 +1020,9 @@ class RedisUpdateBuffer: Parses the list of transactions from Redis """ if isinstance(list_of_transactions, list): - return [json.loads(transaction) for transaction in list_of_transactions] + return [_load_buffered_transaction(transaction) for transaction in list_of_transactions] else: - return [json.loads(list_of_transactions)] + return [_load_buffered_transaction(list_of_transactions)] @staticmethod def _combine_list_of_transactions( diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index baf53604ea5..13a14e95e56 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -72,6 +72,22 @@ def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]: return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException)) +def _prisma_error_code(e: Exception) -> str | None: + import prisma.engine.errors + + if isinstance(e, _exception_types(prisma.errors.DataError)): + code: Final[object] = getattr(e, "code", None) + return code if isinstance(code, str) else None + if not isinstance(e, _exception_types(prisma.engine.errors.EngineRequestError)): + return None + try: + body: Final = _DATABASE_ERROR_META.validate_json(str(e).partition(": ")[2]) + except ValidationError: + return None + engine_code: Final = body.get("error_code") + return engine_code if isinstance(engine_code, str) else None + + class PrismaDBExceptionHandler: """ Class to handle DB Exceptions or Connection Errors @@ -206,6 +222,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return False + if _prisma_error_code(e) == "P1017": + return True if isinstance(e, _exception_types(prisma.errors.PrismaError)): error_message: Final = str(e).lower() connection_keywords: Final = ( diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 52a0be3d49e..50e6aa54da4 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -4,6 +4,7 @@ to a writer client. Used when DATABASE_URL_READ_REPLICA is configured; otherwise PrismaClient uses the writer-only PrismaWrapper directly. """ +import asyncio import os from collections.abc import Callable from datetime import timedelta @@ -233,8 +234,13 @@ class RoutingPrismaWrapper: await self._reader.start_token_refresh_task() async def stop_token_refresh_task(self) -> None: - await self._writer.stop_token_refresh_task() - await self._reader.stop_token_refresh_task() + results: Final = await asyncio.gather( + self._writer.stop_token_refresh_task(), + self._reader.stop_token_refresh_task(), + return_exceptions=True, + ) + if (error := next((result for result in results if isinstance(result, BaseException)), None)) is not None: + raise error async def recreate_prisma_client( self, diff --git a/litellm/proxy/deps/__init__.py b/litellm/proxy/deps/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/deps/database.py b/litellm/proxy/deps/database.py new file mode 100644 index 00000000000..adeab38aadb --- /dev/null +++ b/litellm/proxy/deps/database.py @@ -0,0 +1,22 @@ +from typing import Annotated, Final, TypeAlias + +from fastapi import Depends, Request + +from litellm.repositories.database import Database + +_MISSING_DATABASE: Final = object() + + +class DatabaseStateMissingError(RuntimeError): + def __init__(self) -> None: + super().__init__("The proxy lifespan did not publish a Database in request state") + + +async def get_database(request: Request) -> Database | None: + database: Final[object] = getattr(request.state, "database", _MISSING_DATABASE) + if database is None or isinstance(database, Database): + return database + raise DatabaseStateMissingError() + + +DatabaseDep: TypeAlias = Annotated[Database | None, Depends(get_database)] diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py new file mode 100644 index 00000000000..e37ecfd582a --- /dev/null +++ b/litellm/proxy/deps/repositories.py @@ -0,0 +1,36 @@ +from typing import Annotated, TypeAlias + +from fastapi import Depends + +from litellm.models.deployment_warning import DeploymentWarningIgnoreStore +from litellm.models.gateway_activity import GatewayActivityStore +from litellm.models.user_banner import UserBannerStore +from litellm.proxy.deps.database import DatabaseDep +from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository +from litellm.repositories.gateway_activity_repository import GatewayActivityRepository +from litellm.repositories.user_banner_repository import UserBannerRepository + + +async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None: + return None if database is None else UserBannerRepository(database.prisma) + + +UserBannerRepositoryDep: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)] + + +async def get_deployment_warning_ignore_repository(database: DatabaseDep) -> DeploymentWarningIgnoreStore | None: + return None if database is None else DeploymentWarningIgnoreRepository(database.prisma, use_writer=True) + + +DeploymentWarningIgnoreRepositoryDep: TypeAlias = Annotated[ + DeploymentWarningIgnoreStore | None, Depends(get_deployment_warning_ignore_repository) +] + + +async def get_gateway_activity_repository(database: DatabaseDep) -> GatewayActivityStore | None: + return None if database is None else GatewayActivityRepository(database.prisma) + + +GatewayActivityRepositoryDep: TypeAlias = Annotated[ + GatewayActivityStore | None, Depends(get_gateway_activity_repository) +] diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index a5822bfb0e3..363b9162e00 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX +from litellm.proxy.common_utils.openai_error_payload import status_typed_proxy_exception from litellm.proxy.common_utils.path_utils import is_within, safe_join from litellm.proxy.guardrails.content_filter_data import CATEGORIES_DIR, DATA_ROOTS, category_dirs, find_category_file from litellm.proxy.guardrails.guardrail_hooks.custom_code.bounded_execution import ( @@ -2533,7 +2534,7 @@ async def apply_guardrail( e = transformed_exception except Exception: verbose_proxy_logger.exception("apply_guardrail: post_call_failure_hook failed") - raise handle_exception_on_proxy(e) + raise handle_exception_on_proxy(status_typed_proxy_exception(e) if isinstance(e, HTTPException) else e) # Success logging outside except so a hook error never triggers failure handlers. response = await _emit_guardrail_success_logs( diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index dac802a1ddc..0b762b2d5fe 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -1,10 +1,10 @@ from __future__ import annotations +import hashlib import json import math import re import time -import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, TypeGuard @@ -28,6 +28,8 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( group_tool_exchanges, has_tool_with_name, ) +from litellm.litellm_core_utils.provider_affinity import get_caller_session_id +from litellm.litellm_core_utils.routing_common_utils import request_user_api_key_auth from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] httpxSpecialProvider, @@ -47,6 +49,11 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, as_converted_stream, ) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HEADROOM_REQUEST_CONTEXT_KEY, + HeadroomRequestContext, +) from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -357,15 +364,18 @@ def _build_headroom_retrieve_tool() -> dict[str, object]: } -def _resolve_call_id(logging_obj: object, request_state: dict[str, object]) -> str | None: - """Resolve the litellm_call_id shared by a request's pre-call hook and its - agentic-loop hooks, so CCR hash validation can be scoped per call instead - of trusting any hash-shaped string that shows up in message text.""" - logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None) - if isinstance(logging_call_id, str) and logging_call_id: - return logging_call_id - kwargs_call_id: Final = request_state.get("litellm_call_id") - return kwargs_call_id if isinstance(kwargs_call_id, str) else None +def _session_affinity(request_data: Mapping[str, object]) -> str | None: + auth: Final = request_user_api_key_auth(request_data) + session_id: Final = get_caller_session_id(request_data) + if not isinstance(auth, UserAPIKeyAuth) or not session_id: + return None + tenant: Final = ( + ("key", auth.api_key) if auth.api_key else ("principal", auth.org_id, auth.team_id, auth.user_id, auth.agent_id) + ) + if not any(tenant[1:]): + return None + identity: Final = json.dumps(["headroom-affinity-v1", tenant, session_id], separators=(",", ":")) + return hashlib.sha256(identity.encode("utf-8")).hexdigest() def has_headroom_retrieve_tool(tools: object) -> bool: @@ -521,7 +531,6 @@ class HeadroomGuardrail(CustomGuardrail): self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) - self._issued_hashes_by_call_id: dict[str, tuple[frozenset[str], float]] = {} super().__init__( # pyright: ignore[reportUnknownMemberType] guardrail_name=guardrail_name, event_hook=event_hook, @@ -540,10 +549,12 @@ class HeadroomGuardrail(CustomGuardrail): value: Final = headers.get(BYPASS_HEADER) return str(value).lower() == "true" - def _request_headers(self) -> dict[str, str]: + def _request_headers(self, affinity: str | None) -> dict[str, str]: headers: Final[dict[str, str]] = {"Content-Type": "application/json"} if self.headroom_api_key: headers["Authorization"] = f"Bearer {self.headroom_api_key}" + if affinity is not None: + headers["x-headroom-affinity"] = affinity return headers @staticmethod @@ -569,14 +580,6 @@ class HeadroomGuardrail(CustomGuardrail): super().update_in_memory_litellm_params(litellm_params) self.timeout = self._resolve_timeout(litellm_params.timeout) - def _prune_expired_hashes(self) -> None: - now: Final = time.monotonic() - self._issued_hashes_by_call_id = { - call_id: (hashes, expiry) - for call_id, (hashes, expiry) in self._issued_hashes_by_call_id.items() - if expiry > now - } - def _handle_compress_failure( self, messages: list[dict[str, object]], @@ -596,16 +599,19 @@ class HeadroomGuardrail(CustomGuardrail): self, messages: list[dict[str, object]], model: str | None, + affinity: str | None, ) -> _CompressResult: payload: Final[dict[str, object]] = {"messages": messages} if model: payload["model"] = model + if self.ccr_retrieval: + payload["config"] = {"mode": "ccr"} try: raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped url=f"{self.headroom_api_base}/v1/compress", json=payload, - headers=self._request_headers(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except httpx.HTTPStatusError as e: @@ -734,7 +740,7 @@ class HeadroomGuardrail(CustomGuardrail): stats["tokens_saved"] = tokens_before - tokens_after return _CompressResult(filtered, True, stats, _read_ccr_hashes(body)) - async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str: + async def _call_retrieve(self, hash_value: str, affinity: str | None, query: str | None = None) -> str: params: Final[dict[str, str]] = {} if query: params["query"] = query @@ -743,7 +749,7 @@ class HeadroomGuardrail(CustomGuardrail): raw_response: HttpxResponse = await self.async_handler.get( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.get is untyped url=f"{self.headroom_api_base}/v1/retrieve/{hash_value}", params=params, - headers=self._request_headers(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: @@ -777,13 +783,14 @@ class HeadroomGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: if input_type != "request": return inputs + request_data.pop(HEADROOM_REQUEST_CONTEXT_KEY, None) if self._should_bypass(request_data): verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER) return inputs @@ -809,9 +816,7 @@ class HeadroomGuardrail(CustomGuardrail): # /v1/compress grows a field for sending the live turn as the retrieval # query without compressing it: query-aware compression reads the newest # user message, so it is withheld here at some cost to history ranking. - # request_data is a bare dict on the base signature; narrow it before - # reading the untranslated messages so long tool names can be recovered. - raw_messages: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data).get("messages") + raw_messages: Final = request_data.get("messages") raw_retrieve_call_ids: Final = _raw_retrieve_call_ids(raw_messages) protected_indices: Final = _protected_indices(messages, raw_retrieve_call_ids) compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices] @@ -820,9 +825,11 @@ class HeadroomGuardrail(CustomGuardrail): model: Final = self.headroom_model or request_data.get("model") start_time: Final = time.time() + affinity: Final = _session_affinity(request_data) result: Final = await self._call_compress( messages=_flatten_messages_for_compression(compressible), model=model if isinstance(model, str) else None, + affinity=affinity, ) end_time: Final = time.time() @@ -868,12 +875,9 @@ class HeadroomGuardrail(CustomGuardrail): if not hashes: return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType] - self._prune_expired_hashes() - call_id = _resolve_call_id(logging_obj, request_data) - if not call_id: - call_id = str(uuid.uuid4()) - request_data["litellm_call_id"] = call_id - self._issued_hashes_by_call_id[call_id] = (frozenset(hashes), time.monotonic() + _HASH_CACHE_TTL_SECONDS) + request_data[HEADROOM_REQUEST_CONTEXT_KEY] = HeadroomRequestContext( + frozenset(hashes), affinity, time.monotonic() + _HASH_CACHE_TTL_SECONDS + ) existing_tools: Final = inputs.get("tools") retrieve_tool: Final = _build_headroom_retrieve_tool() @@ -934,9 +938,11 @@ class HeadroomGuardrail(CustomGuardrail): ) -> AgenticLoopPlan: tool_calls: Final[list[dict[str, object]]] = tools.get("tool_calls", []) - self._prune_expired_hashes() - call_id: Final = _resolve_call_id(logging_obj, kwargs) - valid_hashes = self._issued_hashes_by_call_id.get(call_id, (frozenset(), 0.0))[0] if call_id else frozenset() + context: Final = _REQUEST_DATA_ADAPTER.validate_python(kwargs).get(HEADROOM_REQUEST_CONTEXT_KEY) + valid_context: Final = ( + context if isinstance(context, HeadroomRequestContext) and context.expires_at > time.monotonic() else None + ) + valid_hashes: Final = valid_context.issued_hashes if valid_context is not None else frozenset() retrieved: Final[list[tuple[dict[str, object], str]]] = [] for tc in tool_calls: @@ -944,11 +950,6 @@ class HeadroomGuardrail(CustomGuardrail): raw_hash = arguments.get("hash", "") if isinstance(arguments, dict) else "" hash_value = str(raw_hash).lower() query = arguments.get("query") if isinstance(arguments, dict) else None - # A hash is only honored if it was issued by *this request's own* - # Headroom /v1/compress call, scoped by litellm_call_id. Scoping by - # message text alone is forgeable -- an attacker can plant a - # hash-shaped string in their own prompt, and a hash issued for one - # request would validate for any other request that echoes it back. if hash_value not in valid_hashes: verbose_proxy_logger.warning( "Headroom CCR: rejecting hash=%s not produced by current request compression", @@ -958,6 +959,7 @@ class HeadroomGuardrail(CustomGuardrail): else: content = await self._call_retrieve( hash_value=hash_value, + affinity=valid_context.affinity if valid_context is not None else None, query=str(query) if query else None, ) verbose_proxy_logger.debug("Headroom CCR: retrieved hash=%s (%d chars)", hash_value, len(content)) @@ -996,7 +998,12 @@ class HeadroomGuardrail(CustomGuardrail): max_tokens=max_tokens, optional_params=optional_params_without_max_tokens, kwargs={ - k: v for k, v in kwargs.items() if not k.startswith("_headroom") and k != "litellm_logging_obj" + **{ + k: v + for k, v in _REQUEST_DATA_ADAPTER.validate_python(kwargs).items() + if not k.startswith("_headroom") and k != "litellm_logging_obj" + }, + **({HEADROOM_REQUEST_CONTEXT_KEY: valid_context} if valid_context is not None else {}), }, ), metadata={"tool_type": "headroom_ccr"}, diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py index 342e5b64714..ace0e194ee9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py @@ -14,6 +14,7 @@ from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD +from litellm.router_api import RoutingContext if TYPE_CHECKING: from semantic_router.routers import SemanticRouter @@ -134,7 +135,7 @@ class SemanticGuardRouteLoader: ) encoder: Final = LiteLLMRouterEncoder( - litellm_router_instance=litellm_router, + ctx=RoutingContext(litellm_router), model_name=embedding_model, score_threshold=global_threshold, ) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 20d920f1104..de2ed4fd4ba 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -27,12 +27,12 @@ from litellm.constants import ( HEALTH_CHECK_TIMEOUT_SECONDS, ) from litellm.litellm_core_utils.health_check_helpers import native_health_check_mode +from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin, is_model_capability_field from litellm.router_strategy.auto_router.auto_router_model_naming import ( StrategyRouterDependency, classify_strategy_router_model, strategy_router_dependencies, ) -from litellm.router_utils.common_utils import deployment_model_id_pin, is_model_capability_field from litellm.types.utils import secret_bearing_wif_litellm_params, server_owned_wif_litellm_params # Provider routing and server-owned federation or OAuth token-exchange fields. Allowed for proxy diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index d689ab2a537..cf090d82990 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -13,10 +13,10 @@ from litellm.caching.redis_cache import RedisCache from litellm.constants import ( DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, + HEALTH_CHECKS_TARGET, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.health_check import perform_health_check -from litellm.router_utils.health_state_cache import HEALTH_CHECKS_TARGET if TYPE_CHECKING: from litellm.router import Router diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 59c63a18f0f..333687b07cd 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -7,13 +7,15 @@ import secrets import time import traceback from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from itertools import chain, groupby from types import MappingProxyType from typing import Final, Literal, TypedDict import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status -from pydantic import TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter from typing_extensions import ReadOnly import litellm @@ -24,14 +26,21 @@ from litellm.integrations.SlackAlerting.ms_teams import ( build_ms_teams_payload, get_ms_teams_webhook_url, ) +from litellm.litellm_core_utils.clientside_credential_handler import ( + _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the router path + clientside_credential_keys, +) from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, CallInfo, EnterpriseLicenseData, Litellm_EntityType, + LiteLLM_UserTable, LitellmUserRoles, + ModelAccessDeniedProxyException, ProxyException, SpecialModelNames, UserAPIKeyAuth, @@ -39,13 +48,28 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import ( _resolve_key_models_for_auth_check, # pyright: ignore[reportPrivateUsage] # the auth layer's sentinel resolution, reused so /health scopes exactly like a request + agent_model_grants, + can_agent_call_model, + can_key_call_resolved_model, + can_user_call_model, + get_user_object, + key_model_aliases_for_auth_check, + model_matches_any_wildcard_pattern_in_list, + team_model_aliases_for_auth_check, +) +from litellm.proxy.auth.auth_exception_handler import ( + _as_proxy_exception, # pyright: ignore[reportPrivateUsage] # one auth-failure mapping, shared with user_api_key_auth ) from litellm.proxy.auth.auth_utils import ( _BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check ) -from litellm.proxy.auth.model_checks import get_key_models from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils import http_parsing_utils +from litellm.proxy.common_utils.model_listing_utils import ( + CallerAliases, + dispatched_model, + requestable_alias_names, +) from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -75,10 +99,6 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.router import Router from litellm.router_api import deployment_usable_by_team -from litellm.router_utils.clientside_credential_handler import ( - _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the router path - clientside_credential_keys, -) from litellm.secret_managers.main import get_secret_bool from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields from litellm.types.router import reject_server_owned_wif_fields @@ -253,6 +273,28 @@ def get_callback_identifier(callback): router: Final = APIRouter() _OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) +_DEPLOYMENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +class _ModelGrants(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + models: tuple[object, ...] = () + + +class _ModelInfoAccessGroups(BaseModel): + model_config = ConfigDict(frozen=True) + + access_groups: tuple[object, ...] + + +class _DeploymentAccessGroups(BaseModel): + model_config = ConfigDict(frozen=True) + + model_info: _ModelInfoAccessGroups + + services = ( Literal[ "slack_budget_alerts", @@ -1045,42 +1087,339 @@ def _strip_admin_only_fields_from_health_result(result: dict) -> dict: return out -def _health_accessible_model_names( - user_api_key_dict: UserAPIKeyAuth, llm_router: Router | None -) -> frozenset[str] | None: - """Model names the caller may health-check, or None when the key is unrestricted.""" +def _key_models_unrestricted(user_api_key_dict: UserAPIKeyAuth) -> bool: granted_models: Final = _resolve_key_models_for_auth_check(user_api_key_dict) - if not granted_models or SpecialModelNames.all_proxy_models.value in granted_models: + return not granted_models or SpecialModelNames.all_proxy_models.value in granted_models + + +def _own_team_names(deployment: object, team_id: str | None) -> tuple[str, str] | None: + match deployment: + case { + "model_name": str(internal_name), + "model_info": {"team_id": str(owner_team_id), "team_public_model_name": str(public_name)}, + } if owner_team_id == team_id: + return (internal_name, public_name) + case _: + return None + + +def _names_caller_may_request(deployment: object, team_id: str | None) -> frozenset[str]: + team_names: Final = _own_team_names(deployment, team_id) + if team_names is not None: + return frozenset(team_names) + match deployment: + case {"model_name": str(name)}: + return frozenset((name,)) + case _: + return frozenset() + + +def _public_names_by_internal_name(llm_model_list: Sequence[object], team_id: str | None) -> Mapping[str, str]: + team_names: Final = (_own_team_names(deployment, team_id) for deployment in llm_model_list) + return MappingProxyType(dict(names for names in team_names if names is not None)) + + +async def _key_owner(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_UserTable | None: + if user_api_key_dict.team_id is not None or user_api_key_dict.user_id is None: return None + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + try: + return await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: # noqa: BLE001 # auth reads an unloadable user as no user limit, so /health does too + return None + + +async def _key_may_call_model( + model_name: str, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, + llm_router: Router | None, +) -> bool: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + await asyncio.sleep(0) + try: + await can_key_call_resolved_model( + model=model_name, llm_model_list=None, valid_token=user_api_key_dict, llm_router=llm_router + ) + await can_agent_call_model( + model=model_name, + valid_token=user_api_key_dict, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await can_user_call_model( + model=model_name, + llm_router=llm_router, + user_object=owner, + key_model_aliases=key_model_aliases_for_auth_check(user_api_key_dict), + ) + except ModelAccessDeniedProxyException: + return False + except ProxyException: + raise + except Exception as failure: + if isinstance(failure, HTTPException) and failure.status_code == status.HTTP_403_FORBIDDEN: + return False + raise _as_proxy_exception(failure) from failure + return True + + +async def _caller_aliases(user_api_key_dict: UserAPIKeyAuth) -> CallerAliases: + from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj + + router_settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict, prisma_client, proxy_logging_obj + ) + return CallerAliases( + team_model_aliases_for_auth_check(user_api_key_dict), + key_model_aliases_for_auth_check(user_api_key_dict), + router_settings.get("model_group_alias") if router_settings is not None else None, + ) + + +def _model_auth_checks(model_name: str, aliases: CallerAliases) -> str: + return resolve_model_group_alias(aliases.router, model_name) or model_name + + +def _deployment_id(deployment: object) -> str | None: + match deployment: + case {"model_info": {"id": str(model_id)}}: + return model_id + case _: + return None + + +def _granted_names(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: + grants: Final = ( + *_ModelGrants.model_validate(user_api_key_dict).models, + *_OBJECTS.validate_python(user_api_key_dict.team_models), + ) + return frozenset(name for name in grants if isinstance(name, str)) + + +async def _names_any_limit_grants(user_api_key_dict: UserAPIKeyAuth, owner: LiteLLM_UserTable | None) -> frozenset[str]: + owner_grants: Final = _ModelGrants.model_validate(owner).models if owner is not None else () + return ( + _granted_names(user_api_key_dict) + | frozenset(name for name in owner_grants if isinstance(name, str)) + | await agent_model_grants(user_api_key_dict) + ) + + +def _alias_names_caller_may_request(aliases: CallerAliases, llm_router: Router | None) -> frozenset[str]: + router_aliases: Final = frozenset(llm_router.registry().aliases()) if llm_router is not None else frozenset[str]() + return requestable_alias_names(aliases) | router_aliases + + +def _served_group(model_name: str, aliases: CallerAliases, llm_router: Router | None) -> str: + dispatched: Final = dispatched_model(model_name, aliases) + return llm_router.registry().resolve(dispatched) if llm_router is not None else dispatched + + +def _name_id_pairs(deployment: object, team_id: str | None) -> tuple[tuple[str, str], ...]: + model_id: Final = _deployment_id(deployment) + if model_id is None: + return () + return tuple((name, model_id) for name in _names_caller_may_request(deployment, team_id) | {model_id}) + + +def _ids_by_requested_name(llm_model_list: Sequence[object], team_id: str | None) -> Mapping[str, frozenset[str]]: + pairs: Final = sorted(chain.from_iterable(_name_id_pairs(deployment, team_id) for deployment in llm_model_list)) + return MappingProxyType( + {name: frozenset(model_id for _, model_id in group) for name, group in groupby(pairs, key=lambda pair: pair[0])} + ) + + +def _ids_of(deployments: Iterable[object]) -> frozenset[str]: + return frozenset(model_id for model_id in map(_deployment_id, deployments) if model_id is not None) + + +def _access_groups_of(deployment: object) -> frozenset[str]: + match deployment: + case {"model_info": {"access_groups": list()}}: + access_groups: Final = _DeploymentAccessGroups.model_validate(deployment).model_info.access_groups + return frozenset(access_group for access_group in access_groups if isinstance(access_group, str)) + case _: + return frozenset() + + +@dataclass(frozen=True, slots=True) +class _AccessGroupGrants: + granted_names: frozenset[str] + ids_in_granted_groups: frozenset[str] + + def requested_name(self, name: str, public_names: Mapping[str, str], admitted_names: frozenset[str]) -> str: + public_name: Final = public_names.get(name) + if public_name is None or public_name not in admitted_names or self._grants(name): + return name + return public_name + + def _grants(self, name: str) -> bool: + return name in self.granted_names or model_matches_any_wildcard_pattern_in_list(name, tuple(self.granted_names)) + + def narrow(self, serving_ids: frozenset[str], served_names: frozenset[str]) -> frozenset[str]: + in_granted_groups: Final = serving_ids & self.ids_in_granted_groups + granted_directly: Final = not self.granted_names.isdisjoint(served_names) + return serving_ids if granted_directly or not in_granted_groups else in_granted_groups + + +def _access_group_grants(llm_model_list: Sequence[object], user_api_key_dict: UserAPIKeyAuth) -> _AccessGroupGrants: + granted: Final = _granted_names(user_api_key_dict) + in_granted_groups: Final = ( + deployment for deployment in llm_model_list if not granted.isdisjoint(_access_groups_of(deployment)) + ) + return _AccessGroupGrants(granted, _ids_of(in_granted_groups)) + + +def _ids_serving_group( + group: str, + ids_by_name: Mapping[str, frozenset[str]], + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router | None, + access_group_grants: _AccessGroupGrants, +) -> frozenset[str]: + named: Final = ids_by_name.get(group) + if named is not None: + return access_group_grants.narrow(named, frozenset((group,))) if llm_router is None: - return frozenset(granted_models) - registry: Final = llm_router.registry() + return frozenset() + team_id: Final = user_api_key_dict.team_id + wildcard_served: Final = llm_router.registry().deployments(group, source="wildcard", team_id=team_id) + served_names: Final = frozenset((group,)).union( + *(_names_caller_may_request(deployment, team_id) for deployment in wildcard_served) + ) + return access_group_grants.narrow(_ids_of(wildcard_served), served_names) + + +async def _names_key_may_request( + names: Iterable[str], + aliases: CallerAliases, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, + llm_router: Router | None, +) -> frozenset[str]: return frozenset( - get_key_models( - user_api_key_dict=user_api_key_dict, - proxy_model_list=registry.names(team_id=user_api_key_dict.team_id), - model_access_groups=registry.access_groups(), + [ + name + for name in names + if await _key_may_call_model(_model_auth_checks(name, aliases), user_api_key_dict, owner, llm_router) + ] + ) + + +def _ids_serving( + groups: Iterable[str], + ids_by_name: Mapping[str, frozenset[str]], + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router | None, + access_group_grants: _AccessGroupGrants, +) -> frozenset[str]: + return frozenset[str]().union( + *( + _ids_serving_group(group, ids_by_name, user_api_key_dict, llm_router, access_group_grants) + for group in groups ) ) -def _caller_may_probe_deployment( - deployment: Mapping[str, object], - allowed_models: frozenset[str] | None, +def _shadowed_ids( + llm_model_list: Sequence[object], admitted_names: frozenset[str], reached_ids: frozenset[str], team_id: str | None +) -> frozenset[str]: + ids_of_admitted_names: Final = ( + _deployment_id(deployment) + for deployment in llm_model_list + if not admitted_names.isdisjoint(_names_caller_may_request(deployment, team_id)) + ) + return frozenset(model_id for model_id in ids_of_admitted_names if model_id is not None) - reached_ids + + +@dataclass(frozen=True, slots=True) +class _ReachableDeployments: + model_ids: frozenset[str] + groups: frozenset[str] + + def include(self, deployment: object, team_id: str | None) -> bool: + model_id: Final = _deployment_id(deployment) + if model_id is None: + return not self.groups.isdisjoint(_names_caller_may_request(deployment, team_id)) + return model_id in self.model_ids + + +_NOTHING_REACHABLE: Final = _ReachableDeployments(frozenset(), frozenset()) + + +async def _deployments_key_may_reach( + llm_model_list: Sequence[object], + aliases: CallerAliases, + user_api_key_dict: UserAPIKeyAuth, + owner: LiteLLM_UserTable | None, llm_router: Router | None, +) -> _ReachableDeployments: + candidate_names: Final = (await _names_any_limit_grants(user_api_key_dict, owner)).union( + _alias_names_caller_may_request(aliases, llm_router), + *(_names_caller_may_request(deployment, user_api_key_dict.team_id) for deployment in llm_model_list), + ) + admitted_names: Final = await _names_key_may_request(candidate_names, aliases, user_api_key_dict, owner, llm_router) + access_group_grants: Final = _access_group_grants(llm_model_list, user_api_key_dict) + public_names: Final = _public_names_by_internal_name(llm_model_list, user_api_key_dict.team_id) + groups: Final = frozenset( + _served_group(access_group_grants.requested_name(name, public_names, admitted_names), aliases, llm_router) + for name in admitted_names + ) + ids_by_name: Final = _ids_by_requested_name(llm_model_list, user_api_key_dict.team_id) + reached_ids: Final = _ids_serving(groups, ids_by_name, user_api_key_dict, llm_router, access_group_grants) + admitted_ids: Final = await _names_key_may_request( + _shadowed_ids(llm_model_list, admitted_names, reached_ids, user_api_key_dict.team_id), + aliases, + user_api_key_dict, + owner, + llm_router, + ) + id_groups: Final = frozenset(_served_group(model_id, aliases, llm_router) for model_id in admitted_ids) + return _ReachableDeployments( + reached_ids | _ids_serving(id_groups, ids_by_name, user_api_key_dict, llm_router, access_group_grants), groups + ) + + +def _model_param_target( + model: str, + aliases: CallerAliases, + llm_model_list: Sequence[Mapping[str, object]], team_id: str | None, - caller_is_admin: bool, + llm_router: Router | None, +) -> str: + group: Final = _served_group(model, aliases, llm_router) + if llm_router is None or deployments_targeted_by_name(llm_model_list, group, team_id): + return group + patterns: Final = frozenset( + deployment.get("model_name") + for deployment in llm_router.registry().deployments(group, source="wildcard", team_id=team_id) + ) + match tuple(patterns): + case (str(pattern),): + return pattern + case _: + return group + + +def _caller_may_probe_deployment( + deployment: Mapping[str, object], reachable: _ReachableDeployments, team_id: str | None, caller_is_admin: bool ) -> bool: """Same deployment visibility rule as routing: another team's deployment is never in scope, team-less callers included.""" if not caller_is_admin and not deployment_usable_by_team(deployment, team_id): return False - if allowed_models is None: - return True - if llm_router is None: - return deployment.get("model_name") in allowed_models - model: Final = dict(deployment) - registry: Final = llm_router.registry() - return any(registry.includes(model, model_name=name, team_id=team_id) for name in allowed_models) + return reachable.include(deployment, team_id) def _resolve_targeted_model_ids( @@ -1137,6 +1476,7 @@ def _filter_health_check_results_by_model_ids(results: dict, allowed_model_ids: async def _perform_health_check_and_save( model_list, target_model, + recorded_model_name: str | None, cli_model, details, prisma_client, @@ -1159,8 +1499,7 @@ async def _perform_health_check_and_save( # Optionally save health check result to database (non-blocking) if prisma_client is not None: - # For CLI model, use cli_model name; for router models, use target_model - model_name_for_db: Final = cli_model if cli_model is not None else target_model + model_name_for_db: Final = cli_model if cli_model is not None else recorded_model_name if model_name_for_db is not None: asyncio.create_task( _save_health_check_to_db( @@ -1290,6 +1629,7 @@ async def health_endpoint( cli_result: Final = await _perform_health_check_and_save( model_list=[], target_model=None, + recorded_model_name=None, cli_model=user_model, details=health_check_details, prisma_client=prisma_client, @@ -1304,15 +1644,39 @@ async def health_endpoint( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "Model list not initialized"}, ) - allowed_models: Final = _health_accessible_model_names(user_api_key_dict, llm_router) - restrict_to_allowed_models: Final = not is_admin or allowed_models is not None + restrict_to_allowed_models: Final = not is_admin or not _key_models_unrestricted(user_api_key_dict) + aliases: Final = await _caller_aliases(user_api_key_dict) if restrict_to_allowed_models else None + reachable: Final = ( + await _deployments_key_may_reach( + _OBJECTS.validate_python(llm_model_list), + aliases, + user_api_key_dict, + await _key_owner(user_api_key_dict), + llm_router, + ) + if aliases is not None + else _NOTHING_REACHABLE + ) + requested_model: Final = ( + _model_param_target( + model, + aliases, + _DEPLOYMENTS.validate_python(llm_model_list), + user_api_key_dict.team_id, + llm_router, + ) + if model is not None and aliases is not None + else model + ) _llm_model_list: Final = [ m for m in copy.deepcopy(llm_model_list) if not restrict_to_allowed_models - or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id, is_admin) + or _caller_may_probe_deployment(m, reachable, user_api_key_dict.team_id, is_admin) ] - targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id, user_api_key_dict.team_id) + targeted_ids: Final = _resolve_targeted_model_ids( + _llm_model_list, requested_model, model_id, user_api_key_dict.team_id + ) if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -1370,7 +1734,8 @@ async def health_endpoint( else: router_result: Final = await _perform_health_check_and_save( model_list=_llm_model_list, - target_model=target_model, + target_model=target_model if model_id else requested_model, + recorded_model_name=target_model, cli_model=None, details=health_check_details, prisma_client=prisma_client, diff --git a/litellm/proxy/health_endpoints/deployment_warning_endpoints.py b/litellm/proxy/health_endpoints/deployment_warning_endpoints.py index 115068f79f8..3936eaba73f 100644 --- a/litellm/proxy/health_endpoints/deployment_warning_endpoints.py +++ b/litellm/proxy/health_endpoints/deployment_warning_endpoints.py @@ -1,4 +1,3 @@ -import json import logging from collections.abc import Mapping from dataclasses import dataclass @@ -7,12 +6,14 @@ from types import MappingProxyType from typing import Annotated, Final, Literal, Protocol, TypeAlias from fastapi import APIRouter, Depends, HTTPException, status -from pydantic import AwareDatetime, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, model_validator +from pydantic import ConfigDict, Field, JsonValue, TypeAdapter, model_validator -from litellm._logging import verbose_logger, verbose_proxy_logger +from litellm._logging import verbose_logger from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.models.deployment_warning import AdminIgnore, DeploymentWarningIgnoreStore from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import DeploymentWarningIgnoreRepositoryDep from litellm.proxy.health_endpoints._health_endpoints import ( DISABLE_NO_REDIS_WARNING_ENV_VAR, env_credential_login_enabled, @@ -22,7 +23,6 @@ from litellm.proxy.health_endpoints._health_endpoints import ( ) from litellm.repositories.deployment_warning_ignore_repository import ( DEPLOYMENT_WARNING_IGNORES_ROW_ID, - DeploymentWarningIgnoreRepository, ) from litellm.types.llms.base import LiteLLMBaseModel @@ -50,18 +50,6 @@ class WarningCondition: env_var_ignore: str | None = None -class AdminIgnore(LiteLLMBaseModel): - model_config = ConfigDict(frozen=True) - - source: Literal["admin"] = "admin" - ignored_by: str - ignored_at: AwareDatetime - ignored_until: AwareDatetime | None = Field( - default=None, - description="When the snooze ends. Null means the warning is ignored until an admin restores it.", - ) - - class EnvVarIgnore(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) @@ -109,12 +97,6 @@ class DeploymentWarningIgnoreResponse(LiteLLMBaseModel): ignore: AdminIgnore | None -class DeploymentWarningIgnoreStore(Protocol): - async def get_raw_ignores(self) -> object: ... - - async def upsert_ignores(self, payload: str) -> None: ... - - class ConfigAuditLogger(Protocol): async def __call__( self, @@ -127,7 +109,6 @@ class ConfigAuditLogger(Protocol): ) -> None: ... -_IGNORES_ADAPTER: Final[TypeAdapter[Mapping[str, AdminIgnore]]] = TypeAdapter(Mapping[str, AdminIgnore]) _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _AUDIT_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[DeploymentWarningIgnoreResponse, ...]]] = TypeAdapter( tuple[DeploymentWarningIgnoreResponse, ...] @@ -185,18 +166,8 @@ def is_known_warning_id(warning_id: str) -> bool: return False -def _parse_stored_ignores(raw: object) -> Mapping[str, AdminIgnore]: - if raw is None: - return MappingProxyType({}) - try: - return MappingProxyType(_IGNORES_ADAPTER.validate_python(json.loads(raw) if isinstance(raw, str) else raw)) - except (json.JSONDecodeError, ValidationError): - verbose_proxy_logger.warning("Ignoring unreadable %s row", DEPLOYMENT_WARNING_IGNORES_ROW_ID) - return MappingProxyType({}) - - async def _read_unexpired_ignores(store: DeploymentWarningIgnoreStore, now: datetime) -> Mapping[str, AdminIgnore]: - stored: Final = _parse_stored_ignores(await store.get_raw_ignores()) + stored: Final = await store.get_ignores() return MappingProxyType( { warning_id: ignore @@ -225,14 +196,6 @@ def utc_now() -> datetime: return datetime.now(timezone.utc) -def deployment_warning_ignore_store() -> DeploymentWarningIgnoreStore | None: - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - return None - return DeploymentWarningIgnoreRepository(prisma_client, use_writer=True) - - def config_audit_logger() -> ConfigAuditLogger: from litellm.proxy.proxy_server import create_config_audit_log @@ -281,7 +244,7 @@ async def current_warning_conditions(now: Annotated[datetime, Depends(utc_now)]) async def get_deployment_warnings( _: Annotated[UserAPIKeyAuth, Depends(_deployment_warning_reader)], conditions: Annotated[tuple[WarningCondition, ...], Depends(current_warning_conditions)], - store: Annotated[DeploymentWarningIgnoreStore | None, Depends(deployment_warning_ignore_store)], + store: DeploymentWarningIgnoreRepositoryDep, now: Annotated[datetime, Depends(utc_now)], ) -> DeploymentWarningsResponse: """ @@ -309,7 +272,7 @@ async def update_deployment_warning_ignore( warning_id: str, update: DeploymentWarningIgnoreUpdate, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_deployment_warning_admin)], - store: Annotated[DeploymentWarningIgnoreStore | None, Depends(deployment_warning_ignore_store)], + store: DeploymentWarningIgnoreRepositoryDep, audit_logger: Annotated[ConfigAuditLogger, Depends(config_audit_logger)], now: Annotated[datetime, Depends(utc_now)], ) -> DeploymentWarningIgnoreResponse: @@ -338,7 +301,7 @@ async def update_deployment_warning_ignore( others: Final = {other_id: other for other_id, other in before.items() if other_id != warning_id} after: Final = MappingProxyType(others if ignore is None else {**others, warning_id: ignore}) - await store.upsert_ignores(_IGNORES_ADAPTER.dump_json(dict(after)).decode()) + await store.save_ignores(after) await audit_logger( param_name=DEPLOYMENT_WARNING_IGNORES_ROW_ID, action="updated", diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 0011f5cd082..dac2d219b38 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -15,6 +15,10 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, @@ -34,10 +38,6 @@ from litellm.proxy.hooks.rate_limiter_utils import ( resolve_llm_provider_for_rate_limit, ) from litellm.proxy.utils import InternalUsageCache -from litellm.router_utils.add_retry_fallback_headers import ( - ensure_response_additional_headers, - response_has_hidden_params, -) from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral, LLMResponseTypes diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 1c5d6e05f09..1c10826124e 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -204,7 +204,7 @@ def _team_public_names_by_group(router: "Router", team_id: str | None) -> Mappin groups_by_name: Final = MappingProxyType( { name: model_access_groups_of(router=router, model=name, team_id=team_id) - for name in router.team_public_model_names_of(team_id) + for name in router.registry().names(kind="team_public", team_id=team_id) } ) return MappingProxyType( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 15282679622..a515c051bbd 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -49,6 +49,10 @@ from litellm.constants import ( ) from litellm.cost_calculator import get_usage_object from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -56,6 +60,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( split_media_for_tpm, ) from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled +from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.token_counter import estimate_image_tokens, offload_token_count from litellm.llms.azure.ptu_capacity import normalized_tokens from litellm.proxy._types import UserAPIKeyAuth @@ -86,12 +92,6 @@ from litellm.proxy.spend_tracking.input_tokens import ( openai_chat_pdf_file_tokens, ) from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction -from litellm.router_utils.add_retry_fallback_headers import ( - ensure_response_additional_headers, - response_has_hidden_params, -) -from litellm.router_utils.common_utils import resolve_model_group_alias -from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.extraction import ExtractionResponse from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage @@ -3561,9 +3561,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns True if any deployment has failures in the current minute. """ from litellm.proxy.proxy_server import llm_router - from litellm.router_utils.router_callbacks.track_deployment_metrics import ( - get_deployment_failures_for_current_minute, - ) if llm_router is None: return False @@ -3581,10 +3578,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): continue # Use router's existing failure tracking - failure_count = get_deployment_failures_for_current_minute( - litellm_router_instance=llm_router, - deployment_id=deployment_id, - ) + failure_count = llm_router.health().failures_this_minute(deployment_id) if failure_count > DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE: verbose_proxy_logger.debug( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 9c4a4bc775b..2951bc02d8c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,10 +12,12 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final +from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, + finite_cost, get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, get_parent_otel_span_from_kwargs, @@ -133,6 +135,8 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) ) +FAILURE_SPEND_COUNTER_WAIT_SECONDS: Final = 1.0 + _OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) @@ -234,6 +238,12 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +def _proxy_spend_counters() -> "_IncrementSpendCounters": + from litellm.proxy.proxy_server import increment_spend_counters + + return increment_spend_counters + + DeploymentLookup: TypeAlias = Callable[[str], Deployment | None] @@ -360,11 +370,16 @@ class _ProxyDBLogger(CustomLogger): message_logging: bool = True, spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer, deployment_lookup: DeploymentLookup = _proxy_deployment, + spend_counters: Callable[[], "_IncrementSpendCounters"] = _proxy_spend_counters, + failure_counter_wait_seconds: float = FAILURE_SPEND_COUNTER_WAIT_SECONDS, ) -> None: super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging) self.spend_event_producer = spend_event_producer self._spend_writer: Final = spend_writer self._deployment_lookup: Final = deployment_lookup + self._spend_counters: Final = spend_counters + self._failure_counter_wait_seconds: Final = failure_counter_wait_seconds + self._pending_increments: Final[set[asyncio.Task[None]]] = set() # mutable-ok: tasks leave when done async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime @@ -589,9 +604,11 @@ class _ProxyDBLogger(CustomLogger): completed_cost: bool = False, request_tags: list[str] | None = None, ) -> None: - """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair.""" - for row_cost, row_kwargs in spend_rows: - if completed_cost: + """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair, then the + live spend counters raised by every charged row that cost something. A completed cost settles through the + success path's write, which raises the counters against the kept reservation.""" + if completed_cost: + for row_cost, row_kwargs in spend_rows: from litellm.proxy.proxy_server import increment_spend_counters, proxy_logging_obj await _update_database_and_spend_counters( @@ -622,12 +639,16 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ), ) - continue + return + + token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) + end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id + charged: Final = [ await self._spend_writer().update_database( - token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), + token=token, response_cost=row_cost, user_id=user_api_key_dict.user_id, - end_user_id=None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id, + end_user_id=end_user_id, team_id=user_api_key_dict.team_id, kwargs=dict(row_kwargs), completion_response=original_exception, @@ -636,6 +657,91 @@ class _ProxyDBLogger(CustomLogger): org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, ) + for row_cost, row_kwargs in spend_rows + ] + counted_rows: Final = tuple( + row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0 + ) + if counted_rows and not ProxyUpdateSpend.disable_spend_updates(): + await self._increment_failure_spend_counters( + user_api_key_dict, token, end_user_id, counted_rows, start_time + ) + + async def _increment_failure_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + counted_rows: Sequence[tuple[float, Mapping[str, object]]], + start_time: datetime, + ) -> None: + """The counters the success path raises for the same spend, sent before the client sees the error. A + counter store slower than ``failure_counter_wait_seconds`` finishes them in the background instead of + holding the error back. The reservation was already released, so every counter takes the full row cost.""" + increment: Final = asyncio.create_task( + self._raise_failure_spend_counters(user_api_key_dict, token, end_user_id, counted_rows, start_time) + ) + self._pending_increments.add(increment) + increment.add_done_callback(self._pending_increments.discard) + finished, _ = await asyncio.wait((increment,), timeout=self._failure_counter_wait_seconds) + if not finished: + verbose_proxy_logger.warning( + "Spend counters for failure spend rows took over %ss, finishing them in the background", + self._failure_counter_wait_seconds, + ) + + async def _raise_failure_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + counted_rows: Sequence[tuple[float, Mapping[str, object]]], + start_time: datetime, + ) -> None: + for row_cost, row_kwargs in counted_rows: + await self._raise_row_spend_counters( + user_api_key_dict, token, end_user_id, row_cost, row_kwargs, start_time + ) + await flush_post_call_redis_batches() + + async def _raise_row_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + row_cost: float, + row_kwargs: Mapping[str, object], + start_time: datetime, + ) -> None: + standard_logging_object: Final = row_kwargs.get("standard_logging_object") + sl_object: Final = ( + cast("StandardLoggingPayload", standard_logging_object) # cast-ok: the failure snapshot dict + if isinstance(standard_logging_object, dict) + else None + ) + metadata: Final = _str_keyed_fields(_str_keyed_fields(row_kwargs.get("litellm_params")).get("metadata")) + served_model_id: Final = _str_keyed_fields(metadata.get("model_info")).get("id") or _str_keyed_fields( + standard_logging_object + ).get("model_id") + try: + await self._spend_counters()( + token=token, + team_id=user_api_key_dict.team_id, + user_id=user_api_key_dict.user_id, + response_cost=row_cost, + org_id=user_api_key_dict.org_id, + end_user_id=end_user_id, + tags=_get_request_tags_for_cost_tracking(sl_object=sl_object, metadata=dict(metadata)), + request_started_at=start_time, + model_access_groups=debitable_model_access_groups( + attributed=get_request_model_access_groups(row_kwargs), + served_model_id=served_model_id if isinstance(served_model_id, str) else None, + router=get_llm_router(), + ), + project_id=user_api_key_dict.project_id, + ) + except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing + verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row") async def _PROXY_track_cost_callback( self, @@ -703,7 +809,7 @@ class _ProxyDBLogger(CustomLogger): raw_cost: Final = _str_keyed_fields( sl_object if sl_object is not None else kwargs # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload ).get("response_cost") - priced_cost: Final = raw_cost if isinstance(raw_cost, (int, float)) else None + priced_cost: Final = finite_cost(raw_cost) response_cost: Final = ( await apply_search_free_tier( kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload @@ -1059,7 +1165,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]: return tuple(sorted(str(key) for key in metadata)) -def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: +def _get_budget_reservation_from_metadata(metadata: dict) -> dict[str, object] | None: return budget_reservation_from_metadata(metadata) @@ -1094,6 +1200,7 @@ class _IncrementSpendCounters(Protocol): tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ) -> None: ... @@ -1110,7 +1217,7 @@ async def _update_database_and_spend_counters( start_time: datetime | None, end_time: datetime | None, response_cost: float, - budget_reservation: dict | None, + budget_reservation: dict[str, object] | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, @@ -1173,7 +1280,7 @@ async def _update_database_and_spend_counters_in_batch( start_time: datetime | None, end_time: datetime | None, response_cost: float, - budget_reservation: dict | None, + budget_reservation: dict[str, object] | None, request_tags: list[str] | None, model_access_groups: Sequence[str] | None, project_id: str | None, diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 28fb9baf2aa..6244e5025ce 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -8,8 +8,10 @@ from typing import Final, get_type_hints import orjson from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile, status from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter import litellm +from litellm.constants import MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -17,6 +19,7 @@ from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, + hierarchical_router_settings_updates, log_llm_api_exception, resolve_litellm_call_id, ) @@ -36,6 +39,7 @@ from litellm.types.images.main import ImageEditRequestParams from litellm.types.llms.openai import ChatCompletionUserMessage router: Final = APIRouter() +_IMAGE_GENERATION_MODEL_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object) IMAGE_EDIT_NUMERIC_FORM_FIELDS: Final = numeric_form_fields(get_type_hints(ImageEditRequestParams)) IMAGE_EDIT_OPTIONAL_FIELD_DEFAULTS: Final = MappingProxyType({"prompt": None, "image": None}) @@ -132,6 +136,19 @@ async def image_generation( if data["model"] in litellm.model_alias_map: data["model"] = litellm.model_alias_map[data["model"]] + if llm_router is not None: + requested_model: Final = _IMAGE_GENERATION_MODEL_ADAPTER.validate_python(data.get("model")) + data.update( + await hierarchical_router_settings_updates( + requested_model=requested_model, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + proxy_config=proxy_config, + proxy_logging_obj=proxy_logging_obj, + apply_model_group_alias=request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is not True, + ) + ) + ### CALL HOOKS ### - modify incoming data / reject request before calling the model prompt_value: Final = data.get("prompt") if prompt_value is not None: @@ -156,6 +173,7 @@ async def image_generation( route_type="aimage_generation", llm_router=llm_router, user_model=user_model, + user_api_key_dict=user_api_key_dict, ) response = await llm_call diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9b52f2f8a8d..4444d6d25f9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -48,7 +48,7 @@ from litellm.constants import ( VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) -from litellm.integrations.custom_guardrail import without_server_streaming_classification +from litellm.integrations.custom_guardrail import merged_guardrail_entries, without_server_streaming_classification from litellm.litellm_core_utils.core_helpers import is_codex_user_agent, normalize_drop_params from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -112,6 +112,7 @@ _EXPLICIT_SESSION_HEADERS: Final = frozenset({"x-litellm-trace-id", "x-litellm-s # ``session-id``/``thread-id``; builds before the codex-api split sent # ``session_id``/``conversation_id``. Ordered session before thread. _CODEX_SESSION_ID_HEADERS: Final = ("session-id", "session_id", "thread-id", "conversation_id") +_BEDROCK_SESSION_ID_METADATA_KEYS: Final = ("sessionId", "threadId", "conversationId") # Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores # (covers UUIDs and most common session-id formats). _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") @@ -1055,6 +1056,22 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None: return session_id +def _get_bedrock_session_id_from_request_metadata( + normalized_headers: Mapping[str, str], request_body: Mapping[str, object] +) -> str | None: + """Read the conversation id a Bedrock SDK client tags its calls with: Converse's ``requestMetadata`` + body field or InvokeModel's ``X-Amzn-Bedrock-Request-Metadata`` JSON header. Only the Bedrock + pass-through route nests the client's body under ``data``.""" + if request_body.get("custom_llm_provider") != "bedrock" or not isinstance(request_body.get("data"), Mapping): + return None + header: Final = normalized_headers.get("x-amzn-bedrock-request-metadata") + body_tags: Final = _request_fields(_request_fields(request_body.get("data")).get("requestMetadata")) + header_tags: Final = _request_fields(safe_json_loads(header) if header else None) + tag_sources: Final = itertools.product((body_tags, header_tags), _BEDROCK_SESSION_ID_METADATA_KEYS) + values: Final = (tags.get(key) for tags, key in tag_sources) + return next((value for value in values if isinstance(value, str) and value), None) + + def _is_llm_inference_route(request: Request) -> bool: route: Final = get_request_route(request) return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access( @@ -1967,6 +1984,14 @@ class LiteLLMProxyRequestSetup: body_metadata["user_id"] = session_id verbose_proxy_logger.debug("Extracted session_id from Anthropic metadata.user_id") + if "litellm_session_id" not in data: + bedrock_session_id: Final = _get_bedrock_session_id_from_request_metadata(normalized_headers, request_body) + metadata_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id") + if bedrock_session_id and metadata_session_id in (None, bedrock_session_id): + metadata_from_headers["session_id"] = bedrock_session_id + data["litellm_session_id"] = bedrock_session_id # rebind-ok: data is an out-param + verbose_proxy_logger.debug("Extracted session_id from Bedrock request metadata") + # Last-resort fallback: the W3C standards for trace/session propagation # (https://www.w3.org/TR/trace-context/, https://www.w3.org/TR/baggage/). # Lower priority than everything above - only fires when neither the @@ -2566,6 +2591,7 @@ def refresh_proxy_server_request_body_snapshot( VECTOR_STORE_CALLER_METADATA_KWARG, STREAM_USAGE_INJECTED_FIELD, "_litellm_strip_stream_usage", + "_headroom_request_context", } ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS @@ -3265,7 +3291,7 @@ def _fallback_edges( and its provider-stripped suffix, plus the router matcher run over the fuzzy entries that have not fired yet (tracked in ``fired_fuzzy`` by ``(list_index, entry_index)``), so each entry's targets are expanded once.""" - from litellm.router_utils.fallback_event_handlers import get_fallback_model_group + from litellm.router_api import fallback_chain stripped: Final = model_group.split("/", 1)[1] if "/" in model_group else None targets: Final[set[str]] = set() # mutable-ok: accumulates fallback targets @@ -3276,7 +3302,7 @@ def _fallback_edges( pending = tuple(index for index in range(len(fuzzy)) if (list_index, index) not in fired_fuzzy) if not pending: continue - resolved = get_fallback_model_group(fallbacks=[fuzzy[index] for index in pending], model_group=model_group)[0] + resolved = fallback_chain([fuzzy[index] for index in pending], model_group)[0] values = resolved if isinstance(resolved, list) else ([resolved] if resolved is not None else []) targets.update(name for value in values if (name := _fallback_target_name(value)) is not None) if resolved is not None: @@ -3463,7 +3489,7 @@ async def _resolve_user_provider_credentials_for_request( status_code=400, detail=(f"fallback lists in the request body cannot exceed {_FALLBACK_DISCOVERY_LIMIT} entries"), ) - from litellm.router_utils.common_utils import resolve_model_group_alias + from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias # Discovery must see the group the router actually routes to, so include the # post-alias name from every alias map that may still rewrite data["model"] @@ -3947,30 +3973,14 @@ def add_guardrails_from_key_or_team_metadata( """ from litellm.proxy.utils import premium_user_check - # Initialize guardrails set (avoiding duplicates) - combined_guardrails: Final = set() - - # Add key-level guardrails first - if key_metadata and "guardrails" in key_metadata: - if isinstance(key_metadata["guardrails"], list) and len(key_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(key_metadata["guardrails"]) - - # Add team-level guardrails (set automatically handles duplicates) - if team_metadata and "guardrails" in team_metadata: - if isinstance(team_metadata["guardrails"], list) and len(team_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(team_metadata["guardrails"]) - - # Add project-level guardrails (set automatically handles duplicates) - if project_metadata and "guardrails" in project_metadata: - if isinstance(project_metadata["guardrails"], list) and len(project_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(project_metadata["guardrails"]) - - # Set combined guardrails in metadata as list - if combined_guardrails: - data[metadata_variable_name]["guardrails"] = list(combined_guardrails) + configured_guardrails: Final = tuple( + metadata["guardrails"] + for metadata in (key_metadata, team_metadata, project_metadata) + if metadata and isinstance(metadata.get("guardrails"), list) and len(metadata["guardrails"]) > 0 + ) + if configured_guardrails: + premium_user_check() + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries(*configured_guardrails) _add_guardrails_from_key_or_team_metadata: Final = add_guardrails_from_key_or_team_metadata @@ -4065,14 +4075,10 @@ def _add_guardrails_from_policies_in_metadata( if metadata_variable_name not in data: data[metadata_variable_name] = {} - existing_guardrails = data[metadata_variable_name].get("guardrails", []) - if not isinstance(existing_guardrails, list): - existing_guardrails = [] - - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = set(existing_guardrails) - combined.update(resolved_guardrails) - data[metadata_variable_name]["guardrails"] = list(combined) + existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries( + existing_guardrails if isinstance(existing_guardrails, list) else (), resolved_guardrails + ) # Store applied policies in metadata for tracking if "applied_policies" not in data[metadata_variable_name]: @@ -4307,10 +4313,11 @@ def _apply_resolved_guardrails_to_metadata( return existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) - existing_guardrails_list: Final = existing_guardrails if isinstance(existing_guardrails, list) else [] + existing_guardrails_list: Final = ( + _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () + ) - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = list(dict.fromkeys((*existing_guardrails_list, *resolved_guardrails))) + combined: Final = merged_guardrail_entries(existing_guardrails_list, resolved_guardrails) data[metadata_variable_name]["guardrails"] = combined verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) @@ -4327,7 +4334,7 @@ def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_nam async def add_guardrails_from_policy_engine( - data: dict, + data: dict[str, object], # mutable-ok: pops the request's policies and records the resolved guardrails in place metadata_variable_name: str, user_api_key_dict: UserAPIKeyAuth, ) -> None: @@ -4370,11 +4377,13 @@ async def add_guardrails_from_policy_engine( all_tags: Final = _policy_match_tags(data, metadata_variable_name) _team_alias: Final = user_api_key_dict.team_alias _key_alias: Final = user_api_key_dict.key_alias - context: Final = PolicyMatchContext( - team_alias=_team_alias if isinstance(_team_alias, str) else None, - key_alias=_key_alias if isinstance(_key_alias, str) else None, - model=data.get("model"), - tags=all_tags, + context: Final = PolicyMatchContext.model_validate( + { + "team_alias": _team_alias if isinstance(_team_alias, str) else None, + "key_alias": _key_alias if isinstance(_key_alias, str) else None, + "model": data.get("model"), + "tags": all_tags, + } ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index f4d307a6cda..42321fe1b1d 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -60,6 +60,7 @@ from litellm.repositories.table_repositories import PrismaTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository +from litellm.router_api import RoutingContext from litellm.router_strategy.auto_router.auto_router_model_naming import ( StrategyRouterDependencyRole, classify_strategy_router_model, @@ -614,7 +615,7 @@ async def preview_auto_router_routing( try: complexity_strategy: Final = ComplexityRouter( model_name=resolved.router_name, - litellm_router_instance=llm_router, + ctx=RoutingContext(llm_router), complexity_router_config=resolved.complexity_router_config.model_dump(exclude_none=True), default_model=resolved.complexity_router_config.resolve_default_model(resolved.default_model), derive_savings_baseline=False, diff --git a/litellm/proxy/management_endpoints/daily_activity_scopes.py b/litellm/proxy/management_endpoints/daily_activity_scopes.py index ca1a96f9d84..7ff73ac6e2e 100644 --- a/litellm/proxy/management_endpoints/daily_activity_scopes.py +++ b/litellm/proxy/management_endpoints/daily_activity_scopes.py @@ -216,7 +216,11 @@ def _user_query( start_date: str | None = Query(default=None, description="Start date in YYYY-MM-DD format"), end_date: str | None = Query(default=None, description="End date in YYYY-MM-DD format"), model: str | None = Query(default=None, description="Filter by specific model"), - api_key: str | None = Query(default=None, description="Filter by specific API key"), + api_key: str | None = Query( + default=None, + description="Filter by reporting API key ID. Missing keys are reported as unattributed. " + "For a literal key unattributed or a key starting with raw:, add an extra raw: prefix", + ), user_id: str | None = Query( default=None, description="Filter by specific user ID. Admins can filter by any user or omit for global view. " diff --git a/litellm/proxy/management_endpoints/gateway_request_endpoints.py b/litellm/proxy/management_endpoints/gateway_request_endpoints.py index 4052e6f714e..01c9ad137c2 100644 --- a/litellm/proxy/management_endpoints/gateway_request_endpoints.py +++ b/litellm/proxy/management_endpoints/gateway_request_endpoints.py @@ -13,15 +13,15 @@ and the endpoint is restricted to proxy admin roles. from collections.abc import Sequence from datetime import datetime, timedelta, timezone -from typing import Annotated, Final, Literal +from typing import Annotated, Final from fastapi import APIRouter, Depends, HTTPException, Query -from pydantic import ConfigDict, Field, TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.proxy.deps.repositories import GatewayActivityRepositoryDep from litellm.types.proxy.gateway_requests import ( GatewayRequestActivityResponse, GatewayRequestBreakdownEntry, @@ -33,50 +33,6 @@ router: Final = APIRouter() _DEFAULT_LOOKBACK_DAYS: Final = 30 -_ACTIVITY_SQL: Final = """ - SELECT - 'route' AS kind, date, category, route, - SUM(successful_requests)::bigint AS successful_requests, - SUM(failed_requests)::bigint AS failed_requests, - NULL::integer AS status_code - FROM "LiteLLM_DailyGatewayRequests" - WHERE date >= $1 AND date <= $2 - GROUP BY date, category, route - UNION ALL - SELECT - 'status' AS kind, NULL, NULL, NULL, - 0::bigint, - SUM(failed_requests)::bigint, - status_code - FROM "LiteLLM_DailyGatewayFailedRequests" - WHERE date >= $1 AND date <= $2 - GROUP BY status_code -""" - - -class _AggregateRow(LiteLLMBaseModel): - """Validates one query_raw row so the handler works with typed values, not Any.""" - - kind: Literal["route"] - date: str - category: str - route: str - successful_requests: int - failed_requests: int - - -class _StatusCodeAggregateRow(LiteLLMBaseModel): - model_config = ConfigDict(frozen=True) - - kind: Literal["status"] - status_code: int - failed_requests: int - - -_ACTIVITY_ROWS_ADAPTER: Final = TypeAdapter( - tuple[Annotated[_AggregateRow | _StatusCodeAggregateRow, Field(discriminator="kind")], ...] -) - def _default_range() -> tuple[str, str]: end: Final = datetime.now(timezone.utc) @@ -84,7 +40,7 @@ def _default_range() -> tuple[str, str]: return start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d") -def _fold_by_date(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestDailyEntry, ...]: +def _fold_by_date(rows: Sequence[GatewayRouteActivity]) -> tuple[GatewayRequestDailyEntry, ...]: dates: Final = sorted(frozenset(row.date for row in rows)) return tuple( GatewayRequestDailyEntry( @@ -96,7 +52,7 @@ def _fold_by_date(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestDailyEnt ) -def _fold_by_route(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestBreakdownEntry, ...]: +def _fold_by_route(rows: Sequence[GatewayRouteActivity]) -> tuple[GatewayRequestBreakdownEntry, ...]: pairs: Final = sorted(frozenset((row.category, row.route) for row in rows)) entries: Final = tuple( GatewayRequestBreakdownEntry( @@ -113,7 +69,7 @@ def _fold_by_route(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestBreakdo def _fold_by_status_code( - rows: Sequence[_StatusCodeAggregateRow], + rows: Sequence[GatewayStatusCodeActivity], ) -> tuple[GatewayRequestStatusCodeEntry, ...]: return tuple( GatewayRequestStatusCodeEntry(status_code=row.status_code, failed_requests=row.failed_requests) @@ -128,6 +84,7 @@ def _fold_by_status_code( ) async def get_gateway_daily_activity( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + activity: GatewayActivityRepositoryDep, start_date: str | None = Query(default=None, description="Start date in YYYY-MM-DD format"), end_date: str | None = Query(default=None, description="End date in YYYY-MM-DD format"), ) -> GatewayRequestActivityResponse: @@ -137,8 +94,6 @@ async def get_gateway_daily_activity( Deployment-wide: the underlying table has no per-key or per-user dimension, so this is admin-only. """ - from litellm.proxy.proxy_server import prisma_client - if user_api_key_dict.user_role not in ( LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, @@ -148,22 +103,15 @@ async def get_gateway_daily_activity( detail="Only proxy admin roles can view gateway request counts across the deployment", ) - if prisma_client is None: + if activity is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) default_start, default_end = _default_range() selected_start: Final = start_date or default_start selected_end: Final = end_date or default_end - raw_rows: Final = await prisma_client.db.query_raw( # pyright: ignore[reportAny] # untyped prisma client - _ACTIVITY_SQL, - selected_start, - selected_end, - ) - activity_rows: Final = _ACTIVITY_ROWS_ADAPTER.validate_python(raw_rows or ()) - rows: Final[tuple[_AggregateRow, ...]] = tuple(row for row in activity_rows if isinstance(row, _AggregateRow)) - status_code_rows: Final[tuple[_StatusCodeAggregateRow, ...]] = tuple( - row for row in activity_rows if isinstance(row, _StatusCodeAggregateRow) - ) + activity_rows: Final = await activity.get_activity(selected_start, selected_end) + rows: Final = tuple(row for row in activity_rows if isinstance(row, GatewayRouteActivity)) + status_code_rows: Final = tuple(row for row in activity_rows if isinstance(row, GatewayStatusCodeActivity)) verbose_proxy_logger.debug("/gateway/daily/activity - aggregated %d rows", len(rows)) return GatewayRequestActivityResponse( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6950cf08a25..54a2d4062d2 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2922,7 +2922,8 @@ async def get_user_daily_activity( ), api_key: str | None = fastapi.Query( default=None, - description="Filter by specific API key", + description="Filter by reporting API key ID. Missing keys are reported as unattributed. " + "For a literal key unattributed or a key starting with raw:, add an extra raw: prefix", ), user_id: str | None = fastapi.Query( default=None, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a6814b593ca..61da5599886 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo _delete_cache_key_object, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export can_team_access_model, delete_cache_key_object, + delete_cache_key_objects, get_jwt_key_mapping_cache_keys_for_token, get_key_end_user_budget_id, get_org_object, @@ -78,7 +79,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time -from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache, is_user_key_cache_key from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.management.teams.authz import TEAM_ADMIN_ONLY, TEAM_OR_ORG_ADMIN, is_team_admin @@ -161,7 +162,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # _is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export is_master_key, ) -from litellm.proxy.utils import ( # noqa: F401 # legacy module exports +from litellm.proxy.utils import ( PrismaClient, ProxyLogging, _hash_token_if_needed, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -279,6 +280,10 @@ class _BudgetRowSoftBudgetCreate(TypedDict): updated_by: ReadOnly[str] +class _KeyRowDumpable(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + class _KeyUpdateTx(Protocol): @property def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ... @@ -286,6 +291,9 @@ class _KeyUpdateTx(Protocol): @property def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ... + @property + def litellm_objectpermissiontable(self) -> "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]": ... + class _ConfigTableActions(Protocol): """Config table surface this module needs; the shared repository seam exposes no ``update``.""" @@ -2071,6 +2079,56 @@ async def _check_project_key_limits( ) +async def _validate_project_assignment( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Validate assigning a project to a key that does not have one yet.""" + if data.project_id is None or data.project_id == existing_key_row.project_id: + return + if existing_key_row.project_id is not None: + raise HTTPException( + status_code=400, detail="Project reassignment is not supported. Use null to detach the key." + ) + project_obj: Final = await get_project_object( + project_id=data.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if project_obj is None: + raise HTTPException( + status_code=404, + detail={"error": f"Project not found, project_id={data.project_id}"}, + ) + team: Final = data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id + if team is None or team != project_obj.team_id: + raise HTTPException( + status_code=400, + detail={ + "error": f"Project {data.project_id} belongs to team {project_obj.team_id}. Keys can only be assigned to a project owned by their own team (key team: {team})" + }, + ) + await _check_project_key_limits( + project_id=data.project_id, + data=data.model_copy( + update={ + "models": ( + data.models + if "models" in data.model_fields_set + else cast( # cast-ok: LiteLLM_VerificationToken.models is a bare list + "list[str] | None", existing_key_row.models + ) + ), + "max_budget": data.max_budget if "max_budget" in data.model_fields_set else existing_key_row.max_budget, + } + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + async def _check_key_project_team( project_id: str, key_team_id: str | None, @@ -2770,6 +2828,115 @@ async def _apply_soft_budget_update( return remaining +async def _apply_object_permission_update( + data_json: Mapping[str, object], + existing_object_permission_id: str | None, + prisma_client: PrismaClient, + table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]", +) -> Mapping[str, object]: + """Upsert the requested object permission row inside the caller's writer transaction.""" + new_object_permission: Final = data_json.get("object_permission") + if new_object_permission is None: + return MappingProxyType({k: v for k, v in data_json.items() if k != "object_permission"}) + loaded: Final[object] = ( + json.loads(new_object_permission) if isinstance(new_object_permission, str) else new_object_permission + ) + grants: Final[Mapping[str, object]] = ( + cast( # cast-ok: object_permission payloads are grant-name to grant-list mappings + "Mapping[str, object]", loaded + ) + if isinstance(loaded, dict) + else MappingProxyType({}) + ) + upsert: Final = await prepare_object_permission_upsert( + new_object_permission=grants, + existing_object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) + row: Final = await table.upsert( + where={"object_permission_id": upsert.object_permission_id}, + data={"create": upsert.record, "update": upsert.record}, + ) + return MappingProxyType( + { + **{k: v for k, v in data_json.items() if k != "object_permission"}, + "object_permission_id": row.object_permission_id, + } + ) + + +async def _write_guarded_project_assignment( + table: "TableActions[_KeyRowDumpable]", + hashed_token: str, + existing_key_row: LiteLLM_VerificationToken, + data: Mapping[str, object], +) -> "_KeyRowDumpable | None": + validated_models: Final[Sequence[str]] = ( + cast( # cast-ok: LiteLLM_VerificationToken.models is a bare list + "Sequence[str] | None", existing_key_row.models + ) + or [] + ) + models_guard: Final[Mapping[str, object]] = ( + {"OR": [{"models": {"equals": list[str]()}}, {"models": {"equals": None}}]} + if not validated_models + else {"models": {"equals": validated_models}} + ) + updated_count: Final = await table.update_many( + where={ + "token": hashed_token, + "project_id": None, + "team_id": existing_key_row.team_id, + **models_guard, + "max_budget": existing_key_row.max_budget, + "object_permission_id": existing_key_row.object_permission_id, + }, + data=data, + ) + if updated_count == 0: + raise HTTPException( + status_code=400, + detail="Project assignment failed because the key was updated concurrently. Retry the request.", + ) + return await table.find_unique( + where={"token": hashed_token}, + include={"object_permission": True}, + ) + + +async def _update_key_row_assigning_project( + prisma_client: PrismaClient, + key: str, + non_default_values: Mapping[str, object], + existing_key_row: LiteLLM_VerificationToken, +) -> _KeyUpdateResult: + """Assign a project in one writer transaction, guarded on the key's validated state.""" + hashed_token: Final = _hash_token_if_needed(key) + async with prisma_client.tx() as tx_ctx: + tx: Final[_KeyUpdateTx] = cast( # cast-ok: the transaction object exposes the same table actions + "_KeyUpdateTx", tx_ctx + ) + update_values: Final[Mapping[str, object]] = await _apply_object_permission_update( + data_json=non_default_values, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + table=tx.litellm_objectpermissiontable, + ) + updated_row: Final = await _write_guarded_project_assignment( + table=tx.litellm_verificationtoken, + hashed_token=hashed_token, + existing_key_row=existing_key_row, + data=with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ), + ) + updated_data: Final[Mapping[str, object]] = ( + updated_row.model_dump() if updated_row is not None else MappingProxyType({}) + ) + result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data} + return result + + async def _update_key_row_with_soft_budget( prisma_client: PrismaClient, key: str, @@ -2777,25 +2944,48 @@ async def _update_key_row_with_soft_budget( non_default_values: Mapping[str, object], existing_key_row: LiteLLM_VerificationToken, changed_by: str, + expect_unassigned_project: bool = False, ) -> _KeyUpdateResult: hashed_token: Final = hash_token_if_needed(key) key_where: Final[_KeyRowWhere] = {"token": hashed_token} - tx: _KeyUpdateTx - async with prisma_client.tx() as tx: + async with prisma_client.tx() as tx_ctx: + tx: Final[_KeyUpdateTx] = cast( # cast-ok: the transaction object exposes the same table actions + "_KeyUpdateTx", tx_ctx + ) + update_input: Final[Mapping[str, object]] = ( + await _apply_object_permission_update( + data_json=non_default_values, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + table=tx.litellm_objectpermissiontable, + ) + if expect_unassigned_project + else non_default_values + ) update_values: Final = await _apply_soft_budget_update( data=data, - non_default_values=non_default_values, + non_default_values=update_input, db=tx, existing_key_row=existing_key_row, changed_by=changed_by, ) include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True} - updated_row: Final = await tx.litellm_verificationtoken.update( - where=key_where, - data=with_settings_updated_at( - prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) - ), - include=include_object_permission, + update_payload: Final = with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ) + updated_row: Final = ( + await _write_guarded_project_assignment( + table=tx.litellm_verificationtoken, + hashed_token=hashed_token, + existing_key_row=existing_key_row, + data=update_payload, + ) + if expect_unassigned_project + else await tx.litellm_verificationtoken.update( + where=key_where, + data=update_payload, + include=include_object_permission, + ) ) updated_data: Final[Mapping[str, object]] = ( updated_row.model_dump() if updated_row is not None else MappingProxyType({}) @@ -3562,10 +3752,6 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, ) - if data.project_id is not None and data.project_id != existing_key_row.project_id: - raise HTTPException( - status_code=400, detail="Project reassignment is not supported. Use null to detach the key." - ) is_project_change: Final = "project_id" in data.model_fields_set and data.project_id != existing_key_row.project_id acting_as_team_admin: Final = await _acting_as_team_admin_for_key_update( @@ -3697,6 +3883,12 @@ async def _validate_update_key_data( user_api_key_cache=user_api_key_cache, route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"), ) + await _validate_project_assignment( + data=data, + existing_key_row=existing_key_row, + prisma_client=checked_prisma_client, + user_api_key_cache=user_api_key_cache, + ) check_denied_passthrough_routes_caller_permission( data, user_api_key_dict, @@ -3846,7 +4038,7 @@ async def update_key_fn( - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key - agent_id: Optional[str] - The agent id associated with the key. - - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. + - project_id: Optional[str] - Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is rejected. - organization_id: Optional[str] - The organization id of the key. - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it. @@ -3994,23 +4186,31 @@ async def update_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") + is_project_assignment: Final = data.project_id is not None and existing_key_row.project_id is None + update_fields: Final[Mapping[str, object]] = cast( # cast-ok: prepare_key_update_data returns a bare dict + "Mapping[str, object]", non_default_values + ) object_permission_upsert: Final = await _prepare_key_update_object_permission( - object_permission_data=non_default_values.get("object_permission"), + object_permission_data=update_fields.get("object_permission"), existing_key_row=existing_key_row, prisma_client=prisma_client, ) key_update_values: Final = MappingProxyType( - {field: value for field, value in non_default_values.items() if field != "object_permission"} + {field: value for field, value in update_fields.items() if field != "object_permission"} ) await _check_key_project_team_on_mutation( data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, ) - update_values: Final = await _write_prepared_key_update_object_permission( - data_json=key_update_values, - upsert=object_permission_upsert, - prisma_client=prisma_client, + update_values: Final = ( + update_fields + if is_project_assignment + else await _write_prepared_key_update_object_permission( + data_json=key_update_values, + upsert=object_permission_upsert, + prisma_client=prisma_client, + ) ) changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name response: Final = ( @@ -4021,8 +4221,16 @@ async def update_key_fn( non_default_values=update_values, existing_key_row=existing_key_row, changed_by=changed_by, + expect_unassigned_project=is_project_assignment, ) if "soft_budget" in data.model_fields_set + else await _update_key_row_assigning_project( + prisma_client=prisma_client, + key=key, + non_default_values=update_fields, + existing_key_row=existing_key_row, + ) + if is_project_assignment else await _update_key_row( prisma_client=prisma_client, key=key, @@ -5483,7 +5691,7 @@ async def delete_verification_tokens( - List of keys being deleted, this contains information about the key_alias, token, and user_id being deleted, this is passed down to the KeyManagementEventHooks to delete the keys from the secret manager and handle audit logs """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj failed_tokens: list = [] try: @@ -5567,11 +5775,11 @@ async def delete_verification_tokens( verbose_proxy_logger.debug(traceback.format_exc()) raise e - for key in tokens: - user_api_key_cache.delete_cache(key) - # remove hash token from cache - hashed_token = hash_token(cast(str, key)) - user_api_key_cache.delete_cache(hashed_token) + await delete_cache_key_objects( + hashed_tokens=tuple(token for token in hashed_tokens if is_user_key_cache_key(token)), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) # After credential invalidation, so a failure here can never keep a deleted key alive. for deleted_key in authorized_keys: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 74fa8e3b3c1..2c7453ebb1c 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -486,6 +486,7 @@ async def _raise_on_invalid_credential_name( AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 _CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_MODEL_ROW_LOCK_SQL: Final = 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = $1 FOR UPDATE' _STORED_LITELLM_PARAMS_SQL: Final = ( "(CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)" ) @@ -633,6 +634,7 @@ async def _auto_router_capability_slot( effective_params: Mapping[str, object], model_id: str | None, member_write: MemberAutoRouterWrite | None = None, + lock_model_row: bool = False, ) -> AsyncGenerator[_ProxyModelTable, None]: """Hand out the model table to write through while the row's claim on a licensed capability is settled. @@ -649,6 +651,9 @@ async def _auto_router_capability_slot( A heuristic-v1 router whose tuning has moved off its recorded baseline is judged the same way under the same lock, against the DB rows plus this proxy's config.yaml routers. + + ``lock_model_row`` runs the write in a transaction holding the row's ``FOR UPDATE`` lock, + so a caller that re-reads the row through the yielded table writes onto its committed state. """ from litellm.proxy.proxy_server import ( _license_check, # pyright: ignore[reportPrivateUsage] # existing capability slot reads the proxy license singleton @@ -662,7 +667,8 @@ async def _auto_router_capability_slot( baselines: Final = heuristic_v1_tuning_baselines tuning_candidate: Final = _tuning_candidate(effective_params, model_id=model_id) judges_tuning: Final = baselines is not None and is_mutable_tuned_candidate(tuning_candidate, baselines) - if member_write is None and (limit is None or (not capabilities and not judges_tuning)): + checks_license: Final = member_write is not None or (limit is not None and (bool(capabilities) or judges_tuning)) + if not checks_license and not lock_model_row: yield _proxy_model_table(prisma_client) return transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db) @@ -673,14 +679,15 @@ async def _auto_router_capability_slot( ) async with transaction as tx_ctx: tables: Final[_TxModelTables] = tx_ctx - await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) + if checks_license: + await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) + locked_model_id: Final = ( + member_write.model_id if member_write is not None else (model_id if lock_model_row else None) + ) + if locked_model_id is not None: + await tx_ctx.query_raw(_MODEL_ROW_LOCK_SQL, locked_model_id) config_rows: Final = () if llm_router is None else tuple(llm_router.registry().deployments(source="config")) if member_write is not None: - if member_write.model_id is not None: - await tx_ctx.query_raw( - 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = $1 FOR UPDATE', - member_write.model_id, - ) pinned_client: Final = _TransactionClient(tx_ctx) team_where: Final[prisma_types.LiteLLM_TeamTableWhereUniqueInput] = {"team_id": member_write.team_id} team_include: Final[prisma_types.LiteLLM_TeamTableInclude] = {"litellm_model_table": True} @@ -752,7 +759,8 @@ async def _auto_router_capability_slot( prisma_client=pinned_client, llm_router=llm_router, ) - for capability in capabilities: + judged_capabilities: Final = capabilities if checks_license else () + for capability in judged_capabilities: if ( violation := await _claimed_capability_violation( tables, capability=capability, config_rows=config_rows, model_id=model_id, limit=limit @@ -761,7 +769,7 @@ async def _auto_router_capability_slot( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" ) - if judges_tuning and baselines is not None: + if checks_license and judges_tuning and baselines is not None: model_rows: Final = await ModelRepository(_TransactionClient(tx_ctx)).find_all_except(model_id or "") _raise_on_tuning_quota_violation( candidate=tuning_candidate, @@ -1049,6 +1057,16 @@ def _raise_if_ptu_cost_attribution_disabled(incoming_model_info: Mapping[str, ob ) +_ACCESS_GROUPS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None) + + +def _raise_if_access_groups_malformed(model_info_writes: Mapping[str, object]) -> None: + try: + _ACCESS_GROUPS.validate_python(model_info_writes.get("access_groups")) + except ValidationError: + raise HTTPException(status_code=400, detail="model_info.access_groups must be a list of strings") from None + + def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: """Enforce the PTU cross-field invariant on the effective model_info. @@ -1291,6 +1309,109 @@ def _loaded_catalog_entry( return loaded_catalog().get(catalog_key, MappingProxyType({})) +_LEGACY_UPDATE_IGNORED_MODEL_INFO_FIELDS: Final = frozenset( + { + "id", + "db_model", + "created_at", + "created_by", + "updated_at", + "updated_by", + "blocked", + "member_auto_router", + "team_id", + "team_public_model_name", + "direct_access", + "access_via_team_ids", + COST_MAP_LOOKUP_KEY, + *PTU_MODEL_INFO_FIELDS, + } +) + + +def legacy_update_model_info_changes( + db_model: Deployment, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object] = MappingProxyType({}), + loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, +) -> Mapping[str, object]: + """The ``model_info`` fields ``POST /model/update`` persists from its request. + + Only fields the caller sent with a value count, so ``ModelInfo`` defaults and the blank + strings typed clients send for unset fields never overwrite stored values. Server-owned, + display-only, team, PTU and pricing fields keep their dedicated write paths. As on + ``update_db_model``, only a body carrying the ``key`` that ``/model/info`` emits is an echo: + its values equal to the cost map entry of the base model the row has after the write, the + entry ``key`` names, or a limit discovered from the upstream stay display-only, and an echoed + value that differs from the stored override resets that override (``None``). + """ + stored: Final = db_model.model_info.model_dump(exclude_none=True) + requested: Final = incoming_model_info.model_dump(exclude_unset=True, exclude_none=True) + sent: Final = {k: v for k, v in requested.items() if k not in _LEGACY_UPDATE_IGNORED_MODEL_INFO_FIELDS and v != ""} + echoed: Final = frozenset( + echoed_cost_map_fields( + requested, + _cost_map_entry(db_model, {**stored, **sent}), + _loaded_catalog_entry(requested, loaded_catalog), + discovered_model_info, + ) + ) + kept: Final = without_server_derived_pricing(sent) + return MappingProxyType( + { + **{k: getattr(incoming_model_info, k) for k in kept if k not in echoed}, + **{k: None for k, v in kept.items() if k in echoed and stored.get(k) not in (None, v)}, + } + ) + + +def _legacy_update_model_info_writes( + current: Deployment, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object], + member_marker: bool | None, +) -> Mapping[str, object]: + changes: Final = { + **legacy_update_model_info_changes(current, incoming_model_info, discovered_model_info), + **({"member_auto_router": member_marker} if member_marker is not None else {}), + } + return MappingProxyType({k: v for k, v in changes.items() if getattr(current.model_info, k, None) != v}) + + +async def _locked_legacy_update_model_info_json( + table: _ProxyModelTable, + model_id: str, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object], + member_marker: bool | None, +) -> str | None: + row: Final = await table.find_unique(where={"model_id": model_id}) + if row is None: + return None + payload: Final = row.model_dump() + writes: Final = _legacy_update_model_info_writes( + Deployment(**payload), incoming_model_info, discovered_model_info, member_marker + ) + if not writes: + return None + stored_model_info: Final[object] = payload.get("model_info") + stored: Final = ( + TypeAdapter(dict[str, object]).validate_json(stored_model_info) + if isinstance(stored_model_info, str) + else TypeAdapter(dict[str, object]).validate_python(stored_model_info or {}) + ) + return ( + TypeAdapter(dict[str, object]) + .dump_json( + { + **{k: v for k, v in stored.items() if k not in writes}, + **{k: v for k, v in writes.items() if v is not None}, + } + ) + .decode() + ) + + def _param_keys_to_remove(patch: updateDeployment) -> tuple[str, ...]: return patch.litellm_params_to_remove if isinstance(patch, PatchModelRequest) else () @@ -1394,14 +1515,10 @@ def update_db_model( # convert to prisma compatible format - for key, value in merged_model_info.items(): - if isinstance(value, datetime.datetime): - merged_model_info[key] = value.isoformat() - prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( model_name=merged_model_name, litellm_params=json.dumps(merged_litellm_params), - model_info=json.dumps(merged_model_info), + model_info=TypeAdapter(dict[str, object]).dump_json(merged_model_info).decode(), ) if updated_patch.blocked is not None: @@ -3279,22 +3396,32 @@ async def update_model( renamed_update: Final[PrismaCompatibleUpdateDBModel] = ( {**base_update, "model_name": renamed_to} if renamed_to is not None else base_update ) - _data: Final[PrismaCompatibleUpdateDBModel] = ( - { - **renamed_update, - "model_info": deployment.model_info.model_copy( - update=MappingProxyType({"member_auto_router": member_marker}) - ).model_dump_json(exclude_none=True), - } - if member_marker is not None - else renamed_update + discovered_model_info: Final[Mapping[str, int]] = ( + llm_router.registry().discovered_model_info(_model_id) + if llm_router is not None + else MappingProxyType({}) ) + requested_model_info: Final = _legacy_update_model_info_writes( + deployment, _model_info, discovered_model_info, member_marker + ) + _raise_if_access_groups_malformed(requested_model_info) async with _auto_router_capability_slot( prisma_client, effective_params=effective_params, model_id=_model_id, member_write=member_write, + lock_model_row=bool(requested_model_info), ) as table: + model_info_json: Final = ( + await _locked_legacy_update_model_info_json( + table, _model_id, _model_info, discovered_model_info, member_marker + ) + if requested_model_info + else None + ) + _data: Final[PrismaCompatibleUpdateDBModel] = ( + {**renamed_update, "model_info": model_info_json} if model_info_json is not None else renamed_update + ) model_response: Final = await table.update( where={"model_id": _model_id}, data=_data, diff --git a/litellm/proxy/management_endpoints/ptu_consumption.py b/litellm/proxy/management_endpoints/ptu_consumption.py index b09c07a0443..86a35fbbb0c 100644 --- a/litellm/proxy/management_endpoints/ptu_consumption.py +++ b/litellm/proxy/management_endpoints/ptu_consumption.py @@ -12,14 +12,14 @@ from types import MappingProxyType from typing import Final from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled -from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours -from litellm.router import Router -from litellm.router_utils.common_utils import resolve_model_group_alias -from litellm.router_utils.ptu_shares import ( +from litellm.litellm_core_utils.ptu_shares import ( model_group_ptu_capacity, routed_deployments, team_servable_deployments, ) +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias +from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours +from litellm.router import Router from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendData, MetricWithMetadata, diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 221a6221fab..d4fd7a31d13 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -93,5 +93,6 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: class CustomOpenID(OpenID): team_ids: list[str] + team_id_to_upsert: str | None = None user_role: LitellmUserRoles | None = None extra_fields: dict[str, object] | None = None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 59c2df91d77..53079d84fdd 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -87,7 +87,7 @@ from litellm.proxy._types import ( TeamMemberAddRequest, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_team_object, get_user_object from litellm.proxy.auth.auth_utils import ( # noqa: F401 # legacy module exports _get_request_ip_address, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export get_request_ip_address, @@ -1266,6 +1266,7 @@ def generic_response_convertor( sso_jwt_handler: JWTHandler | None = None, role_mappings: Optional["RoleMappings"] = None, team_mappings: Optional["TeamMappings"] = None, + jwt_auth_enabled: bool = False, ) -> CustomOpenID: generic_user_id_attribute_name: Final = os.getenv("GENERIC_USER_ID_ATTRIBUTE", "preferred_username") generic_user_display_name_attribute_name: Final = os.getenv("GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub") @@ -1286,14 +1287,15 @@ def generic_response_convertor( generic_user_email_attribute_name, ) - all_teams: Final = [] + token: Final = cast(dict[str, object], response) # cast-ok: the IdP userinfo JSON object fastapi-sso hands back + all_teams: Final = list[str]() if sso_jwt_handler is not None: - team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response)) + team_ids = sso_jwt_handler.get_all_jwt_team_ids(token) all_teams.extend(team_ids) if team_mappings is not None and team_mappings.team_ids_jwt_field is not None: team_ids_from_db_mapping: Final[list[str] | None] = get_nested_value( - data=cast(dict, response), + data=token, key_path=team_mappings.team_ids_jwt_field, default=[], ) @@ -1305,8 +1307,17 @@ def generic_response_convertor( team_ids_from_db_mapping, ) else: - team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response)) + team_ids = jwt_handler.get_all_jwt_team_ids(token) all_teams.extend(team_ids) + team_id_jwt_field: Final = jwt_handler.litellm_jwtauth.team_id_jwt_field + upsert_candidate: Final[object] = ( + get_nested_value(data=token, key_path=team_id_jwt_field) + if jwt_auth_enabled and jwt_handler.litellm_jwtauth.team_id_upsert and team_id_jwt_field is not None + else None + ) + team_id_to_upsert: Final = ( + upsert_candidate if isinstance(upsert_candidate, str) and upsert_candidate.strip() else None + ) # Determine user role based on role_mappings if available # Only apply role_mappings for GENERIC SSO provider @@ -1375,6 +1386,7 @@ def generic_response_convertor( last_name=get_nested_value(response, generic_user_last_name_attribute_name), provider=get_nested_value(response, generic_provider_attribute_name), team_ids=all_teams, + team_id_to_upsert=team_id_to_upsert, user_role=user_role, extra_fields=extra_fields, ) @@ -1614,6 +1626,8 @@ async def get_generic_sso_response( # make generic sso provider from fastapi_sso.sso.base import DiscoveryDocument + from litellm.proxy import proxy_server + received_response: dict | None = None sso_assertion: SSOIdentityAssertion | None = None @@ -1665,6 +1679,7 @@ async def get_generic_sso_response( sso_jwt_handler=sso_jwt_handler, role_mappings=role_mappings, team_mappings=team_mappings, + jwt_auth_enabled=proxy_server.general_settings.get("enable_jwt_auth", False) is True, ) SSOProvider: Final = create_generic_sso_provider( @@ -1884,6 +1899,29 @@ async def find_existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | No return frozenset(row.team_id for row in rows) +async def create_sso_team(team_id: str) -> None: + """Create the ``team_id_jwt_field`` team on UI SSO login the way the API JWT path does under ``team_id_upsert``.""" + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + return + await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + team_id_upsert=True, + ) + + +async def _create_sso_team_or_warn(team_id: str, create_team: Callable[[str], Awaitable[object]]) -> None: + """A concurrent login may create the team first, so a failed create is not fatal: the join still runs.""" + try: + await create_team(team_id) + except Exception as e: # noqa: BLE001 # fail-safe: a failed team create must not fail the login or skip the join + verbose_proxy_logger.warning("SSO login could not create team %s, trying to join it anyway: %s", team_id, e) + + async def add_missing_team_member( user_info: NewUserResponse | LiteLLM_UserTable, sso_teams: Sequence[str], @@ -1892,21 +1930,30 @@ async def add_missing_team_member( [str, NewUserResponse | LiteLLM_UserTable], Awaitable[object] ] = create_team_member_add_task, existing_team_ids: Callable[[Sequence[str]], Awaitable[frozenset[str] | None]] = find_existing_team_ids, + team_id_to_upsert: str | None = None, + create_team: Callable[[str], Awaitable[object]] = create_sso_team, ) -> None: """ - Get missing teams (diff b/w user_info.team_ids and sso_teams) - Drop the IdP groups that are not teams with one lookup, since group claims often carry many of them + - Join the ``team_id_to_upsert`` team too, creating it when it is missing (``team_id_upsert``), the way the API JWT path does; + it stays out of ``result.team_ids`` so restricted_sso_group and the access-token group fallback see only the claimed groups - Add missing user to missing teams, one at a time in sso_teams order, so user.teams keeps the claim's order """ held_teams: Final = frozenset(user_info.teams or ()) - missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(sso_teams) if team_id not in held_teams) + claimed_teams: Final = sso_teams if team_id_to_upsert is None else (*sso_teams, team_id_to_upsert) + missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(claimed_teams) if team_id not in held_teams) if not missing_teams: return existing: Final = await existing_team_ids(missing_teams) teams_to_add: Final = ( - missing_teams if existing is None else tuple(team_id for team_id in missing_teams if team_id in existing) + missing_teams + if existing is None + else tuple(team_id for team_id in missing_teams if team_id in existing or team_id == team_id_to_upsert) ) for team_id in teams_to_add: + if team_id == team_id_to_upsert and (existing is None or team_id not in existing): + await _create_sso_team_or_warn(team_id, create_team) await add_team_member(team_id, user_info) @@ -3663,7 +3710,11 @@ class SSOAuthenticationHandler: verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition") return sso_teams: Final = result.team_ids if isinstance(result, CustomOpenID) else [] - await add_missing_team_member(user_info=user_info, sso_teams=sso_teams) + await add_missing_team_member( + user_info=user_info, + sso_teams=sso_teams, + team_id_to_upsert=result.team_id_to_upsert if isinstance(result, CustomOpenID) else None, + ) @staticmethod def verify_user_in_restricted_sso_group( diff --git a/litellm/proxy/management_helpers/model_quota_policy.py b/litellm/proxy/management_helpers/model_quota_policy.py index 7e66bfaf0ad..7cf2d7677e7 100644 --- a/litellm/proxy/management_helpers/model_quota_policy.py +++ b/litellm/proxy/management_helpers/model_quota_policy.py @@ -9,10 +9,10 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator import litellm from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider, is_openai_compatible_endpoint +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils +from litellm.litellm_core_utils.routing_common_utils import provider_for_generic_call from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.llms.openai_like.json_loader import JSONProviderRegistry -from litellm.router_utils.common_utils import provider_for_generic_call -from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils from litellm.secret_managers.main import get_secret_str if TYPE_CHECKING: diff --git a/litellm/proxy/management_helpers/ptu_share_guard.py b/litellm/proxy/management_helpers/ptu_share_guard.py index 0c8b21d02c6..529ad60c353 100644 --- a/litellm/proxy/management_helpers/ptu_share_guard.py +++ b/litellm/proxy/management_helpers/ptu_share_guard.py @@ -5,8 +5,8 @@ from typing import Final from fastapi import HTTPException +from litellm.litellm_core_utils.ptu_shares import deployments_shared_with from litellm.router import Router -from litellm.router_utils.ptu_shares import deployments_shared_with def raise_if_teams_hold_ptu_shares(team_ids: Sequence[str], llm_router: Router | None) -> None: diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index db640b6d689..be24b1630b0 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -438,6 +438,14 @@ def is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bo _is_mcp_admitted_user_subject: Final = is_mcp_admitted_user_subject +def has_team_unbound_mcp_admin_view(user_api_key_auth: UserAPIKeyAuth | None) -> bool: + return ( + user_api_key_auth is not None + and user_api_key_has_admin_view(user_api_key_auth) + and (not user_api_key_auth.team_id or user_api_key_auth.team_id == UI_TEAM_ID) + ) + + def _gateway_dcr_challenge_target( route: str, mcp_servers: list[str] | None, @@ -4290,19 +4298,19 @@ class MCPRequestHandler: """Whether this principal's admin-view role grants the unscoped MCP resolution, whatever credential carries it (admin key, dashboard session, or OAuth-admitted session subject). - Two bounds disqualify, one per ownership of the row. A CREDENTIAL's explicit - ``object_permission.mcp_servers`` scope wins even for admins, including the empty list. An - admitted subject's object_permission is the user's own row, whose ``mcp_servers`` column is - [] by DB default, so for that shape the row binds through the entitlement ceiling instead - (any non-empty entitlement, or an unresolved one, disqualifies), exactly as - ``operator_open_server_ids`` reads the same row. The one owner of this predicate: the - server-axis registry resolution in ``get_allowed_mcp_servers`` and the tools-axis open - channel in ``resolve_admitted_subject_tools`` both consult it, so the two axes cannot - disagree.""" + A credential on a real team never qualifies, a virtual key or a ``lite login`` session alike: its team + bounds its MCP access whoever owns it, as it bounds the credential's model access. Two more bounds disqualify, one per ownership of the row. A CREDENTIAL's + explicit ``object_permission.mcp_servers`` scope wins even for admins, including the empty list. An + admitted subject's object_permission is the user's own row, whose ``mcp_servers`` column is [] by DB + default, so for that shape the row binds through the entitlement ceiling instead (any non-empty + entitlement, or an unresolved one, disqualifies), exactly as ``operator_open_server_ids`` reads the same + row. The one owner of this predicate: the server-axis registry resolution in ``get_allowed_mcp_servers`` + and the tools-axis open channel in ``resolve_admitted_subject_tools`` both consult it, so the two axes + cannot disagree.""" if ( user_api_key_auth is None or user_api_key_auth.mcp_explicit_grants_only - or not user_api_key_has_admin_view(user_api_key_auth) + or not has_team_unbound_mcp_admin_view(user_api_key_auth) ): return False object_permission: Final = user_api_key_auth.object_permission diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index c772e2e9d09..eacba25257e 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -3705,7 +3705,7 @@ class MCPServerManager: Priority: 1. If object_permission.mcp_servers is explicitly set, use it (even for admins) - 2. If admin and no object_permission, return all servers + 2. If admin, not bound to a real team, and no object_permission, return all servers 3. Otherwise, use standard permission checks """ if managed_agent_policy(user_api_key_auth) is not None: diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index d486a69ef1a..66e46c73f98 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -17,6 +17,7 @@ import anyio from fastapi import HTTPException from mcp import MCPError, ReadResourceResult, Resource from mcp.types import ( + METHOD_NOT_FOUND, CallToolRequest, CallToolRequestParams, CallToolResult, @@ -222,6 +223,7 @@ __all__ = ( "mcp_get_prompt", "mcp_read_resource", "raise_denied_scoped_mcp_access", + "raise_generic_mcp_access_denial", ) @@ -762,8 +764,14 @@ async def raise_denied_scoped_mcp_access( ) } raise HTTPException(status_code=403, detail=group_denial) + raise_generic_mcp_access_denial(requested_names) + + +def raise_generic_mcp_access_denial(requested_names: Sequence[str]) -> NoReturn: generic_denial: Final[_McpDeniedDetail] = { - "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}" + "error": ( + f"The key is not allowed to access the requested MCP servers: {', '.join(dict.fromkeys(requested_names))}" + ) } raise HTTPException(status_code=403, detail=generic_denial) @@ -2806,6 +2814,58 @@ async def _fire_mcp_tool_call_logging( return result +async def log_mcp_sdk_rejection( + *, + method: str, + request_id: str | int | None, + error_code: int, + start_time: datetime, + user_api_key_auth: UserAPIKeyAuth, +) -> None: + from litellm.proxy.proxy_server import proxy_logging_obj + from litellm.utils import Rules, function_setup + + message: Final = "Method not found" if error_code == METHOD_NOT_FOUND else "Invalid request parameters" + rejection: Final = { + "method": "".join(char for char in method[:256] if char.isprintable()), + "request_id": ( + "".join(char for char in str(request_id)[:128] if char.isprintable()) if request_id is not None else None + ), + "error_code": error_code, + "error_message": message, + "stage": "sdk_dispatch", + "route": user_api_key_auth.request_route or "/mcp", + } + metadata: Final = { + **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_auth), + "spend_logs_metadata": {"mcp_sdk_rejection": rejection}, + } + logging_obj, request_data = function_setup( + original_function=CallTypes.call_mcp_tool.value, + rules_obj=Rules(), + start_time=start_time, + model="MCP: request_rejected", + litellm_call_id=str(uuid.uuid4()), + metadata=metadata, + ) + logging_obj.update_environment_variables( # pyright: ignore[reportUnknownMemberType] # untyped legacy signature + litellm_params={"metadata": metadata}, + optional_params={}, + user=user_api_key_auth.user_id, + ) + safe_error: Final = HTTPException(status_code=400, detail=message) + end_time: Final = datetime.now() # noqa: DTZ005 # matches the failure logging start time + await asyncio.to_thread(logging_obj.failure_handler, safe_error, "", start_time, end_time) + await logging_obj.async_failure_handler(safe_error, "", start_time, end_time) + await proxy_logging_obj.post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # legacy signature + request_data={**request_data, "litellm_logging_obj": logging_obj}, + original_exception=safe_error, + user_api_key_dict=user_api_key_auth, + route="/mcp", + traceback_str="", + ) + + async def fire_mcp_tool_call_failure_logging( logging_obj: LiteLLMLoggingObj | None, exception: Exception, diff --git a/litellm/proxy/mcp_server/semantic_tool_filter.py b/litellm/proxy/mcp_server/semantic_tool_filter.py index 4fe723d6b24..8666c05079d 100644 --- a/litellm/proxy/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/mcp_server/semantic_tool_filter.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers from litellm.proxy.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy.mcp_server.faults import iter_exception_tree from litellm.proxy.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR +from litellm.router_api import RoutingContext if TYPE_CHECKING: from semantic_router.routers import SemanticRouter @@ -172,7 +173,7 @@ class SemanticMCPToolFilter: self.tool_router = SemanticRouter( routes=routes, encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.router_instance, + ctx=RoutingContext(self.router_instance), model_name=self.embedding_model, score_threshold=self.similarity_threshold, ), @@ -247,7 +248,7 @@ class SemanticMCPToolFilter: router: Final = SemanticRouter( routes=[], encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.router_instance, + ctx=RoutingContext(self.router_instance), model_name=self.embedding_model, score_threshold=self.similarity_threshold, ), diff --git a/litellm/proxy/mcp_server/server.py b/litellm/proxy/mcp_server/server.py index 2105d875b58..aeb640a7c8b 100644 --- a/litellm/proxy/mcp_server/server.py +++ b/litellm/proxy/mcp_server/server.py @@ -14,6 +14,7 @@ import time import types from collections import Counter from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence, Set +from datetime import datetime from typing import TYPE_CHECKING, Final, Literal, NoReturn, Protocol import httpx @@ -519,6 +520,7 @@ if MCP_AVAILABLE: "mcp_get_prompt", "mcp_read_resource", "raise_denied_scoped_mcp_access", + "raise_generic_mcp_access_denial", "redact_mcp_resource_url", ) from mcp.server import Server @@ -530,13 +532,15 @@ if MCP_AVAILABLE: ) from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from mcp.server.auth.provider import AccessToken - from mcp.server.context import ServerRequestContext + from mcp.server.context import CallNext, HandlerResult, ServerRequestContext from mcp.server.lowlevel.server import NotificationOptions from mcp.server.lowlevel.server import Server as MCPProtocolServer from mcp.server.models import InitializationOptions from mcp.shared.exceptions import MCPError from mcp.shared.inbound import InboundLadderRejection, classify_inbound_request, find_duplicated_routing_header from mcp.types import ( + INVALID_PARAMS, + METHOD_NOT_FOUND, CallToolRequest, GetPromptRequest, JSONRPCRequest, @@ -546,6 +550,7 @@ if MCP_AVAILABLE: ListToolsRequest, ReadResourceRequest, ) + from mcp_types.methods import CLIENT_REQUESTS, SPEC_CLIENT_METHODS from litellm.proxy.mcp_server import operations from litellm.proxy.mcp_server.contracts import OperationContext @@ -643,7 +648,56 @@ if MCP_AVAILABLE: name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, ) - server.middleware.extend((finish_scope_response, GatewayVersionPolicy())) + + def _sdk_rejection_code(ctx: ServerRequestContext, error: MCPError | ValidationError) -> int | None: + if ctx.request_id is None: + return None + entry: Final = server.get_request_handler(ctx.method) + if entry is None: + if ctx.method == "initialize": + return None + if isinstance(error, ValidationError): + return INVALID_PARAMS + if error.code == METHOD_NOT_FOUND: + return METHOD_NOT_FOUND + return None + if not isinstance(error, ValidationError): + unavailable_for_version: Final = ( + ctx.method in SPEC_CLIENT_METHODS and (ctx.method, ctx.protocol_version) not in CLIENT_REQUESTS + ) + return METHOD_NOT_FOUND if error.code == METHOD_NOT_FOUND and unavailable_for_version else None + try: + entry.params_type.model_validate({} if ctx.params is None else ctx.params, by_name=False) + except ValidationError: + return INVALID_PARAMS + return None + + async def _log_sdk_rejection(ctx: ServerRequestContext, call_next: CallNext) -> HandlerResult: + start_time: Final = datetime.now() # noqa: DTZ005 # failure logging uses naive timestamps + try: + return await call_next(ctx) + except (MCPError, ValidationError) as error: + try: + code: Final = _sdk_rejection_code(ctx, error) + if code is not None: + token: Final = active_mcp_request_ctx_var.set(ctx) + try: + auth, *_ = await get_or_extract_auth_context() + finally: + active_mcp_request_ctx_var.reset(token) + if auth is not None: + await operations.log_mcp_sdk_rejection( + method=ctx.method, + request_id=ctx.request_id, + error_code=code, + start_time=start_time, + user_api_key_auth=auth, + ) + except Exception: # noqa: BLE001 # logging must not replace the SDK response or expose rejected input + verbose_logger.warning("Failed to log MCP SDK rejection") + raise + + server.middleware.extend((finish_scope_response, GatewayVersionPolicy(), _log_sdk_rejection)) server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server) sse: Final[SseServerTransport] = SseServerTransport("/sse/messages") @@ -1144,6 +1198,7 @@ if MCP_AVAILABLE: apply_display_name_overrides, filter_tools_by_allowed_tools, raise_denied_scoped_mcp_access, + raise_generic_mcp_access_denial, ) @contextlib.asynccontextmanager @@ -1771,10 +1826,10 @@ if MCP_AVAILABLE: row capping it. Raises HTTPException(403) unless the principal holds toolset_id through one - of its grant sources. Admins always pass. + of its grant sources. Admins pass unless the key is bound to a real team. """ from litellm.proxy._types import LiteLLM_ObjectPermissionTable - from litellm.proxy.management_endpoints.common_utils import user_api_key_has_admin_view + from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import has_team_unbound_mcp_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 @@ -1788,7 +1843,7 @@ if MCP_AVAILABLE: ) acting: Final = await acting_user(user_api_key_auth) - is_admin: Final = user_api_key_has_admin_view(acting) + is_admin: Final = has_team_unbound_mcp_admin_view(acting) if not is_admin and toolset_id not in await granted(acting): raise HTTPException( status_code=403, diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 907c2ca303e..c0d81abeba7 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1,4 +1,5 @@ import base64 +import itertools import logging import mimetypes import re @@ -541,7 +542,8 @@ def get_team_provider_credentials( listings on the team's own provider account/key instead of a shared global one. 2. Fallback: any deployment the caller is granted access to for this - provider, expanding wildcard routes and the all-proxy-models sentinel. + provider, matching wildcard grants against deployed model names without + listing any provider's catalog, and expanding the all-proxy-models sentinel. Skipped when ``include_accessible_deployments`` is False. Credential lookup is scoped to both the team's allowlist and the key's own @@ -559,6 +561,7 @@ def get_team_provider_credentials( return None from litellm.proxy._types import SpecialModelNames + from litellm.proxy.auth.auth_checks import is_model_allowed_by_pattern from litellm.proxy.auth.model_checks import get_complete_model_list, get_key_models team_id: Final = user_api_key_dict.team_id @@ -645,28 +648,36 @@ def get_team_provider_credentials( grants_all_models: Final = SpecialModelNames.all_proxy_models.value in team_models effective_team_models: Final = [] if grants_all_models else team_models - models_to_try: Final = list( - dict.fromkeys( - get_complete_model_list( - key_models=list(key_model_allowlist), - team_models=effective_team_models, - proxy_model_list=proxy_model_list, - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=True, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=True, - team_id=team_id, - ) + granted_models: Final = get_complete_model_list( + key_models=list(key_model_allowlist), + team_models=effective_team_models, + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=True, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=True, + team_id=team_id, + expand_wildcard_routes=False, + ) + deployments_granted_by_wildcard: Final = ( + model_name + for model_name in proxy_model_list + if any( + is_model_allowed_by_pattern(model=model_name, allowed_model_pattern=granted) for granted in granted_models ) ) - for model_name in models_to_try: - credentials = _provider_credentials(model_name) - if credentials is not None: - return credentials - - return None + return next( + ( + credentials + for credentials in map( + _provider_credentials, dict.fromkeys(itertools.chain(granted_models, deployments_granted_by_wildcard)) + ) + if credentials is not None + ), + None, + ) HTTP_HEADER_VALUE: Final = re.compile(r"[\x21-\x7e]+(?:[ \t]+[\x21-\x7e]+)*") diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index c1d4534ffa2..e613d744fb0 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1786,7 +1786,7 @@ def _routed_model_name_on_provider( def _stored_model_name_routes_to_provider( llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str ) -> bool: - routing_group: Final = llm_router.get_routing_group(model_name) + routing_group: Final = llm_router.registry().routing_group(model_name) routed_model_names: Final = ( tuple(routing_group.models) if routing_group is not None else (llm_router.registry().resolve(model_name),) ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 21078dfc848..33b18c3d939 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -32,6 +32,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, AWS_ERROR_TYPE_HEADER, @@ -72,6 +73,7 @@ from litellm.llms.openai.common_utils import is_openai_agents_api_path, openai_a from litellm.llms.openai.workload_identity import get_workload_identity_bearer_token_for_api_base from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, validate_oss_request from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.llms.vllm.common_utils import endpoint_under_api_base from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import enforced_model_allowlists @@ -148,7 +150,7 @@ from litellm.types.passthrough_endpoints.tinyfish import ( ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials from litellm.types.router import LiteLLMParamsTypedDict -from litellm.types.utils import LlmProviders +from litellm.types.utils import CallTypes, LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager @@ -600,7 +602,7 @@ async def vllm_proxy_route( ) return await llm_passthrough_factory_proxy_route( - endpoint=endpoint, + endpoint=endpoint_under_api_base(endpoint, get_secret_str("VLLM_API_BASE")), request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -4104,7 +4106,12 @@ async def vertex_ai_live_websocket_passthrough( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data={}, + request_data={ + "model": model or "", + "custom_llm_provider": litellm.LlmProviders.VERTEX_AI.value, + "call_type": CallTypes.pass_through.value, + "litellm_call_id": str(uuid.uuid4()), + }, ) if websocket.client_state != WebSocketState.DISCONNECTED: await websocket.close(code=1011, reason=VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON) @@ -4139,6 +4146,8 @@ async def vertex_ai_live_websocket_passthrough( forward_headers=False, endpoint="/vertex_ai/live", accept_websocket=False, + builtin_provider_route=True, + model=model, setup_model_rewriter=_build_vertex_live_setup_model_rewriter( vertex_project=resolved_project, vertex_location=resolved_location, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index 7a5d5a1e259..d449f62af2c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -1,4 +1,3 @@ -import re from collections.abc import Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final @@ -9,6 +8,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.gemini.videos.transformation import GeminiVideoConfig +from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( GeminiSseChunkParser, @@ -207,11 +207,7 @@ class GeminiPassthroughLoggingHandler: @staticmethod def extract_model_from_url(url: str) -> str: - pattern: Final = r"/models/([^:]+)" - match: Final = re.search(pattern, url) - if match: - return match.group(1) - return "unknown" + return get_vertex_model_id_from_url(url) or "unknown" @staticmethod def _create_gemini_response_logging_payload_for_generate_content( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index b3cbad91624..d5c162296d3 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -3,7 +3,7 @@ import re from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast -from urllib.parse import urlparse +from urllib.parse import parse_qs, urlparse import httpx from pydantic import ConfigDict, TypeAdapter @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import from litellm.llms.vertex_ai.common_utils import ( get_vertex_ai_lyria_generation_cost, get_vertex_location_from_url, + get_vertex_model_id_from_url, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as VertexModelResponseIterator, @@ -47,6 +48,19 @@ from litellm.types.utils import ( ) vertex_search_api_config: Final = VertexSearchAPIVectorStoreConfig() +GEMINI_API_HOST_SUFFIX: Final = "generativelanguage.googleapis.com" +GOOGLE_AI_HOST_SUFFIXES: Final = ( + GEMINI_API_HOST_SUFFIX, + "aiplatform.googleapis.com", + "discoveryengine.googleapis.com", +) +_VERTEX_MULTI_REGION_HOST: Final = re.compile(r"aiplatform\.[a-z0-9-]+\.rep\.googleapis\.com") + + +def _is_host_or_subdomain(hostname: str, suffix: str) -> bool: + return hostname == suffix or hostname.endswith((f".{suffix}", f"-{suffix}")) + + if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch @@ -684,11 +698,18 @@ class VertexPassthroughLoggingHandler: @staticmethod def extract_model_from_url(url: str) -> str: - pattern: Final = r"/models/([^:]+)" - match: Final = re.search(pattern, url) - if match: - return match.group(1) - return "unknown" + return get_vertex_model_id_from_url(url) or "unknown" + + @staticmethod + def model_from_url_route(url_route: str) -> str | None: + parsed_url: Final = urlparse(url_route) + model: Final = get_vertex_model_id_from_url(parsed_url.path) + if model is None: + return "vertex_ai/search_api" if ":search" in parsed_url.path else None + streamed: Final = "sse" in parse_qs(parsed_url.query).get("alt", ()) + if parsed_url.path.endswith((":rawPredict", ":streamRawPredict")) and not streamed: + return f"vertex_ai/{model}" + return model @staticmethod def extract_model_name_from_vertex_path(vertex_model_path: str) -> str: @@ -734,18 +755,32 @@ class VertexPassthroughLoggingHandler: return None @staticmethod - def _get_custom_llm_provider(url: str, logging_obj: LiteLLMLoggingObj) -> str: + def is_google_ai_url(url: str) -> bool: hostname: Final = urlparse(url).hostname - if logging_obj.model_call_details.get("custom_llm_provider") == litellm.LlmProviders.GEMINI.value or ( - hostname is not None - and ( - hostname == "generativelanguage.googleapis.com" - or hostname.endswith(".generativelanguage.googleapis.com") - ) - ): + return hostname is not None and ( + any(_is_host_or_subdomain(hostname, suffix) for suffix in GOOGLE_AI_HOST_SUFFIXES) + or _VERTEX_MULTI_REGION_HOST.fullmatch(hostname) is not None + ) + + @staticmethod + def is_gemini_api_host(url: str) -> bool: + hostname: Final = urlparse(url).hostname + return hostname is not None and ( + hostname == GEMINI_API_HOST_SUFFIX or hostname.endswith(f".{GEMINI_API_HOST_SUFFIX}") + ) + + @staticmethod + def custom_llm_provider_from_url(url: str) -> str: + if VertexPassthroughLoggingHandler.is_gemini_api_host(url): return litellm.LlmProviders.GEMINI.value return litellm.LlmProviders.VERTEX_AI.value + @staticmethod + def _get_custom_llm_provider(url: str, logging_obj: LiteLLMLoggingObj) -> str: + if logging_obj.model_call_details.get("custom_llm_provider") == litellm.LlmProviders.GEMINI.value: + return litellm.LlmProviders.GEMINI.value + return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url) + @staticmethod def _is_multimodal_embedding_response(json_response: dict) -> bool: """ diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index bda2bfb9e04..647e145aac1 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3,6 +3,7 @@ import asyncio import copy import json import posixpath +import re import traceback from base64 import b64encode from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence @@ -77,6 +78,7 @@ from litellm.llms.base_llm.managed_resources.utils import ( from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.eu_jev.common_utils import eu_jev_scoped_model, eu_jev_upstream_body from litellm.llms.oss_decision import validate_oss_request +from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url from litellm.passthrough import BasePassthroughUtils from litellm.proxy._lazy_features import lazy_owned_routes from litellm.proxy._types import ( @@ -146,12 +148,15 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, EndpointType, PassthroughStandardLoggingPayload, + WebSocketSessionMeter, + WebSocketSessionUsage, ) -from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage +from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( is_tinyfish_agent_url, ) +from .llm_provider_handlers.vertex_passthrough_logging_handler import VertexPassthroughLoggingHandler from .streaming_handler import PassThroughStreamingHandler from .success_handler import PassThroughEndpointLogging from .upstream_usage_headers import ( @@ -874,13 +879,45 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non metadata.setdefault("standard_logging_guardrail_information", list(entries)) +_GOOGLE_PASSTHROUGH_PROVIDERS: Final = frozenset( + (litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value) +) + + +def _passthrough_google_provider(url_route: str, custom_llm_provider: str | None) -> str | None: + if custom_llm_provider is not None: + return custom_llm_provider if custom_llm_provider in _GOOGLE_PASSTHROUGH_PROVIDERS else None + if VertexPassthroughLoggingHandler.is_google_ai_url(url_route): + return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route) + return None + + +_OPERATION_POLL_SUFFIX: Final = re.compile(r"/operations/[^/]+$") + + +def _without_operation_poll_suffix(model: str | None) -> str | None: + return _OPERATION_POLL_SUFFIX.sub("", model) if model else model + + +def _passthrough_google_model(url_route: str, google_provider: str | None) -> str | None: + match google_provider: + case "gemini": + return _without_operation_poll_suffix(get_vertex_model_id_from_url(urlparse(url_route).path)) + case "vertex_ai": + return _without_operation_poll_suffix(VertexPassthroughLoggingHandler.model_from_url_route(url_route)) + case _: + return None + + def _build_passthrough_failure_request_payload( parsed_body: dict | None, kwargs: dict | None, logging_obj: LiteLLMLoggingObj | None, custom_llm_provider: str | None, upstream_usage: UpstreamReportedUsage | None = None, + url: httpx.URL | None = None, upstream_model: str | None = None, + builtin_provider_route: bool = False, ) -> dict: """Build the ``request_data`` dict passed to ``post_call_failure_hook``. @@ -893,16 +930,25 @@ def _build_passthrough_failure_request_payload( comes paired with a usage object, so both keys are written together. ``upstream_model`` is the model name the upstream reported, which replaces the body's ``model`` on the failure row just as it does on a success row. + ``builtin_provider_route`` lets a Gemini or Vertex URL name the failure's + model and provider. A user-defined route never gets that, since its URL + names an upstream model rather than a LiteLLM-managed one. """ request_payload: Final[dict] = dict(cast(Mapping[str, object], parsed_body or {})) # cast-ok: json body if kwargs: request_payload.update(kwargs) if logging_obj is not None: request_payload["litellm_logging_obj"] = logging_obj - if "model" not in request_payload and parsed_body and isinstance(parsed_body, dict): - request_payload["model"] = parsed_body.get("model", "") - if "custom_llm_provider" not in request_payload and custom_llm_provider: - request_payload["custom_llm_provider"] = custom_llm_provider + url_route: Final = str(url) if url is not None else "" + google_provider: Final = ( + _passthrough_google_provider(url_route, custom_llm_provider) if builtin_provider_route else None + ) + if "model" not in request_payload: + request_payload["model"] = _passthrough_google_model(url_route, google_provider) or "" + resolved_provider: Final = custom_llm_provider or google_provider + request_payload.pop("custom_llm_provider", None) + if resolved_provider: + request_payload["custom_llm_provider"] = resolved_provider if upstream_usage is not None: request_payload["response_cost"] = upstream_usage.response_cost or 0.0 request_payload["combined_usage_object"] = Usage(total_tokens=upstream_usage.total_tokens or 0) @@ -1203,6 +1249,7 @@ async def pass_through_request( # Initialize variables ######################################################### litellm_call_id: Final = str(uuid.uuid4()) + builtin_provider_route: Final = not request_dispatched_to_pass_through_endpoint(request) url: httpx.URL | None = None # parsed request body @@ -1612,6 +1659,8 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, + builtin_provider_route=builtin_provider_route, upstream_model=upstream_model, ), logging_obj=logging_obj, @@ -1651,6 +1700,8 @@ async def pass_through_request( kwargs=kwargs, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + url=url, + builtin_provider_route=builtin_provider_route, ), ), managed_id_provider=_managed_id_provider, @@ -1710,6 +1761,8 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, + builtin_provider_route=builtin_provider_route, upstream_model=upstream_model, ), logging_obj=logging_obj, @@ -1749,6 +1802,8 @@ async def pass_through_request( kwargs=kwargs, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + url=url, + builtin_provider_route=builtin_provider_route, ), ), managed_id_provider=_managed_id_provider, @@ -1811,6 +1866,8 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, + builtin_provider_route=builtin_provider_route, upstream_model=upstream_model, ) failure_request_payload["response_body"] = response_body @@ -1991,7 +2048,14 @@ async def pass_through_request( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data=e.request_data, + request_data=_build_passthrough_failure_request_payload( + parsed_body=e.request_data, + kwargs={"call_type": CallTypes.pass_through.value, "litellm_call_id": litellm_call_id}, + logging_obj=None, + custom_llm_provider=custom_llm_provider, + url=url, + builtin_provider_route=builtin_provider_route, + ), ) except Exception: verbose_proxy_logger.warning( @@ -2010,6 +2074,7 @@ async def pass_through_request( content=json.dumps(error_body), status_code=200, media_type="application/json", + headers={"x-litellm-call-id": litellm_call_id}, ) except Exception as e: custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -2033,19 +2098,14 @@ async def pass_through_request( # Monitoring: Trigger post_call_failure_hook # for pass through endpoint failure ######################################################### - request_payload: Final[dict] = _parsed_body or {} - # add user_api_key_dict, litellm_call_id, passthrough_logging_payloa for logging - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - if "model" not in request_payload and _parsed_body and isinstance(_parsed_body, dict): - request_payload["model"] = _parsed_body.get("model", "") - if "custom_llm_provider" not in request_payload and custom_llm_provider: - request_payload["custom_llm_provider"] = custom_llm_provider - + request_payload: Final = _build_passthrough_failure_request_payload( + parsed_body=_parsed_body, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + url=url, + builtin_provider_route=builtin_provider_route, + ) _carry_guardrail_logging_info(request_payload, post_call_guardrail_data) await proxy_logging_obj.post_call_failure_hook( @@ -2537,6 +2597,13 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +def _log_session_usage(logging_obj: LiteLLMLoggingObj, usage: WebSocketSessionUsage) -> None: + logging_obj.model = usage.model # rebind-ok: the spend logger reads model and cost off the shared logging object + logging_obj.model_call_details["model"] = usage.model # rebind-ok: same shared logging object + logging_obj.model_call_details["custom_llm_provider"] = usage.custom_llm_provider # rebind-ok: same + logging_obj.model_call_details["response_cost"] = usage.response_cost # rebind-ok: same + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2547,6 +2614,9 @@ async def websocket_passthrough_request( cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + builtin_provider_route: bool = False, + model: str | None = None, + meter: WebSocketSessionMeter | None = None, ): """ WebSocket passthrough request handler. @@ -2560,6 +2630,9 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider + model: The model the caller named before connecting, logged on a failure before the setup frame arrives + meter: Prices the session from the client text frames delivered upstream, for an upstream that reports no usage """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -2570,6 +2643,7 @@ async def websocket_passthrough_request( # Initialize tracking variables start_time: Final = datetime.now() websocket_messages: Final[list[dict[str, object]]] = [] + metered_total = 0 # rebind-ok: forward_client_to_upstream adds the measure of each client text frame it delivers litellm_call_id: Final = str(uuid.uuid4()) verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target) @@ -2688,6 +2762,7 @@ async def websocket_passthrough_request( async def forward_client_to_upstream() -> None: """Forward messages from client to upstream WebSocket""" + nonlocal metered_total try: while True: message = await websocket.receive() @@ -2757,6 +2832,8 @@ async def websocket_passthrough_request( # Not a JSON message or doesn't contain setup data await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) + if meter is not None: + metered_total += meter.measure(text_data) elif bytes_data is not None: await upstream_ws.send(bytes_data) except asyncio.CancelledError: @@ -2883,6 +2960,8 @@ async def websocket_passthrough_request( self.url = target_url mock_response: Final = MockWebSocketResponse(target) + if meter is not None: + _log_session_usage(logging_obj, meter.price(metered_total)) # Use the same success handler as HTTP passthrough endpoints GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( @@ -2912,19 +2991,17 @@ async def websocket_passthrough_request( except InvalidStatus as exc: verbose_proxy_logger.exception("WebSocket passthrough (%s): upstream rejected WebSocket connection", endpoint) - # Prepare request payload for logging - request_payload = {} - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - # Log the connection failure using the same pattern as HTTP await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=exc, - request_data=request_payload, + request_data=_build_passthrough_failure_request_payload( + parsed_body=None if model is None else {"model": model}, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=None, + url=httpx.URL(target), + builtin_provider_route=builtin_provider_route, + ), traceback_str=traceback.format_exc( limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, ), @@ -2940,19 +3017,17 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): unexpected error while proxying WebSocket", endpoint ) - # Prepare request payload for logging - request_payload = {} - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - # Log the unexpected error using the same pattern as HTTP await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data=request_payload, + request_data=_build_passthrough_failure_request_payload( + parsed_body=None if model is None else {"model": model}, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=None, + url=httpx.URL(target), + builtin_provider_route=builtin_provider_route, + ), traceback_str=traceback.format_exc( limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, ), diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index 1c5d270542f..bd8390d704e 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -10,6 +10,7 @@ Handles guardrail execution for passthrough endpoints with: from typing import Any, Final from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import merged_guardrail_entries from litellm.proxy._types import ( PassThroughGuardrailsConfig, PassThroughGuardrailSettings, @@ -205,7 +206,7 @@ class PassthroughGuardrailHandler: def collect_guardrails( user_api_key_dict: UserAPIKeyAuth, passthrough_guardrails_config: PassThroughGuardrailsConfigInput | None, - ) -> dict[str, bool] | None: + ) -> list[object] | None: """ Collect guardrails for a passthrough endpoint. @@ -225,7 +226,7 @@ class PassthroughGuardrailHandler: passthrough_guardrails_config: List or Dict of guardrail names/settings Returns: - Dict of guardrail names to run (format: {guardrail_name: True}), or None + Guardrail entries to run, each a name or a ``{name: settings}`` entry, or None """ from litellm.proxy.litellm_pre_call_utils import ( add_guardrails_from_key_or_team_metadata, @@ -242,15 +243,6 @@ class PassthroughGuardrailHandler: verbose_proxy_logger.debug("Passthrough guardrails config is empty, skipping") return None - # Passthrough is enabled - collect guardrails - guardrails_to_run: Final[dict[str, bool]] = {} - - # Add passthrough-specific guardrails - for guardrail_name in normalized_config: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added passthrough-specific guardrail: %s", guardrail_name) - - # Add org/team/key level guardrails using shared helper temp_data: Final[dict[str, Any]] = {"metadata": {}} add_guardrails_from_key_or_team_metadata( key_metadata=user_api_key_dict.metadata, @@ -258,20 +250,12 @@ class PassthroughGuardrailHandler: data=temp_data, metadata_variable_name="metadata", ) + inherited_guardrails: Final[list[object]] = temp_data["metadata"].get("guardrails", []) + guardrails_to_run: Final = merged_guardrail_entries(normalized_config, inherited_guardrails) - # Merge inherited guardrails into guardrails_to_run - inherited_guardrails: Final = temp_data["metadata"].get("guardrails", []) - for guardrail_name in inherited_guardrails: - if guardrail_name not in guardrails_to_run: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added inherited guardrail (key/team level): %s", guardrail_name) + verbose_proxy_logger.debug("Collected guardrails for passthrough endpoint: %s", guardrails_to_run) - verbose_proxy_logger.debug( - "Collected guardrails for passthrough endpoint: %s", - list(guardrails_to_run.keys()), - ) - - return guardrails_to_run if guardrails_to_run else None + return guardrails_to_run or None @staticmethod def get_field_targeted_text( diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 47572dd7ceb..b9a60d9e37f 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -329,6 +329,48 @@ def _prepare_hook_input( return hook_input, scans_raw_request +async def _run_pre_call_hook( + target: CustomLogger, + callback: CustomGuardrail, + data: dict[str, object], # mutable-ok: the request payload the run is recorded on + hook_input: dict[str, object], # mutable-ok: same request-payload shape the hooks mutate + user_api_key_dict: "UserAPIKeyAuth", + call_type: str, +) -> dict[str, object] | None: + try: + response: Final = await target.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=None, + data=hook_input, + call_type=call_type, + ) + except Exception: + callback.mark_pre_call_hook_ran(data, raised=True) + raise + callback.mark_pre_call_hook_ran(data) + if not isinstance(response, dict): + return None + payload: Final = cast(dict[str, object], response) # cast-ok: a pre_call hook returns the request payload + callback.mark_pre_call_hook_ran(payload) + return payload + + +def _name_step_that_ran( + request_data: dict[str, object], # mutable-ok: the request payload the step's name is recorded on + metadata_key: Literal["metadata", "litellm_metadata"], + guardrail_name: str, +) -> None: + """Name a step that ran in the applied guardrails of the metadata bucket the pipeline's request started with.""" + existing: Final = request_data.get(metadata_key) + bucket: Final[dict[str, object]] = ( # mutable-ok: the request's own metadata bucket, written in place + cast("dict[str, object]", existing) # cast-ok: metadata buckets are str-keyed + if isinstance(existing, dict) + else {} + ) + request_data[metadata_key] = bucket # rebind-ok: the step's name is recorded on the request payload + add_guardrail_to_applied_guardrails_header(request_data={metadata_key: bucket}, guardrail_name=guardrail_name) + + def _release_original_chunks( guardrail_name: str, reason: str, @@ -389,6 +431,9 @@ class PipelineExecutor: PipelineExecutionResult with terminal action and step results """ step_results: Final[list[PipelineStepResult]] = [] + metadata_key: Final = get_metadata_variable_name_from_kwargs( + cast(Mapping[str, object], data) # cast-ok: the request payload is str-keyed + ) working_data = data.copy() if "metadata" in working_data: working_data["metadata"] = working_data["metadata"].copy() @@ -407,6 +452,7 @@ class PipelineExecutor: data=working_data, user_api_key_dict=user_api_key_dict, call_type=call_type, + metadata_key=metadata_key, raw_request_snapshot=raw_request_snapshot, streaming_chunks=streaming_chunks, endpoint_translation=endpoint_translation, @@ -480,6 +526,7 @@ class PipelineExecutor: user_api_key_dict: "UserAPIKeyAuth", litellm_logging_obj: "LiteLLMLoggingObj | None", stream_already_sent: bool, + metadata_key: Literal["metadata", "litellm_metadata"], ) -> None: """Run one streaming post_call step through the endpoint translation, delivering text and tool-call rewrites on translations that support ended-stream write-back. A @@ -489,8 +536,8 @@ class PipelineExecutor: live, one that drops or adds a tool call, or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the originals and the step passes, so the - client gets the stream the merge base sent, and the guardrail stays out of the - applied-guardrails header since its output never reached the client. The response an + client gets the stream the merge base sent. The guardrail is named in the + applied-guardrails header either way, since it ran. The response an earlier step's translation stored under ``request_data["response"]`` is dropped first, so this step's hook sees the stream as the steps before it left it.""" scanner: Final = ( @@ -525,12 +572,11 @@ class PipelineExecutor: except UndeliverableStreamRewrite as undeliverable: _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals) return + finally: + _name_step_that_ran(hook_input, metadata_key, step.guardrail) discard_reason: Final = observer.discard_reason(rewrite_undeliverable_reason) if discard_reason is not None: _release_original_chunks(step.guardrail, discard_reason, streaming_chunks, originals) - return - if not callback.records_own_guardrail_information: - add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) @staticmethod async def _run_step( @@ -539,6 +585,7 @@ class PipelineExecutor: data: dict, user_api_key_dict: "UserAPIKeyAuth", call_type: str, + metadata_key: Literal["metadata", "litellm_metadata"], raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step endpoint_translation: "BaseTranslation | None" = None, @@ -582,16 +629,7 @@ class PipelineExecutor: try: if mode == "pre_call": - response = await target.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=None, - data=hook_input, - call_type=call_type, - ) - if isinstance(callback, CustomGuardrail): - callback.mark_pre_call_hook_ran(data) - if isinstance(response, dict): - callback.mark_pre_call_hook_ran(response) + response = await _run_pre_call_hook(target, callback, data, hook_input, user_api_key_dict, call_type) elif mode == "post_call" and streaming_chunks is not None: if endpoint_translation is None: return ( @@ -609,9 +647,15 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, litellm_logging_obj=data.get("litellm_logging_obj"), stream_already_sent=stream_already_sent, + metadata_key=metadata_key, ) response = None elif mode == "post_call": + _name_step_that_ran( + cast(dict[str, object], data), # cast-ok: the request payload is str-keyed + metadata_key, + step.guardrail, + ) response = await target.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 2de3f01d640..514f05d0bcd 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -7,6 +7,7 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, add_policy_sources_to_metadata, @@ -18,7 +19,6 @@ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.types.proxy.policy_engine import PolicyMatchContext from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f078267c6b2..a767f5423ed 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -89,6 +89,11 @@ from litellm.constants import ( OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) +from litellm.litellm_core_utils.access_windows import access_windows_config_error +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + get_fallback_errors_from_headers, + get_hidden_params_dict, +) from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.bug_report import ( allowlisted, @@ -100,6 +105,8 @@ from litellm.litellm_core_utils.litellm_logging import ( initialize_sentry, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias +from litellm.litellm_core_utils.routing_groups import parse_routing_groups from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import ( @@ -174,6 +181,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields +from litellm.repositories import database as database_runtime from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, @@ -188,17 +196,11 @@ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( snapshot_tuning_baselines, tuning_limit_violation, ) -from litellm.router_utils.access_windows import access_windows_config_error -from litellm.router_utils.add_retry_fallback_headers import ( - get_fallback_errors_from_headers, - get_hidden_params_dict, -) -from litellm.router_utils.common_utils import resolve_model_group_alias -from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( PRICING_OVERRIDES_KEY, CountContextManagement, + LlmProviders, ModelResponse, ModelResponseStream, TextCompletionResponse, @@ -443,6 +445,7 @@ from litellm.proxy.auth.model_checks import ( get_key_models, get_mcp_server_ids, get_team_models, + run_model_listing, ) from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401, RUF100 # legacy module exports @@ -996,7 +999,7 @@ from litellm.router import ( LiteLLM_Params, ModelGroupInfo, ) -from litellm.router_api import ChangeSet, deployment_id +from litellm.router_api import ChangeSet, deployment_id, router_init_args from litellm.scheduler import FlowItem, Scheduler from litellm.secret_managers.aws_secret_manager import load_aws_kms from litellm.secret_managers.google_kms import load_google_kms @@ -1348,13 +1351,7 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N if worker_heartbeat is not None and prisma_client: await worker_heartbeat.deregister() if prisma_client: - # Drain the SGR fold first: it lives in memory, so an un-drained interval - # is lost, and a write attempted after disconnect raises - # ClientNotConnectedError rather than persisting anything. Ordering this - # inside the same guard is what keeps the two from drifting apart. await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client)) - verbose_proxy_logger.debug("Disconnecting from Prisma") - await prisma_client.disconnect() else: await _drain_audit_tasks_on_shutdown() @@ -1477,23 +1474,38 @@ async def initialize_shared_aiohttp_session() -> "ClientSession | None": _initialize_shared_aiohttp_session: Final = initialize_shared_aiohttp_session -async def _connect_to_count_stored_values() -> SupportsRawQueries: +async def _connect_to_count_stored_values(resources: AsyncExitStack) -> SupportsRawQueries: client: Final = prisma_client or PrismaClient( database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj ) + if prisma_client is None: + resources.push_async_callback(client.aclose) await client.connect() return client.writer_db +async def _own_database(resources: AsyncExitStack, client: PrismaClient | None) -> database_runtime.Database | None: + if client is None: + return None + database: Final = await database_runtime.Database.open(client) + resources.push_async_callback(database.aclose) + return database + + +def _forget_prisma_client() -> None: + global prisma_client # noqa: PLW0603 # routes still read this global until they move to DatabaseDep + prisma_client = None # rebind-ok: a closed client must not survive into the next lifespan entry + + @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: with GracefulShutdownManager.capture_sigterm(): - async with _proxy_startup_event(app) as state: + async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state: yield state @asynccontextmanager -async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: +async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGenerator[ProxyLifespanState, None]: global \ prisma_client, \ master_key, \ @@ -1512,6 +1524,9 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat shared_aiohttp_session import json + resources.callback(_forget_prisma_client) + existing_database: Final = await _own_database(resources, prisma_client) + init_verbose_loggers() running_loop: Final = asyncio.get_running_loop() @@ -1590,21 +1605,24 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat if decoded_worker_config is not None: await initialize_from_worker_config(decoded_worker_config) - enforce_master_key_boot_verdict( - await with_stored_secrets_counted( - master_key_boot_verdict( - master_key=master_key, - environment_master_key=os.getenv(MASTER_KEY_ENV_VAR), - general_settings=general_settings, - config_file_path=user_config_file_path, - override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True, - salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, - database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None, + async with AsyncExitStack() as boot_database: + enforce_master_key_boot_verdict( + await with_stored_secrets_counted( + master_key_boot_verdict( + master_key=master_key, + environment_master_key=os.getenv(MASTER_KEY_ENV_VAR), + general_settings=general_settings, + config_file_path=user_config_file_path, + override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True, + salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, + database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None, + ), + count_values_encrypted_with=partial( + count_values_encrypted_with_or_none, partial(_connect_to_count_stored_values, boot_database) + ), ), - count_values_encrypted_with=partial(count_values_encrypted_with_or_none, _connect_to_count_stored_values), - ), - announce=announce_on_stderr_at_exit, - ) + announce=announce_on_stderr_at_exit, + ) # check if DATABASE_URL in environment - load from there if prisma_client is None: @@ -1617,6 +1635,8 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat ) ) + database: Final = existing_database or await _own_database(resources, prisma_client) + await migrate_if_requested( environ=os.environ, master_key=master_key, @@ -1626,15 +1646,16 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat ) if prisma_client is not None: + startup_prisma: Final = prisma_client - async def _run_pw_migration(): + async def _run_pw_migration() -> None: try: - result: Final = await migrate_passwords_to_scrypt_async(prisma_client) + result: Final = await migrate_passwords_to_scrypt_async(startup_prisma) verbose_proxy_logger.info("Password migration: %s", result) except Exception as e: verbose_proxy_logger.warning("Password migration skipped: %s", e) - asyncio.create_task(_run_pw_migration()) + await resources.enter_async_context(database_runtime.background_task(_run_pw_migration())) async def _run_agent_grant_id_migration() -> None: from litellm.proxy.agent_endpoints.agent_registry import ( @@ -1645,7 +1666,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat for attempt in range(3): try: result = await global_agent_registry.migrate_legacy_grant_ids( - table=object_permission_table(prisma_client) + table=object_permission_table(startup_prisma) ) if result.rewritten: verbose_proxy_logger.info( @@ -1665,7 +1686,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat if attempt < 2: await asyncio.sleep(5) - asyncio.create_task(_run_agent_grant_id_migration()) + await resources.enter_async_context(database_runtime.background_task(_run_agent_grant_id_migration())) ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before @@ -1851,7 +1872,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat enabled=tracing_enabled, settings=tracing_settings, ) as receiver: - state: Final[ProxyLifespanState] = {"tracing_receiver": receiver} + state: Final[ProxyLifespanState] = {"tracing_receiver": receiver, "database": database} from litellm.proxy.admin_mcp import admin_mcp_lifespan telemetry_settings: Final = load_telemetry_settings() @@ -1905,30 +1926,6 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat except Exception as e: verbose_proxy_logger.error("Error closing shared aiohttp session: %s", e) - # Shutdown event - stop RDS IAM token refresh background task - if ( - prisma_client is not None - and hasattr(prisma_client, "db") - and hasattr(prisma_client.db, "stop_token_refresh_task") - ): - try: - await prisma_client.db.stop_token_refresh_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping token refresh task: %s", e) - - # Shutdown event - stop Prisma DB health watchdog task - if prisma_client is not None and hasattr(prisma_client, "stop_db_health_watchdog_task"): - try: - await prisma_client.stop_db_health_watchdog_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) - - if prisma_client is not None and hasattr(prisma_client, "stop_view_setup_task"): - try: - await prisma_client.stop_view_setup_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping the spend view setup task: %s", e) - await _drain_spend_event_producer_on_shutdown() if telemetry_start is not None: telemetry_start.cancel() @@ -3508,7 +3505,7 @@ async def increment_spend_counters( user_id: str | None, response_cost: float | None, org_id: str | None = None, - budget_reservation: dict | None = None, + budget_reservation: dict[str, object] | None = None, end_user_id: str | None = None, tags: list[str] | None = None, request_started_at: datetime | None = None, @@ -4663,10 +4660,6 @@ def _write_health_state_to_router_cache( is on. """ from litellm.proxy.health_check import build_deployment_health_states - from litellm.router_utils.cooldown_handlers import set_cooldown_deployments - from litellm.router_utils.router_callbacks.track_deployment_metrics import ( - increment_deployment_failures_for_current_minute, - ) _exceptions: Final[dict] = exceptions_by_model_id or {} @@ -4718,18 +4711,7 @@ def _write_health_state_to_router_cache( ): continue - increment_deployment_failures_for_current_minute( - litellm_router_instance=llm_router, - deployment_id=model_id, - ) - - set_cooldown_deployments( - litellm_router_instance=llm_router, - original_exception=original_exception, - exception_status=exception_status, - deployment=model_id, - time_to_cooldown=router_settings.cooldown_time, - ) + llm_router.health().record_failure(model_id, original_exception, exception_status) except Exception as e: verbose_proxy_logger.warning("Failed to write health state to router cache: %s", str(e)) @@ -7373,9 +7355,7 @@ class ProxyConfig: router_settings: Final = config.get("router_settings", None) if router_settings and isinstance(router_settings, dict): - available_args: Final = [ - x for x in litellm.Router.get_valid_args() if x not in ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG - ] + available_args: Final = [x for x in router_init_args() if x not in ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG] for k, v in router_settings.items(): if k in available_args: @@ -9599,7 +9579,6 @@ class ProxyConfig: SearchToolRegistry, keep_loaded_search_tools_that_do_not_decrypt, ) - from litellm.router_utils.search_api_router import SearchAPIRouter try: db_search_tools: Final = keep_loaded_search_tools_that_do_not_decrypt( @@ -9623,7 +9602,21 @@ class ProxyConfig: ) if llm_router is not None: - await SearchAPIRouter.update_router_search_tools(router_instance=llm_router, search_tools=search_tools) + llm_router.apply( + ChangeSet( + assign={ + "search_tools": [ + { + "search_tool_id": tool.get("search_tool_id"), + "search_tool_name": tool.get("search_tool_name"), + "litellm_params": tool.get("litellm_params", {}), + "search_tool_info": tool.get("search_tool_info"), + } + for tool in search_tools + ] + } + ) + ) verbose_proxy_logger.info("Successfully loaded %s search tool(s) into router", len(search_tools)) else: verbose_proxy_logger.debug( @@ -12122,17 +12115,20 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, user_api_key_cache: UserApiKeyCache, ) -> PrismaClient | None: - """ - - Sets up prisma client - - Adds necessary views to proxy - """ + if database_url is None: + return None + client: Final = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) + async with AsyncExitStack() as setup_resources: + setup_resources.push_async_callback(client.aclose) + connected: Final = await cls._initialize_prisma_client(client) + if connected is not None: + setup_resources.pop_all() + return connected + + @staticmethod + async def _initialize_prisma_client(prisma_client: PrismaClient) -> PrismaClient | None: connected_client: PrismaClient | None = None try: - if database_url is None: - return None - - prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) - try: await prisma_client.connect() except Exception as e: @@ -12487,20 +12483,23 @@ async def model_list( # Get complete model list including wildcard routes if requested from litellm.proxy.auth.model_checks import get_complete_model_list - all_models = get_complete_model_list( - key_models=[], - team_models=[], - proxy_model_list=[ - *proxy_model_list, - *(typesafe_allowlist_model_names() if asks_for_allowlist_names else ()), - ], - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=include_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups or False, - only_model_access_groups=only_model_access_groups or False, + all_models = await run_model_listing( + partial( + get_complete_model_list, + key_models=[], + team_models=[], + proxy_model_list=[ + *proxy_model_list, + *(typesafe_allowlist_model_names() if asks_for_allowlist_names else ()), + ], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=include_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups or False, + only_model_access_groups=only_model_access_groups or False, + ) ) expanded_undiscoverable_names: Final = undiscoverable_model_names( @@ -14747,6 +14746,8 @@ from litellm.responses.litellm_completion_transformation.count_request import ( with_provider_file_ids, ) +_PROVIDERS_WITH_LISTING_ONLY_MODEL_INFO: Final = frozenset({LlmProviders.OPENROUTER}) + def _with_client_oauth_token( counter: BaseTokenCounter | None, client_oauth_token: str | None @@ -14782,14 +14783,13 @@ def _get_provider_token_counter( ) # Switch case pattern using existing get_provider_model_info - from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager # Convert string provider to LlmProviders enum llm_provider_enum: Final = LlmProviders(custom_llm_provider) # Add more provider mappings as needed - if llm_provider_enum: + if llm_provider_enum not in _PROVIDERS_WITH_LISTING_ONLY_MODEL_INFO: provider_model_info: Final = ProviderConfigManager.get_provider_model_info( model=full_model, provider=llm_provider_enum ) @@ -16148,7 +16148,7 @@ def _enrich_model_info_with_litellm_data( if debug is True: _openai_client = "None" if llm_router is not None: - _openai_client = llm_router._get_client(deployment=model, kwargs={}, client_type="async") or "None" + _openai_client = llm_router.registry().client(model, client_type="async") or "None" else: _openai_client = "llm_router_is_None" openai_client: Final = str(_openai_client) @@ -17941,7 +17941,7 @@ async def model_info_v1( alias_models: Final = copy.deepcopy(registry.deployments(source="alias")) all_models.extend(alias_models) - all_models = expand_wildcard_deployments_for_model_info(all_models, llm_router) + all_models = await run_model_listing(partial(expand_wildcard_deployments_for_model_info, all_models, llm_router)) allowed_model_names: Final = await _get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=user_api_key_dict, @@ -18279,14 +18279,17 @@ async def model_group_info( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ) all_models_str: Final = ( - get_complete_model_list( - key_models=(), - team_models=(), - proxy_model_list=llm_router.registry().names(), - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), - return_wildcard_routes=False, - llm_router=llm_router, + await run_model_listing( + partial( + get_complete_model_list, + key_models=(), + team_models=(), + proxy_model_list=llm_router.registry().names(), + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=False, + llm_router=llm_router, + ) ) if is_proxy_admin else await get_available_models_for_user( @@ -18595,7 +18598,7 @@ async def async_queue_request( router: Final = llm_router data.update( - router.priority_header_update( + router.health().priority_header_fields( model=data["model"], kwargs=data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER) ) ) @@ -19744,7 +19747,7 @@ async def update_config( raw_router_settings: Final = request_body.get("router_settings") if isinstance(raw_router_settings, dict): supported_router_settings: Final = RUNTIME_UPDATABLE_ROUTER_SETTINGS | ( - frozenset(litellm.Router.get_valid_args()) - ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG + frozenset(router_init_args()) - ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG ) unsupported_router_settings: Final = sorted(set(raw_router_settings) - supported_router_settings) if unsupported_router_settings: @@ -21901,6 +21904,7 @@ async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "Streami while True: chunk = await body_queue.get() if chunk is None: + await handler_task break yield chunk finally: @@ -21981,9 +21985,9 @@ async def toolset_mcp_route(toolset_name: str, request: Request): Namespace a toolset as its own MCP endpoint. Connecting to /toolset//mcp exposes exactly the tools defined in - the toolset. Access is enforced: non-admin API keys must have the toolset - listed in their object_permission.mcp_toolsets grant list, or the request - will be rejected with a 403. + the toolset. Access is enforced: API keys other than an admin's key with no + team must hold the toolset through one of their grant sources, or the + request will be rejected with a 403. """ try: from litellm.proxy.mcp_server.mcp_server_manager import ( diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index df3af4ad487..01a822c5d93 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2377,6 +2377,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "META", + "provider_display_name": "Meta Model API", + "litellm_provider": "meta", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "meta/muse-spark-1.1" + }, { "provider": "MILVUS", "provider_display_name": "Milvus", @@ -4257,6 +4285,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "ZAI", + "provider_display_name": "Z.AI (Zhipu AI)", + "litellm_provider": "zai", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.z.ai/api/paas/v4", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "zai/glm-4.5" + }, { "provider": "ZeroGPU", "provider_display_name": "ZeroGPU", diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index ef59deb2e5e..089cb75a795 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,11 +1,13 @@ import asyncio import contextlib +import copy import json import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum from functools import partial +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -35,7 +37,7 @@ from litellm.anthropic_interface.exceptions.exception_mapping_utils import Anthr from litellm.caching.dual_cache import DualCache from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.exceptions import EndpointNotSupportedByModelError -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import ModifyResponseException, guardrail_request_data_with_streaming from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -47,6 +49,11 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( build_blocked_response, ) from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import ( + get_team_object, + requests_guardrail_modification, + require_guardrail_modification_permission, +) from litellm.proxy.auth.auth_utils import carries_client_side_overrides from litellm.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, @@ -57,6 +64,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, cost_headers, create_response, + proxy_exception_from_route_error, resolve_counted_tool_fields, resolve_litellm_call_id, ) @@ -79,6 +87,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry from litellm.proxy.litellm_pre_call_utils import ( _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + add_guardrails_from_policy_engine, is_chat_completions_body, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -89,7 +98,12 @@ from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_ from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD from litellm.responses.retrieve_stream_query import RETRIEVE_INCLUDE_QUERY_KEYS, RETRIEVE_STREAM_QUERY_KEYS -from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task +from litellm.responses.streaming_iterator import ( + GuardedResponsesWsFrame, + ResponsesWsFrameGuard, + ResponsesWsFrameRejection, + spawn_logging_task, +) from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -2149,11 +2163,15 @@ _PII_TOKENS: Final = TypeAdapter(dict[str, str]) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" _GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") -# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") +_METADATA_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config") +_PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response") +_UNFORWARDED_FRAME_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS, "litellm_metadata") +_OBJECTS: Final = TypeAdapter(list[object]) +_TURN_COPIED_CONTAINERS: Final = (list, dict) -class _PreCallGuardrailWalk(Protocol): +class _GuardrailWalk(Protocol): async def pre_call_hook( self, *, @@ -2163,15 +2181,49 @@ class _PreCallGuardrailWalk(Protocol): guardrails_only: bool, ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns the guardrails' own dict + async def during_call_hook( + self, + *, + data: dict[str, object], # mutable-ok: ProxyLogging.during_call_hook hands the guardrails a dict they edit + user_api_key_dict: UserAPIKeyAuth | None, + call_type: Literal["aresponses"], + guardrails_only: bool, + ) -> object: ... + + +async def _key_team_object(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + team_id: Final = user_api_key_dict.team_id + if team_id is None or team_id == UI_TEAM_ID: + return None + team_from_key: Final = LiteLLM_TeamTable(team_id=team_id, metadata=user_api_key_dict.team_metadata) + if prisma_client is None: + return team_from_key + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + return team_from_key + def _responses_ws_frame_guard( admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - proxy_logging_obj: _PreCallGuardrailWalk, + proxy_logging_obj: _GuardrailWalk, record_refused_turn: Callable[[tuple[Mapping[str, object], ...], Exception], None], + *, + team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]] = _key_team_object, ) -> ResponsesWsFrameGuard: """Socket admission skips guardrails (the first frame's ``input`` is not part of the admitted request), so every - ``response.create`` frame gets the pre_call guardrail walk the HTTP ``/v1/responses`` path runs on its body.""" + ``response.create`` frame gets the pre_call and during_call guardrail walks the HTTP ``/v1/responses`` path runs on + its body. during_call runs before the frame is forwarded, since a forwarded frame cannot be recalled, and like HTTP, + which sends the model call beside during_call, the frame carries the pre_call output and not during_call rewrites.""" metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted) admitted_bucket: Final = admitted.get(metadata_key) bucket: Final = { @@ -2186,32 +2238,57 @@ def _responses_ws_frame_guard( # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) + async def rejection( + exception: Exception, + data: dict[str, object], # mutable-ok: billed + ) -> ResponsesWsFrameRejection: + await _bill_rejected_frame(rejected_frame_auth, exception, data) + proxy_exception: Final = proxy_exception_from_route_error(exception) + return ResponsesWsFrameRejection(proxy_exception.message, proxy_exception.provider_specific_fields) + async def guard(frame: str) -> GuardedResponsesWsFrame: frame_obj: Final = _WS_FRAME.validate_json(frame) nested: Final = frame_obj.get("response") - if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS): - raise ValueError( - "response.create carries input, instructions or tools both at the top level and under response" - ) body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj - frame_data: Final = { - **{ - key: _without_untrusted_metadata(key, value) - for key, value in body.items() - if key != "type" and key not in _CLIENT_GUARDRAIL_SELECTORS + selection_sources: Final = (frame_obj, body) if isinstance(nested, dict) else (body,) + turn_data: Final = guardrail_request_data_with_streaming( + { + **{ + key: _without_untrusted_metadata(key, value) + for key, value in body.items() + if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS + }, + "model": model, + "litellm_call_id": str(uuid.uuid4()), + metadata_key: {**bucket}, }, - "model": model, - "litellm_call_id": str(uuid.uuid4()), - # shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied - metadata_key: {**bucket}, - } + is_streaming=True, + ) + try: + _reject_guarded_field_beside_nested_response(frame_obj, nested) + await _check_frame_may_modify_guardrails(selection_sources, user_api_key_dict, team_object) + frame_data: Final = { + **turn_data, + metadata_key: await _metadata_with_frame_guardrails( + bucket, selection_sources, model, metadata_key, user_api_key_dict + ), + } + except Exception as e: + raise await rejection(e, turn_data) from e try: guarded_data: Final = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True ) + # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them + guarded_fields: Final = { + key: copy.deepcopy(guarded_data[key]) for key in _GUARDED_FRAME_FIELDS if key in guarded_data + } + pii_tokens: Final = _recorded_pii_tokens(guarded_data.get("metadata")) + await proxy_logging_obj.during_call_hook( + data=guarded_data, user_api_key_dict=user_api_key_dict, call_type="aresponses", guardrails_only=True + ) except Exception as e: - await _bill_rejected_frame(rejected_frame_auth, e, frame_data) - raise + raise await rejection(e, frame_data) from e if guarded_data.get("model") != model: rerouted: Final = ValueError( f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket" @@ -2219,25 +2296,186 @@ def _responses_ws_frame_guard( ) await _bill_rejected_frame(rejected_frame_auth, rerouted, frame_data) raise rerouted - # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them - guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data} - kept: Final = {key: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS} rewritten: Final = ( - {**frame_obj, "response": {**kept, **guarded_fields}} + {**_forwarded_frame_fields(frame_obj), "response": {**_forwarded_frame_fields(body), **guarded_fields}} if isinstance(nested, dict) - else {**kept, **guarded_fields} + else {**_forwarded_frame_fields(body), **guarded_fields} ) guardrail_information: Final = _recorded_guardrail_information(guarded_data.get(metadata_key)) return GuardedResponsesWsFrame( frame=json.dumps(rewritten), guardrail_information=guardrail_information, - pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + pii_tokens=pii_tokens, bill_refusal=partial(record_refused_turn, guardrail_information), ) return guard +def _reject_guarded_field_beside_nested_response(frame: Mapping[str, object], nested: object) -> None: + """Only the nested ``response`` body runs through the guardrails, so a guarded field beside it would go upstream + unguarded""" + stray: Final = ( + next((key for key in _GUARDED_FRAME_FIELDS if key in frame), None) if isinstance(nested, dict) else None + ) + if stray is not None: + raise ValueError(f"response.create sends {stray} at the top level next to response; send it inside response") + + +async def _check_frame_may_modify_guardrails( + sources: tuple[Mapping[str, object], ...], + user_api_key_dict: UserAPIKeyAuth, + team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]], +) -> None: + """A frame that picks guardrails or policies needs the team's permission to modify guardrails""" + if any( + "policies" in source or requests_guardrail_modification(_without_client_metadata(source)) for source in sources + ): + require_guardrail_modification_permission(await team_object(user_api_key_dict)) + + +def _without_client_metadata(source: Mapping[str, object]) -> Mapping[str, object]: + metadata: Final = source.get("metadata") + return {**source, "metadata": _metadata_picks("metadata", metadata)} if isinstance(metadata, dict) else source + + +def _forwarded_frame_fields(source: Mapping[str, object]) -> dict[str, object]: + return { + key: _without_guardrail_selection(value) if key == "metadata" else value + for key, value in source.items() + if key not in _UNFORWARDED_FRAME_FIELDS + } + + +class _FrameGuardrailSelection(TypedDict, total=False): + guardrails: ReadOnly[Sequence[str | Mapping[str, Mapping[str, object]]]] + guardrail_config: ReadOnly[Mapping[str, object]] + policies: ReadOnly[Sequence[str]] + + +_FRAME_GUARDRAIL_SELECTION: Final = TypeAdapter(_FrameGuardrailSelection) + + +async def _metadata_with_frame_guardrails( + bucket: Mapping[str, object], + sources: tuple[Mapping[str, object], ...], + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + """A frame selects guardrails the way an HTTP ``/v1/responses`` body does, adding to the key, team and policy + guardrails the socket was admitted with, so a frame can never shrink that list. A nested frame selects at its top + level and in its ``response`` body, the body last so its ``guardrail_config`` wins""" + selections: Final = tuple(chain.from_iterable(_frame_guardrail_selections(source) for source in sources)) + picked_policies: Final = tuple(selection["policies"] for selection in selections if "policies" in selection) + metadata: Final = ( + await _metadata_with_frame_policies( + bucket, tuple(chain.from_iterable(picked_policies)), model, metadata_key, user_api_key_dict + ) + if picked_policies + else bucket + ) + selected_guardrails: Final = tuple(chain.from_iterable(selection.get("guardrails", ()) for selection in selections)) + selected_config: Final = dict( + chain.from_iterable( + selection["guardrail_config"].items() for selection in selections if "guardrail_config" in selection + ) + ) + existing_guardrails: Final = metadata.get("guardrails") + existing_config: Final = metadata.get("guardrail_config") + return { + **metadata, + **( + { + "guardrails": [ + *(_OBJECTS.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else ()), + *selected_guardrails, + ] + } + if selected_guardrails + else {} + ), + **( + { + "guardrail_config": { + **(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}), + **selected_config, + } + } + if selected_config + else {} + ), + } + + +def _frame_guardrail_selections(body: Mapping[str, object]) -> tuple[_FrameGuardrailSelection, ...]: + """Like HTTP ``/v1/responses``, a frame selects guardrails at its top level and in its ``metadata`` and + ``litellm_metadata``, the top level last so its ``guardrail_config`` wins""" + return ( + *( + _frame_guardrail_selection((key,), _metadata_picks(key, value), _METADATA_GUARDRAIL_SELECTORS) + for key, value in body.items() + if key in _CLIENT_METADATA_KEYS and isinstance(value, dict) + ), + _frame_guardrail_selection((), body, _CLIENT_GUARDRAIL_SELECTORS), + ) + + +def _metadata_picks(key: str, metadata: object) -> Mapping[str, object]: + """OpenAI ``metadata`` values are strings, so a string there is the client's own metadata and never a pick""" + validated: Final = _WS_FRAME.validate_python(metadata) + if key != "metadata": + return validated + return { + name: item + for name, item in validated.items() + if name not in _METADATA_GUARDRAIL_SELECTORS or not isinstance(item, str) + } + + +def _frame_guardrail_selection( + location: tuple[str, ...], source: Mapping[str, object], selectors: tuple[str, ...] +) -> _FrameGuardrailSelection: + try: + return _FRAME_GUARDRAIL_SELECTION.validate_python({key: source[key] for key in selectors if key in source}) + except ValidationError as e: + raise ValueError(f"response.create {first_validation_error(e, prefix=location)}") from e + + +def _without_guardrail_selection(metadata: object) -> object: + if not isinstance(metadata, dict): + return metadata + return { + name: item + for name, item in _WS_FRAME.validate_python(metadata).items() + if name not in _METADATA_GUARDRAIL_SELECTORS or isinstance(item, str) + } + + +async def _metadata_with_frame_policies( + bucket: Mapping[str, object], + policies: Sequence[str], + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + selection: Final = { + "model": model, + "policies": list(policies), + metadata_key: {key: _turn_copy(value) for key, value in bucket.items()}, + } + await add_guardrails_from_policy_engine( + data=selection, metadata_variable_name=metadata_key, user_api_key_dict=user_api_key_dict + ) + return _WS_FRAME.validate_python(selection[metadata_key]) + + +def _turn_copy(value: object) -> object: + """Shallow, since the bucket carries proxy-owned objects (the OTEL parent span) that cannot be copied, and the + policy engine writes its lists in place""" + return copy.copy(value) if type(value) in _TURN_COPIED_CONTAINERS else value + + def _without_untrusted_metadata(key: str, value: object) -> object: """Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge.""" if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict): diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index a59f1967fc1..a17b7fee8c9 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -16,8 +16,9 @@ from pydantic import BeforeValidator, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request from litellm.proxy._types import ProxyException, UserAPIKeyAuth -from litellm.router_utils.common_utils import is_proxy_admin_request +from litellm.router_api import router_init_args # Client-supplied params that make the router or the call path fabricate a # failure or a delay instead of calling the provider. The ``mock_testing_*`` @@ -93,7 +94,7 @@ async def _route_user_config_request(data: dict, route_type: str): # Filter router_config to only include valid Router.__init__ arguments # This prevents TypeError when invalid parameters are stored in the database - valid_args: Final = litellm.Router.get_valid_args() + valid_args: Final = router_init_args() filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args} logging_obj: Final = _request_logging_obj(JSON_OBJECT_ADAPTER.validate_python(data), route_type) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 59c7fe2a41b..9f37d2965bc 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Final import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import OPENAI_CHAT_PDF_FILE_TOKENS, PDF_DATA_URL_PREFIX +from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models, str_keyed from litellm.litellm_core_utils.token_counter import ( deployment_openai_model, pdf_page_pricing, @@ -31,7 +32,6 @@ from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction, holds_compaction_summary, ) -from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.token_counter import ( @@ -162,7 +162,7 @@ def models_counted_for(model: str, llm_router: Router | None, team_id: str | Non def _routed_deployments(model: str, llm_router: Router, team_id: str | None) -> Sequence[object]: registry: Final = llm_router.registry() resolved: Final = registry.resolve(model, kind="alias") or model - routing_group: Final = llm_router.get_routing_group(resolved) + routing_group: Final = registry.routing_group(resolved) members: Final = routing_group.models if routing_group is not None else (resolved,) named: Final = tuple( chain.from_iterable(registry.deployments(member, source="all", team_id=team_id) for member in members) diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 28602fbbefc..8e79ea5255f 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -689,9 +689,13 @@ def compute_savings_spend( cost_breakdown=cost_breakdown, ) ) - return SavingsSpend( + savings: Final = SavingsSpend( compression=compression, prompt_caching=prompt_caching, autorouter=0.0 if autorouter is None else autorouter, gateway_injected_caching=gateway_injected_caching, ) + if all(isfinite(value) for value in savings): + return savings + verbose_proxy_logger.warning("Spend tracking - recording non-finite savings %s as 0 for model=%s", savings, model) + return SavingsSpend(*(value if isfinite(value) else 0.0 for value in savings)) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 09b7ce7ee9a..306e70f7029 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,6 +11,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, cast, runtime_checkable from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -194,6 +195,7 @@ _STAMPED_METADATA_KEYS: Final = frozenset( "autorouter_savings_estimate", "autorouter_baseline_observation", "used_client_oauth_token", + "upstream_provider", KUBERNETES_POD_ROUTING_KEY, "litellm_roi_estimator", "is_silent_experiment", @@ -203,6 +205,10 @@ _STAMPED_METADATA_KEYS: Final = frozenset( ) +class _UpstreamProviderMetadata(TypedDict, total=False): + upstream_provider: ReadOnly[str] + + def _stored_user_email(metadata: Mapping[str, object]) -> str | None: email: Final = metadata.get("user_api_key_user_email") return email if isinstance(email, str) and email.strip() else None @@ -221,6 +227,7 @@ def _get_spend_logs_metadata( model_map_information: StandardLoggingModelInformation | None = None, cold_storage_object_key: str | None = None, litellm_overhead_time_ms: float | None = None, + upstream_provider: str | None = None, cost_breakdown: CostBreakdown | None = None, litellm_call_id: str | None = None, autorouter_savings: float | None = None, @@ -232,6 +239,9 @@ def _get_spend_logs_metadata( kubernetes_pod_routing: StandardLoggingKubernetesPodRouting | None = None, is_silent_experiment: bool = False, ) -> SpendLogsMetadata: + upstream_provider_metadata: Final[_UpstreamProviderMetadata] = ( + {"upstream_provider": upstream_provider} if upstream_provider is not None else {} + ) if metadata is None: return SpendLogsMetadata( user_api_key=None, @@ -265,6 +275,7 @@ def _get_spend_logs_metadata( eval_information=None, cold_storage_object_key=cold_storage_object_key, litellm_overhead_time_ms=None, + **upstream_provider_metadata, attempted_retries=None, max_retries=None, attempted_fallbacks=None, @@ -290,6 +301,7 @@ def _get_spend_logs_metadata( **MappingProxyType( {key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key not in _STAMPED_METADATA_KEYS} ), + **upstream_provider_metadata, autorouter_savings=autorouter_savings, autorouter_savings_estimate=autorouter_savings_estimate, autorouter_baseline_observation=autorouter_baseline_observation, @@ -706,6 +718,14 @@ def get_logging_payload( if standard_logging_payload is not None: hidden_params: Final = standard_logging_payload.get("hidden_params", {}) litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms") + upstream_provider_value: Final[object] = ( + standard_logging_payload["hidden_params"]["upstream_provider"] + if standard_logging_payload is not None + and "hidden_params" in standard_logging_payload + and "upstream_provider" in standard_logging_payload["hidden_params"] + else None + ) + upstream_provider: Final = upstream_provider_value if isinstance(upstream_provider_value, str) else None logged_provider: Final = ( kwargs.get("custom_llm_provider") @@ -807,6 +827,7 @@ def get_logging_payload( else None ), litellm_overhead_time_ms=litellm_overhead_time_ms, + upstream_provider=upstream_provider, cost_breakdown=( standard_logging_payload.get("cost_breakdown", None) if standard_logging_payload is not None else None ), @@ -912,7 +933,7 @@ def get_logging_payload( user=metadata.get("user_api_key_user_id", "") or "", team_id=metadata.get("user_api_key_team_id", "") or "", organization_id=metadata.get("user_api_key_org_id") or "", - metadata=safe_dumps(clean_metadata), + metadata=safe_dumps(clean_metadata, non_finite_as_null=True), cache_key=cache_key, spend=kwargs.get("response_cost", 0), total_tokens=usage.get("total_tokens", standard_logging_total_tokens), @@ -1224,7 +1245,7 @@ def _get_messages_for_spend_logs_payload( messages: Final = standard_logging_payload.get("messages") if messages is not None: try: - return safe_dumps(messages) + return safe_dumps(messages, non_finite_as_null=True) except Exception: return "{}" return "{}" @@ -1919,7 +1940,8 @@ def _get_proxy_server_request_for_spend_logs_payload( _request_body_json_str: Final = safe_dumps( _stored_request_without_requester_ip_address(stored_request_body) if drop_requester_ip_address - else stored_request_body + else stored_request_body, + non_finite_as_null=True, ) if LITELLM_TRUNCATED_PAYLOAD_FIELD in _request_body_json_str: verbose_proxy_logger.info( @@ -2000,7 +2022,7 @@ def _get_response_for_spend_logs_payload( if isinstance(sanitized_response, str): result_str = strip_null_bytes(sanitized_response) else: - result_str = safe_dumps(sanitized_response) + result_str = safe_dumps(sanitized_response, non_finite_as_null=True) if LITELLM_TRUNCATED_PAYLOAD_FIELD in result_str: verbose_proxy_logger.info( "Spend Log: response was truncated before storing in DB. %s", @@ -2071,7 +2093,7 @@ def spend_log_row_with_retained_metadata( return row metadata: Final = _SPEND_LOGS_METADATA_ADAPTER.validate_json(metadata_json) retained: Final = {name: value for name, value in metadata.items() if fields.keeps(name)} - return {**row, "metadata": safe_dumps(retained)} + return {**row, "metadata": safe_dumps(retained, non_finite_as_null=True)} def _get_status_for_spend_log( diff --git a/litellm/proxy/spend_tracking/usage_meter.py b/litellm/proxy/spend_tracking/usage_meter.py new file mode 100644 index 00000000000..3bc7572b476 --- /dev/null +++ b/litellm/proxy/spend_tracking/usage_meter.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +import asyncio +from collections import deque +from collections.abc import Awaitable, Callable +from concurrent.futures import Future +from dataclasses import dataclass +from itertools import accumulate, chain, count, takewhile +from threading import Lock +from typing import Final + +from pydantic import BaseModel, ConfigDict, Field + +from litellm._logging import verbose_proxy_logger +from litellm.accounting import ( + DeliveryTarget, + Dimension, + DimensionName, + InboundUsageEvent, + OutboundUsageEvent, +) +from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite + +_DAILY_TARGETS: Final = ( + (DimensionName.KEY, DeliveryTarget.DAILY_USER_SPEND), + (DimensionName.USER, DeliveryTarget.DAILY_USER_SPEND), + (DimensionName.TEAM, DeliveryTarget.DAILY_TEAM_SPEND), + (DimensionName.ORGANIZATION, DeliveryTarget.DAILY_ORGANIZATION_SPEND), + (DimensionName.END_USER, DeliveryTarget.DAILY_END_USER_SPEND), + (DimensionName.AGENT, DeliveryTarget.DAILY_AGENT_SPEND), + (DimensionName.TAG, DeliveryTarget.DAILY_TAG_SPEND), + (DimensionName.MODEL, DeliveryTarget.DAILY_MODEL_USAGE), + (DimensionName.TOOL, DeliveryTarget.DAILY_TOOL_SPEND), +) + + +def usage_delivery_targets(dimensions: tuple[Dimension, ...]) -> tuple[DeliveryTarget, ...]: + names: Final = frozenset(dimension.name for dimension in dimensions) + return tuple(dict.fromkeys(target for dimension, target in _DAILY_TARGETS if dimension in names)) + + +class UsageCaptureSettings(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + batch_size: int = Field(default=1000, ge=2, le=1000) + flush_interval_seconds: float = Field(default=5.0, gt=0, le=10) + max_pending_bytes: int = Field(default=16 * 1024 * 1024, ge=1024) + max_pending_records: int = Field(default=10000, ge=2) + + +class UsageCaptureError(RuntimeError): + pass + + +@dataclass(frozen=True, slots=True) +class _Pending: + writes: tuple[UsageWrite, ...] + size: int + committed: Future[None] + finished: asyncio.Event + prerequisites: tuple[asyncio.Future[None], ...] + + +async def _wait_for_batch(wakeup: asyncio.Event, timeout: float) -> None: + try: + await asyncio.wait_for(wakeup.wait(), timeout) + except asyncio.TimeoutError: + return + + +def _report_failure(receipt: Future[None]) -> None: + if not receipt.cancelled() and receipt.exception() is not None: + verbose_proxy_logger.error("Usage capture did not acknowledge commit", exc_info=receipt.exception()) + + +def _retryable_persistence_failure(error: Exception) -> bool: + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain + + outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(error) + return (outage is not None and not PrismaDBExceptionHandler.is_permanent_database_fault(outage)) or any( + isinstance(cause, Exception) + and ( + PrismaDBExceptionHandler.is_deadlock_error(cause) + or PrismaDBExceptionHandler.is_lock_timeout_error(cause) + or PrismaDBExceptionHandler.postgres_sqlstate(cause) == "40001" + ) + for cause in exception_chain(error) + ) + + +def _isolatable_persistence_failure(error: Exception) -> bool: + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain + + return isinstance(error, AccountingConflictError) or any( + isinstance(cause, Exception) and (PrismaDBExceptionHandler.postgres_sqlstate(cause) or "").startswith("22") + for cause in exception_chain(error) + ) + + +class DurableUsageMeter: + def __init__( + self, + write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], + settings: UsageCaptureSettings, + *, + wait_for_batch: Callable[[asyncio.Event, float], Awaitable[None]] = _wait_for_batch, + retry_delay: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + self._write_batch: Final = write_batch + self.settings: Final = settings + self._wait_for_batch: Final = wait_for_batch + self._retry_delay: Final = retry_delay + self._loop: Final = asyncio.get_running_loop() + self._admission: Final = Lock() + self._queue: Final[deque[_Pending]] = deque() # mutable-ok: one bounded producer/consumer queue + self._wakeup: Final = asyncio.Event() + self._task: asyncio.Task[None] | None = None + self._closed = False + self.pending_bytes = 0 + self.pending_records = 0 + + def submit( + self, writes: tuple[UsageWrite, ...], *, prerequisites: tuple[asyncio.Future[None], ...] = () + ) -> asyncio.Future[None]: + size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes) + with self._admission: + if ( + self._closed + or not writes + or len(writes) > self.settings.batch_size + or self.pending_bytes + size > self.settings.max_pending_bytes + or self.pending_records + len(writes) > self.settings.max_pending_records + ): + raise UsageCaptureError("Usage capture admission rejected: closed, empty or capacity exhausted") + receipt: Final = Future[None]() + receipt.set_running_or_notify_cancel() + receipt.add_done_callback(_report_failure) + self._queue.append(_Pending(writes, size, receipt, asyncio.Event(), prerequisites)) + self.pending_bytes += size + self.pending_records += len(writes) + self._loop.call_soon_threadsafe(self._ensure_worker) + wrapped: Final = asyncio.wrap_future(receipt, loop=self._loop) + wrapped.add_done_callback(lambda done: None if done.cancelled() else done.exception()) + return wrapped + + def _ensure_worker(self) -> None: + if self.pending_records >= self.settings.batch_size or self._queue and self._queue[0].prerequisites: + self._wakeup.set() + if self._task is None or self._task.done(): + self._task = asyncio.create_task(self._run(), name="accounting usage capture") + + async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: + envelope: Final = '{"event":' + event.model_dump_json() + "}" + await asyncio.shield(self.submit((UsageWrite(event, envelope, usage_delivery_targets(event.dimensions)),))) + + async def flush(self) -> None: + self._ensure_worker() + self._wakeup.set() + if self._task is not None: + await asyncio.shield(self._task) + + async def close(self) -> None: + with self._admission: + self._closed = True + await self.flush() + + async def stop(self, timeout: float) -> None: + with self._admission: + self._closed = True + try: + await asyncio.wait_for(self.close(), timeout) + except asyncio.TimeoutError: + if self._task is not None: + self._task.cancel() + await asyncio.gather(self._task, return_exceptions=True) + while self._queue: + self._finish( + self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement") + ) + verbose_proxy_logger.error("Usage capture shutdown deadline exhausted; commit acknowledgement incomplete") + + def _take_batch(self) -> tuple[_Pending, ...]: + with self._admission: + if self._queue[0].prerequisites: + return (self._queue.popleft(),) + independent: Final = takewhile(lambda item: not item.prerequisites, self._queue) + sizes: Final = accumulate(len(item.writes) for item in independent) + count: Final = sum(1 for _ in takewhile(lambda size: size <= self.settings.batch_size, sizes)) + return tuple(self._queue.popleft() for _ in range(count)) + + async def _persist(self, writes: tuple[UsageWrite, ...]) -> None: + for attempt in count(): + try: + await self._write_batch(writes) + except Exception as error: # noqa: BLE001 # retry only classified database failures, preserving frozen writes + if not _retryable_persistence_failure(error): + raise + verbose_proxy_logger.warning("Usage persistence unavailable; retaining admitted batch for retry") + await self._retry_delay(min(10.0, 0.2 * 2 ** min(attempt, 6))) + else: + return + + async def _commit(self, batch: tuple[_Pending, ...]) -> None: + writes: Final = tuple(chain.from_iterable(item.writes for item in batch)) + try: + if batch[0].prerequisites: + outcomes: Final = await asyncio.gather( + *(asyncio.shield(receipt) for receipt in batch[0].prerequisites), return_exceptions=True + ) + for outcome in outcomes: + if isinstance(outcome, BaseException): + raise UsageCaptureError("Usage dependency did not acknowledge commit") from outcome + await self._persist(writes) + except asyncio.CancelledError: + self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted")) + raise + except Exception as error: # noqa: BLE001 # isolate data faults, surface systemic failures to all producers + if len(batch) > 1 and _isolatable_persistence_failure(error): + try: + for item in batch: + await self._commit((item,)) + except asyncio.CancelledError: + self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted")) + raise + return + self._finish( + batch, + UsageCaptureError("Usage identity conflicts with accepted content") + if isinstance(error, AccountingConflictError) + else error, + ) + else: + self._finish(batch, None) + + def _finish(self, batch: tuple[_Pending, ...], error: Exception | None) -> None: + for item in batch: + if item.finished.is_set(): + continue + item.finished.set() + with self._admission: + self.pending_bytes -= item.size + self.pending_records -= len(item.writes) + if item.committed.done(): + continue + if error is None: + item.committed.set_result(None) + else: + item.committed.set_exception(error) + + async def _run(self) -> None: + try: + while self._queue: + await self._wait_for_batch(self._wakeup, self.settings.flush_interval_seconds) + self._wakeup.clear() + await self._commit(self._take_batch()) + if ( + self._closed + or self.pending_records >= self.settings.batch_size + or self._queue + and self._queue[0].prerequisites + ): + self._wakeup.set() + except (Exception, asyncio.CancelledError): # noqa: BLE001 # reject queued receipts if the writer exits abnormally + while self._queue: + self._finish( + self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement") + ) + raise diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py new file mode 100644 index 00000000000..6d5b43504df --- /dev/null +++ b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py @@ -0,0 +1,31 @@ +from collections.abc import Coroutine +from typing import Annotated, Protocol, TypeAlias + +from fastapi import Depends + +from litellm.models.user_banner import UserBanner +from litellm.proxy._types import LitellmTableNames, UserAPIKeyAuth +from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID + + +class UserBannerAudit(Protocol): + def __call__(self, before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> Coroutine[None, None, None]: ... + + +async def get_user_banner_audit() -> UserBannerAudit: + from litellm.proxy.proxy_server import create_config_audit_log + + async def audit(before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> None: + await create_config_audit_log( + param_name=USER_BANNER_ROW_ID, + action="updated", + before_value=before.model_dump(), + after_value=after.model_dump(), + user_api_key_dict=user, + table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME, + ) + + return audit + + +UserBannerAuditDep: TypeAlias = Annotated[UserBannerAudit, Depends(get_user_banner_audit)] diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py index 56339202d46..a839827667c 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py @@ -1,88 +1,43 @@ -import asyncio -import json -from typing import Annotated, Final, Literal +from typing import Annotated, Final -from fastapi import APIRouter, Depends, HTTPException -from pydantic import Field, ValidationError, model_validator +from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException +from pydantic import ConfigDict from litellm._uuid import uuid4 -from litellm.proxy._types import LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth +from litellm.models.user_banner import UserBanner, UserBannerUpdate +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository +from litellm.proxy.deps.repositories import UserBannerRepositoryDep +from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import ( + UserBannerAuditDep, +) from litellm.types.llms.base import LiteLLMBaseModel router: Final = APIRouter() -USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000 - -UserBannerSeverity = Literal["info", "warning", "error"] - - -class UserBannerUpdate(LiteLLMBaseModel): - enabled: bool = Field( - default=False, - description="If true, the banner is shown to all authenticated dashboard users.", - ) - message: str = Field( - default="", - max_length=USER_BANNER_MAX_MESSAGE_LENGTH, - description="Banner text shown to dashboard users. Markdown is supported.", - ) - severity: UserBannerSeverity = Field( - default="info", - description="Visual style of the banner.", - ) - - @model_validator(mode="after") - def _require_message_when_enabled(self) -> "UserBannerUpdate": - if self.enabled and not self.message.strip(): - raise ValueError("message must be non-empty when the banner is enabled") - return self - - -class UserBanner(UserBannerUpdate): - revision: str = Field( - default="", - description=( - "Server-stamped opaque publish identity; a fresh value is generated on every " - "update so clients re-surface dismissed banners on republish." - ), - ) - class UpdateUserBannerResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + message: str banner: UserBanner -def parse_user_banner(raw_settings: object) -> UserBanner: - if raw_settings is None: - return UserBanner() - try: - parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings - return UserBanner.model_validate(parsed) - except (json.JSONDecodeError, ValidationError): - return UserBanner() - - @router.get( "/get/user_banner", tags=["UI Settings"], dependencies=[Depends(user_api_key_auth)], response_model=UserBanner, ) -async def get_user_banner() -> UserBanner: +async def get_user_banner(banners: UserBannerRepositoryDep) -> UserBanner: """ Get the admin-published dashboard banner. Readable by any authenticated user; rendered on every dashboard page. """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: + if banners is None: return UserBanner() - raw_settings: Final = await UserBannerRepository(prisma_client).get_raw_settings() - return parse_user_banner(raw_settings) + return await banners.get_banner() @router.patch( @@ -93,21 +48,21 @@ async def get_user_banner() -> UserBanner: async def update_user_banner( banner_update: UserBannerUpdate, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + banners: UserBannerRepositoryDep, + audit: UserBannerAuditDep, + background_tasks: BackgroundTasks, ) -> UpdateUserBannerResponse: """ Publish, edit, or unpublish the dashboard banner. Only proxy admins are allowed to modify it. """ - from litellm.proxy.proxy_server import create_config_audit_log, prisma_client - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException(status_code=403, detail="Only proxy admins can update the user banner.") - if prisma_client is None: + if banners is None: raise HTTPException(status_code=500, detail="Database not connected. Please connect a database.") - repository: Final = UserBannerRepository(prisma_client) - before: Final = parse_user_banner(await repository.get_raw_settings()) + before: Final = await banners.get_banner() banner: Final = UserBanner( enabled=banner_update.enabled, message=banner_update.message, @@ -115,17 +70,8 @@ async def update_user_banner( revision=uuid4().hex, ) - await repository.upsert_settings(json.dumps(banner.model_dump())) + await banners.save_banner(banner) - asyncio.create_task( - create_config_audit_log( - param_name=USER_BANNER_ROW_ID, - action="updated", - before_value=before.model_dump(), - after_value=banner.model_dump(), - user_api_key_dict=user_api_key_dict, - table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME, - ) - ) + background_tasks.add_task(audit, before, banner, user_api_key_dict) return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b1055bb9124..9e237ffcb57 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -85,12 +85,13 @@ from litellm.proxy._types import ( from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.openai_error_payload import ( attribute_of, - litellm_call_id_headers, - openai_error_param, + error_status_code, + openai_shaped_proxy_exception, with_litellm_call_id, ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction +from litellm.responses.mcp.parked_stream_dispatch import held_round_spend from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage @@ -146,6 +147,7 @@ from litellm.exceptions import ( from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, + merged_guardrail_entries, ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger @@ -158,10 +160,12 @@ from litellm.litellm_core_utils.core_helpers import ( independent_snapshot, is_expected_client_error, ) +from litellm.litellm_core_utils.deployment_pricing import pricing_models from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.served_output_texts import ( @@ -276,8 +280,6 @@ from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) -from litellm.router_utils.common_utils import resolve_model_group_alias -from litellm.router_utils.deployment_pricing import pricing_models from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES from litellm.types.llms.openai import ResponsesAPIResponse @@ -632,6 +634,8 @@ def _is_client_error_exception(exc: Exception) -> bool: return exc.status_code < 500 if isinstance(exc, ProxyException): return not (exc.code.isdigit() and int(exc.code) >= 500) + if isinstance(exc, litellm.BudgetExceededError): + return exc.status_code < 500 return False @@ -710,6 +714,120 @@ def pipeline_managed_guardrail_names( ) +def _ran_or_runs_outside_pipelines( + callback: CustomGuardrail | None, + data: Mapping[str, object], + pipeline_modes: frozenset[str], +) -> bool: + return callback is not None and ( + callback.pre_call_hook_ran_or_raised(data) + or callback.deployment_hook_ran(data) + or any( + callback._event_hook_is_event_type(hook) # pyright: ignore[reportPrivateUsage] # no public hook check + for hook in GuardrailEventHooks + if hook.value not in pipeline_modes + ) + ) + + +def _runs_only_in_pipelines( + name: str, data: Mapping[str, object], pre_call: frozenset[str], post_call: frozenset[str] +) -> bool: + return not _ran_or_runs_outside_pipelines( + PipelineExecutor.find_guardrail_callback(name), + data, + frozenset(mode for mode, names in (("pre_call", pre_call), ("post_call", post_call)) if name in names), + ) + + +def _request_steps_that_have_not_run(data: Mapping[str, object]) -> frozenset[str]: + """Pre_call policy pipeline steps a pipeline skipped, whose guardrail nothing else ran or runs for this request.""" + pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") + return frozenset(name for name in pre_call if _runs_only_in_pipelines(name, data, pre_call, frozenset())) + + +def _steps_that_ran(result: PipelineExecutionResult) -> frozenset[str]: + return frozenset(step.guardrail_name for step in result.step_results if step.outcome != "skip") + + +def _whole_response_log_waits_for_its_pipelines(data: Mapping[str, object], call_type: str) -> bool: + """ + The success log is parked until the response hooks finish, and the headers are built + after them. The LLM pass-through route builds its headers first, and the generateContent + adapter logs its nested completion before the response hooks run. + """ + logging_obj: Final = data.get("litellm_logging_obj") + return ( + call_type not in (CallTypes.allm_passthrough_route.value, CallTypes.agenerate_content.value) + and isinstance(logging_obj, Logging) + and logging_obj.enqueue_deferred_logging is not None + ) + + +def _stream_log_waits_for_its_pipelines(data: Mapping[str, object]) -> bool: + return getattr(data.get("litellm_logging_obj"), "_on_deferred_stream_complete", None) is not None + + +def _pipelines_settled_by( + pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult +) -> Sequence[tuple[str, "GuardrailPipeline"]]: + """This pipeline, plus every later response pipeline when its verdict ends the request before they run.""" + if result.terminal_action == "allow": + return pipelines[index : index + 1] + return ( + pipelines[index], + *((policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call"), + ) + + +def _steps_later_pipelines_still_run( + pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult +) -> frozenset[str]: + """The steps of the later response pipelines, which still run when this pipeline's verdict lets the request on.""" + if result.terminal_action != "allow": + return frozenset() + later: Final = tuple( + (policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call" + ) + return _pipeline_step_guardrail_names(later) + + +def _drop_response_steps_that_did_not_run( + data: dict[str, object], # mutable-ok: the request's applied guardrails are trimmed in place + pipelines: Sequence[tuple[str, "GuardrailPipeline"]], + steps_to_keep: frozenset[str], +) -> None: + """ + Drop from the applied guardrails the steps of these response pipelines outside ``steps_to_keep``, + unless their guardrail also runs outside pipelines. Every step is named up front, so a + log built before the pipeline finished, such as a cache hit's, still lists the steps + that ran. The list is trimmed in place, since a route's logging object may hold a + shallow copy of the metadata. + """ + pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") + post_call: Final = _pipeline_step_guardrail_names(pipelines) + not_run: Final = frozenset( + name for name in post_call - steps_to_keep if _runs_only_in_pipelines(name, data, pre_call, post_call) + ) + if not not_run: + return + for bucket_key in ("metadata", "litellm_metadata"): + _drop_applied_guardrails_in_place(data.get(bucket_key), not_run) + + +def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None: + if not isinstance(bucket, dict): + return + metadata: Final = cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed + applied: Final = metadata.get("applied_guardrails") + if not isinstance(applied, list): + return + entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names + entries[:] = [name for name in entries if name not in names] + if not entries: + metadata.pop("applied_guardrails") + + def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: resolved: Final = tuple( litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -2707,8 +2825,12 @@ class ProxyLogging: data.pop("prompt_version", None) data.pop("prompt_environment", None) - def _process_guardrail_metadata(self, data: dict) -> None: - """Process guardrails from metadata and add to applied_guardrails.""" + def _process_guardrail_metadata(self, data: dict, *, skip_guardrails: bool = False) -> None: + """ + Process guardrails from metadata and add to applied_guardrails, leaving out the pre_call + policy pipeline steps that have not run. With ``skip_guardrails`` the pipelines run on + another request, so every step is added. + """ from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) @@ -2732,8 +2854,17 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] + not_run: Final = ( + frozenset[str]() + if skip_guardrails + else _request_steps_that_have_not_run(cast(dict[str, object], data)) # cast-ok: request payload + ) for guardrail_name in guardrails_in_metadata: - if isinstance(guardrail_name, str) and guardrail_name not in applied_guardrails: + if ( + isinstance(guardrail_name, str) + and guardrail_name not in applied_guardrails + and guardrail_name not in not_run + ): add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) async def _maybe_execute_pipelines( @@ -2765,7 +2896,10 @@ class ProxyLogging: return data, None current_response = response # rebind-ok: chains each pipeline's replacement response into the next - for policy_name, pipeline in pipelines: + steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran + request: Final = cast(dict[str, object], data) # cast-ok: request payload, updated in place by every pipeline + log_waits: Final = event_hook == "post_call" and _whole_response_log_waits_for_its_pipelines(request, call_type) + for index, (policy_name, pipeline) in enumerate(pipelines): if pipeline.mode != event_hook: continue @@ -2781,6 +2915,14 @@ class ProxyLogging: raw_request_snapshot=raw_request_snapshot, ) + steps_that_ran = steps_that_ran | _steps_that_ran(result) + if log_waits: + _drop_response_steps_that_did_not_run( + request, + _pipelines_settled_by(pipelines, index, result), + steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result), + ) + data = self._handle_pipeline_result( result=result, data=data, @@ -3023,7 +3165,7 @@ class ProxyLogging: and (guardrails_only or not caps.has_pre_call_override) ): if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) return data parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( @@ -3127,17 +3269,17 @@ class ProxyLogging: ) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) return data except SensitiveDataRouteException as e: data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) return data except Exception: if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) raise async def _run_parallel_pre_call_guardrails( @@ -3503,6 +3645,7 @@ class ProxyLogging: data: dict, user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ): caps: Final = ProxyLogging.callback_capabilities() if not caps.has_guardrail and not caps.has_moderation_override: @@ -3514,6 +3657,7 @@ class ProxyLogging: if ( isinstance(callback, CustomLogger) and not isinstance(callback, CustomGuardrail) + and (not guardrails_only or callback.enforces_request_content) and _overrides_moderation_hook(callback) and user_api_key_dict is not None ): @@ -4730,7 +4874,8 @@ class ProxyLogging: call_type, endpoint_translation = translation - for policy_name, pipeline in pipelines: + steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran + for index, (policy_name, pipeline) in enumerate(pipelines): result: PipelineExecutionResult = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode="post_call", @@ -4741,6 +4886,13 @@ class ProxyLogging: streaming_chunks=buffered, endpoint_translation=endpoint_translation, ) + steps_that_ran = steps_that_ran | _steps_that_ran(result) + if _stream_log_waits_for_its_pipelines(request_data): + _drop_response_steps_that_did_not_run( + request_data, + _pipelines_settled_by(pipelines, index, result), + steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result), + ) try: ProxyLogging._handle_pipeline_result( result, data=request_data, policy_name=policy_name, original_response=buffered @@ -4943,6 +5095,8 @@ class ProxyLogging: if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS): return False match _args: + case (ModelResponse(), *_) if (held := held_round_spend(logging_obj)) is not None: + logging_obj.record_billed_rounds_for_failure(held) case (ModelResponse() as assembled, *_): logging_obj.record_assembled_response_for_failure(assembled) case (BilledRoundsDispatch() as billed,): @@ -5189,6 +5343,8 @@ async def prefetch_config_params(prisma_client: "PrismaClient | None", param_nam _WRITER_WRITABILITY_PROBE_SQL: Final = "SELECT current_setting('transaction_read_only') AS transaction_read_only" _WRITER_WRITABILITY_PROBE_ROWS: Final = TypeAdapter(list[dict[str, object]]) _READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS: Final = 600 +_SHUTDOWN_RECONNECT_WAIT_SECONDS: Final = 5.0 +_SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS: Final = 5.0 class _ForcedRecreateDeclined(Exception): @@ -5273,6 +5429,7 @@ class PrismaClient: database_url: str, proxy_logging_obj: ProxyLogging, http_client: "HttpConfig | None" = None, + engine_exit_wait_seconds: float = _SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS, ): ## init logging object self.baseline_accounting_transactions: list[BaselineAccountingRecord] = [] @@ -5374,6 +5531,9 @@ class PrismaClient: else: self.db = writer_wrapper # Client to connect to Prisma db self._db_reconnect_lock = asyncio.Lock() + self._closing: bool = False + self._engine_exit_wait_seconds: Final = engine_exit_wait_seconds + self._engine_tasks: frozenset[asyncio.Task[object]] = frozenset() self._db_health_watchdog_task: asyncio.Task | None = None self._view_setup_task: asyncio.Task[_ViewSetupOutcome] | None = None self._db_last_reconnect_attempt_ts: float = 0.0 @@ -5415,6 +5575,7 @@ class PrismaClient: self._watching_engine: bool = False self._engine_confirmed_dead: bool = False self._engine_wait_thread: threading.Thread | None = None + self._engine_wait_threads: tuple[threading.Thread, ...] = () verbose_proxy_logger.debug("Success - Created Prisma Client") @property @@ -6623,6 +6784,14 @@ class PrismaClient: start_time: Final = time.time() try: await self.db.disconnect() + if self._closing: + await asyncio.gather( + *( + asyncio.to_thread(thread.join, self._engine_exit_wait_seconds) + for thread in self._engine_wait_threads + ) + ) + self._engine_wait_threads = () except Exception as e: import traceback @@ -6641,6 +6810,54 @@ class PrismaClient: ) raise e + def _spawn_engine_task(self, work: Coroutine[object, object, object]) -> None: + if self._closing: + work.close() + return + task: Final = asyncio.create_task(work) + self._engine_tasks = self._engine_tasks | {task} + task.add_done_callback(self._engine_task_finished) + + def _engine_task_finished(self, task: asyncio.Task[object]) -> None: + self._engine_tasks = self._engine_tasks - {task} + if not task.cancelled() and (error := task.exception()) is not None: + verbose_proxy_logger.error("Prisma engine background task failed: %s", error) + + async def _stop_engine_tasks(self) -> None: + tasks: Final = tuple(self._engine_tasks) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await self._wait_for_inflight_reconnect() + + async def _wait_for_inflight_reconnect(self) -> None: + waiter: Final = asyncio.ensure_future(self._db_reconnect_lock.acquire()) + waiter.add_done_callback(lambda acquired: acquired.cancelled() or self._db_reconnect_lock.release()) + done, _pending = await asyncio.wait({waiter}, timeout=_SHUTDOWN_RECONNECT_WAIT_SECONDS) + if waiter not in done: + waiter.cancel() + verbose_proxy_logger.warning("Shutdown stopped waiting for an in-flight DB reconnect") + + async def aclose(self) -> None: + verbose_proxy_logger.debug("Disconnecting from Prisma") + try: + await self.stop_background_tasks() + except Exception: + verbose_proxy_logger.exception("Stopping database background tasks failed") + try: + await self.disconnect() + except Exception: + verbose_proxy_logger.exception("Database disconnect failed") + + async def stop_background_tasks(self) -> None: + self._closing = True + self._stop_engine_watcher() + async with contextlib.AsyncExitStack() as cleanup: + cleanup.push_async_callback(self._stop_engine_tasks) + cleanup.push_async_callback(self.db.stop_token_refresh_task) + cleanup.push_async_callback(self.stop_view_setup_task) + cleanup.push_async_callback(self.stop_db_health_watchdog_task) + def _get_engine_pid(self) -> int: """Get the PID of the writer's engine subprocess, or 0 if unavailable. @@ -6734,7 +6951,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6754,6 +6971,9 @@ class PrismaClient: name=f"prisma-engine-waitpid-{pid}", ) thread.start() + self._engine_wait_threads = tuple(thread for thread in self._engine_wait_threads if thread.is_alive()) + ( + thread, + ) self._engine_wait_thread = thread return True @@ -6815,7 +7035,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6877,7 +7097,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6950,7 +7170,7 @@ class PrismaClient: waitpid thread nor pidfd are available. """ - if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + if self._closing or self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: return pid: Final = self._get_engine_pid() if pid == 0: @@ -6977,7 +7197,7 @@ class PrismaClient: pid, ) self._watching_engine = True - asyncio.create_task(self._poll_engine_proc()) + self._spawn_engine_task(self._poll_engine_proc()) def _stop_engine_watcher(self) -> None: """Stop watching the engine process and clean up all resources.""" @@ -6999,7 +7219,7 @@ class PrismaClient: """ self._engine_confirmed_dead = False self._cleanup_engine_watcher() - asyncio.create_task(self._start_engine_watcher()) + self._spawn_engine_task(self._start_engine_watcher()) async def _run_reconnect_cycle( self, @@ -7206,6 +7426,8 @@ class PrismaClient: force_recreate: bool = False, stale_read_engine: "_StaleReadEngine | None" = None, ) -> bool: + if self._closing: + return False now: Final = time.time() if ( force is False @@ -7396,7 +7618,7 @@ class PrismaClient: - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling. """ - if self._db_health_watchdog_enabled is not True: + if self._closing or self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug("Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED") return if self._db_health_watchdog_task is not None: @@ -7430,7 +7652,7 @@ class PrismaClient: verbose_proxy_logger.info("Stopped Prisma DB health watchdog") def start_view_setup_task(self) -> None: - if self._view_setup_task is not None: + if self._closing or self._view_setup_task is not None: return self._view_setup_task = asyncio.create_task(self._run_view_setup()) @@ -9095,8 +9317,7 @@ def _merge_guardrails_with_existing(data: dict, model_level_guardrails: object) if not isinstance(model_level_guardrails, list): model_level_guardrails = [model_level_guardrails] if model_level_guardrails else [] - # Combine existing and model-level guardrails - metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails)) + metadata["guardrails"] = merged_guardrail_entries(existing_guardrails, model_level_guardrails) return modified_data @@ -9186,39 +9407,30 @@ def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | Non asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction")) +def _client_message(e: Exception) -> str: + status_is_none: Final = hasattr(e, "status_code") and attribute_of(e, "status_code") is None + return "Internal server error" if status_is_none else strip_bug_report_notice(str(e)) + + def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) -> ProxyException: """ Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible """ - from fastapi import status - verbose_proxy_logger.exception("Exception: %s", e) if PrismaDBExceptionHandler.is_read_only_transaction_error(e): from litellm.proxy.proxy_server import prisma_client _recreate_writer_on_read_only_transaction(prisma_client) - headers: Final = litellm_call_id_headers(litellm_call_id) - if isinstance(e, HTTPException): - return ProxyException( - message=getattr(e, "detail", f"error({e})"), - type=ProxyErrorTypes.internal_server_error, - param=openai_error_param(e), - headers=headers, - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): + if isinstance(e, ProxyException): return with_litellm_call_id(e, litellm_call_id) - _status_code: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) + if isinstance(e, HTTPException): + return with_litellm_call_id(openai_shaped_proxy_exception(e, str(e.detail), e.status_code), litellm_call_id) if should_report_bug(e): verbose_proxy_logger.error(bug_report_notice(build_proxy_bug_report(e))) - return ProxyException( - message=strip_bug_report_notice(str(e)), - type=ProxyErrorTypes.internal_server_error, - param=openai_error_param(e), - headers=headers, - code=_status_code, - ) + message: Final = _client_message(e) + status_code: Final = error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR) + return with_litellm_call_id(openai_shaped_proxy_exception(e, message, status_code), litellm_call_id) def require_enterprise_license(feature: str | None = None) -> None: @@ -9506,7 +9718,7 @@ async def _agent_access_group_visible_models( ) -> frozenset[str] | None: """Models an agent key may still list once its attached access groups cap it, ``None`` when nothing caps it, so ``/v1/models`` never advertises a model the same key would be denied on.""" - from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models + from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models, run_model_listing if not user_api_key_dict.agent_id: return None @@ -9527,17 +9739,20 @@ async def _agent_access_group_visible_models( if not granted: return frozenset() return frozenset( - get_complete_model_list( - key_models=granted, - team_models=(), - proxy_model_list=proxy_model_list, - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=return_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - team_id=team_id, + await run_model_listing( + partial( + get_complete_model_list, + key_models=granted, + team_models=(), + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=return_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + team_id=team_id, + ) ) ) @@ -9582,6 +9797,7 @@ async def get_available_models_for_user( get_complete_model_list, get_key_models, get_team_models, + run_model_listing, ) # Get proxy model list and access groups @@ -9662,18 +9878,21 @@ async def get_available_models_for_user( ) # Get complete model list - complete_models: Final = get_complete_model_list( - key_models=granted_key_models, - team_models=granted_team_models, - proxy_model_list=proxy_model_list, - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), - return_wildcard_routes=return_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - only_model_access_groups=only_model_access_groups, - team_id=effective_team_id, + complete_models: Final = await run_model_listing( + partial( + get_complete_model_list, + key_models=granted_key_models, + team_models=granted_team_models, + proxy_model_list=proxy_model_list, + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=return_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + only_model_access_groups=only_model_access_groups, + team_id=effective_team_id, + ) ) all_models: Final = ( list( diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py new file mode 100644 index 00000000000..d60a2099477 --- /dev/null +++ b/litellm/repositories/database.py @@ -0,0 +1,41 @@ +import asyncio +from collections.abc import AsyncGenerator, Coroutine +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final + +from litellm.rust_bridge._native import NativeDatabase + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +@dataclass(frozen=True, slots=True) +class Database: + prisma: "PrismaClient" + native: NativeDatabase + + @classmethod + async def open(cls, prisma: "PrismaClient") -> "Database": + try: + native: Final = await NativeDatabase.connect() + except BaseException: + await prisma.aclose() + raise + return cls(prisma=prisma, native=native) + + async def aclose(self) -> None: + try: + await self.prisma.aclose() + finally: + await self.native.close() + + +@asynccontextmanager +async def background_task(work: Coroutine[object, object, None]) -> AsyncGenerator[None, None]: + task: Final = asyncio.create_task(work) + try: + yield + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) diff --git a/litellm/repositories/deployment_warning_ignore_repository.py b/litellm/repositories/deployment_warning_ignore_repository.py index da02d424f80..95c9121d4e1 100644 --- a/litellm/repositories/deployment_warning_ignore_repository.py +++ b/litellm/repositories/deployment_warning_ignore_repository.py @@ -1,5 +1,12 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.models.deployment_warning import AdminIgnore from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: @@ -7,15 +14,28 @@ if TYPE_CHECKING: DEPLOYMENT_WARNING_IGNORES_ROW_ID: Final = "deployment_warning_ignores" +_IGNORES_ADAPTER: Final[TypeAdapter[Mapping[str, AdminIgnore]]] = TypeAdapter(Mapping[str, AdminIgnore]) + + +def parse_ignores(raw: object) -> Mapping[str, AdminIgnore]: + if raw is None: + return MappingProxyType({}) + try: + return MappingProxyType(_IGNORES_ADAPTER.validate_python(json.loads(raw) if isinstance(raw, str) else raw)) + except (json.JSONDecodeError, ValidationError): + verbose_proxy_logger.warning("Ignoring unreadable %s row", DEPLOYMENT_WARNING_IGNORES_ROW_ID) + return MappingProxyType({}) + class DeploymentWarningIgnoreRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]): table_name = "litellm_uisettings" - async def get_raw_ignores(self) -> object: + async def get_ignores(self) -> Mapping[str, AdminIgnore]: db_record: Final = await self.table.find_unique(where={"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID}) - return db_record.ui_settings if db_record is not None else None + return parse_ignores(db_record.ui_settings if db_record is not None else None) - async def upsert_ignores(self, payload: str) -> None: + async def save_ignores(self, ignores: Mapping[str, AdminIgnore]) -> None: + payload: Final = _IGNORES_ADAPTER.dump_json(dict(ignores)).decode() row: Final = {"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID, "ui_settings": payload} await self.table.upsert( where={"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID}, diff --git a/litellm/repositories/gateway_activity_repository.py b/litellm/repositories/gateway_activity_repository.py new file mode 100644 index 00000000000..2950668a931 --- /dev/null +++ b/litellm/repositories/gateway_activity_repository.py @@ -0,0 +1,41 @@ +from typing import TYPE_CHECKING, Annotated, Final + +from pydantic import Field, TypeAdapter + +from litellm.models.gateway_activity import GatewayActivity + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +_ACTIVITY_SQL: Final = """ + SELECT + 'route' AS kind, date, category, route, + SUM(successful_requests)::bigint AS successful_requests, + SUM(failed_requests)::bigint AS failed_requests, + NULL::integer AS status_code + FROM "LiteLLM_DailyGatewayRequests" + WHERE date >= $1 AND date <= $2 + GROUP BY date, category, route + UNION ALL + SELECT + 'status' AS kind, NULL, NULL, NULL, + 0::bigint, + SUM(failed_requests)::bigint, + status_code + FROM "LiteLLM_DailyGatewayFailedRequests" + WHERE date >= $1 AND date <= $2 + GROUP BY status_code +""" + +_ACTIVITY_ROWS_ADAPTER: Final = TypeAdapter(tuple[Annotated[GatewayActivity, Field(discriminator="kind")], ...]) + + +class GatewayActivityRepository: + def __init__(self, prisma_client: "PrismaClient") -> None: + self._prisma_client = prisma_client + + async def get_activity(self, start_date: str, end_date: str) -> tuple[GatewayActivity, ...]: + raw_rows: Final = await self._prisma_client.db.query_raw( # pyright: ignore[reportAny] # untyped Prisma client + _ACTIVITY_SQL, start_date, end_date + ) + return _ACTIVITY_ROWS_ADAPTER.validate_python(raw_rows or ()) diff --git a/litellm/repositories/user_banner_repository.py b/litellm/repositories/user_banner_repository.py index 4e113dff988..bc25d3af1fe 100644 --- a/litellm/repositories/user_banner_repository.py +++ b/litellm/repositories/user_banner_repository.py @@ -1,5 +1,9 @@ +import json from typing import TYPE_CHECKING, Final +from pydantic import ValidationError + +from litellm.models.user_banner import UserBanner from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: @@ -8,9 +12,25 @@ if TYPE_CHECKING: USER_BANNER_ROW_ID: Final = "user_banner" +def parse_user_banner(raw_settings: object) -> UserBanner: + if raw_settings is None: + return UserBanner() + try: + parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings + return UserBanner.model_validate(parsed) + except (json.JSONDecodeError, ValidationError): + return UserBanner() + + class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]): table_name = "litellm_uisettings" + async def get_banner(self) -> UserBanner: + return parse_user_banner(await self.get_raw_settings()) + + async def save_banner(self, banner: UserBanner) -> None: + await self.upsert_settings(json.dumps(banner.model_dump())) + async def get_raw_settings(self) -> object: db_record: Final = await self.table.find_unique(where={"id": USER_BANNER_ROW_ID}) return db_record.ui_settings if db_record is not None else None diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 1f56ddf863e..e13151eff1b 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -33,6 +33,14 @@ from litellm.types.utils import ModelResponse _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) +def _chat_verbosity(responses_api_request: Mapping[str, object]) -> Mapping[str, object]: + match responses_api_request.get("text"): + case {"verbosity": str() as verbosity}: + return {"verbosity": verbosity} + case _: + return {} + + @dataclass(frozen=True, slots=True) class _BridgedRequest: input: str | ResponseInputParam @@ -95,8 +103,8 @@ class LiteLLMCompletionTransformationHandler: bridged: Final = _bridged_request(hoisted, responses_api_request, ()) bridged_input: Final = bridged.input bridged_request: Final = bridged.responses_api_request - litellm_completion_request: Final[dict] = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + litellm_completion_request: Final[dict] = { + **LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, input=bridged_input, responses_api_request=bridged_request, @@ -104,8 +112,9 @@ class LiteLLMCompletionTransformationHandler: stream=stream, extra_headers=extra_headers, **kwargs, - ) - ) + ), + **_chat_verbosity(bridged_request), + } if _is_async: return self.async_response_api_handler( @@ -169,8 +178,8 @@ class LiteLLMCompletionTransformationHandler: return await self.async_response_api_handler( previous_response_client_ids, stored_turns, - litellm_completion_request=( - LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + litellm_completion_request={ + **LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, input=bridged.input, responses_api_request=bridged.responses_api_request, @@ -178,8 +187,9 @@ class LiteLLMCompletionTransformationHandler: stream=stream, extra_headers=extra_headers, **request_kwargs, - ) - ), + ), + **_chat_verbosity(bridged.responses_api_request), + }, request_input=bridged.input, responses_api_request=bridged.responses_api_request, **request_kwargs, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index ce8841c9190..8f0a85d67de 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -398,23 +398,25 @@ def tool_output_to_chat_content( if single_media is not None: return [single_media] if isinstance(output, list): + parts: Final = cast(Sequence[object], output) # cast-ok: checked list shape, elements remain unvalidated normalized_blocks: Final[list[ToolMessageContentPart]] = [] text_acc: Final[list[str]] = [] - for part in output: + for part in parts: if not isinstance(part, dict): continue part_type = part.get("type") if part_type in ("input_text", "output_text", "text"): txt = part.get("text") - if isinstance(txt, str) and txt: + if isinstance(txt, str): text_acc.append(txt) - normalized_blocks.append({"type": "text", "text": txt}) + if txt: + normalized_blocks.append({"type": "text", "text": txt}) elif (media := decode_responses_media(part)) is not None: normalized_blocks.append(media) if any(b.get("type") != "text" for b in normalized_blocks): return normalized_blocks - if text_acc: + if any(text_acc) or len(text_acc) == len(parts): return "".join(text_acc) try: return json.dumps(output) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5c8350589b8..b7b93a20b4e 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars import json -from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterable, Mapping, Sequence +from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass from functools import partial @@ -42,6 +42,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( update_responses_tools_with_model_file_ids, ) from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header +from litellm.litellm_core_utils.reasoning_effort_capability import ( + deployment_reasoning_effort_levels, + model_group_from_metadata, + resolve_deployment_reasoning_effort, +) from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.common_utils import get_openai_credentials @@ -63,11 +68,6 @@ from litellm.responses.mcp.round_billing import ( ) from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.router_utils.reasoning_effort_capability import ( - deployment_reasoning_effort_levels, - model_group_from_metadata, - resolve_deployment_reasoning_effort, -) from litellm.types.llms.openai import ( PromptObject, Reasoning, @@ -81,7 +81,7 @@ from litellm.types.llms.openai import ( # Handle ResponseText import with fallback if TYPE_CHECKING: - from litellm.responses.streaming_iterator import ResponsesWsFrameGuard + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard, ResponsesWsTurnAdmission from litellm.types.llms.openai import ResponseText else: ResponseText = str # Fallback for ResponseText import @@ -2867,7 +2867,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, - turn_admission: Callable[[], Awaitable[None]] | None = None, + turn_admission: "ResponsesWsTurnAdmission | None" = None, responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs, ) -> Exception | None: diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index cd072825ef5..1e07b768f62 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -1,15 +1,18 @@ """Helpers for handling MCP-aware `/chat/completions` requests.""" +import asyncio import logging from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +import anyio from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._internal_context import internal_sub_call from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.chat_stream_round_billing import ChatStreamRoundBilling, engaged_round_billing from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) @@ -19,7 +22,7 @@ from litellm.responses.mcp.round_billing import ( first_round_billed_on_failure, summed_cost_breakdown, ) -from litellm.types.utils import Message, ModelResponse +from litellm.types.utils import Message, ModelResponse, Usage from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: @@ -258,10 +261,20 @@ async def acompletion_with_mcp( self.follow_up_stream = None self.follow_up_iterator = None self.follow_up_exhausted = False + self._round_billing: ChatStreamRoundBilling | None = engaged_round_billing(initial_stream.logging_obj) + self._unsent_usage_chunk: ModelResponseStream | None = None def __aiter__(self): return self + def __del__(self) -> None: + self.release_abandoned() + + def release_abandoned(self) -> None: + round_billing: Final[object] = getattr(self, "_round_billing", None) + if isinstance(round_billing, ChatStreamRoundBilling): + round_billing.release_abandoned() + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( @@ -321,18 +334,54 @@ async def acompletion_with_mcp( return chunk - async def _drain_inner_stream(self): + async def _drain_inner_stream(self) -> tuple[ModelResponseStream, ...]: try: - while True: - await self._stream_iterator.__anext__() - except StopAsyncIteration: - pass + return tuple([chunk async for chunk in initial_stream]) except Exception: logging.getLogger("LiteLLM").exception( "Error draining inner MCP stream after final chunk; spend logging may be incomplete" ) + return () - async def __anext__(self): + async def _complete_first_round(self) -> ModelResponseStream | None: + drained: Final = await self._drain_inner_stream() + if self._round_billing is not None: + self._round_billing.complete_first_round() + usage_chunks: Final = tuple( + chunk for chunk in drained if isinstance(getattr(chunk, "usage", None), Usage) + ) + return usage_chunks[-1] if usage_chunks else None + + async def _follow_up_or_release_first_round(self, usage_chunk: ModelResponseStream | None) -> None: + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + if self.follow_up_stream is not None: + return + self._unsent_usage_chunk = usage_chunk + if self._round_billing is not None: + self._round_billing.finish() + + async def __anext__(self) -> ModelResponseStream: + round_billing: Final = self._round_billing + if round_billing is None: + return await self._next_chunk() + with round_billing.armed_while_reading(): + try: + return await self._next_chunk() + except StopAsyncIteration: + usage_chunk: Final = round_billing.finish() + if usage_chunk is None: + raise + return usage_chunk + except asyncio.CancelledError: + with anyio.CancelScope(shield=True): + await round_billing.bill_in_flight_round() + raise + except Exception: + round_billing.fail() + raise + + async def _next_chunk(self) -> ModelResponseStream: # Phase 1: Collect and yield initial stream chunks if not self.stream_exhausted: # Get the iterator from the stream wrapper @@ -362,29 +411,22 @@ async def acompletion_with_mcp( if is_final: self.stream_exhausted = True + usage_chunk: Final = await self._complete_first_round() await self._process_tool_calls() chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - # Drain inner stream so CustomStreamWrapper fires its - # end-of-stream handler (dispatch_success_handlers → - # _ProxyDBLogger → LiteLLM_SpendLogs). The CSW may - # yield one usage chunk before raising StopAsyncIteration. - await self._drain_inner_stream() + await self._follow_up_or_release_first_round(usage_chunk) return chunk except StopAsyncIteration: self.stream_exhausted = True + drained_usage_chunk: Final = await self._complete_first_round() # Process tool calls after stream is exhausted await self._process_tool_calls() # If we have chunks, yield the final one with metadata if self.collected_chunks: final_chunk = self.collected_chunks[-1] final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) - # If we have tool results, prepare follow-up call - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - await self._drain_inner_stream() + await self._follow_up_or_release_first_round(drained_usage_chunk) return final_chunk # Phase 2: Yield follow-up stream chunks if available @@ -400,7 +442,6 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", chunk) - return chunk except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger @@ -408,6 +449,10 @@ async def acompletion_with_mcp( verbose_logger.debug("Follow-up stream exhausted") # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration raise StopAsyncIteration + client_chunk: Final = ( + chunk if self._round_billing is None else self._round_billing.client_chunk(chunk) + ) + return await self._next_chunk() if client_chunk is None else client_chunk # If we're here and follow_up_stream is None but we expected it, log a warning if ( @@ -420,7 +465,11 @@ async def acompletion_with_mcp( verbose_logger.warning("Follow-up stream was not created despite having tool results") - raise StopAsyncIteration + unsent_usage_chunk: Final = self._unsent_usage_chunk + if unsent_usage_chunk is None: + raise StopAsyncIteration + self._unsent_usage_chunk = None + return unsent_usage_chunk async def _process_tool_calls(self): """Process tool calls after streaming completes.""" @@ -488,10 +537,16 @@ async def acompletion_with_mcp( # This ensures the patch works correctly in tests import litellm + earlier_rounds: Final = None if self._round_billing is None else self._round_billing.bank_first_round() follow_up_response: Final = await litellm.acompletion(**follow_up_call_args) # Ensure follow-up response is a CustomStreamWrapper if isinstance(follow_up_response, CustomStreamWrapper): + if earlier_rounds is not None: + follow_up_response.bill_failure_after_earlier_rounds(earlier_rounds) + follow_up_response.messages = follow_up_messages + if self._round_billing is not None: + self._round_billing.follow_up_started(follow_up_response.partial_response) self.follow_up_stream = follow_up_response from litellm._logging import verbose_logger @@ -548,11 +603,13 @@ async def acompletion_with_mcp( def __aiter__(self): return self._custom_iterator + async def aclose(self) -> None: + self._custom_iterator.release_abandoned() + await super().aclose() + def __iter__(self): # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: - import asyncio - try: self._sync_loop = asyncio.get_event_loop() except RuntimeError: diff --git a/litellm/responses/mcp/chat_stream_round_billing.py b/litellm/responses/mcp/chat_stream_round_billing.py new file mode 100644 index 00000000000..370de663606 --- /dev/null +++ b/litellm/responses/mcp/chat_stream_round_billing.py @@ -0,0 +1,281 @@ +import asyncio +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.litellm_core_utils.asyncify import asyncify, is_event_loop_running +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + HELD_ROUND_SPEND, + IN_FLIGHT_ROUND, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + DeferredStreamRelease, + parked_success, +) +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown +from litellm.responses.streaming_iterator import spawn_logging_task +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage + +Parked: TypeAlias = tuple[object, ...] +PartialAnswer: TypeAlias = Callable[[], object] + + +@dataclass(frozen=True, slots=True) +class _AssembledRound: + response: ModelResponse + usage: Usage + cache_hit: bool + parked: Parked + + +@dataclass(frozen=True, slots=True) +class _PricedRound: + response: ModelResponse + usage: Usage + price: UncachedPrice + cache_hit: bool + + @property + def spend(self) -> BilledRoundsSpend: + return BilledRoundsSpend(self.usage, self.price.cost or 0.0, self.price.cost_breakdown) + + +@dataclass(frozen=True, slots=True) +class _HeldUsageChunk: + chunk: ModelResponseStream + usage: Usage + + +@dataclass(frozen=True, slots=True) +class _FirstRound: + priced: _PricedRound + parked: Parked + logged_alone: Parked + banked: LiteLLMLoggingObj | None = None + usage_chunk: _HeldUsageChunk | None = None + + def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None: + if self.usage_chunk is None: + return None + usage: Final = ( + self.priced.usage.model_copy(update=MappingProxyType({"cost": self.priced.price.cost})) + if billed is None + else merged_chat_usage(self.priced.usage, self.usage_chunk.usage).model_copy( + update=MappingProxyType({"cost": _cost(billed)}) + ) + ) + return self.usage_chunk.chunk.model_copy( + update=MappingProxyType({"id": self.priced.response.id, "usage": usage}) + ) + + +def _cost(response: ModelResponse) -> object: + hidden_params: Final = get_hidden_params(response) + return None if hidden_params is None else hidden_params.get("response_cost") + + +def _assembled_round(parked: Parked | None) -> _AssembledRound | None: + match parked: + case (ModelResponse() as response, cache_hit) if isinstance(usage := getattr(response, "usage", None), Usage): + return _AssembledRound(response, usage, cache_hit is True, parked) + case _: + return None + + +def engaged_round_billing(logging_obj: object) -> "ChatStreamRoundBilling | None": + return ChatStreamRoundBilling(logging_obj) if isinstance(logging_obj, LiteLLMLoggingObj) else None + + +class ChatStreamRoundBilling: + def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: + self._logging_obj: Final = logging_obj + self._first: _FirstRound | None = None + self._partial_follow_up: PartialAnswer | None = None + self._reading_loop: asyncio.AbstractEventLoop | None = None + + @contextmanager + def armed_while_reading(self) -> Generator[None]: + armed: Final = getattr(self._logging_obj, DISPATCH_RELEASE, None) is None + if armed: + setattr(self._logging_obj, DISPATCH_RELEASE, self._log_success) + self._reading_loop = asyncio.get_running_loop() + try: + yield + finally: + if armed and self._peek_parked() is None: + self._disarm() + + def complete_first_round(self) -> None: + assembled: Final = _assembled_round(self._peek_parked()) + if assembled is not None: + self._first = _FirstRound(self._priced(assembled), assembled.parked, assembled.parked) + setattr(self._logging_obj, HELD_ROUND_SPEND, self._held_round_one_spend) + + def bank_first_round(self) -> BilledRoundsSpend | None: + self._bank_apart_from_the_follow_up() + return None if self._first is None else self._first.priced.spend + + def client_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: + usage: Final[object] = getattr(chunk, "usage", None) + if self._first is None or not isinstance(usage, Usage): + return chunk + self._first = replace(self._first, usage_chunk=_HeldUsageChunk(chunk, usage)) + return None + + def finish(self) -> ModelResponseStream | None: + first: Final = self._first + billed: Final = self._release_rounds() + return None if first is None else first.client_usage_chunk(billed) + + def fail(self) -> None: + self._bank_apart_from_the_follow_up() + first: Final = self._first + if first is None or first.banked is None or not self._holds_round_one_alone(first): + self.finish() + return + self._first = None + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._log_through_the_release((*first.parked, first.banked.guardrail_scoped_copy())) + + def follow_up_started(self, partial_follow_up: PartialAnswer) -> None: + self._partial_follow_up = partial_follow_up + setattr(self._logging_obj, IN_FLIGHT_ROUND, self.bill_in_flight_round) + + async def bill_in_flight_round(self) -> None: + partial_follow_up: Final = self._partial_follow_up + if partial_follow_up is None or self._first_round_held_alone() is None: + return + final_round: Final = _assembled_round((await asyncify(partial_follow_up)(), False)) + first: Final = self._first_round_held_alone() + if final_round is None or first is None: + return + self._first = None + setattr( + self._logging_obj, PARKED_DISPATCH, self._billed_for_both_rounds(first.priced, self._priced(final_round)) + ) + + def release_abandoned(self) -> None: + if not self._owns_release(): + return + if is_event_loop_running() or self._idle_reading_loop() is not None: + self.finish() + return + setattr(self._logging_obj, PARKED_DISPATCH, None) + self._first = None + self._disarm() + + def _bank_apart_from_the_follow_up(self) -> None: + first: Final = self._first + if first is None or first.banked is not None or self._peek_parked() is not first.parked: + return + banked: Final = self._logging_obj.attempt_scoped_copy() + if not self._owns_release(): + self._first = replace(first, banked=banked) + return + logged_alone: Final = (*first.parked, banked) + self._first = replace(first, logged_alone=logged_alone, banked=banked) + setattr(self._logging_obj, PARKED_DISPATCH, logged_alone) + + def _release_rounds(self) -> ModelResponse | None: + parked: Final = self._take_parked() + first: Final = self._first + self._first = None + if parked is None: + self._disarm() + return None + if first is None or parked is first.logged_alone: + self._release(parked) + return None + final_round: Final = _assembled_round(parked) + if final_round is None: + self._release(first.logged_alone) + return None + billed, cache_hit = self._billed_for_both_rounds(first.priced, self._priced(final_round)) + self._release((billed, cache_hit)) + return billed + + def _billed_for_both_rounds(self, first: _PricedRound, final: _PricedRound) -> tuple[ModelResponse, bool]: + first_breakdown: Final = first.price.cost_breakdown + final_breakdown: Final = final.price.cost_breakdown + rounds: Final = billed_for_every_round(first.response, first_breakdown, final.response, final_breakdown) + usage: Final[object] = getattr(rounds, "usage", None) + charged_usage: Final = ( + usage.model_copy(update=MappingProxyType({"cost": _cost(rounds)})) if isinstance(usage, Usage) else usage + ) + billed: Final = rounds.model_copy(update=MappingProxyType({"id": first.response.id, "usage": charged_usage})) + cache_hit: Final = first.cache_hit and final.cache_hit + set_hidden_params(billed, {**(get_hidden_params(rounds) or {}), "cache_hit": cache_hit}) + self._logging_obj.cost_breakdown = summed_cost_breakdown( + first.response, first_breakdown, final.response, final_breakdown + ) + return billed, cache_hit + + def _priced(self, assembled: _AssembledRound) -> _PricedRound: + served: Final = assembled.cache_hit or served_from_cache(assembled.response) + price: Final = SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response) + priced: Final = assembled.response.model_copy() + set_hidden_params( + priced, + {**(get_hidden_params(assembled.response) or {}), "response_cost": price.cost, "cache_hit": served}, + ) + return _PricedRound(priced, assembled.usage, price, served) + + def _peek_parked(self) -> Parked | None: + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) + return parked if isinstance(parked, tuple) else None + + def _take_parked(self) -> Parked | None: + parked: Final = self._peek_parked() + setattr(self._logging_obj, PARKED_DISPATCH, None) + return parked + + def _first_round_held_alone(self) -> _FirstRound | None: + first: Final = self._first + return first if first is not None and self._holds_round_one_alone(first) else None + + def _held_round_one_spend(self) -> BilledRoundsSpend | None: + first: Final = self._first_round_held_alone() + return None if first is None else first.priced.spend + + def _holds_round_one_alone(self, first: _FirstRound) -> bool: + parked: Final = self._peek_parked() + return parked is first.parked or parked is first.logged_alone + + def _owns_release(self) -> bool: + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) + return release == self._log_success + + def _release(self, parked: Parked) -> None: + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) + return + self._log_through_the_release(parked) + + def _log_through_the_release(self, parked: Parked) -> None: + armed: Final[DeferredStreamRelease | None] = getattr(self._logging_obj, DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) + logged: Final = (self._log_success if armed is None else armed)(*parked) + idle_reading_loop: Final = None if is_event_loop_running() else self._idle_reading_loop() + if idle_reading_loop is not None: + idle_reading_loop.run_until_complete(logged) + return + spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log") + + def _idle_reading_loop(self) -> asyncio.AbstractEventLoop | None: + loop: Final = self._reading_loop + return loop if loop is not None and not loop.is_closed() and not loop.is_running() else None + + def _disarm(self) -> None: + if self._owns_release(): + setattr(self._logging_obj, DISPATCH_RELEASE, None) + + async def _log_success(self, *parked: object) -> None: + success: Final = parked_success(parked) + if success is not None: + await success.dispatch(self._logging_obj) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index c88a8c1cc28..1104aee353e 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -23,6 +23,7 @@ from litellm.responses.mcp.tool_precheck import ( LITELLM_PROXY_MCP_SERVER_URL_PREFIX, may_name_litellm_mcp_gateway, names_gateway_explicitly, + names_whole_gateway, proxy_path_mcp_name, ) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator @@ -182,8 +183,8 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _apply_toolset_permissions( - resolved_toolset_ids: list[str], - resolved_mcp_servers: list[str], + resolved_toolset_ids: Sequence[str], + resolved_mcp_servers: Sequence[str], user_api_key_auth: "UserAPIKeyAuth", ) -> "UserAPIKeyAuth": """Apply resolved toolset permissions to user_api_key_auth and return updated auth.""" @@ -195,7 +196,7 @@ class LiteLLM_Proxy_MCP_Handler: ) tool_permissions: Final = await global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=resolved_toolset_ids + toolset_ids=list(resolved_toolset_ids) ) all_server_ids: Final = list(set(tool_permissions.keys()) | set(resolved_mcp_servers)) existing_op: Final = user_api_key_auth.object_permission @@ -224,6 +225,62 @@ class LiteLLM_Proxy_MCP_Handler: verbose_logger.debug("Could not apply toolset permissions: %s", _e) return user_api_key_auth + @staticmethod + async def _resolve_proxy_mcp_name( + name: str, + user_api_key_auth: "UserAPIKeyAuth | None", + granted_toolsets: "GrantedToolsetIds | None", + ) -> tuple[Literal["server", "toolset", "denied_toolset"], str]: + from litellm.proxy.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + if global_mcp_server_manager.get_mcp_server_by_name(name): + return ("server", name) + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return ("server", name) + toolset: Final = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) + if toolset is None: + return ("server", name) + from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import ( + has_team_unbound_mcp_admin_view, + ) + from litellm.proxy.mcp_server.ui_session_utils import ( + granted_toolset_ids, + ) + + denied: Final = ( + user_api_key_auth is not None + and not has_team_unbound_mcp_admin_view(user_api_key_auth) + and toolset.toolset_id not in await (granted_toolsets or granted_toolset_ids)(user_api_key_auth) + ) + except Exception as _e: + verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, _e) + return ("server", name) + if not denied: + return ("toolset", toolset.toolset_id) + if await LiteLLM_Proxy_MCP_Handler._holds_access_group_named(name, user_api_key_auth): + return ("server", name) + verbose_logger.debug("Key does not have access to toolset '%s', denying it.", name) + return ("denied_toolset", name) + + @staticmethod + async def _holds_access_group_named(name: str, user_api_key_auth: "UserAPIKeyAuth | None") -> bool: + from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + group_server_ids: Final = await MCPRequestHandler.get_mcp_servers_from_access_groups([name]) + return bool(group_server_ids) and not frozenset(group_server_ids).isdisjoint( + await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) + ) + @staticmethod @catalog_operation(global_manager) async def _get_mcp_tools_from_manager( @@ -256,6 +313,7 @@ class LiteLLM_Proxy_MCP_Handler: from litellm.proxy.mcp_server.server import ( _get_allowed_mcp_servers_from_mcp_server_names, _get_tools_from_mcp_servers, + raise_generic_mcp_access_denial, ) mcp_servers: Final = [ @@ -268,36 +326,15 @@ class LiteLLM_Proxy_MCP_Handler: # Resolve toolset names: collect all toolset IDs first, then apply their # combined permissions in a single pass so multiple toolsets are unioned # rather than the last one overwriting the others. - resolved_mcp_servers: Final[list[str]] = [] - resolved_toolset_ids: Final[list[str]] = [] - for name in mcp_servers: - if not global_mcp_server_manager.get_mcp_server_by_name(name): - try: - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is not None: - toolset = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) - if toolset is not None: - if user_api_key_auth is not None: - from litellm.proxy.management_endpoints.common_utils import ( - user_api_key_has_admin_view, - ) - from litellm.proxy.mcp_server.ui_session_utils import ( - granted_toolset_ids, - ) - - if not user_api_key_has_admin_view(user_api_key_auth) and toolset.toolset_id not in ( - await (granted_toolsets or granted_toolset_ids)(user_api_key_auth) - ): - verbose_logger.debug("Key does not have access to toolset '%s', skipping.", name) - continue - resolved_toolset_ids.append(toolset.toolset_id) - # Don't add to resolved_mcp_servers — toolset scope - # restricts via object_permission, not server name filter. - continue - except Exception as _e: - verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, _e) - resolved_mcp_servers.append(name) + resolutions: Final = tuple( + [ + await LiteLLM_Proxy_MCP_Handler._resolve_proxy_mcp_name(name, user_api_key_auth, granted_toolsets) + for name in mcp_servers + ] + ) + resolved_mcp_servers: Final = tuple(value for kind, value in resolutions if kind == "server") + resolved_toolset_ids: Final = tuple(value for kind, value in resolutions if kind == "toolset") + denied_toolset_names: Final = tuple(value for kind, value in resolutions if kind == "denied_toolset") # Apply all resolved toolsets at once (union), avoiding permission overwrite. if resolved_toolset_ids and user_api_key_auth is not None: @@ -311,7 +348,10 @@ class LiteLLM_Proxy_MCP_Handler: # full union (toolset server IDs + direct server names). Passing a name-based # filter here would exclude those toolset server IDs (which are UUIDs, not # names), so use None and let the auth object's mcp_servers do the filtering. - effective_server_filter: Final = None if resolved_toolset_ids else (resolved_mcp_servers or None) + asks_whole_gateway: Final = any(map(names_whole_gateway, mcp_tools_with_litellm_proxy or ())) + if denied_toolset_names and not (asks_whole_gateway or resolved_mcp_servers or resolved_toolset_ids): + raise_generic_mcp_access_denial(denied_toolset_names) + effective_server_filter: Final = None if resolved_toolset_ids else (list(resolved_mcp_servers) or None) listing: Final = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, diff --git a/litellm/responses/mcp/parked_stream_dispatch.py b/litellm/responses/mcp/parked_stream_dispatch.py new file mode 100644 index 00000000000..792363f0bab --- /dev/null +++ b/litellm/responses/mcp/parked_stream_dispatch.py @@ -0,0 +1,72 @@ +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final, Protocol + +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +IN_FLIGHT_ROUND: Final = "_bill_in_flight_round_on_disconnect" +HELD_ROUND_SPEND: Final = "_held_round_spend_for_failure" +SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +class DeferredStreamRelease(Protocol): + def __call__(self, *parked: object) -> Coroutine[object, object, None]: ... + + +class InFlightRoundBiller(Protocol): + def __call__(self) -> Coroutine[object, object, None]: ... + + +class HeldRoundSpend(Protocol): + def __call__(self) -> BilledRoundsSpend | None: ... + + +def success_dispatcher(logger: LiteLLMLoggingObj) -> DispatchesSuccessHandlers: + return logger + + +@dataclass(frozen=True, slots=True) +class ParkedSuccess: + assembled: object + cache_hit: object + logged_on: LiteLLMLoggingObj | None = None + + async def dispatch(self, request_logger: LiteLLMLoggingObj) -> None: + logger: Final = request_logger if self.logged_on is None else self.logged_on + await success_dispatcher(logger).dispatch_success_handlers( + self.assembled, cache_hit=self.cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + + +def in_flight_round_biller(logging_obj: object) -> InFlightRoundBiller | None: + biller: Final[InFlightRoundBiller | None] = getattr(logging_obj, IN_FLIGHT_ROUND, None) + return biller if isinstance(logging_obj, LiteLLMLoggingObj) and callable(biller) else None + + +def held_round_spend(logging_obj: object) -> BilledRoundsSpend | None: + held: Final[HeldRoundSpend | None] = getattr(logging_obj, HELD_ROUND_SPEND, None) + return held() if isinstance(logging_obj, LiteLLMLoggingObj) and callable(held) else None + + +def parked_success(parked: tuple[object, ...]) -> ParkedSuccess | None: + match parked: + case (assembled, cache_hit): + return ParkedSuccess(assembled, cache_hit) + case (assembled, cache_hit, LiteLLMLoggingObj() as logged_on): + return ParkedSuccess(assembled, cache_hit, logged_on) + case _: + return None diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py index 443147230bf..5cb256a8ac0 100644 --- a/litellm/responses/mcp/stream_round_billing.py +++ b/litellm/responses/mcp/stream_round_billing.py @@ -2,11 +2,18 @@ import asyncio from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass from datetime import datetime -from typing import Final, Protocol, TypeAlias +from typing import Final, TypeAlias from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.parked_stream_dispatch import ( + DISPATCH_RELEASE, + PARKED_DISPATCH, + SERVED_FROM_CACHE, + parked_success, + success_dispatcher, +) from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task from litellm.responses.utils import ResponseAPILoggingUtils @@ -16,25 +23,6 @@ from litellm.types.utils import CostBreakdown, ModelResponse, Usage Parked: TypeAlias = tuple[object, ...] RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None] -_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" -_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" -_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) - - -class _DispatchesSuccessHandlers(Protocol): - async def dispatch_success_handlers( - self, - result: object = None, - start_time: object = None, - end_time: object = None, - cache_hit: object = None, - prefer_async_handlers: bool = False, - ) -> None: ... - - -def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: - return logger - def _event_loop_running() -> bool: try: @@ -198,8 +186,8 @@ class StreamedRoundBilling: self._billed: _BilledRounds | None = None self._banked: Parked | None = None self._banked_row: _RoundRow | None = None - if getattr(logging_obj, _DISPATCH_RELEASE, None) is None: - setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked) + if getattr(logging_obj, DISPATCH_RELEASE, None) is None: + setattr(logging_obj, DISPATCH_RELEASE, self._run_parked) def complete_round( self, @@ -286,7 +274,7 @@ class StreamedRoundBilling: def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: if served_from_cache(response): - return _SERVED_FROM_CACHE + return SERVED_FROM_CACHE match round_parked: case (ModelResponse() as assembled, bool()): return self._logging_obj.uncached_price(assembled) @@ -297,7 +285,7 @@ class StreamedRoundBilling: self._banked = parked self._banked_row = row if parked is not None: - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + setattr(self._logging_obj, PARKED_DISPATCH, parked) def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: row: Final = self._banked_row @@ -320,7 +308,7 @@ class StreamedRoundBilling: @staticmethod async def _log_success(success: _SuccessLog) -> None: success.carry_bridged_usage() - await _success_dispatcher(success.logger).dispatch_success_handlers( + await success_dispatcher(success.logger).dispatch_success_handlers( success.row.logged, start_time=success.logger.start_time, end_time=success.row.end_time, @@ -329,10 +317,10 @@ class StreamedRoundBilling: ) def _take_parked(self) -> Parked | None: - parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None) if not isinstance(parked, tuple): return None - setattr(self._logging_obj, _PARKED_DISPATCH, None) + setattr(self._logging_obj, PARKED_DISPATCH, None) return parked async def _run_parked(self, *parked: object) -> None: @@ -341,29 +329,27 @@ class StreamedRoundBilling: await dispatch case (dispatch,) if asyncio.iscoroutine(dispatch): await dispatch - case (assembled, cache_hit): - await _success_dispatcher(self._logging_obj).dispatch_success_handlers( - assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True - ) + case _ if (success := parked_success(parked)) is not None: + await success.dispatch(self._logging_obj) case _: return def _owns_release(self) -> bool: - release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None) return release == self._run_parked def _release(self, parked: Parked) -> None: - if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): - setattr(self._logging_obj, _PARKED_DISPATCH, parked) + if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, PARKED_DISPATCH, parked) return - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) spawn_logging_task( asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging" ) def _disarm(self) -> None: if self._owns_release(): - setattr(self._logging_obj, _DISPATCH_RELEASE, None) + setattr(self._logging_obj, DISPATCH_RELEASE, None) @staticmethod def _close(parked: Parked | None) -> None: diff --git a/litellm/responses/mcp/tool_precheck.py b/litellm/responses/mcp/tool_precheck.py index a29e6895f9a..a67d81e70ae 100644 --- a/litellm/responses/mcp/tool_precheck.py +++ b/litellm/responses/mcp/tool_precheck.py @@ -22,6 +22,13 @@ def names_gateway_explicitly(tool: object) -> bool: return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) +def names_whole_gateway(tool: object) -> bool: + server_url: Final = _mcp_server_url(tool) or "" + return server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL) and not server_url.startswith( + LITELLM_PROXY_MCP_SERVER_URL_PREFIX + ) + + def proxy_path_mcp_name(tool: object) -> str | None: server_url: Final = _mcp_server_url(tool) match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index f0bb2201162..a8e55a5e94d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -312,7 +312,7 @@ def map_stream_error_to_exception( error_body: Final = {"message": error_message, "type": error_type, "code": error_code} provider_exception: Final = BaseLLMException( status_code=status_code, - message=f"Error code: {status_code} - {{'error': {error_body}}}", + message=f"Error code: {status_code} - {json.dumps({'error': error_body})}", body=error_body, headers=headers, ) @@ -2145,18 +2145,32 @@ async def _enforce_frame_project_quota( ) -async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None, raw_message: str) -> Exception | None: +@dataclass(frozen=True, slots=True) +class _AdmittedTurn: + reservation: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + + +@dataclass(frozen=True, slots=True) +class _RefusedTurn: + rejection: Exception + + +ResponsesWsTurnAdmission: TypeAlias = Callable[[], Awaitable[Mapping[str, object]]] + + +async def _admit_frame_turn( + turn_admission: ResponsesWsTurnAdmission | None, raw_message: str +) -> _AdmittedTurn | _RefusedTurn: frame: Final = _json_object_or_none(raw_message) if turn_admission is None or frame is None or frame.get("type") != "response.create": - return None + return _AdmittedTurn() try: - await turn_admission() + return _AdmittedTurn(await turn_admission()) except RateLimitError as refusal: - return refusal + return _RefusedTurn(refusal) except Exception as e: verbose_logger.exception("Responses WS turn admission failed: %s", e) - return e - return None + return _RefusedTurn(e) def _turn_refusal_frame(refusal: RateLimitError) -> str: @@ -2191,6 +2205,7 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) _PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) +_NO_IO_TOKEN_RESERVATION: Final[Mapping[str, object]] = MappingProxyType({}) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" @@ -2262,6 +2277,29 @@ class GuardedResponsesWsFrame: ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] +class ResponsesWsFrameRejection(ValueError): + def __init__(self, message: str, provider_specific_fields: Mapping[str, object] | None = None) -> None: + super().__init__(message) + self.provider_specific_fields: Final = provider_specific_fields + + +def _frame_rejection_event(rejection: Exception) -> str: + provider_fields: Final = ( + rejection.provider_specific_fields if isinstance(rejection, ResponsesWsFrameRejection) else None + ) + return json.dumps( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": str(rejection), + **({"provider_specific_fields": dict(provider_fields)} if provider_fields else {}), + }, + }, + default=str, + ) + + async def _reject_guarded_turn( websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception ) -> None: @@ -2318,6 +2356,17 @@ def _set_turn_guardrail_information( _replace_guardrail_information(litellm_params.get(key), guardrail_information) +def _set_turn_io_token_reservation(turn_logging_obj: LiteLLMLoggingObj, reservation: Mapping[str, object]) -> None: + litellm_params: Final = _json_object_or_none( + turn_logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if litellm_params is None or not reservation: + return + for bucket in (litellm_params.get(key) for key in _PER_TURN_METADATA_KEYS): + if _is_json_object(bucket): + bucket.update(reservation) + + def _replace_guardrail_information(bucket: object, guardrail_information: tuple[Mapping[str, object], ...]) -> None: if not _is_json_object(bucket): return @@ -2458,7 +2507,7 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, - turn_admission: Callable[[], Awaitable[None]] | None = None, + turn_admission: ResponsesWsTurnAdmission | None = None, frame_guard: ResponsesWsFrameGuard | None = None, ): self.websocket = websocket @@ -2480,10 +2529,13 @@ class ResponsesWebSocketStreaming: # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults - self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission + self.turn_admission: ResponsesWsTurnAdmission | None = turn_admission def _start_turn( - self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = () + self, + client_frame: str | None, + guardrail_information: tuple[Mapping[str, object], ...] = (), + io_token_reservation: Mapping[str, object] = _NO_IO_TOKEN_RESERVATION, ) -> None: start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time turn_logging_obj: Final = self.logging_obj.copy_for_new_call( @@ -2491,6 +2543,7 @@ class ResponsesWebSocketStreaming: ) self._detach_budget_reservation() _set_turn_guardrail_information(turn_logging_obj, guardrail_information) + _set_turn_io_token_reservation(turn_logging_obj, io_token_reservation) turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) turn_logging_obj.pre_call(input=client_frame, api_key="") self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) @@ -2510,11 +2563,14 @@ class ResponsesWebSocketStreaming: } def _record_client_frame( - self, client_frame: str, guardrail_information: tuple[Mapping[str, object], ...] = () + self, + client_frame: str, + guardrail_information: tuple[Mapping[str, object], ...] = (), + io_token_reservation: Mapping[str, object] = _NO_IO_TOKEN_RESERVATION, ) -> None: frame: Final = _json_object_or_none(client_frame) if frame is not None and frame.get("type") == "response.create": - self._start_turn(client_frame, guardrail_information) + self._start_turn(client_frame, guardrail_information, io_token_reservation) async def _guard_frame(self, message: str) -> GuardedResponsesWsFrame | None: """Run the proxy's pre_call guardrails on one ``response.create`` frame. @@ -2529,19 +2585,23 @@ class ResponsesWebSocketStreaming: return await self.frame_guard(message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error_event("invalid_request_error", str(e)) + with contextlib.suppress(Exception): + await self.websocket.send_text(_frame_rejection_event(e)) return None async def _forward_guarded_frame( - self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + self, message: str, turn_admission: ResponsesWsTurnAdmission | None = None ) -> None: guarded: Final = await self._guard_frame(message) - if guarded is None or not await self._admit_or_reject_turn(guarded, message, turn_admission): + if guarded is None: + return + admitted: Final = await self._admit_or_reject_turn(guarded, message, turn_admission) + if admitted is None: return if guarded.pii_tokens: self._remember_pii_tokens(guarded.pii_tokens) masked: Final = await self._mask_response_create(guarded.frame) - self._record_client_frame(masked, guarded.guardrail_information) + self._record_client_frame(masked, guarded.guardrail_information, admitted.reservation) await self.backend_ws.send(masked) def _remember_pii_tokens(self, pii_tokens: Mapping[str, str]) -> None: @@ -2941,15 +3001,17 @@ class ResponsesWebSocketStreaming: return True async def _admit_or_reject_turn( - self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: Callable[[], Awaitable[None]] | None - ) -> bool: - rejection: Final = await _admit_frame_turn(turn_admission, message) - if rejection is None: - return True - if self._failure is None: - self._failure = rejection - await _reject_guarded_turn(self.websocket, guarded, rejection) - return False + self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: ResponsesWsTurnAdmission | None + ) -> _AdmittedTurn | None: + verdict: Final = await _admit_frame_turn(turn_admission, message) + match verdict: + case _AdmittedTurn(): + return verdict + case _RefusedTurn(rejection=rejection): + if self._failure is None: + self._failure = rejection + await _reject_guarded_turn(self.websocket, guarded, rejection) + return None async def _send_turn_refusal(self, refusal: RateLimitError) -> None: try: @@ -2957,14 +3019,6 @@ class ResponsesWebSocketStreaming: except Exception: # noqa: BLE001, S110 # client may already be gone pass - async def _send_error_event(self, error_type: str, message: str) -> None: - try: - await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) - ) - except Exception: # noqa: BLE001, S110 # client may already be gone - pass - async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: @@ -3052,7 +3106,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, - turn_admission: Callable[[], Awaitable[None]] | None = None, + turn_admission: ResponsesWsTurnAdmission | None = None, frame_guard: ResponsesWsFrameGuard | None = None, guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, **kwargs: object, @@ -3076,7 +3130,7 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () - self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission + self.turn_admission: ResponsesWsTurnAdmission | None = turn_admission self._failure: Exception | None = None # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} @@ -3378,7 +3432,7 @@ class ManagedResponsesWebSocketHandler: return await self.frame_guard(raw_message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error(str(e), error_type="invalid_request_error") + await self._send_frame(_frame_rejection_event(e)) return None @staticmethod @@ -3473,7 +3527,7 @@ class ManagedResponsesWebSocketHandler: # ------------------------------------------------------------------ async def _process_response_create( - self, raw_message: str, turn_admission: Callable[[], Awaitable[None]] | None = None + self, raw_message: str, turn_admission: ResponsesWsTurnAdmission | None = None ) -> None: """ Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``, @@ -3520,11 +3574,11 @@ class ManagedResponsesWebSocketHandler: guarded: Final = await self._guard_frame(raw_message) if guarded is None: return - rejection: Final = await _admit_frame_turn(turn_admission, raw_message) - if rejection is not None: + verdict: Final = await _admit_frame_turn(turn_admission, raw_message) + if isinstance(verdict, _RefusedTurn): if self._failure is None: - self._failure = rejection - await _reject_guarded_turn(self.websocket, guarded, rejection) + self._failure = verdict.rejection + await _reject_guarded_turn(self.websocket, guarded, verdict.rejection) return if guarded.pii_tokens: self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} @@ -3555,6 +3609,11 @@ class ManagedResponsesWebSocketHandler: ) self._update_proxy_request(call_kwargs, requested_model or self.model_group or model) call_kwargs.update(self.extra_kwargs) + turn_litellm_metadata: Final = call_kwargs.get("litellm_metadata") + call_kwargs["litellm_metadata"] = { + **(turn_litellm_metadata if _is_json_object(turn_litellm_metadata) else {}), + **verdict.reservation, + } try: terminal_event: Final = await self._stream_and_forward(model, call_kwargs) diff --git a/litellm/router.py b/litellm/router.py index 846cec8f05b..bc1a0f5ce57 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -106,6 +106,20 @@ from litellm.integrations.custom_guardrail import CustomGuardrail, is_guardrail_ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.routing import routing_decision_attributes from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span +from litellm.litellm_core_utils.access_windows import access_windows_config_error, filter_reserved_deployments +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, + add_retry_headers_to_response, + apply_quality_router_decision_headers, + apply_remaining_usage_headers, + apply_response_model_id, + complexity_router_decision_headers, + ensure_response_additional_headers, + get_hidden_params_dict, + prepare_response_for_header_attachment, + replace_complexity_router_headers, + response_total_token_count, +) from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.batch_utils import ( batch_create_routing, @@ -116,6 +130,10 @@ from litellm.litellm_core_utils.batch_utils import ( should_replace_model_in_jsonl, without_deployments_that_cannot_batch, ) +from litellm.litellm_core_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + is_clientside_credential, +) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, @@ -127,12 +145,17 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference from litellm.litellm_core_utils.dd_tracing import tracer +from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models from litellm.litellm_core_utils.get_llm_provider_logic import ( declared_authenticating_provider, inferred_provider, is_registered_custom_provider, resolve_model_and_provider_for_metadata, ) +from litellm.litellm_core_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, +) +from litellm.litellm_core_utils.get_retry_from_policy import resolve_retry_policy from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR as _HIDDEN_PARAMS_ATTR, ) @@ -142,6 +165,7 @@ from litellm.litellm_core_utils.hidden_params import ( ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter from litellm.litellm_core_utils.ptu_pricing import ( PTU_COST_ATTRIBUTION_ENV_VAR, declares_ptu, @@ -152,10 +176,52 @@ from litellm.litellm_core_utils.ptu_pricing import ( ptu_terms, zeroed_ptu_pricing, ) +from litellm.litellm_core_utils.ptu_shares import ( + PTUTeamCeiling, + filter_ptu_shared_deployments, + ptu_capacity_warning, + ptu_share_teams, + share_holder_request_kwargs, + team_ptu_ceiling, +) +from litellm.litellm_core_utils.reasoning_effort_capability import ( + deployment_is_catalog_mapped, + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, get_service_tier_timeout, ) +from litellm.litellm_core_utils.routing_common_utils import ( + MID_STREAM_CONTINUATION_KWARG, + as_retry_skipped_deployment_ids, + deployment_model_id_pin, + deployments_the_caller_may_use, + filter_team_based_models, + filter_web_search_deployments, + format_content_policy_fallback_hint, + format_content_policy_no_fallback_note, + format_fallback_outcome_message, + format_no_fallback_group_message, + get_request_team_id, + is_proxy_admin_request, + provider_for_generic_call, + request_user_api_key_auth, + resolve_model_group_alias, + team_may_use_deployment, + truncate_fallback_error_detail, + warn_on_ignored_model_capability_flags, + warn_on_provider_credential_mismatch, + with_admin_api_key_secret_reference, + without_model_capability_flags, + without_router_only_kwargs, +) +from litellm.litellm_core_utils.routing_groups import ( + apply_routing_group_priority, + parse_routing_groups, + validate_routing_strategy, +) from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, @@ -196,6 +262,7 @@ from litellm.router_api import ( aapply_changes, apply_changes, ) +from litellm.router_api import RoutingContext as RouterContext from litellm.router_api import deployment_id as _deployment_id from litellm.router_deprecated import DeprecatedRouterMembers from litellm.router_host import router_host @@ -241,49 +308,7 @@ from litellm.router_strategy.tag_based_routing import ( is_valid_deployment_tag, tags_with_serving_deployment_tags, ) -from litellm.router_utils.access_windows import access_windows_config_error, filter_reserved_deployments -from litellm.router_utils.add_retry_fallback_headers import ( - add_fallback_headers_to_response, - add_retry_headers_to_response, - apply_quality_router_decision_headers, - apply_remaining_usage_headers, - apply_response_model_id, - complexity_router_decision_headers, - ensure_response_additional_headers, - get_hidden_params_dict, - prepare_response_for_header_attachment, - replace_complexity_router_headers, - response_total_token_count, -) from litellm.router_utils.client_initalization_utils import InitalizeCachedClient, MaxParallelRequestsLimit -from litellm.router_utils.clientside_credential_handler import ( - get_dynamic_litellm_params, - is_clientside_credential, -) -from litellm.router_utils.common_utils import ( - MID_STREAM_CONTINUATION_KWARG, - as_retry_skipped_deployment_ids, - deployment_model_id_pin, - deployments_the_caller_may_use, - filter_team_based_models, - filter_web_search_deployments, - format_content_policy_fallback_hint, - format_content_policy_no_fallback_note, - format_fallback_outcome_message, - format_no_fallback_group_message, - get_request_team_id, - is_proxy_admin_request, - provider_for_generic_call, - request_user_api_key_auth, - resolve_model_group_alias, - team_may_use_deployment, - truncate_fallback_error_detail, - warn_on_ignored_model_capability_flags, - warn_on_provider_credential_mismatch, - with_admin_api_key_secret_reference, - without_model_capability_flags, - without_router_only_kwargs, -) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( DEFAULT_COOLDOWN_TIME_SECONDS, @@ -300,7 +325,6 @@ from litellm.router_utils.cooldown_handlers import ( mark_deployment_selection_failure, set_cooldown_deployments, ) -from litellm.router_utils.deployment_pricing import deployment_pricing_models from litellm.router_utils.embedding_batch_size_check import ( async_log_rejection_before_provider_call, clear_rejected_before_provider_call, @@ -345,10 +369,6 @@ from litellm.router_utils.fallback_event_handlers import ( routed_deployment_id, run_async_fallback, ) -from litellm.router_utils.get_retry_from_policy import ( - get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, -) -from litellm.router_utils.get_retry_from_policy import resolve_retry_policy from litellm.router_utils.handle_error import ( async_raise_no_deployment_exception, send_llm_exception_alert, @@ -386,6 +406,7 @@ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check impor EncryptedContentAffinityCheck, ) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + async_admit_io_token_turn, async_io_token_refund_failure, build_io_token_rate_limit_headers, deployment_has_io_token_limits, @@ -399,19 +420,6 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, ) -from litellm.router_utils.ptu_shares import ( - PTUTeamCeiling, - filter_ptu_shared_deployments, - ptu_capacity_warning, - ptu_share_teams, - share_holder_request_kwargs, - team_ptu_ceiling, -) -from litellm.router_utils.reasoning_effort_capability import ( - deployment_is_catalog_mapped, - intersect_supported_reasoning_efforts, - resolve_supported_reasoning_efforts, -) from litellm.router_utils.request_priority import ( InvalidPriority, parse_default_priority, @@ -427,11 +435,6 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, ) -from litellm.router_utils.routing_groups import ( - apply_routing_group_priority, - parse_routing_groups, - validate_routing_strategy, -) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage from litellm.scheduler import FlowItem, Scheduler @@ -511,6 +514,7 @@ from litellm.utils import ( ProviderConfigManager, Rules, TextCompletionResponse, + _restore_correlation_context_if_supported, # pyright: ignore[reportPrivateUsage] # same task-local cleanup as the SDK wrapper aresolve_api_key_secret_reference, function_setup, get_llm_provider, @@ -526,8 +530,6 @@ from litellm.utils import ( token_counter, ) -from .router_utils.pattern_match_deployments import PatternMatchRouter - if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -904,6 +906,23 @@ def _stream_pre_content_error( ) +def _context_window_stream_fallback_error( + error: "MidStreamFallbackError | litellm.ContextWindowExceededError", + model: object, + llm_provider: str, + has_output: bool, +) -> "MidStreamFallbackError": + if isinstance(error, litellm.ContextWindowExceededError): + if has_output: + raise error + return _stream_pre_content_error(error, model if isinstance(model, str) else "", llm_provider) + if isinstance(error.original_exception, litellm.ContextWindowExceededError) and ( + has_output or error.generated_content or error.emitted_disqualifying_content + ): + raise error.original_exception + return error + + def _deployment_num_retries(deployment: "Deployment | None") -> int | None: """The deployment's own num_retries litellm_param, an int or a digit string the way the config loader leaves it.""" configured: Final = getattr(deployment.litellm_params, "num_retries", None) if deployment is not None else None @@ -1600,6 +1619,7 @@ class Router(DeprecatedRouterMembers): router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}]) ``` """ + self._core_ctx: Final = RouterContext(self) self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments @@ -1878,7 +1898,7 @@ class Router(DeprecatedRouterMembers): if self._normalize_strategy(routing_strategy) == "lar1": from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - apply_lar1_routing_strategy(self, routing_strategy_args) + apply_lar1_routing_strategy(self._core_ctx, routing_strategy_args) else: self.routing_strategy_init( routing_strategy=routing_strategy, @@ -2254,7 +2274,7 @@ class Router(DeprecatedRouterMembers): if strategy == "lar1": from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - apply_lar1_routing_strategy(self, self.routing_strategy_args) + apply_lar1_routing_strategy(self._core_ctx, self.routing_strategy_args) return attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") @@ -2579,7 +2599,7 @@ class Router(DeprecatedRouterMembers): ) deployments: Final = self._core_get_model_list(model_name=model) if deployments: - RoutingPrefetch.arm(self, usage_selector, deployments) + RoutingPrefetch.arm(self._core_ctx, usage_selector, deployments) except Exception as e: # noqa: BLE001 # a prefetch is an optimisation, never a reason to fail the request verbose_router_logger.debug( "routing read prefetch not armed for %s: %s", _without_line_breaks(model), _without_line_breaks(e) @@ -3134,7 +3154,7 @@ class Router(DeprecatedRouterMembers): break if existing_ec_callback is not None: - existing_ec_callback.router = self + existing_ec_callback.ctx = self._core_ctx existing_ec_callback.enable_global_affinity = ( existing_ec_callback.enable_global_affinity or enable_global_affinity ) @@ -3142,7 +3162,7 @@ class Router(DeprecatedRouterMembers): ec_callback = existing_ec_callback else: ec_callback = EncryptedContentAffinityCheck( - router=self, + ctx=self._core_ctx, enable_global_affinity=enable_global_affinity, model_group_affinity_config=self.model_group_affinity_config, ) @@ -3692,7 +3712,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -3865,7 +3885,19 @@ class Router(DeprecatedRouterMembers): try: async for item in _achunks_with_first(model_response, first_chunk): yield item - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + bool(getattr(model_response, "response_uptil_now", "")) + or _stream_chunks_have_generated_content(model_response.chunks), + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) with anyio.CancelScope(shield=True): await close_model_response() await held_slot.aclose() @@ -3937,7 +3969,9 @@ class Router(DeprecatedRouterMembers): initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = await self.async_function_with_fallbacks_common_utils( - e=e, + e=e.original_exception + if isinstance(e.original_exception, litellm.ContextWindowExceededError) + else e, disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, @@ -4005,10 +4039,10 @@ class Router(DeprecatedRouterMembers): if fallback_response is not None and hasattr(fallback_response, "aclose"): try: await fallback_response.aclose() - except BaseException as e: + except BaseException as close_error: verbose_router_logger.debug( "stream_with_fallbacks: error closing fallback_response: %s", - e, + close_error, ) wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) @@ -4368,6 +4402,10 @@ class Router(DeprecatedRouterMembers): def __aiter__(self): return self + @property + def yielded_first_chunk(self) -> bool: + return self._yielded_first_chunk + async def __anext__(self): try: chunk: Final = await self._async_generator.__anext__() @@ -4395,6 +4433,7 @@ class Router(DeprecatedRouterMembers): # the very key that created the container (#30210). if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: self.completed_response = chunk + self._yielded_first_chunk = True return chunk async def aclose(self): @@ -4421,7 +4460,18 @@ class Router(DeprecatedRouterMembers): yield item for held_event in held_lifecycle_events: yield held_event - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + source_iterator.model or "", + source_iterator.custom_llm_provider or "", + wrapper.yielded_first_chunk, + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) if response_id_exposed: self._raise_released_stream_error(e, initial_kwargs) async with contextlib.aclosing( @@ -4592,7 +4642,11 @@ class Router(DeprecatedRouterMembers): # The content-policy dispatch branch matches on the trigger's own type, so a refusal's # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. fallback_trigger: Final[Exception] = ( - e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e + e.original_exception + if isinstance( + e.original_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError) + ) + else e ) fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success e=fallback_trigger, @@ -4678,7 +4732,19 @@ class Router(DeprecatedRouterMembers): try: for item in _chunks_with_first(model_response, first_chunk): yield item - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + bool(getattr(model_response, "response_uptil_now", "")) + or _stream_chunks_have_generated_content(model_response.chunks), + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) if fallbacks_disabled_for_request(initial_kwargs) or ( not e.is_pre_first_chunk and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) @@ -4703,7 +4769,9 @@ class Router(DeprecatedRouterMembers): router_self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = run_async_function( router_self.async_function_with_fallbacks_common_utils, - e=e, + e=e.original_exception + if isinstance(e.original_exception, litellm.ContextWindowExceededError) + else e, disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, @@ -4887,7 +4955,7 @@ class Router(DeprecatedRouterMembers): kwargs=kwargs, client_type="max_parallel_requests", ) - compacted_input: Final = await compact_to_fit(self, deployment, input_kwargs, "chat") + compacted_input: Final = await compact_to_fit(self._core_ctx, deployment, input_kwargs, "chat") async with contextlib.AsyncExitStack() as deployment_slot: if isinstance(max_parallel_requests_limit, MaxParallelRequestsLimit): deployment_slot.enter_context(max_parallel_requests_limit) @@ -5134,6 +5202,7 @@ class Router(DeprecatedRouterMembers): kwargs: dict, function_name: str | None = None, splits_stream_timeout: bool = False, + opens_responses_socket: bool = False, ) -> None: """ 5 jobs: @@ -5223,7 +5292,12 @@ class Router(DeprecatedRouterMembers): # here before it's wiped below, instead of relying on that attempt's # (possibly still-pending) failure event to do it. refund_stale_reservation_before_retry(self._core_cache, kwargs) - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=deployment_has_io_token_limits(deployment)) + set_io_token_rate_limit_request_kwargs( + kwargs, + store_in_context=deployment_has_io_token_limits(deployment), + opens_responses_socket=opens_responses_socket, + metadata_channel=metadata_variable_name, + ) kwargs[metadata_variable_name].setdefault( ROUTING_REQUEST_TAGS_METADATA_KEY, @@ -5344,7 +5418,7 @@ class Router(DeprecatedRouterMembers): hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) if _is_fallback_hop(hop_kwargs): EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( - self, + self._core_ctx, hop_kwargs.get("input"), hop_kwargs.get("messages"), (_FALLBACK_HOP_ADAPTER.validate_python(deployment),), @@ -5951,7 +6025,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -5961,15 +6035,42 @@ class Router(DeprecatedRouterMembers): async def _aimage_generation(self, prompt: str, model: str, **kwargs): model_name = model + start_time: Final = datetime.now() # noqa: DTZ005 # SDK logging uses naive local timestamps try: verbose_router_logger.debug("Inside _image_generation()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + try: + deployment: Final = await self._core_async_get_available_deployment( + model=model, + messages=[{"role": "user", "content": "prompt"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + except Exception as selection_error: + # Admitted attempts keep SDK-owned logging, including deployment callbacks and retry state + request_kwargs: Final = cast( # cast-ok: Python kwargs have string keys and opaque values + Mapping[str, object], kwargs + ) + if request_kwargs.get("litellm_logging_obj") is None: + refusal_kwargs: Final = {**request_kwargs, "model": model, "prompt": prompt} + refusal_kwargs.setdefault("litellm_call_id", str(uuid.uuid4())) + logging_obj, _ = function_setup( + "aimage_generation", + Rules(), + start_time, + **refusal_kwargs, # pyright: ignore[reportArgumentType] # preserve opaque SDK options, including is_async_call + ) + try: + asyncio.create_task( + logging_obj.dispatch_failure_handlers( + exception=selection_error, + traceback_exception=traceback.format_exc(), + prefer_async_handlers=True, + ) + ) + finally: + _restore_correlation_context_if_supported(logging_obj) + raise self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) data: Final = deployment["litellm_params"].copy() @@ -6036,7 +6137,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -6129,7 +6230,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -6191,7 +6292,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -6308,7 +6409,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -6378,7 +6479,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -6437,7 +6538,7 @@ class Router(DeprecatedRouterMembers): return _OPAQUE_VALUE_ADAPTER.validate_python( await SearchAPIRouter.async_search_with_fallbacks( - router_instance=self, + litellm_routing_ctx=self._core_ctx, original_function=original_function, **kwargs, ) @@ -6451,7 +6552,7 @@ class Router(DeprecatedRouterMembers): from litellm.router_utils.search_api_router import SearchAPIRouter return await SearchAPIRouter.async_search_with_fallbacks_helper( - router_instance=self, + litellm_routing_ctx=self._core_ctx, model=model, original_generic_function=original_generic_function, **kwargs, @@ -6569,7 +6670,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -6638,6 +6739,7 @@ class Router(DeprecatedRouterMembers): kwargs=kwargs, function_name=function_name, splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + opens_responses_socket=handler_name == "_aresponses_websocket", ) data: Final = deployment["litellm_params"].copy() @@ -6668,7 +6770,7 @@ class Router(DeprecatedRouterMembers): response_kwargs["custom_llm_provider"] = custom_llm_provider compacted_input: Final = await compact_to_fit( - self, + self._core_ctx, deployment, response_kwargs, surface_for_call(generic_function_name), @@ -7551,7 +7653,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -7630,7 +7732,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -7644,7 +7746,7 @@ class Router(DeprecatedRouterMembers): **kwargs, ) -> OpenAIFileObject: try: - from litellm.router_utils.common_utils import add_model_file_id_mappings + from litellm.litellm_core_utils.routing_common_utils import add_model_file_id_mappings verbose_router_logger.debug("Inside _atext_completion()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) @@ -7873,7 +7975,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -8041,7 +8143,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -8072,7 +8174,7 @@ class Router(DeprecatedRouterMembers): except Exception as e: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=self, + ctx=self._core_ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -8863,7 +8965,7 @@ class Router(DeprecatedRouterMembers): # that fails with RouterRateLimitError whenever the "remaining" entries # are all in cooldown — the inner async_get_healthy_deployments call # would find an empty list and raise immediately. - cooldown_ids = set(await async_get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=None)) + cooldown_ids = set(await async_get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=None)) remaining: Final = (all_ids - cooldown_ids) - excluded if not remaining: return None @@ -8943,7 +9045,7 @@ class Router(DeprecatedRouterMembers): raise e input_kwargs: Final = { - "litellm_router": self, + "litellm_routing_ctx": self._core_ctx, "original_exception": original_exception, **kwargs, "attempted_targets": fallback_walk, @@ -9159,7 +9261,7 @@ class Router(DeprecatedRouterMembers): parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) fallback_failure_exception_str = truncate_fallback_error_detail(redact_string(str(new_exception))) cooldown_info: Final = await async_get_cooldown_deployments_with_debug_info( - litellm_router_instance=self, + ctx=self._core_ctx, parent_otel_span=parent_otel_span, ) verbose_router_logger.error( @@ -9781,7 +9883,7 @@ class Router(DeprecatedRouterMembers): # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits increment_deployment_successes_for_current_minute( - litellm_router_instance=self, + ctx=self._core_ctx, deployment_id=str(deployment_id), ) @@ -9961,7 +10063,7 @@ class Router(DeprecatedRouterMembers): if id is not None: key: Final = increment_deployment_successes_for_current_minute( - litellm_router_instance=self, + ctx=self._core_ctx, deployment_id=id, ) return key @@ -10072,11 +10174,11 @@ class Router(DeprecatedRouterMembers): if deployment_id is None: return False increment_deployment_failures_for_current_minute( - litellm_router_instance=self, + ctx=self._core_ctx, deployment_id=deployment_id, ) result: Final = set_cooldown_deployments( - litellm_router_instance=self, + ctx=self._core_ctx, exception_status=exception_status, original_exception=exception, deployment=deployment_id, @@ -10382,9 +10484,7 @@ class Router(DeprecatedRouterMembers): except Exception: pass - unhealthy_deployments: Final = get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span - ) + unhealthy_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) unhealthy_set: Final = set(unhealthy_deployments) healthy_deployments: list = [d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_set] healthy_deployments = self._filter_blocked_deployments(healthy_deployments) @@ -10412,7 +10512,7 @@ class Router(DeprecatedRouterMembers): pass unhealthy_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span + ctx=self._core_ctx, parent_otel_span=parent_otel_span ) # Convert to set for O(1) lookup instead of O(n) unhealthy_deployments_set: Final = set(unhealthy_deployments) @@ -10470,7 +10570,7 @@ class Router(DeprecatedRouterMembers): ) ) set_cooldown_deployments( - litellm_router_instance=self, + ctx=self._core_ctx, exception_status=e.status_code, original_exception=e, deployment=deployment["model_info"]["id"], @@ -10498,9 +10598,20 @@ class Router(DeprecatedRouterMembers): return deployment if current is None else current.model_dump(exclude_none=True) async def async_admit_deployment_turn( - self, model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> None: + self, + model_group: str, + opened_on: Mapping[str, object], + request_kwargs: Mapping[str, object], + ) -> Mapping[str, object]: deployment: Final = self._deployment_as_configured_now(opened_on) + rate_limit_check: Final = next( + (callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)), None + ) + io_token_admission: Final[Mapping[str, object]] = ( + MappingProxyType({}) + if rate_limit_check is None + else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) + ) budget_limiter: Final = self._core_get_router_deployment_budget_limiter() if budget_limiter is not None: try: @@ -10521,6 +10632,7 @@ class Router(DeprecatedRouterMembers): await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None ) + return io_token_admission async def _admit_turn_through_routing_strategy( self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] @@ -10543,7 +10655,7 @@ class Router(DeprecatedRouterMembers): ) if selected is None: no_deployment: Final = await async_raise_no_deployment_exception( - litellm_router_instance=self, model=model_group, parent_otel_span=None + ctx=self._core_ctx, model=model_group, parent_otel_span=None ) raise litellm.RateLimitError( message=str(no_deployment), @@ -10974,7 +11086,7 @@ class Router(DeprecatedRouterMembers): auto_router_config=auto_router_config, default_model=default_model, embedding_model=embedding_model, - litellm_router_instance=self, + ctx=self._core_ctx, max_input_chars=(max_input_chars if max_input_chars is not None else DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS), ) self._register_pre_routing_strategy( @@ -11055,7 +11167,7 @@ class Router(DeprecatedRouterMembers): complexity_router: Final[ComplexityRouter] = ComplexityRouter( model_name=deployment.model_name, default_model=default_model, - litellm_router_instance=self, + ctx=self._core_ctx, complexity_router_config=complexity_router_config, ) self._register_pre_routing_strategy( @@ -11355,7 +11467,7 @@ class Router(DeprecatedRouterMembers): quality_router: Final[QualityRouter] = QualityRouter( model_name=deployment.model_name, default_model=default_model, - litellm_router_instance=self, + ctx=self._core_ctx, quality_router_config=quality_router_config, ) self._register_pre_routing_strategy( @@ -13576,7 +13688,7 @@ class Router(DeprecatedRouterMembers): while remaining: try: picked = await get_deployments_for_tag( - llm_router_instance=self, + ctx=self._core_ctx, model=model, healthy_deployments=remaining, request_kwargs=tag_request, @@ -14459,7 +14571,7 @@ class Router(DeprecatedRouterMembers): if routing_strategy == "lar1": from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - apply_lar1_routing_strategy(self, kwargs.get("routing_strategy_args")) + apply_lar1_routing_strategy(self._core_ctx, kwargs.get("routing_strategy_args")) return self.routing_strategy_init( routing_strategy=routing_strategy, routing_strategy_args=kwargs.get("routing_strategy_args", {}) @@ -14552,7 +14664,7 @@ class Router(DeprecatedRouterMembers): cache_key = f"{model_id}_max_parallel_requests_client" client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) if client is None: - InitalizeCachedClient.set_max_parallel_requests_client(litellm_router_instance=self, model=deployment) + InitalizeCachedClient.set_max_parallel_requests_client(ctx=self._core_ctx, model=deployment) client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) return client elif client_type == "async": @@ -15483,10 +15595,10 @@ class Router(DeprecatedRouterMembers): routing_read_batch: Final = RoutingReadBatch.active() cooldown_deployments: Final = ( - await async_get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) + await async_get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) if routing_read_batch is None else await routing_read_batch.async_get_cooldown_deployments( - litellm_router_instance=self, + ctx=self._core_ctx, healthy_deployments=healthy_deployments, parent_otel_span=parent_otel_span, ) @@ -15541,7 +15653,7 @@ class Router(DeprecatedRouterMembers): ) # check if user wants to do tag based routing healthy_deployments = await get_deployments_for_tag( - llm_router_instance=self, + ctx=self._core_ctx, model=model, request_kwargs=request_kwargs, healthy_deployments=healthy_deployments, @@ -15620,7 +15732,7 @@ class Router(DeprecatedRouterMembers): if len(healthy_deployments) == 0: exception: Final = await async_raise_no_deployment_exception( - litellm_router_instance=self, + ctx=self._core_ctx, model=model, parent_otel_span=parent_otel_span, route_model_ids=await self._async_routable_cooled_deployment_ids( @@ -15869,7 +15981,7 @@ class Router(DeprecatedRouterMembers): ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( - litellm_router_instance=self, + ctx=self._core_ctx, model=model, parent_otel_span=parent_otel_span, ) @@ -16040,7 +16152,7 @@ class Router(DeprecatedRouterMembers): if deployment is None: exception: Final = await async_raise_no_deployment_exception( - litellm_router_instance=self, + ctx=self._core_ctx, model=model, parent_otel_span=parent_otel_span, ) @@ -16408,7 +16520,7 @@ class Router(DeprecatedRouterMembers): if isinstance(selected_strategy.strategy, ComplexityRouter) else (), parent_model=model, - router=self, + ctx=self._core_ctx, allow_escalation=isinstance(selected_strategy.strategy, ComplexityRouter) and selected_strategy.strategy.config.enable_context_window_escalation, messages=messages, @@ -16771,9 +16883,7 @@ class Router(DeprecatedRouterMembers): parent_otel_span=parent_otel_span, ) - cooldown_deployments: Final = get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span - ) + cooldown_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) _pre_cooldown_deployments: Final = healthy_deployments healthy_deployments = self._filter_cooldown_deployments( healthy_deployments=healthy_deployments, @@ -16872,7 +16982,7 @@ class Router(DeprecatedRouterMembers): _cooldown_time = self._core_cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) - _cooldown_list = get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) + _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) raise RouterRateLimitError( model=model, cooldown_time=_cooldown_time, @@ -16880,7 +16990,7 @@ class Router(DeprecatedRouterMembers): cooldown_list=_cooldown_list, model_ids=model_ids, cooldown_deployments=get_cooldown_deployment_info( - self, + self._core_ctx, self._deployment_ids( self._region_routable_deployments( request_deployments, @@ -16920,7 +17030,7 @@ class Router(DeprecatedRouterMembers): _cooldown_time = self._core_cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) - _cooldown_list = get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) + _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) raise RouterRateLimitError( model=model, cooldown_time=_cooldown_time, @@ -17027,9 +17137,7 @@ class Router(DeprecatedRouterMembers): healthy_deployments=pass_through_deployments, parent_otel_span=parent_otel_span, ) - cooldown_deployments: Final = get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span - ) + cooldown_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) pass_through_deployments = self._filter_cooldown_deployments( healthy_deployments=pass_through_deployments, cooldown_deployments=cooldown_deployments, @@ -17058,7 +17166,7 @@ class Router(DeprecatedRouterMembers): _cooldown_time = self._core_cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) - _cooldown_list = get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) + _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) raise RouterRateLimitError( model=model, cooldown_time=_cooldown_time, @@ -17066,7 +17174,7 @@ class Router(DeprecatedRouterMembers): cooldown_list=_cooldown_list, model_ids=pass_through_model_ids, cooldown_deployments=get_cooldown_deployment_info( - self, + self._core_ctx, self._deployment_ids( self._region_routable_deployments( request_deployments, @@ -17105,7 +17213,7 @@ class Router(DeprecatedRouterMembers): _cooldown_time = self._core_cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) - _cooldown_list = get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) + _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) raise RouterRateLimitError( model=model, cooldown_time=_cooldown_time, diff --git a/litellm/router_api.py b/litellm/router_api.py index d62999cf261..e4939254b81 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -33,19 +33,28 @@ as they were. # repeat them unchanged, and typing the members themselves is the core's job. import hashlib +import inspect import json +import warnings from collections.abc import Iterator, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass, field from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, TypeAlias, overload +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, overload if TYPE_CHECKING: from typing import Any # noqa: TID251 # only to spell the cores' untyped (Unknown) containers and one Any type + from opentelemetry.trace import Span + from litellm.caching.caching import DualCache, RedisCache from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter + from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.messages.dispatch import MessagesResult + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.router import Router from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.auto_router.auto_router import AutoRouter @@ -53,12 +62,14 @@ if TYPE_CHECKING: from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.quality_router.quality_router import QualityRouter from litellm.router_utils.cooldown_cache import CooldownCache + from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets as _AttemptedFallbackTargets from litellm.router_utils.health_state_cache import DeploymentHealthCache from litellm.router_utils.kubernetes_pod_discovery import KubernetesPodDiscovery - from litellm.router_utils.pattern_match_deployments import PatternMatchRouter - from litellm.router_utils.ptu_shares import PTUTeamCeiling + from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse + from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import ( AllowedFailsPolicy, + CustomRoutingStrategyBase, Deployment, DeploymentModelListingInfo, DeploymentTypedDict, @@ -74,7 +85,7 @@ if TYPE_CHECKING: SearchToolTypedDict, TaggedPreRoutingStrategy, ) - from litellm.types.utils import BudgetConfig + from litellm.types.utils import BudgetConfig, EmbeddingResponse, ModelResponse from litellm.types.utils import ModelInfo as ModelMapInfo # The cores' untyped containers (bare ``list`` / ``dict``, element types Unknown), spelled with Any: @@ -82,6 +93,9 @@ if TYPE_CHECKING: UntypedList: TypeAlias = list[Any] # pyright: ignore[reportExplicitAny] # see above UntypedDict: TypeAlias = dict[Any, Any] # pyright: ignore[reportExplicitAny] # see above + # Re-exported for litellm.types, which may not import router_utils. + AttemptedFallbackTargets: TypeAlias = _AttemptedFallbackTargets + # What registry().plugins(kind) returns, one member per kind. RegistryPlugins: TypeAlias = ( dict[str, list[TaggedPreRoutingStrategy[AutoRouter]]] @@ -184,6 +198,73 @@ def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str return hashlib.sha256("".join(parts).encode()).hexdigest() +def router_init_args() -> tuple[str, ...]: + """The arguments ``Router.__init__`` accepts (was ``Router.get_valid_args``). Needs no router instance.""" + from litellm.router import Router + + return tuple(Router.get_valid_args()) + + +def mark_advisor_orchestration_failure(exception: BaseException) -> None: + """Tag an advisor orchestration failure, so the router does not cool down the deployment it picked.""" + from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure as mark + + mark(exception) + + +def fallback_chain( + fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: the fallbacks config shape + model_group: str, +) -> tuple[list[str] | None, int | None]: # mutable-ok: get_fallback_model_group's type, unchanged + """ + The fallbacks configured for a model group, and the index of the generic ``*`` entry (was + ``get_fallback_model_group``). Needs no router instance. + """ + from litellm.router_utils.fallback_event_handlers import get_fallback_model_group + + return get_fallback_model_group(fallbacks=fallbacks, model_group=model_group) + + +def request_fallback_chain( + fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: the fallbacks config shape + request_kwargs: Mapping[str, object], + model_group: str | None, +) -> tuple[list[str] | None, int | None]: # mutable-ok: get_fallback_model_group's type, unchanged + """ + ``fallback_chain`` for a request: the group a pre-routing hook picked first, then the routed, + requested and original groups (was ``get_fallback_model_group_for_lookup_groups``). + """ + from litellm.router_utils.fallback_event_handlers import ( + fallback_lookup_groups, + get_fallback_model_group_for_lookup_groups, + ) + + return get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, lookup_groups=fallback_lookup_groups(kwargs=request_kwargs, model_group=model_group) + ) + + +def fallbacks_disabled(request_kwargs: Mapping[str, object]) -> bool: + """Whether the request opted out of fallbacks (was ``fallbacks_disabled_for_request``).""" + from litellm.router_utils.fallback_event_handlers import fallbacks_disabled_for_request + + return fallbacks_disabled_for_request(kwargs=request_kwargs) + + +def references_provider_scoped_resource(request_kwargs: Mapping[str, object]) -> bool: + """Whether the request names a file, batch or fine-tuning job that only its own provider knows.""" + from litellm.router_utils.fallback_event_handlers import PROVIDER_SCOPED_RESOURCE_KEYS + + return any(request_kwargs.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS) + + +def deployment_metadata(request_kwargs: Mapping[str, object]) -> Mapping[str, object] | None: + """The metadata bucket the router stamped the picked deployment into (was ``find_deployment_metadata``).""" + from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata + + return find_deployment_metadata(request_kwargs) + + def _json_default_stable_id(obj: object) -> str: from litellm.router import Router @@ -550,7 +631,8 @@ class RouterRegistry: - ``group_set``: the set of configured group names (was the ``model_names`` attribute) - ``ids``: deployment ids, of one group or all (was ``get_model_ids``) - ``deployment_names``: (was the ``deployment_names`` attribute) - - ``team_public``: team public model names (was ``team_public_model_names``) + - ``team_public``: team public model names (was ``team_public_model_names``), or only the + ones ``team_id`` serves its own deployments under (was ``team_public_model_names_of``) - ``litellm_models``: upstream model names a group resolves to (was ``resolved_litellm_models``) - ``litellm_base_models``: the same, with each deployment's ``base_model`` where it sets one (was ``resolved_litellm_models(..., use_base_model=True)``) @@ -568,7 +650,7 @@ class RouterRegistry: if kind == "deployment_names": return r._core_deployment_names if kind == "team_public": - return r._core_team_public_model_names + return r._core_team_public_model_names if team_id is None else r.team_public_model_names_of(team_id) if model_name is None: raise TypeError(f"registry().names(kind={kind!r}) needs a model_name") if kind == "litellm_models": @@ -832,6 +914,14 @@ class RouterRegistry: """Routing groups by name (was the ``routing_groups`` attribute).""" return self._router._core_routing_groups + def routing_group(self, name: str) -> "RoutingGroup | None": + """The routing group callable as ``name``, or None (was ``get_routing_group``).""" + return self._router.get_routing_group(name) + + def client(self, deployment: Mapping[str, object], *, client_type: str | None = None) -> object: + """The SDK client the router cached for a deployment, or None (was ``_get_client``).""" + return self._router._get_client(deployment=deployment, kwargs={}, client_type=client_type) + def access_groups( self, model_name: str | None = None, @@ -1128,6 +1218,41 @@ class RouterHealth: """The request priority and the kwargs without it (was ``split_request_priority``).""" return self._router._core_split_request_priority(model, kwargs) + def failures_this_minute(self, deployment_id: str) -> int: + """Failures counted against a deployment this minute (was ``get_deployment_failures_for_current_minute``).""" + from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + get_deployment_failures_for_current_minute, + ) + + return get_deployment_failures_for_current_minute(self._router._core_ctx, deployment_id) + + def record_failure(self, deployment_id: str, exception: Exception, exception_status: str | int) -> bool: + """ + Count a failure against a deployment this minute and cool it down once it is over its allowed + fails (was ``increment_deployment_failures_for_current_minute`` then ``set_cooldown_deployments``). + Returns whether the deployment was cooled down. + """ + from litellm.router_utils.cooldown_handlers import set_cooldown_deployments + from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + increment_deployment_failures_for_current_minute, + ) + + ctx: Final = self._router._core_ctx + increment_deployment_failures_for_current_minute(ctx, deployment_id) + return set_cooldown_deployments( + ctx, + original_exception=exception, + exception_status=exception_status, + deployment=deployment_id, + time_to_cooldown=self._router._core_cooldown_time, + ) + + def priority_header_fields( + self, model: str, kwargs: Mapping[str, object], header: str | None + ) -> Mapping[str, int | str]: + """The request fields a priority header sets (was ``priority_header_update``).""" + return self._router.priority_header_update(model=model, kwargs=kwargs, header=header) + @dataclass class ChangeSet: @@ -1202,6 +1327,40 @@ async def aapply_changes(router: "Router", changes: ChangeSet) -> ApplyResult: return result +def _takes_legacy_llm_router(check: object) -> bool: + """Whether a fallback check still declares ``llm_router`` instead of ``ctx`` (the pre-RoutingContext shape).""" + try: + params: Final = inspect.signature(check).parameters # pyright: ignore[reportArgumentType] # any callable + except (TypeError, ValueError): + return False + return "ctx" not in params and "llm_router" in params + + +async def _run_fallback_check( + check: "FallbackAccessCheck | FallbackBudgetCheck", + ctx: "RoutingContext", + *, + model: str, + request_kwargs: Mapping[str, object], +) -> bool: + if _takes_legacy_llm_router(check): + warnings.warn( + f"{type(check).__name__} takes llm_router=; fallback checks now receive ctx: RoutingContext. " + "Passing llm_router is deprecated and will be removed", + DeprecationWarning, + stacklevel=2, + ) + legacy: Final = check # pyright: ignore[reportAny] # the old protocol shape, checked above + return await legacy(model=model, request_kwargs=request_kwargs, llm_router=ctx._router) # pyright: ignore[reportCallIssue, reportPrivateUsage] # see above + return await check(model=model, request_kwargs=request_kwargs, ctx=ctx) + + +class RegistryReader(Protocol): + """Anything that answers ``registry()``: a Router or a RoutingContext. What proxy auth helpers read.""" + + def registry(self) -> "RouterRegistry": ... + + class RoutingContext: """ What a routing plugin (strategy, pre-call check, sub-router) may use of the router. Plugins get @@ -1216,12 +1375,107 @@ class RoutingContext: def registry(self) -> RouterRegistry: return self._router.registry() + def deployment_params(self, request: Mapping[str, object], deployment: "Deployment | None") -> Mapping[str, object]: + return self._router._params_sent_to_deployment(request, deployment) + def settings(self) -> RouterSettings: return self._router.settings() def health(self) -> RouterHealth: return self._router.health() + def observers(self) -> RoutingObservers: + return self._router.observe() + + def belongs_to(self, router: "Router") -> bool: + return self._router is router + + @property + def explicit_allowed_fails(self) -> int | None: + """``allowed_fails`` as the caller passed it, ``None`` when it came from the global default.""" + return self._router.explicit_allowed_fails + + async def select( + self, + model: str, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict + *, + kind: SelectKind = "pick", + messages: list[dict[str, str]] | None = None, # mutable-ok: the selection cores take a list + input: str | list[object] | None = None, # mutable-ok: the selection cores take a list + specific_deployment: bool | None = False, + parent_otel_span: "Span | None" = None, + health_check_probe: bool = False, + ) -> object: + return await self._router.select( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportCallIssue, reportArgumentType] # the selection cores return untyped deployment dicts + model, + request_kwargs, + kind=kind, # pyright: ignore[reportArgumentType] # forwarded as given, the overloads narrow it + messages=messages, + input=input, + specific_deployment=specific_deployment, + parent_otel_span=parent_otel_span, + health_check_probe=health_check_probe, + ) + + async def with_fallbacks( + self, + *args: object, + **kwargs: object, # kwargs-ok: the router's fallback loop takes the request's own keyword payload + ) -> object: + """Run a call through the router's retry and fallback loop (was ``async_function_with_fallbacks``).""" + return await self._router.async_function_with_fallbacks(*args, **kwargs) + + async def acompletion( + self, + **kwargs: object, # kwargs-ok: forwarded to the router endpoint as given + ) -> "ModelResponse | CustomStreamWrapper": + return await self._router.acompletion(**kwargs) # pyright: ignore[reportCallIssue, reportArgumentType] # forwarded as given + + async def aresponses( + self, + **kwargs: object, # kwargs-ok: forwarded to the router endpoint as given + ) -> "ResponsesAPIResponse | BaseResponsesAPIStreamingIterator": + return await self._router.aresponses(**kwargs) # pyright: ignore[reportReturnType, reportArgumentType] # factory-built endpoint, returns what litellm.aresponses returns + + async def aembedding( + self, + **kwargs: object, # kwargs-ok: forwarded to the router endpoint as given + ) -> "EmbeddingResponse": + return await self._router.aembedding(**kwargs) # pyright: ignore[reportCallIssue, reportArgumentType] # forwarded as given + + def embedding( + self, + **kwargs: object, # kwargs-ok: forwarded to the router endpoint as given + ) -> "EmbeddingResponse": + return self._router.embedding(**kwargs) # pyright: ignore[reportCallIssue, reportArgumentType] # forwarded as given + + async def adecisions( + self, + **kwargs: object, # kwargs-ok: forwarded to the router endpoint as given + ) -> "OpenAIDecisionResponse | DecisionsResponse": + return await self._router.adecisions(**kwargs) # pyright: ignore[reportReturnType, reportArgumentType] # factory-built endpoint, returns what litellm.adecisions returns + + async def aanthropic_messages( + self, + **kwargs: object, # kwargs-ok: forwarded to the router endpoint as given + ) -> "MessagesResult": + return await self._router.aanthropic_messages(**kwargs) # pyright: ignore[reportReturnType, reportArgumentType] # factory-built endpoint, returns what litellm.anthropic_messages returns + + def apply(self, changes: ChangeSet) -> ApplyResult: + return self._router.apply(changes) + + def register_plugin(self, plugin: "CustomRoutingStrategyBase") -> None: + self._router.register_plugin(plugin) + + async def fallback_access_allowed(self, *, model: str, request_kwargs: Mapping[str, object]) -> bool: + check: Final = self._router.settings().fallback_access_check + return check is None or await _run_fallback_check(check, self, model=model, request_kwargs=request_kwargs) + + async def fallback_budget_allowed(self, *, model: str, request_kwargs: Mapping[str, object]) -> bool: + check: Final = self._router.settings().fallback_budget_check + return check is None or await _run_fallback_check(check, self, model=model, request_kwargs=request_kwargs) + def pre_call_checks( self, model: str, @@ -1289,7 +1543,7 @@ class RoutingContext: self, kwargs: dict[str, object], # mutable-ok: Router._core_log_retry records the retry into the request kwargs e: Exception, - ) -> Mapping[str, object]: + ) -> dict: # mutable-ok: the fallback loop edits it; Router._core_log_retry's type return self._router._core_log_retry(kwargs=kwargs, e=e) diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index b7df0677a10..756d54595da 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -41,11 +41,11 @@ if TYPE_CHECKING: from litellm.caching.caching import RedisCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import Router from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_utils.kubernetes_pod_discovery import KubernetesPodDiscovery - from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.router import ( CustomRoutingStrategyBase, Deployment, @@ -852,6 +852,109 @@ class DeprecatedRouterMembers: return _as_router(self)._core_log_retry(kwargs=kwargs, e=e) +# Router-family functions and classes that took the Router (``litellm_router_instance=``, +# ``router=``, ...) take its RoutingContext now. These keep the old spellings working. + +_RoutingCallable: Final = TypeVar("_RoutingCallable", bound=Callable[..., object]) + + +@functools.cache +def _router_class() -> "type[Router]": + from litellm.router import Router + + return Router + + +def _note_renamed(label: str, old: str, instead: str, *, stacklevel: int = 3) -> None: + DEPRECATED_ROUTER_MEMBER_USES[label] += 1 + warnings.warn( + f"{label}: {old} is deprecated; {instead} (the Router's RoutingContext) instead", + DeprecationWarning, + stacklevel=stacklevel, # the caller of the wrapper or descriptor that noted it + ) + + +def _routing_context_of(value: object) -> object: + """The Router's RoutingContext when ``value`` is a Router, else ``value`` as given.""" + return value._core_ctx if isinstance(value, _router_class()) else value + + +def renamed_router_param(old: str, new: str = "ctx") -> Callable[[_RoutingCallable], _RoutingCallable]: + """ + Keep a router parameter that became a RoutingContext (``new``) callable the old way: ``old=`` + still works, and so does a Router passed where ``new`` goes. Either gets the Router's + RoutingContext and a DeprecationWarning that names ``new``. + """ + + def decorate(fn: _RoutingCallable) -> _RoutingCallable: + index: Final = list(inspect.signature(fn).parameters).index(new) + label: Final = f"{fn.__module__}.{fn.__qualname__}" + + def convert( + args: tuple[object, ...], + kwargs: dict[str, object], # mutable-ok: the call's own kwargs + ) -> tuple[object, ...]: + if old in kwargs: + if new in kwargs: + raise TypeError(f"{fn.__qualname__}() got both {old}= and {new}=") + _note_renamed(label, f"{old}=", f"pass {new}=", stacklevel=4) + kwargs[new] = _routing_context_of(kwargs.pop(old)) + elif new in kwargs: + if isinstance(kwargs[new], _router_class()): + _note_renamed(label, f"a Router as {new}=", f"pass {new}=", stacklevel=4) + kwargs[new] = _routing_context_of(kwargs[new]) + elif len(args) > index and isinstance(args[index], _router_class()): + _note_renamed(label, "passing a Router", f"pass {new}", stacklevel=4) + return (*args[:index], _routing_context_of(args[index]), *args[index + 1 :]) + return args + + # An async fn keeps an async wrapper, so iscoroutinefunction (and mock.patch's AsyncMock pick) still sees it + if inspect.iscoroutinefunction(fn): + + @functools.wraps(fn) + async def acall(*args: object, **kwargs: object) -> object: # kwargs-ok: forwarded to fn as given + return await fn(*convert(args, kwargs), **kwargs) # pyright: ignore[reportAny] # fn's own result + + return acall # pyright: ignore[reportReturnType] # same call shape as fn, see functools.wraps + + @functools.wraps(fn) + def call(*args: object, **kwargs: object) -> object: # kwargs-ok: forwarded to fn as given + return fn(*convert(args, kwargs), **kwargs) + + return call # pyright: ignore[reportReturnType] # same call shape as fn, see functools.wraps + + return decorate + + +class RenamedRouterAttribute: + """ + A class attribute that keeps an old router attribute (``self.router``, ``self.litellm_router_instance``) + readable after it became the RoutingContext attribute ``new``: reading it returns that + RoutingContext, writing a Router stores the Router's RoutingContext, each with a DeprecationWarning. + """ + + def __init__(self, new: str = "ctx") -> None: + self.new: Final = new + self.label = "" + + def __set_name__(self, owner: type, name: str) -> None: + self.label = f"{owner.__module__}.{owner.__qualname__}.{name}" + + @overload + def __get__(self, obj: None, owner: type | None = None) -> "RenamedRouterAttribute": ... + @overload + def __get__(self, obj: object, owner: type | None = None) -> object: ... + def __get__(self, obj: object | None, owner: type | None = None) -> object: + if obj is None: + return self + _note_renamed(self.label, "reading this attribute", f"read .{self.new}") + return getattr(obj, self.new) # pyright: ignore[reportAny] # the attribute named ``new`` holds the RoutingContext + + def __set__(self, obj: object, value: object) -> None: + _note_renamed(self.label, "setting this attribute", f"set .{self.new}") + setattr(obj, self.new, _routing_context_of(value)) + + def _name_members_after_router() -> None: # The old members were Router methods, so errors that name the method (a wrong-arity TypeError, # or "litellm.router.Router.() argument after ** must be a mapping") keep naming it the diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index ee620790282..b92c53fa93c 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.internal_call_metadata import ( forwarded_internal_call_metadata, parent_session_kwargs, ) +from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, @@ -29,11 +30,11 @@ if TYPE_CHECKING: from semantic_router.routers import SemanticRouter from semantic_router.routers.base import Route - from litellm.router import Router + from litellm.router_api import RoutingContext from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder from litellm.types.router import PreRoutingHookResponse else: - Router = Any + RoutingContext = Any PreRoutingHookResponse = Any Route = Any SemanticRouter = Any @@ -55,13 +56,15 @@ class _SemanticMatchOutcome(NamedTuple): class AutoRouter(CustomLogger): DEFAULT_AUTO_SYNC_VALUE = "local" + litellm_router_instance: Final = RenamedRouterAttribute() # deprecated: the router attribute before ``ctx`` + @renamed_router_param("litellm_router_instance") def __init__( self, model_name: str, default_model: str, embedding_model: str, - litellm_router_instance: "Router", + ctx: "RoutingContext", auto_router_config_path: str | None = None, auto_router_config: str | None = None, max_input_chars: int = DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, @@ -75,7 +78,7 @@ class AutoRouter(CustomLogger): auto_router_config: The config to use for the auto-router. You can either use this or auto_router_config_path, not both. default_model: The default model to use if no route is found. embedding_model: The embedding model to use for the auto-router. - litellm_router_instance: The instance of the LiteLLM Router. + ctx: The routing context of the owning router. max_input_chars: Longest prompt, in characters, handed to the embedding model. Longer prompts are cut to this length so an embedding context window far smaller than the routed model's never fails the request. @@ -95,9 +98,9 @@ class AutoRouter(CustomLogger): self.default_model = default_model self.embedding_model: str = embedding_model self.max_input_chars: int = max_input_chars - self.litellm_router_instance: Router = litellm_router_instance + self.ctx: Final = ctx self.encoder: LiteLLMRouterEncoder = LiteLLMRouterEncoder( - litellm_router_instance=litellm_router_instance, + ctx=ctx, model_name=embedding_model, max_input_chars=max_input_chars, ) diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py index b7430062f1c..340efb509eb 100644 --- a/litellm/router_strategy/auto_router/litellm_encoder.py +++ b/litellm/router_strategy/auto_router/litellm_encoder.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional from pydantic import ConfigDict from semantic_router.encoders import DenseEncoder @@ -6,11 +6,12 @@ from semantic_router.encoders.base import AsymmetricDenseMixin import litellm from litellm._logging import verbose_router_logger +from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext else: - Router = Any + RoutingContext = Any def litellm_to_list(embeds: litellm.EmbeddingResponse) -> list[list[float]]: @@ -19,21 +20,22 @@ def litellm_to_list(embeds: litellm.EmbeddingResponse) -> list[list[float]]: :param embeds: The LiteLLM embedding response. :return: A list of embeddings. """ - if not embeds or not isinstance(embeds, litellm.EmbeddingResponse) or not embeds.data: + if not embeds or not embeds.data: raise ValueError("No embeddings found in LiteLLM embedding response.") return [x["embedding"] for x in embeds.data] class CustomDenseEncoder(DenseEncoder): model_config = ConfigDict(extra="allow") + litellm_router_instance: ClassVar = RenamedRouterAttribute() # deprecated: the router attribute before ``ctx`` - def __init__(self, litellm_router_instance: Optional["Router"] = None, **kwargs): - # Extract litellm_router_instance from kwargs if passed there - if "litellm_router_instance" in kwargs: - litellm_router_instance = kwargs.pop("litellm_router_instance") + @renamed_router_param("litellm_router_instance") + def __init__(self, ctx: Optional["RoutingContext"] = None, **kwargs): + if "ctx" in kwargs: + ctx = kwargs.pop("ctx") super().__init__(**kwargs) - self.litellm_router_instance = litellm_router_instance + self.ctx = ctx class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): @@ -46,17 +48,18 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): type: str = "internal_litellm_router" + @renamed_router_param("litellm_router_instance") def __init__( self, - litellm_router_instance: "Router", + ctx: "RoutingContext", model_name: str, score_threshold: float | None = None, max_input_chars: int = 0, ): """Initialize the LiteLLMEncoder. - :param litellm_router_instance: The instance of the LiteLLM Router. - :type litellm_router_instance: Router + :param ctx: The routing context of the owning router. + :type ctx: RoutingContext :param model_name: The name of the embedding model to use. Must use LiteLLM naming convention (e.g. "openai/text-embedding-3-small" or "mistral/mistral-embed"). :type model_name: str @@ -75,7 +78,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): score_threshold=score_threshold if score_threshold is not None else 0.3, ) self.model_name = model_name - self.litellm_router_instance = litellm_router_instance + self.ctx = ctx self.max_input_chars = max_input_chars def _clamp(self, docs: list[str]) -> list[str]: # mutable-ok: embedding() takes `input: str | list` @@ -109,45 +112,37 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): return await self.aencode_queries(docs, **kwargs) def encode_queries(self, docs: list[str], **kwargs) -> list[list[float]]: - if self.litellm_router_instance is None: - raise ValueError("litellm_router_instance is not set") + if self.ctx is None: + raise ValueError("ctx is not set") try: - embeds: Final = self.litellm_router_instance.embedding( - input=self._clamp(docs), model=self.model_name, **kwargs - ) + embeds: Final = self.ctx.embedding(input=self._clamp(docs), model=self.model_name, **kwargs) return litellm_to_list(embeds) except Exception as e: raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e def encode_documents(self, docs: list[str], **kwargs) -> list[list[float]]: - if self.litellm_router_instance is None: - raise ValueError("litellm_router_instance is not set") + if self.ctx is None: + raise ValueError("ctx is not set") try: - embeds: Final = self.litellm_router_instance.embedding( - input=self._clamp(docs), model=self.model_name, **kwargs - ) + embeds: Final = self.ctx.embedding(input=self._clamp(docs), model=self.model_name, **kwargs) return litellm_to_list(embeds) except Exception as e: raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e async def aencode_queries(self, docs: list[str], **kwargs) -> list[list[float]]: - if self.litellm_router_instance is None: - raise ValueError("litellm_router_instance is not set") + if self.ctx is None: + raise ValueError("ctx is not set") try: - embeds: Final = await self.litellm_router_instance.aembedding( - input=self._clamp(docs), model=self.model_name, **kwargs - ) + embeds: Final = await self.ctx.aembedding(input=self._clamp(docs), model=self.model_name, **kwargs) return litellm_to_list(embeds) except Exception as e: raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e async def aencode_documents(self, docs: list[str], **kwargs) -> list[list[float]]: - if self.litellm_router_instance is None: - raise ValueError("litellm_router_instance is not set") + if self.ctx is None: + raise ValueError("ctx is not set") try: - embeds: Final = await self.litellm_router_instance.aembedding( - input=self._clamp(docs), model=self.model_name, **kwargs - ) + embeds: Final = await self.ctx.aembedding(input=self._clamp(docs), model=self.model_name, **kwargs) return litellm_to_list(embeds) except Exception as e: raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 30f1021251f..4f4ed76c170 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -64,6 +64,7 @@ from litellm.llms.anthropic.common_utils import is_claude_code_user_agent from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG +from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.context_compaction import compaction_pending from litellm.router_strategy.complexity_router.tier_predictor import ( @@ -126,10 +127,13 @@ from .jev_classifier import ( DEFAULT_JEV_INSTRUCTIONS, HttpJevClassifierClient, JevClassifierClient, + JevSystemOneRequest, JevVerdict, RouterJevClassifierClient, build_jev_request, + fit_jev_request, jev_classifier_cost, + jev_input_targets, ) from .llm_v2 import LLM_V2_PROMPT_VERSION, LLMV2Decision, LLMV2TaskContext, LLMV2Verdict, llm_v2_response_format from .stall_detector import detect_stalled_task @@ -137,12 +141,12 @@ from .stall_detector import detect_stalled_task if TYPE_CHECKING: from semantic_router.routers import SemanticRouter - from litellm.router import Router + from litellm.router_api import RoutingContext from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.savings_baseline import Baseline from litellm.types.router import PreRoutingHookResponse else: - Router = Any + RoutingContext = Any PreRoutingHookResponse = Any SemanticRouter = Any @@ -1166,6 +1170,7 @@ class ClassificationOutcome(NamedTuple): heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None classifier_failure_reason: ClassifierFailureReason | None = None classifier_error_type: str | None = None + classifier_input_telemetry: StandardLoggingRoutingDecision | None = None def _with_classifier_failure( @@ -1207,23 +1212,24 @@ def _with_classifier_forecast( decision: StandardLoggingRoutingDecision, outcome: ClassificationOutcome ) -> StandardLoggingRoutingDecision: """Attach validated forecasts and their applied policy to the routing decision.""" + with_input: Final[StandardLoggingRoutingDecision] = {**decision, **(outcome.classifier_input_telemetry or {})} if outcome.heuristic_v2_forecast is not None: - return {**decision, "heuristic_v2_forecast": outcome.heuristic_v2_forecast} + return {**with_input, "heuristic_v2_forecast": outcome.heuristic_v2_forecast} if outcome.jev_verdict is not None: forecasted_decision: Final[StandardLoggingRoutingDecision] = { - **decision, + **with_input, "classifier_probabilities": outcome.jev_verdict.probabilities, "classifier_confidence": outcome.jev_verdict.confidence, } return forecasted_decision if outcome.llm_v2_forecast is not None: - return _with_llm_v2_forecast(decision, outcome.llm_v2_forecast) + return _with_llm_v2_forecast(with_input, outcome.llm_v2_forecast) forecast: Final = outcome.capability_forecast if forecast is None: - return decision + return with_input verdict: Final = forecast.verdict enriched: Final[StandardLoggingRoutingDecision] = { - **decision, + **with_input, "classifier_crux": verdict.crux, "classifier_primary_rule": verdict.primary_rule, "classifier_capability_boundary": verdict.capability_boundary, @@ -1411,9 +1417,11 @@ class ComplexityRouter(CustomLogger): - Question complexity (multiple questions) """ + litellm_router_instance: Final = RenamedRouterAttribute() # deprecated: the router attribute before ``ctx`` + def _build_jev_client(self, config: OpenSourceClassifierConfig) -> JevClassifierClient | _JevConnectionError: if config.deployment_name is not None: - return RouterJevClassifierClient(self.litellm_router_instance, config.deployment_name) + return RouterJevClassifierClient(self.ctx, config.deployment_name) if config.provider in ("laya", "bespoke"): from litellm.llms.oss_decision import oss_connection @@ -1454,10 +1462,11 @@ class ComplexityRouter(CustomLogger): http_client=get_async_httpx_client(httpxSpecialProvider.PassThroughEndpoint), ) + @renamed_router_param("litellm_router_instance") def __init__( self, model_name: str, - litellm_router_instance: Router, + ctx: RoutingContext, complexity_router_config: Mapping[str, object] | None = None, default_model: str | None = None, derive_savings_baseline: bool = True, @@ -1468,7 +1477,7 @@ class ComplexityRouter(CustomLogger): Args: model_name: The name of the model/deployment using this router. - litellm_router_instance: The LiteLLM Router instance. + ctx: The routing context of the owning router. complexity_router_config: Optional configuration dict from proxy config. default_model: Optional default model to use if tier cannot be determined. derive_savings_baseline: False for callers whose decisions are never spend @@ -1476,7 +1485,7 @@ class ComplexityRouter(CustomLogger): would leak deployment mappings the caller was not authorized for. """ self.model_name = model_name - self.litellm_router_instance = litellm_router_instance + self.ctx: Final = ctx self._derive_savings_baseline = derive_savings_baseline # Parse config - always create a new instance to avoid singleton mutation @@ -1693,7 +1702,7 @@ class ComplexityRouter(CustomLogger): if not self._derive_savings_baseline: return None if not self._savings_baseline_derived: - self._savings_baseline = resolve_baseline(self.litellm_router_instance, self._hardest_tier_models()) + self._savings_baseline = resolve_baseline(self.ctx, self._hardest_tier_models()) self._savings_baseline_derived = True return self._savings_baseline @@ -1974,19 +1983,21 @@ class ComplexityRouter(CustomLogger): `score` because the score is only interpretable against the boundaries that mapped it to a tier. """ - decision: Final = StandardLoggingRoutingDecision( - router_model_name=self.model_name, - router_type="complexity", - routed_model=routed_model, - cause=cause, - conversation_continuing=conversation_continuing, - **( - {"classifier_failure_reason": classifier_failure_reason} - if classifier_failure_reason is not None - else {} - ), - **({"classifier_error_type": classifier_error_type} if classifier_error_type is not None else {}), + failure: Final[StandardLoggingRoutingDecision] = ( + {"classifier_failure_reason": classifier_failure_reason} if classifier_failure_reason is not None else {} ) + error: Final[StandardLoggingRoutingDecision] = ( + {"classifier_error_type": classifier_error_type} if classifier_error_type is not None else {} + ) + decision: Final[StandardLoggingRoutingDecision] = { + "router_model_name": self.model_name, + "router_type": "complexity", + "routed_model": routed_model, + "cause": cause, + "conversation_continuing": conversation_continuing, + **failure, + **error, + } if (baseline := self.savings_baseline) is not None: decision["savings_baseline_model"] = baseline.model if baseline.deployment_id is not None: @@ -2371,16 +2382,6 @@ class ComplexityRouter(CustomLogger): system_prompt, failure_reason="unsupported_input", ) - breaker: Final = self._classifier_circuit_breaker - permit: Final = breaker.acquire_permit() if breaker is not None else None - if breaker is not None and permit is None: - return self._classifier_failure_outcome( - "jev classifier circuit is open", - prompt, - system_prompt, - failure_reason="circuit_open", - signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL, - ) criteria: Final[Mapping[str, str]] = ( MappingProxyType( { @@ -2395,13 +2396,40 @@ class ComplexityRouter(CustomLogger): ) ) timeout_s: Final = config.timeout_ms / 1000 - request: Final = build_jev_request( - prompt=self._classifier_context_payload(prompt, system_prompt, request_kwargs, messages), - system_prompt=None, - model=config.model, - instructions=config.instructions or DEFAULT_JEV_INSTRUCTIONS, - criteria=criteria, + from litellm.litellm_core_utils.token_counter import offload_token_count + + def build_request(current_prompt: str, caller_constraints: str | None) -> JevSystemOneRequest: + return build_jev_request( + prompt=self._classifier_context_payload( + current_prompt, caller_constraints, request_kwargs, messages, current_ask_for_history=prompt + ), + system_prompt=None, + model=config.model, + instructions=config.instructions or DEFAULT_JEV_INSTRUCTIONS, + criteria=criteria, + ) + + budgeted: Final = await offload_token_count(fit_jev_request)( + build_request, prompt, system_prompt, jev_input_targets(self.ctx, config), _truncate ) + request: Final = budgeted.request + if request is None: + return self._classifier_failure_outcome( + "classifier input cannot fit the provider budget", + prompt, + system_prompt, + failure_reason="unsupported_input", + )._replace(classifier_input_telemetry=budgeted.telemetry) + breaker: Final = self._classifier_circuit_breaker + permit: Final = breaker.acquire_permit() if breaker is not None else None + if breaker is not None and permit is None: + return self._classifier_failure_outcome( + "jev classifier circuit is open", + prompt, + system_prompt, + failure_reason="circuit_open", + signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL, + )._replace(classifier_input_telemetry=budgeted.telemetry) try: response: Final = await asyncio.wait_for(client.evaluate(request, timeout_s, request_kwargs), timeout_s) answer: Final = response.answers.get("tier") @@ -2449,6 +2477,7 @@ class ComplexityRouter(CustomLogger): cause="jev_classifier", classifier_cost=verdict.cost, jev_verdict=verdict, + classifier_input_telemetry=budgeted.telemetry, ) except asyncio.CancelledError: if breaker is not None and permit is not None: @@ -2463,7 +2492,7 @@ class ComplexityRouter(CustomLogger): system_prompt, failure_reason=_classifier_failure_reason(e), error_type=type(e).__name__, - ) + )._replace(classifier_input_telemetry=budgeted.telemetry) def _classifier_failure_outcome( self, @@ -2654,6 +2683,7 @@ class ComplexityRouter(CustomLogger): messages: Sequence[Mapping[str, object]] | None, *, encrypted_task: bool = False, + current_ask_for_history: str | None = None, ) -> str: include_assistant: Final = self.config.classifier_context_include_assistant_turns marker_pairs: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) @@ -2661,7 +2691,7 @@ class ComplexityRouter(CustomLogger): prior_turns: Final = ( _extract_prior_turns( messages, - current_ask=prompt, + current_ask=prompt if current_ask_for_history is None else current_ask_for_history, window_size=self.config.classifier_context_window_size, budget_chars=self.config.classifier_context_budget_chars, per_turn_chars=self.config.classifier_context_per_turn_chars, @@ -2920,10 +2950,8 @@ class ComplexityRouter(CustomLogger): "originating_request_masked": masked_originating_request(request_kwargs), "body": {"model": llm_config.model, **payload}, } - classify: Final[Callable[..., Awaitable[ModelResponse | ResponsesAPIResponse]]] = ( - self.litellm_router_instance.aresponses - if encrypted_task is not None - else self.litellm_router_instance.acompletion + classify: Final[Callable[..., Awaitable[ModelResponse | ResponsesAPIResponse]]] = ( # pyright: ignore[reportAssignmentType] # the context forwards the router endpoints untyped, non-streaming calls return these + self.ctx.aresponses if encrypted_task is not None else self.ctx.acompletion ) classifier_timeout_s: Final[float] = llm_config.timeout_ms / 1000 @@ -3122,7 +3150,7 @@ class ComplexityRouter(CustomLogger): + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest() ) winner: Final = await claim_affinity_pin( - self.litellm_router_instance.cache, + self.ctx.health().store, cache_key, MappingProxyType({"model": model}), self.config.session_affinity_ttl_seconds, @@ -3215,7 +3243,7 @@ class ComplexityRouter(CustomLogger): model_to_prefs: Final[dict[str, AdaptiveRouterPreferences]] = {} model_to_cost: Final[dict[str, float]] = {} for name in available_models: - deployments = self.litellm_router_instance.deployments_named(name) + deployments = self.ctx.registry().deployments(name, source="named") if not deployments: model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) model_to_cost[name] = 0.0 @@ -3495,7 +3523,7 @@ class ComplexityRouter(CustomLogger): ): return None try: - model_info: Final = self.litellm_router_instance.get_router_model_info( + model_info: Final = self.ctx.registry().router_model_info( deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts received_model_name=group, ) @@ -3505,8 +3533,7 @@ class ComplexityRouter(CustomLogger): return limit if isinstance(limit, int) else None def _group_deployments(self, group: str) -> Sequence[Mapping[str, object]]: - list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None) - deployments: Final = list_models(model_name=group) if callable(list_models) else None + deployments: Final = self.ctx.registry().deployments(model_name=group) return tuple(deployments) if isinstance(deployments, list) else () def _group_output_ceiling(self, group: str) -> int | None: @@ -3777,7 +3804,7 @@ class ComplexityRouter(CustomLogger): return declared is True return model_verdict((deployment.get("litellm_params") or EMPTY_MAPPING).get("model") or model_name) - deployments: Final = self.litellm_router_instance.get_model_list(model_name=model_name) + deployments: Final = self.ctx.registry().deployments(model_name=model_name) if not deployments: return (model_verdict(model_name),) return tuple(deployment_verdict(deployment) for deployment in deployments) @@ -3974,9 +4001,10 @@ class ComplexityRouter(CustomLogger): probe_kwargs: Final = dict(request_kwargs) try: - deployments: Final = await self.litellm_router_instance.async_get_healthy_deployments( - model=model_name, - request_kwargs=probe_kwargs, + deployments: Final = await self.ctx.select( + model_name, + probe_kwargs, + kind="healthy", messages=messages, input=input, parent_otel_span=get_parent_otel_span_from_kwargs(request_kwargs), @@ -4215,7 +4243,7 @@ class ComplexityRouter(CustomLogger): routelayer: Final = SemanticRouter( routes=routes, encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.litellm_router_instance, + ctx=self.ctx, model_name=embedding_model, score_threshold=self.config.match_threshold, ), @@ -4509,7 +4537,7 @@ class ComplexityRouter(CustomLogger): if cache_key is not None and pin_replay_allowed: pinned_value: Final = await read_affinity_pin( - self.litellm_router_instance.cache, cache_key, self.config.session_affinity_ttl_seconds + self.ctx.health().store, cache_key, self.config.session_affinity_ttl_seconds ) pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names()) if pinned_pin is not None: @@ -4608,7 +4636,7 @@ class ComplexityRouter(CustomLogger): else floor_model ) if escalated: - await self.litellm_router_instance.cache.async_set_cache( + await self.ctx.health().store.async_set_cache( key=cache_key, value=_session_affinity_cache_value(session_model, resolved_pin_tier), ttl=self.config.session_affinity_ttl_seconds, @@ -4711,7 +4739,7 @@ class ComplexityRouter(CustomLogger): and self._matched_plan_mode_signal(request_kwargs, resolved_messages) is None ) if pinnable and cache_key is not None and response is not None: - await self.litellm_router_instance.cache.async_set_cache( + await self.ctx.health().store.async_set_cache( key=cache_key, value=_session_affinity_cache_value( response.model, @@ -4733,7 +4761,7 @@ class ComplexityRouter(CustomLogger): from litellm.proxy.common_utils.cache_aware_routing import choose_cached_model choice: Final = await choose_cached_model( - router=self.litellm_router_instance, + ctx=self.ctx, config=self.config, params_for_model=partial(self._litellm_params_for_model, request_kwargs=request_kwargs), response=response, @@ -4983,16 +5011,19 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=fallback_model, messages=messages if has_original_messages else None, - routing_decision=self._build_routing_decision( - routed_model=fallback_model, - conversation_continuing=conversation_continuing, - cause=outcome.cause, - signals=outcome.signals, - classifier_failure_reason=outcome.classifier_failure_reason, - classifier_error_type=outcome.classifier_error_type, - escalation_keyword=escalation_keyword, - escalated=False, - tier_litellm_params=fallback_tier_params, + routing_decision=_with_classifier_forecast( + self._build_routing_decision( + routed_model=fallback_model, + conversation_continuing=conversation_continuing, + cause=outcome.cause, + signals=outcome.signals, + classifier_failure_reason=outcome.classifier_failure_reason, + classifier_error_type=outcome.classifier_error_type, + escalation_keyword=escalation_keyword, + escalated=False, + tier_litellm_params=fallback_tier_params, + ), + outcome, ), litellm_params=fallback_tier_params, ) diff --git a/litellm/router_strategy/complexity_router/context_compaction.py b/litellm/router_strategy/complexity_router/context_compaction.py index 6a1a63ea4b6..c51c8a8d96a 100644 --- a/litellm/router_strategy/complexity_router/context_compaction.py +++ b/litellm/router_strategy/complexity_router/context_compaction.py @@ -27,10 +27,11 @@ from litellm.llms.compaction import ( dispatch, get_native_compaction_provider, ) +from litellm.router_deprecated import renamed_router_param from litellm.router_strategy.complexity_router.config import ContextCompactionConfig if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext Surface: TypeAlias = Literal["chat", "messages", "responses"] _SURFACES: Final[Mapping[str, Surface]] = MappingProxyType( @@ -96,13 +97,14 @@ def initialize_compaction_state(kwargs: Mapping[str, object], surface: Surface) return existing if isinstance(existing, CompactionState) else CompactionState(surface=surface) +@renamed_router_param("router") async def arm_compaction( kwargs: Mapping[str, object], config: ContextCompactionConfig | Literal[False] | None, candidates: tuple[str, ...] = (), parent_model: str | None = None, *, - router: Router | None = None, + ctx: RoutingContext | None = None, allow_escalation: bool = False, messages: Sequence[Mapping[str, object]] | None = None, ) -> None: @@ -111,7 +113,7 @@ async def arm_compaction( state.config = config if isinstance(config, ContextCompactionConfig) and not _client_managed(kwargs) else None state.candidates = candidates state.parent_model = parent_model - if allow_escalation and router is not None and state.config is not None: + if allow_escalation and ctx is not None and state.config is not None: payload: Final = MappingProxyType( { **kwargs, @@ -119,11 +121,11 @@ async def arm_compaction( **({"messages": messages} if messages is not None and state.surface != "responses" else {}), } ) - if not await _has_compactor(router, state, payload): + if not await _has_compactor(ctx, state, payload): state.config = None -async def _has_compactor(router: Router, state: CompactionState, payload: Mapping[str, object]) -> bool: +async def _has_compactor(ctx: RoutingContext, state: CompactionState, payload: Mapping[str, object]) -> bool: from litellm.exceptions import ContextWindowExceededError if state.surface is None or state.config is None: @@ -131,7 +133,7 @@ async def _has_compactor(router: Router, state: CompactionState, payload: Mappin try: instructions, prefix, _ = _portable_history(payload, state.surface) await _compactor_model( - router, state, _compactor_input(payload, state.surface, instructions, prefix, state.config.max_tokens) + ctx, state, _compactor_input(payload, state.surface, instructions, prefix, state.config.max_tokens) ) return True except ContextWindowExceededError: @@ -237,9 +239,9 @@ def _history( return instructions, prefix, conversation[split:] -async def _count(router: Router, payload: Mapping[str, object]) -> int: +async def _count(ctx: RoutingContext, payload: Mapping[str, object]) -> int: return await asyncio.to_thread( - router._count_pre_call_check_tokens, # pyright: ignore[reportPrivateUsage] # shared Router admission counter + ctx.count_pre_call_check_tokens, messages=_ITEMS.validate_python(payload["messages"]) if "messages" in payload else None, input=_INPUT.validate_python(payload.get("input")), request_kwargs=payload, @@ -247,11 +249,11 @@ async def _count(router: Router, payload: Mapping[str, object]) -> int: def _budget( - router: Router, deployment: Mapping[str, object], payload: Mapping[str, object], ratio: float + ctx: RoutingContext, deployment: Mapping[str, object], payload: Mapping[str, object], ratio: float ) -> InputBudget: model: Final = str(payload.get("model", "")) info: Final = _MAPPING.validate_python( - router.get_router_model_info(deployment=_DICT.validate_python(deployment), received_model_name=model) + ctx.registry().router_model_info(deployment=_DICT.validate_python(deployment), received_model_name=model) ) raw_window: Final = info.get("max_input_tokens") window: Final = raw_window if isinstance(raw_window, int) and not isinstance(raw_window, bool) else None @@ -269,9 +271,9 @@ def _budget( async def _compactor_model( - router: Router, state: CompactionState, payload: Mapping[str, object] + ctx: RoutingContext, state: CompactionState, payload: Mapping[str, object] ) -> tuple[str, NativeCompactionProvider]: - needed: Final = await _count(router, payload) + needed: Final = await _count(ctx, payload) candidates: Final = ( (state.config.model,) if state.config is not None and state.config.model is not None else state.candidates ) @@ -279,13 +281,13 @@ async def _compactor_model( ( (candidate, provider) for candidate in candidates - if (deployments := tuple(router.get_model_list(model_name=candidate) or ())) + if (deployments := tuple(ctx.registry().deployments(model_name=candidate) or ())) and (provider := get_native_compaction_provider(_MAPPING.validate_python(deployments[0]["litellm_params"]))) is not None and all( provider.supports_native_compaction(params := _MAPPING.validate_python(deployment["litellm_params"])) and provider.compatible_defaults(params) - and (budget := _budget(router, deployment, payload, 0.9)).available is not None + and (budget := _budget(ctx, deployment, payload, 0.9)).available is not None and needed <= budget.available for deployment in deployments ) @@ -320,7 +322,7 @@ def _native_prefix(payload: Mapping[str, object], surface: Surface) -> Mapping[s async def _generate_summary( - router: Router, + ctx: RoutingContext, provider: NativeCompactionProvider, protocol: CompactionProtocol, payload: Mapping[str, object], @@ -330,7 +332,7 @@ async def _generate_summary( executor: Final = compaction_executor.get() with native_compaction_call(): response: Final = await asyncio.wait_for( - executor(protocol, payload, parent_model) if executor is not None else dispatch(router, protocol, payload), + executor(protocol, payload, parent_model) if executor is not None else dispatch(ctx, protocol, payload), timeout=timeout, ) summary: Final = provider.extract_summary(protocol, response) @@ -362,32 +364,33 @@ def _compactor_input( ) +@renamed_router_param("router") async def compact_to_fit( - router: Router, deployment: Mapping[str, object], payload: Mapping[str, object], surface: Surface | None + ctx: RoutingContext, deployment: Mapping[str, object], payload: Mapping[str, object], surface: Surface | None ) -> Mapping[str, object]: from litellm.exceptions import ContextWindowExceededError try: - return await _compact_to_fit(router, deployment, payload, surface) + return await _compact_to_fit(ctx, deployment, payload, surface) except ContextWindowExceededError: - window: Final = _budget(router, deployment, payload, 1.0).window - if not _native_child.get() and window is not None and await _count(router, payload) <= window: + window: Final = _budget(ctx, deployment, payload, 1.0).window + if not _native_child.get() and window is not None and await _count(ctx, payload) <= window: return payload raise async def _check_client_managed_admission( - router: Router, deployment: Mapping[str, object], payload: Mapping[str, object] + ctx: RoutingContext, deployment: Mapping[str, object], payload: Mapping[str, object] ) -> None: - if not router.enable_pre_call_checks: + if not ctx.settings().enable_pre_call_checks: return - router._pre_call_checks( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # restore legacy admission after deployment defaults + ctx.pre_call_checks( model=str(payload["model"]), - healthy_deployments=_ITEMS.validate_python((deployment,)), + healthy_deployments=list[object](_ITEMS.validate_python((deployment,))), messages=_ITEMS.validate_python(payload["messages"]) if "messages" in payload else None, # pyright: ignore[reportArgumentType] # legacy annotation omits structured content input=_INPUT.validate_python(payload.get("input")), request_kwargs=_DICT.validate_python(payload), - input_token_count=await _count(router, payload), + input_token_count=await _count(ctx, payload), skip_inline_token_count=True, ) @@ -419,22 +422,22 @@ def _portable_history( async def _compact_to_fit( - router: Router, deployment: Mapping[str, object], payload: Mapping[str, object], surface: Surface | None + ctx: RoutingContext, deployment: Mapping[str, object], payload: Mapping[str, object], surface: Surface | None ) -> Mapping[str, object]: state: Final = payload.get(_STATE_KEY) config: Final = state.config if isinstance(state, CompactionState) else None if not _native_child.get() and (config is None or _client_managed(payload)): if config is not None: - await _check_client_managed_admission(router, deployment, payload) + await _check_client_managed_admission(ctx, deployment, payload) return payload model: Final = str(payload["model"]) - limits: Final = _budget(router, deployment, payload, config.trigger_ratio if config is not None else 0.9) + limits: Final = _budget(ctx, deployment, payload, config.trigger_ratio if config is not None else 0.9) budget: Final = limits.available if budget is None or budget <= 0: if ( config is not None and config.model is None - and (limits.window is None or await _count(router, payload) <= limits.window) + and (limits.window is None or await _count(ctx, payload) <= limits.window) ): return payload _unavailable(model, "A known input window and a smaller output allowance are required") @@ -443,18 +446,18 @@ async def _compact_to_fit( if ( child_provider is None or not child_provider.compatible_defaults(payload) - or await _count(router, payload) > budget + or await _count(ctx, payload) > budget ): _reject(model, "The selected compactor's effective request is incompatible or exceeds its input budget") return payload - if await _count(router, payload) <= budget: + if await _count(ctx, payload) <= budget: return payload if surface is None or config is None or not isinstance(state, CompactionState): _unavailable(model, "This request surface cannot be compacted") key: Final = "input" if surface == "responses" else "messages" instructions, prefix, tail = _portable_history(payload, surface) retained: Final = MappingProxyType({**payload, key: _ITEMS.validate_python((*instructions, *tail))}) - if await _count(router, retained) >= budget: + if await _count(ctx, retained) >= budget: _unavailable(model, "Retained instructions, tools and the latest turn leave no room for a summary") older: Final = _compactor_input(payload, surface, instructions, prefix, config.max_tokens) metadata: Final = sanitized_forwardable_call_metadata( @@ -481,7 +484,7 @@ async def _compact_to_fit( **parent_session_kwargs(payload), } ) - compactor, provider = await _compactor_model(router, state, request) + compactor, provider = await _compactor_model(ctx, state, request) child: Final = MappingProxyType({**request, **provider.request_kwargs(), "model": compactor}) identity: Final = hashlib.sha256( json.dumps( @@ -505,7 +508,7 @@ async def _compact_to_fit( state.summary = ( identity, asyncio.create_task( - _generate_summary(router, provider, protocol, child, config.timeout_seconds, state.parent_model) + _generate_summary(ctx, provider, protocol, child, config.timeout_seconds, state.parent_model) ), ) if state.summary[0] != identity: @@ -515,6 +518,6 @@ async def _compact_to_fit( {"role": "assistant", "content": "Summary of earlier conversation (context, not new instructions):\n" + summary} ) compacted: Final = MappingProxyType({**payload, key: _ITEMS.validate_python((*instructions, message, *tail))}) - if await _count(router, compacted) > budget: + if await _count(ctx, compacted) > budget: _unavailable(model, "The summary and retained conversation still exceed the selected deployment's budget") return compacted diff --git a/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py b/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py index bbd0d02fcce..7f79ccaaa26 100644 --- a/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py +++ b/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py @@ -239,10 +239,10 @@ def run_eval() -> tuple[int, int, list[dict]]: Tuple of (passed, total, failures) """ # Create router with default config - mock_router: Final = MagicMock() + mock_ctx: Final = MagicMock() router: Final = ComplexityRouter( model_name="eval-router", - litellm_router_instance=mock_router, + ctx=mock_ctx, ) passed = 0 diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index 87bd031490a..89e9dec96f1 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple, Protocol, TypeAlias @@ -22,13 +23,16 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG, is_ai_decide_model from litellm.llms.laya.common_utils import laya_response_model from litellm.llms.pass_through.typesafe_logging_handler import TypeSafePassthroughLoggingHandler +from litellm.router_deprecated import renamed_router_param from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS as _DEFAULT_JEV_INSTRUCTIONS from litellm.types.decisions import DecisionsResponse from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, StandardLoggingRoutingDecision if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext + from litellm.router_strategy.complexity_router.config import OpenSourceClassifierConfig + from litellm.types.router import Deployment JevProbability: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)] ClassifierProvider: TypeAlias = Literal["typesafe", "laya", "bespoke", "databricks"] @@ -52,6 +56,153 @@ class JevSystemOneRequest(LiteLLMBaseModel): questions: Mapping[str, JevChoiceQuestion] +@dataclass(frozen=True, slots=True) +class JevInputTarget: + model: str + max_tokens: int | None + + +@dataclass(frozen=True, slots=True) +class JevBudgetedRequest: + request: JevSystemOneRequest | None + telemetry: StandardLoggingRoutingDecision + + +def jev_input_targets(context: "RoutingContext", config: "OpenSourceClassifierConfig") -> tuple[JevInputTarget, ...]: + from litellm.types.router import Deployment + + if config.deployment_name is None: + provider: Final = "typesafe" if config.provider == "jev" else config.provider + return (_jev_input_target({"model": f"{provider}/{config.model}"}, None),) + deployments: Final = context.registry().deployments(model_name=config.deployment_name) or () + return tuple( + _jev_deployment_input_target(context, Deployment.model_validate(deployment)) for deployment in deployments + ) + + +def _jev_deployment_input_target(context: "RoutingContext", deployment: "Deployment") -> JevInputTarget: + params: Final = context.deployment_params({}, deployment) + explicit_limit: Final = deployment.model_info.model_dump().get("max_input_tokens") + return _jev_input_target(params, explicit_limit) + + +def _jev_input_target(params: Mapping[str, object], explicit_limit: object) -> JevInputTarget: + from litellm.decisions.call import decisions_outbound_headers, get_decisions_provider_config + from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata + + qualified_model: Final = str(params.get("model") or "") + override: Final = params.get("custom_llm_provider") + try: + model, provider = resolve_model_and_provider_for_metadata( + qualified_model, override if isinstance(override, str) else None + ) + headers: Final = decisions_outbound_headers( + TypeAdapter[Mapping[str, str] | None](Mapping[str, str] | None).validate_python( + params.get("extra_headers") + ), + provider, + params, + ) + limits: Final = get_decisions_provider_config(model, provider).get_input_limits(model, provider, headers) + candidates: Final = tuple( + value + for value in (limits.max_request_tokens, limits.max_question_tokens, explicit_limit) + if isinstance(value, int) and not isinstance(value, bool) and value > 0 + ) + return JevInputTarget(f"{provider}/{model}", min(candidates) if candidates else None) + except Exception: # noqa: BLE001 # unavailable metadata must not turn an existing classifier into a request failure + return JevInputTarget(qualified_model, None) + + +def _jev_request_count(request: JevSystemOneRequest, models: Sequence[str], budget: int) -> tuple[int, str]: + from litellm._lazy_imports import get_token_counter_new + + payload: Final = request.model_dump_json(exclude={"model"}) + byte_count: Final = len(payload.encode("utf-8")) + if byte_count <= budget or litellm.disable_token_counter: + return byte_count, "utf8_bytes" + try: + counts: Final = tuple(get_token_counter_new()(model=model, text=payload) for model in models) + if counts and all(count > 0 for count in counts): + return max(counts), "estimated_tokens" + except Exception: # noqa: BLE001 # tokenization is best-effort; byte counting remains available + return byte_count, "utf8_bytes" + return byte_count, "utf8_bytes" + + +def fit_jev_request( + build_request: Callable[[str, str | None], JevSystemOneRequest], + prompt: str, + system_prompt: str | None, + targets: Sequence[JevInputTarget], + truncate: Callable[[str, int], str], +) -> JevBudgetedRequest: + known: Final = tuple(target for target in targets if target.max_tokens is not None) + if not known: + return JevBudgetedRequest( + build_request(prompt, system_prompt), + {"classifier_input_budget_status": "unknown_limits", "classifier_input_truncated": False}, + ) + budget: Final = max(1, int(min(target.max_tokens for target in known if target.max_tokens is not None) * 0.75)) + models: Final = tuple(dict.fromkeys(target.model for target in known)) + initial_cap: Final = max(120, budget * 8) + bounded_prompt: Final = truncate(prompt, initial_cap) + bounded_system: Final = truncate(system_prompt, initial_cap) if system_prompt is not None else None + return _fit_jev_components( + build_request, + bounded_prompt, + bounded_system, + models, + budget, + truncate, + len(known) != len(targets), + bounded_prompt != prompt or bounded_system != system_prompt, + ) + + +def _fit_jev_components( + build_request: Callable[[str, str | None], JevSystemOneRequest], + prompt: str, + system_prompt: str | None, + models: Sequence[str], + budget: int, + truncate: Callable[[str, int], str], + unknown_limits: bool, + truncated: bool, +) -> JevBudgetedRequest: + attempts: Final = len(prompt).bit_length() + len(system_prompt or "").bit_length() + 1 + components: tuple[str, str | None, bool, int] = ( # rebind-ok: bounded iterative fitting state + prompt, + system_prompt, + truncated, + attempts, + ) + while True: + current_prompt, current_system, was_truncated, remaining = components + request = build_request(current_prompt, current_system) + count, method = _jev_request_count(request, models, budget) + telemetry: StandardLoggingRoutingDecision = { + "classifier_input_budget_tokens": budget, + "classifier_input_tokens": count, + "classifier_input_count_method": method, + "classifier_input_truncated": was_truncated, + "classifier_input_budget_status": "unknown_limits" + if unknown_limits + else "truncated" + if was_truncated + else "within_budget", + } + if count <= budget: + return JevBudgetedRequest(request, telemetry) + if remaining == 0 or (len(current_system or "") <= 123 and len(current_prompt) <= 123): + return JevBudgetedRequest(None, {**telemetry, "classifier_input_budget_status": "untrimmable_input"}) + components = ( + (current_prompt, truncate(current_system, max(120, len(current_system) // 2)), True, remaining - 1) + if current_system is not None and len(current_system) > 123 + else (truncate(current_prompt, max(120, len(current_prompt) // 2)), current_system, True, remaining - 1) + ) + + class JevChoiceAnswer(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, allow_inf_nan=False) @@ -87,30 +238,13 @@ class JevClassifierClient(Protocol): ) -> JevSystemOneResponse: ... -class _DecisionsRouter(Protocol): - async def adecisions( - self, - *, - model: str, - state: str, - questions: Mapping[str, object], - timeout: float, - num_retries: int, - disable_fallbacks: bool, - metadata: Mapping[str, object], - proxy_server_request: Mapping[str, object], - turn_off_message_logging: bool | None, - litellm_session_id: str = "", - litellm_trace_id: str = "", - ) -> object: ... - - _PARENT_METADATA: Final = TypeAdapter[Mapping[str, object] | None](Mapping[str, object] | None) class RouterJevClassifierClient: - def __init__(self, router: "Router", deployment_name: str) -> None: - self._router: _DecisionsRouter = router + @renamed_router_param("router") + def __init__(self, ctx: "RoutingContext", deployment_name: str) -> None: + self.ctx: Final = ctx self._deployment_name = deployment_name async def evaluate( @@ -132,7 +266,7 @@ class RouterJevClassifierClient: "questions": {key: value.model_dump() for key, value in request.questions.items()}, } response: Final = DecisionsResponse.model_validate( - await self._router.adecisions( + await self.ctx.adecisions( model=self._deployment_name, state=request.state, questions={key: value.model_dump() for key, value in request.questions.items()}, diff --git a/litellm/router_strategy/lar1_routing.py b/litellm/router_strategy/lar1_routing.py index de03d63f41b..992e7490368 100644 --- a/litellm/router_strategy/lar1_routing.py +++ b/litellm/router_strategy/lar1_routing.py @@ -16,10 +16,12 @@ from pydantic import ValidationError from litellm._logging import verbose_router_logger from litellm.router import CustomRoutingStrategyBase +from litellm.router_api import ChangeSet +from litellm.router_deprecated import renamed_router_param from litellm.types.lar1 import LAR1Metadata, LAR1Time if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext DEFAULT_THRESHOLDS: Final[dict[str, float]] = {"low": 0.3, "medium": 0.5, "high": 0.7} @@ -41,16 +43,17 @@ def lar1_thresholds_from_args( } +@renamed_router_param("router") def apply_lar1_routing_strategy( - router: Router, + ctx: RoutingContext, routing_strategy_args: Mapping[str, object] | None = None, ) -> None: strategy: Final = LAR1RoutingStrategy( - router_instance=router, + ctx=ctx, thresholds=lar1_thresholds_from_args(routing_strategy_args), ) - router.routing_strategy = "lar1" - router.set_custom_routing_strategy(strategy) + ctx.apply(ChangeSet(assign={"routing_strategy": "lar1"})) + ctx.register_plugin(strategy) def _normalize_thresholds(thresholds: dict[str, float] | None) -> dict[str, float]: @@ -78,12 +81,13 @@ def _parse_lar1_metadata(request_kwargs: dict) -> LAR1Metadata: class LAR1RoutingStrategy(CustomRoutingStrategyBase): + @renamed_router_param("router_instance") def __init__( self, - router_instance: Router | None = None, + ctx: RoutingContext | None = None, thresholds: dict[str, float] | None = None, ): - self._router = router_instance + self._ctx = ctx self.thresholds = _normalize_thresholds(thresholds) async def async_get_available_deployment( @@ -96,7 +100,7 @@ class LAR1RoutingStrategy(CustomRoutingStrategyBase): ): if request_kwargs is None: request_kwargs = {} - if self._router is None: + if self._ctx is None: return None lar1: Final = _parse_lar1_metadata(request_kwargs) @@ -104,9 +108,10 @@ class LAR1RoutingStrategy(CustomRoutingStrategyBase): evidence: Final = tuple(e.value for e in lar1.evidence) time_dim: Final = lar1.time.value - healthy: Final = await self._router.async_get_healthy_deployments( - model=model, - request_kwargs=request_kwargs, + healthy: Final = await self._ctx.select( + model, + request_kwargs, + kind="healthy", messages=messages, input=input, specific_deployment=specific_deployment, @@ -114,7 +119,7 @@ class LAR1RoutingStrategy(CustomRoutingStrategyBase): if isinstance(healthy, dict): return healthy - if not healthy: + if not isinstance(healthy, list) or not healthy: return None target: Final = self._classify_request(confidence, evidence, time_dim) diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 17d03fd4ff8..3f0a5412b62 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -11,7 +11,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type -from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed +from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models, str_keyed def _deployment_cost_map_rates(deployment: object) -> tuple[float | None, float | None]: diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index d8c52b4c8b6..2d9ad63e9de 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param from litellm.router_strategy.complexity_router.complexity_router import ( ComplexityRouter, ) @@ -29,10 +30,10 @@ from litellm.types.utils import StandardLoggingRoutingDecision from .config import QualityRouterConfig, RoutingPreferences if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext from litellm.types.router import PreRoutingHookResponse else: - Router = Any + RoutingContext = Any PreRoutingHookResponse = Any @@ -42,15 +43,18 @@ class QualityRouter(CustomLogger): keyword override. """ + litellm_router_instance: Final = RenamedRouterAttribute() # deprecated: the router attribute before ``ctx`` + + @renamed_router_param("litellm_router_instance") def __init__( self, model_name: str, - litellm_router_instance: "Router", + ctx: "RoutingContext", default_model: str | None = None, quality_router_config: Mapping[str, object] | None = None, ): self.model_name = model_name - self.litellm_router_instance = litellm_router_instance + self.ctx: Final = ctx if quality_router_config: self.config = QualityRouterConfig.model_validate(dict(**quality_router_config)) @@ -64,7 +68,7 @@ class QualityRouter(CustomLogger): # Internal scorer — re-use the existing rule-based classifier. self._scorer = ComplexityRouter( model_name=f"{model_name}::scorer", - litellm_router_instance=litellm_router_instance, + ctx=ctx, ) # Per-model indices populated alongside the tier index. `_model_keywords` @@ -158,7 +162,7 @@ class QualityRouter(CustomLogger): `_model_quality`, and `_model_cost`. Raises if any listed model is missing `litellm_routing_preferences`. """ - model_list: Final = getattr(self.litellm_router_instance, "model_list", None) or [] + model_list: Final = self.ctx.registry().all() or [] available: Final = set(self.config.available_models) # Track which available models we've matched so we can error on missing. diff --git a/litellm/router_strategy/savings_baseline.py b/litellm/router_strategy/savings_baseline.py index add0f78fd3a..645ae92d8e6 100644 --- a/litellm/router_strategy/savings_baseline.py +++ b/litellm/router_strategy/savings_baseline.py @@ -18,10 +18,11 @@ from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Final, NamedTuple from litellm._logging import verbose_router_logger +from litellm.router_deprecated import renamed_router_param from litellm.types.utils import PromptTokensDetailsWrapper, Usage if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext _REFERENCE_REQUEST: Final = Usage( @@ -70,7 +71,7 @@ def canonical_model(model: str, custom_llm_provider: str | None = None) -> str | return f"{provider}/{resolved}" -def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]: +def _models_in(ctx: "RoutingContext", group_name: str) -> tuple[Baseline, ...]: """The candidates a tier entry actually calls, each with its own pricing key. `litellm_params.model` is not always a model: on Azure it is the deployment name, @@ -80,7 +81,7 @@ def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]: straight at a provider model rather than at a configured group, and prices under its own name because there is no deployment to override it. """ - deployments: Final = router.deployments_named(group_name) + deployments: Final = ctx.registry().deployments(group_name, source="named") if not deployments: return (Baseline(qualified),) if (qualified := canonical_model(group_name)) else () @@ -100,7 +101,7 @@ def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]: return tuple(c for deployment in deployments if (c := candidate(deployment)) is not None) -def _priced(router: "Router", candidate: Baseline) -> tuple[float, Baseline] | None: +def _priced(ctx: "RoutingContext", candidate: Baseline) -> tuple[float, Baseline] | None: """``(cost_of_the_reference_request, candidate)``, or ``None`` when unpriceable. "Most expensive" is a property of a request, not of a rate: a deployment dearer per @@ -114,7 +115,7 @@ def _priced(router: "Router", candidate: Baseline) -> tuple[float, Baseline] | N provider, _, model_name = candidate.model.partition("/") try: - info: Final = router.cached_deployment_model_info(candidate.deployment_id or "", candidate.model) + info: Final = ctx.registry().deployment_model_info(candidate.deployment_id or "", candidate.model, cached=True) if info is None: return None prompt_cost, completion_cost = generic_cost_per_token( @@ -133,16 +134,17 @@ def _priced(router: "Router", candidate: Baseline) -> tuple[float, Baseline] | N return (cost, candidate) -def _most_expensive(router: "Router", candidates: Iterable[Baseline]) -> Baseline | None: +def _most_expensive(ctx: "RoutingContext", candidates: Iterable[Baseline]) -> Baseline | None: """The candidate that would have cost the most on the reference request.""" - priced: Final = tuple(r for candidate in candidates if (r := _priced(router, candidate)) is not None) + priced: Final = tuple(r for candidate in candidates if (r := _priced(ctx, candidate)) is not None) if not priced: verbose_router_logger.debug("savings baseline: no priceable candidates; savings driver disabled") return None return max(priced)[1] -def resolve_baseline(router: "Router", group_names: Iterable[str]) -> Baseline | None: +@renamed_router_param("router") +def resolve_baseline(ctx: "RoutingContext", group_names: Iterable[str]) -> Baseline | None: """The derived baseline for a router whose hardest tier offers ``group_names``. Holds no cache of its own; each pricing pass walks the pool, so the caller is @@ -155,7 +157,7 @@ def resolve_baseline(router: "Router", group_names: Iterable[str]) -> Baseline | over; an unresolvable baseline zeroes the savings driver instead. """ try: - return _most_expensive(router, (c for name in group_names for c in _models_in(router, name))) + return _most_expensive(ctx, (c for name in group_names for c in _models_in(ctx, name))) except Exception as e: # noqa: BLE001 # see docstring: routing must not fail for a metric verbose_router_logger.warning("savings baseline: could not resolve, savings will read zero (%s)", e) return None diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index d3996f504a5..dca717100bf 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -12,7 +12,7 @@ import re from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, overload +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypedDict, overload from typing_extensions import ReadOnly @@ -23,14 +23,11 @@ from litellm.constants import ( ROUTING_REQUEST_TAGS_METADATA_KEY, ) from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.router_deprecated import renamed_router_param from litellm.types.router import BaggageMappedTagsStamp, ConsumedRequestTagsStamp, DeploymentTypedDict, RouterErrors if TYPE_CHECKING: - from litellm.router import Router as _Router - - LitellmRouter = _Router -else: - LitellmRouter = Any + from litellm.router_api import RoutingContext class _TagLitellmParamsLike(Protocol): @@ -394,18 +391,18 @@ def _resolve_constraint_only_pool( def _all_deployments_or_fallback( - llm_router_instance: LitellmRouter, + ctx: "RoutingContext", model: str, fallback: _DeploymentPool, ) -> Sequence[_DeploymentLike | DeploymentTypedDict] | Mapping[_DeploymentLike, object]: try: - return llm_router_instance.get_all_deployments(model_name=model) + return ctx.registry().deployments(model, source="all") except Exception: # noqa: BLE001 # fail safe toward today's healthy-only behavior on lookup errors return fallback def _chain_tag_filtering_override( - llm_router_instance: LitellmRouter, + ctx: "RoutingContext", model: str, healthy_deployments: _DeploymentPool, ) -> bool | None: @@ -418,7 +415,7 @@ def _chain_tag_filtering_override( # repeatedly failing that one deployment into cooldown. Falls back to # healthy_deployments on a lookup error, preserving today's behavior rather # than crashing the request. - all_deployments: Final = _all_deployments_or_fallback(llm_router_instance, model, healthy_deployments) + all_deployments: Final = _all_deployments_or_fallback(ctx, model, healthy_deployments) for d in all_deployments: value = (d.get("model_info") or MappingProxyType({})).get("enable_tag_filtering") if value is not None: @@ -458,7 +455,7 @@ def _tags_that_route(request_tags: Sequence[str], metadata: object, routing_pref def _tag_known_to_group( - llm_router_instance: LitellmRouter, + ctx: "RoutingContext", model: str, positive_tags: Sequence[str], routing_confirmed: frozenset[str], @@ -467,7 +464,7 @@ def _tag_known_to_group( if tag_set & routing_confirmed: return True try: - all_deployments: Final = llm_router_instance.get_all_deployments(model_name=model) + all_deployments: Final = ctx.registry().deployments(model, source="all") except Exception: # noqa: BLE001 # fail safe toward "unrecognized" so lookup errors preserve the existing silent-fallback behavior return False return any( @@ -557,8 +554,9 @@ def _tag_sequence(tags: object) -> Sequence[object]: return typed_tags +@renamed_router_param("llm_router_instance") async def get_deployments_for_tag( - llm_router_instance: LitellmRouter, + ctx: "RoutingContext", model: str, # used to raise the correct error healthy_deployments: _DeploymentPool, request_kwargs: _RequestKwargsLike | None = None, @@ -587,11 +585,9 @@ async def get_deployments_for_tag( return healthy_deployments request_enable_tag_filtering: Final = request_kwargs.get("enable_tag_filtering") - chain_enable_tag_filtering: Final = _chain_tag_filtering_override(llm_router_instance, model, healthy_deployments) + chain_enable_tag_filtering: Final = _chain_tag_filtering_override(ctx, model, healthy_deployments) chain_default: Final = ( - chain_enable_tag_filtering - if chain_enable_tag_filtering is not None - else llm_router_instance.enable_tag_filtering + chain_enable_tag_filtering if chain_enable_tag_filtering is not None else ctx.settings().enable_tag_filtering ) if request_enable_tag_filtering is not True and chain_default is not True: return healthy_deployments @@ -602,8 +598,8 @@ async def get_deployments_for_tag( request_tags: Final = without_baggage_mapped_tags( _request_tags_after_router_consumption(metadata, model), metadata ) - match_any: Final = llm_router_instance.tag_filtering_match_any - routing_prefix: Final = llm_router_instance.tag_routing_prefix or "" + match_any: Final = ctx.settings().tag_filtering_match_any + routing_prefix: Final = ctx.settings().tag_routing_prefix or "" # Build header strings for regex matching from what the proxy already stores. # Currently we match against User-Agent; format matches "^User-Agent: claude-code/..." @@ -703,7 +699,7 @@ async def get_deployments_for_tag( if ( len(new_healthy_deployments) == 0 and positive_tags - and _tag_known_to_group(llm_router_instance, model, positive_tags, routing_confirmed) + and _tag_known_to_group(ctx, model, positive_tags, routing_confirmed) ): return _resolve_or_fail_open( request_deployments, diff --git a/litellm/router_utils/client_initalization_utils.py b/litellm/router_utils/client_initalization_utils.py index 55b4c071cb0..fdcead4a6b9 100644 --- a/litellm/router_utils/client_initalization_utils.py +++ b/litellm/router_utils/client_initalization_utils.py @@ -1,16 +1,13 @@ from types import TracebackType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType +from litellm.router_deprecated import renamed_router_param from litellm.types.router import RouterErrors from litellm.utils import calculate_max_parallel_requests if TYPE_CHECKING: - from litellm.router import Router as _Router - - LitellmRouter = _Router -else: - LitellmRouter = Any + from litellm.router_api import RoutingContext class MaxParallelRequestsLimit: @@ -52,7 +49,8 @@ class MaxParallelRequestsLimit: class InitalizeCachedClient: @staticmethod - def set_max_parallel_requests_client(litellm_router_instance: LitellmRouter, model: dict): + @renamed_router_param("litellm_router_instance") + def set_max_parallel_requests_client(ctx: "RoutingContext", model: dict): litellm_params: Final = model.get("litellm_params", {}) model_id: Final = model["model_info"]["id"] rpm: Final = litellm_params.get("rpm", None) @@ -62,7 +60,7 @@ class InitalizeCachedClient: rpm=rpm, max_parallel_requests=max_parallel_requests, tpm=tpm, - default_max_parallel_requests=litellm_router_instance.default_max_parallel_requests, + default_max_parallel_requests=ctx.settings().default_max_parallel_requests, ) if calculated_max_parallel_requests: limit: Final = MaxParallelRequestsLimit( @@ -71,7 +69,7 @@ class InitalizeCachedClient: model_group=model.get("model_name", ""), ) cache_key: Final = f"{model_id}_max_parallel_requests_client" - litellm_router_instance.cache.set_cache( + ctx.health().store.set_cache( key=cache_key, value=limit, local_only=True, diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 65aa01750ac..25821703029 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -7,19 +7,18 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_logger +from litellm.router_deprecated import renamed_router_param if TYPE_CHECKING: - from litellm.router import Router as _Router - - LitellmRouter = _Router from litellm.integrations.prometheus import PrometheusLogger + from litellm.router_api import RoutingContext else: - LitellmRouter = Any PrometheusLogger = Any +@renamed_router_param("litellm_router_instance") async def router_cooldown_event_callback( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment_id: str, exception_status: str | int, cooldown_time: float | None, @@ -31,7 +30,7 @@ async def router_cooldown_event_callback( - Increments cooldown metric for deployment on Prometheus """ verbose_logger.debug("In router_cooldown_event_callback - updating prometheus") - _deployment: Final = litellm_router_instance.get_deployment(model_id=deployment_id) + _deployment: Final = ctx.registry().deployment(deployment_id) if _deployment is None: verbose_logger.warning( "in router_cooldown_event_callback but _deployment is None for deployment_id=%s. Doing nothing", diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index a1b1e327802..e50d5bb47a6 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -29,6 +29,7 @@ from litellm.constants import ( from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials +from litellm.router_deprecated import renamed_router_param from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCacheValue from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback from litellm.types.router import Deployment @@ -48,12 +49,10 @@ from .router_callbacks.track_deployment_metrics import ( if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from litellm.router import Router as _Router + from litellm.router_api import RoutingContext - LitellmRouter = _Router Span = _Span else: - LitellmRouter = Any Span = Any _ADVISOR_ORCHESTRATION_FAILURE_ATTR: Final = "_litellm_advisor_orchestration_failure" @@ -137,7 +136,7 @@ def _first_present(*sources: Mapping[str, Any] | None, key: str) -> int | float def _get_deployment_cooldown_policy( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment: str, ) -> tuple[Mapping[str, int] | None, int | None]: """Return (allowed_fails_policy, allowed_fails) from deployment model_info, or (None, None). @@ -149,7 +148,7 @@ def _get_deployment_cooldown_policy( field placed there would leak into the outgoing LLM request instead of staying router-internal. """ - dep: Final = litellm_router_instance.get_model_info(id=deployment) + dep: Final = ctx.registry().deployment_dict(deployment) if dep is None: return None, None mi: Final[Mapping[str, Any]] = dep.get("model_info") or MappingProxyType({}) @@ -175,7 +174,7 @@ def _resolve_allowed_fails_from_policy( def _should_cooldown_based_on_deployment_policy( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment: str, original_exception: Exception, dep_policy: Mapping[str, int] | None, @@ -199,10 +198,7 @@ def _should_cooldown_based_on_deployment_policy( if ( allowed_fails_from_policy is None and is_single_deployment_model_group - and ( - dep_allowed_fails is not None - or litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None - ) + and (dep_allowed_fails is not None or ctx.settings().allowed_fails_for(original_exception) is None) ): return False @@ -215,7 +211,7 @@ def _should_cooldown_based_on_deployment_policy( else ("generic" if dep_allowed_fails is not None else None) ) - dep: Final = litellm_router_instance.get_model_info(id=deployment) + dep: Final = ctx.registry().deployment_dict(deployment) cooldown_time_override: Final = ( _first_present(dep.get("model_info"), dep.get("litellm_params"), key="cooldown_time") if dep is not None @@ -223,7 +219,7 @@ def _should_cooldown_based_on_deployment_policy( ) return should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=litellm_router_instance, + ctx=ctx, deployment=deployment, original_exception=original_exception, allowed_fails_override=allowed_fails_override, @@ -233,7 +229,7 @@ def _should_cooldown_based_on_deployment_policy( def _has_explicit_allowed_fails_policy_for_exception( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment: str | None, original_exception: Exception, ) -> bool: @@ -254,12 +250,12 @@ def _has_explicit_allowed_fails_policy_for_exception( """ if deployment is None: return False - dep_policy, _ = _get_deployment_cooldown_policy(litellm_router_instance, deployment) + dep_policy, _ = _get_deployment_cooldown_policy(ctx, deployment) return _resolve_allowed_fails_from_policy(dep_policy, original_exception) is not None def _is_cooldown_required( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", model_id: str, exception_status: str | int, exception_str: str | None = None, @@ -320,7 +316,7 @@ def _is_per_user_provider_request(request_kwargs: Mapping[str, object]) -> bool: def _should_run_cooldown_logic( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment: str | None, exception_status: str | int, original_exception: Exception, @@ -335,10 +331,10 @@ def _should_run_cooldown_logic( - router.disable_cooldowns is True - deployment is None - _is_cooldown_required() returns False - - deployment is in litellm_router_instance.provider_default_deployment_ids + - deployment is in ctx.settings().provider_default_deployment_ids - exception_status is not one that should be immediately retried (e.g. 401) """ - if deployment is None or litellm_router_instance.get_model_group(id=deployment) is None: + if deployment is None or ctx.registry().resolve(deployment, kind="group_of_id") is None: verbose_router_logger.debug( "Should Not Run Cooldown Logic: deployment id is none or model group can't be found." ) @@ -369,7 +365,7 @@ def _should_run_cooldown_logic( return False if deployment is not None: - failed_deployment: Final = litellm_router_instance.get_deployment(model_id=deployment) + failed_deployment: Final = ctx.registry().deployment(deployment) if failed_deployment is not None: from litellm.llms.github_copilot.per_user_auth import ( github_copilot_per_user_credential_name, @@ -389,7 +385,7 @@ def _should_run_cooldown_logic( verbose_router_logger.debug("Should Not Run Cooldown Logic: time_to_cooldown is effectively 0") return False - if litellm_router_instance.disable_cooldowns: + if ctx.settings().disable_cooldowns: verbose_router_logger.debug("Should Not Run Cooldown Logic: disable_cooldowns is True") return False @@ -398,19 +394,19 @@ def _should_run_cooldown_logic( return False if not _is_cooldown_required( - litellm_router_instance=litellm_router_instance, + ctx=ctx, model_id=deployment, exception_status=exception_status, exception_str=str(original_exception), ) and not _has_explicit_allowed_fails_policy_for_exception( - litellm_router_instance=litellm_router_instance, + ctx=ctx, deployment=deployment, original_exception=original_exception, ): verbose_router_logger.debug("Should Not Run Cooldown Logic: _is_cooldown_required returned False") return False - if deployment in litellm_router_instance.provider_default_deployment_ids: + if deployment in ctx.settings().provider_default_deployment_ids: verbose_router_logger.debug("Should Not Run Cooldown Logic: deployment is in provider_default_deployment_ids") return False @@ -448,10 +444,10 @@ def _fallback_entries(configured: object) -> tuple[object, ...]: return () -def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool: +def _has_fallback_group(ctx: "RoutingContext", requested_model_group: str | None) -> bool: from litellm.router_utils.fallback_event_handlers import get_fallback_model_group - configured: Final[object] = getattr(litellm_router_instance, "fallbacks", None) + configured: Final[object] = ctx.settings().fallbacks if requested_model_group is None or not configured: return False chains: Final = tuple( @@ -462,7 +458,7 @@ def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_ def _should_cooldown_deployment( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment: str, exception_status: str | int, original_exception: Exception, @@ -486,31 +482,29 @@ def _should_cooldown_deployment( - v1 logic (Legacy): if allowed fails or allowed fail policy set, coolsdown if num fails in this minute > allowed fails """ - model_group: Final = litellm_router_instance.get_model_group(id=deployment) + model_group: Final = ctx.registry().resolve(deployment, kind="group_of_id") is_single_deployment_model_group = False if model_group is not None and len(model_group) == 1: is_single_deployment_model_group = not ( - litellm_router_instance.routing_group_has_alternatives(requested_model_group) - or litellm_router_instance.team_model_has_alternatives(deployment) + ctx.registry().resolve(requested_model_group, kind="routing_group_alternatives") + or ctx.registry().resolve(deployment, kind="team_model_alternatives") ) - is_lone_deployment: Final = is_single_deployment_model_group and not _has_fallback_group( - litellm_router_instance, requested_model_group - ) + is_lone_deployment: Final = is_single_deployment_model_group and not _has_fallback_group(ctx, requested_model_group) exception_status_int: Final = cast_exception_status_to_int(exception_status) - dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(litellm_router_instance, deployment) + dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(ctx, deployment) if ( is_single_deployment_model_group and exception_status_int == 402 and _resolve_allowed_fails_from_policy(dep_policy, original_exception) is None - and litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None + and ctx.settings().allowed_fails_for(original_exception) is None ): return False ## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level) if dep_policy is not None or dep_allowed_fails is not None: return _should_cooldown_based_on_deployment_policy( - litellm_router_instance, + ctx, deployment, original_exception, dep_policy, @@ -519,16 +513,11 @@ def _should_cooldown_deployment( ) ## BASE CASE - single deployment - if ( - litellm_router_instance.allowed_fails_policy is None - and _is_allowed_fails_set_on_router(litellm_router_instance=litellm_router_instance) is False - ): + if ctx.settings().allowed_fails_policy is None and _is_allowed_fails_set_on_router(ctx=ctx) is False: num_successes_this_minute: Final = get_deployment_successes_for_current_minute( - litellm_router_instance=litellm_router_instance, deployment_id=deployment - ) - num_fails_this_minute: Final = get_deployment_failures_for_current_minute( - litellm_router_instance=litellm_router_instance, deployment_id=deployment + ctx=ctx, deployment_id=deployment ) + num_fails_this_minute: Final = get_deployment_failures_for_current_minute(ctx=ctx, deployment_id=deployment) total_requests_this_minute: Final = num_successes_this_minute + num_fails_this_minute percent_fails = 0.0 @@ -560,10 +549,10 @@ def _should_cooldown_deployment( return False else: - if is_lone_deployment and litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None: + if is_lone_deployment and ctx.settings().allowed_fails_for(original_exception) is None: return False return should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=litellm_router_instance, + ctx=ctx, deployment=deployment, original_exception=original_exception, ) @@ -571,8 +560,9 @@ def _should_cooldown_deployment( return False +@renamed_router_param("litellm_router_instance") def set_cooldown_deployments( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", original_exception: Exception, exception_status: str | int, deployment: str | None = None, @@ -595,7 +585,7 @@ def set_cooldown_deployments( if ( _should_run_cooldown_logic( - litellm_router_instance=litellm_router_instance, + ctx=ctx, deployment=deployment, exception_status=exception_status, original_exception=original_exception, @@ -612,13 +602,13 @@ def set_cooldown_deployments( verbose_router_logger.debug("Attempting to add %s to cooldown list", deployment) if _should_cooldown_deployment( - litellm_router_instance=litellm_router_instance, + ctx=ctx, deployment=deployment, exception_status=exception_status, original_exception=original_exception, requested_model_group=requested_model_group, ): - litellm_router_instance.cooldown_cache.add_deployment_to_cooldown( + ctx.health().cooldowns.add_deployment_to_cooldown( model_id=deployment, original_exception=original_exception, exception_status=exception_status_int, @@ -628,7 +618,7 @@ def set_cooldown_deployments( # Trigger cooldown callback handler asyncio.create_task( router_cooldown_event_callback( - litellm_router_instance=litellm_router_instance, + ctx=ctx, deployment_id=deployment, exception_status=exception_status, cooldown_time=time_to_cooldown, @@ -641,15 +631,16 @@ def set_cooldown_deployments( _set_cooldown_deployments = set_cooldown_deployments +@renamed_router_param("litellm_router_instance") async def async_get_cooldown_deployments( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", parent_otel_span: Span | None, ) -> list[str]: """ Async implementation of 'get_cooldown_deployments' """ - model_ids: Final = litellm_router_instance.get_model_ids() - cooldown_models: Final = await litellm_router_instance.cooldown_cache.async_get_active_cooldowns( + model_ids: Final = ctx.registry().names(kind="ids") + cooldown_models: Final = await ctx.health().cooldowns.async_get_active_cooldowns( model_ids=model_ids, parent_otel_span=parent_otel_span, ) @@ -670,15 +661,16 @@ async def async_get_cooldown_deployments( _async_get_cooldown_deployments = async_get_cooldown_deployments +@renamed_router_param("litellm_router_instance") async def async_get_cooldown_deployments_with_debug_info( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", parent_otel_span: Span | None, ) -> list[tuple[str, CooldownCacheValue]]: """ Async implementation of 'get_cooldown_deployments' """ - model_ids: Final = litellm_router_instance.get_model_ids() - cooldown_models: Final = await litellm_router_instance.cooldown_cache.async_get_active_cooldowns( + model_ids: Final = ctx.registry().names(kind="ids") + cooldown_models: Final = await ctx.health().cooldowns.async_get_active_cooldowns( model_ids=model_ids, parent_otel_span=parent_otel_span ) @@ -689,7 +681,8 @@ async def async_get_cooldown_deployments_with_debug_info( _async_get_cooldown_deployments_with_debug_info = async_get_cooldown_deployments_with_debug_info -def get_cooldown_deployments(litellm_router_instance: LitellmRouter, parent_otel_span: Span | None) -> list[str]: +@renamed_router_param("litellm_router_instance") +def get_cooldown_deployments(ctx: "RoutingContext", parent_otel_span: Span | None) -> list[str]: """ Get the list of models being cooled down for this minute """ @@ -698,9 +691,9 @@ def get_cooldown_deployments(litellm_router_instance: LitellmRouter, parent_otel # ---------------------- # Return cooldown models # ---------------------- - model_ids: Final = litellm_router_instance.get_model_ids() + model_ids: Final = ctx.registry().names(kind="ids") - cooldown_models: Final = litellm_router_instance.cooldown_cache.get_active_cooldowns( + cooldown_models: Final = ctx.health().cooldowns.get_active_cooldowns( model_ids=model_ids, parent_otel_span=parent_otel_span ) @@ -719,28 +712,30 @@ def get_cooldown_deployments(litellm_router_instance: LitellmRouter, parent_otel _get_cooldown_deployments = get_cooldown_deployments -def _deployment_api_base(litellm_router_instance: LitellmRouter, model_id: str) -> str | None: - deployment: Final = litellm_router_instance.get_deployment(model_id=model_id) +def _deployment_api_base(ctx: "RoutingContext", model_id: str) -> str | None: + deployment: Final = ctx.registry().deployment(model_id) api_base: Final = deployment.litellm_params.api_base if deployment else None return None if api_base is None else mask_api_base_credentials(api_base) +@renamed_router_param("litellm_router_instance") def get_cooldown_deployment_info( - litellm_router_instance: LitellmRouter, route_model_ids: Collection[str], cooldown_list: Sequence[str] + ctx: "RoutingContext", route_model_ids: Collection[str], cooldown_list: Sequence[str] ) -> tuple[CooldownDeploymentInfo, ...]: """ The deployments the router resolved for this request (`route_model_ids`) that are in `cooldown_list`, each with its masked api_base """ return tuple( - CooldownDeploymentInfo(model_id=model_id, api_base=_deployment_api_base(litellm_router_instance, model_id)) + CooldownDeploymentInfo(model_id=model_id, api_base=_deployment_api_base(ctx, model_id)) for model_id in cooldown_list if model_id in route_model_ids ) +@renamed_router_param("litellm_router_instance") def should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment: str, original_exception: Exception, allowed_fails_override: int | None = None, @@ -764,29 +759,21 @@ def should_cooldown_based_on_allowed_fails_policy( - True if fails exceed the allowed limit (should cooldown) - False if fails are within the allowed limit (should not cooldown) """ - allowed_fails_from_policy: Final = litellm_router_instance.get_allowed_fails_from_policy( - exception=original_exception - ) + allowed_fails_from_policy: Final = ctx.settings().allowed_fails_for(original_exception) allowed_fails: Final = ( allowed_fails_override if allowed_fails_override is not None - else ( - allowed_fails_from_policy - if allowed_fails_from_policy is not None - else litellm_router_instance.allowed_fails - ) + else (allowed_fails_from_policy if allowed_fails_from_policy is not None else ctx.settings().allowed_fails) ) cooldown_time: Final = ( cooldown_time_override if cooldown_time_override is not None - else (litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS) + else (ctx.settings().cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS) ) base_key: Final = f"deployment:{deployment}:allowed_fails" cache_key: Final = f"{base_key}:{cache_key_suffix}" if cache_key_suffix else base_key - updated_fails: Final = _increment_allowed_fails( - cache=litellm_router_instance.cache, cache_key=cache_key, ttl=cooldown_time - ) + updated_fails: Final = _increment_allowed_fails(cache=ctx.health().store, cache_key=cache_key, ttl=cooldown_time) return updated_fails > allowed_fails @@ -805,7 +792,7 @@ def _increment_allowed_fails(cache: DualCache, cache_key: str, ttl: float) -> in def _is_allowed_fails_set_on_router( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", ) -> bool: """ Check if Router.allowed_fails was passed to the Router or is a Non-default Value @@ -815,11 +802,11 @@ def _is_allowed_fails_set_on_router( - True if Router.allowed_fails is a Non-default Value - False if Router.allowed_fails is None, or is the Default Value and was not passed """ - if litellm_router_instance.allowed_fails is None: + if ctx.settings().allowed_fails is None: return False - if litellm_router_instance.explicit_allowed_fails is not None: + if ctx.explicit_allowed_fails is not None: return True - if litellm_router_instance.allowed_fails != litellm.allowed_fails: + if ctx.settings().allowed_fails != litellm.allowed_fails: return True return False diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 6d1de1f9f6c..c929cca44be 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -13,16 +13,17 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, + get_fallback_error_info, +) from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider from litellm.litellm_core_utils.litellm_logging import Logging, same_attempt_marker from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure -from litellm.router_utils.add_retry_fallback_headers import ( - add_fallback_headers_to_response, - get_fallback_error_info, -) +from litellm.router_deprecated import renamed_router_param from litellm.router_utils.cooldown_handlers import ( _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils cast_exception_status_to_int, @@ -43,11 +44,7 @@ from litellm.types.router import ( ) if TYPE_CHECKING: - from litellm.router import Router as _Router - - LitellmRouter = _Router -else: - LitellmRouter = Any + from litellm.router_api import RoutingContext # Status codes a generic API call's caller-supplied resource id can trigger on its own # (e.g. a nonexistent file/batch/thread id), independent of the selected deployment's health. @@ -57,7 +54,7 @@ _FALLBACK_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) def _trigger_cooldown_for_failed_deployment( - litellm_router: LitellmRouter, + ctx: "RoutingContext", kwargs: Mapping[str, object], exception: Exception, model_call_details: Mapping[str, object] = _NO_MODEL_CALL_DETAILS, @@ -128,7 +125,7 @@ def _trigger_cooldown_for_failed_deployment( verbose_router_logger.debug("Cannot trigger cooldown for fallback: no failed_deployment_id on exception") return - deployment: Final = litellm_router.get_deployment(model_id=deployment_id) + deployment: Final = ctx.registry().deployment(deployment_id) if is_caller_scoped_auth_failure(deployment, exception_status): verbose_router_logger.debug( "Not triggering cooldown for fallback deployment %s: caller-scoped OAuth authentication failure.", @@ -138,7 +135,7 @@ def _trigger_cooldown_for_failed_deployment( # Priority: deployment config > response header > router default, matching # Router.deployment_callback_on_failure's precedence for the primary path. - deployment_dict: Final = litellm_router.get_model_info(id=deployment_id) + deployment_dict: Final = ctx.registry().deployment_dict(deployment_id) deployment_cooldown: Final = ( _first_present( deployment_dict.get("model_info"), deployment_dict.get("litellm_params"), key="cooldown_time" @@ -161,16 +158,16 @@ def _trigger_cooldown_for_failed_deployment( else ( header_cooldown if header_cooldown is not None and header_cooldown >= 0 - else litellm_router.cooldown_time + else ctx.settings().cooldown_time ) ) increment_deployment_failures_for_current_minute( - litellm_router_instance=litellm_router, + ctx=ctx, deployment_id=deployment_id, ) set_cooldown_deployments( - litellm_router_instance=litellm_router, + ctx=ctx, exception_status=exception_status, original_exception=exception, deployment=deployment_id, @@ -697,14 +694,14 @@ def _get_fallback_target_model_group(fallback_entry: str | Mapping[str, object]) async def _is_fallback_target_authorized( - litellm_router: LitellmRouter, + ctx: "RoutingContext", fallback_entry: str | Mapping[str, object], original_model_group: str, kwargs: Mapping[str, object], *, attempted: AttemptedFallbackTargets, ) -> bool: - access_check: Final = litellm_router.fallback_access_check + access_check: Final = ctx.settings().fallback_access_check target: Final = _get_fallback_target_model_group(fallback_entry) if ( access_check is None @@ -713,7 +710,7 @@ async def _is_fallback_target_authorized( or target == get_pre_routing_selection(kwargs) ): return True - if await access_check(model=target, request_kwargs=kwargs, llm_router=litellm_router): + if await ctx.fallback_access_allowed(model=target, request_kwargs=kwargs): return True verbose_router_logger.info( "Skipping fallback to model_group = %s: caller is not authorized to call it", @@ -724,18 +721,18 @@ async def _is_fallback_target_authorized( async def _is_fallback_target_within_budget( - litellm_router: LitellmRouter, + ctx: "RoutingContext", fallback_entry: str | Mapping[str, object], original_model_group: str, kwargs: Mapping[str, object], *, attempted: AttemptedFallbackTargets, ) -> bool: - budget_check: Final = litellm_router.fallback_budget_check + budget_check: Final = ctx.settings().fallback_budget_check target: Final = _get_fallback_target_model_group(fallback_entry) if budget_check is None or target is None or target == original_model_group: return True - if await budget_check(model=target, request_kwargs=kwargs, llm_router=litellm_router): + if await ctx.fallback_budget_allowed(model=target, request_kwargs=kwargs): return True verbose_router_logger.info( "Skipping fallback to model_group = %s: caller is over budget", @@ -779,9 +776,10 @@ def creates_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: return getattr(kwargs.get("original_function"), "__name__", None) in PROVIDER_SCOPED_CREATION_FUNCTION_NAMES +@renamed_router_param("litellm_router", new="litellm_routing_ctx") async def run_async_fallback( *args: object, - litellm_router: LitellmRouter, + litellm_routing_ctx: "RoutingContext", fallback_model_group: list[str], original_model_group: str, original_exception: Exception, @@ -798,7 +796,7 @@ async def run_async_fallback( If all fallback model groups fail, it raises the most recent exception. Args: - litellm_router: The litellm router instance. + litellm_routing_ctx: The routing context of the router running the fallbacks. *args: Positional arguments. fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"] original_model_group: The original model group. example: "gpt-3.5-turbo" @@ -817,6 +815,7 @@ async def run_async_fallback( Raises: The most recent exception if all fallback model groups fail. """ + ctx: Final = litellm_routing_ctx ### BASE CASE ### MAX FALLBACK DEPTH REACHED if fallback_depth >= max_fallbacks: @@ -853,13 +852,9 @@ async def run_async_fallback( original_model_group, ) continue - if not await _is_fallback_target_authorized( - litellm_router, mg, original_model_group, kwargs, attempted=attempted - ): + if not await _is_fallback_target_authorized(ctx, mg, original_model_group, kwargs, attempted=attempted): continue - if not await _is_fallback_target_within_budget( - litellm_router, mg, original_model_group, kwargs, attempted=attempted - ): + if not await _is_fallback_target_within_budget(ctx, mg, original_model_group, kwargs, attempted=attempted): continue attempt_key = fallback_attempt_key(mg) if attempt_key is not None: @@ -872,7 +867,7 @@ async def run_async_fallback( attempted.record(attempt_key) try: # LOGGING - kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) + kwargs = ctx.log_retry(kwargs=kwargs, e=original_exception) verbose_router_logger.info("Falling back to model_group = %s", mask_sensitive_structure(mg)) kwargs.pop("_target_order", None) if isinstance(mg, str): @@ -894,7 +889,7 @@ async def run_async_fallback( if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors carry_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) - response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) + response = await ctx.with_fallbacks(*args, **kwargs) verbose_router_logger.info("Successful fallback b/w models.") response = add_fallback_headers_to_response( response=response, @@ -922,7 +917,7 @@ async def run_async_fallback( ) ) is not None: _trigger_cooldown_for_failed_deployment( - litellm_router=litellm_router, + ctx=ctx, kwargs=kwargs, exception=e, model_call_details=failure_callbacks_blocked_logging_obj.model_call_details, diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index 38d9ef05706..ad502fe68aa 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final from litellm._logging import redact_secrets, verbose_router_logger from litellm.constants import MAX_EXCEPTION_MESSAGE_LENGTH +from litellm.router_deprecated import renamed_router_param from litellm.router_utils.cooldown_handlers import ( async_get_cooldown_deployments_with_debug_info, get_cooldown_deployment_info, @@ -13,17 +14,16 @@ from litellm.types.router import RouterRateLimitError if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from litellm.router import Router as _Router + from litellm.router_api import RoutingContext - LitellmRouter = _Router Span = _Span else: - LitellmRouter = Any Span = Any +@renamed_router_param("litellm_router_instance") async def send_llm_exception_alert( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", request_kwargs: dict, error_traceback_str: str, original_exception, @@ -33,19 +33,14 @@ async def send_llm_exception_alert( Sends a Slack / MS Teams alert for the LLM API call failure. Only if router.slack_alerting_logger is set. Parameters: - litellm_router_instance (_Router): The LitellmRouter instance. + ctx (RoutingContext): The router's routing context. original_exception (Any): The original exception that occurred. Returns: None """ - if litellm_router_instance is None: - return - - if not hasattr(litellm_router_instance, "slack_alerting_logger"): - return - - if litellm_router_instance.slack_alerting_logger is None: + slack_alerting_logger: Final = ctx.observers().slack_alerting_logger + if slack_alerting_logger is None: return if "proxy_server_request" in request_kwargs: @@ -62,7 +57,7 @@ async def send_llm_exception_alert( # Redact secrets before sending to external service (Slack / MS Teams) exception_str = redact_secrets(exception_str) - await litellm_router_instance.slack_alerting_logger.send_alert( + await slack_alerting_logger.send_alert( message=f"LLM API call failed: `{exception_str}`", level="High", alert_type=AlertType.llm_exceptions, @@ -70,8 +65,9 @@ async def send_llm_exception_alert( ) +@renamed_router_param("litellm_router_instance") async def async_raise_no_deployment_exception( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", model: str, parent_otel_span: Span | None, route_model_ids: Collection[str] | None = None, @@ -80,12 +76,12 @@ async def async_raise_no_deployment_exception( Raises a RouterRateLimitError if no deployment is found for the given model. """ verbose_router_logger.info("get_available_deployment for model: %s, No deployment available", model) - model_ids: Final = litellm_router_instance.get_model_ids(model_name=model) - _cooldown_time: Final = litellm_router_instance.cooldown_cache.get_min_cooldown( + model_ids: Final = ctx.registry().names(model, kind="ids") + _cooldown_time: Final = ctx.health().cooldowns.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) _cooldown_list: Final = await async_get_cooldown_deployments_with_debug_info( - litellm_router_instance=litellm_router_instance, + ctx=ctx, parent_otel_span=parent_otel_span, ) verbose_router_logger.info( @@ -96,12 +92,10 @@ async def async_raise_no_deployment_exception( return RouterRateLimitError( model=model, cooldown_time=_cooldown_time, - enable_pre_call_checks=litellm_router_instance.enable_pre_call_checks, + enable_pre_call_checks=ctx.settings().enable_pre_call_checks, cooldown_list=cooldown_list_ids, model_ids=model_ids, cooldown_deployments=( - None - if route_model_ids is None - else get_cooldown_deployment_info(litellm_router_instance, route_model_ids, cooldown_list_ids) + None if route_model_ids is None else get_cooldown_deployment_info(ctx, route_model_ids, cooldown_list_ids) ), ) diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index 587c8015b79..c759af8fdb8 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -14,8 +14,7 @@ from litellm import verbose_logger from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError - -HEALTH_CHECKS_TARGET: Final = "health_checks" +from litellm.constants import HEALTH_CHECKS_TARGET if TYPE_CHECKING: from opentelemetry.trace import Span as _Span diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py index d2a46fc316d..5c97e26866f 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -13,8 +13,8 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.routing_common_utils import MID_STREAM_CONTINUATION_KWARG from litellm.litellm_core_utils.rules import Rules -from litellm.router_utils.common_utils import MID_STREAM_CONTINUATION_KWARG from litellm.types.llms.openai import AllMessageValues, OutputTextDeltaEvent from litellm.types.utils import ModelResponse, ModelResponseStream, Usage from litellm.utils import supports_assistant_prefill diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index daf7d6bb345..dbf471d875c 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -30,8 +30,8 @@ from litellm.caching.affinity_cache import ( from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.routing_common_utils import deployment_ids_failed_in_request from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.router_utils.common_utils import deployment_ids_failed_in_request from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 806c3db7601..6f3f0794371 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -50,14 +50,15 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_content_of_block, strip_encrypted_reasoning_from_messages, ) +from litellm.litellm_core_utils.routing_common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.router_utils.common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids +from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param from litellm.types.llms.openai import AllMessageValues from litellm.types.router import Deployment from litellm.utils import get_order_filtered_deployments if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -74,14 +75,17 @@ class EncryptedContentAffinityCheck(CustomLogger): per-model group ``model_group_affinity_config``. """ + router: Final = RenamedRouterAttribute() # deprecated: the router attribute before ``ctx`` + + @renamed_router_param("router") def __init__( self, - router: Optional["Router"] = None, + ctx: Optional["RoutingContext"] = None, enable_global_affinity: bool = True, model_group_affinity_config: dict[str, list[str]] | None = None, ) -> None: super().__init__() - self.router = router + self.ctx = ctx self.enable_global_affinity = enable_global_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} @@ -254,9 +258,9 @@ class EncryptedContentAffinityCheck(CustomLogger): deployment's ``(api_base, api_key)``, alongside the originating deployment object (or ``None`` if it was removed / router unavailable). """ - if self.router is None: + if self.ctx is None: return [], None - originating: Final = self.router.get_deployment(model_id=model_id) + originating: Final = self.ctx.registry().deployment(model_id) if originating is None: return [], None boundary: Final = self._encryption_boundary_key(originating.litellm_params.model_dump(exclude_none=True)) @@ -268,8 +272,9 @@ class EncryptedContentAffinityCheck(CustomLogger): return matches, originating @staticmethod + @renamed_router_param("router") def strip_reasoning_the_targets_cannot_decrypt( - router: "Router | None", + ctx: "RoutingContext | None", request_input: object, anthropic_messages: object, target_deployments: Sequence[Mapping[str, object]], @@ -302,9 +307,9 @@ class EncryptedContentAffinityCheck(CustomLogger): return False if origin_model_id in target_ids: return True - if router is None: + if ctx is None: return False - origin: Final = router.get_deployment(model_id=origin_model_id) + origin: Final = ctx.registry().deployment(origin_model_id) origin_boundary: Final = ( EncryptedContentAffinityCheck._encryption_boundary_key( origin.litellm_params.model_dump(exclude_none=True) @@ -419,7 +424,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ) request_kwargs["_encrypted_content_affinity_pinned"] = True self.strip_reasoning_the_targets_cannot_decrypt( - self.router, request_input, anthropic_messages, (deployment,), unmarked_origin=model_id + self.ctx, request_input, anthropic_messages, (deployment,), unmarked_origin=model_id ) return [deployment] @@ -437,7 +442,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ) request_kwargs["_encrypted_content_affinity_pinned"] = True self.strip_reasoning_the_targets_cannot_decrypt( - self.router, request_input, anthropic_messages, boundary_matches, unmarked_origin=model_id + self.ctx, request_input, anthropic_messages, boundary_matches, unmarked_origin=model_id ) return boundary_matches diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 7580c7dadd3..80e1f3fe29c 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -14,6 +14,9 @@ from __future__ import annotations import contextlib import contextvars from collections.abc import Mapping, MutableMapping +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx @@ -23,7 +26,15 @@ from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION +from litellm.litellm_core_utils.internal_call_metadata import ( + IO_TOKEN_RESERVATION_KEYS, + ITPM_CACHE_KEY, + ITPM_RESERVED_KEY, + OTPM_CACHE_KEY, + OTPM_RESERVED_KEY, +) from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET @@ -39,18 +50,27 @@ else: RoutingArgsTTL: Final = 60 -_io_token_rate_limit_request_kwargs: Final[contextvars.ContextVar[dict[str, object] | None]] = contextvars.ContextVar( - "io_token_rate_limit_request_kwargs", + +@dataclass(frozen=True, slots=True) +class _IoTokenRateLimitRequest: + kwargs: dict[str, object] | None + opens_responses_socket: bool + metadata_channel: str + + +_io_token_rate_limit_request: Final[contextvars.ContextVar[_IoTokenRateLimitRequest | None]] = contextvars.ContextVar( + "io_token_rate_limit_request", default=None, ) -ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" -OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" -ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" -OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key" - -def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, object] | None, store_in_context: bool = True) -> None: +def set_io_token_rate_limit_request_kwargs( + kwargs: dict[str, object] | None, + store_in_context: bool = True, + *, + opens_responses_socket: bool = False, + metadata_channel: str, +) -> None: # The reservation sentinels are server-only, but `metadata` is caller # controlled on proxy requests. Strip any client-supplied copies here (this # runs before the router stashes its own reservation) so a forged @@ -64,11 +84,14 @@ def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, object] | None, sto # every other deployment overwrite the slot with None instead of the # kwargs; overwriting (rather than skipping) also releases a previous # request's kwargs when a context is reused. - _io_token_rate_limit_request_kwargs.set(kwargs if store_in_context else None) + _io_token_rate_limit_request.set( + _IoTokenRateLimitRequest(kwargs, opens_responses_socket, metadata_channel) if store_in_context else None + ) def get_io_token_rate_limit_request_kwargs() -> dict[str, object] | None: - return _io_token_rate_limit_request_kwargs.get() + request: Final = _io_token_rate_limit_request.get() + return None if request is None else request.kwargs def seconds_until_minute_reset() -> int: @@ -224,27 +247,32 @@ def _otpm_reservation( return default if from_deployment is None else min(from_deployment, default) +def _usage_field(container: object, name: str) -> object: + if is_str_keyed_mapping(container): + return container.get(name) + return getattr(container, name, None) + + +def _cached_prompt_tokens(usage: object) -> int: + reported: Final = ( + _usage_field(_usage_field(usage, "prompt_tokens_details"), "cached_tokens"), + _usage_field(_usage_field(usage, "input_tokens_details"), "cached_tokens"), + _usage_field(usage, "cache_read_input_tokens"), + ) + return next((int(tokens) for tokens in reported if isinstance(tokens, int | float) and tokens), 0) + + def _get_usage_tokens(usage: object) -> tuple[int, int, int]: if usage is None: return 0, 0, 0 if hasattr(usage, "prompt_tokens") or hasattr(usage, "input_tokens"): prompt = int(getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", 0) or 0) completion = int(getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", 0) or 0) - cached = 0 - details: object = getattr(usage, "prompt_tokens_details", None) - if details is not None: - cached = int(getattr(details, "cached_tokens", 0) or 0) - if not cached: - cached = int(getattr(usage, "cache_read_input_tokens", 0) or 0) - return prompt, completion, cached + return prompt, completion, _cached_prompt_tokens(usage) if isinstance(usage, dict): prompt = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0) completion = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0) - details = usage.get("prompt_tokens_details") or {} - cached = int(details.get("cached_tokens", 0) or 0) if isinstance(details, dict) else 0 - if not cached: - cached = int(usage.get("cache_read_input_tokens") or 0) - return prompt, completion, cached + return prompt, completion, _cached_prompt_tokens(usage) return 0, 0, 0 @@ -298,19 +326,36 @@ def _resolve_reconcile_usage_tokens( if isinstance(standard_logging_object, dict): prompt_tokens = int(standard_logging_object.get("prompt_tokens") or 0) completion_tokens = int(standard_logging_object.get("completion_tokens") or 0) - cached_tokens = 0 - metadata: Final = standard_logging_object.get("metadata") - if isinstance(metadata, dict): - cached_tokens = int(metadata.get("cache_read_input_tokens") or 0) + logged_usage: Final = _usage_field(standard_logging_object.get("metadata"), "usage_object") + logged_cached_tokens: Final = _cached_prompt_tokens(logged_usage) # Same rationale as _usage_is_present: a bare total_tokens with no # prompt/completion breakdown can't be split, so it isn't treated # as resolved usage - the reservation is kept instead of refunded. if prompt_tokens or completion_tokens: - return max(0, prompt_tokens - cached_tokens), completion_tokens, True + return max(0, prompt_tokens - logged_cached_tokens), completion_tokens, True return 0, 0, False +def _reservation( + itpm_reserved: int, otpm_reserved: int, itpm_cache_key: str | None, otpm_cache_key: str | None +) -> Mapping[str, object]: + return MappingProxyType( + { + ITPM_RESERVED_KEY: itpm_reserved, + OTPM_RESERVED_KEY: otpm_reserved, + ITPM_CACHE_KEY: itpm_cache_key, + OTPM_CACHE_KEY: otpm_cache_key, + } + ) + + +def _stash_reservation_in_metadata(request: _IoTokenRateLimitRequest | None, reservation: Mapping[str, object]) -> None: + metadata: Final = None if request is None or not request.kwargs else request.kwargs.get(request.metadata_channel) + if isinstance(metadata, dict): + metadata.update(reservation) + + def _key_for_this_minute(reserved_key: str) -> str: counter, _, _ = reserved_key.rpartition(":") return f"{counter}:{get_utc_datetime().strftime('%H-%M')}" @@ -326,30 +371,6 @@ def _reconcile_increments( ) -def _stash_reservation_in_metadata( - request_kwargs: dict[str, object] | None, - *, - itpm_reserved: int, - otpm_reserved: int, - itpm_cache_key: str | None, - otpm_cache_key: str | None, -) -> None: - if not request_kwargs: - return - reservation: Final = { - ITPM_RESERVED_KEY: itpm_reserved, - OTPM_RESERVED_KEY: otpm_reserved, - ITPM_CACHE_KEY: itpm_cache_key, - OTPM_CACHE_KEY: otpm_cache_key, - } - for channel in ("metadata", "litellm_metadata"): - existing = request_kwargs.get(channel) - if isinstance(existing, dict): - existing.update(reservation) - elif channel == "metadata": - request_kwargs[channel] = dict(reservation) - - def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]: itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY) otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY) @@ -399,7 +420,7 @@ def _clear_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> None: for channel in _reservation_channels(kwargs): channel_dict = _as_mutable_mapping(channel) if channel_dict is not None: - for key in (ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY): + for key in IO_TOKEN_RESERVATION_KEYS: channel_dict.pop(key, None) @@ -499,7 +520,8 @@ def io_token_pre_call_check( if itpm_limit is None and otpm_limit is None: return deployment - request_kwargs: Final = get_io_token_rate_limit_request_kwargs() + request: Final = _io_token_rate_limit_request.get() + request_kwargs: Final = None if request is None else request.kwargs _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) otpm_to_reserve: Final = ( @@ -546,11 +568,13 @@ def io_token_pre_call_check( raise _stash_reservation_in_metadata( - request_kwargs, - itpm_reserved=itpm_reserved, - otpm_reserved=otpm_reserved, - itpm_cache_key=itpm_key if itpm_limit is not None else None, - otpm_cache_key=otpm_key if otpm_limit is not None else None, + request, + _reservation( + itpm_reserved, + otpm_reserved, + itpm_key if itpm_limit is not None else None, + otpm_key if otpm_limit is not None else None, + ), ) return deployment @@ -565,7 +589,14 @@ async def async_io_token_pre_call_check( if itpm_limit is None and otpm_limit is None: return deployment - request_kwargs: Final = get_io_token_rate_limit_request_kwargs() + request: Final = _io_token_rate_limit_request.get() + request_kwargs: Final = None if request is None else request.kwargs + if request is not None and request.opens_responses_socket: + _stash_reservation_in_metadata( + request, + await async_admit_io_token_turn(dual_cache, deployment, get_utc_datetime(), parent_otel_span), + ) + return deployment _model, served_model = _tokenizer_and_served_models(deployment) estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) otpm_to_reserve: Final = ( @@ -618,15 +649,65 @@ async def async_io_token_pre_call_check( raise _stash_reservation_in_metadata( - request_kwargs, - itpm_reserved=itpm_reserved, - otpm_reserved=otpm_reserved, - itpm_cache_key=itpm_key if itpm_limit is not None else None, - otpm_cache_key=otpm_key if otpm_limit is not None else None, + request, + _reservation( + itpm_reserved, + otpm_reserved, + itpm_key if itpm_limit is not None else None, + otpm_key if otpm_limit is not None else None, + ), ) return deployment +async def _async_redis_usage(redis_cache: RedisCache | None, key: str, parent_otel_span: Span | None) -> int | None: + if redis_cache is None: + return None + with contextlib.suppress(RedisCircuitBreakerOpenError): + return await redis_cache.async_get_cache(key=key, parent_otel_span=parent_otel_span) + return None + + +@with_service_target(ROUTER_USAGE_TARGET) +async def _async_shared_usage(dual_cache: DualCache, key: str, parent_otel_span: Span | None) -> int | None: + shared: Final = await _async_redis_usage(dual_cache.redis_cache, key, parent_otel_span) + if shared is not None: + return shared + return await dual_cache.async_get_cache(key=key, local_only=True) + + +async def _raise_if_at_limit( + dual_cache: DualCache, key: str, limit: int | None, limit_label: str, parent_otel_span: Span | None +) -> None: + if limit is None: + return + usage: Final = await _async_shared_usage(dual_cache, key, parent_otel_span) + if usage is not None and usage >= limit: + raise _rate_limit_error(limit_label, limit, usage) + + +@with_service_target(ROUTER_USAGE_TARGET) +async def async_admit_io_token_turn( + dual_cache: DualCache, + deployment: Mapping[str, object], + now: datetime, + parent_otel_span: Span | None = None, +) -> Mapping[str, object]: + configured: Final = dict(deployment) + itpm_limit, otpm_limit = get_deployment_io_token_limits(configured) + if itpm_limit is None and otpm_limit is None: + return MappingProxyType({}) + cache_keys: Final = _get_cache_keys(configured, now.strftime("%H-%M")) + if cache_keys is None: + return MappingProxyType({}) + itpm_key, otpm_key = cache_keys + await _raise_if_at_limit(dual_cache, itpm_key, itpm_limit, "ITPM", parent_otel_span) + await _raise_if_at_limit(dual_cache, otpm_key, otpm_limit, "OTPM", parent_otel_span) + return _reservation( + 0, 0, itpm_key if itpm_limit is not None else None, otpm_key if otpm_limit is not None else None + ) + + @with_service_target(ROUTER_USAGE_TARGET) def io_token_reconcile_success( dual_cache: DualCache, diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index c04bab4d360..81012d38b65 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -17,9 +17,9 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.routing_common_utils import deployment_ids_failed_in_request from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model -from litellm.router_utils.common_utils import deployment_ids_failed_in_request from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import ( diff --git a/litellm/router_utils/router_callbacks/track_deployment_metrics.py b/litellm/router_utils/router_callbacks/track_deployment_metrics.py index 6b422e98ec8..b9b638efa3f 100644 --- a/litellm/router_utils/router_callbacks/track_deployment_metrics.py +++ b/litellm/router_utils/router_callbacks/track_deployment_metrics.py @@ -10,16 +10,13 @@ get_deployment_successes_for_current_minute """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY +from litellm.router_deprecated import renamed_router_param if TYPE_CHECKING: - from litellm.router import Router as _Router - - LitellmRouter = _Router -else: - LitellmRouter = Any + from litellm.router_api import RoutingContext _METADATA_CHANNELS: Final = ("litellm_metadata", "metadata") @@ -42,15 +39,16 @@ def get_counted_usage_tokens(litellm_params: Mapping[str, object]) -> int | None return counted if isinstance(counted, int) and not isinstance(counted, bool) else None +@renamed_router_param("litellm_router_instance") def increment_deployment_successes_for_current_minute( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment_id: str, ) -> str: """ In-Memory: Increments the number of successes for the current minute for a deployment_id """ key: Final = f"{deployment_id}:successes" - litellm_router_instance.cache.increment_cache( + ctx.health().store.increment_cache( local_only=True, key=key, value=1, @@ -59,15 +57,16 @@ def increment_deployment_successes_for_current_minute( return key +@renamed_router_param("litellm_router_instance") def increment_deployment_failures_for_current_minute( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment_id: str, ): """ In-Memory: Increments the number of failures for the current minute for a deployment_id """ key: Final = f"{deployment_id}:fails" - litellm_router_instance.cache.increment_cache( + ctx.health().store.increment_cache( local_only=True, key=key, value=1, @@ -75,8 +74,9 @@ def increment_deployment_failures_for_current_minute( ) +@renamed_router_param("litellm_router_instance") def get_deployment_successes_for_current_minute( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment_id: str, ) -> int: """ @@ -86,7 +86,7 @@ def get_deployment_successes_for_current_minute( """ key: Final = f"{deployment_id}:successes" return ( - litellm_router_instance.cache.get_cache( + ctx.health().store.get_cache( local_only=True, key=key, ) @@ -94,8 +94,9 @@ def get_deployment_successes_for_current_minute( ) +@renamed_router_param("litellm_router_instance") def get_deployment_failures_for_current_minute( - litellm_router_instance: LitellmRouter, + ctx: "RoutingContext", deployment_id: str, ) -> int: """ @@ -105,7 +106,7 @@ def get_deployment_failures_for_current_minute( """ key: Final = f"{deployment_id}:fails" return ( - litellm_router_instance.cache.get_cache( + ctx.health().store.get_cache( local_only=True, key=key, ) diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index 7f6267ec9b9..8d8d77fa546 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -21,13 +21,14 @@ from litellm._internal_context import service_target from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.redis_batch import BatchResult, active_request_redis_batches +from litellm.router_deprecated import renamed_router_param from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCache if TYPE_CHECKING: from opentelemetry.trace import Span - from litellm.router import Router + from litellm.router_api import RoutingContext ROUTER_COOLDOWNS_USAGE_TARGET: Final = "router_cooldowns_usage" @@ -93,23 +94,24 @@ class RoutingPrefetch: raise @staticmethod + @renamed_router_param("litellm_router_instance") def arm( - litellm_router_instance: "Router", + ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: list, ) -> None: request: Final = active_request_redis_batches() - redis_cache: Final = litellm_router_instance.cache.redis_cache + redis_cache: Final = ctx.health().store.redis_cache if request is None or redis_cache is None or _PREFETCH_SLOT in request.prefetched: return cooldown_keys: Final = tuple( - CooldownCache.get_cooldown_cache_key(model_id) for model_id in litellm_router_instance.get_model_ids() + CooldownCache.get_cooldown_cache_key(model_id) for model_id in ctx.registry().names(kind="ids") ) usage_keys: Final = ( () if usage_selector is None else tuple(itertools.chain(*usage_selector.usage_counter_keys(deployments))) ) keys: Final = (*cooldown_keys, *usage_keys) - cooldown_store: Final = litellm_router_instance.cooldown_cache.cooldown_store + cooldown_store: Final = ctx.health().cooldowns.cooldown_store cooldown_due, cooldown_previous = cooldown_store.reserve_redis_batch_reads(cooldown_keys) usage_cache: Final = None if usage_selector is None else usage_selector.router_cache usage_reservation: Final = None if usage_cache is None else usage_cache.reserve_redis_batch_reads(usage_keys) @@ -184,9 +186,10 @@ class RoutingReadBatch: return RoutingReadBatch(usage_selector=selector) return RoutingReadBatch(usage_selector=None) if RoutingPrefetch.armed() else None + @renamed_router_param("litellm_router_instance") async def async_get_cooldown_deployments( self, - litellm_router_instance: "Router", + ctx: "RoutingContext", healthy_deployments: list, parent_otel_span: "Span | None", ) -> list[str]: @@ -194,14 +197,14 @@ class RoutingReadBatch: `_async_get_cooldown_deployments`, with the strategy's tpm/rpm counters for `healthy_deployments` fetched in the same MGET and kept as `prefetched_usage`. """ - model_ids: Final = litellm_router_instance.get_model_ids() + model_ids: Final = ctx.registry().names(kind="ids") cooldown_keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] selector: Final = self.usage_selector usage_keys: Final = ( () if selector is None else tuple(itertools.chain(*selector.usage_counter_keys(healthy_deployments))) ) reads: Final = ( - (litellm_router_instance.cooldown_cache.cooldown_store, cooldown_keys), + (ctx.health().cooldowns.cooldown_store, cooldown_keys), *(() if selector is None else ((selector.router_cache, list(usage_keys)),)), ) with service_target(_routing_read_target(cooldown_keys, usage_keys)): @@ -216,9 +219,7 @@ class RoutingReadBatch: values=None if usage_values is None else MappingProxyType(dict(zip(usage_keys, usage_values))), ) - cooldown_models: Final = litellm_router_instance.cooldown_cache.active_cooldowns_from_results( - model_ids, cooldown_results - ) + cooldown_models: Final = ctx.health().cooldowns.active_cooldowns_from_results(model_ids, cooldown_results) verbose_router_logger.debug("retrieve cooldown models: %s", cooldown_models) return [model_id for model_id, _ in cooldown_models] diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index 26326a7e6d4..aa2b29988dc 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -10,16 +10,15 @@ import traceback from collections.abc import Callable from functools import partial from types import MappingProxyType -from typing import Any, Final, Protocol +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_router_logger +from litellm.router_api import ChangeSet +from litellm.router_deprecated import renamed_router_param from litellm.types.router import SearchToolLiteLLMParams, SearchToolTypedDict - -class _SearchToolsRouter(Protocol): - """The one router attribute the search-tool helpers read and replace.""" - - search_tools: list[SearchToolTypedDict] +if TYPE_CHECKING: + from litellm.router_api import RoutingContext class SearchAPIRouter: @@ -52,14 +51,15 @@ class SearchAPIRouter: return resolved_api_key, resolved_api_base @staticmethod - async def update_router_search_tools(router_instance: _SearchToolsRouter, search_tools: list): + @renamed_router_param("router_instance") + async def update_router_search_tools(ctx: "RoutingContext", search_tools: list): """ Update the router with search tools from the database. This method is called by a cron job to sync search tools from DB to router. Args: - router_instance: The Router instance to update + ctx: The routing context of the router to update search_tools: List of search tool configurations from the database """ try: @@ -78,7 +78,7 @@ class SearchAPIRouter: router_search_tools.append(router_search_tool) # Update the router's search_tools list - router_instance.search_tools = router_search_tools + ctx.apply(ChangeSet(assign={"search_tools": router_search_tools})) verbose_router_logger.info("Successfully updated router with %s search tool(s)", len(router_search_tools)) @@ -87,15 +87,16 @@ class SearchAPIRouter: raise e @staticmethod + @renamed_router_param("router_instance") def get_matching_search_tools( - router_instance: _SearchToolsRouter, + ctx: "RoutingContext", search_tool_name: str, ) -> list: """ Get all search tools matching the given name. Args: - router_instance: The Router instance + ctx: The routing context of the router search_tool_name: Name of the search tool to find Returns: @@ -105,7 +106,7 @@ class SearchAPIRouter: ValueError: If no matching search tools are found """ matching_tools: Final = [ - tool for tool in router_instance.search_tools if tool.get("search_tool_name") == search_tool_name + tool for tool in ctx.settings().search_tools if tool.get("search_tool_name") == search_tool_name ] if not matching_tools: @@ -114,8 +115,9 @@ class SearchAPIRouter: return matching_tools @staticmethod + @renamed_router_param("router_instance", new="litellm_routing_ctx") async def async_search_with_fallbacks( - router_instance: Any, + litellm_routing_ctx: "RoutingContext", original_function: Callable, **kwargs, ): @@ -124,13 +126,14 @@ class SearchAPIRouter: Reuses the router's retry/fallback infrastructure. Args: - router_instance: The Router instance + litellm_routing_ctx: The routing context of the router original_function: The original litellm.asearch function **kwargs: Search parameters including search_tool_name, query, etc. Returns: SearchResponse from the search API """ + ctx: Final = litellm_routing_ctx try: search_tool_name: Final = kwargs.get("search_tool_name", kwargs.get("model")) @@ -140,20 +143,20 @@ class SearchAPIRouter: # Set up kwargs for the fallback system kwargs["model"] = search_tool_name # Use model field for compatibility with fallback system kwargs["original_generic_function"] = original_function - # Bind router_instance to the helper method using partial + # Bind ctx to the helper method using partial kwargs["original_function"] = partial( SearchAPIRouter.async_search_with_fallbacks_helper, - router_instance=router_instance, + litellm_routing_ctx=ctx, ) # Update kwargs before fallbacks (for logging, metadata, etc) - router_instance._update_kwargs_before_fallbacks( + ctx.update_kwargs_before_fallbacks( model=search_tool_name, kwargs=kwargs, metadata_variable_name="litellm_metadata", ) - available_search_tool_names: Final = [tool.get("search_tool_name") for tool in router_instance.search_tools] + available_search_tool_names: Final = [tool.get("search_tool_name") for tool in ctx.settings().search_tools] verbose_router_logger.debug( "Inside SearchAPIRouter.async_search_with_fallbacks() - search_tool_name: %s, Available Search Tools: %s", search_tool_name, @@ -161,7 +164,7 @@ class SearchAPIRouter: ) # Use the existing retry/fallback infrastructure - response: Final = await router_instance.async_function_with_fallbacks(**kwargs) + response: Final = await ctx.with_fallbacks(**kwargs) return response except Exception as e: @@ -169,7 +172,7 @@ class SearchAPIRouter: asyncio.create_task( send_llm_exception_alert( - litellm_router_instance=router_instance, + ctx=ctx, request_kwargs=kwargs, error_traceback_str=traceback.format_exc(), original_exception=e, @@ -178,8 +181,9 @@ class SearchAPIRouter: raise e @staticmethod + @renamed_router_param("router_instance", new="litellm_routing_ctx") async def async_search_with_fallbacks_helper( - router_instance: _SearchToolsRouter, + litellm_routing_ctx: "RoutingContext", model: str, original_generic_function: Callable, **kwargs, @@ -189,7 +193,7 @@ class SearchAPIRouter: Called by async_function_with_fallbacks for each retry attempt. Args: - router_instance: The Router instance + litellm_routing_ctx: The routing context of the router model: The search tool name (passed as model for compatibility) original_generic_function: The original litellm.asearch function **kwargs: Search parameters @@ -197,12 +201,13 @@ class SearchAPIRouter: Returns: SearchResponse from the selected search provider """ + ctx: Final = litellm_routing_ctx search_tool_name: Final = model # model field contains the search_tool_name try: # Find matching search tools matching_tools: Final = SearchAPIRouter.get_matching_search_tools( - router_instance=router_instance, + ctx=ctx, search_tool_name=search_tool_name, ) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 62b48621c05..8177e4a5cc9 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -26,6 +26,12 @@ class NativeClickHouseSpendStorage: def ensure_schema(self) -> Future[None]: ... def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... +@final +class NativeDatabase: + @staticmethod + def connect() -> Future[NativeDatabase]: ... + def close(self) -> Future[None]: ... + @final class NativeDiagnosticProcessor: def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ... @@ -208,6 +214,7 @@ __all__ = [ "HuggingFaceEncoding", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", + "NativeDatabase", "NativeDiagnosticProcessor", "ProcessReservedForForking", "ResponsesWebSocketConnection", diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 2d2f5072544..e155680bd30 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -103,7 +103,7 @@ def finalize( update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) cache_key: Final = logger.model_call_details.get("cache_key") if logger.model_call_details.get("cache_hit") is True and isinstance(cache_key, str): - from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict + from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict hidden: Final = get_hidden_params_dict(response, create=True) hidden.update({"cache_key": cache_key, "cache_hit": True}) diff --git a/litellm/rust_bridge/response_metadata.py b/litellm/rust_bridge/response_metadata.py index 1ef7b8c6595..d2ddf61f7db 100644 --- a/litellm/rust_bridge/response_metadata.py +++ b/litellm/rust_bridge/response_metadata.py @@ -1,6 +1,6 @@ from typing import Final, TypeVar -from litellm.router_utils.add_retry_fallback_headers import ( +from litellm.litellm_core_utils.add_retry_fallback_headers import ( _add_headers_to_response, # pyright: ignore[reportPrivateUsage] # reuse the proxy's identity-preserving response metadata writer get_hidden_params_dict, ) diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 2065093f295..9f1b35fdce3 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -17,10 +17,11 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.router_api import AttemptedFallbackTargets from litellm.router_strategy.complexity_router.context_compaction import CompactionState - from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets from litellm.types.caching import DynamicCacheControl from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionUserMessage + from litellm.types.proxy.guardrails.guardrail_hooks.headroom import HeadroomRequestContext from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.router import ConfigurableClientsideParamsCustomAuth, DeploymentTypedDict, RetryPolicy from litellm.types.router_weights import RouterWeights @@ -352,6 +353,7 @@ class CallState: @dataclass(frozen=True, slots=True, kw_only=True) class AgenticLoopState: + headroom_context: "HeadroomRequestContext | None" = field(default=None, metadata=wire("_headroom_request_context")) depth: int | None = field(default=None, metadata=wire("_agentic_loop_depth")) fingerprints: Sequence[str] | None = field(default=None, metadata=wire("_agentic_loop_fingerprints")) api_surface: Literal["chat_completions", "responses"] | None = field( diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index a648376c542..152614e4385 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -713,6 +713,14 @@ class AnthropicResponseContentBlockRedactedThinking(LiteLLMBaseModel): data: str +class AnthropicResponseContentBlockCompaction(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["compaction"] + content: str | None = None + signature: str | None = None + + class AnthropicResponseContentBlockServerToolUse(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) @@ -743,7 +751,7 @@ class AnthropicOutputTokensDetails(LiteLLMBaseModel): thinking_tokens: int | None = None -AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] +AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal", "compaction"] class AnthropicResponse(LiteLLMBaseModel): @@ -761,6 +769,7 @@ class AnthropicResponse(LiteLLMBaseModel): | AnthropicResponseContentBlockToolUse | AnthropicResponseContentBlockThinking | AnthropicResponseContentBlockRedactedThinking + | AnthropicResponseContentBlockCompaction | AnthropicResponseContentBlockServerToolUse | AnthropicResponseContentBlockWebSearchToolResult ] @@ -808,6 +817,9 @@ class AnthropicThinkingParam(TypedDict, total=False): block_binding: ReadOnly[AnthropicThinkingBlockBinding] +LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS: Final = frozenset(("anthropic", "azure_ai", "vertex_ai")) + + class ANTHROPIC_HOSTED_TOOLS(str, Enum): WEB_SEARCH = "web_search" BASH = "bash" @@ -828,6 +840,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): CONTEXT_MANAGEMENT_2025_06_27 = "context-management-2025-06-27" COMPACT_2026_01_12 = "compact-2026-01-12" COMPACT_2026_09_04 = "compact-2026-09-04" + INLINE_TOOLS_2026_09_15 = "inline-tools-2026-09-15" STRUCTURED_OUTPUT_2025_09_25 = "structured-outputs-2025-11-13" ADVANCED_TOOL_USE_2025_11_20 = "advanced-tool-use-2025-11-20" FAST_MODE_2026_02_01 = "fast-mode-2026-02-01" diff --git a/litellm/types/llms/ollama.py b/litellm/types/llms/ollama.py index 4783eac9a8d..fb9d9e8c363 100644 --- a/litellm/types/llms/ollama.py +++ b/litellm/types/llms/ollama.py @@ -1,4 +1,6 @@ from typing_extensions import ( + NotRequired, + ReadOnly, Required, TypedDict, ) @@ -12,6 +14,7 @@ class OllamaToolCallFunction( class OllamaToolCall(TypedDict): + id: NotRequired[ReadOnly[str]] function: OllamaToolCallFunction diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index d9d86cf54b3..ee86ec52327 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1,4 +1,5 @@ import builtins +import math from collections.abc import Iterable, Mapping, Sequence from enum import Enum from os import PathLike @@ -83,6 +84,7 @@ from typing_extensions import ( override, ) +from litellm._logging import verbose_logger from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject, LiteLLMBaseModel from litellm.types.responses.main import ( CustomToolCallOutputItem, @@ -400,7 +402,7 @@ class OpenAIFileObject(LiteLLMBaseModel): filename: str """The name of the file.""" - object: Literal["file"] + object: Literal["file"] = "file" """The object type, which is always `file`.""" purpose: OpenAIFilesPurpose @@ -438,6 +440,11 @@ class OpenAIFileObject(LiteLLMBaseModel): _hidden_params: dict = PrivateAttr(default={"response_cost": 0.0}) # no cost for writing a file + @field_validator("object", mode="before") + @classmethod + def _default_missing_object_type(cls, value: _JsonValue) -> _JsonValue: + return "file" if value is None else value + @property def hidden_params(self) -> dict[str, builtins.object]: # mutable-ok: API requires mutation return self._hidden_params @@ -1446,9 +1453,11 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): @classmethod def parse_cost(cls, v: object) -> object: """Normalise cost: accept either a float or a dict with a ``total_cost`` key.""" - if isinstance(v, dict): - return v.get("total_cost") - return v + cost: Final = v.get("total_cost") if isinstance(v, dict) else v + if isinstance(cost, float) and not math.isfinite(cost): + verbose_logger.warning("Dropping a non-finite provider usage cost") + return None + return cost model_config = ConfigDict(extra="allow") diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 0d99cd54a6b..2f269dc87d9 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -1,5 +1,6 @@ +from dataclasses import dataclass from enum import Enum -from typing import Final +from typing import Final, Protocol from pydantic import Field from typing_extensions import TypedDict @@ -75,6 +76,19 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): """ +@dataclass(frozen=True, slots=True) +class WebSocketSessionUsage: + model: str + custom_llm_provider: str + response_cost: float | None + + +class WebSocketSessionMeter(Protocol): + def measure(self, client_text: str) -> int: ... + + def price(self, total: int) -> WebSocketSessionUsage: ... + + class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase): """ Settings for a specific guardrail on a passthrough endpoint. diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index 0df0f0c80ef..76e806e3899 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -1,9 +1,19 @@ -from typing import Literal +from dataclasses import dataclass +from typing import Final, Literal from pydantic import BaseModel, Field from .base import GuardrailConfigModel +HEADROOM_REQUEST_CONTEXT_KEY: Final = "_headroom_request_context" + + +@dataclass(frozen=True, slots=True) +class HeadroomRequestContext: + issued_hashes: frozenset[str] + affinity: str | None + expires_at: float + class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): api_base: str | None = Field( diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 3f1a7980b27..37f59711050 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -68,29 +68,34 @@ class MetricWithMetadata(MetricBase): class GuardrailSpendMetrics(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - spend: float = 0.0 - untracked_units: int = 0 + spend: float = Field(default=0.0, description="Guardrail spend in USD, already included in key and model spend") + untracked_units: int = Field(default=0, description="Guardrail units with no known price, excluded from spend") class GuardrailKeyMetric(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - metrics: GuardrailSpendMetrics - metadata: KeyMetadata + metrics: GuardrailSpendMetrics = Field(description="Guardrail spend and unpriced units for this API key") + metadata: KeyMetadata = Field(description="API key alias and existence information") class GuardrailMetadata(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - guardrail_name: str + guardrail_name: str = Field(description="Guardrail display name, falling back to the guardrail ID") class GuardrailMetric(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - metrics: GuardrailSpendMetrics - metadata: GuardrailMetadata - api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field(default_factory=dict) + metrics: GuardrailSpendMetrics = Field( + description="Guardrail spend and unpriced units across all matching API keys" + ) + metadata: GuardrailMetadata = Field(description="Guardrail name information") + api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field( + default_factory=dict, + description="Per-key guardrail spend and unpriced units, limited to keys in the day's API key breakdown", + ) class BreakdownMetrics(LiteLLMBaseModel): @@ -103,7 +108,11 @@ class BreakdownMetrics(LiteLLMBaseModel): endpoints: dict[str, MetricWithMetadata] = Field(default_factory=dict) # endpoint -> {metrics, metadata} api_keys: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} entities: dict[str, MetricWithMetadata] = Field(default_factory=dict) # entity -> {metrics, metadata} - guardrails: Mapping[str, GuardrailMetric] = Field(default_factory=dict) # guardrail_id -> spend + guardrails: Mapping[str, GuardrailMetric] = Field( + default_factory=dict, + description="Spend by guardrail ID, already included in key and model spend. " + "Populated on user, team and organization aggregated activity only, without a model filter", + ) class DailySpendData(LiteLLMBaseModel): diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index b49cb63c8e0..f42a412cf2d 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -118,7 +118,7 @@ class UpdateKeyRequest(KeyRequestBase): project_id: str | None = Field( default=None, - description="Omit to retain the project, or send null to detach. Assigning a different project is not supported.", + description="Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is not supported.", ) @model_validator(mode="before") diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index ffdcf53fd7e..ca2767a7174 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,10 +4,10 @@ https://docs.cohere.com/reference/rerank """ -from collections.abc import Sequence -from typing import Literal, overload +from collections.abc import Mapping, Sequence +from typing import Final, Literal, overload -from pydantic import ConfigDict, PrivateAttr +from pydantic import ConfigDict, PrivateAttr, TypeAdapter, ValidationError from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMBaseModel @@ -57,6 +57,8 @@ class OptionalRerankParams(TypedDict, total=False): class RerankBilledUnits(TypedDict, total=False): search_units: int | None total_tokens: int | None + input_tokens: ReadOnly[int | None] + output_tokens: ReadOnly[int | None] class RerankTokens(TypedDict, total=False): @@ -65,11 +67,47 @@ class RerankTokens(TypedDict, total=False): class RerankResponseMeta(TypedDict, total=False): - api_version: dict | None + api_version: dict[str, object] | None billed_units: RerankBilledUnits | None tokens: RerankTokens | None +class RerankUsage(TypedDict): + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + + +_RERANK_META: Final = TypeAdapter(RerankResponseMeta) + + +def get_rerank_usage(response: Mapping[str, object]) -> RerankUsage | None: + if "results" not in response or response.get("meta") is None: + return None + try: + meta: Final = _RERANK_META.validate_python(response.get("meta")) + except ValidationError: + return None + tokens: Final = meta.get("tokens") or RerankTokens() + billed: Final = meta.get("billed_units") or RerankBilledUnits() + input_tokens: Final = next( + (count for count in (tokens.get("input_tokens"), billed.get("input_tokens")) if count is not None), None + ) + output_tokens: Final = next( + (count for count in (tokens.get("output_tokens"), billed.get("output_tokens")) if count is not None), None + ) + total_tokens: Final = billed.get("total_tokens") + if input_tokens is None and output_tokens is None and total_tokens is None: + return None + completion_tokens: Final = output_tokens or 0 + prompt_tokens: Final = input_tokens if input_tokens is not None else max((total_tokens or 0) - completion_tokens, 0) + return RerankUsage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens if total_tokens is not None else prompt_tokens + completion_tokens, + ) + + class RerankResponseDocument(TypedDict): text: str diff --git a/litellm/types/router.py b/litellm/types/router.py index 726a328e913..cab78cee832 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -37,7 +37,7 @@ from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router_weights import RouterWeights if TYPE_CHECKING: - from litellm.router import Router + from litellm.router_api import RoutingContext as RouterContext from .completion import CompletionRequest from .embedding import EmbeddingRequest @@ -1276,7 +1276,7 @@ class FallbackAccessCheck(Protocol): rejects, so a fallback can never reach a model the caller could not have requested directly. """ - async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: "RouterContext") -> bool: ... class FallbackBudgetCheck(Protocol): @@ -1289,7 +1289,7 @@ class FallbackBudgetCheck(Protocol): and skips targets it rejects, leaving the free attempt itself untouched. """ - async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: "RouterContext") -> bool: ... class AutoRouterCapabilityLimit(Protocol): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 780c9ef6824..2f6918d0744 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1,4 +1,5 @@ import json +import math import re import time from collections.abc import Collection, Mapping, Sequence @@ -74,6 +75,7 @@ from litellm.types.workload_identity import ( ) from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers +from ..litellm_core_utils.hidden_params import set_extra_attribute from . import litellm_params as _litellm_params from .agents import LiteLLMSendMessageResponse from .guardrails import GuardrailEventHooks @@ -104,6 +106,7 @@ from .llms.openai import ( OpenAIRealtimeStreamList, ResponsesAPIResponse, WebSearchOptions, + _JsonValue, ) from .rerank import RerankResponse as RerankResponse @@ -462,6 +465,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): annotation_cost_per_page: float | None # for OCR models annotation_cost_per_page_batches: ReadOnly[float | None] search_context_cost_per_query: SearchContextCostPerQuery | None # Cost for using web search tool + search_queries_cost_per_query: ReadOnly[float | SearchContextCostPerQuery | None] web_search_billing_unit: ( Literal["per_query", "per_prompt"] | None ) # "per_query" (Gemini 3.x) or "per_prompt" (Gemini 2.x) @@ -1971,6 +1975,19 @@ class ServerToolUse(LiteLLMBaseModel): return getattr(self, key) +_COST_BREAKDOWN_ADAPTER: Final = TypeAdapter(Mapping[object, object]) + + +def _is_non_finite_float(value: object) -> bool: + return isinstance(value, float) and not math.isfinite(value) + + +def _carries_non_finite_number(cost: object) -> bool: + if not isinstance(cost, Mapping): + return _is_non_finite_float(cost) + return any(_is_non_finite_float(member) for member in _COST_BREAKDOWN_ADAPTER.validate_python(cost).values()) + + class Usage(SafeAttributeModel, CompletionUsage): _cache_creation_input_tokens: int = PrivateAttr( 0 @@ -1978,6 +1995,7 @@ class Usage(SafeAttributeModel, CompletionUsage): _cache_read_input_tokens: int = PrivateAttr( 0 ) # hidden param for prompt caching. Might change, once openai introduces their equivalent. + _dropped_non_finite_cost: bool = PrivateAttr(False) server_tool_use: ServerToolUse | None = None cost: float | None = None @@ -2084,10 +2102,14 @@ class Usage(SafeAttributeModel, CompletionUsage): else: # maintain openai compatibility in usage object if possible del self.server_tool_use - if cost is not None: - self.cost = cost - else: + dropped_non_finite_cost: Final = cost is not None and _carries_non_finite_number(cost) + if cost is None: del self.cost + elif dropped_non_finite_cost: + verbose_logger.warning("Dropping a non-finite provider usage cost") + del self.cost + else: + self.cost = cost ## ANTHROPIC MAPPING ## if "cache_creation_input_tokens" in params and isinstance(params["cache_creation_input_tokens"], int): @@ -2102,6 +2124,11 @@ class Usage(SafeAttributeModel, CompletionUsage): for k, v in params.items(): setattr(self, k, v) + self._dropped_non_finite_cost = dropped_non_finite_cost + + @staticmethod + def dropped_non_finite_cost(usage: "Usage") -> bool: + return usage._dropped_non_finite_cost def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator @@ -2125,6 +2152,13 @@ class UnreportedUsage(Usage): class StreamingChoices(OpenAIObject): + if TYPE_CHECKING: + finish_reason: str | None + index: int + delta: Delta + enhancements: object + logprobs: ChoiceLogprobs | JsonValue + def __init__( self, finish_reason=None, @@ -2138,25 +2172,22 @@ class StreamingChoices(OpenAIObject): # https://github.com/BerriAI/litellm/issues/8455 params.pop("message", None) super().__init__(**params) - if finish_reason: - self.finish_reason = map_finish_reason(finish_reason) - else: - self.finish_reason = None - self.index = index + set_extra_attribute(self, "finish_reason", map_finish_reason(finish_reason) if finish_reason else None) + set_extra_attribute(self, "index", index) if delta is not None: if isinstance(delta, Delta): - self.delta = delta + set_extra_attribute(self, "delta", delta) elif isinstance(delta, dict): - self.delta = Delta(**delta) + set_extra_attribute(self, "delta", Delta(**delta)) else: - self.delta = Delta() + set_extra_attribute(self, "delta", Delta()) if enhancements is not None: - self.enhancements = enhancements + set_extra_attribute(self, "enhancements", enhancements) if isinstance(logprobs, (dict, BaseModel)): - self.logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs) + set_extra_attribute(self, "logprobs", CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs)) else: - self.logprobs = logprobs + set_extra_attribute(self, "logprobs", logprobs) def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator @@ -2210,7 +2241,10 @@ class ModelResponseBase(OpenAIObject): @property def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation - return self._hidden_params + private: Final = self.__pydantic_private__ + if private is None or "_hidden_params" not in private: + return self._hidden_params + return private["_hidden_params"] @hidden_params.setter def hidden_params(self, hidden_params: dict[str, object]) -> None: # mutable-ok: API requires mutation @@ -3291,6 +3325,11 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): classifier_cost: float classifier_failure_reason: ReadOnly[ClassifierFailureReason] classifier_error_type: ReadOnly[str] + classifier_input_budget_tokens: ReadOnly[int] + classifier_input_tokens: ReadOnly[int] + classifier_input_count_method: ReadOnly[str] + classifier_input_truncated: ReadOnly[bool] + classifier_input_budget_status: ReadOnly[str] classifier_probabilities: ReadOnly[Mapping[str, float]] classifier_confidence: ReadOnly[float] heuristic_v2_forecast: ReadOnly[StandardLoggingHeuristicV2Forecast] @@ -3349,6 +3388,11 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "classifier_cost", "classifier_failure_reason", "classifier_error_type", + "classifier_input_budget_tokens", + "classifier_input_tokens", + "classifier_input_count_method", + "classifier_input_truncated", + "classifier_input_budget_status", "classifier_probabilities", "classifier_confidence", "heuristic_v2_forecast", @@ -3428,6 +3472,7 @@ class StandardLoggingHiddenParams(TypedDict): batch_models: list[str] | None batch_successful_requests: ReadOnly[int | None] batch_failed_requests: ReadOnly[int | None] + upstream_provider: ReadOnly[NotRequired[str | None]] litellm_model_name: str | None # the model name sent to the provider by litellm usage_object: dict | None @@ -4125,6 +4170,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_video_per_second: float | None = None output_cost_per_audio_per_second: float | None = None search_context_cost_per_query: dict[str, Any] | None = None + search_queries_cost_per_query: float | SearchContextCostPerQuery | None = None google_maps_grounding_cost_per_query: float | None = None citation_cost_per_token: float | None = None cache_read_input_token_cost_above_272k_tokens: float | None = None @@ -4918,29 +4964,50 @@ def _batch_usage_from_chat_usage(usage: Usage) -> BatchUsage: ) +OPENAI_BATCH_STATUS_FOR_VLLM_ROUTER_STATUS: Final = MappingProxyType( + {"pending": "validating", "running": "in_progress"} +) + + +def openai_batch_status(status: _JsonValue) -> _JsonValue: + return OPENAI_BATCH_STATUS_FOR_VLLM_ROUTER_STATUS.get(status, status) if isinstance(status, str) else status + + class LiteLLMBatch(Batch): _hidden_params: dict = {} + object: Literal["batch"] = "batch" + + @field_validator("object", mode="before") + @classmethod + def _default_missing_object_type(cls, value: _JsonValue) -> _JsonValue: + return "batch" if value is None else value + + @field_validator("status", mode="before") + @classmethod + def _map_vllm_router_status(cls, value: _JsonValue) -> _JsonValue: + return openai_batch_status(value) + @property - def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation + def hidden_params(self) -> dict[str, _JsonValue]: # mutable-ok: API requires mutation return self._hidden_params @hidden_params.setter - def hidden_params(self, hidden_params: dict[str, object]) -> None: # mutable-ok: API requires mutation + def hidden_params(self, hidden_params: dict[str, _JsonValue]) -> None: # mutable-ok: API requires mutation self._hidden_params = hidden_params usage: Usage | None = None @field_validator("usage", mode="before") @classmethod - def _usage_from_batch_token_counts(cls, value: object) -> object: + def _usage_from_batch_token_counts(cls, value: _JsonValue) -> _JsonValue: counts: Final = _parse_batch_token_counts(value) return value if counts is None else _usage_with_batch_token_counts(Usage.model_validate(value), counts) @field_serializer("usage") def _serialize_usage_with_batch_fields( self, usage: Usage | None, info: FieldSerializationInfo - ) -> Mapping[str, object] | None: + ) -> Mapping[str, _JsonValue] | None: if not isinstance(usage, Usage): return usage chat_fields: Final = usage.model_dump( diff --git a/litellm/utils.py b/litellm/utils.py index 6f56cc65f45..8fe3fd84fbd 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -94,6 +94,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetModelCostMap from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases +from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens @@ -362,6 +363,10 @@ if TYPE_CHECKING: get_llm_provider, is_non_openai_azure_model, ) + from litellm.litellm_core_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy, + reset_retry_policy, + ) from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) @@ -454,10 +459,6 @@ if TYPE_CHECKING: ) from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - from litellm.router_utils.get_retry_from_policy import ( - get_num_retries_from_retry_policy, - reset_retry_policy, - ) from litellm.types.llms.anthropic import ( ANTHROPIC_API_ONLY_HEADERS, AnthropicThinkingParam, @@ -5257,7 +5258,7 @@ def pre_process_optional_params(passed_params: dict, non_default_params: dict, c def _declared_reasoning_efforts(model_info: Mapping[str, object] | None) -> tuple[str, ...]: - from litellm.router_utils.reasoning_effort_capability import deployment_reasoning_effort_levels + from litellm.litellm_core_utils.reasoning_effort_capability import deployment_reasoning_effort_levels return deployment_reasoning_effort_levels(model_info) or () @@ -6888,6 +6889,8 @@ def cached_get_model_info_helper( model: str, custom_llm_provider: str | None, api_base: str | None = None, + *, + preserve_unset_costs: bool = False, ) -> ModelInfoBase: """ _get_model_info_helper wrapped with lru_cache @@ -6898,6 +6901,7 @@ def cached_get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, + preserve_unset_costs=preserve_unset_costs, ) @@ -6961,6 +6965,7 @@ def get_model_info_helper( api_key: str | None = None, *, catalog_only: bool = False, + preserve_unset_costs: bool = False, ) -> ModelInfoBase: """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's @@ -7078,7 +7083,7 @@ def get_model_info_helper( verbose_logger.debug("model not found in litellm.model_cost, tried: %s", potential_model_names) raise ModelNotMappedError(_model_not_mapped_message(model, custom_llm_provider)) _input_cost_per_token: float | None = _model_info.get("input_cost_per_token") - if _input_cost_per_token is None: + if _input_cost_per_token is None and not preserve_unset_costs: # default value to 0, be noisy about this verbose_logger.debug( "model=%s, custom_llm_provider=%s has no input_cost_per_token in model_cost_map. Defaulting to 0.", @@ -7088,7 +7093,7 @@ def get_model_info_helper( _input_cost_per_token = 0 _output_cost_per_token: float | None = _model_info.get("output_cost_per_token") - if _output_cost_per_token is None: + if _output_cost_per_token is None and not preserve_unset_costs: # default value to 0, be noisy about this verbose_logger.debug( "model=%s, custom_llm_provider=%s has no output_cost_per_token in model_cost_map. Defaulting to 0.", @@ -7370,6 +7375,7 @@ def get_model_info_helper( supports_regex_lookaround=_model_info.get("supports_regex_lookaround", None), supports_computer_use=_model_info.get("supports_computer_use", None), search_context_cost_per_query=_model_info.get("search_context_cost_per_query", None), + search_queries_cost_per_query=_model_info.get("search_queries_cost_per_query"), web_search_billing_unit=_model_info.get("web_search_billing_unit", None), search_free_requests=_model_info.get("search_free_requests", None), google_maps_grounding_cost_per_query=_model_info.get("google_maps_grounding_cost_per_query", None), @@ -8723,11 +8729,13 @@ class AvailableModelsCache(InMemoryCache): cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) - result: Final = cast(list[str] | None, self.get_cache(cache_key)) + result: Final = self._cached(cache_key) - if result is not None: - return copy.deepcopy(result) - return result + if isinstance(result, Sequence): + return list(result) + if self._cached(_failed_listing_key(cache_key)) is not None: + return [] + return None def set_cached_model_info( self, @@ -8737,12 +8745,46 @@ class AvailableModelsCache(InMemoryCache): ): """Set cached model info""" cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) - self.set_cache(cache_key, copy.deepcopy(available_models)) + self._store(cache_key, copy.deepcopy(available_models)) + + def set_failed_model_listing( + self, custom_llm_provider: str, litellm_params: LiteLLM_Params | None, ttl_seconds: float + ) -> None: + cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) + self._store(_failed_listing_key(cache_key), True, ttl=ttl_seconds) + + def listing_lock(self, custom_llm_provider: str, litellm_params: LiteLLM_Params | None) -> threading.RLock: + return _provider_listing_lock(self._get_cache_key(custom_llm_provider, litellm_params)) + + def _cached(self, key: str) -> Sequence[str] | bool | None: + return cast(Sequence[str] | bool | None, self.get_cache(key)) + + def _store(self, key: str, value: object, ttl: float | None = None) -> None: + self.set_cache(key, value, ttl=ttl) + + +def _failed_listing_key(cache_key: str) -> str: + return f"{cache_key}:failed" + + +_PROVIDER_LISTING_LOCK_GUARD: Final = threading.Lock() + + +@lru_cache(maxsize=1024) +def _cached_provider_listing_lock(cache_key: str) -> threading.RLock: + return threading.RLock() + + +def _provider_listing_lock(cache_key: str) -> threading.RLock: + with _PROVIDER_LISTING_LOCK_GUARD: + return _cached_provider_listing_lock(cache_key) # Global cache instance _model_cache: Final = AvailableModelsCache() +_FAILED_PROVIDER_LISTING_TTL_SECONDS: Final = 30 + def _infer_valid_provider_from_env_vars( custom_llm_provider: str | None = None, @@ -8790,20 +8832,33 @@ def _get_valid_models_from_provider_api( custom_llm_provider: str, litellm_params: LiteLLM_Params | None = None, ) -> list[str]: - try: - cached_result: Final = _model_cache.get_cached_model_info(custom_llm_provider, litellm_params) + with _model_cache.listing_lock(custom_llm_provider, litellm_params): + try: + cached_result: Final = _model_cache.get_cached_model_info(custom_llm_provider, litellm_params) - if cached_result is not None: - return cached_result - models: Final = provider_config.discover_models( - litellm_params=_provider_listing_params(custom_llm_provider, litellm_params) - ) + if cached_result is not None: + return cached_result + models: Final = provider_config.discover_models( + litellm_params=_provider_listing_params(custom_llm_provider, litellm_params) + ) - _model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models) - return models - except Exception as e: - verbose_logger.warning("Error getting valid models: %s", e) - return [] + _model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models) + return models + except httpx.TimeoutException as e: + verbose_logger.warning("Error getting valid models: %s", e) + _model_cache.set_failed_model_listing( + custom_llm_provider, litellm_params, ttl_seconds=_FAILED_PROVIDER_LISTING_TTL_SECONDS + ) + return [] + except Exception as e: + verbose_logger.warning("Error getting valid models: %s", e) + return [] + + +def bundled_provider_models(provider: str) -> tuple[str, ...]: + bundled: Final[tuple[str, ...]] = tuple(litellm.models_by_provider.get(provider, ())) + cost_map: Final = _builtin_cost_map() + return tuple(model for model in bundled if not is_hidden_as_retired(cost_map.get(model, {}))) def get_valid_models( @@ -8871,8 +8926,7 @@ def get_valid_models( ) ) else: - models_for_provider = copy.deepcopy(litellm.models_by_provider.get(provider, [])) - valid_models.extend(models_for_provider) + valid_models.extend(bundled_provider_models(provider)) return valid_models except Exception as e: @@ -10464,6 +10518,10 @@ class ProviderConfigManager: from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo return AzureFoundryModelInfo(model=model) + elif LlmProviders.OPENROUTER == provider: + from litellm.llms.openrouter.common_utils import OpenRouterModelInfo + + return OpenRouterModelInfo() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0799a298942..cbe502e5cae 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -8715,7 +8715,7 @@ "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -8742,7 +8742,11 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, @@ -8763,7 +8767,7 @@ "search_context_size_low": 0.014, "search_context_size_medium": 0.014 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -8790,7 +8794,11 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 }, "azure/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, @@ -23070,6 +23078,24 @@ "output_cost_per_token_batches": 1e-06, "source": "https://developers.openai.com/api/docs/pricing" }, + "deepgram/aura": { + "input_cost_per_character": 1.5e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0150/1000 characters = $0.000015 per character", + "note": "Aura-1 text-to-speech voices (aura--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, + "deepgram/aura-2": { + "input_cost_per_character": 3e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.030/1000 characters = $0.00003 per character", + "note": "Aura-2 text-to-speech voices (aura-2--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, "deepgram/base": { "input_cost_per_second": 0.00020833, "litellm_provider": "deepgram", @@ -25276,6 +25302,38 @@ "cache_creation_input_token_cost": 1e-06, "prompt_cache_min_tokens": 2048 }, + "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_1hr": 2.2e-06, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "deprecation_date": "2026-10-15", + "litellm_provider": "bedrock_converse", + "supports_tool_search": true, + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_native_structured_output": true, + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06, + "token_count_model": { + "model": "anthropic.claude-haiku-4-5-20251001-v1:0", + "litellm_provider": "bedrock" + } + }, "eu.anthropic.claude-3-5-sonnet-20240620-v1:0": { "input_cost_per_token": 3e-06, "litellm_provider": "bedrock", @@ -29553,6 +29611,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30159,6 +30218,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_pdf_input": true, @@ -30202,6 +30262,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30266,6 +30327,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30326,6 +30388,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30393,6 +30456,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -30856,8 +30920,9 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/embeddings#model-versions", - "tpm": 10000000 + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "tpm": 10000000, + "input_cost_per_token_batches": 7.5e-08 }, "gemini/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, @@ -31209,6 +31274,9 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, "rpm": 100000, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ @@ -31389,6 +31457,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "rpm": 1000, @@ -31713,7 +31782,10 @@ "input_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, - "supports_url_context": true + "supports_url_context": true, + "cache_read_input_token_cost_above_200k_tokens_batches": 2.5e-07, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "deprecation_date": "2026-07-28", @@ -31871,7 +31943,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "supports_computer_use": true }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -31937,7 +32010,8 @@ "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, - "supports_audio_input": true + "supports_audio_input": true, + "supports_computer_use": true }, "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -32184,7 +32258,8 @@ }, "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014 + "google_maps_grounding_cost_per_query": 0.014, + "supports_computer_use": true }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -32280,7 +32355,10 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "output_cost_per_token_flex": 6e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06 }, "gemini/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -32343,7 +32421,10 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "output_cost_per_token_flex": 6e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06 }, "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -32384,6 +32465,7 @@ "text" ], "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -32458,6 +32540,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32522,6 +32605,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32582,6 +32666,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32649,6 +32734,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -47603,6 +47689,24 @@ "output_cost_per_token": 0.0, "source": "https://cohere.com/pricing" }, + "rerank-v5.0-fast": { + "input_cost_per_token": 2e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, + "rerank-v5.0-pro": { + "input_cost_per_token": 5e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -51283,6 +51387,11 @@ "max_input_tokens": 200000, "max_output_tokens": 64000, "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -51311,6 +51420,11 @@ "max_input_tokens": 200000, "max_output_tokens": 64000, "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -52057,6 +52171,11 @@ "max_input_tokens": 200000, "max_output_tokens": 64000, "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -52162,6 +52281,11 @@ "max_input_tokens": 200000, "max_output_tokens": 64000, "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -52692,6 +52816,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -59693,12 +59818,13 @@ "cache_read_input_token_cost": 1.375e-06, "output_cost_per_token": 8.25e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -72228,6 +72354,7 @@ }, "aiand/deepseek-ai/deepseek-v4-flash": { "cache_read_input_token_cost": 8e-08, + "deprecation_date": "2026-10-27", "input_cost_per_token": 1.5e-07, "litellm_provider": "aiand", "max_input_tokens": 1048576, @@ -72257,6 +72384,22 @@ "supports_vision": false }, "aiand/deepseek-ai/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2026-11-05", + "input_cost_per_token": 3e-07, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "aiand/deepseek-ai/deepseek-v4.1-flash-1017": { "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 3e-07, "litellm_provider": "aiand", @@ -72286,8 +72429,24 @@ "supports_tool_choice": true, "supports_vision": true }, + "aiand/meta/muse-spark-open": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "aiand", + "max_input_tokens": 1048576, + "max_output_tokens": 943718, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://api.aiand.com/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "aiand/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2026-10-27", "input_cost_per_token": 7.5e-07, "litellm_provider": "aiand", "max_input_tokens": 262144, @@ -72378,6 +72537,7 @@ }, "aiand/zai-org/glm-5.2": { "cache_read_input_token_cost": 3e-07, + "deprecation_date": "2026-10-27", "input_cost_per_token": 1e-06, "litellm_provider": "aiand", "max_input_tokens": 1048576, @@ -75343,6 +75503,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "rpm": 1000.0, @@ -76443,6 +76604,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76488,6 +76650,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76540,6 +76703,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76598,6 +76762,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -81100,5 +81265,61 @@ "output_cost_per_token": 0.0, "source": "https://inference.baseten.co/v1/models", "supports_vision": true + }, + "gpt-6-rosalind": { + "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_12hr": 3.75e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 7.5e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.75e-05, + "cache_read_input_token_cost": 1.5e-06, + "cache_read_input_token_cost_above_272k_tokens": 3e-06, + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_above_272k_tokens": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6-rosalind": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } diff --git a/scripts/benchmark_ocr_callbacks.py b/scripts/benchmark_ocr_callbacks.py index 85cd4ffae49..e2282356356 100644 --- a/scripts/benchmark_ocr_callbacks.py +++ b/scripts/benchmark_ocr_callbacks.py @@ -103,7 +103,7 @@ def percentile(values: Sequence[float], quantile: float) -> float: def verify(response: object, expected_transport: str) -> None: - from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict + from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict pages: Final = getattr(response, "pages", ()) if len(pages) != 1 or getattr(pages[0], "markdown", None) != EXPECTED_MARKDOWN: diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 597e0c6d4d8..4a1838a7b60 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -18,7 +18,7 @@ longer signal it. - **credential, model, mcp_server**: Write-only secret attributes so secrets never land in Terraform state (Terraform 1.11 or later). `litellm_credential` gains `credential_values_wo` (a JSON object, `jsonencode({...})`, since the SDK cannot mark a map write-only) with `credential_values_wo_version`, and `credential_values` becomes optional with exactly one of the two required. `litellm_model` gains `model_api_key_wo`, `aws_secret_access_key_wo` and `vertex_credentials_wo`, each with a `_wo_version` companion and each conflicting with its in-state twin. `litellm_mcp_server` gains `client_id`, `client_secret_wo` and `auth_value_wo` (with `_wo_version` companions), sent to the proxy as `credentials`, and `auth_type` now accepts every value the proxy does (`oauth2`, `api_key`, `bearer_token`, ...). Every `_wo` value is sent on each apply from the configuration and produces no diff, so bump its `_wo_version` to push a rotated secret. `litellm_vector_store.litellm_params` warns at plan time when it carries a secret-looking key such as `api_key` - **team**: Optional `team_member_model_max_budget` argument on `litellm_team`, a JSON string of per-model budgets applied to every team member and sent as `team_member_model_max_budget` on `/team/new` and `/team/update`; read back from `team_member_budget_table.model_max_budget` on `/team/info` so out-of-band changes show as drift, with `budget_limit`/`time_period` treated as equal to the `max_budget`/`budget_duration` spelling the proxy stores; removing it sends `{}`, which clears every per-model member budget; `tpm_limit`/`rpm_limit` are rejected at plan time because the proxy rejects them on team budgets. Needs a proxy that includes [#43887](https://github.com/BerriAI/litellm/pull/43887): v1.105.0 and v1.106.0-rc.1 accept the field but silently drop it, so every plan shows a diff -- **model**: Optional `display_name` argument on `litellm_model`, sent as `model_info.display_name` and returned as `display_name` by `/v1/models`, so client model pickers show a readable name; changes are persisted through `/model/{id}/update` since `/model/update` ignores `model_info`; also exported by the `litellm_model` and `litellm_models` data sources +- **model**: Optional `display_name` argument on `litellm_model`, sent as `model_info.display_name` and returned as `display_name` by `/v1/models`, so client model pickers show a readable name; changes are also persisted through `/model/{id}/update`, since proxies older than this provider drop the `model_info` sent to `/model/update`; also exported by the `litellm_model` and `litellm_models` data sources - **key**: Computed `server_metadata` attribute on `litellm_key` exposing every metadata entry the proxy stores, so metadata created outside Terraform is visible in state and drift on it shows on refresh, while `metadata` keeps tracking only the declared entries and updates keep preserving undeclared ones - **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them - **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement @@ -70,6 +70,7 @@ longer signal it. ### Changed +- **model** (breaking): on a proxy that includes [#46902](https://github.com/BerriAI/litellm/pull/46902), `litellm_model` updates now save the `model_info` the provider sends to `/model/update`, so `pricing_base_model`, `tier` and `mode` changes take effect. The first apply after upgrading the proxy overwrites values set outside Terraform with the configuration's, including `tier`, which defaults to `free`, and `base_model`, which falls back to the deployment name when `pricing_base_model` is unset. Run `terraform plan` before that apply and copy any value to keep into the configuration - **model** (breaking): a key named in `additional_litellm_params.additional_drop_params` is no longer deleted from the deployment's `litellm_params`. A configuration that also sets that key, such as `vertex_location`, stores it on the deployment at its next apply, and the proxy applies the drop list to requests instead. For `vertex_location` this changes routing. `/v1/messages` calls now go to the configured region, while chat completions drop a listed `vertex_location` and use the default region. To keep such a key off a deployment, remove its attribute and replace the model with `terraform apply -replace`, which gives it a new model id, since an update keeps a stored value the provider leaves out. `model`, `custom_llm_provider` and `merge_reasoning_content_in_choices` are always sent. A value that is not a JSON list of names, such as a bare `"store"` or an empty string, now fails the plan, and with it `terraform destroy`, instead of being stored as a string the proxy ignores. Rewrite it as `jsonencode([...])`, or remove the model without fixing it through `terraform destroy -refresh=false` - **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 3db6f68760d..1f5aa62560f 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -357,7 +357,7 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e return retryModelRead(d, m, 5) } -// /model/update only merges litellm_params, so model_info changes go through the PATCH endpoint. +// Older proxies drop the model_info sent to /model/update, so display_name also goes through the PATCH endpoint. func patchModelDisplayName(client *Client, modelID, displayName string) error { resp, err := MakeRequest(client, "PATCH", fmt.Sprintf(endpointModelPatch, url.PathEscape(modelID)), ModelInfoPatch{ ModelInfo: ModelInfoPatchFields{ID: modelID, DisplayName: displayName}, diff --git a/tests/_support/native_database.py b/tests/_support/native_database.py new file mode 100644 index 00000000000..29b1fff399f --- /dev/null +++ b/tests/_support/native_database.py @@ -0,0 +1,11 @@ +import asyncio + +from litellm.rust_bridge._native import NativeDatabase + + +async def _connect() -> NativeDatabase: + return await NativeDatabase.connect() + + +def native_database() -> NativeDatabase: + return asyncio.run(_connect()) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 7e05e98d917..69ad0c91ee9 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -198,11 +198,11 @@ def main(): router_file = [ "./litellm/router.py", "./litellm/litellm_core_utils/batch_utils.py", - "./litellm/router_utils/pattern_match_deployments.py", + "./litellm/litellm_core_utils/pattern_match_deployments.py", ] # router_file = [ # "../../litellm/router.py", - # "../../litellm/router_utils/pattern_match_deployments.py", + # "../../litellm/litellm_core_utils/pattern_match_deployments.py", # "../../litellm/litellm_core_utils/batch_utils.py", # ] ## LOCAL TESTING tests_dir = "./tests/" # Update this path if your tests directory is located elsewhere diff --git a/tests/code_coverage_tests/router_surface_allowlist.json b/tests/code_coverage_tests/router_surface_allowlist.json index 54758b7ec03..3a55e8fa41c 100644 --- a/tests/code_coverage_tests/router_surface_allowlist.json +++ b/tests/code_coverage_tests/router_surface_allowlist.json @@ -1,7 +1,6 @@ { "_doc": "Router members and litellm.router_utils modules used outside the router. Shrink-only: see tests/code_coverage_tests/router_surface_check.py.", "members_outside": [ - "_get_client", "a_add_message", "aadapter_completion", "aanthropic_messages", @@ -38,79 +37,14 @@ "avector_store_search", "completion", "embedding", - "get_routing_group", - "get_valid_args", "health", "observe", - "priority_header_update", "registry", "retire", "schedule_acompletion", "select", - "settings", - "team_public_model_names_of" + "settings" ], - "members_router_family": [ - "_count_pre_call_check_tokens", - "_pre_call_checks", - "_update_kwargs_before_fallbacks", - "acompletion", - "adecisions", - "aembedding", - "allowed_fails", - "allowed_fails_policy", - "aresponses", - "async_function_with_fallbacks", - "async_get_healthy_deployments", - "cache", - "cached_deployment_model_info", - "cooldown_cache", - "cooldown_time", - "default_max_parallel_requests", - "deployments_named", - "disable_cooldowns", - "embedding", - "enable_pre_call_checks", - "enable_tag_filtering", - "explicit_allowed_fails", - "fallback_access_check", - "fallback_budget_check", - "fallbacks", - "get_all_deployments", - "get_allowed_fails_from_policy", - "get_deployment", - "get_model_group", - "get_model_ids", - "get_model_info", - "get_model_list", - "get_router_model_info", - "log_retry", - "model_list", - "provider_default_deployment_ids", - "routing_group_has_alternatives", - "routing_strategy", - "search_tools", - "set_custom_routing_strategy", - "slack_alerting_logger", - "tag_filtering_match_any", - "tag_routing_prefix", - "team_model_has_alternatives" - ], - "router_utils_imports_outside": [ - "litellm.router_utils.access_windows", - "litellm.router_utils.add_retry_fallback_headers", - "litellm.router_utils.clientside_credential_handler", - "litellm.router_utils.common_utils", - "litellm.router_utils.cooldown_handlers", - "litellm.router_utils.deployment_pricing", - "litellm.router_utils.fallback_event_handlers", - "litellm.router_utils.get_retry_from_policy", - "litellm.router_utils.health_state_cache", - "litellm.router_utils.pattern_match_deployments", - "litellm.router_utils.ptu_shares", - "litellm.router_utils.reasoning_effort_capability", - "litellm.router_utils.router_callbacks.track_deployment_metrics", - "litellm.router_utils.routing_groups", - "litellm.router_utils.search_api_router" - ] + "members_router_family": [], + "router_utils_imports_outside": [] } diff --git a/tests/e2e/ui/fixtures/menuMappings.ts b/tests/e2e/ui/fixtures/menuMappings.ts index d6e7ea86982..4cf90c00cd0 100644 --- a/tests/e2e/ui/fixtures/menuMappings.ts +++ b/tests/e2e/ui/fixtures/menuMappings.ts @@ -32,7 +32,6 @@ export const menuLabelToPage: Record = { Budgets: Page.Budgets, "API Playground": Page.TransformRequest, "Tag Management": Page.TagManagement, - "Old Usage": Page.Usage, // Tools submenu items "MCP Servers": Page.McpServers, "Vector Stores": Page.VectorStores, diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index 5fa1f9d2444..6fa8d0c2587 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -118,12 +118,6 @@ export const MIGRATED_E2E_PAGES: Readonly> = { content: { role: "heading", name: "AI Hub" }, }, new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage" } }, - usage: { - segment: "old-usage", - linkName: "Old Usage", - group: "Experimental", - content: { role: "tab", name: "All Up" }, - }, agents: { segment: "agents", linkName: "Agents", group: "Agentic", content: { role: "heading", name: "Agents" } }, "router-settings": { segment: "router-settings", diff --git a/tests/e2e/ui/fixtures/pages.ts b/tests/e2e/ui/fixtures/pages.ts index 8210334c166..a850dcb4755 100644 --- a/tests/e2e/ui/fixtures/pages.ts +++ b/tests/e2e/ui/fixtures/pages.ts @@ -30,5 +30,4 @@ export enum Page { TagManagement = "tag-management", VectorStores = "vector-stores", NewUsage = "new_usage", - Usage = "usage", } diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts index 585eb31c37c..a611f760e8b 100644 --- a/tests/e2e/ui/tests/logs/logsPagination.spec.ts +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -84,7 +84,7 @@ test.describe("Logs pagination with Aria in the header", () => { await expect(aria).toBeFocused(); await visibleTestId(page, "pagination-first").click(); await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); - for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { + for (const destination of [Page.NewUsage, Page.LlmPlayground, Page.ApiKeys]) { await navigateToPage(page, destination); await aria.click(); await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); diff --git a/tests/e2e/ui/tests/usage/usagePage.spec.ts b/tests/e2e/ui/tests/usage/usagePage.spec.ts index d532042c7af..2a58a1ef59d 100644 --- a/tests/e2e/ui/tests/usage/usagePage.spec.ts +++ b/tests/e2e/ui/tests/usage/usagePage.spec.ts @@ -11,7 +11,7 @@ import { waitForSpendLog, } from "../../helpers/traffic"; -/** Covers /ui/usage. The legacy /ui/old-usage view is deprecated and deliberately not covered. */ +/** Covers /ui/usage. */ const topKeysCard = (page: PlaywrightPage): Locator => page.locator("section").filter({ has: page.getByRole("heading", { name: "Top Virtual Keys", exact: true }) }); diff --git a/tests/integration/_support/hook_count_probe.py b/tests/integration/_support/hook_count_probe.py index 0d4056e8fb9..3cd72f1a6ba 100644 --- a/tests/integration/_support/hook_count_probe.py +++ b/tests/integration/_support/hook_count_probe.py @@ -75,6 +75,24 @@ class BlockingPostCallGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail") +class UsageRecordingPostCallGuardrail(CustomGuardrail): + """Records the token usage of the response it is handed in the request's guardrail information, so the spend row + shows which response the guardrail saw.""" + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: object, response: object + ) -> None: + usage: Final = getattr(response, "usage", None) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ + "prompt_tokens": getattr(usage, "prompt_tokens", None), + "completion_tokens": getattr(usage, "completion_tokens", None), + }, + request_data=data, + guardrail_status="success", + ) + + probe: Final = HookCountProbe() sync_only_probe: Final = SyncOnlyHookCountProbe() diff --git a/tests/integration/authorization/test_guardrail_submissions.py b/tests/integration/authorization/test_guardrail_submissions.py index 16b9cbf31dc..fb095010eb2 100644 --- a/tests/integration/authorization/test_guardrail_submissions.py +++ b/tests/integration/authorization/test_guardrail_submissions.py @@ -432,7 +432,7 @@ def test_team_guardrail_submissions_require_admin_approval( f"Guardrail '{first_name}' not found. " "Please ensure the guardrail is configured in your LiteLLM proxy." ), - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": "404", } @@ -481,7 +481,7 @@ def test_team_guardrail_submissions_require_admin_approval( f"Guardrail '{second_name}' not found. " "Please ensure the guardrail is configured in your LiteLLM proxy." ), - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": "404", } diff --git a/tests/integration/authorization/test_user_self_update_guard.py b/tests/integration/authorization/test_user_self_update_guard.py index 2c495beb77e..4294e8a088d 100644 --- a/tests/integration/authorization/test_user_self_update_guard.py +++ b/tests/integration/authorization/test_user_self_update_guard.py @@ -279,7 +279,7 @@ def test_new_user_with_admin_role_requires_proxy_admin(gateway: Gateway, role: s assert refused.json() == { "error": { "message": f"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). Attempted to create user with role: {role}. Your role: internal_user", - "type": "internal_server_error", + "type": "auth_error", "param": None, "code": "403", } diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index 005e581c643..d47503e7f08 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -1199,7 +1199,7 @@ def test_anthropic_messages_explicit_null_reaches_upstream(gateway: Gateway) -> "anthropic_messages", deployment_params={"max_tokens": 32}, provider_model="openai/gpt-4o-mini", - response_route="acompletion", + response_route="aresponses", ) response: Final = _post( gateway, diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 2fb005dbeca..8590a19646b 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -3,7 +3,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" import json -from typing import Final, cast +from typing import Final, Literal, cast from pydantic import JsonValue @@ -41,6 +41,7 @@ from integration.cost_calculation.cases.anthropic import ( from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, + Deployment, ExactExpected, JsonResponse, RecountExpected, @@ -3387,6 +3388,232 @@ VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS: Final = bridge_pairs( ) +VERTEX_AI_CLAUDE_DEPLOYMENT_RATES: Final = Deployment( + input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01 +) + + +def _vertex_claude_web_search_message(web_search_requests: int) -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "scripted answer for cost fixtures"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "server_tool_use": {"web_search_requests": web_search_requests}, + }, + }, + ) + + +def _vertex_claude_web_search_stream(web_search_requests: int) -> SseResponse: + return SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1840}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer for cost fixtures"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 412, "server_tool_use": {"web_search_requests": web_search_requests}}, + }, + {"type": "message_stop"}, + ), + ) + + +def _vertex_claude_deployment_web_search_case( + name: str, + endpoint: Literal["/v1/chat/completions", "/v1/responses", "/v1/messages"], + request: dict[str, JsonValue], + response: JsonResponse | SseResponse, + tool_usage_cost: float, + deployment: Deployment = VERTEX_AI_CLAUDE_DEPLOYMENT_RATES, +) -> CostTrackingTestCase: + return CostTrackingTestCase( + name=name, + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="vertex_ai/claude-sonnet-5-5", + endpoint=endpoint, + deployment=deployment, + request={"model": "$MODEL", **request}, + response=response, + expected=ExactExpected( + spend=round(0.0039 + tool_usage_cost, 10), + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=tool_usage_cost, + ), + ) + + +VERTEX_AI_CLAUDE_WEB_SEARCH_TOOL: Final[dict[str, JsonValue]] = { + "type": "web_search_20250305", + "name": "web_search", + "max_uses": 3, +} + +VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_MESSAGES_WEB_SEARCH: Final = _vertex_claude_deployment_web_search_case( + "vertex_ai/claude-sonnet-5-5-deployment_rates_messages_web_search_two_requests", + "/v1/messages", + { + "max_tokens": 412, + "messages": [{"role": "user", "content": "vertex claude deployment rates messages web search"}], + "tools": [VERTEX_AI_CLAUDE_WEB_SEARCH_TOOL], + }, + _vertex_claude_web_search_message(2), + 0.02, +) +VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_STREAM: Final = _vertex_claude_deployment_web_search_case( + "vertex_ai/claude-sonnet-5-5-deployment_rates_messages_web_search_two_requests_stream", + "/v1/messages", + { + "max_tokens": 412, + "messages": [{"role": "user", "content": "vertex claude deployment rates messages web search stream"}], + "stream": True, + "tools": [VERTEX_AI_CLAUDE_WEB_SEARCH_TOOL], + }, + _vertex_claude_web_search_stream(2), + 0.02, +) +VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_ZERO_RATE: Final = ( + _vertex_claude_deployment_web_search_case( + "vertex_ai/claude-sonnet-5-5-deployment_zero_search_rate_messages_two_requests", + "/v1/messages", + { + "max_tokens": 412, + "messages": [{"role": "user", "content": "vertex claude deployment zero search rate messages"}], + "tools": [VERTEX_AI_CLAUDE_WEB_SEARCH_TOOL], + }, + _vertex_claude_web_search_message(2), + 0.0, + Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.0), + ) +) +VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_TOKEN_RATES_ONLY: Final = ( + _vertex_claude_deployment_web_search_case( + "vertex_ai/claude-sonnet-5-5-deployment_token_rates_only_messages_web_search_two_requests", + "/v1/messages", + { + "max_tokens": 412, + "messages": [{"role": "user", "content": "vertex claude deployment token rates only messages web search"}], + "tools": [VERTEX_AI_CLAUDE_WEB_SEARCH_TOOL], + }, + _vertex_claude_web_search_message(2), + 0.034, + Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06), + ) +) +VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_CHAT_WEB_SEARCH: Final = _vertex_claude_deployment_web_search_case( + "vertex_ai/claude-sonnet-5-5-deployment_rates_chat_web_search_three_requests", + "/v1/chat/completions", + { + "messages": [{"role": "user", "content": "vertex claude deployment rates chat web search"}], + "tools": [VERTEX_AI_CLAUDE_WEB_SEARCH_TOOL], + "allowed_openai_params": ["web_search_options"], + }, + _vertex_claude_web_search_message(3), + 0.03, +) +VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_CHAT_WEB_SEARCH_STREAM: Final = _vertex_claude_deployment_web_search_case( + "vertex_ai/claude-sonnet-5-5-deployment_rates_chat_web_search_two_requests_stream", + "/v1/chat/completions", + { + "messages": [{"role": "user", "content": "vertex claude deployment rates chat web search stream"}], + "stream": True, + "tools": [VERTEX_AI_CLAUDE_WEB_SEARCH_TOOL], + "allowed_openai_params": ["web_search_options"], + }, + _vertex_claude_web_search_stream(2), + 0.02, +) +VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_RESPONSES_WEB_SEARCH: Final = _vertex_claude_deployment_web_search_case( + "vertex_ai/claude-sonnet-5-5-deployment_rates_responses_web_search_two_requests", + "/v1/responses", + {"input": "vertex claude deployment rates responses web search", "tools": [{"type": "web_search"}]}, + _vertex_claude_web_search_message(2), + 0.02, +) + +GEMINI_3_1_PRO_PREVIEW_DEPLOYMENT_MESSAGES_WEB_SEARCH_STREAM: Final = CostTrackingTestCase( + name="vertex_ai/gemini-3.1-pro-preview-deployment_rates_messages_web_search_two_queries_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="vertex_ai/gemini-3.1-pro-preview", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.03), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "vertex gemini deployment rates messages web search stream"}], + "stream": True, + "tools": [VERTEX_AI_CLAUDE_WEB_SEARCH_TOOL], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "candidates": [ + { + "content": {"parts": [{"text": "scripted answer for cost fixtures"}], "role": "model"}, + "index": 0, + } + ], + "modelVersion": "gemini-3.1-pro-preview", + }, + { + "candidates": [ + { + "content": {"parts": [{"text": ""}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + "groundingMetadata": {"webSearchQueries": ["query 0", "query 1"]}, + } + ], + "usageMetadata": { + "promptTokenCount": 1840, + "candidatesTokenCount": 412, + "totalTokenCount": 2252, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + }, + "modelVersion": "gemini-3.1-pro-preview", + }, + ), + ), + expected=ExactExpected( + spend=0.0339, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.03, + ), +) + GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( GEMINI_3_8_FLASH_INPUT_TEXT, "vertex_ai/gemini-3.1-pro-preview-input_text", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412), @@ -3725,6 +3952,14 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM, *(pair.bridged for pair in VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_MESSAGES_WEB_SEARCH, + VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_STREAM, + VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_ZERO_RATE, + VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_TOKEN_RATES_ONLY, + VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_CHAT_WEB_SEARCH, + VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_CHAT_WEB_SEARCH_STREAM, + VERTEX_AI_CLAUDE_SONNET_5_5_DEPLOYMENT_RESPONSES_WEB_SEARCH, + GEMINI_3_1_PRO_PREVIEW_DEPLOYMENT_MESSAGES_WEB_SEARCH_STREAM, GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_3_1_PRO_PREVIEW_STREAM, GEMINI_3_1_PRO_PREVIEW_CACHE_READ, diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index d06f0b13cff..5741c85470b 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -1033,13 +1033,14 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( create_index: Final = sql.SQL('CREATE INDEX "entry_probe_idx" ON "entry_probe" ({})').format( sql.SQL(", ").join(sql.Identifier(column) for column in index_columns) ) + fitting_key: Final = secrets.token_hex(longest_fitting_key)[:longest_fitting_key] with psycopg.connect(scratch_database, autocommit=True) as conn: conn.execute('CREATE TABLE "entry_probe" ("s" TEXT, "k" TEXT, "t" TIMESTAMP(3))') conn.execute( 'INSERT INTO "entry_probe" ("s", "k", "t") VALUES (%s, %s, %s), (%s, %s, %s)', ( s_value, - secrets.token_hex(longest_fitting_key)[:longest_fitting_key], + fitting_key, t_value, s_value, secrets.token_hex(longest_fitting_key)[: longest_fitting_key + 1], @@ -1051,7 +1052,8 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( ) with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="exceeds btree version 4 maximum 2704"): conn.execute(create_index) - conn.execute('DELETE FROM "entry_probe" WHERE octet_length("k") > %s', (longest_fitting_key,)) + conn.execute('TRUNCATE "entry_probe"') + conn.execute('INSERT INTO "entry_probe" ("s", "k", "t") VALUES (%s, %s, %s)', (s_value, fitting_key, t_value)) conn.execute(create_index) (finding,) = findings @@ -1075,7 +1077,8 @@ def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random ) with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"): conn.execute(create_index) - conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,)) + conn.execute('TRUNCATE "compress_probe"') + conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s)', ("a" * 3000,)) conn.execute(create_index) stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone() diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py index 0c2ceb62eaf..86c2b325226 100644 --- a/tests/integration/database/test_accounting_ledger.py +++ b/tests/integration/database/test_accounting_ledger.py @@ -32,7 +32,7 @@ from litellm.accounting import ( PricingSnapshot, ) from litellm.proxy.db import accounting_ledger -from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger +from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger, UsageWrite from litellm.types.utils import CallTypes from tests.integration._support.database import scratch_database from tests.integration._support.prisma_clients import connected_proxy_prisma @@ -532,3 +532,74 @@ def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition) for name, condition in conditions ), conditions + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "conflict"]) +async def test_usage_batch_is_atomic_and_replay_preserves_receipts( + writer: _Writer, ledger_database: str, failure: str +) -> None: + client, store = writer + events: Final = (_usage("caller", inbound=True), _usage("attempt")) + writes: Final = tuple(UsageWrite(event, _envelope(event), _TARGETS) for event in events) + + @asynccontextmanager + async def interrupted() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + if failure == "before_commit": + raise RuntimeError("injected pre-commit failure") + if failure == "after_commit": + raise RuntimeError("injected lost acknowledgement") + + if failure == "conflict": + await store.insert_usage_batch((writes[1],)) + _complete_and_age(ledger_database) + rejected: Final = (writes[0], UsageWrite(events[1], _envelope(events[1], context={"changed": True}), _TARGETS)) + with pytest.raises(AccountingConflictError): + await store.insert_usage_batch(rejected) + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("attempt",)] + else: + with pytest.raises(RuntimeError, match="injected"): + await AccountingLedger(interrupted).insert_usage_batch(writes) + count: Final = 2 if failure == "after_commit" else 0 + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(count,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(count * len(_TARGETS),)] + await store.insert_usage_batch(writes) + recovered: Final = await store.load_usage_batch(tuple(event.event_id for event in events)) + assert len(recovered) == 1 + assert recovered[0].event == events[1] + assert json.loads(recovered[0].envelope) == json.loads(writes[1].envelope) + assert set(recovered[0].targets) == set(_TARGETS) + _complete_and_age(ledger_database) + before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') + await store.insert_usage_batch((*writes, *writes)) + assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') == before + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(2,)] + + +@pytest.mark.asyncio +async def test_invalid_jsonb_submission_cannot_reject_another_producers_commit( + writer: _Writer, ledger_database: str +) -> None: + from prisma.errors import DataError + + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + + _, store = writer + invalid: Final = _usage("invalid") + valid: Final = _usage("valid") + meter: Final = DurableUsageMeter(store.insert_usage_batch, UsageCaptureSettings(batch_size=2)) + rejected: Final = meter.submit( + (UsageWrite(invalid, _envelope(invalid, context={"parent": "bad\x00text"}), _TARGETS),) + ) + accepted: Final = meter.submit((UsageWrite(valid, _envelope(valid), _TARGETS),)) + try: + with pytest.raises(DataError): + await rejected + await accepted + finally: + await meter.close() + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("valid",)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] + assert (meter.pending_bytes, meter.pending_records) == (0, 0) diff --git a/tests/integration/database/test_proxy_database_lifespan.py b/tests/integration/database/test_proxy_database_lifespan.py new file mode 100644 index 00000000000..b80ff8acb30 --- /dev/null +++ b/tests/integration/database/test_proxy_database_lifespan.py @@ -0,0 +1,160 @@ +import asyncio +import os +import subprocess +import sys +import threading +import uuid +from pathlib import Path +from typing import Final, Literal, TypeAlias + +import psutil +import pytest + +from litellm.rust_bridge._native import NativeDatabase +from tests.integration._support.client import eventually +from tests.integration._support.database import read_rows, scratch_database + +pytestmark: Final = pytest.mark.timeout(180) +ROOT: Final = Path(__file__).resolve().parents[3] +PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234" +Mode: TypeAlias = Literal["late", "unsafe"] +CONFIG: Final = ( + "model_list: []\n" + "general_settings:\n" + " disable_prisma_schema_update: true\n" + " disable_spend_logs: true\n" + " disable_model_info_refresh: true\n" +) + + +def _assert_database_idle() -> None: + assert ( + eventually( + lambda: read_rows( + "SELECT pid FROM pg_stat_activity WHERE datname = current_database() AND pid <> pg_backend_pid()", + (), + ), + lambda rows: not rows, + seconds=5, + ) + == [] + ) + assert psutil.Process().children(recursive=True) == [] + assert tuple(thread.name for thread in threading.enumerate() if thread.name.startswith("prisma-engine-")) == () + + +async def _assert_lifespan_closed() -> None: + from litellm.proxy import proxy_server + + await asyncio.sleep(0) + assert proxy_server.prisma_client is None + assert ( + tuple( + task.get_coro().__qualname__ + for task in asyncio.all_tasks() + if task.get_coro().__qualname__.startswith( + ("PrismaClient.", "PrismaWrapper.", "_proxy_lifespan.._run_") + ) + ) + == () + ) + await asyncio.to_thread(_assert_database_idle) + + +async def _exercise_lifespans(mode: Mode) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError + + expected: Final = UnsafeMasterKeyError if mode == "unsafe" else ValueError + match: Final = "master key" if mode == "unsafe" else "use_redis_transaction_buffer" + with pytest.raises(expected, match=match): + async with proxy_server.proxy_startup_event(proxy_server.app): + pytest.fail("Invalid configuration started the proxy") + await _assert_lifespan_closed() + + Path(os.environ["CONFIG_FILE_PATH"]).write_text(CONFIG) + with pytest.MonkeyPatch.context() as environment: + environment.setenv("LITELLM_MASTER_KEY", f"sk-{uuid.uuid4().hex}") + for cancel in (False, True, False): + try: + async with proxy_server.proxy_startup_event(proxy_server.app) as state: + database: Final = state["database"] + assert database is not None + assert isinstance(database.native, NativeDatabase) + assert await database.prisma.writer_db.query_raw("SELECT 42 AS answer") == [{"answer": 42}] + assert await database.prisma.read_db.query_raw("SELECT 43 AS answer") == [{"answer": 43}] + assert psutil.Process().children(recursive=True) + if cancel: + raise asyncio.CancelledError + except asyncio.CancelledError: + if not cancel: + raise + await _assert_lifespan_closed() + + +@pytest.mark.parametrize("mode", ("late", "unsafe")) +@pytest.mark.parametrize("replica", (False, True), ids=("writer", "writer-and-reader")) +def test_failed_startup_then_retry_releases_database_resources(tmp_path: Path, mode: Mode, replica: bool) -> None: + config: Final = tmp_path / "config.yaml" + config.write_text(CONFIG + (" use_redis_transaction_buffer: true\n" if mode == "late" else "")) + inherited: Final = { + key: value + for key, value in os.environ.items() + if not key.startswith("REDIS_") + and key + not in ( + "WORKER_CONFIG", + "DATABASE_URL_READ_REPLICA", + "LITELLM_SALT_KEY", + "LITELLM_MIGRATE_FROM_MASTER_KEY", + "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", + "INTEGRATION_UPSTREAM_URL", + ) + } + with scratch_database("lifespan") as database_url: + environment: Final = { + **inherited, + "DATABASE_URL": database_url, + "DIRECT_URL": database_url, + "CONFIG_FILE_PATH": str(config), + "LITELLM_MASTER_KEY": PUBLICLY_KNOWN_KEY if mode == "unsafe" else f"sk-{uuid.uuid4().hex}", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_TELEMETRY_ENABLED": "false", + "PRISMA_HEALTH_WATCHDOG_ENABLED": "true", + "PYTHONPATH": os.pathsep.join((str(ROOT), str(ROOT / "enterprise"))), + **({"DATABASE_URL_READ_REPLICA": database_url} if replica else {}), + } + schema: Final = subprocess.run( + ( + sys.executable, + "-P", + "-m", + "prisma", + "db", + "push", + "--skip-generate", + "--schema", + str(ROOT / "litellm/proxy/schema.prisma"), + ), + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=60, + ) + assert schema.returncode == 0, schema.stderr + probe: Final = subprocess.run( + ( + sys.executable, + "-P", + "-c", + "import asyncio; from tests.integration.database.test_proxy_database_lifespan import _exercise_lifespans; " + f"asyncio.run(_exercise_lifespans({mode!r}))", + ), + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=90, + ) + assert probe.returncode == 0, probe.stdout + probe.stderr diff --git a/tests/integration/management/test_auto_router_databricks_classifier_writes.py b/tests/integration/management/test_auto_router_databricks_classifier_writes.py index d17ce947671..ecb37261790 100644 --- a/tests/integration/management/test_auto_router_databricks_classifier_writes.py +++ b/tests/integration/management/test_auto_router_databricks_classifier_writes.py @@ -278,6 +278,8 @@ def test_test_routing_classifies_through_databricks_without_calling_the_tier_it_ ], "classifier_model": f"databricks/{model}", **cost, + "classifier_input_truncated": False, + "classifier_input_budget_status": "unknown_limits", "classifier_probabilities": {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91}, "classifier_confidence": 0.91, "conversation_continuing": False, diff --git a/tests/integration/management/test_key_delete_revocation.py b/tests/integration/management/test_key_delete_revocation.py index 5b3527a52ae..bdc91dfe5f6 100644 --- a/tests/integration/management/test_key_delete_revocation.py +++ b/tests/integration/management/test_key_delete_revocation.py @@ -171,11 +171,8 @@ def test_deleting_a_key_revokes_it(gateway: Gateway, spelling: str) -> None: assert deleted_rows == (_DeletedKeyRow(token=key_hash),) -def test_deleting_a_key_revokes_it_on_a_warmed_peer(gateway: Gateway, peer: Gateway) -> None: - pytest.skip( - "BUG: POST /key/delete does not broadcast the key eviction, so a warmed peer keeps answering 200 for about " - "60 seconds until its in-memory cache TTL expires" - ) +@pytest.mark.parametrize("spelling", ("keys", "hashed_keys", "key_aliases")) +def test_deleting_a_key_revokes_it_on_a_warmed_peer(gateway: Gateway, peer: Gateway, spelling: str) -> None: backend: Final = "integration-keys-auth-delete" prompts: Final = ( "delete warm primary " + uuid.uuid4().hex, @@ -208,9 +205,13 @@ def test_deleting_a_key_revokes_it_on_a_warmed_peer(gateway: Gateway, peer: Gate ) assert len(warm_responses) == 2 assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,) * 2 - deleted: Final = gateway.request("POST", "/key/delete", {"keys": [key]}) + identifier: Final = {"keys": key, "hashed_keys": sha256(key.encode()).hexdigest(), "key_aliases": alias}[ + spelling + ] + delete_body: Final = {"key_aliases" if spelling == "key_aliases" else "keys": [identifier]} + deleted: Final = gateway.request("POST", "/key/delete", delete_body) assert deleted.status_code == 200, deleted.text - assert _DeleteResponse.model_validate_json(deleted.content).deleted_keys == (key,), deleted.text + assert _DeleteResponse.model_validate_json(deleted.content).deleted_keys == (identifier,), deleted.text peer_refusal: Final = eventually( lambda: _chat(peer, model, key, prompts[2]), lambda response: response.status_code == 401, diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index d0b6e2579c8..93cfa00ea13 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/tests/integration/management/test_sso_sign_in.py @@ -16,7 +16,7 @@ import yaml from pydantic import JsonValue from tests.integration._support.client import Gateway, gateway_from_environment -from tests.integration._support.database import read_rows +from tests.integration._support.database import read_rows, write_rows from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState, owned_proxy from tests.integration._support.wire import Reply, Request, wire_server @@ -52,14 +52,29 @@ def _idp_reply(request: Request) -> Reply: "preferred_username": signed_in, "email": f"{signed_in}@example.com", "provider": IDP_PROVIDER, + "dept": _dept(signed_in), + "groups": [_group(signed_in)], } ).encode() ) return Reply(status=404, body=b'{"error": "not_found"}') +def _dept(subject: str) -> str: + return f"dept-{subject}" + + +def _group(subject: str) -> str: + return f"group-{subject}" + + @contextmanager -def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> Iterator[Gateway]: +def _sso_gateway( + directory: Path, + litellm_settings: Mapping[str, JsonValue], + general_settings: Mapping[str, JsonValue] | None = None, + db_sso_settings: Mapping[str, JsonValue] | None = None, +) -> Iterator[Gateway]: config: Final = directory / "sso_proxy_config.yaml" config.write_text( yaml.safe_dump( @@ -68,6 +83,7 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> "general_settings": { "master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL", + **(general_settings or {}), }, "litellm_settings": dict(litellm_settings), } @@ -83,7 +99,33 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> "OAUTHLIB_INSECURE_TRANSPORT": "1", } with owned_proxy(rig, directory, environment, config=config, remove_environment=("PROXY_BASE_URL",)) as gw: - yield gw + if db_sso_settings is None: + yield gw + return + with contextlib.ExitStack() as cleanups: + cleanups.callback( + write_rows, + 'DELETE FROM "LiteLLM_SSOConfig" WHERE id = %s', + ("sso_config",), + database_url=gw.database_url, + ) + stored: Final = _store_sso_settings(gw, {**_generic_sso_settings(environment), **db_sso_settings}) + assert stored.status_code == 200, stored.text + yield gw + + +def _generic_sso_settings(environment: Mapping[str, str]) -> Mapping[str, JsonValue]: + return { + "generic_client_id": environment["GENERIC_CLIENT_ID"], + "generic_client_secret": environment["GENERIC_CLIENT_SECRET"], + "generic_authorization_endpoint": environment["GENERIC_AUTHORIZATION_ENDPOINT"], + "generic_token_endpoint": environment["GENERIC_TOKEN_ENDPOINT"], + "generic_userinfo_endpoint": environment["GENERIC_USERINFO_ENDPOINT"], + } + + +def _store_sso_settings(gateway: Gateway, settings: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("PATCH", "/update/sso_settings", settings) def _sign_in(gateway: Gateway, subject: str) -> httpx.Response: @@ -165,3 +207,65 @@ def test_free_sso_admission_preserves_existing_users_and_serializes_the_last_sea refused: Final = _sign_in(gateway, "free-sso-extra") assert (above_capacity.status_code, refused.status_code) == (303, 403) assert gateway.rows('SELECT COUNT(*) AS count FROM "LiteLLM_UserTable"', ()) == [{"count": 6}] + + +@pytest.mark.parametrize( + ("enable_jwt_auth", "team_id_upsert", "team_exists_before", "expect_member"), + [(True, True, False, True), (False, True, False, False), (True, False, False, False), (True, False, True, True)], + ids=[ + "upsert-creates-missing-team", + "upsert-without-jwt-auth-skips-missing-team", + "no-upsert-skips-missing-team", + "no-upsert-joins-existing-team", + ], +) +def test_sso_sign_in_upserts_the_team_id_jwt_field_team_only_when_team_id_upsert_is_on( + tmp_path: Path, enable_jwt_auth: bool, team_id_upsert: bool, team_exists_before: bool, expect_member: bool +) -> None: + subject: Final = f"sso-dept-user-{uuid.uuid4().hex[:12]}" + team_id: Final = _dept(subject) + jwtauth: Final = { + "enable_jwt_auth": enable_jwt_auth, + "litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": team_id_upsert}, + } + with _sso_gateway(tmp_path, {}, jwtauth) as gateway, contextlib.ExitStack() as cleanups: + cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id]}) + cleanups.callback(_delete_user, gateway, subject) + if team_exists_before: + created: Final = gateway.request("POST", "/team/new", {"team_id": team_id}) + assert created.status_code == 200, created.text + + callback: Final = _sign_in(gateway, subject) + + assert callback.status_code == 303, f"{callback.status_code} {callback.text}" + teams: Final = gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)) + user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) + assert (len(teams), user) == (int(expect_member), [{"teams": [team_id] if expect_member else []}]) + + +def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_and_only_joins_groups( + tmp_path: Path, +) -> None: + subject: Final = f"sso-mapped-user-{uuid.uuid4().hex[:12]}" + team_id: Final = _dept(subject) + group_team_id: Final = _group(subject) + jwtauth: Final = { + "enable_jwt_auth": True, + "litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": True}, + } + team_mappings: Final = {"team_mappings": {"team_ids_jwt_field": "groups"}} + with ( + _sso_gateway(tmp_path, {}, jwtauth, db_sso_settings=team_mappings) as gateway, + contextlib.ExitStack() as cleanups, + ): + cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id, group_team_id]}) + cleanups.callback(_delete_user, gateway, subject) + + callback: Final = _sign_in(gateway, subject) + + assert callback.status_code == 303, f"{callback.status_code} {callback.text}" + teams: Final = gateway.rows( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id IN (%s, %s)', (team_id, group_team_id) + ) + user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) + assert (teams, user) == ([{"team_id": team_id}], [{"teams": [team_id]}]) diff --git a/tests/integration/management/test_user_delete_revokes_keys.py b/tests/integration/management/test_user_delete_revokes_keys.py index 513a80b7319..2a1bd210b49 100644 --- a/tests/integration/management/test_user_delete_revokes_keys.py +++ b/tests/integration/management/test_user_delete_revokes_keys.py @@ -170,7 +170,7 @@ def test_user_delete_revokes_keys_and_memberships(gateway: Gateway, peer: Gatewa assert info.json() == { "error": { "message": f"User {user_id} not found", - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": "404", } @@ -337,7 +337,7 @@ def test_scim_user_delete_revokes_keys(gateway: Gateway, peer: Gateway, monkeypa expected_info: Final = { "error": { "message": f"User {user_id} not found", - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": "404", } diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 4b967802893..f6743a8751a 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -46,9 +46,10 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state] -@pytest.fixture(scope="session") -def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: - directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") +@contextmanager +def _private_gateway( + directory: Path, guardrails: Sequence[Mapping[str, JsonValue]] = () +) -> Iterator[tuple[Gateway, Mapping[str, str]]]: stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = directory / "stored-responses.yaml" (directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE) @@ -68,7 +69,8 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat { "guardrail_name": STREAM_BLOCK, "litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"}, - } + }, + *guardrails, ], } ) @@ -81,15 +83,61 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat yield private, services -@pytest.fixture -def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: - private, services = _llm_gateway +@contextmanager +def _serving(served: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + private, services = served with pytest.MonkeyPatch.context() as environment: for name, value in services.items(): environment.setenv(name, value) yield private +@pytest.fixture(scope="session") +def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-gateway")) as served: + yield served + + +@pytest.fixture +def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_llm_gateway) as private: + yield private + + +ROUND_USAGE_GUARDRAIL: Final = "mcp-llm-round-usage" +CONTENT_FILTER: Final = f"mcp-llm-content-filter-{uuid.uuid4().hex[:8]}" + + +@pytest.fixture(scope="module") +def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + round_usage: Final[Mapping[str, JsonValue]] = { + "guardrail_name": ROUND_USAGE_GUARDRAIL, + "litellm_params": { + "guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail", + "mode": "post_call", + "default_on": True, + }, + } + content_filter: Final[Mapping[str, JsonValue]] = { + "guardrail_name": CONTENT_FILTER, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "blocked_words": [{"keyword": "forbidden-keeper", "action": "BLOCK"}], + }, + } + guardrails: Final = (round_usage, content_filter) + with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), guardrails) as served: + yield served + + +@pytest.fixture +def guarded_gateway(_guarded_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + with _serving(_guarded_llm_gateway) as private: + yield private + + Surface = Literal["chat", "responses", "messages", "messages_bridge", "responses_bridge"] SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge") ADD: Final = {"a": 2, "b": 3} @@ -1032,6 +1080,184 @@ def test_nonstream_failed_follow_up_bills_round_one_on_the_failure_row( ) +def _streamed_chunks(response: httpx.Response) -> tuple[dict[str, object], ...]: + return tuple( + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def test_streaming_chat_auto_execute_logs_the_answer_and_bills_both_rounds(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert response.status_code == 200, response.text + chunks: Final = _streamed_chunks(response) + text: Final = "".join( + str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices") + ) + assert text == ANSWER, response.text + usages: Final = tuple(chunk["usage"] for chunk in chunks if chunk.get("usage")) + assert usages, response.text + rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") + assert len(rows) == 1, rows + row: Final = rows[0] + assert row["request_id"] == chunks[0]["id"], (row, chunks[0]) + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + FIRST_ROUND[0] + FOLLOW_UP_ROUND[0], + FIRST_ROUND[1] + FOLLOW_UP_ROUND[1], + ), row + assert float(str(row["spend"])) == pytest.approx(_price(FIRST_ROUND, FOLLOW_UP_ROUND)), row + final_usage: Final = usages[-1] + assert (final_usage["prompt_tokens"], final_usage["completion_tokens"], final_usage["total_tokens"]) == ( + row["prompt_tokens"], + row["completion_tokens"], + sum((*FIRST_ROUND, *FOLLOW_UP_ROUND)), + ), usages + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_bills_round_one_once(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert float(str(succeeded[0]["spend"])) == pytest.approx(_price(FIRST_ROUND)), succeeded + failed: Final = _spend_rows(call_id, "acompletion", "failure") + assert [float(str(row["spend"])) for row in failed] == [0.0], failed + + +def _guardrail_information(metadata: JsonValue) -> tuple[JsonValue, ...]: + entries: Final = object_value(metadata).get("guardrail_information") + return tuple(entries) if isinstance(entries, list) else () + + +def _round_usage_seen_by_the_guardrail(call_id: str) -> list[JsonValue]: + rows: Final = read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "acompletion", "success"), + ) + entries: Final = tuple(chain.from_iterable(_guardrail_information(row["metadata"]) for row in rows)) + return [ + entry.get("guardrail_response") + for entry in entries + if isinstance(entry, dict) and entry.get("guardrail_name") == ROUND_USAGE_GUARDRAIL + ] + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_guardrail_on_round_one( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + succeeded: Final = _spend_rows(call_id, "acompletion", "success") + assert len(succeeded) == 1, succeeded + assert (succeeded[0]["prompt_tokens"], succeeded[0]["completion_tokens"]) == FIRST_ROUND, succeeded + assert _round_usage_seen_by_the_guardrail(call_id) == [ + {"prompt_tokens": FIRST_ROUND[0], "completion_tokens": FIRST_ROUND[1]} + ] + + +def _chat_rows(call_id: str, statuses: Sequence[str]) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s', + (call_id, "acompletion"), + ), + lambda found: _statuses(found) == sorted(statuses), + seconds=70, + ) + + +def _entries_of(guardrail_name: str, row: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + entry + for entry in _guardrail_information(row["metadata"]) + if isinstance(entry, dict) and entry.get("guardrail_name") == guardrail_name + ) + + +def test_streaming_chat_auto_execute_whose_follow_up_fails_counts_its_post_call_content_filter_once( + guarded_gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + stream=True, + stream_options={"include_usage": True}, + guardrails=[CONTENT_FILTER], + ) + assert FOLLOW_UP_OUTAGE in response.text, response.text + rows: Final = _chat_rows(response.headers["x-litellm-call-id"], ("success", "failure")) + filtered: Final = tuple( + (row["status"], entry.get("guardrail_status")) for row in rows for entry in _entries_of(CONTENT_FILTER, row) + ) + assert filtered == (("failure", "guardrail_failed_to_respond"),), rows + (guardrail_id,) = frozenset( + str(entry.get("guardrail_id") or CONTENT_FILTER) + for row in rows + for entry in _entries_of(CONTENT_FILTER, row) + ) + counted: Final = eventually( + lambda: read_rows( + 'SELECT requests_evaluated, flagged_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s', + (guardrail_id,), + ), + lambda found: bool(found), + seconds=70, + ) + assert counted == [{"requests_evaluated": 1, "flagged_count": 1}], counted + + +@pytest.mark.parametrize("include_usage", (True, False), ids=("include-usage", "no-usage")) +def test_streaming_chat_auto_execute_answered_in_one_round_ends_with_the_usage_it_billed( + gateway: Gateway, include_usage: bool +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "chat", billing=billing, prices=PRICES) as rig: + response: Final = rig.send( + _granted_key(rig), + [AUTO], + prompt=f"{DIRECT} {rig.alias}", + stream=True, + stream_options={"include_usage": include_usage}, + ) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + chunks: Final = _streamed_chunks(response) + text: Final = "".join( + str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices") + ) + assert text == ANSWER, response.text + rows: Final = _spend_rows(response.headers["x-litellm-call-id"], "acompletion", "success") + assert len(rows) == 1 and (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == FIRST_ROUND, rows + usage_at: Final = tuple(index for index, chunk in enumerate(chunks) if chunk.get("usage")) + if not include_usage: + assert usage_at == (), response.text + return + assert usage_at == (len(chunks) - 1,), response.text + finished_at: Final = tuple( + index + for index, chunk in enumerate(chunks) + if chunk.get("choices") and chunk["choices"][0].get("finish_reason") is not None + ) + assert finished_at and finished_at[-1] < usage_at[0], response.text + usage: Final = chunks[-1]["usage"] + assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == ( + *FIRST_ROUND, + sum(FIRST_ROUND), + ), usage + + NO_CACHE: Final = {"no-cache": True} @@ -1269,6 +1495,27 @@ def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_an assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows +def test_streaming_chat_follow_up_broken_mid_answer_bills_round_one_and_its_partial_answer_with_one_fixed_margin( + gateway: Gateway, +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True) + with _rig(gateway, "chat", upstream="messages_bridge", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True}) + assert "the sum" in response.text, response.text + assert ANSWER not in response.text, response.text + assert len(rig.upstream_tools()) == 2 + round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN + rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one)) + assert _statuses(rows) == ["failure", "success"], rows + (success,) = (row for row in rows if row["status"] == "success") + (failure,) = (row for row in rows if row["status"] == "failure") + assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows + assert float(str(success["spend"])) == pytest.approx(round_one), rows + partial_answer: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"]))) + assert min(partial_answer) > 0, rows + assert float(str(failure["spend"])) == pytest.approx(_price(partial_answer)), rows + + @pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS) def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once( gateway: Gateway, upstream: StreamedUpstream diff --git a/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py b/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py new file mode 100644 index 00000000000..f2622da1bf6 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py @@ -0,0 +1,191 @@ +import base64 +import json +import uuid +from typing import Final + +import anthropic +import httpx +import pytest +from integration._support import claude_code as cc +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL_ID: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" +_TOKEN: Final = "synthetic-bedrock-bearer" +_PROMPT: Final = "Reply with exactly the word PONG" +_INPUT_TOKENS: Final = 10 +_OUTPUT_TOKENS: Final = 4 +_CACHE_READ_TOKENS: Final = 900 +_CACHE_WRITE_TOKENS: Final = 300 +_INPUT_RATE: Final = 0.001 +_OUTPUT_RATE: Final = 0.002 +_CACHE_READ_RATE: Final = 0.0001 +_CACHE_WRITE_RATE: Final = 0.0015 +_EXPECTED_SPEND: Final = ( + _INPUT_TOKENS * _INPUT_RATE + + _CACHE_READ_TOKENS * _CACHE_READ_RATE + + _CACHE_WRITE_TOKENS * _CACHE_WRITE_RATE + + _OUTPUT_TOKENS * _OUTPUT_RATE +) +_PROMPT_TOKENS: Final = _INPUT_TOKENS + _CACHE_READ_TOKENS + _CACHE_WRITE_TOKENS +_USAGE: Final = { + "input_tokens": _INPUT_TOKENS, + "output_tokens": _OUTPUT_TOKENS, + "cache_read_input_tokens": _CACHE_READ_TOKENS, + "cache_creation_input_tokens": _CACHE_WRITE_TOKENS, +} +_USER_CONTENT: Final = [{"type": "text", "text": _PROMPT, "cache_control": cc.CACHE}] +_CACHED_SYSTEM_BLOCKS: Final = [block for block in cc.system_blocks() if "cache_control" in block] + + +def _assert_claude_code_breakpoints_reach_bedrock(request: Request, target: str) -> dict[str, JsonValue]: + assert request.method == "POST" and request.target == target, request.target + assert request.headers["authorization"] == f"Bearer {_TOKEN}" + body: Final = cc.JSON_OBJECT.validate_json(request.body) + assert body["anthropic_version"] == "bedrock-2023-05-31", body + assert body["system"] == _CACHED_SYSTEM_BLOCKS, f"billing-header block dropped, breakpoints kept: {body}" + assert body["messages"] == [{"role": "user", "content": _USER_CONTENT}], body + return body + + +def _invoke_chunk(event: dict[str, JsonValue]) -> bytes: + return _aws_event_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}, "sc", "u") + + +def _stream_frames(message_id: str) -> tuple[bytes, ...]: + return ( + _invoke_chunk( + { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": "claude-sonnet-4-5-20250929", + "stop_reason": None, + "stop_sequence": None, + "usage": {**_USAGE, "output_tokens": 1}, + }, + } + ), + _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "PONG"}}), + _invoke_chunk({"type": "content_block_stop", "index": 0}), + _invoke_chunk( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": _OUTPUT_TOKENS}, + } + ), + _invoke_chunk({"type": "message_stop"}), + ) + + +def _spend_row(message_id: str) -> dict[str, object]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (message_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0] + + +def _client(gateway: Gateway, transport: httpx.Client) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=transport + ) + + +def test_bedrock_claude_cache_control_reaches_invoke_and_cache_tokens_are_reported_and_charged( + gateway: Gateway, +) -> None: + message_id: Final = f"msg_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + _assert_claude_code_breakpoints_reach_bedrock(request, f"/model/{_MODEL_ID}/invoke") + return Reply( + body=cc.message_reply( + message_id, "claude-sonnet-4-5-20250929", ({"type": "text", "text": "PONG"},), dict(_USAGE) + ) + ) + + with ( + wire_server(respond) as wire, + gateway.scenario() as scenario, + httpx.Client(timeout=15, trust_env=False) as transport, + ): + model: Final = scenario.model( + model=f"bedrock/invoke/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + api_base=wire.url, + input_cost_per_token=_INPUT_RATE, + output_cost_per_token=_OUTPUT_RATE, + cache_read_input_token_cost=_CACHE_READ_RATE, + cache_creation_input_token_cost=_CACHE_WRITE_RATE, + ) + final: Final = _client(gateway, transport).messages.create( + model=model, + max_tokens=32, + system=list(cc.system_blocks()), + messages=[{"role": "user", "content": _USER_CONTENT}], + ) + assert len(wire.drain()) == 1 + assert final.id == message_id, final + assert [(block.type, block.text) for block in final.content] == [("text", "PONG")], final + assert final.usage.model_dump(exclude_none=True, include=set(_USAGE)) == _USAGE, final.usage + row: Final = _spend_row(message_id) + + assert float(str(row["spend"])) == pytest.approx(_EXPECTED_SPEND), row + assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _OUTPUT_TOKENS), row + + +def test_bedrock_claude_stream_cache_tokens_from_message_start_are_reported_and_charged(gateway: Gateway) -> None: + message_id: Final = f"msg_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _assert_claude_code_breakpoints_reach_bedrock( + request, f"/model/{_MODEL_ID}/invoke-with-response-stream" + ) + assert "stream" not in body, body + return Reply(chunks=_stream_frames(message_id), content_type="application/vnd.amazon.eventstream") + + with ( + wire_server(respond) as wire, + gateway.scenario() as scenario, + httpx.Client(timeout=15, trust_env=False) as transport, + ): + model: Final = scenario.model( + model=f"bedrock/invoke/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + api_base=wire.url, + input_cost_per_token=_INPUT_RATE, + output_cost_per_token=_OUTPUT_RATE, + cache_read_input_token_cost=_CACHE_READ_RATE, + cache_creation_input_token_cost=_CACHE_WRITE_RATE, + ) + with _client(gateway, transport).messages.stream( + model=model, + max_tokens=32, + system=list(cc.system_blocks()), + messages=[{"role": "user", "content": _USER_CONTENT}], + ) as stream: + deltas: Final = tuple(stream.text_stream) + final: Final = stream.get_final_message() + assert len(wire.drain()) == 1 + assert deltas == ("PONG",) + assert final.id == message_id, final + assert final.usage.model_dump(exclude_none=True, include=set(_USAGE)) == _USAGE, final.usage + row: Final = _spend_row(message_id) + + assert float(str(row["spend"])) == pytest.approx(_EXPECTED_SPEND), row + assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _OUTPUT_TOKENS), row diff --git a/tests/integration/observability/test_apply_guardrail_endpoint.py b/tests/integration/observability/test_apply_guardrail_endpoint.py index 9fa7d1a74db..f50679bd54e 100644 --- a/tests/integration/observability/test_apply_guardrail_endpoint.py +++ b/tests/integration/observability/test_apply_guardrail_endpoint.py @@ -156,7 +156,7 @@ def test_apply_guardrail_returns_a_block_error_for_generic_guardrails(gateway: G assert error.model_dump() == { "error": { "message": _BLOCK_REASON, - "type": "internal_server_error", + "type": "invalid_request_error", "param": None, "code": 400, } diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 4f5c04bddcc..1220a65528c 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -23,7 +23,7 @@ import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.mcp import mcp_peer, register_mcp, tool_names -from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from openai import APIError, AsyncOpenAI, OpenAI @@ -3263,7 +3263,10 @@ def test_client_disconnect_during_the_body_upload_runs_no_endpoint_and_writes_no with ( wire_server(_openai_upstream) as upstream, owned_proxy_process( - gateway, tmp_path, {"OPENAI_BASE_URL": upstream.url + "/v1", "OPENAI_API_KEY": "synthetic-openai-key"} + gateway, + tmp_path, + {"OPENAI_BASE_URL": upstream.url + "/v1", "OPENAI_API_KEY": "synthetic-openai-key"}, + state=OwnedState.PRIVATE, ) as owned, owned.gateway.scenario() as scenario, ): diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py new file mode 100644 index 00000000000..ae8fd1c350b --- /dev/null +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -0,0 +1,1227 @@ +from __future__ import annotations + +import ast +import asyncio +import itertools +import json +import uuid +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import pytest +import yaml + +from tests.integration._support.client import Gateway, JsonValue, eventually, gateway_from_environment, object_value +from tests.integration._support.otlp_sink import Span, SpanSinks, recorded_spans +from tests.integration._support.process import owned_proxy +from tests.integration._support.upstream import aws_event_stream_frame +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +_HEADER: Final = "x-litellm-applied-guardrails" +_VENDOR_PATH: Final = "/beta/litellm_basic_guardrail_api" +_PROMPT_SHIELD_PATH: Final = "/contentsafety/text:shieldPrompt" +_PROMPT_SHIELD: Final = "prompt-shield-step" +_REQUEST_FIRST: Final = "request-first-step" +_REQUEST_SECOND: Final = "request-second-step" +_REQUEST_BLOCKER: Final = "request-blocking-step" +_REQUEST_ERRORING: Final = "request-erroring-step" +_RESPONSE_FIRST: Final = "response-first-step" +_RESPONSE_SECOND: Final = "response-second-step" +_RESPONSE_BLOCKER: Final = "response-blocking-step" +_LIVE_FIRST: Final = "live-first-step" +_LIVE_SECOND: Final = "live-second-step" +_UNGATED_STEP: Final = "ungated-step" +_JUDGE: Final = "judge-step" +_JUDGE_MODEL: Final = "pipeline-judge" +_STOPS_AFTER_FIRST: Final = "stops-after-first" +_RUNS_EVERY_STEP: Final = "runs-every-step" +_BLOCKS_AT_FIRST: Final = "blocks-at-first" +_RESPONSE_STOPS_AFTER_FIRST: Final = "response-stops-after-first" +_RESPONSE_RUNS_EVERY_STEP: Final = "response-runs-every-step" +_RESPONSE_STOPS_AFTER_SECOND: Final = "response-stops-after-second" +_SHIELDS_THEN_CHECKS: Final = "shields-then-checks" +_MOVES_PAST_FIRST: Final = "moves-past-first" +_RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first" +_SCANS_AFTER_RELEASE: Final = "scans-after-release" +_CANNOT_HOLD_STREAM: Final = "cannot-hold-stream" +_ERRORS_PAST_FIRST: Final = "errors-past-first" +_JUDGES_THEN_CHECKS: Final = "judges-then-checks" +_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/responses", "/v1/messages") +_BEDROCK_MODEL: Final = "pipeline-bedrock" +_BEDROCK_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + vendor: Wire + upstream: Wire + model: str + traces: str + sync_success_log: Path + + +_GEMINI_RESPONSE: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "guarded"}]}, "finishReason": "STOP", "index": 0}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 1, "totalTokenCount": 4}, + "modelVersion": "gemini-2.0-flash", +} +_FINE_TUNING_JOB: Final[dict[str, JsonValue]] = { + "id": "ftjob-pipeline", + "object": "fine_tuning.job", + "created_at": 1700000000, + "error": None, + "fine_tuned_model": None, + "finished_at": None, + "hyperparameters": {"n_epochs": 3}, + "model": "gpt-4o-mini-2024-07-18", + "organization_id": "org-pipeline", + "result_files": [], + "seed": 1, + "status": "queued", + "trained_tokens": None, + "training_file": "file-pipeline", + "validation_file": None, +} +_VIDEO: Final = { + "object": "video", + "model": "sora-2", + "status": "queued", + "created_at": 1, + "progress": 0, + "seconds": "4", + "size": "720x1280", +} +_SYNC_SUCCESS_RECORDER: Final = """ +import json +from pathlib import Path + +from litellm.integrations.custom_logger import CustomLogger + + +class SyncSuccessRecorder(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + metadata = (kwargs.get("standard_logging_object") or {{}}).get("metadata") or {{}} + row = {{"team": metadata.get("user_api_key_team_id"), "applied": metadata.get("applied_guardrails")}} + with Path({path!r}).open("a") as sink: + sink.write(json.dumps(row) + "\\n") + + +recorder = SyncSuccessRecorder() +""" + + +def _completion_chunk(identity: str, delta: dict[str, str], finish_reason: str | None) -> bytes: + chunk: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(chunk)}\n\n".encode() + + +def _event(event: dict[str, object]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _response(identity: str) -> dict[str, object]: + return { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": f"msg_{identity}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "guarded", "annotations": []}], + } + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + } + + +def _responses_reply(request: Request) -> Reply: + identity: Final = f"resp_{uuid.uuid4().hex}" + response: Final = _response(identity) + if json.loads(request.body).get("stream") is not True: + return Reply(body=json.dumps(response).encode()) + return Reply( + content_type="text/event-stream", + chunks=( + _event( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + } + ), + _event( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": "guarded", + } + ), + _event({"type": "response.completed", "sequence_number": 2, "response": response}), + ), + ) + + +def _bedrock_event(event_type: str, payload: dict[str, object]) -> bytes: + return aws_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload).encode(), + ) + + +def _bedrock_reply(request: Request) -> Reply: + action: Final = request.target.rsplit("/", 1)[1] + if action == "converse-stream": + return Reply( + content_type="application/vnd.amazon.eventstream", + chunks=( + _bedrock_event("messageStart", {"role": "assistant"}), + _bedrock_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "guarded"}}), + _bedrock_event("contentBlockStop", {"contentBlockIndex": 0}), + _bedrock_event("messageStop", {"stopReason": "end_turn"}), + _bedrock_event( + "metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}, "metrics": {}} + ), + ), + ) + if action == "invoke": + return Reply( + body=json.dumps( + { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": _BEDROCK_MODEL_ID, + "content": [{"type": "text", "text": "guarded"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + ).encode() + ) + assert action == "converse", request.target + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "guarded"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +def _completion(request: Request) -> Reply: + if request.target.startswith("/model/"): + return _bedrock_reply(request) + if request.target == "/v1/responses": + return _responses_reply(request) + if request.target in ("/v1/images/generations", "/v1/images/edits"): + return Reply(body=json.dumps({"created": 1, "data": [{"b64_json": "aW1n", "revised_prompt": None}]}).encode()) + if request.target == "/v1/videos": + return Reply(body=json.dumps({**_VIDEO, "id": f"video_{uuid.uuid4().hex}"}).encode()) + if request.target == "/v1/fine_tuning/jobs": + return Reply(body=json.dumps(_FINE_TUNING_JOB).encode()) + if request.target.split("?", 1)[0].endswith(":generateContent"): + return Reply(body=json.dumps(_GEMINI_RESPONSE).encode()) + assert request.target == "/v1/chat/completions", request.target + identity: Final = f"chatcmpl-{uuid.uuid4().hex}" + if json.loads(request.body).get("stream") is True: + return Reply( + content_type="text/event-stream", + chunks=( + _completion_chunk(identity, {"role": "assistant", "content": "guarded"}, None), + _completion_chunk(identity, {}, "stop"), + b"data: [DONE]\n\n", + ), + ) + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "guarded"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +def _vendor(request: Request) -> Reply: + if request.target == f"/{_JUDGE}/v1/decisions": + return Reply( + body=json.dumps( + { + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "unsafe_output", "probability": 0.01}], + "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, + } + ).encode() + ) + if request.target.startswith(f"/{_PROMPT_SHIELD}{_PROMPT_SHIELD_PATH}?"): + return Reply( + body=json.dumps({"userPromptAnalysis": {"attackDetected": False}, "documentsAnalysis": []}).encode() + ) + if request.target.startswith(f"/{_UNGATED_STEP}/"): + return Reply(body=json.dumps({"finalAction": "NONE", "appliedPolicies": []}).encode()) + assert request.target.endswith(_VENDOR_PATH), request.target + if request.target.startswith(f"/{_REQUEST_ERRORING}/"): + return Reply(status=503, body=json.dumps({"error": "guardrail service unavailable"}).encode()) + if request.target.startswith((f"/{_REQUEST_BLOCKER}/", f"/{_RESPONSE_BLOCKER}/")): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "first step refused"}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _guardrail(name: str, mode: str, vendor: Wire, **params: object) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": f"{vendor.url}/{name}", + "api_key": "synthetic-guardrail-key", + **params, + }, + } + + +def _policy(mode: str, steps: tuple[tuple[str, str, str], ...], *, on_error: str | None = None) -> dict[str, object]: + return { + "guardrails": {"add": [guardrail for guardrail, _, _ in steps]}, + "pipeline": { + "mode": mode, + "steps": [ + {"guardrail": guardrail, "on_pass": on_pass, "on_fail": on_fail, "on_error": on_error} + for guardrail, on_pass, on_fail in steps + ], + }, + } + + +def _config(directory: Path, vendor: Wire, upstream: Wire) -> Path: + (directory / "sync_success_recorder.py").write_text( + _SYNC_SUCCESS_RECORDER.format(path=str(directory / "sync-success-log.jsonl")) + ) + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["litellm_settings"] = { + **configuration.get("litellm_settings", {}), + "callbacks": ["otel", "sync_success_recorder.recorder"], + } + configuration["finetune_settings"] = [ + {"custom_llm_provider": "openai", "api_key": "sk-fixture", "api_base": f"{upstream.url}/v1"} + ] + configuration["model_list"] = [ + *configuration.get("model_list", []), + { + "model_name": _BEDROCK_MODEL, + "litellm_params": { + "model": f"bedrock/{_BEDROCK_MODEL_ID}", + "api_base": upstream.url, + "aws_access_key_id": "AKIASYNTHETICPROVIDER", + "aws_secret_access_key": "synthetic-provider-secret", + "aws_region_name": "us-east-1", + }, + }, + { + "model_name": _JUDGE_MODEL, + "litellm_params": { + "model": "openai/gpt-6-luna", + "api_base": f"{vendor.url}/{_JUDGE}/v1", + "api_key": "synthetic-judge-key", + }, + }, + ] + configuration["guardrails"] = [ + *( + _guardrail(name, "pre_call", vendor) + for name in (_REQUEST_FIRST, _REQUEST_SECOND, _REQUEST_BLOCKER, _REQUEST_ERRORING) + ), + *(_guardrail(name, "post_call", vendor) for name in (_RESPONSE_FIRST, _RESPONSE_SECOND, _RESPONSE_BLOCKER)), + *( + _guardrail(name, "post_call", vendor, streaming_buffer_until_moderated=False) + for name in (_LIVE_FIRST, _LIVE_SECOND) + ), + { + "guardrail_name": _UNGATED_STEP, + "litellm_params": { + "guardrail": "dynamoai", + "mode": "post_call", + "default_on": False, + "api_base": f"{vendor.url}/{_UNGATED_STEP}", + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": _JUDGE, + "litellm_params": { + "guardrail": "decision_model", + "mode": "post_call", + "default_on": False, + "decision_model": _JUDGE_MODEL, + "checks": [ + { + "name": "unsafe_output", + "instructions": "Does the response say anything unsafe?", + "action": "block", + "threshold": 0.5, + } + ], + }, + }, + { + "guardrail_name": _PROMPT_SHIELD, + "litellm_params": { + "guardrail": "azure/prompt_shield", + "mode": "pre_call", + "default_on": False, + "api_base": f"{vendor.url}/{_PROMPT_SHIELD}", + "api_key": "synthetic-guardrail-key", + }, + }, + ] + configuration["policies"] = { + _STOPS_AFTER_FIRST: _policy( + "pre_call", ((_REQUEST_FIRST, "allow", "block"), (_REQUEST_SECOND, "next", "block")) + ), + _RUNS_EVERY_STEP: _policy("pre_call", ((_REQUEST_FIRST, "next", "block"), (_REQUEST_SECOND, "next", "block"))), + _BLOCKS_AT_FIRST: _policy( + "pre_call", ((_REQUEST_BLOCKER, "next", "block"), (_REQUEST_SECOND, "next", "block")) + ), + _RESPONSE_STOPS_AFTER_FIRST: _policy( + "post_call", ((_RESPONSE_FIRST, "allow", "block"), (_RESPONSE_SECOND, "next", "block")) + ), + _RESPONSE_RUNS_EVERY_STEP: _policy( + "post_call", ((_RESPONSE_FIRST, "next", "block"), (_RESPONSE_SECOND, "next", "block")) + ), + _RESPONSE_STOPS_AFTER_SECOND: _policy( + "post_call", ((_RESPONSE_SECOND, "allow", "block"), (_RESPONSE_FIRST, "next", "block")) + ), + _SHIELDS_THEN_CHECKS: _policy( + "pre_call", ((_PROMPT_SHIELD, "next", "block"), (_REQUEST_FIRST, "next", "block")) + ), + _MOVES_PAST_FIRST: _policy( + "pre_call", ((_REQUEST_BLOCKER, "next", "next"), (_REQUEST_SECOND, "next", "block")) + ), + _RESPONSE_MOVES_PAST_FIRST: _policy( + "post_call", ((_RESPONSE_BLOCKER, "next", "next"), (_RESPONSE_SECOND, "next", "block")) + ), + _SCANS_AFTER_RELEASE: _policy("post_call", ((_LIVE_FIRST, "next", "next"), (_LIVE_SECOND, "next", "next"))), + _CANNOT_HOLD_STREAM: _policy( + "post_call", ((_UNGATED_STEP, "next", "block"), (_RESPONSE_SECOND, "next", "block")) + ), + _JUDGES_THEN_CHECKS: _policy("post_call", ((_JUDGE, "next", "block"), (_RESPONSE_SECOND, "next", "block"))), + _ERRORS_PAST_FIRST: _policy( + "pre_call", ((_REQUEST_ERRORING, "next", "block"), (_REQUEST_SECOND, "next", "block")), on_error="next" + ), + } + configuration["policy_attachments"] = [{"policy": policy, "tags": [policy]} for policy in configuration["policies"]] + path: Final = directory / "guardrail-pipeline-applied-header.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("guardrail-pipeline-applied-header") + with ( + wire_server(_vendor) as vendor, + wire_server(_completion) as upstream, + gateway_from_environment() as shared, + owned_proxy( + shared, + directory, + {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, "OTEL_BSP_SCHEDULE_DELAY": "100"}, + config=_config(directory, vendor, upstream), + workers=2, + ) as owned, + owned.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture") + yield _Rig(owned, vendor, upstream, model, audit_sinks.operator, directory / "sync-success-log.jsonl") + + +def _body(path: str, model: str, stream: bool) -> dict[str, object]: + prompt: Final = f"hello {uuid.uuid4().hex}" + if path == "/v1/responses": + return {"model": model, "input": prompt, "stream": stream} + messages: Final = [{"role": "user", "content": prompt}] + if path == "/v1/messages": + return {"model": model, "max_tokens": 20, "messages": messages, "stream": stream} + return {"model": model, "messages": messages, "stream": stream} + + +def _ask_as_new_team( + rig: _Rig, + tag: str, + *, + stream: bool = False, + path: str = "/v1/chat/completions", + extra: Mapping[str, object] = MappingProxyType({}), +) -> tuple[str, httpx.Response]: + rig.vendor.drain() + rig.upstream.drain() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [tag]}) + key: Final = scenario.key(team_id=team) + return team, rig.gateway.request("POST", path, {**_body(path, rig.model, stream), **extra}, key=key) + + +def _ask_as_team_tagged( + rig: _Rig, + tag: str, + *, + stream: bool = False, + path: str = "/v1/chat/completions", + extra: Mapping[str, object] = MappingProxyType({}), +) -> httpx.Response: + return _ask_as_new_team(rig, tag, stream=stream, path=path, extra=extra)[1] + + +def _steps_that_ran(rig: _Rig) -> tuple[str, ...]: + return tuple(request.target.split("/")[1] for request in rig.vendor.drain()) + + +def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert len(rig.upstream.drain()) == 1 + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +def test_a_streamed_request_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert len(rig.upstream.drain()) == 1 + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +def test_a_request_pipeline_that_runs_every_step_names_every_step_in_order(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RUNS_EVERY_STEP) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST, _REQUEST_SECOND) + assert response.headers.get(_HEADER) == f"{_REQUEST_FIRST},{_REQUEST_SECOND}", dict(response.headers) + + +def test_a_request_pipeline_names_a_prompt_shield_step_that_ran_before_a_generic_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _SHIELDS_THEN_CHECKS) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_PROMPT_SHIELD, _REQUEST_FIRST) + assert set(response.headers.get(_HEADER, "").split(",")) == {_PROMPT_SHIELD, _REQUEST_FIRST}, dict(response.headers) + + +def test_a_request_pipeline_that_blocks_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _BLOCKS_AT_FIRST) + assert response.status_code == 400, response.text + assert _steps_that_ran(rig) == (_REQUEST_BLOCKER,) + assert rig.upstream.drain() == () + assert response.headers.get(_HEADER) == _REQUEST_BLOCKER, dict(response.headers) + + +def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_streamed_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_request_pipeline_names_a_step_that_flagged_the_request_before_moving_on(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _MOVES_PAST_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_BLOCKER, _REQUEST_SECOND) + assert len(rig.upstream.drain()) == 1 + assert set(response.headers.get(_HEADER, "").split(",")) == {_REQUEST_BLOCKER, _REQUEST_SECOND}, dict( + response.headers + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +def test_a_response_pipeline_names_a_step_that_flagged_the_response_before_moving_on(rig: _Rig, stream: bool) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_MOVES_PAST_FIRST, stream=stream) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_BLOCKER, _RESPONSE_SECOND) + assert set(response.headers.get(_HEADER, "").split(",")) == {_RESPONSE_BLOCKER, _RESPONSE_SECOND}, dict( + response.headers + ) + + +def test_a_stream_released_before_its_pipeline_scans_it_names_the_steps_in_its_headers(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _SCANS_AFTER_RELEASE, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_LIVE_FIRST, _LIVE_SECOND) + assert set(response.headers.get(_HEADER, "").split(",")) == {_LIVE_FIRST, _LIVE_SECOND}, dict(response.headers) + + +def test_a_stream_its_pipeline_cannot_hold_names_the_steps_that_run_on_their_own(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _CANNOT_HOLD_STREAM, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _RESPONSE_SECOND in _steps_that_ran(rig) + assert set(response.headers.get(_HEADER, "").split(",")) == {_UNGATED_STEP, _RESPONSE_SECOND}, dict( + response.headers + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step_on_every_endpoint( + rig: _Rig, path: str, stream: bool +) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, stream=stream, path=path) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert len(rig.upstream.drain()) == 1 + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step_on_every_endpoint( + rig: _Rig, path: str, stream: bool +) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=stream, path=path) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_request_pipeline_names_a_step_whose_service_errored_before_moving_on(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _ERRORS_PAST_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_ERRORING, _REQUEST_SECOND) + assert len(rig.upstream.drain()) == 1 + assert set(response.headers.get(_HEADER, "").split(",")) == {_REQUEST_ERRORING, _REQUEST_SECOND}, dict( + response.headers + ) + + +@pytest.mark.parametrize( + "forged", + [ + pytest.param( + { + "metadata": { + "applied_guardrails": [_REQUEST_SECOND], + "_pre_call_executed_guardrails": [f"deadbeef:{_REQUEST_SECOND}", _REQUEST_SECOND], + } + }, + id="names-in-its-metadata", + ), + pytest.param({"guardrails": [_REQUEST_SECOND]}, id="asks-for-the-skipped-step"), + ], +) +def test_a_request_cannot_get_a_step_that_never_ran_named(rig: _Rig, forged: Mapping[str, object]) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, extra=forged) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +def _spans_naming_guardrails(rig: _Rig, since: int, team: str) -> tuple[Span, ...]: + spans: Final = recorded_spans(rig.traces, since)[1] + return tuple( + span + for span in spans + if span["attributes"].get("metadata.user_api_key_team_id") == team + and "metadata.applied_guardrails" in span["attributes"] + ) + + +@pytest.mark.parametrize( + ("tag", "path", "stream", "expected"), + [ + pytest.param(_STOPS_AFTER_FIRST, "/v1/chat/completions", False, [_REQUEST_FIRST], id="request-pipeline"), + pytest.param(_STOPS_AFTER_FIRST, "/v1/chat/completions", True, [_REQUEST_FIRST], id="request-pipeline-stream"), + pytest.param(_STOPS_AFTER_FIRST, "/v1/responses", True, [_REQUEST_FIRST], id="request-pipeline-responses"), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/chat/completions", False, [_RESPONSE_FIRST], id="response-pipeline" + ), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/chat/completions", True, [_RESPONSE_FIRST], id="response-pipeline-stream" + ), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/messages", True, [_RESPONSE_FIRST], id="response-pipeline-messages" + ), + ], +) +def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran( + rig: _Rig, tag: str, path: str, stream: bool, expected: list[str] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + team, response = _ask_as_new_team(rig, tag, stream=stream, path=path) + assert response.status_code == 200, response.text + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == 1, + seconds=70, + ) + spans: Final = eventually(lambda: _spans_naming_guardrails(rig, since, team), bool, seconds=30) + assert object_value(rows[0]["metadata"]).get("applied_guardrails") == expected, rows + assert [ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"])) for span in spans] == [ + expected + ] * len(spans), spans + + +@pytest.mark.parametrize( + ("tag", "model_guardrail", "expected"), + [ + pytest.param(_STOPS_AFTER_FIRST, _REQUEST_SECOND, [_REQUEST_FIRST, _REQUEST_SECOND], id="request-pipeline"), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, _RESPONSE_SECOND, [_RESPONSE_FIRST, _RESPONSE_SECOND], id="response-pipeline" + ), + ], +) +def test_a_model_level_guardrail_the_pipeline_skipped_is_named_everywhere_once_the_model_runs_it( + rig: _Rig, tag: str, model_guardrail: str, expected: list[str] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + with rig.gateway.scenario() as scenario: + model: Final = scenario.model( + api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture", guardrails=[model_guardrail] + ) + team: Final = scenario.team(metadata={"tags": [tag]}) + key: Final = scenario.key(team_id=team) + rig.vendor.drain() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _body("/v1/chat/completions", model, False), key=key + ) + header, rows, spans = _applied_in_header_spend_log_and_trace(rig, since, team, response) + assert response.status_code == 200, response.text + assert sorted(_steps_that_ran(rig)) == expected + assert (header, rows, spans) == (expected, [expected], [expected] * len(spans)) + + +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +def test_a_response_pipeline_names_a_judge_step_that_records_its_own_results(rig: _Rig, stream: bool) -> None: + team, response = _ask_as_new_team(rig, _JUDGES_THEN_CHECKS, stream=stream) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_JUDGE, _RESPONSE_SECOND) + assert sorted(response.headers.get(_HEADER, "").split(",")) == [_JUDGE, _RESPONSE_SECOND], dict(response.headers) + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == 1, + seconds=70, + ) + assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [_JUDGE, _RESPONSE_SECOND] + + +def _bedrock_body(action: str) -> dict[str, object]: + prompt: Final = f"hello {uuid.uuid4().hex}" + if action == "invoke": + return { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 20, + "messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}], + } + return {"messages": [{"role": "user", "content": [{"text": prompt}]}]} + + +@pytest.mark.parametrize("action", ["converse", "invoke", "converse-stream"]) +def test_a_pass_through_route_names_its_response_pipeline_up_front_because_its_headers_are_built_first( + rig: _Rig, action: str +) -> None: + rig.vendor.drain() + rig.upstream.drain() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_STOPS_AFTER_FIRST]}) + response: Final = rig.gateway.request( + "POST", f"/bedrock/model/{_BEDROCK_MODEL}/{action}", _bedrock_body(action), key=scenario.key(team_id=team) + ) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert [request.target.rsplit("/", 1)[1] for request in rig.upstream.drain()] == [action] + assert sorted(response.headers.get(_HEADER, "").split(",")) == [_RESPONSE_FIRST, _RESPONSE_SECOND], dict( + response.headers + ) + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == 1, + seconds=70, + ) + assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [ + _RESPONSE_FIRST, + _RESPONSE_SECOND, + ], rows + + +def test_an_image_generation_names_its_response_pipeline_up_front_because_its_spend_log_is_built_first( + rig: _Rig, +) -> None: + rig.upstream.drain() + with rig.gateway.scenario() as scenario: + image_model: Final = scenario.model( + model="openai/dall-e-3", api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture" + ) + team: Final = scenario.team(metadata={"tags": [_RESPONSE_STOPS_AFTER_FIRST]}) + response: Final = rig.gateway.request( + "POST", + "/v1/images/generations", + {"model": image_model, "prompt": f"a cat {uuid.uuid4().hex}", "size": "1024x1024"}, + key=scenario.key(team_id=team), + ) + assert response.status_code == 200, response.text + assert [request.target for request in rig.upstream.drain()] == ["/v1/images/generations"] + assert sorted(response.headers.get(_HEADER, "").split(",")) == [_RESPONSE_FIRST, _RESPONSE_SECOND], dict( + response.headers + ) + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == 1, + seconds=70, + ) + assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [ + _RESPONSE_FIRST, + _RESPONSE_SECOND, + ], rows + + +def _applied_in_header_spend_log_and_trace( + rig: _Rig, since: int, team: str, response: httpx.Response, *, rows_expected: int = 1 +) -> tuple[list[str], list[list[str]], list[list[str]]]: + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == rows_expected, + seconds=70, + ) + spans: Final = eventually( + lambda: _spans_naming_guardrails(rig, since, team), lambda found: len(found) >= rows_expected, seconds=30 + ) + return ( + sorted(name for name in response.headers.get(_HEADER, "").split(",") if name), + [sorted(object_value(row["metadata"]).get("applied_guardrails") or []) for row in rows], + [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans], + ) + + +def _cache_hits_and_metadata_logged(rig: _Rig, team: str, count: int) -> list[dict[str, JsonValue]]: + return eventually( + lambda: rig.gateway.rows( + 'SELECT cache_hit, metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s ORDER BY "startTime"', (team,) + ), + lambda found: len(found) == count, + seconds=70, + ) + + +def test_a_cache_hit_names_every_response_step_that_ran_in_the_spend_log_and_the_trace(rig: _Rig) -> None: + since: Final = recorded_spans(rig.traces)[0] + rig.vendor.drain() + body: Final = {"model": rig.model, "messages": [{"role": "user", "content": f"cache me {uuid.uuid4().hex}"}]} + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_RUNS_EVERY_STEP]}) + key: Final = scenario.key(team_id=team) + first: Final = rig.gateway.request("POST", "/v1/chat/completions", body, key=key) + _cache_hits_and_metadata_logged(rig, team, 1) + cached: Final = rig.gateway.request("POST", "/v1/chat/completions", body, key=key) + rows: Final = _cache_hits_and_metadata_logged(rig, team, 2) + spans: Final = eventually(lambda: _spans_naming_guardrails(rig, since, team), lambda found: len(found) >= 2) + expected: Final = [_RESPONSE_FIRST, _RESPONSE_SECOND] + assert [first.status_code, cached.status_code] == [200, 200], (first.text, cached.text) + assert _steps_that_ran(rig) == (_RESPONSE_FIRST, _RESPONSE_SECOND) * 2 + assert [row["cache_hit"] == "True" for row in rows] == [False, True], rows + assert [sorted(response.headers.get(_HEADER, "").split(",")) for response in (first, cached)] == [expected] * 2 + assert [sorted(object_value(row["metadata"]).get("applied_guardrails") or []) for row in rows] == [expected] * 2 + assert [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans] == [ + expected + ] * len(spans), spans + + +def test_a_sync_success_callback_sees_every_response_step_that_ran(rig: _Rig) -> None: + team, response = _ask_as_new_team(rig, _RESPONSE_RUNS_EVERY_STEP) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST, _RESPONSE_SECOND) + logged: Final = eventually( + lambda: ( + [row for row in map(json.loads, rig.sync_success_log.read_text().splitlines()) if row["team"] == team] + if rig.sync_success_log.exists() + else [] + ), + lambda found: len(found) == 1, + seconds=30, + ) + assert sorted(logged[0]["applied"] or []) == [_RESPONSE_FIRST, _RESPONSE_SECOND], logged + + +def _image_edit(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + return rig.gateway.request_multipart( + "/v1/images/edits", + {"model": scenario_model(model="openai/gpt-image-1"), "prompt": f"a cat {uuid.uuid4().hex}"}, + {"image": ("cat.png", _PNG, "image/png")}, + key=key, + ) + + +def _video(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + return rig.gateway.request( + "POST", + "/v1/videos", + {"model": scenario_model(model="openai/sora-2"), "prompt": f"a cat {uuid.uuid4().hex}"}, + key=key, + ) + + +def _generate_content(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + return rig.gateway.request( + "POST", + f"/v1beta/models/{rig.model}:generateContent", + {"contents": [{"role": "user", "parts": [{"text": f"hello {uuid.uuid4().hex}"}]}]}, + key=key, + ) + + +def _stream_generate_content(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + return rig.gateway.request( + "POST", + f"/v1beta/models/{rig.model}:streamGenerateContent", + {"contents": [{"role": "user", "parts": [{"text": f"hello {uuid.uuid4().hex}"}]}]}, + key=key, + params={"alt": "sse"}, + ) + + +def _native_generate_content(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + model: Final = scenario_model(model="gemini/gemini-2.0-flash", api_base=rig.upstream.url) + return rig.gateway.request( + "POST", + f"/v1beta/models/{model}:generateContent", + {"contents": [{"role": "user", "parts": [{"text": f"hello {uuid.uuid4().hex}"}]}]}, + key=key, + ) + + +_PNG: Final = bytes.fromhex( + "89504e470d0a1a0a0000000d4948445200000001000000010806000000" + "1f15c4890000000d49444154789c6360000002000154a24f5d0000000049454e44ae426082" +) +_ROUTES: Final = ( + pytest.param(_image_edit, id="image-edit"), + pytest.param(_video, id="video"), + pytest.param(_generate_content, id="generate-content"), + pytest.param(_stream_generate_content, id="stream-generate-content"), + pytest.param(_native_generate_content, id="native-generate-content"), +) + + +@pytest.mark.parametrize("send", _ROUTES) +def test_a_response_pipeline_that_runs_every_step_names_both_steps_on_every_route( + rig: _Rig, send: Callable[[_Rig, Callable[..., str], str], httpx.Response] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_RUNS_EVERY_STEP]}) + response: Final = send( + rig, + partial(scenario.model, api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture"), + scenario.key(team_id=team), + ) + assert response.status_code == 200, response.text + expected: Final = [_RESPONSE_FIRST, _RESPONSE_SECOND] + header, rows, spans = _applied_in_header_spend_log_and_trace(rig, since, team, response) + assert (header, rows, spans) == (expected, [expected], [expected] * len(spans)) + + +@pytest.mark.parametrize("send", _ROUTES) +def test_a_response_pipeline_that_stops_early_names_the_same_steps_in_the_header_and_the_spend_log_on_every_route( + rig: _Rig, send: Callable[[_Rig, Callable[..., str], str], httpx.Response] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_STOPS_AFTER_FIRST]}) + response: Final = send( + rig, + partial(scenario.model, api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture"), + scenario.key(team_id=team), + ) + assert response.status_code == 200, response.text + header, rows, spans = _applied_in_header_spend_log_and_trace(rig, since, team, response) + assert _RESPONSE_FIRST in header, header + assert (rows, spans) == ([header], [header] * len(spans)) + + +def test_a_fine_tuning_job_with_a_method_object_and_a_guardrail_is_created(rig: _Rig) -> None: + rig.upstream.drain() + method: Final = {"type": "dpo", "dpo": {"hyperparameters": {"beta": 0.1}}} + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(team_id=scenario.team(metadata={"tags": [_STOPS_AFTER_FIRST]})) + response: Final = rig.gateway.request( + "POST", + "/v1/fine_tuning/jobs", + { + "model": "gpt-4o-mini-2024-07-18", + "training_file": "file-pipeline", + "custom_llm_provider": "openai", + "method": method, + "guardrails": [_REQUEST_FIRST], + }, + key=key, + ) + assert response.status_code == 200, response.text + assert [request.target for request in rig.upstream.drain()] == ["/v1/fine_tuning/jobs"] + + +def _agent_card(name: str, url: str) -> dict[str, object]: + return { + "protocolVersion": "0.3", + "name": name, + "description": "Scripted agent that answers every message with one text part", + "version": "1.0.0", + "url": url, + "capabilities": {"streaming": True}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + + +@pytest.fixture +def agent(rig: _Rig) -> Iterator[str]: + name: Final = f"pipeline-agent-{uuid.uuid4().hex}" + + def answer(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps(_agent_card(name, f"{peer.url}/")).encode()) + message: Final = { + "kind": "message", + "role": "agent", + "messageId": uuid.uuid4().hex, + "parts": [{"kind": "text", "text": "guarded"}], + } + return Reply( + body=json.dumps({"jsonrpc": "2.0", "id": json.loads(request.body)["id"], "result": message}).encode() + ) + + with wire_server(answer) as peer: + identity: Final = str( + rig.gateway.post( + "/v1/agents", {"agent_name": name, "agent_card_params": _agent_card(name, f"{peer.url}/")} + )["agent_id"] + ) + yield identity + rig.gateway.request("DELETE", f"/v1/agents/{identity}") + + +@pytest.mark.parametrize( + ("tags", "expected"), + [ + pytest.param((_SCANS_AFTER_RELEASE,), [_LIVE_FIRST, _LIVE_SECOND], id="every-step-passes"), + pytest.param( + (_RESPONSE_MOVES_PAST_FIRST,), [_RESPONSE_BLOCKER, _RESPONSE_SECOND], id="first-step-flags-and-moves-on" + ), + pytest.param((_RESPONSE_STOPS_AFTER_FIRST,), [_RESPONSE_FIRST], id="first-step-allows"), + pytest.param( + (_RESPONSE_STOPS_AFTER_FIRST, _RESPONSE_STOPS_AFTER_SECOND), + [_RESPONSE_FIRST, _RESPONSE_SECOND], + id="two-policies-each-run-the-step-the-other-skipped", + ), + ], +) +def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_and_the_trace( + rig: _Rig, agent: str, tags: tuple[str, ...], expected: list[str] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + rig.vendor.drain() + message: Final = { + "kind": "message", + "role": "user", + "messageId": uuid.uuid4().hex, + "parts": [{"kind": "text", "text": f"hello {uuid.uuid4().hex}"}], + } + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": list(tags)}) + response: Final = rig.gateway.request( + "POST", + f"/a2a/{agent}", + {"jsonrpc": "2.0", "id": uuid.uuid4().hex, "method": "message/send", "params": {"message": message}}, + key=scenario.key(team_id=team), + headers={"Accept": "application/json"}, + ) + assert response.status_code == 200, response.text + assert sorted(_steps_that_ran(rig)) == expected + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == 1, + seconds=70, + ) + spans: Final = eventually(lambda: _spans_naming_guardrails(rig, since, team), bool, seconds=30) + assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == expected, rows + assert [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans] == [ + expected + ] * len(spans), spans + + +def _openai_chat(base: str, key: str, model: str) -> str | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client, + ): + response: Final = client.chat.completions.with_raw_response.create( + model=model, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}] + ) + return response.headers.get(_HEADER) + + +def _openai_responses(base: str, key: str, model: str) -> str | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client, + ): + return client.responses.with_raw_response.create(model=model, input=f"hello {uuid.uuid4().hex}").headers.get( + _HEADER + ) + + +def _async_openai_streamed_chat(base: str, key: str, model: str) -> str | None: + async def run() -> str | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client, + client.chat.completions.with_streaming_response.create( + model=model, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}], stream=True + ) as response, + ): + lines: Final = [line async for line in response.iter_lines()] + assert "data: [DONE]" in lines, lines + return response.headers.get(_HEADER) + + return asyncio.run(run()) + + +def _async_openai_streamed_responses(base: str, key: str, model: str) -> str | None: + async def run() -> str | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client, + client.responses.with_streaming_response.create( + model=model, input=f"hello {uuid.uuid4().hex}", stream=True + ) as response, + ): + lines: Final = [line async for line in response.iter_lines()] + assert any("response.completed" in line for line in lines), lines + return response.headers.get(_HEADER) + + return asyncio.run(run()) + + +def _anthropic_streamed_messages(base: str, key: str, model: str) -> str | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + anthropic.Anthropic(api_key=key, base_url=base, max_retries=0, http_client=transport) as client, + client.messages.with_streaming_response.create( + model=model, + max_tokens=20, + messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}], + stream=True, + ) as response, + ): + lines: Final = tuple(response.iter_lines()) + assert any("message_stop" in line for line in lines), lines + return response.headers.get(_HEADER) + + +def _async_anthropic_messages(base: str, key: str, model: str) -> str | None: + async def run() -> str | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + anthropic.AsyncAnthropic(api_key=key, base_url=base, max_retries=0, http_client=transport) as client, + ): + response: Final = await client.messages.with_raw_response.create( + model=model, max_tokens=20, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}] + ) + return response.headers.get(_HEADER) + + return asyncio.run(run()) + + +@pytest.mark.parametrize( + "client", + [ + pytest.param(_openai_chat, id="openai-chat"), + pytest.param(_async_openai_streamed_chat, id="async-openai-chat-stream"), + pytest.param(_openai_responses, id="openai-responses"), + pytest.param(_async_openai_streamed_responses, id="async-openai-responses-stream"), + pytest.param(_anthropic_streamed_messages, id="anthropic-messages-stream"), + pytest.param(_async_anthropic_messages, id="async-anthropic-messages"), + ], +) +def test_every_sdk_client_sees_only_the_step_that_ran(rig: _Rig, client: Callable[[str, str, str], str | None]) -> None: + rig.vendor.drain() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(team_id=scenario.team(metadata={"tags": [_STOPS_AFTER_FIRST]})) + header: Final = client(str(rig.gateway.client.base_url).rstrip("/"), key, rig.model) + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert header == _REQUEST_FIRST + + +async def _send_concurrently( + rig: _Rig, keys: Mapping[str, str], cases: tuple[tuple[str, str, bool], ...] +) -> tuple[httpx.Response, ...]: + async with httpx.AsyncClient(base_url=rig.gateway.client.base_url, timeout=60, trust_env=False) as client: + return tuple( + await asyncio.gather( + *( + client.post( + path, json=_body(path, rig.model, stream), headers={"Authorization": f"Bearer {keys[tag]}"} + ) + for tag, path, stream in cases + ) + ) + ) + + +def test_a_concurrent_burst_names_only_each_requests_own_steps(rig: _Rig) -> None: + expected: Final = MappingProxyType( + { + _STOPS_AFTER_FIRST: _REQUEST_FIRST, + _RUNS_EVERY_STEP: f"{_REQUEST_FIRST},{_REQUEST_SECOND}", + _RESPONSE_STOPS_AFTER_FIRST: _RESPONSE_FIRST, + } + ) + cases: Final = tuple(itertools.product(expected, _ENDPOINTS, (False, True))) * 2 + with rig.gateway.scenario() as scenario: + keys: Final = MappingProxyType( + {tag: scenario.key(team_id=scenario.team(metadata={"tags": [tag]})) for tag in expected} + ) + responses: Final = asyncio.run(_send_concurrently(rig, keys, cases)) + rig.vendor.drain() + rig.upstream.drain() + assert [ + (tag, path, stream, response.status_code, response.headers.get(_HEADER)) + for (tag, path, stream), response in zip(cases, responses, strict=True) + ] == [(tag, path, stream, 200, expected[tag]) for tag, path, stream in cases] diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 854da6fce0b..4441582fb6b 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -1037,7 +1037,7 @@ def test_s07_a_5kb_caller_model_under_the_wildcard_is_a_bounded_4xx_and_liveline error: Final = _JSON.validate_python(_JSON.validate_python(json.loads(outcome.raw))["error"]) assert "a" * 5120 in str(error["message"]), outcome.raw[:300] assert_within_factor_of_control(oversized.seconds, control=control.seconds, factor=20) - assert rig.wire.drain() == () + assert len(rig.wire.drain()) == 1 _failure_row(outcome.call_id) @@ -1325,7 +1325,9 @@ def test_x01_flipping_the_flag_to_true_through_a_patch_update_turns_cache_points @pytest.mark.timeout(600) -def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone(gateway: Gateway) -> None: +def test_x06_flipping_the_flag_to_true_through_the_legacy_post_update_turns_cache_points_back_on( + gateway: Gateway, +) -> None: arn: Final = _arn("legacy") with wire_server(_peer()) as wire, gateway.scenario() as scenario: name: Final = _deployment( @@ -1333,6 +1335,7 @@ def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone ) before, before_received = _observe_at(gateway, wire, "chat", _chat_body(name, _prompt(), _SYSTEM_AND_USER)) _assert_answered(before) + assert before_received.cache_points == 0, before_received.body identity: Final = _model_id(name) updated: Final = gateway.request( "POST", @@ -1344,13 +1347,14 @@ def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone }, ) assert updated.status_code == 200, updated.text - assert _stored_flag(identity) is False - _settled(gateway, name, wire) - after, after_received = _observe_at(gateway, wire, "chat", _chat_body(name, _prompt(), _SYSTEM_AND_USER)) - _assert_answered(after) - assert after_received.cache_points == before_received.cache_points, (before_received.body, after_received.body) + assert _stored_flag(identity) is True + points: Final = eventually( + lambda: _six_cache_point_counts(gateway, wire, name), + lambda counts: len(counts) == 6 and all(count == 2 for count in counts), + seconds=60, + ) + assert points == (2,) * 6, points _success_row(before.response_id) - _success_row(after.response_id) @dataclass(frozen=True, slots=True) diff --git a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py index 37aa96c2467..23b75a5a8cf 100644 --- a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py @@ -419,12 +419,11 @@ def test_context_window_fallback_is_not_consulted_on_streamed_messages(cwf_rig: _assert_no_fallback(cwf_rig, "overflow", prompt) -def test_context_window_fallback_does_not_fire_on_streamed_responses(cwf_rig: Rig) -> None: +def test_context_window_fallback_fires_on_streamed_responses(cwf_rig: Rig) -> None: prompt: Final = _prompt("cwf responses stream") response, text = _consumed(cwf_rig.gateway, "/v1/responses", _responses(cwf_rig.models["cwf"], prompt, True)) assert response.status_code == 200, text - assert _GENERIC in str(_only_failed(text)["message"]), text - _assert_no_fallback(cwf_rig, "overflow", prompt) + _assert_fell_back(cwf_rig, "overflow", prompt, text) def test_plain_fallback_fires_for_the_overflow_on_chat_completions(fb_rig: Rig) -> None: @@ -467,12 +466,11 @@ def test_streamed_messages_overflow_on_a_plain_fallback_deployment_returns_400_w _assert_no_fallback(fb_rig, "overflow", prompt) -def test_plain_fallback_does_not_fire_on_streamed_responses(fb_rig: Rig) -> None: +def test_plain_fallback_fires_on_streamed_responses(fb_rig: Rig) -> None: prompt: Final = _prompt("fb responses stream") response, text = _consumed(fb_rig.gateway, "/v1/responses", _responses(fb_rig.models["fb"], prompt, True)) assert response.status_code == 200, text - assert _GENERIC in str(_only_failed(text)["message"]), text - _assert_no_fallback(fb_rig, "overflow", prompt) + _assert_fell_back(fb_rig, "overflow", prompt, text) def test_streamed_messages_overflow_on_an_openai_fallback_deployment_returns_the_error_event_without_falling_back( diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py index 90ff32e01dd..20472dc2399 100644 --- a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py @@ -501,7 +501,7 @@ def test_vertex_batch_file_rows_carry_the_mapped_voice(gateway: Gateway) -> None assert tuple(text_in(row) for row in rows) == (first, second) -def test_vertex_live_default_mode_sends_one_setup_without_the_client_voice(gateway: Gateway) -> None: +def test_vertex_live_sends_one_setup_carrying_the_mapped_session_update_voice(gateway: Gateway) -> None: with gateway.scenario() as scenario: handle: Final = live_scenario(scenario) key: Final = scenario.key() @@ -511,7 +511,7 @@ def test_vertex_live_default_mode_sends_one_setup_without_the_client_voice(gatew assert types[0] == "session.created" and types[-1] == "response.done", frames setups: Final = live_setups(gateway, handle.scenario_id) assert len(setups) == 1, setups - assert voice_in(setups[0]) == NO_VOICE, setups[0] + assert voice_in(setups[0]) == "Umbriel", setups[0] @dataclass(frozen=True, slots=True) diff --git a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py index fad7898871d..018c6a653d1 100644 --- a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py +++ b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py @@ -279,7 +279,11 @@ def cooldown_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Cooldown upstream, directory, {"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "2"}, - config=cooldown_config(directory, router_settings(cache, allowed_fails=0), prometheus=False), + config=cooldown_config( + directory, + router_settings(cache, allowed_fails=0, allowed_fails_policy={"TimeoutErrorAllowedFails": 0}), + prometheus=False, + ), ) ) yield _CooldownRig(owned.gateway, cache) diff --git a/tests/integration/providers/test_responses_bridge_verbosity_wire.py b/tests/integration/providers/test_responses_bridge_verbosity_wire.py new file mode 100644 index 00000000000..037450ba820 --- /dev/null +++ b/tests/integration/providers/test_responses_bridge_verbosity_wire.py @@ -0,0 +1,413 @@ +import uuid +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support import responses_vendor as rv +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.openai_wire import answering_model_discovery, chat_reply, responses_reply +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(120) + +_GPT: Final = "gpt-5.6" +_DATABRICKS: Final = "databricks-meta-llama-4-maverick" +_PROVIDER_KEY: Final = "synthetic-provider-key" +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}, "celsius": {"type": "number"}}, + "required": ["city", "celsius"], + "additionalProperties": False, +} +_TEXT_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "name": "weather", + "schema": _WEATHER_SCHEMA, + "strict": True, +} +_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "weather", "schema": _WEATHER_SCHEMA, "strict": True}, +} +_STREAM_OPTIONS: Final[dict[str, JsonValue]] = {"stream_options": {"include_usage": True}} + +Verbosity = Literal["low", "high"] + + +def _prompt(marker: str) -> str: + return f"verbosity probe marker-{marker}" + + +def _answer(marker: str) -> str: + return f"scripted answer marker-{marker}" + + +def _marker_of(request: Request) -> str: + body: Final = request.body.decode() + start: Final = body.rfind("marker-") + return body[start + len("marker-") : start + len("marker-") + 32] if start >= 0 else "health" + + +def _vendor(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body or b"{}") + marker: Final = _marker_of(request) + stream: Final = body.get("stream") is True + if request.target.endswith("/responses"): + return responses_reply(f"resp_{marker}", _GPT, _answer(marker), stream=stream) + return chat_reply(f"chatcmpl-{marker}", str(body.get("model")), _answer(marker), stream=stream) + + +@contextmanager +def _upstream() -> Iterator[Wire]: + with wire_server(answering_model_discovery(_vendor)) as wire: + yield wire + + +def _bridged_gpt(scenario: Scenario, wire: Wire, **parameters: JsonValue) -> str: + return scenario.model( + model=f"openai/{_GPT}", + api_base=f"{wire.url}/v1", + api_key=_PROVIDER_KEY, + use_chat_completions_api=True, + **parameters, + ) + + +def _native_gpt(scenario: Scenario, wire: Wire) -> str: + return scenario.model(model=f"openai/{_GPT}", api_base=f"{wire.url}/v1", api_key=_PROVIDER_KEY) + + +def _databricks(scenario: Scenario, wire: Wire, **parameters: JsonValue) -> str: + return scenario.model( + model=f"databricks/{_DATABRICKS}", + api_base=f"{wire.url}/serving-endpoints", + api_key=_PROVIDER_KEY, + **parameters, + ) + + +def _responses(gateway: Gateway, model: str, marker: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", "/v1/responses", {"model": model, "input": _prompt(marker), "cache": _NO_CACHE, **fields} + ) + + +def _streamed_events(gateway: Gateway, model: str, marker: str, **fields: JsonValue) -> list[dict[str, JsonValue]]: + with gateway.client.stream( + "POST", + "/v1/responses", + json={"model": model, "input": _prompt(marker), "cache": _NO_CACHE, "stream": True, **fields}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read() + return [ + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: {") + ] + + +def _output_text(response: Mapping[str, JsonValue]) -> str: + output: Final = response["output"] + assert isinstance(output, list), response + return "".join( + str(part["text"]) + for item in output + if isinstance(item, dict) and item.get("type") == "message" + for part in item["content"] + if isinstance(part, dict) and part.get("type") == "output_text" + ) + + +def _answered(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + body: Final = _JSON_OBJECT.validate_json(response.text) + assert body["status"] == "completed" and _output_text(body) == _answer(marker), body + + +def _streamed_answer(events: list[dict[str, JsonValue]], marker: str) -> None: + completed: Final = [event for event in events if event["type"] == "response.completed"] + assert len(completed) == 1, [event["type"] for event in events] + final: Final = completed[0]["response"] + assert isinstance(final, dict) and _output_text(final) == _answer(marker), final + deltas: Final = "".join(str(event["delta"]) for event in events if event["type"] == "response.output_text.delta") + assert deltas == _answer(marker), events + + +def _posted(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if request.method == "POST") + + +def _only_post(wire: Wire, target: str) -> dict[str, JsonValue]: + posted: Final = _posted(wire) + assert [request.target for request in posted] == [target], posted + return _JSON_OBJECT.validate_json(posted[0].body) + + +def _chat_messages(marker: str) -> list[JsonValue]: + return [{"role": "user", "content": _prompt(marker)}] + + +def _unsupported_param_error(response: httpx.Response) -> None: + assert response.status_code == 400, response.text + error: Final = _JSON_OBJECT.validate_json(response.text)["error"] + assert isinstance(error, dict) and "verbosity" in str(error["message"]), response.text + + +@pytest.mark.parametrize("verbosity", ["low", "high"]) +def test_regression_explicit_bridge_sends_text_verbosity_as_top_level_verbosity( + gateway: Gateway, verbosity: Verbosity +) -> None: + marker: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _bridged_gpt(scenario, wire) + _answered(_responses(gateway, model, marker, text={"verbosity": verbosity}), marker) + sent: Final = _only_post(wire, "/v1/chat/completions") + assert sent == {"model": _GPT, "messages": _chat_messages(marker), "verbosity": verbosity}, sent + + +@pytest.mark.parametrize("verbosity", ["low", "high"]) +def test_regression_explicit_bridge_stream_sends_text_verbosity_as_top_level_verbosity( + gateway: Gateway, verbosity: Verbosity +) -> None: + marker: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _bridged_gpt(scenario, wire) + _streamed_answer(_streamed_events(gateway, model, marker, text={"verbosity": verbosity}), marker) + sent: Final = _only_post(wire, "/v1/chat/completions") + assert sent == { + "model": _GPT, + "messages": _chat_messages(marker), + "stream": True, + **_STREAM_OPTIONS, + "verbosity": verbosity, + }, sent + + +def test_regression_explicit_bridge_sends_verbosity_beside_the_json_schema_response_format(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _bridged_gpt(scenario, wire) + _answered(_responses(gateway, model, marker, text={"format": _TEXT_FORMAT, "verbosity": "low"}), marker) + sent: Final = _only_post(wire, "/v1/chat/completions") + assert sent == { + "model": _GPT, + "messages": _chat_messages(marker), + "response_format": _RESPONSE_FORMAT, + "verbosity": "low", + }, sent + + +def test_regression_automatic_bridge_refuses_unsupported_verbosity_before_the_wire(gateway: Gateway) -> None: + refused: Final = uuid.uuid4().hex + control: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _databricks(scenario, wire, drop_params=False) + _unsupported_param_error(_responses(gateway, model, refused, text={"verbosity": "low"})) + _answered(_responses(gateway, model, control), control) + posted: Final = _posted(wire) + assert [(request.target, _marker_of(request)) for request in posted] == [ + ("/serving-endpoints/chat/completions", control) + ], posted + + +def test_regression_automatic_bridge_stream_refuses_unsupported_verbosity_before_the_wire(gateway: Gateway) -> None: + refused: Final = uuid.uuid4().hex + control: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _databricks(scenario, wire, drop_params=False) + _unsupported_param_error(_responses(gateway, model, refused, text={"verbosity": "high"}, stream=True)) + _answered(_responses(gateway, model, control), control) + posted: Final = _posted(wire) + assert [(request.target, _marker_of(request)) for request in posted] == [ + ("/serving-endpoints/chat/completions", control) + ], posted + + +@pytest.mark.parametrize("source", ["deployment", "request"]) +def test_regression_allowed_openai_params_forwards_verbosity_through_the_automatic_bridge( + gateway: Gateway, source: Literal["deployment", "request"] +) -> None: + marker: Final = uuid.uuid4().hex + allowed: Final[dict[str, JsonValue]] = {"allowed_openai_params": ["verbosity"]} + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _databricks(scenario, wire, drop_params=False, **(allowed if source == "deployment" else {})) + extra: Final = allowed if source == "request" else {} + _answered(_responses(gateway, model, marker, text={"verbosity": "low"}, **extra), marker) + sent: Final = _only_post(wire, "/serving-endpoints/chat/completions") + assert sent == {"model": _DATABRICKS, "messages": _chat_messages(marker), "verbosity": "low"}, sent + + +@pytest.mark.parametrize("source", ["deployment", "request"]) +def test_guard_drop_params_drops_unsupported_verbosity_and_keeps_the_json_schema_format( + gateway: Gateway, source: Literal["deployment", "request"] +) -> None: + marker: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _databricks(scenario, wire, drop_params=source == "deployment") + extra: Final[dict[str, JsonValue]] = {"drop_params": True} if source == "request" else {} + _answered( + _responses(gateway, model, marker, text={"format": _TEXT_FORMAT, "verbosity": "low"}, **extra), marker + ) + sent: Final = _only_post(wire, "/serving-endpoints/chat/completions") + assert sent == {"model": _DATABRICKS, "messages": _chat_messages(marker), "response_format": _RESPONSE_FORMAT}, sent + + +def test_guard_drop_params_stream_drops_unsupported_verbosity(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _databricks(scenario, wire, drop_params=True) + _streamed_answer(_streamed_events(gateway, model, marker, text={"verbosity": "high"}), marker) + sent: Final = _only_post(wire, "/serving-endpoints/chat/completions") + assert sent == {"model": _DATABRICKS, "messages": _chat_messages(marker), "stream": True}, sent + + +def test_guard_explicit_bridge_without_verbosity_sends_only_the_json_schema_format(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _bridged_gpt(scenario, wire) + _answered(_responses(gateway, model, marker, text={"format": _TEXT_FORMAT}), marker) + sent: Final = _only_post(wire, "/v1/chat/completions") + assert sent == {"model": _GPT, "messages": _chat_messages(marker), "response_format": _RESPONSE_FORMAT}, sent + + +def test_guard_explicit_bridge_without_text_sends_no_verbosity(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _bridged_gpt(scenario, wire) + _answered(_responses(gateway, model, marker), marker) + sent: Final = _only_post(wire, "/v1/chat/completions") + assert sent == {"model": _GPT, "messages": _chat_messages(marker)}, sent + + +@pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) +def test_guard_native_responses_route_forwards_text_verbosity_unchanged(gateway: Gateway, stream: bool) -> None: + marker: Final = uuid.uuid4().hex + text: Final[dict[str, JsonValue]] = {"format": _TEXT_FORMAT, "verbosity": "high"} + with _upstream() as wire, gateway.scenario() as scenario: + model: Final = _native_gpt(scenario, wire) + if stream: + _streamed_answer(_streamed_events(gateway, model, marker, text=text), marker) + else: + _answered(_responses(gateway, model, marker, text=text), marker) + sent: Final = _only_post(wire, "/v1/responses") + assert sent == {"model": _GPT, "input": _prompt(marker), "text": text, **({"stream": True} if stream else {})}, sent + + +_CLIENT_TOOL_SEARCH: Final[dict[str, JsonValue]] = {"type": "tool_search", "execution": "client"} +_CLIENT_TOOL_SEARCH_ON_THE_WIRE: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "tool_search", + "description": ( + "Search the deferred tools by keyword. The matching tools are loaded and become callable on the next turn." + ), + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query for deferred tools."}, + "limit": {"type": "integer", "description": "Maximum number of tools to return. Defaults to 8."}, + }, + "required": ["query"], + "additionalProperties": False, + }, + }, +} + + +@dataclass(frozen=True, slots=True) +class _FollowUpWire: + sent: JsonValue + expected: dict[str, JsonValue] + + +def _stored_prompts_config(directory: Path) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + general: Final = _JSON_OBJECT.validate_python(stock["general_settings"]) + path: Final = directory / "verbosity-stored-turns.yaml" + path.write_text(yaml.safe_dump({**stock, "general_settings": {**general, "store_prompts_in_spend_logs": True}})) + return path + + +@pytest.fixture(scope="module") +def stored_turns_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[OwnedProxy, Wire]]: + directory: Final = tmp_path_factory.mktemp("verbosity-stored-turns") + with _upstream() as wire, gateway_from_environment() as gateway: + with owned_proxy_process(gateway, directory, {}, config=_stored_prompts_config(directory)) as owned: + yield owned, wire + + +def _stored_first_turn(proxy: Gateway, model: str, marker: str) -> str: + first: Final = _responses(proxy, model, marker, tools=[_CLIENT_TOOL_SEARCH]) + _answered(first, marker) + response_id: Final = str(_JSON_OBJECT.validate_json(first.text)["id"]) + eventually( + lambda: proxy.rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda rows: any(rv.same_response(str(row["request_id"]), response_id) for row in rows), + seconds=70, + ) + return response_id + + +def _client_tool_search_follow_up( + stored_turns_proxy: tuple[OwnedProxy, Wire], stream: bool, **fields: JsonValue +) -> _FollowUpWire: + owned, wire = stored_turns_proxy + proxy: Final = owned.gateway + first: Final = uuid.uuid4().hex + second: Final = uuid.uuid4().hex + with proxy.scenario() as scenario: + model: Final = _bridged_gpt(scenario, wire) + previous: Final = _stored_first_turn(proxy, model, first) + wire.drain() + follow_up: Final[dict[str, JsonValue]] = { + "previous_response_id": previous, + "tools": [_CLIENT_TOOL_SEARCH], + **fields, + } + if stream: + _streamed_answer(_streamed_events(proxy, model, second, **follow_up), second) + else: + _answered(_responses(proxy, model, second, **follow_up), second) + sent: Final = _only_post(wire, "/v1/chat/completions") + return _FollowUpWire( + sent, + { + "model": _GPT, + "messages": [ + *_chat_messages(first), + {"role": "assistant", "content": _answer(first)}, + *_chat_messages(second), + ], + "tools": [_CLIENT_TOOL_SEARCH_ON_THE_WIRE], + **({"stream": True, **_STREAM_OPTIONS} if stream else {}), + }, + ) + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) +@pytest.mark.parametrize("verbosity", ["low", "high"]) +def test_regression_client_tool_search_follow_up_sends_text_verbosity_as_top_level_verbosity( + stored_turns_proxy: tuple[OwnedProxy, Wire], stream: bool, verbosity: Verbosity +) -> None: + observed: Final = _client_tool_search_follow_up(stored_turns_proxy, stream, text={"verbosity": verbosity}) + assert observed.sent == {**observed.expected, "verbosity": verbosity}, observed.sent + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) +def test_guard_client_tool_search_follow_up_without_verbosity_sends_no_verbosity( + stored_turns_proxy: tuple[OwnedProxy, Wire], stream: bool +) -> None: + observed: Final = _client_tool_search_follow_up(stored_turns_proxy, stream) + assert observed.sent == observed.expected, observed.sent diff --git a/tests/integration/routing/test_key_router_settings_image_generation.py b/tests/integration/routing/test_key_router_settings_image_generation.py new file mode 100644 index 00000000000..84415c72f2f --- /dev/null +++ b/tests/integration/routing/test_key_router_settings_image_generation.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import json +import uuid +from types import MappingProxyType +from typing import Final + +from integration._support.client import Gateway, Scenario, gateway_from_environment +from integration._support.wire import Reply, Request, Wire, wire_server + +FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"}) +IMAGE_RESPONSE: Final = { + "created": 1, + "background": None, + "data": [ + { + "b64_json": None, + "revised_prompt": None, + "url": "https://images.example/generated.png", + "provider_specific_fields": None, + } + ], + "output_format": None, + "quality": None, + "size": None, + "usage": { + "total_tokens": 0, + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + }, +} + + +def _primary_unavailable(_: Request) -> Reply: + return Reply( + status=503, + body=json.dumps({"error": {"message": "primary unavailable", "type": "server_error"}}).encode(), + ) + + +def _image_reply(_: Request) -> Reply: + return Reply( + body=json.dumps( + { + "created": 1, + "data": [{"url": "https://images.example/generated.png"}], + } + ).encode() + ) + + +def _image_model(scenario: Scenario, group: str, wire: Wire) -> str: + return scenario.model( + name=group, + model="openai/dall-e-3", + api_base=f"{wire.url}/v1", + api_key="integration-provider-key", + num_retries=0, + ) + + +def _assert_image_fallback( + gateway: Gateway, + key: str, + primary: Wire, + fallback: Wire, + primary_model: str, +) -> None: + response: Final = gateway.request( + "POST", + "/v1/images/generations", + {"model": primary_model, "prompt": "draw a tree"}, + key=key, + headers=FRESH_CONNECTION, + ) + + assert response.status_code == 200, response.text + assert response.json() == IMAGE_RESPONSE, response.text + assert tuple((request.method, request.target) for request in primary.drain()) == ( + ("POST", "/v1/images/generations"), + ), response.text + assert tuple((request.method, request.target) for request in fallback.drain()) == ( + ("POST", "/v1/images/generations"), + ), response.text + + +def test_key_router_settings_fallback_routes_image_generation() -> None: + with ( + wire_server(_primary_unavailable) as primary_wire, + wire_server(_image_reply) as fallback_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + fallback_model: Final = _image_model(scenario, f"image-fallback-{uuid.uuid4().hex}", fallback_wire) + key: Final = scenario.key(router_settings={"fallbacks": [{primary_model: [fallback_model]}], "num_retries": 0}) + + _assert_image_fallback(gateway, key, primary_wire, fallback_wire, primary_model) + + +def test_team_router_settings_fallback_routes_image_generation() -> None: + with ( + wire_server(_primary_unavailable) as primary_wire, + wire_server(_image_reply) as fallback_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + fallback_model: Final = _image_model(scenario, f"image-fallback-{uuid.uuid4().hex}", fallback_wire) + team: Final = scenario.team( + router_settings={"fallbacks": [{primary_model: [fallback_model]}], "num_retries": 0} + ) + key: Final = scenario.key(team_id=team) + + _assert_image_fallback(gateway, key, primary_wire, fallback_wire, primary_model) + + +def test_key_model_group_alias_resolves_once_for_image_generation() -> None: + with ( + wire_server(_image_reply) as primary_wire, + wire_server(_image_reply) as target_wire, + wire_server(_image_reply) as chained_target_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + target_model: Final = _image_model(scenario, f"image-target-{uuid.uuid4().hex}", target_wire) + chained_target_model: Final = _image_model( + scenario, f"image-chained-target-{uuid.uuid4().hex}", chained_target_wire + ) + alias_name: Final = f"image-alias-{uuid.uuid4().hex}" + key: Final = scenario.key( + router_settings={ + "model_group_alias": { + alias_name: target_model, + target_model: chained_target_model, + } + } + ) + + response: Final = gateway.request( + "POST", + "/v1/images/generations", + {"model": alias_name, "prompt": "draw a tree"}, + key=key, + headers=FRESH_CONNECTION, + ) + primary_requests: Final = tuple((request.method, request.target) for request in primary_wire.drain()) + target_requests: Final = tuple((request.method, request.target) for request in target_wire.drain()) + chained_target_requests: Final = tuple( + (request.method, request.target) for request in chained_target_wire.drain() + ) + + assert alias_name not in {primary_model, target_model, chained_target_model}, response.text + assert response.status_code == 200, response.text + assert response.json() == IMAGE_RESPONSE, response.text + assert primary_requests == (), response.text + assert target_requests == (("POST", "/v1/images/generations"),), ( + f"{response.text}; chained target requests: {chained_target_requests}" + ) + assert chained_target_requests == (), response.text diff --git a/tests/integration/routing/test_observed_routing.py b/tests/integration/routing/test_observed_routing.py index ca5f248e0de..64ab799fe3a 100644 --- a/tests/integration/routing/test_observed_routing.py +++ b/tests/integration/routing/test_observed_routing.py @@ -593,7 +593,7 @@ def test_health_routing_fallbacks_preserve_error_specific_targets_and_request_co "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}], }, response_text assert {key: events[1]["error"][key] for key in events[1]["error"] if key != "message"} == { - "type": None, + "type": "internal_server_error", "param": None, "code": "500", }, response_text diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index 87b56a92979..7b0e14d044f 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -518,9 +518,7 @@ def test_a_tpm_limit_lowered_while_the_socket_is_open_applies_to_its_next_turn( assert deployment.forwarded_turns() == 1 -def test_a_socket_on_an_otpm_limited_deployment_reserves_once_and_serves_every_turn( - limited: Gateway, cert: tuple[Path, Path] -) -> None: +def test_a_socket_on_an_otpm_limited_deployment_serves_every_turn(limited: Gateway, cert: tuple[Path, Path]) -> None: with responses_peer(cert) as peer, limited.scenario() as scenario: deployment: Final = _native(scenario, peer, otpm=10_000) key: Final = scenario.key(models=[deployment.name]) @@ -536,6 +534,115 @@ def test_a_socket_on_an_otpm_limited_deployment_reserves_once_and_serves_every_t assert deployment.forwarded_turns() == 5 +IO_TOKEN_LIMIT: Final = 10_000 +TOKENS_PER_TURN: Final = MappingProxyType({"itpm": USAGE["input_tokens"], "otpm": USAGE["output_tokens"]}) + + +def _io_tokens_counted(identity: str, kind: str, minute: str) -> int: + with _redis() as redis: + return sum(int(redis.get(key) or 0) for key in redis.scan_iter(f"global_router:{identity}:*:{kind}:{minute}")) + + +def _add_io_tokens(deployment: Deployment, kind: str, minute: str, tokens: int) -> None: + with _redis() as redis: + redis.incrby(f"global_router:{deployment.identity}:openai/{PROVIDER_MODEL}:{kind}:{minute}", tokens) + + +async def _say_until_counted( + connection: websockets.ClientConnection, deployment: Deployment, counted: Callable[[], float], expected: int +) -> tuple[dict[str, JsonValue], ...]: + received: Final = await _say(connection, deployment.name) + await _until(counted, lambda tokens: tokens == expected) + return received + + +@pytest.mark.parametrize("kind", ("otpm", "itpm")) +def test_a_turn_is_refused_once_earlier_turns_fill_the_deployment_io_token_limit( + limited: Gateway, cert: tuple[Path, Path], kind: str +) -> None: + per_turn: Final = TOKENS_PER_TURN[kind] + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, **{kind: IO_TOKEN_LIMIT}) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + counted: Final = partial(_io_tokens_counted, deployment.identity, kind, minute) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + first: Final = await _say_until_counted(connection, deployment, counted, per_turn) + await asyncio.to_thread(_add_io_tokens, deployment, kind, minute, IO_TOKEN_LIMIT - 5 * per_turn) + answered: Final = tuple( + [ + await _say_until_counted( + connection, deployment, counted, IO_TOKEN_LIMIT - (5 - number) * per_turn + ) + for number in range(2, 6) + ] + ) + return (first, *answered, await _say(connection, deployment.name)) + + turns: Final = asyncio.run(session()) + + assert [completed(received)["status"] for received in turns[:5]] == ["completed"] * 5 + _assert_refused(turns[5]) + assert deployment.forwarded_turns() == 5 + assert held(counted, lambda tokens: tokens == IO_TOKEN_LIMIT, holding=3, seconds=20) == IO_TOKEN_LIMIT + + +def test_a_socket_at_half_its_deployment_otpm_serves_ten_turns(limited: Gateway, cert: tuple[Path, Path]) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + deployment: Final = _native(scenario, peer, otpm=IO_TOKEN_LIMIT) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + _add_io_tokens(deployment, "otpm", minute, IO_TOKEN_LIMIT // 2) + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + return tuple([await _say(connection, deployment.name) for _ in range(10)]) + + turns: Final = asyncio.run(session()) + + assert [completed(received)["status"] for received in turns] == ["completed"] * 10 + assert deployment.forwarded_turns() == 10 + expected: Final = IO_TOKEN_LIMIT // 2 + 10 * USAGE["output_tokens"] + counted: Final = partial(_io_tokens_counted, deployment.identity, "otpm", minute) + assert held(counted, lambda tokens: tokens == expected, holding=3, seconds=20) == expected + + +@pytest.mark.parametrize("relay", ("native", "managed")) +def test_every_turn_counts_its_actual_tokens_toward_the_deployment_itpm_and_otpm( + limited: Gateway, cert: tuple[Path, Path], relay: str +) -> None: + with responses_peer(cert) as peer, limited.scenario() as scenario: + limits: Final = {"itpm": 1_000_000, "otpm": 1_000_000} + deployment: Final = _native(scenario, peer, **limits) if relay == "native" else _managed(scenario, **limits) + key: Final = scenario.key(models=[deployment.name]) + minute: Final = _fresh_minute() + + async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]: + async with _socket(limited, key, deployment.name) as connection: + return tuple([await _say(connection, deployment.name) for _ in range(3)]) + + turns: Final = asyncio.run(session()) + + usages: Final = tuple(object_value(completed(received)["usage"]) for received in turns) + expected: Final = { + "itpm": sum(int(str(usage["input_tokens"])) for usage in usages), + "otpm": sum(int(str(usage["output_tokens"])) for usage in usages), + } + assert any(tokens > 0 for tokens in expected.values()), usages + counted: Final = { + kind: held( + partial(_io_tokens_counted, deployment.identity, kind, minute), + lambda tokens, kind=kind: tokens == expected[kind], + holding=3, + seconds=20, + ) + for kind in expected + } + assert counted == expected + + def test_usage_based_routing_refuses_a_turn_only_once_earlier_turns_used_up_the_deployment_tpm( usage_routed: Gateway, cert: tuple[Path, Path] ) -> None: diff --git a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py index 0a5bfd339fc..955d9008a9c 100644 --- a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py +++ b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py @@ -143,6 +143,6 @@ async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(fallb model="primary", messages=[{"role": "user", "content": f"cooldown {uuid.uuid4().hex}"}], ) - cooled: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + cooled: Final = await async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) assert tuple(cooled) == ("primary-deployment",) diff --git a/tests/integration/security/test_config_overrides.py b/tests/integration/security/test_config_overrides.py index 70f9226ea79..6b8e8791b93 100644 --- a/tests/integration/security/test_config_overrides.py +++ b/tests/integration/security/test_config_overrides.py @@ -715,8 +715,11 @@ def test_cyberark_incremental_update_encrypts_reloads_on_restart_and_delete_clea }, saved.text stored_full: Final = _stored_override(database_url, "cyberark") assert set(stored_full) == set(full_config), stored_full + assert all(stored_full[key] != value for key, value in full_config.items()), stored_full serialized_full: Final = json.dumps(stored_full) - assert all(value not in serialized_full for value in full_config.values()), serialized_full + assert all(value not in serialized_full for value in (cyberark_wire.url, account, username, api_key)), ( + serialized_full + ) model_one: Final = "integration-cyberark-" + uuid.uuid4().hex _new_model(first, model_one, provider_wire.url, secret_names[0]) diff --git a/tests/integration/security/test_hide_secrets_tool_compatibility.py b/tests/integration/security/test_hide_secrets_tool_compatibility.py index d8b900de6d0..a89678aef8e 100644 --- a/tests/integration/security/test_hide_secrets_tool_compatibility.py +++ b/tests/integration/security/test_hide_secrets_tool_compatibility.py @@ -140,9 +140,15 @@ def _key(rig: Rig, *, permissions: dict[str, JsonValue] | None = None) -> Iterat ) +def _without_cooldowns(configuration: dict[str, JsonValue]) -> None: + settings: Final = configuration["router_settings"] + assert isinstance(settings, dict) + settings["disable_cooldowns"] = True + + @pytest.fixture(scope="module") def caller(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Caller]: - with tool_rig(tmp_path_factory.mktemp("hide-secrets-compat"), _peer) as rig: + with tool_rig(tmp_path_factory.mktemp("hide-secrets-compat"), _peer, configure=_without_cooldowns) as rig: with _key(rig) as key: yield Caller(rig, key) diff --git a/tests/integration/spend/_partial_stream.py b/tests/integration/spend/_partial_stream.py new file mode 100644 index 00000000000..e35320a8d8e --- /dev/null +++ b/tests/integration/spend/_partial_stream.py @@ -0,0 +1,173 @@ +import json +import threading +import uuid +from collections.abc import Callable, Iterable, Iterator +from hashlib import sha256 +from typing import Final + +import httpx +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.wire import Reply, Request, Wire +from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +BACKEND: Final = "integration-partial-stream" +PIECE: Final = "partial stream billed tokens " +DELIVERED_PIECES: Final = 12 +_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"' +ROWS: Final = 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"' + + +class _SpendRow(BaseModel): + model_config = ConfigDict(frozen=True) + + spend: float + + +class StatusRow(BaseModel): + model_config = ConfigDict(frozen=True) + + status: str + spend: float + + +class _BudgetError(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str + + +class _ErrorResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + error: _BudgetError + + +def _chunk(stream_id: str, delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes: + body: Final = { + "id": stream_id, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": BACKEND, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(body)}\n\n".encode() + + +def completion() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": BACKEND, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ).encode() + + +def peer(delivered: int | None, gate: threading.Event | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if JSON_OBJECT.validate_json(request.body).get("stream") is True: + stream_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + *(_chunk(stream_id, {"role": "assistant", "content": PIECE}) for _ in range(DELIVERED_PIECES)), + _chunk(stream_id, {}, "stop"), + b"data: [DONE]\n\n", + ), + abort_after=delivered, + pause_between_chunks=0.02, + gate_before_first=gate, + ) + return Reply(body=completion()) + + return respond + + +def priced_model(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"openai/{BACKEND}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + + +def stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "stream": True, + "messages": [{"role": "user", "content": "partial stream"}], + **fields, + }, + key=key, + ) + + +def broken_stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> None: + response: Final = stream(gateway, model, key, **fields) + assert response.status_code == 200, response.text + events: Final = tuple( + line.removeprefix("data: ") + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert "error" in JSON_OBJECT.validate_json(events[-1]), response.text + delivered: Final = "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1])) + assert delivered == PIECE * DELIVERED_PIECES, response.text + + +def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]: + for chunk in chunks: + yield from (choice.delta.content or "" for choice in chunk.choices) + + +def chat(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}], + **fields, + }, + key=key, + ) + + +def failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70) + found: Final = held( + lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2 + ) + return tuple(_SpendRow.model_validate(row).spend for row in found) + + +def chats_sent(wire: Wire) -> int: + return sum(1 for request in wire.drain() if request.method == "POST") + + +def recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float: + key: Final = scenario.key(models=[model], max_budget=100) + broken_stream(gateway, model, key) + (spend,) = failure_spend(gateway, key) + assert spend > 0 + return spend + + +def assert_budget_refusal(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py index 58fcd4e9de2..fca848fbb43 100644 --- a/tests/integration/spend/test_agent_spend.py +++ b/tests/integration/spend/test_agent_spend.py @@ -269,9 +269,4 @@ def test_a2a_send_bills_input_and_output_tokens_at_the_agent_per_token_prices(ga def test_a2a_stream_bills_input_and_output_tokens_at_the_agent_per_token_prices(gateway: Gateway) -> None: - pytest.skip( - "BUG: message/stream to an agent priced with input_cost_per_token/output_cost_per_token logs spend 0.0 " - "with prompt_tokens 0 and completion_tokens 0, because A2AStreamingIterator reads text from the request " - "Part RootModels without dumping them and ignores artifact-update text, while message/send bills 1.125" - ) _assert_a2a_calls_are_billed(gateway, _PER_TOKEN, _PER_TOKEN_COST, (("message/stream", "/a2a/{agent}"),)) diff --git a/tests/integration/spend/test_failure_spend_live_counters_chaos.py b/tests/integration/spend/test_failure_spend_live_counters_chaos.py new file mode 100644 index 00000000000..ecdd5aadb8e --- /dev/null +++ b/tests/integration/spend/test_failure_spend_live_counters_chaos.py @@ -0,0 +1,239 @@ +import json +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.timing import assert_not_delayed_by, timed +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.spend._partial_stream import ( + DELIVERED_PIECES, + assert_budget_refusal, + broken_stream, + chat, + chats_sent, + failure_spend, + peer, + priced_model, + recovered_cost, +) +from pydantic import JsonValue +from redis import Redis + +from litellm.constants import REDIS_SOCKET_TIMEOUT + +_BURST: Final = 12 +_SPREAD: Final = 6 +_ATTACK_MARKER: Final = "synthetic-attack-marker" +_POLICY_ID: Final = "synthetic-priced-policy" +_POLICY_APPLY_PATH: Final = f"/guardrail/{_POLICY_ID}/version/DRAFT/apply" + + +def _proxy_config(directory: Path, name: str, sections: dict[str, JsonValue]) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = directory / f"{name}.yaml" + path.write_text(yaml.safe_dump({**base, **sections})) + return path + + +def _redis_counter_config(directory: Path) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + return _proxy_config( + directory, + "redis-counters", + { + "router_settings": { + **base["router_settings"], + "redis_host": "os.environ/REDIS_HOST", + "redis_port": "os.environ/REDIS_PORT", + } + }, + ) + + +def _spread_requests(gateway: Gateway, model: str, key: str) -> tuple[httpx.Response, ...]: + return tuple( + gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"spread {uuid.uuid4().hex}"}]}, + key=key, + headers={"connection": "close"}, + ) + for _ in range(_SPREAD) + ) + + +@pytest.mark.timeout(600) +def test_broken_streams_through_a_redis_stall_log_once_and_the_budget_holds_on_every_worker( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + with ( + owned_redis(tmp_path) as counters, + wire_server(peer(DELIVERED_PIECES)) as wire, + wire_server(peer(DELIVERED_PIECES, gate=release)) as held_wire, + owned_proxy_process( + gateway, + tmp_path, + {"REDIS_HOST": counters.host, "REDIS_PORT": str(counters.port)}, + config=_redis_counter_config(tmp_path), + workers=2, + ) as owned, + owned.gateway.scenario() as scenario, + ): + scenario.cleanups.callback(release.set) + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + stalled_model: Final = priced_model(scenario, held_wire) + cost: Final = recovered_cost(candidate, scenario, model) + control: Final = timed(lambda: broken_stream(candidate, model, scenario.key(models=[model]))) + burst_key: Final = scenario.key(models=[stalled_model]) + chats_sent(held_wire) + + with ThreadPoolExecutor(max_workers=_BURST) as pool: + pending: Final = tuple( + pool.submit(broken_stream, candidate, stalled_model, burst_key) for _ in range(_BURST) + ) + eventually(held_wire.received.qsize, lambda held: held >= _BURST, seconds=30) + with counters.paused(): + release.set() + drained: Final = timed(lambda: tuple(future.result() for future in pending)) + + assert_not_delayed_by((drained.seconds,), control=(control.seconds,), injected_delay=REDIS_SOCKET_TIMEOUT) + assert failure_spend(candidate, burst_key, rows=_BURST) == pytest.approx((cost,) * _BURST) + spent: Final = cost * _BURST + with Redis(host=counters.host, port=counters.port) as live_counters: + landed: Final = eventually( + lambda: float(live_counters.get(f"spend:key:{sha256(burst_key.encode()).hexdigest()}") or 0), + lambda counted: counted >= spent - 1e-9, + seconds=30, + return_last_on_timeout=True, + ) + assert spent - 1e-9 <= landed <= 2 * spent + 1e-9, ( + f"one broken stream costs {cost}. The cold counter is seeded from the database, and a spend flush" + " during the stall can land these rows there first, so up to twice their cost is expected, as for" + " completed streams held past a flush" + ) + + budgeted: Final = scenario.key(models=[model], max_budget=cost / 2) + broken_stream(candidate, model, budgeted) + before_kill: Final = _spread_requests(candidate, model, budgeted) + owned.kill_worker(min(owned.worker_pids())) + after_kill: Final = _spread_requests(candidate, model, budgeted) + + for refused in (*before_kill, *after_kill): + assert_budget_refusal(refused) + + +def _policy(request: Request) -> Reply: + if request.method != "POST" or request.target != _POLICY_APPLY_PATH: + return Reply(status=404) + if _ATTACK_MARKER.encode() not in request.body: + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "blocked by synthetic policy"}], + "assessments": [ + {"contentPolicy": {"filters": [{"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}]}} + ], + "usage": {"contentPolicyUnits": 1}, + } + ).encode() + ) + + +def _policy_config(directory: Path, policy: Wire) -> Path: + return _proxy_config( + directory, + "priced-policy", + { + "guardrails": [ + { + "guardrail_name": "priced-bedrock-policy", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "default_on": True, + "guardrailIdentifier": _POLICY_ID, + "guardrailVersion": "DRAFT", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIASYNTHETICGUARDRAIL", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_bedrock_runtime_endpoint": policy.url, + }, + } + ] + }, + ) + + +def _attack(gateway: Gateway, model: str, key: str) -> httpx.Response: + blocked: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"{_ATTACK_MARKER} {uuid.uuid4().hex}"}]}, + key=key, + ) + assert blocked.status_code == 400, blocked.text + return blocked + + +@pytest.mark.timeout(480) +def test_guardrail_cost_on_blocked_requests_counts_toward_the_budget(gateway: Gateway, tmp_path: Path) -> None: + with ( + wire_server(_policy) as policy, + wire_server(peer(None)) as wire, + owned_proxy_process(gateway, tmp_path, {}, config=_policy_config(tmp_path, policy), workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + meter: Final = scenario.key(models=[model], max_budget=100) + _attack(candidate, model, meter) + (cost,) = failure_spend(candidate, meter) + assert cost > 0 + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + chats_sent(wire) + + _attack(candidate, model, key) + _attack(candidate, model, key) + after: Final = chat(candidate, model, key) + + assert_budget_refusal(after) + assert chats_sent(wire) == 0 + + +@pytest.mark.timeout(480) +def test_broken_streams_leave_the_budget_alone_while_spend_updates_are_disabled( + gateway: Gateway, tmp_path: Path +) -> None: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = _proxy_config( + tmp_path, "spend-updates-off", {"general_settings": {**base["general_settings"], "disable_spend_updates": True}} + ) + with ( + wire_server(peer(DELIVERED_PIECES)) as wire, + gateway.scenario() as metering, + owned_proxy_process(gateway, tmp_path, {}, config=config) as owned, + owned.gateway.scenario() as scenario, + ): + cost: Final = recovered_cost(gateway, metering, priced_model(metering, wire)) + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + + broken_stream(candidate, model, key) + after: Final = chat(candidate, model, key) + + assert after.status_code == 200, after.text diff --git a/tests/integration/spend/test_non_finite_cost_spend_writes.py b/tests/integration/spend/test_non_finite_cost_spend_writes.py new file mode 100644 index 00000000000..9a4d2dd45a4 --- /dev/null +++ b/tests/integration/spend/test_non_finite_cost_spend_writes.py @@ -0,0 +1,166 @@ +import json +from hashlib import sha256 +from pathlib import Path +from typing import Final, Literal +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +PROMPT_TOKENS: Final = 5 +COMPLETION_TOKENS: Final = 3 +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +TOKEN_PRICE: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE +OVERFLOWING_INPUT_RATE: Final = 1e308 +POISONED_UPSTREAM: Final = "poisoned-upstream" +CONTROL_UPSTREAM: Final = "control-upstream" +DAILY_TOTALS: Final = "SELECT model_group, SUM(spend)::float8 AS spend, SUM(api_requests)::int8 AS api_requests" +DAILY_USER_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyUserSpend" WHERE api_key=%s GROUP BY model_group' +DAILY_TEAM_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyTeamSpend" WHERE api_key=%s GROUP BY model_group' +DAILY_TAG_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s AND tag=%s GROUP BY model_group' + +Poison = Literal["reported_nan_cost", "overflowing_price"] + + +def _frame(payload: dict[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _stream(upstream_model: str, reported_cost: float | None) -> tuple[bytes, ...]: + chunk: Final = {"id": f"chatcmpl-{uuid4().hex}", "object": "chat.completion.chunk", "created": 1} + usage: Final[dict[str, JsonValue]] = { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + } + reported_usage: Final = usage if reported_cost is None else {**usage, "cost": reported_cost} + return ( + _frame({**chunk, "model": upstream_model, "choices": [{"index": 0, "delta": {"content": "ok"}}]}), + _frame( + { + **chunk, + "model": upstream_model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": reported_usage, + } + ), + b"data: [DONE]\n\n", + ) + + +def _deployment(alias: str, upstream_model: str, wire_url: str, input_rate: float) -> dict[str, JsonValue]: + return { + "model_name": alias, + "litellm_params": { + "model": f"openai/{upstream_model}", + "api_base": f"{wire_url}/v1", + "api_key": "integration-provider-key", + "input_cost_per_token": input_rate, + "output_cost_per_token": OUTPUT_RATE, + }, + } + + +def _proxy_config(directory: Path, deployments: tuple[dict[str, JsonValue], ...]) -> Path: + config: Final = directory / "non_finite_cost_config.yaml" + config.write_text( + json.dumps( + { + "model_list": list(deployments), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_spend_logs": False, + "proxy_batch_write_at": 1, + "proxy_batch_polling_interval": 1, + }, + "router_settings": {"disable_cooldowns": True}, + } + ) + ) + return config + + +def _by_model_group(rows: list[dict[str, JsonValue]]) -> dict[str, float]: + return {str(row["model_group"]): float(str(row["spend"])) for row in rows} + + +def _both_requests_rows(query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]: + return eventually(lambda: read_rows(query, parameters), lambda rows: len(rows) == 2, seconds=70) + + +@pytest.mark.parametrize("poison", ["reported_nan_cost", "overflowing_price"]) +@pytest.mark.timeout(240) +def test_a_non_finite_cost_never_blocks_the_next_requests_spend_rows( + gateway: Gateway, tmp_path: Path, poison: Poison +) -> None: + poisoned_alias: Final = f"non-finite-cost-{uuid4().hex}" + control_alias: Final = f"finite-cost-{uuid4().hex}" + tag: Final = f"non-finite-cost-{uuid4().hex}" + poisoned_spend: Final = TOKEN_PRICE if poison == "reported_nan_cost" else 0.0 + poisoned_input_rate: Final = INPUT_RATE if poison == "reported_nan_cost" else OVERFLOWING_INPUT_RATE + + def respond(request: Request) -> Reply: + assert request.target.endswith("/chat/completions"), request.target + upstream_model: Final = json.loads(request.body)["model"] + reported_cost: Final = ( + float("nan") if upstream_model == POISONED_UPSTREAM and poison == "reported_nan_cost" else None + ) + return Reply(content_type="text/event-stream", chunks=_stream(upstream_model, reported_cost)) + + with wire_server(answering_model_discovery(respond)) as wire: + config: Final = _proxy_config( + tmp_path, + ( + _deployment(poisoned_alias, POISONED_UPSTREAM, wire.url, poisoned_input_rate), + _deployment(control_alias, CONTROL_UPSTREAM, wire.url, INPUT_RATE), + ), + ) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + team: Final = scenario.team() + key: Final = scenario.key(user_id=f"non-finite-cost-{uuid4().hex}", team_id=team) + digest: Final = sha256(key.encode()).hexdigest() + responses: Final = tuple( + candidate.request( + "POST", + "/v1/chat/completions", + { + "model": alias, + "messages": [{"role": "user", "content": f"non-finite cost {uuid4().hex}"}], + "stream": True, + "metadata": {"tags": [tag]}, + }, + key=key, + ) + for alias in (poisoned_alias, control_alias) + ) + assert all(response.status_code == 200 for response in responses), [response.text for response in responses] + + expected: Final = {poisoned_alias: pytest.approx(poisoned_spend), control_alias: pytest.approx(TOKEN_PRICE)} + spend_logs: Final = _both_requests_rows( + 'SELECT model_group, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,) + ) + assert _by_model_group(spend_logs) == expected, spend_logs + daily_rows: Final = ( + _both_requests_rows(DAILY_USER_QUERY, (digest,)), + _both_requests_rows(DAILY_TEAM_QUERY, (digest,)), + _both_requests_rows(DAILY_TAG_QUERY, (digest, tag)), + ) + assert [_by_model_group(rows) for rows in daily_rows] == [expected] * 3, daily_rows + assert [{row["api_requests"] for row in rows} for rows in daily_rows] == [{1}] * 3, daily_rows + key_total: Final = poisoned_spend + TOKEN_PRICE + key_spend: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(key_total), + seconds=70, + return_last_on_timeout=True, + ) + assert float(str(key_spend[0]["spend"])) == pytest.approx(key_total), key_spend diff --git a/tests/integration/spend/test_openrouter_upstream_provider.py b/tests/integration/spend/test_openrouter_upstream_provider.py new file mode 100644 index 00000000000..fa17a364aea --- /dev/null +++ b/tests/integration/spend/test_openrouter_upstream_provider.py @@ -0,0 +1,216 @@ +import json +import uuid +from collections.abc import Callable +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +_BACKEND: Final = "openai/gpt-4o-mini" +_PROMPT: Final = "upstream provider trace" +_PROVIDER: Final = "serving-site-a" +_ANSWER: Final = "hello" +_USAGE: Final[dict[str, JsonValue]] = { + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + "cost": 0.00015, +} + + +def _reply( + request_id: str, + *, + response_provider: bool, + stream_provider_mode: Literal["all", "last", "none"], +) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target.split("?", 1)[0].endswith("/chat/completions"), request.target + body: Final = JSON_OBJECT.validate_json(request.body) + assert body["messages"] == [{"role": "user", "content": _PROMPT}], body + if body.get("stream") is True: + common: Final = { + "id": request_id, + "object": "chat.completion.chunk", + "created": 1, + "model": _BACKEND, + } + all_provider: Final[dict[str, JsonValue]] = {"provider": _PROVIDER} + first_provider: Final = all_provider if stream_provider_mode == "all" else {} + finish_provider: Final = ( + all_provider if stream_provider_mode in ("all", "last") else {} + ) + usage_provider: Final = all_provider if stream_provider_mode == "all" else {} + chunks: Final[tuple[dict[str, JsonValue], ...]] = ( + { + **common, + **first_provider, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": _ANSWER}, "finish_reason": None}], + }, + { + **common, + **finish_provider, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + {**common, **usage_provider, "choices": [], "usage": _USAGE}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(chunk)}\n\n".encode() for chunk in chunks) + + (b"data: [DONE]\n\n",), + ) + + provider: Final[dict[str, JsonValue]] = {"provider": _PROVIDER} if response_provider else {} + return Reply( + body=json.dumps( + { + "id": request_id, + "object": "chat.completion", + "created": 1, + "model": _BACKEND, + **provider, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": _ANSWER}, + "finish_reason": "stop", + } + ], + "usage": _USAGE, + } + ).encode() + ) + + return respond + + +def _stream_chunks(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _send( + gateway: Gateway, + scenario: Scenario, + wire_url: str, + request_id: str, + *, + openrouter: bool, + stream: bool, +) -> httpx.Response: + model: Final = scenario.model( + model=f"{'openrouter' if openrouter else 'openai'}/{_BACKEND}", + api_base=wire_url, + api_key="synthetic-provider-key", + ) + key: Final = scenario.key(models=[model]) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": stream, + **({"stream_options": {"include_usage": True}} if stream else {}), + }, + key=key, + ) + assert response.status_code == 200, response.text + returned_id: Final = ( + string_value(_stream_chunks(response)[0]["id"]) + if stream + else string_value(JSON_OBJECT.validate_json(response.content)["id"]) + ) + assert returned_id == request_id, response.text + return response + + +def _spend_metadata(request_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return object_value(rows[0]["metadata"]) + + +@pytest.mark.parametrize( + ("stream", "stream_provider_mode"), + [ + pytest.param(False, "none", id="nonstream"), + pytest.param(True, "all", id="stream-every-chunk"), + pytest.param(True, "last", id="stream-later-chunk"), + ], +) +def test_openrouter_provider_is_persisted_in_spend_metadata_regression( + gateway: Gateway, stream: bool, stream_provider_mode: Literal["all", "last", "none"] +) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=True, + stream_provider_mode=stream_provider_mode, + ) + ) as wire, + gateway.scenario() as scenario, + ): + _send(gateway, scenario, wire.url, request_id, openrouter=True, stream=stream) + metadata: Final = _spend_metadata(request_id) + + assert metadata.get("upstream_provider") == _PROVIDER, metadata + + +def test_openrouter_client_response_keeps_provider_in_body_guard(gateway: Gateway) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=True, + stream_provider_mode="none", + ) + ) as wire, + gateway.scenario() as scenario, + ): + response: Final = _send(gateway, scenario, wire.url, request_id, openrouter=True, stream=False) + + body: Final = JSON_OBJECT.validate_json(response.content) + choices: Final = body["choices"] + assert isinstance(choices, list) + message: Final = object_value(object_value(choices[0])["message"]) + assert body["provider"] == _PROVIDER + assert "upstream_provider" not in body + assert message["content"] == _ANSWER + + +def test_non_openrouter_spend_metadata_has_no_provider_guard(gateway: Gateway) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=False, + stream_provider_mode="none", + ) + ) as wire, + gateway.scenario() as scenario, + ): + _send(gateway, scenario, wire.url, request_id, openrouter=False, stream=False) + metadata: Final = _spend_metadata(request_id) + + assert "upstream_provider" not in metadata, metadata diff --git a/tests/integration/spend/test_partial_stream_live_budget.py b/tests/integration/spend/test_partial_stream_live_budget.py new file mode 100644 index 00000000000..89fb66c9bbf --- /dev/null +++ b/tests/integration/spend/test_partial_stream_live_budget.py @@ -0,0 +1,350 @@ +import json +import threading +import uuid +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, held +from integration._support.wire import Reply, Request, wire_server +from integration.spend._partial_stream import ( + BACKEND, + DELIVERED_PIECES, + JSON_OBJECT, + PIECE, + ROWS, + StatusRow, + assert_budget_refusal, + broken_stream, + chat, + chats_sent, + completion, + failure_spend, + peer, + priced_model, + recovered_cost, + stream, +) +from pydantic import JsonValue + +_CONCURRENT_STREAMS: Final = 6 +_MESSAGES_BACKEND: Final = "claude-integration-partial-stream" +_UNCACHED: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}} + + +@pytest.mark.timeout(240) +def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + chats_sent(wire) + + broken_stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0)) + assert chats_sent(wire) == 1 + + +@pytest.mark.timeout(240) +def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + broken_stream(gateway, model, key) + broken_stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0)) + + +@pytest.mark.timeout(240) +def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None: + with wire_server(peer(0)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=0.05) + + refused: Final = stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert refused.status_code == 500, refused.text + + assert after.status_code == 200, after.text + assert failure_spend(gateway, key) == (0.0,) + + +def _broken_stream_then_refusal( + gateway: Gateway, model: str, spender: str, other: str, **fields: JsonValue +) -> httpx.Response: + broken_stream(gateway, model, spender, **fields) + return chat(gateway, model, other, **fields) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_on_one_team_key_blocks_the_teams_other_keys(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + team: Final = scenario.team(models=[model], max_budget=cost / 2) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model]) + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_on_one_of_a_users_keys_blocks_their_other_keys(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + user: Final = scenario.user(max_budget=cost / 2) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model]) + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_for_an_end_user_blocks_their_next_request_on_any_key(gateway: Gateway) -> None: + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + gateway.post("/customer/new", {"user_id": end_user, "max_budget": cost / 2}) + scenario.cleanups.callback(gateway.post, "/customer/delete", {"user_ids": [end_user]}) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(models=[model]), scenario.key(models=[model]), user=end_user + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_under_a_tag_blocks_the_next_request_with_that_tag(gateway: Gateway) -> None: + tag: Final = f"integration-partial-stream-{uuid.uuid4().hex}" + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + gateway.post("/tag/new", {"name": tag, "max_budget": cost / 2}) + scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": tag}) + key: Final = scenario.key(models=[model]) + + after: Final = _broken_stream_then_refusal(gateway, model, key, key, metadata={"tags": [tag]}) + untagged: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert untagged.status_code == 200, untagged.text + + +@pytest.mark.timeout(240) +def test_a_broken_stream_fills_the_keys_budget_window(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], budget_limits=[{"budget_duration": "1d", "max_budget": cost / 2}]) + + after: Final = _broken_stream_then_refusal(gateway, model, key, key) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_concurrent_broken_streams_each_count_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * (_CONCURRENT_STREAMS + 0.5)) + + with ThreadPoolExecutor(max_workers=_CONCURRENT_STREAMS) as pool: + tuple(pool.map(lambda _: broken_stream(gateway, model, key), range(_CONCURRENT_STREAMS))) + within_budget: Final = chat(gateway, model, key) + broken_stream(gateway, model, key) + over_budget: Final = chat(gateway, model, key) + + assert within_budget.status_code == 200, within_budget.text + assert_budget_refusal(over_budget) + assert sorted(failure_spend(gateway, key, rows=_CONCURRENT_STREAMS + 2)) == pytest.approx( + [0.0] + [cost] * (_CONCURRENT_STREAMS + 1) + ) + + +@pytest.mark.timeout(240) +def test_completed_streams_count_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(None)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + meter: Final = scenario.key(models=[model], max_budget=100) + assert stream(gateway, model, meter, **_UNCACHED).status_code == 200 + (metered,) = eventually( + lambda: gateway.rows(ROWS, (sha256(meter.encode()).hexdigest(),)), + lambda found: len(found) == 1, + seconds=70, + ) + cost: Final = StatusRow.model_validate(metered).spend + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + first: Final = stream(gateway, model, key, **_UNCACHED) + second: Final = stream(gateway, model, key, **_UNCACHED) + assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) + assert _rows(gateway, key, 2) == (StatusRow(status="success", spend=cost),) * 2 + third: Final = stream(gateway, model, key, **_UNCACHED) + + assert_budget_refusal(third) + + +def _hung_then_served(release: threading.Event, hung_model: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if JSON_OBJECT.validate_json(request.body).get("model") == hung_model: + return Reply(chunks=(completion(),), gate_before_first=release) + return Reply(body=completion()) + + return respond + + +def _timed_out_spend(gateway: Gateway, model: str, key: str) -> float: + timed_out: Final = chat(gateway, model, key, messages=[{"role": "user", "content": "hung"}]) + assert timed_out.status_code in (408, 500), timed_out.text + digest: Final = sha256(key.encode()).hexdigest() + landed: Final = eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= 3, seconds=70) + rows: Final = tuple( + StatusRow.model_validate(row) + for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == len(landed), holding=2) + ) + assert {row.status for row in rows} == {"failure"}, rows + return sum(row.spend for row in rows) + + +@pytest.mark.timeout(240) +def test_retries_the_provider_billed_before_a_timeout_count_toward_the_budget(gateway: Gateway) -> None: + hung_backend: Final = f"{BACKEND}-hung" + release: Final = threading.Event() + with wire_server(_hung_then_served(release, hung_backend)) as wire, gateway.scenario() as scenario: + scenario.cleanups.callback(release.set) + hung: Final = scenario.model( + model=f"openai/{hung_backend}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + timeout=1, + num_retries=2, + ) + served: Final = priced_model(scenario, wire) + billed: Final = _timed_out_spend(gateway, hung, scenario.key(models=[hung, served], max_budget=100)) + assert billed > 0 + key: Final = scenario.key(models=[hung, served], max_budget=billed / 2) + + assert _timed_out_spend(gateway, hung, key) == pytest.approx(billed) + after: Final = chat(gateway, served, key) + + assert_budget_refusal(after) + + +def _messages_event(kind: str, fields: dict[str, JsonValue]) -> bytes: + return f"event: {kind}\ndata: {json.dumps({'type': kind, **fields})}\n\n".encode() + + +def _messages_peer(request: Request) -> Reply: + if request.method != "POST" or request.target != "/v1/messages": + return Reply(status=404) + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": _MESSAGES_BACKEND, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + } + if JSON_OBJECT.validate_json(request.body).get("stream") is not True: + return Reply( + body=json.dumps( + {**message, "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn"} + ).encode() + ) + return Reply( + content_type="text/event-stream", + chunks=( + _messages_event("message_start", {"message": {**message, "content": []}}), + _messages_event("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}), + *( + _messages_event("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": PIECE}}) + for _ in range(DELIVERED_PIECES) + ), + _messages_event("content_block_stop", {"index": 0}), + _messages_event( + "message_delta", + {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 60}}, + ), + _messages_event("message_stop", {}), + ), + abort_after=DELIVERED_PIECES + 2, + pause_between_chunks=0.02, + ) + + +def _messages(gateway: Gateway, model: str, key: str, *, streamed: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 5, + "stream": streamed, + "messages": [{"role": "user", "content": f"partial stream {uuid.uuid4().hex}"}], + }, + key=key, + ) + + +def _broken_messages_stream(gateway: Gateway, model: str, key: str) -> None: + response: Final = _messages(gateway, model, key, streamed=True) + assert response.status_code == 200, response.text + assert response.text.count(PIECE) == DELIVERED_PIECES, response.text + assert "message_stop" not in response.text, response.text + + +def _rows(gateway: Gateway, key: str, count: int) -> tuple[StatusRow, ...]: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= count, seconds=70) + return tuple( + StatusRow.model_validate(row) + for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == count, holding=2) + ) + + +@pytest.mark.timeout(240) +def test_a_broken_messages_stream_counts_toward_the_budget(gateway: Gateway) -> None: + with wire_server(_messages_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MESSAGES_BACKEND}", + api_base=wire.url, + api_key="synthetic-anthropic-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + meter: Final = scenario.key(models=[model], max_budget=100) + _broken_messages_stream(gateway, model, meter) + (metered,) = _rows(gateway, meter, 1) + assert metered.spend > 0, metered + key: Final = scenario.key(models=[model], max_budget=metered.spend / 2) + + _broken_messages_stream(gateway, model, key) + after: Final = _messages(gateway, model, key, streamed=False) + + assert_budget_refusal(after) + assert _rows(gateway, key, 2) == (metered, StatusRow(status="failure", spend=0.0)) diff --git a/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py b/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py new file mode 100644 index 00000000000..91269e6557f --- /dev/null +++ b/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py @@ -0,0 +1,181 @@ +"""Two Vertex deployments share a base model but live in different projects with different +buckets. An unmanaged gs:// batch created on the second one must be polled through that +deployment (LIT-10326). Before the fix the cost poller matched the base model name in the +gs:// path to the first deployment in the list and asked the wrong project for the job.""" + +import json +import re +from collections.abc import Callable +from contextlib import ExitStack +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.process import owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +LOCATION: Final = "us-central1" +BASE_MODEL: Final = "gemini-2.5-flash" +TEAMS: Final = ("a", "b") +_JOB_PATH: Final = re.compile( + r"^/v1/projects/(?P[^/]+)/locations/(?P[^/]+)/batchPredictionJobs(?:/(?P\d+))?$" +) + + +def _project(team: str) -> str: + return f"scripted-project-{team}" + + +def _deployment(team: str) -> str: + return f"{BASE_MODEL}-team-{team}" + + +def _input_file(team: str) -> str: + return f"gs://scripted-bucket-{team}/litellm-vertex-files/publishers/google/models/{BASE_MODEL}/{uuid4().hex}/input.jsonl" + + +def _job_number(team: str) -> str: + return f"{TEAMS.index(team) + 1}{uuid4().int % 10**17:017d}" + + +def _vertex_peer(jobs: dict[str, tuple[str, str]]) -> Callable[[Request], Reply]: + """A Vertex double that remembers which project each job was created in. A retrieve + from any other project answers the 404 real Vertex gives (the job does not exist + there), a retrieve from the right one reports the job still running.""" + + def respond(request: Request) -> Reply: + path: Final = _JOB_PATH.match(request.target.split("?", 1)[0]) + assert path is not None, request.target + assert request.headers.get("authorization") == "Bearer scripted-token", request.headers + project, location, job = path.group("project"), path.group("location"), path.group("job") + if request.method == "POST": + team: Final = project.rsplit("-", 1)[-1] + uri: Final = re.search(r'"(gs://[^"]+)"', request.body.decode()) + assert uri is not None and uri.group(1).startswith(f"gs://scripted-bucket-{team}/"), (project, request.body) + created, created_uri = jobs.setdefault(project, (_job_number(team), uri.group(1))) + return Reply(body=json.dumps(_job(project, location, created, created_uri, "JOB_STATE_PENDING")).encode()) + assert request.method == "GET" and job is not None, (request.method, request.target) + known: Final = jobs.get(project) + if known is None or known[0] != job: + return Reply( + status=404, + body=json.dumps( + { + "error": { + "code": 404, + "message": f"The BatchPredictionJob does not exist: {job}", + "status": "NOT_FOUND", + } + } + ).encode(), + ) + return Reply(body=json.dumps(_job(project, location, job, known[1], "JOB_STATE_RUNNING")).encode()) + + return respond + + +def _job(project: str, location: str, job: str, input_uri: str, state: str) -> dict[str, JsonValue]: + return { + "name": f"projects/{project}/locations/{location}/batchPredictionJobs/{job}", + "displayName": "litellm-vertex-batch-scripted", + "model": f"publishers/google/models/{BASE_MODEL}", + "inputConfig": {"gcsSource": {"uris": [input_uri]}, "instancesFormat": "jsonl"}, + "outputConfig": {"predictionsFormat": "jsonl", "gcsDestination": {"outputUriPrefix": input_uri.rsplit("/", 1)[0]}}, + "state": state, + "createTime": "2026-10-30T20:00:00.000000Z", + "updateTime": "2026-10-30T20:00:00.000000Z", + } + + +def _owned_config(path: Path, gateway: Gateway, wire_url: str) -> Path: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + { + "model_name": _deployment(team), + "litellm_params": { + "model": f"vertex_ai/{BASE_MODEL}", + "api_base": wire_url, + "vertex_project": _project(team), + "vertex_location": LOCATION, + "vertex_credentials": service_account_json(_project(team), gateway.upstream_url.rstrip("/")), + "gcs_bucket_name": f"scripted-bucket-{team}", + }, + } + for team in TEAMS + ], + "general_settings": { + **object_value(config.get("general_settings") or {}), + "track_unmanaged_batch_cost": True, + }, + "litellm_settings": { + **object_value(config.get("litellm_settings") or {}), + "enable_loadbalancing_on_batch_endpoints": True, + "user_url_allowed_hosts": ["127.0.0.1"], + }, + } + ) + ) + return path + + +def _create_batch(gateway: Gateway, key: str, team: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": _input_file(team), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": _deployment(team), + }, + key=key, + ) + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _retrieve_targets(requests: tuple[Request, ...], job: str) -> tuple[str, ...]: + return tuple(request.target for request in requests if request.method == "GET" and request.target.endswith(f"/{job}")) + + +@pytest.mark.timeout(240) +def test_unmanaged_batch_on_the_second_deployment_sharing_a_base_model_is_polled_through_its_own_project( + gateway: Gateway, tmp_path: Path +) -> None: + jobs: Final[dict[str, tuple[str, str]]] = {} # mutable-ok: the wire double records which project created each job + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_vertex_peer(jobs))) + config: Final = _owned_config(tmp_path / "two-vertex-teams.yaml", gateway, wire.url) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + scenario: Final = stack.enter_context(owned.gateway.scenario()) + key: Final = scenario.key(models=[_deployment(team) for team in TEAMS]) + + batch_a: Final = _create_batch(owned.gateway, key, "a") + batch_b: Final = _create_batch(owned.gateway, key, "b") + assert (batch_a, batch_b) == (jobs[_project("a")][0], jobs[_project("b")][0]), (batch_a, batch_b, jobs) + + seen: Final[list[Request]] = [] # mutable-ok: drain empties the wire queue, so samples accumulate across retries + + def polled_so_far() -> tuple[Request, ...]: + seen.extend(wire.drain()) + return tuple(seen) + + polled: Final = eventually( + polled_so_far, + lambda requests: bool(_retrieve_targets(requests, batch_b)) and bool(_retrieve_targets(requests, batch_a)), + seconds=90, + ) + for team, batch in (("a", batch_a), ("b", batch_b)): + assert all( + target.startswith(f"/v1/projects/{_project(team)}/") for target in _retrieve_targets(polled, batch) + ), (team, _retrieve_targets(polled, batch)) + diff --git a/tests/test_litellm_rust/cache/support.py b/tests/test_litellm_rust/cache/support.py index a30f6b68cc6..f83e2a8b9da 100644 --- a/tests/test_litellm_rust/cache/support.py +++ b/tests/test_litellm_rust/cache/support.py @@ -7,7 +7,7 @@ from typing import Final, Literal from pydantic import BaseModel, TypeAdapter import litellm -from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.chat_completions.entrypoints import NATIVE_ACOMPLETION diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index d855d212fec..722cdbcbaca 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -10,6 +10,7 @@ from litellm.caching.caching import Cache, CacheMode from litellm.caching.caching_handler import ( _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request ) +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.proxy.hooks.model_max_budget_limiter import ( PROXY_VirtualKeyModelMaxBudgetLimiter, @@ -17,7 +18,6 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( ) from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache -from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout diff --git a/tests/test_litellm_rust/cache/v2/test_cache.py b/tests/test_litellm_rust/cache/v2/test_cache.py index 0ba0a17007f..1461e6eac61 100644 --- a/tests/test_litellm_rust/cache/v2/test_cache.py +++ b/tests/test_litellm_rust/cache/v2/test_cache.py @@ -11,7 +11,7 @@ from litellm._v2.cache import NativeBackend from litellm.caching.caching_handler import ( _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request ) -from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.types.utils import ModelResponse from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE, RESPONSES_TEXT diff --git a/tests/test_litellm_rust/support/served.py b/tests/test_litellm_rust/support/served.py index edf492ea7e6..0890077f745 100644 --- a/tests/test_litellm_rust/support/served.py +++ b/tests/test_litellm_rust/support/served.py @@ -1,6 +1,6 @@ from typing import Final -from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from tests.test_litellm_rust.support.recording_server import RecordingServer diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py index 7ced3f8bdff..134184c7c01 100644 --- a/tests/test_litellm_rust/test_fork_guard.py +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -66,7 +66,7 @@ _SDK_CONTRACT = textwrap.dedent( import litellm from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted - from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict + from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict class Handler(BaseHTTPRequestHandler): def do_POST(self): diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py index 38b896723f7..2eaa04a767e 100644 --- a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py @@ -148,6 +148,8 @@ async def test_handle_streaming_runs_the_agent_and_ignores_unrelated_keywords(re "kind": "artifact-update", "taskId": ANY, "artifact": {"artifactId": ANY, "parts": [{"kind": "text", "text": "wxo says hi"}]}, + "append": False, + "lastChunk": True, } assert chunks[3]["result"] == { "contextId": ANY, diff --git a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py index 6d685674984..0d4103ae979 100644 --- a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py +++ b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py @@ -5,7 +5,9 @@ async_success_handler (cross-thread pydantic mutation segfaults pydantic-core). """ import asyncio +import importlib.util import time +from collections.abc import AsyncIterator, Callable, Mapping from datetime import datetime from typing import Final from types import SimpleNamespace @@ -15,6 +17,7 @@ import pytest import litellm from litellm.a2a_protocol import streaming_iterator as a2a_streaming_iterator_module from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator +from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -176,7 +179,9 @@ async def test_stream_logs_its_first_chunk_time_measured_from_the_call_start(): iterator: Final = A2AStreamingIterator( stream=_stream(), - request=SimpleNamespace(params=SimpleNamespace(message={"role": "user", "parts": [{"kind": "text", "text": "hi"}]})), + request=SimpleNamespace( + params=SimpleNamespace(message={"role": "user", "parts": [{"kind": "text", "text": "hi"}]}) + ), logging_obj=logging_obj, agent_name="test-agent", ) @@ -189,3 +194,133 @@ async def test_stream_logs_its_first_chunk_time_measured_from_the_call_start(): assert first_chunk_at is not None assert recorder.logged == [(call_start, first_chunk_at)] + + +_REPLY: Final = ( + "Paris is the capital of France. It sits on the Seine, holds about two million people, " + "and has been the seat of the French government for most of the last thousand years." +) + + +def _artifact_event(text: str, append: bool | None, artifact_id: str = "a1") -> Mapping[str, object]: + result: Final = { + "kind": "artifact-update", + "taskId": "t1", + "contextId": "c1", + "artifact": {"artifactId": artifact_id, "parts": [{"kind": "text", "text": text}]}, + } + return {"jsonrpc": "2.0", "id": "r1", "result": result if append is None else {**result, "append": append}} + + +def _status_event(state: str, text: str | None = None) -> Mapping[str, object]: + message: Final = {"kind": "message", "role": "agent", "messageId": "m1", "parts": [{"kind": "text", "text": text}]} + status: Final = {"state": state} if text is None else {"state": state, "message": message} + return { + "jsonrpc": "2.0", + "id": "r1", + "result": { + "kind": "status-update", + "taskId": "t1", + "contextId": "c1", + "final": state == "completed", + "status": status, + }, + } + + +_FIRST_ARTIFACT: Final = "Paris is the capital of France and sits on the Seine." +_SECOND_ARTIFACT: Final = "It has been the seat of the French government for most of the last thousand years." + + +async def _snapshot_stream() -> AsyncIterator[Mapping[str, object]]: + for end in range(40, len(_REPLY) + 40, 40): + yield _artifact_event(_REPLY[:end], append=None) + + +async def _watsonx_stream() -> AsyncIterator[Mapping[str, object]]: + from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import ( + WatsonxOrchestrateTransformation, + ) + + async for event in WatsonxOrchestrateTransformation.fake_streaming_from_text(_REPLY, "r1", delay_ms=0): + yield event + + +async def _two_artifact_stream() -> AsyncIterator[Mapping[str, object]]: + yield _artifact_event(_FIRST_ARTIFACT, append=None, artifact_id="a1") + yield _artifact_event(_SECOND_ARTIFACT, append=None, artifact_id="a2") + yield _status_event("completed") + + +async def _status_message_stream() -> AsyncIterator[Mapping[str, object]]: + yield _status_event("working", "Looking that up") + yield _status_event("working", _REPLY) + yield _status_event("completed") + + +async def _artifact_without_id_stream() -> AsyncIterator[Mapping[str, object]]: + yield { + "jsonrpc": "2.0", + "id": "r1", + "result": { + "kind": "artifact-update", + "taskId": "t1", + "contextId": "c1", + "artifact": {"parts": [{"kind": "text", "text": _REPLY}]}, + }, + } + yield _status_event("completed") + + +async def _url_agent_sdk_stream() -> AsyncIterator[object]: + from a2a.compat.v0_3.types import SendStreamingMessageResponse + + for event in ( + _status_event("working", "Looking that up"), + _artifact_event(_REPLY, append=False), + _status_event("completed"), + ): + yield SendStreamingMessageResponse.model_validate(event) + + +@pytest.mark.parametrize( + ("agent_stream", "reply"), + ( + pytest.param(_snapshot_stream, _REPLY, id="replacing_snapshots"), + pytest.param(_watsonx_stream, _REPLY, id="watsonx_chunks"), + pytest.param(_two_artifact_stream, f"{_FIRST_ARTIFACT} {_SECOND_ARTIFACT}", id="two_artifacts"), + pytest.param(_status_message_stream, _REPLY, id="status_message_reply"), + pytest.param(_artifact_without_id_stream, _REPLY, id="artifact_without_id"), + pytest.param( + _url_agent_sdk_stream, + _REPLY, + id="url_agent_sdk_events", + marks=pytest.mark.skipif(importlib.util.find_spec("a2a") is None, reason="needs the optional a2a-sdk"), + ), + ), +) +@pytest.mark.asyncio +async def test_stream_counts_the_whole_reply_once( + agent_stream: Callable[[], AsyncIterator[object]], reply: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj: Final = LitellmLogging( + model="a2a_agent/trip-planner", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="asend_message_streaming", + start_time=time.time(), + litellm_call_id="artifact-count", + function_id="artifact-count", + ) + iterator: Final = A2AStreamingIterator( + stream=agent_stream(), + request=SimpleNamespace(params=SimpleNamespace(message={"role": "user", "parts": []})), + logging_obj=logging_obj, + agent_name="trip-planner", + ) + + _ = [chunk async for chunk in iterator] + + assert logging_obj.model_call_details["usage"].completion_tokens == A2ARequestUtils.count_tokens(reply) diff --git a/tests/unit/a2a_protocol/test_main.py b/tests/unit/a2a_protocol/test_main.py index 75fd9eda75b..c0cabee1127 100644 --- a/tests/unit/a2a_protocol/test_main.py +++ b/tests/unit/a2a_protocol/test_main.py @@ -1,12 +1,16 @@ """Tests for litellm/a2a_protocol/main.py non-streaming send behavior.""" import asyncio +import contextlib import datetime import json -from typing import Final +from collections.abc import AsyncGenerator, Callable, Coroutine, Iterator, Mapping +from typing import Final, TypeAlias import httpx import pytest +import respx +from pydantic import BaseModel, ConfigDict pytest.importorskip("a2a.compat.v0_3.conversions") @@ -27,6 +31,7 @@ from litellm.a2a_protocol.main import ( create_a2a_client, ) from litellm.a2a_protocol.exceptions import A2AError +from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.caching.llm_caching_handler import LLMClientCache from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT from litellm.llms.custom_httpx.http_handler import ( @@ -798,6 +803,403 @@ def test_streaming_logging_obj_keeps_agent_credentials_out_of_logging_params(): assert logging_obj.model_call_details["litellm_params"] == expected +class _LoggedSpend(BaseModel): + model_config = ConfigDict(frozen=True) + + call_type: str + model: str + prompt_tokens: int + completion_tokens: int + response_cost: float + + +class _A2ASpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: Final = asyncio.Event() + self.spend: _LoggedSpend | None = None + self.agent_id: object = None + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + spend: Final = _LoggedSpend.model_validate(kwargs["standard_logging_object"]) + if spend.call_type not in ("asend_message", "asend_message_streaming"): + return + self.spend = spend + self.agent_id = kwargs.get("agent_id") + self.logged.set() + + async def logged_spend(self) -> _LoggedSpend: + await asyncio.wait_for(self.logged.wait(), timeout=10) + assert self.spend is not None + return self.spend + + +@pytest.fixture +def a2a_spend_recorder(monkeypatch: pytest.MonkeyPatch) -> _A2ASpendRecorder: + recorder: Final = _A2ASpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + return recorder + + +@pytest.fixture +def respx_reachable_httpx(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +_AgentPrices: TypeAlias = Mapping[str, float] +_ExpectedCost: TypeAlias = Callable[[int, int], float] + +_BRIDGE_PROMPT: Final = "What is the capital of France?" +_BRIDGE_REPLY: Final = "Paris is the capital of France." +_LONG_AGENT_REPLY: Final = ( + "Paris is the capital of France. It sits on the Seine, holds about two million people, " + "and has been the seat of the French government for most of the last thousand years." +) +_PER_TOKEN_PRICES: Final[_AgentPrices] = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +_AGENTCORE_ARN: Final = "arn:aws:bedrock-agentcore:us-west-2:123456789:runtime/trip_planner" +_AGENTCORE_PARAMS: Final[Mapping[str, object]] = { + "custom_llm_provider": "bedrock", + "model": f"bedrock/agentcore/{_AGENTCORE_ARN}", + "api_key": "test-jwt-token", +} +_JSON_RPC_ERROR: Final[Mapping[str, object]] = { + "jsonrpc": "2.0", + "id": "agentcore-call", + "error": {"code": -32603, "message": "agent crashed"}, +} + + +def _per_token_cost(prompt_tokens: int, completion_tokens: int) -> float: + return prompt_tokens * 0.001 + completion_tokens * 0.002 + + +_BRIDGE_AGENT_PRICES: Final = ( + pytest.param({"cost_per_query": 0.25}, lambda prompt_tokens, completion_tokens: 0.25, id="cost_per_query"), + pytest.param(_PER_TOKEN_PRICES, _per_token_cost, id="per_token"), +) + + +def _bridge_agent_params(prices: _AgentPrices) -> dict[str, object]: + return { + "custom_llm_provider": "langgraph", + "model": "agent", + "mock_response": _BRIDGE_REPLY, + **prices, + } + + +def _bridge_message_params() -> MessageSendParams: + return MessageSendParams( + message={ + "messageId": "m1", + "role": "user", + "parts": [{"kind": "text", "text": _BRIDGE_PROMPT}], + } + ) + + +def _assert_billed_at_agent_prices( + spend: _LoggedSpend, recorder: _A2ASpendRecorder, expected_cost: _ExpectedCost, reply: str +) -> None: + assert spend.prompt_tokens == A2ARequestUtils.count_tokens(_BRIDGE_PROMPT) + assert spend.completion_tokens == A2ARequestUtils.count_tokens(reply) + assert spend.response_cost == pytest.approx(expected_cost(spend.prompt_tokens, spend.completion_tokens)) + assert recorder.agent_id == "agent-1" + + +async def _drain_bridge_stream(litellm_params: Mapping[str, object]) -> None: + from litellm.a2a_protocol.main import asend_message_streaming + + _ = [ + chunk + async for chunk in asend_message_streaming( + request=SendStreamingMessageRequest(id="bridge-stream", params=_bridge_message_params()), + api_base="http://127.0.0.1:9", + litellm_params=dict(litellm_params), + agent_id="agent-1", + agent_name="trip-planner", + ) + ] + + +def _sse(*events: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events), + ) + + +def _agentcore_appended_chunks(text: str, chunk_size: int) -> tuple[Mapping[str, object], ...]: + return tuple( + { + "jsonrpc": "2.0", + "id": "agentcore-call", + "result": { + "kind": "artifact-update", + "taskId": "t1", + "contextId": "c1", + "append": start > 0, + "artifact": {"artifactId": "a1", "parts": [{"kind": "text", "text": text[start : start + chunk_size]}]}, + }, + } + for start in range(0, len(text), chunk_size) + ) + + +@pytest.mark.parametrize(("prices", "expected_cost"), _BRIDGE_AGENT_PRICES) +@pytest.mark.asyncio +async def test_completion_bridge_send_bills_the_agents_prices( + prices: _AgentPrices, expected_cost: _ExpectedCost, a2a_spend_recorder: _A2ASpendRecorder +) -> None: + await asend_message( + request=SendMessageRequest(id="bridge-send", params=_bridge_message_params()), + api_base="http://127.0.0.1:9", + litellm_params=_bridge_agent_params(prices), + agent_id="agent-1", + ) + + spend: Final = await a2a_spend_recorder.logged_spend() + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, expected_cost, _BRIDGE_REPLY) + + +@pytest.mark.parametrize(("prices", "expected_cost"), _BRIDGE_AGENT_PRICES) +@pytest.mark.asyncio +async def test_completion_bridge_stream_bills_the_agents_prices( + prices: _AgentPrices, expected_cost: _ExpectedCost, a2a_spend_recorder: _A2ASpendRecorder +) -> None: + await _drain_bridge_stream(_bridge_agent_params(prices)) + + spend: Final = await a2a_spend_recorder.logged_spend() + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, expected_cost, _BRIDGE_REPLY) + assert spend.model == "a2a_agent/trip-planner" + + +@pytest.mark.asyncio +async def test_agentcore_send_bills_a_task_reply_at_the_agents_token_prices( + a2a_spend_recorder: _A2ASpendRecorder, respx_reachable_httpx: None +) -> None: + task_reply: Final = { + "jsonrpc": "2.0", + "id": "bridge-send", + "result": { + "kind": "task", + "id": "t1", + "contextId": "c1", + "status": {"state": "completed"}, + "artifacts": [{"artifactId": "a1", "parts": [{"kind": "text", "text": _LONG_AGENT_REPLY}]}], + }, + } + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock(return_value=httpx.Response(200, json=task_reply)) + await asend_message( + request=SendMessageRequest(id="bridge-send", params=_bridge_message_params()), + litellm_params={**_AGENTCORE_PARAMS, **_PER_TOKEN_PRICES}, + agent_id="agent-1", + ) + + spend: Final = await a2a_spend_recorder.logged_spend() + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, _per_token_cost, _LONG_AGENT_REPLY) + + +@pytest.mark.parametrize( + "extra_fields", + (pytest.param({}, id="result_only"), pytest.param({"error": None}, id="null_error_field")), +) +@pytest.mark.asyncio +async def test_agentcore_stream_bills_every_appended_reply_chunk( + extra_fields: Mapping[str, object], a2a_spend_recorder: _A2ASpendRecorder, respx_reachable_httpx: None +) -> None: + events: Final = tuple({**event, **extra_fields} for event in _agentcore_appended_chunks(_LONG_AGENT_REPLY, 40)) + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock(return_value=_sse(*events)) + await _drain_bridge_stream({**_AGENTCORE_PARAMS, **_PER_TOKEN_PRICES}) + + spend: Final = await a2a_spend_recorder.logged_spend() + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, _per_token_cost, _LONG_AGENT_REPLY) + + +_PYDANTIC_COMPLETED_TASK: Final[Mapping[str, object]] = { + "jsonrpc": "2.0", + "id": "bridge-stream", + "result": { + "kind": "task", + "id": "t1", + "contextId": "c1", + "status": {"state": "completed"}, + "artifacts": [{"artifactId": "a1", "parts": [{"kind": "text", "text": _LONG_AGENT_REPLY}]}], + }, +} + + +@pytest.mark.asyncio +async def test_pydantic_ai_stream_bills_the_whole_chunked_reply( + a2a_spend_recorder: _A2ASpendRecorder, respx_reachable_httpx: None +) -> None: + with respx.mock(assert_all_called=True) as router: + router.post("http://127.0.0.1:9/").mock(return_value=httpx.Response(200, json=_PYDANTIC_COMPLETED_TASK)) + await _drain_bridge_stream({"custom_llm_provider": "pydantic_ai_agents", **_PER_TOKEN_PRICES}) + + spend: Final = await a2a_spend_recorder.logged_spend() + assert len(_LONG_AGENT_REPLY) > 50 + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, _per_token_cost, _LONG_AGENT_REPLY) + + +async def _close_after_first_event(stream: AsyncGenerator[object, None]) -> None: + _ = await stream.__anext__() + await stream.aclose() + + +async def _cancel_reader_after_first_event(stream: AsyncGenerator[object, None]) -> None: + first_event: Final = asyncio.Event() + + async def _read() -> None: + async for _ in stream: + first_event.set() + + reader: Final = asyncio.create_task(_read()) + await first_event.wait() + reader.cancel() + with contextlib.suppress(asyncio.CancelledError): + await reader + + +def _pydantic_bridge_agent_stream() -> AsyncGenerator[object, None]: + from litellm.a2a_protocol.main import asend_message_streaming + + return asend_message_streaming( + request=SendStreamingMessageRequest(id="bridge-stream", params=_bridge_message_params()), + api_base="http://127.0.0.1:9", + litellm_params={"custom_llm_provider": "pydantic_ai_agents", "cost_per_query": 0.25}, + agent_id="agent-1", + agent_name="trip-planner", + ) + + +def _url_agent_stream() -> AsyncGenerator[object, None]: + from litellm.a2a_protocol.main import asend_message_streaming + + return asend_message_streaming( + request=SendStreamingMessageRequest(id="reply", params=_bridge_message_params()), + api_base="http://127.0.0.1:9", + litellm_params={"cost_per_query": 0.25}, + agent_id="agent-1", + ) + + +@pytest.mark.parametrize( + ("open_stream", "billed_model"), + ( + pytest.param(_pydantic_bridge_agent_stream, "a2a_agent/trip-planner", id="bridge_agent"), + pytest.param(_url_agent_stream, f"a2a_agent/{_AGENT_CARD['name']}", id="url_agent"), + ), +) +@pytest.mark.parametrize( + "drop_stream", + ( + pytest.param(_close_after_first_event, id="proxy_closes_the_stream"), + pytest.param(_cancel_reader_after_first_event, id="reader_cancelled_mid_stream"), + ), +) +@pytest.mark.asyncio +async def test_stream_the_client_drops_still_bills_the_agents_price( + open_stream: Callable[[], AsyncGenerator[object, None]], + billed_model: str, + drop_stream: Callable[[AsyncGenerator[object, None]], Coroutine[object, object, None]], + a2a_spend_recorder: _A2ASpendRecorder, + isolated_client_cache: LLMClientCache, + respx_reachable_httpx: None, +) -> None: + _ = await _seed_shared_a2a_client(rpc_stream_events=_LANGGRAPH_STREAM_EVENTS) + with respx.mock(assert_all_called=False) as router: + router.post("http://127.0.0.1:9/").mock(return_value=httpx.Response(200, json=_PYDANTIC_COMPLETED_TASK)) + await drop_stream(open_stream()) + + spend: Final = await a2a_spend_recorder.logged_spend() + assert spend.response_cost == pytest.approx(0.25) + assert spend.model == billed_model + + +class _AgentThatNeverAnswers(_RequestRecorder): + def __init__(self) -> None: + super().__init__() + self.message_sent: Final = asyncio.Event() + + def __call__(self, request: httpx.Request) -> httpx.Response | Coroutine[object, object, httpx.Response]: + if request.method == "GET": + return super().__call__(request) + return self._never_answer() + + async def _never_answer(self) -> httpx.Response: + self.message_sent.set() + reply: Final[asyncio.Future[httpx.Response]] = asyncio.get_running_loop().create_future() + return await reply + + +@pytest.mark.asyncio +async def test_stream_dropped_before_the_agent_answers_bills_nothing( + a2a_spend_recorder: _A2ASpendRecorder, isolated_client_cache: LLMClientCache +) -> None: + agent: Final = _AgentThatNeverAnswers() + _ = await _seed_shared_a2a_client(recorder=agent) + stream: Final = _url_agent_stream() + reader: Final = asyncio.create_task(anext(stream)) + await asyncio.wait_for(agent.message_sent.wait(), timeout=10) + reader.cancel() + with contextlib.suppress(asyncio.CancelledError): + await reader + + with pytest.raises(TimeoutError): + await asyncio.wait_for(a2a_spend_recorder.logged.wait(), timeout=2) + + +@pytest.mark.asyncio +async def test_completion_bridge_stream_with_an_empty_reply_bills_no_reply_tokens( + a2a_spend_recorder: _A2ASpendRecorder, +) -> None: + await _drain_bridge_stream({**_bridge_agent_params(_PER_TOKEN_PRICES), "mock_response": ""}) + + spend: Final = await a2a_spend_recorder.logged_spend() + assert spend.completion_tokens == 0 + assert spend.response_cost == pytest.approx(_per_token_cost(spend.prompt_tokens, 0)) + + +@pytest.mark.parametrize( + "upstream_reply", + ( + pytest.param(httpx.Response(200, json=_JSON_RPC_ERROR), id="send"), + pytest.param(_sse(_JSON_RPC_ERROR), id="stream"), + ), +) +@pytest.mark.asyncio +async def test_agentcore_error_reply_is_not_billed_at_the_agents_price( + upstream_reply: httpx.Response, a2a_spend_recorder: _A2ASpendRecorder, respx_reachable_httpx: None +) -> None: + agent_params: Final = {**_AGENTCORE_PARAMS, "cost_per_query": 0.25} + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock(return_value=upstream_reply) + if upstream_reply.headers["content-type"] == "text/event-stream": + await _drain_bridge_stream(agent_params) + else: + response = await asend_message( + request=SendMessageRequest(id="agentcore-call", params=_bridge_message_params()), + litellm_params=dict(agent_params), + agent_id="agent-1", + ) + assert response.model_dump(mode="json", exclude_none=True)["error"]["message"] == "agent crashed" + + spend: Final = await a2a_spend_recorder.logged_spend() + assert spend.response_cost == 0.0 + + class MockA2AResponse: def __init__(self, text: str): self._payload = { diff --git a/tests/unit/batches/test_main.py b/tests/unit/batches/test_main.py index fb97e0e4919..8b22395f462 100644 --- a/tests/unit/batches/test_main.py +++ b/tests/unit/batches/test_main.py @@ -37,13 +37,15 @@ import litellm.batches.main as bm import asyncio import datetime import json -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import Final import httpx import respx -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm.integrations.custom_logger import CustomLogger +from openai.types.batch import Batch +from litellm.types.utils import LiteLLMBatch # --------------------------------------------------------------------------- # @@ -1204,3 +1206,183 @@ async def test_delete_batch_output_file(respx_mock: respx.MockRouter, monkeypatc assert delete_route.call_count == 1 assert delete_response.id == "file-output123" assert delete_response.deleted is True + + +VLLM_ROUTER_API_BASE: Final = "http://vllm-router.test/v1" +VLLM_ROUTER_BATCH: Final = MappingProxyType( + { + "id": "batch_5f8dac", + "status": "pending", + "input_file_id": "file-b73f4c", + "created_at": 1793239207, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "output_file_id": None, + "total_requests": None, + "completed_requests": 0, + "failed_requests": 0, + "metadata": {}, + } +) + + +async def _create_on_vllm_router() -> LiteLLMBatch: + return await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-b73f4c", + custom_llm_provider="hosted_vllm", + api_key="sk-test", + api_base=VLLM_ROUTER_API_BASE, + ) + + +async def _retrieve_from_vllm_router() -> LiteLLMBatch: + return await litellm.aretrieve_batch( + batch_id="batch_5f8dac", custom_llm_provider="hosted_vllm", api_key="sk-test", api_base=VLLM_ROUTER_API_BASE + ) + + +VLLM_ROUTER_BATCH_CALLS: Final = ( + pytest.param("POST", "/batches", _create_on_vllm_router, id="acreate_batch"), + pytest.param("GET", "/batches/batch_5f8dac", _retrieve_from_vllm_router, id="aretrieve_batch"), +) + + +@pytest.mark.parametrize("method, path, call_router", VLLM_ROUTER_BATCH_CALLS) +@pytest.mark.parametrize( + "router_status, openai_status", + ( + ("pending", "validating"), + ("running", "in_progress"), + ("completed", "completed"), + ("failed", "failed"), + ("cancelled", "cancelled"), + ), +) +async def test_hosted_vllm_router_batch_comes_back_as_an_openai_batch( + respx_mock: respx.MockRouter, + method: str, + path: str, + call_router: Callable[[], Awaitable[LiteLLMBatch]], + router_status: str, + openai_status: str, +): + respx_mock.route(method=method, url=f"{VLLM_ROUTER_API_BASE}{path}").respond( + 200, json={**VLLM_ROUTER_BATCH, "status": router_status} + ) + + batch: Final = await call_router() + + assert (batch.id, batch.object, batch.status) == ("batch_5f8dac", "batch", openai_status) + + +@pytest.mark.parametrize("method, path, call_router", VLLM_ROUTER_BATCH_CALLS) +@pytest.mark.parametrize("field, value", (("status", "queued"), ("object", "file"))) +async def test_hosted_vllm_router_batch_with_an_unknown_value_is_rejected( + respx_mock: respx.MockRouter, + method: str, + path: str, + call_router: Callable[[], Awaitable[LiteLLMBatch]], + field: str, + value: str, +): + respx_mock.route(method=method, url=f"{VLLM_ROUTER_API_BASE}{path}").respond( + 200, json={**VLLM_ROUTER_BATCH, field: value} + ) + + with pytest.raises(ValidationError) as exc_info: + await call_router() + + assert [error["loc"] for error in exc_info.value.errors()] == [(field,)] + + +async def _list_on_vllm_router() -> Sequence[Batch]: + page: Final = await litellm.alist_batches( + custom_llm_provider="hosted_vllm", api_key="sk-test", api_base=VLLM_ROUTER_API_BASE + ) + return page.data + + +async def _list_on_vllm_router_sync() -> Sequence[Batch]: + return litellm.list_batches(custom_llm_provider="hosted_vllm", api_key="sk-test", api_base=VLLM_ROUTER_API_BASE).data + + +@pytest.mark.parametrize("list_router", (_list_on_vllm_router, _list_on_vllm_router_sync)) +async def test_hosted_vllm_router_batch_list_comes_back_as_openai_batches( + respx_mock: respx.MockRouter, list_router: Callable[[], Awaitable[Sequence[Batch]]] +): + respx_mock.get(f"{VLLM_ROUTER_API_BASE}/batches").respond( + 200, + json={ + "object": "list", + "data": [ + {**VLLM_ROUTER_BATCH, "id": "batch_5f8dac", "status": "pending"}, + {**VLLM_ROUTER_BATCH, "id": "batch_c74802", "status": "running"}, + ], + "first_id": "batch_5f8dac", + "last_id": "batch_c74802", + "has_more": False, + }, + ) + + batches: Final = await list_router() + + assert [(batch.id, batch.object, batch.status) for batch in batches] == [ + ("batch_5f8dac", "batch", "validating"), + ("batch_c74802", "batch", "in_progress"), + ] + + +async def test_hosted_vllm_router_batch_list_passes_through_a_status_that_is_not_a_string( + respx_mock: respx.MockRouter, +): + respx_mock.get(f"{VLLM_ROUTER_API_BASE}/batches").respond( + 200, + json={ + "object": "list", + "data": [{**VLLM_ROUTER_BATCH, "status": ["pending"]}], + "first_id": "batch_5f8dac", + "last_id": "batch_5f8dac", + "has_more": False, + }, + ) + + batches: Final = await _list_on_vllm_router() + + assert [(batch.id, batch.object, batch.status) for batch in batches] == [("batch_5f8dac", "batch", ["pending"])] + + +async def test_openai_batch_list_items_come_back_as_openai_sent_them(respx_mock: respx.MockRouter): + respx_mock.get("http://openai.test/v1/batches").respond( + 200, + json={ + "object": "list", + "data": [ + { + **_OPENAI_BATCH_JSON, + "status": "completed", + "usage": { + "input_tokens": 10, + "input_tokens_details": {"cached_tokens": 4}, + "output_tokens": 5, + "output_tokens_details": {"reasoning_tokens": 2}, + "total_tokens": 15, + }, + }, + {**_OPENAI_BATCH_JSON, "id": "batch_def456", "status": "queued"}, + ], + "first_id": "batch_abc123", + "last_id": "batch_def456", + "has_more": False, + }, + ) + + page: Final = await litellm.alist_batches( + custom_llm_provider="openai", api_key="sk-test", api_base="http://openai.test/v1" + ) + + first, second = page.data + assert first.usage is not None + assert (first.usage.input_tokens_details.cached_tokens, first.usage.output_tokens_details.reasoning_tokens) == (4, 2) + assert (second.id, second.object, second.status) == ("batch_def456", "batch", "queued") diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index b1da2a413b6..37169e13a0b 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -29,6 +29,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3, ) from litellm.proxy.utils import InternalUsageCache +from litellm.router_api import RoutingContext from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCache from litellm.router_utils.routing_read_batch import ( ROUTER_COOLDOWNS_USAGE_TARGET, @@ -390,7 +391,7 @@ async def test_a_prefetch_with_incomplete_usage_keys_releases_cooldown_reservati ) with request_redis_batch_scope(): - RoutingPrefetch.arm(router, router.lowesttpm_logger_v2, router.model_list[:1]) + RoutingPrefetch.arm(RoutingContext(router), router.lowesttpm_logger_v2, router.model_list[:1]) deployment: Final = await router.async_get_available_deployment( model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} ) diff --git a/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py b/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py index 5b26dad269e..13fef7a0945 100644 --- a/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py @@ -241,21 +241,21 @@ async def test_router_cooldown_event_callback(): Ensures that the router_cooldown_event_callback function correctly logs the cooldown event to the PrometheusLogger """ - # Mock Router instance - mock_router = MagicMock() + # Mock routing context + mock_ctx = MagicMock() mock_deployment = { "litellm_params": {"model": "gpt-5-mini"}, "model_name": "gpt-5-mini", "model_info": ModelInfo(id="test-model-id"), } - mock_router.get_deployment.return_value = mock_deployment + mock_ctx.registry.return_value.deployment.return_value = mock_deployment # Create a real PrometheusLogger instance prometheus_logger = CustomPrometheusLogger() litellm.callbacks = [prometheus_logger] await router_cooldown_event_callback( - litellm_router_instance=mock_router, + ctx=mock_ctx, deployment_id="test-deployment", exception_status="429", cooldown_time=60.0, @@ -263,8 +263,8 @@ async def test_router_cooldown_event_callback(): await asyncio.sleep(0.5) - # Assert that the router's get_deployment method was called - mock_router.get_deployment.assert_called_once_with(model_id="test-deployment") + # Assert that the router's deployment lookup was called + mock_ctx.registry.return_value.deployment.assert_called_once_with("test-deployment") print( "prometheus_logger.deployment_complete_outages", @@ -301,21 +301,21 @@ async def test_router_cooldown_event_callback_no_prometheus(): Ensures that the router_cooldown_event_callback function does not raise an error when no PrometheusLogger is found """ - # Mock Router instance - mock_router = MagicMock() + # Mock routing context + mock_ctx = MagicMock() mock_deployment = { "litellm_params": {"model": "gpt-5-mini"}, "model_name": "gpt-5-mini", "model_info": ModelInfo(id="test-model-id"), } - mock_router.get_deployment.return_value = mock_deployment + mock_ctx.registry.return_value.deployment.return_value = mock_deployment await router_cooldown_event_callback( - litellm_router_instance=mock_router, + ctx=mock_ctx, deployment_id="test-deployment", exception_status="429", cooldown_time=60.0, ) - # Assert that the router's get_deployment method was called - mock_router.get_deployment.assert_called_once_with(model_id="test-deployment") + # Assert that the router's deployment lookup was called + mock_ctx.registry.return_value.deployment.assert_called_once_with("test-deployment") diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index f2df14cc8a3..89fff365c3f 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -1437,6 +1437,8 @@ def _make_object_store_instance(): mock_prisma = MagicMock() mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock() mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_prisma.db.writer = mock_prisma.db return ( PROXY_LiteLLMManagedFiles( @@ -2161,3 +2163,108 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets assert (cached.storage_backend, cached.storage_url) == ("litellm_db", "litellm_db://content-row-1") create_data = file_table.upsert.await_args.kwargs["data"]["create"] assert (create_data["storage_backend"], create_data["storage_url"]) == ("litellm_db", "litellm_db://content-row-1") + + +@pytest.mark.asyncio +async def test_create_stores_the_creating_deployment_in_file_object(): + """Regression (LIT-10326): an unmanaged gs:// batch stores the raw Vertex job id, which + says nothing about which of two deployments sharing a base model created it. The create + must keep the router's deployment id inside the stored file_object so the cost poller + reads the result through that deployment instead of the first base-model match.""" + managed_files, mock_prisma = _make_object_store_instance() + creator: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="validating"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + create_if_missing=True, + creating_model_id="deploy-team-b", + ) + + created: Final = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"] + assert json.loads(created["file_object"])["model_id"] == "deploy-team-b" + + +@pytest.mark.asyncio +async def test_poll_refresh_keeps_the_stored_creating_deployment(): + """A later observe of the row (a GET through whichever deployment the client named) + rewrites file_object; it must carry the create-time model_id forward, not drop it or + replace it with the deployment that served the retrieve.""" + managed_files, mock_prisma = _make_object_store_instance() + stored_row: Final = MagicMock() + stored_row.file_object = json.dumps({**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"}) + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row) + observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=observer, + persist_attribution=False, + create_if_missing=False, + ) + + written: Final = json.loads(mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"]) + assert written["model_id"] == "deploy-team-b" + assert written["status"] == "completed" + + +@pytest.mark.asyncio +async def test_poll_refresh_reads_the_stored_creating_deployment_from_the_writer(): + """With a read replica configured, a retrieve that lands before the replica has the create + must still find the create-time model_id, so the lookup goes to the writer. Prisma may hand + the Json column back already decoded, so the row here is a dict rather than JSON text.""" + managed_files, mock_prisma = _make_object_store_instance() + stored_row: Final = MagicMock() + stored_row.file_object = {**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"} + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_prisma.db.writer = MagicMock() + mock_prisma.db.writer.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row) + observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=observer, + persist_attribution=False, + create_if_missing=False, + ) + + written: Final = json.loads(mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"]) + assert written["model_id"] == "deploy-team-b" + mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_create_hook_passes_the_router_deployment_as_creating_model_id(): + """The post-call hook on a batch create hands the router's model_id to the object store; + a retrieve of the same batch hands None so it cannot overwrite the creator.""" + from litellm_enterprise.proxy.hooks.managed_files import BATCH_CREATE_HIDDEN_PARAM + + managed_files: Final = _make_managed_files_instance() + created: Final = _make_batch_response(model_id="deploy-team-b", output_file_id=None, status="validating") + created._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True + retrieved: Final = _make_batch_response(model_id="deploy-team-a", output_file_id=None, status="in_progress") + + with patch("litellm.proxy.proxy_server.llm_router", MagicMock()): + await managed_files.async_post_call_success_hook( + data={}, user_api_key_dict=_make_user_api_key_dict(), response=created + ) + await managed_files.async_post_call_success_hook( + data={}, user_api_key_dict=_make_user_api_key_dict(), response=retrieved + ) + + first, second = managed_files.store_unified_object_id.await_args_list + assert first.kwargs["creating_model_id"] == "deploy-team-b" + assert second.kwargs["creating_model_id"] is None diff --git a/tests/unit/files/test_main.py b/tests/unit/files/test_main.py index e91951fb5ba..44aadc9014e 100644 --- a/tests/unit/files/test_main.py +++ b/tests/unit/files/test_main.py @@ -1,3 +1,4 @@ +from collections.abc import Awaitable, Callable from types import MappingProxyType from typing import Final from urllib.parse import parse_qs, urlparse @@ -124,6 +125,61 @@ async def test_afile_retrieve_rejects_a_provider_file_without_its_size(): assert [error["loc"] for error in exc_info.value.errors()] == [("bytes",)] +VLLM_ROUTER_API_BASE: Final = "http://vllm-router.test/v1" +VLLM_ROUTER_FILE: Final = MappingProxyType( + {"id": "file-b73f4c", "bytes": 156, "created_at": 1793238707, "filename": "batch_requests.jsonl", "purpose": "batch"} +) + + +async def _upload_to_vllm_router() -> OpenAIFileObject: + return await litellm.acreate_file( + file=("batch_requests.jsonl", b'{"custom_id": "request-1"}'), + purpose="batch", + custom_llm_provider="hosted_vllm", + api_key="sk-test", + api_base=VLLM_ROUTER_API_BASE, + ) + + +async def _retrieve_from_vllm_router() -> OpenAIFileObject: + return await litellm.afile_retrieve( + file_id="file-b73f4c", custom_llm_provider="hosted_vllm", api_key="sk-test", api_base=VLLM_ROUTER_API_BASE + ) + + +VLLM_ROUTER_FILE_CALLS: Final = ( + pytest.param("POST", "/files", _upload_to_vllm_router, id="acreate_file"), + pytest.param("GET", "/files/file-b73f4c", _retrieve_from_vllm_router, id="afile_retrieve"), +) + + +@pytest.mark.parametrize("method, path, call_router", VLLM_ROUTER_FILE_CALLS) +@respx.mock +async def test_hosted_vllm_file_without_an_object_type_comes_back_as_a_file( + method: str, path: str, call_router: Callable[[], Awaitable[OpenAIFileObject]] +): + respx.route(method=method, url=f"{VLLM_ROUTER_API_BASE}{path}").respond(200, json=dict(VLLM_ROUTER_FILE)) + + file_object: Final = await call_router() + + assert file_object.model_dump(exclude_none=True) == {**VLLM_ROUTER_FILE, "object": "file"} + + +@pytest.mark.parametrize("method, path, call_router", VLLM_ROUTER_FILE_CALLS) +@respx.mock +async def test_hosted_vllm_file_with_a_different_object_type_is_rejected( + method: str, path: str, call_router: Callable[[], Awaitable[OpenAIFileObject]] +): + respx.route(method=method, url=f"{VLLM_ROUTER_API_BASE}{path}").respond( + 200, json={**VLLM_ROUTER_FILE, "object": "batch"} + ) + + with pytest.raises(ValidationError) as exc_info: + await call_router() + + assert [error["loc"] for error in exc_info.value.errors()] == [("object",)] + + _FILE_BODY: Final = b'{"prompt": "Hello", "completion": "Hi"}' _FINE_TUNE_FILE_JSON: Final = MappingProxyType( { diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index 1f9d3af71d9..ac1f5557bf6 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -207,6 +207,11 @@ def test_routing_diagnostics_reach_llm_spans_without_promoting_prompt_text(fail: "classifier_cost": 0.001, "classifier_failure_reason": "timeout", "classifier_error_type": "TimeoutError", + "classifier_input_budget_tokens": 24576, + "classifier_input_tokens": 23000, + "classifier_input_count_method": "estimated_tokens", + "classifier_input_truncated": True, + "classifier_input_budget_status": "truncated", "escalated": False, } payload: Final = _payload( diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 0483888533e..c32c1b71f75 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3800,3 +3800,16 @@ def test_default_on_guardrail() -> None: ) is True ) + + +def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings(): + from litellm.integrations.custom_guardrail import merged_guardrail_entries + + assert merged_guardrail_entries(["a"], [{"a": {}, "b": {}}, {"a": {"x": 1}}], ["b", "c", "c"]) == [ + {"a": {}, "b": {}}, + "c", + ] + assert merged_guardrail_entries(["pii", {"pii": {"extra_body": {"threshold": 1}}}], ["model-guard"]) == [ + {"pii": {"extra_body": {"threshold": 1}}}, + "model-guard", + ] diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index f0fca94c370..5f749e951b1 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -344,6 +344,11 @@ def test_v1_routing_diagnostics_add_scalar_attributes_and_preserve_legacy_metada "cause": "heuristic", "classifier_failure_reason": "timeout", "classifier_error_type": "TimeoutError", + "classifier_input_budget_tokens": 24576, + "classifier_input_tokens": 23000, + "classifier_input_count_method": "estimated_tokens", + "classifier_input_truncated": True, + "classifier_input_budget_status": "truncated", "classifier_cost": 0.001, } with provider.get_tracer(__name__).start_as_current_span("inference") as span: diff --git a/tests/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index bff89b6cbf1..878f19c66f8 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -238,6 +238,53 @@ def test_extract_api_provider_from_request_data_failure_path(): assert extract({"model": "some-unknown-model-xyz"}) is None +_PASS_THROUGH_CALL_TYPE: Final = "pass_through_endpoint" + + +@pytest.mark.parametrize( + "request_data, expected_provider", + [ + ({"model": "gpt-4o-mini", "call_type": "acompletion", "custom_llm_provider": "anthropic"}, "openai"), + ({"model": "gpt-4o-mini", "litellm_params": {"custom_llm_provider": "humanloop"}}, "humanloop"), + ({"model": "gemini-3.8-flash", "call_type": _PASS_THROUGH_CALL_TYPE}, "vertex_ai"), + ( + {"model": "gemini-3.8-flash", "call_type": _PASS_THROUGH_CALL_TYPE, "custom_llm_provider": "gemini"}, + "gemini", + ), + ( + { + "call_type": _PASS_THROUGH_CALL_TYPE, + "custom_llm_provider": "gemini", + "standard_logging_object": {"custom_llm_provider": "vertex_ai"}, + }, + "vertex_ai", + ), + ], +) +def test_extract_api_provider_reads_the_request_declared_provider_only_on_pass_through_failures( + request_data: dict[str, object], expected_provider: str +): + from litellm.integrations.prometheus import PrometheusLogger + + assert PrometheusLogger._extract_api_provider_from_request_data(request_data) == expected_provider + + +@pytest.mark.parametrize( + "declared_provider", + ["azure", "bedrock", "openrouter", "client-supplied-value", "", {"nested": "object"}, 42], +) +def test_extract_api_provider_ignores_pass_through_declared_values_outside_gemini_and_vertex( + declared_provider: object, +): + from litellm.integrations.prometheus import PrometheusLogger + + extract = PrometheusLogger._extract_api_provider_from_request_data + declared = {"call_type": _PASS_THROUGH_CALL_TYPE, "custom_llm_provider": declared_provider} + + assert extract({"model": "gpt-4o-mini", **declared}) == "openai" + assert extract({"model": "no-such-model", **declared}) is None + + def test_extract_api_provider_swallows_unknown_model_but_logs_unexpected_errors(): """ get_llm_provider raises BadRequestError for a model that maps to no @@ -849,6 +896,56 @@ async def test_failure_hook_prefers_request_data_provider_over_exception_provide ) == ["azure"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type, custom_llm_provider, expected_requested_model", + [ + (_PASS_THROUGH_CALL_TYPE, "vertex_ai", "gcp/google/gemini-2.5-flash"), + ("acompletion", "vertex_ai", "other"), + (_PASS_THROUGH_CALL_TYPE, "bedrock", "other"), + ], +) +async def test_failure_hook_labels_a_pass_through_failure_with_the_router_model_the_client_requested( + call_type: str, custom_llm_provider: str, expected_requested_model: str +): + from unittest.mock import patch + + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import UserAPIKeyAuth + + router = litellm.Router( + model_list=[ + { + "model_name": "gcp/google/gemini-2.5-flash", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "p", + "vertex_location": "us-central1", + }, + } + ] + ) + _clear_prometheus_registry() + try: + with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + await PrometheusLogger().async_post_call_failure_hook( + request_data={ + "model": "gemini-2.5-flash", + "call_type": call_type, + "custom_llm_provider": custom_llm_provider, + "metadata": {}, + "litellm_params": {"metadata": {"model_group": "gcp/google/gemini-2.5-flash"}}, + }, + original_exception=Exception("upstream 400"), + user_api_key_dict=UserAPIKeyAuth(token="tok"), + ) + assert { + s.labels.get("requested_model") for s in _collected_samples("litellm_proxy_failed_requests_metric_total") + } == {expected_requested_model} + finally: + _clear_prometheus_registry() + + def test_model_group_in_deployment_metrics(): """ Test that model_group label is present on the deployment-scoped metrics diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index e2452a928a4..9da2632943c 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from typing import Final import pytest @@ -7,7 +8,7 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebSearchOptions -from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams +from litellm.types.utils import ModelResponse, ServerToolUse, StandardBuiltInToolsParams, Usage def test_web_search_cost_low(): @@ -1052,6 +1053,175 @@ def test_azure_ai_non_claude_web_search_keeps_generic_fallback() -> None: assert get_cost_for_web_search_request(custom_llm_provider="azure_ai", usage=usage, model_info=model_info) is None +def _register_priced_deployment( + model: str, search_cost_per_query: float | None, deployment_id: str | None = None +) -> str: + from litellm import Router + + token_rates: Final = {"input_cost_per_token": 3e-6, "output_cost_per_token": 1.5e-5} + search_rates: Final = ( + {} + if search_cost_per_query is None + else { + "search_context_cost_per_query": { + "search_context_size_low": search_cost_per_query, + "search_context_size_medium": search_cost_per_query, + "search_context_size_high": search_cost_per_query, + } + } + ) + router: Final = Router( + model_list=[ + { + "model_name": "priced-deployment", + "litellm_params": {"model": model, "api_key": "k", "vertex_project": "p", "vertex_location": "global"}, + "model_info": {**token_rates, **search_rates, **({} if deployment_id is None else {"id": deployment_id})}, + } + ] + ) + return router.model_list[0]["model_info"]["id"] + + +def _claude_chat_response( + web_search_requests: int, model: str = "claude-sonnet-5-5" +) -> tuple[object, Usage | None]: + usage: Final = Usage( + prompt_tokens=1840, + completion_tokens=412, + total_tokens=2252, + server_tool_use=ServerToolUse(web_search_requests=web_search_requests), + ) + return ModelResponse(model=model, usage=usage), usage + + +def _claude_messages_response( + web_search_requests: int, model: str = "claude-sonnet-5-5" +) -> tuple[object, Usage | None]: + return { + "id": "msg_vertex_claude_web_search", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "scripted answer"}], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "server_tool_use": {"web_search_requests": web_search_requests}, + }, + }, None + + +@pytest.mark.parametrize("build_response", [_claude_chat_response, _claude_messages_response]) +@pytest.mark.parametrize(("web_search_requests", "search_cost_per_query"), [(2, 0.01), (3, 0.01), (2, 0.0)]) +def test_vertex_claude_deployment_priced_web_search_billed_per_request( + local_model_cost_map: None, + build_response: Callable[[int], tuple[object, Usage | None]], + web_search_requests: int, + search_cost_per_query: float, +) -> None: + deployment_id: Final = _register_priced_deployment("vertex_ai/claude-sonnet-5-5", search_cost_per_query) + response_object, usage = build_response(web_search_requests) + + cost: Final = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=deployment_id, + response_object=response_object, + usage=usage, + custom_llm_provider="vertex_ai", + standard_built_in_tools_params=None, + ) + + assert cost == pytest.approx(web_search_requests * search_cost_per_query) + + +@pytest.mark.parametrize("build_response", [_claude_chat_response, _claude_messages_response]) +@pytest.mark.parametrize( + ("deployment_model", "served_model", "custom_llm_provider"), + # Google bills Claude web search on Sonnet 4.5 and Haiku 4.5 too: + # https://cloud.google.com/model-garden/partner-models/pricing (read 2026-10-30) + [ + ("vertex_ai/claude-sonnet-5-5", "claude-sonnet-5-5", "vertex_ai"), + ("vertex_ai/claude-sonnet-4-5@20250929", "claude-sonnet-4-5-20250929", "vertex_ai"), + ("vertex_ai/claude-haiku-4-5@20251001", "claude-haiku-4-5-20251001", "vertex_ai"), + ("anthropic/claude-sonnet-5-5", "claude-sonnet-5-5", "anthropic"), + ("azure_ai/claude-sonnet-5-5", "claude-sonnet-5-5", "azure_ai"), + ], +) +@pytest.mark.parametrize("deployment_id", [None, "claude-token-priced"], ids=["hashed_id", "claude_in_id"]) +def test_claude_deployment_with_token_rates_only_bills_served_model_search_price( + local_model_cost_map: None, + build_response: Callable[[int, str], tuple[object, Usage | None]], + deployment_model: str, + served_model: str, + custom_llm_provider: str, + deployment_id: str | None, +) -> None: + served_search_pricing: Final = ( + litellm.get_model_info(model=deployment_model).get("search_context_cost_per_query") or {} + ) + served_search_price: Final = served_search_pricing.get("search_context_size_medium") or 0.0 + registered_id: Final = _register_priced_deployment(deployment_model, None, deployment_id) + response_object, usage = build_response(2, served_model) + + cost: Final = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=registered_id, + response_object=response_object, + usage=usage, + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=None, + ) + + assert served_search_price > 0 + assert cost == pytest.approx(2 * served_search_price) + + +def test_vertex_claude_deployment_with_token_rates_only_and_unmapped_served_model_bills_no_search( + local_model_cost_map: None, +) -> None: + deployment_id: Final = _register_priced_deployment("vertex_ai/claude-sonnet-5-5", None) + usage: Final = Usage( + prompt_tokens=1840, + completion_tokens=412, + total_tokens=2252, + server_tool_use=ServerToolUse(web_search_requests=2), + ) + + cost: Final = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=deployment_id, + response_object=ModelResponse(model="claude-unmapped-preview", usage=usage), + usage=usage, + custom_llm_provider="vertex_ai", + standard_built_in_tools_params=None, + ) + + assert cost == 0.0 + + +@pytest.mark.parametrize("deployment_id", [None, "claude-code-gemini"], ids=["hashed_id", "claude_in_id"]) +def test_vertex_gemini_deployment_priced_web_search_stays_per_prompt( + local_model_cost_map: None, deployment_id: str | None +) -> None: + search_cost_per_prompt: Final = 0.03 + registered_id: Final = _register_priced_deployment( + "vertex_ai/gemini-2.5-flash", search_cost_per_prompt, deployment_id + ) + usage: Final = Usage( + prompt_tokens=1840, + completion_tokens=412, + total_tokens=2252, + server_tool_use=ServerToolUse(web_search_requests=2), + ) + + cost: Final = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=registered_id, + response_object=ModelResponse(model="gemini-2.5-flash", usage=usage), + usage=usage, + custom_llm_provider="vertex_ai", + standard_built_in_tools_params=None, + ) + + assert cost == pytest.approx(search_cost_per_prompt) + + @pytest.mark.parametrize( "kwargs", [{"tools": 5}, {"web_search_options": None}, {"web_search_options": 5}], diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index de8d6a7a8c5..d0f5bfc6b9f 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -3089,6 +3089,50 @@ def test_convert_to_model_response_object_function_output(): assert result._hidden_params is not None +@pytest.mark.parametrize( + "reported_cost", + [ + float("nan"), + float("inf"), + {"input_tokens_cost": 0.0001, "request_cost": 0.005, "total_cost": float("nan")}, + {"input_tokens_cost": float("inf"), "request_cost": 0.005, "total_cost": 0.0051}, + ], + ids=["nan", "inf", "breakdown-with-nan-total", "breakdown-with-inf-part"], +) +def test_a_non_finite_provider_usage_cost_is_dropped_so_the_response_is_strict_json(reported_cost: object): + response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-non-finite-cost", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4, "cost": reported_cost}, + }, + model_response_object=ModelResponse(), + ) + + usage: Final = json.loads(json.dumps(response.model_dump(), allow_nan=False))["usage"] + assert "cost" not in usage + assert usage["total_tokens"] == 4 + + +def test_a_finite_provider_cost_breakdown_with_non_string_keys_is_kept(): + response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-int-key-cost", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4, "cost": {1: 0.5}}, + }, + model_response_object=ModelResponse(), + ) + + assert response.usage.cost == {1: 0.5} + + _URL_CITATION: Final = { "type": "url_citation", "url_citation": { diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 1c86c4e30d8..93defc0f04f 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -26,15 +26,24 @@ from litellm.litellm_core_utils.core_helpers import ( normalize_drop_params, preserve_upstream_non_openai_attributes, process_response_headers, + provider_reported_usage_cost, qualify_agentic_followup_model, qualify_provider_stripped_model, reconstruct_model_name, redact_nested_match_and_regex_keys, + response_cost_in_hidden_params, set_provider_response_headers_in_hidden_params, unbind_budget_reservation_from_callbacks, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.utils import Delta, ImageResponse, ModelResponseStream, StreamingChoices, TranscriptionResponse +from litellm.types.utils import ( + Delta, + ImageResponse, + ModelResponse, + ModelResponseStream, + StreamingChoices, + TranscriptionResponse, +) def test_parent_otel_span_reads_dict_like_metadata_by_index(): @@ -790,6 +799,23 @@ class TestReceiveStampSurvivesASecondMetadataBucket: assert get_litellm_metadata_from_kwargs(kwargs)["litellm_received_at"] is own +@pytest.mark.parametrize("reported", [float("nan"), float("inf"), {"total_cost": float("nan")}]) +def test_a_non_finite_provider_reported_cost_counts_as_unreported(reported: object) -> None: + assert provider_reported_usage_cost({"cost": reported}, "openrouter") is None + assert provider_reported_usage_cost({"cost": {"total_cost": 0.25}}, "openrouter") == 0.25 + + +@pytest.mark.parametrize("header_cost", [float("nan"), float("inf")]) +def test_a_non_finite_response_cost_header_counts_as_unreported(header_cost: float) -> None: + response: Final = ModelResponse() + response._hidden_params = {"additional_headers": {RESPONSE_COST_HEADER: header_cost}} + finite: Final = ModelResponse() + finite._hidden_params = {"additional_headers": {RESPONSE_COST_HEADER: 0.25}} + + assert response_cost_in_hidden_params(response) is None + assert response_cost_in_hidden_params(finite) == 0.25 + + def test_billed_usage_cost_bills_any_chat_upstream_without_its_own_reader() -> None: from types import SimpleNamespace diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 88501a23235..2b1e4a5eed3 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -60,7 +60,6 @@ def test_openai_cyber_policy_uses_the_structured_error_code(shape: str, code: st assert caught.value.response.headers["x-request-id"] == "cyber-policy-request" - def _debug_info_for_api_base(configured_api_base: str) -> str: request: Final = httpx.Request("POST", "https://gateway.example/v1/chat/completions") original: Final = OpenAIError( @@ -441,6 +440,67 @@ def test_gemini_context_window_error_mapping(error_message, should_raise_context ) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("streaming", [True, False]) +@pytest.mark.parametrize( + ("error_message", "expected_type"), + [ + ( + ( + "Requested token count exceeds the model's maximum context length of 131072 tokens. " + "You requested a total of 140000 tokens: 130000 tokens from the input messages and " + "10000 tokens for the completion. Please reduce the number of tokens in the input " + "messages or the completion to fit within the limit." + ), + litellm.ContextWindowExceededError, + ), + ( + ( + "requested token count exceeds the model's maximum context length of 8192 tokens. " + "you requested a total of 9000 tokens" + ), + litellm.ContextWindowExceededError, + ), + ("max_tokens must be a positive integer", litellm.BadRequestError), + ], + ids=["sglang-context-overflow", "lowercase-context-overflow", "unrelated-bad-request"], +) +@pytest.mark.asyncio +async def test_sglang_context_window_error_mapping( + sync_mode: bool, + streaming: bool, + error_message: str, + expected_type: type[litellm.BadRequestError], +) -> None: + def reject_request(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + request=request, + json={"object": "error", "message": error_message, "type": "BadRequestError", "param": None, "code": 400}, + ) + + transport: Final = httpx.MockTransport(reject_request) + if sync_mode: + with ( + httpx.Client(transport=transport) as http_client, + OpenAI(api_key="sk-test", max_retries=0, http_client=http_client) as client, + pytest.raises(litellm.BadRequestError) as caught, + ): + _call_openai_api_sync("openai/sglang-model", "chat_completion", streaming, client) + else: + async with ( + httpx.AsyncClient(transport=transport) as async_http_client, + AsyncOpenAI(api_key="sk-test", max_retries=0, http_client=async_http_client) as async_client, + ): + with pytest.raises(litellm.BadRequestError) as caught: + await _call_openai_api_async("openai/sglang-model", "chat_completion", streaming, async_client) + + assert type(caught.value) is expected_type + assert caught.value.status_code == 400 + assert error_message in str(caught.value) + assert caught.value.llm_provider == "openai" + + def test_lemonade_context_window_error_mapping(): """Lemonade's llama.cpp backend should map context overflows to LiteLLM's standard error.""" @@ -1602,6 +1662,75 @@ def test_bedrock_mantle_400_maps_to_bad_request(): assert type(excinfo.value) is litellm.BadRequestError +class _ErrorWithoutResponse(Exception): + def __init__(self, message: str, status_code: int) -> None: + super().__init__(message) + self.message = message + self.status_code = status_code + + +def test_bedrock_mantle_context_overflow_without_a_provider_response_answers_400(): + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model="openai.gpt-5.6-sol", + original_exception=_ErrorWithoutResponse( + '{"error":{"code":"context_length_exceeded","message":"too long"}}', 400 + ), + custom_llm_provider="bedrock_mantle", + ) + + assert excinfo.value.response.status_code == 400 + + +def _invalid_model_identifier_from(url: str) -> BedrockError: + return BedrockError( + status_code=400, + message='{"message":"The provided model identifier is invalid."}', + response=httpx.Response(status_code=400, request=httpx.Request(method="POST", url=url)), + ) + + +@pytest.mark.parametrize( + "path", + ["/model/us.newvendor.typo-1/converse", "/model/us.newvendor.typo-1/converse-stream"], +) +def test_bedrock_converse_rejecting_an_id_with_no_invoke_provider_answers_404(path): + with pytest.raises(litellm.NotFoundError) as excinfo: + exception_type( + model="us.newvendor.typo-1", + original_exception=_invalid_model_identifier_from(f"https://bedrock-runtime.us-east-1.amazonaws.com{path}"), + custom_llm_provider="bedrock", + ) + + assert excinfo.value.response.status_code == 404 + + +@pytest.mark.parametrize( + "original_exception", + [ + _invalid_model_identifier_from("https://bedrock-runtime.us-east-1.amazonaws.com/model/claude-typo-9/invoke"), + _ErrorWithoutResponse("The provided model identifier is invalid.", 400), + ], + ids=["invoke", "no-response"], +) +def test_bedrock_invalid_model_identifier_not_from_converse_stays_a_bad_request(original_exception): + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type(model="claude-typo-9", original_exception=original_exception, custom_llm_provider="bedrock") + + assert excinfo.value.status_code == 400 + + +def test_bedrock_invalid_model_identifier_without_a_model_stays_a_bad_request(): + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model=None, + original_exception=_ErrorWithoutResponse("The provided model identifier is invalid.", 400), + custom_llm_provider="bedrock", + ) + + assert excinfo.value.status_code == 400 + + def test_bedrock_mantle_context_overflow_maps_to_context_window_exceeded(): from litellm.llms.base_llm.chat.transformation import BaseLLMException diff --git a/tests/unit/litellm_core_utils/test_hidden_params.py b/tests/unit/litellm_core_utils/test_hidden_params.py index c0485a07696..b32692c67ba 100644 --- a/tests/unit/litellm_core_utils/test_hidden_params.py +++ b/tests/unit/litellm_core_utils/test_hidden_params.py @@ -1,7 +1,11 @@ -from typing import Final +import pickle +from collections.abc import Callable +from typing import ClassVar, Final import pytest +from openai import BaseModel as OpenAIBaseModel from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, ConfigDict, PrivateAttr from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, @@ -9,13 +13,14 @@ from litellm.litellm_core_utils.hidden_params import ( get_hidden_params_storage, get_or_create_hidden_params, served_from_cache, + set_extra_attribute, set_hidden_param, set_hidden_params, ) from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig from litellm.types.decisions import DecisionsResponse from litellm.types.llms.base import HiddenParams -from litellm.types.utils import ModelResponse, TextChoices, TextCompletionResponse, Usage +from litellm.types.utils import ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse, Usage def test_get_and_set_hidden_params_on_plain_object() -> None: @@ -309,3 +314,133 @@ def test_served_from_cache_is_true_only_for_a_response_the_cache_stamped( assert served_from_cache(model_response) is expected assert served_from_cache({HIDDEN_PARAMS_ATTR: dict(hidden_params)}) is expected assert served_from_cache(object()) is False + + +class _ExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow") + + declared: str | None = None + marker: ClassVar[str | None] = None + unset_marker: ClassVar[int] + _secret: str = PrivateAttr(default="hidden") + + @property + def doubled(self) -> str | None: + return self.declared + + @doubled.setter + def doubled(self, value: str) -> None: + self.declared = value * 2 + + +class _FrozenExtraModel(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + +class _ValidatedExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow", validate_assignment=True) + + __pydantic_extra__: dict[str, int] + + +class _IgnoredExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="ignore") + + +class _MirroringModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow") + + def __setattr__(self, name: str, value: object) -> None: + super().__setattr__(name, value) + super().__setattr__(f"{name}_copy", value) + + +class _NoneClassAttribute: + inherited = None + + +class _InheritingExtraModel(_NoneClassAttribute, BaseModel): # frozen-ok: assignment tests need attribute writes + model_config = ConfigDict(extra="allow") + + +class _ForbiddingOpenAIModel(OpenAIBaseModel): # frozen-ok: assignment tests need attribute writes + model_config = ConfigDict(extra="forbid") + + +def _assignment_outcome(model: BaseModel, assign: Callable[[], None], name: str) -> tuple[object, ...]: + try: + assign() + except (AttributeError, ValueError) as error: + return type(error), str(error) + return model.model_dump(), getattr(model, name), model.model_fields_set, model.model_dump_json() + + +@pytest.mark.parametrize( + "build, name, value", + [ + pytest.param(lambda: _ExtraModel(existing=1, later=3), "fresh", 1, id="new-extra"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "existing", 2, id="existing-extra-keeps-order"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "declared", "x", id="declared-field"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "doubled", "ab", id="property"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "_free", 1, id="undeclared-underscore-name"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "_secret", "x", id="private-attribute"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "marker", 1, id="class-var"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "unset_marker", 1, id="class-var-without-value"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "model_config", 1, id="class-attribute"), + pytest.param(_FrozenExtraModel, "fresh", 1, id="frozen"), + pytest.param(_ValidatedExtraModel, "fresh", "5", id="validate-assignment"), + pytest.param(_IgnoredExtraModel, "fresh", 1, id="extra-ignore"), + pytest.param(_MirroringModel, "fresh", 1, id="setattr-override"), + pytest.param(_InheritingExtraModel, "inherited", 1, id="class-attribute-set-to-none"), + pytest.param( + lambda: _ForbiddingOpenAIModel.construct(existing=1), "fresh", 1, id="extra-forbid-built-by-construct" + ), + ], +) +def test_set_extra_attribute_matches_setattr(build: Callable[[], BaseModel], name: str, value: object) -> None: + through_setattr: Final = build() + through_helper: Final = build() + + expected: Final = _assignment_outcome(through_setattr, lambda: setattr(through_setattr, name, value), name) + actual: Final = _assignment_outcome(through_helper, lambda: set_extra_attribute(through_helper, name, value), name) + + assert actual == expected + + +@pytest.mark.parametrize("response_type", [ModelResponse, ModelResponseStream]) +def test_hidden_params_raises_attribute_error_once_deleted( + response_type: type[ModelResponse | ModelResponseStream], +) -> None: + response: Final = response_type() + del response._hidden_params + + assert getattr(response, "hidden_params", "fallback") == "fallback" + + +def _round_trip(response: ModelResponse | ModelResponseStream) -> ModelResponse | ModelResponseStream: + return pickle.loads(pickle.dumps(response)) + + +@pytest.mark.parametrize("response_type", [ModelResponse, ModelResponseStream]) +@pytest.mark.parametrize( + "rebuild", + [ + pytest.param(lambda response: response, id="init"), + pytest.param(lambda response: response.model_copy(), id="model-copy"), + pytest.param(lambda response: response.model_copy(deep=True), id="deep-copy"), + pytest.param(_round_trip, id="pickle"), + ], +) +def test_hidden_params_is_the_stored_dict_on_every_construction_path( + response_type: type[ModelResponse | ModelResponseStream], + rebuild: Callable[[ModelResponse | ModelResponseStream], ModelResponse | ModelResponseStream], +) -> None: + original: Final = response_type() + original.hidden_params["model_id"] = "deployment-1" + response: Final = rebuild(original) + + response.hidden_params["api_base"] = "https://example.com" + + assert response.hidden_params is response._hidden_params + assert get_hidden_params(response) == response._hidden_params + assert response._hidden_params["api_base"] == "https://example.com" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index b3ae3bf6240..e952127a32c 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -59,6 +59,7 @@ from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandle from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, CostBreakdown, @@ -4514,6 +4515,105 @@ def test_logging_preserves_missing_usage_without_accepting_request_metadata( assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (0, 0, 0) +@pytest.mark.parametrize( + ("meta", "expected"), + [ + ({"tokens": {"input_tokens": 7, "output_tokens": 3}}, (7, 3, 10)), + ({"tokens": {}, "billed_units": {"total_tokens": 42}}, (42, 0, 42)), + ({"tokens": {"output_tokens": 3}, "billed_units": {"total_tokens": 12}}, (9, 3, 12)), + ({"tokens": {"input_tokens": 7, "output_tokens": 3}, "billed_units": {"total_tokens": 0}}, (7, 3, 0)), + ({"billed_units": {"input_tokens": 7, "output_tokens": 3}}, (7, 3, 10)), + ( + { + "tokens": {"input_tokens": 7, "output_tokens": 3}, + "billed_units": {"input_tokens": 90, "output_tokens": 10, "total_tokens": 12}, + }, + (7, 3, 12), + ), + ( + { + "tokens": {"input_tokens": 0, "output_tokens": 0}, + "billed_units": {"input_tokens": 90, "output_tokens": 10, "total_tokens": 12}, + }, + (0, 0, 12), + ), + ( + { + "tokens": {"input_tokens": None, "output_tokens": None}, + "billed_units": {"input_tokens": 7, "output_tokens": 3, "total_tokens": None}, + }, + (7, 3, 10), + ), + ({"tokens": {"input_tokens": 0}}, (0, 0, 0)), + (None, None), + ({}, None), + ({"tokens": None, "billed_units": None}, None), + ({"billed_units": {"search_units": 3}}, None), + ], +) +def test_rerank_metadata_reaches_standard_and_spend_logs( + logging_obj: Logging, meta: dict[str, object] | None, expected: tuple[int, int, int] | None +) -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + native: Final = {"id": "rerank-usage", "results": [{"index": 0, "relevance_score": 0.9}], "meta": meta} + response: Final = RerankResponse.model_validate(native) + counts: Final = expected if expected is not None else (0, 0, 0) + expected_usage: Final = dict(zip(("prompt_tokens", "completion_tokens", "total_tokens"), counts)) + assert response.model_dump() == native + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response.model_dump()) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == counts + assert StandardLoggingPayloadSetup.get_usage_as_dict(response.model_dump()) == ( + expected_usage if expected is not None else {} + ) + kwargs: Final = {"model": "rerank-model", "call_type": "arerank", "litellm_params": {"metadata": {}}} + standard: Final = get_standard_logging_object_payload( + kwargs=kwargs, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + assert standard is not None + assert (standard["prompt_tokens"], standard["completion_tokens"], standard["total_tokens"]) == counts + assert standard["metadata"]["usage_object"] == (expected_usage if expected is not None else {}) + spend: Final = get_logging_payload( + kwargs={**kwargs, "standard_logging_object": standard}, response_obj=response, start_time=now, end_time=now + ) + assert (spend["prompt_tokens"], spend["completion_tokens"], spend["total_tokens"]) == counts + + +@pytest.mark.parametrize("combined", [None, Usage(prompt_tokens=11, completion_tokens=2, total_tokens=13)]) +@pytest.mark.parametrize("explicit_usage", [{}, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}]) +def test_rerank_metadata_does_not_override_explicit_or_combined_usage( + combined: Usage | None, explicit_usage: dict[str, int] +) -> None: + response: Final = {"results": [], "meta": {"tokens": {"input_tokens": 42}}, "usage": explicit_usage} + expected: Final = combined.model_dump() if combined is not None else explicit_usage + assert StandardLoggingPayloadSetup.get_usage_as_dict(response, combined) == expected + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response, combined) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + expected.get("prompt_tokens", 0), + expected.get("completion_tokens", 0), + expected.get("total_tokens", 0), + ) + + +@pytest.mark.parametrize( + "response", + [ + {"meta": {"tokens": {"input_tokens": 42}}}, + {"results": [], "meta": {"tokens": {"input_tokens": "invalid"}}}, + ], +) +def test_rerank_metadata_ignores_unrelated_and_invalid_responses(response: dict[str, object]) -> None: + assert StandardLoggingPayloadSetup.get_usage_as_dict(response) == {} + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0) + + def test_append_system_prompt_messages(): """ Test append_system_prompt_messages prepends system message from kwargs to messages list. @@ -6308,6 +6408,60 @@ def test_standard_logging_payload_resolves_kubernetes_pod_routing_from_metadata_ assert payload["metadata"][KUBERNETES_POD_ROUTING_KEY] == expected +@pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), float("-inf")]) +def test_get_standard_logging_object_payload_logs_a_non_finite_response_cost_as_zero( + logging_obj: LitellmLogging, response_cost: float +) -> None: + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload + + now: Final = datetime(2026, 1, 1) + payload: Final = get_standard_logging_object_payload( + kwargs={"model": "gpt-4o", "messages": [], "litellm_params": {}, "response_cost": response_cost}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["response_cost"] == 0 + + +def test_an_overflowing_custom_price_records_a_zero_response_cost(caplog: pytest.LogCaptureFixture) -> None: + messages: Final = [{"role": "user", "content": "hello world"}] + start_time: Final = datetime.datetime.now() + logging_obj: Final = Logging( + model="gpt-4o", + messages=messages, + stream=False, + call_type="completion", + start_time=start_time, + litellm_call_id="overflow\nforged log line", + function_id="overflow", + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + response: Final = litellm.completion( + model="gpt-4o", + mock_response="hi there", + messages=messages, + input_cost_per_token=1e308, + output_cost_per_token=1e308, + litellm_logging_obj=logging_obj, + ) + logging_obj._process_hidden_params_and_response_cost( + response, start_time=start_time, end_time=datetime.datetime.now() + ) + + assert response._hidden_params["response_cost"] == 0.0 + payload: Final = logging_obj.model_call_details["standard_logging_object"] + assert json.loads(json.dumps(payload, allow_nan=False, default=str))["response_cost"] == 0.0 + warnings: Final = [r.getMessage() for r in caplog.records if "non-finite response_cost" in r.getMessage()] + assert warnings == ["Recording a non-finite response_cost as 0 for litellm_call_id=overflowforged log line"] + + def test_get_standard_logging_object_payload_carries_matched_access_groups(logging_obj): """Access groups stamped at auth time reach the logging payload, so integrations see what a request billed.""" from datetime import datetime @@ -6902,6 +7056,95 @@ def test_handle_anthropic_messages_response_logging_reparses_the_served_body_wit assert logged.usage.completion_tokens == 7 # type: ignore[attr-defined] +@pytest.mark.parametrize( + "compaction_block", + [ + {"type": "compaction", "content": "Summary of the chat so far.", "signature": "sig-abc"}, + {"type": "compaction", "signature": "sig-abc"}, + ], + ids=["with_summary", "without_content"], +) +def test_handle_anthropic_messages_response_logging_reparses_a_served_compaction_response( + compaction_block: dict[str, str], +): + logging_obj = _anthropic_messages_logging_obj() + logging_obj.optional_params = {} + + logged = logging_obj._handle_anthropic_messages_response_logging( + result={ + "id": "msg_compacted", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [compaction_block], + "stop_reason": "compaction", + "stop_sequence": None, + "usage": {"input_tokens": 410, "output_tokens": 95}, + } + ) + + assert isinstance(logged, ModelResponse) + assert logged.id == "msg_compacted" + assert logged.usage.prompt_tokens == 410 + assert logged.usage.completion_tokens == 95 + assert logged.choices[0].message.provider_specific_fields["compaction_blocks"] == [compaction_block] + + +_WEB_SEARCH_REPLY: Final = ( + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "weather"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [{"type": "web_search_result", "url": "https://example.com", "title": "Weather"}], + }, + {"type": "text", "text": "It is sunny."}, +) + + +@pytest.mark.parametrize( + ("content", "stop_reason"), + [ + (list(_WEB_SEARCH_REPLY), "end_turn"), + (list(_WEB_SEARCH_REPLY), "pause_turn"), + ([{"type": "text", "text": "The context window is full."}], "model_context_window_exceeded"), + ( + [ + {"type": "mcp_tool_use", "id": "mcptoolu_1", "name": "lookup", "server_name": "docs", "input": {}}, + {"type": "mcp_tool_result", "tool_use_id": "mcptoolu_1", "is_error": False, "content": [{"type": "text", "text": "ok"}]}, + {"type": "text", "text": "Found it."}, + ], + "end_turn", + ), + ([{"type": "container_upload", "file_id": "file_1"}, {"type": "text", "text": "Uploaded."}], "end_turn"), + ], + ids=["web_search", "pause_turn", "model_context_window_exceeded", "mcp_tool_use", "container_upload"], +) +def test_handle_anthropic_messages_response_logging_bills_a_served_server_tool_response( + content: list[dict[str, object]], stop_reason: str +): + logging_obj = _anthropic_messages_logging_obj() + logging_obj.optional_params = {} + + logged = logging_obj._handle_anthropic_messages_response_logging( + result={ + "id": "msg_server_tools", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": content, + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": 1200, "output_tokens": 80}, + } + ) + + assert isinstance(logged, ModelResponse) + assert logged.id == "msg_server_tools" + assert logged.usage.prompt_tokens == 1200 + assert logged.usage.completion_tokens == 80 + assert logged.choices[0].message.content == content[-1]["text"] + + def test_handle_anthropic_messages_response_logging_passes_model_response_through(): """Anthropic-native path already yields a ModelResponse; it must be returned unchanged.""" logging_obj = _anthropic_messages_logging_obj() @@ -12431,9 +12674,10 @@ def test_get_hidden_params(): assert result["api_base"] is None assert result["response_cost"] is None assert result["additional_headers"] is None + assert result.get("upstream_provider") is None # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__ if field != "upstream_provider") # Test with valid params hidden_params = { @@ -12441,6 +12685,7 @@ def test_get_hidden_params(): "cache_key": "test-cache", "api_base": "https://api.test.com", "response_cost": 0.001, + "upstream_provider": "serving-site-a", "additional_headers": { "x-ratelimit-limit-requests": "2000", "x-ratelimit-remaining-requests": "1999", @@ -12453,8 +12698,9 @@ def test_get_hidden_params(): assert result["response_cost"] == 0.001 assert result["additional_headers"] is not None assert result["additional_headers"]["x_ratelimit_limit_requests"] == 2000 + assert result.get("upstream_provider") == "serving-site-a" # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__ if field != "upstream_provider") @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") diff --git a/tests/unit/litellm_core_utils/test_llm_request_utils.py b/tests/unit/litellm_core_utils/test_llm_request_utils.py index e7350fae395..5243d308543 100644 --- a/tests/unit/litellm_core_utils/test_llm_request_utils.py +++ b/tests/unit/litellm_core_utils/test_llm_request_utils.py @@ -1,3 +1,5 @@ +import datetime +import time from typing import Final import httpx @@ -245,3 +247,39 @@ def test_rank_chat_models_flags_a_retired_model_the_provider_lists_without_its_c (ranked_model.model, ranked_model.is_past_deprecation_date) for ranked_model in rank_chat_models_from_llm_provider("vertex_ai") ] == [("fake-retired-vertex-chat", True), ("fake-live-vertex-chat", False)] + + +@pytest.mark.parametrize("host_zone", ["Etc/GMT+12", "Etc/GMT-14"]) +def test_rank_chat_models_flags_retirement_by_the_utc_date_whatever_the_host_zone( + monkeypatch: pytest.MonkeyPatch, host_zone: str +) -> None: + utc_today: Final = datetime.datetime.now(datetime.timezone.utc).date() + retirement_dates: Final = { + "deepseek/fake-retires-today-chat": utc_today, + "deepseek/fake-retires-tomorrow-chat": utc_today + datetime.timedelta(days=1), + } + for name, retirement_date in retirement_dates.items(): + monkeypatch.setitem( + litellm.model_cost, + name, + { + "input_cost_per_token": 1e-12, + "output_cost_per_token": 1e-12, + "mode": "chat", + "litellm_provider": PICKER_PROVIDER, + "deprecation_date": retirement_date.isoformat(), + }, + ) + monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, tuple(retirement_dates)) + monkeypatch.setenv("TZ", host_zone) + time.tzset() + try: + flags: Final = { + ranked_model.model: ranked_model.is_past_deprecation_date + for ranked_model in rank_chat_models_from_llm_provider(PICKER_PROVIDER) + } + finally: + monkeypatch.undo() + time.tzset() + + assert flags == {"deepseek/fake-retires-today-chat": True, "deepseek/fake-retires-tomorrow-chat": False} diff --git a/tests/unit/litellm_core_utils/test_provider_affinity.py b/tests/unit/litellm_core_utils/test_provider_affinity.py index edf4f5169b5..6ccd6b17c01 100644 --- a/tests/unit/litellm_core_utils/test_provider_affinity.py +++ b/tests/unit/litellm_core_utils/test_provider_affinity.py @@ -1,12 +1,16 @@ +from collections.abc import Callable + import pytest from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.litellm_core_utils.provider_affinity import ( add_provider_affinity_header, + get_caller_session_id, get_stable_session_id, ) +@pytest.mark.parametrize("resolve", [get_stable_session_id, get_caller_session_id]) @pytest.mark.parametrize( ("litellm_params", "expected"), [ @@ -16,18 +20,22 @@ from litellm.litellm_core_utils.provider_affinity import ( ({"litellm_metadata": {"session_id": "litellm-metadata-session"}}, "litellm-metadata-session"), ], ) -def test_get_stable_session_id_uses_explicit_session_sources(litellm_params: dict, expected: str): - assert get_stable_session_id(litellm_params) == expected +def test_session_ids(litellm_params: dict[str, object], expected: str, resolve: Callable[[object], str | None]) -> None: + assert resolve(litellm_params) == expected def test_get_stable_session_id_does_not_use_trace_id(): assert get_stable_session_id({"litellm_trace_id": "per-request-trace"}) is None + assert get_caller_session_id( + {"litellm_session_id": "trace", "proxy_server_request": {"headers": {"x-litellm-trace-id": "trace"}}} + ) is None +@pytest.mark.parametrize("resolve", [get_stable_session_id, get_caller_session_id]) @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) -def test_get_stable_session_id_ignores_proxy_generated_session(metadata_key: str): +def test_generated_sessions_are_ignored(metadata_key: str, resolve: Callable[[object], str | None]) -> None: assert ( - get_stable_session_id( + resolve( { "litellm_session_id": "generated-session", metadata_key: { diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 5b64e42d8a8..9e385c588ae 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -57,9 +57,10 @@ def _provider_fields() -> dict[str, object]: "thinking_blocks": [{"type": "thinking", "thinking": SECRET}], "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], "citations": [{"cited_text": SECRET}], - "thought_signatures": [SIGNATURE], - "compaction_blocks": [{"type": "compaction", "content": SECRET}], + "compaction_blocks": [{"type": "compaction", "content": SECRET, "signature": "compaction-signature"}], + "compaction_start": {"type": "compaction", "content": SECRET}, "compaction_delta": {"type": "compaction_delta", "content": SECRET}, + "thought_signatures": [SIGNATURE], } diff --git a/tests/unit/litellm_core_utils/test_safe_json_dumps.py b/tests/unit/litellm_core_utils/test_safe_json_dumps.py index 1f2664f33cb..1b8bdd21326 100644 --- a/tests/unit/litellm_core_utils/test_safe_json_dumps.py +++ b/tests/unit/litellm_core_utils/test_safe_json_dumps.py @@ -169,6 +169,31 @@ def test_null_byte_stripped_in_fallback_str(): assert json.loads(out)["obj"] == "objrepr" +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_non_finite_floats_are_written_as_json_null(value: float) -> None: + out = safe_dumps( + {"cost_breakdown": {"total_cost": value, "output_cost": 0.006}, "costs": [value]}, non_finite_as_null=True + ) + + assert json.loads(out, parse_constant=lambda token: pytest.fail(f"non-JSON token {token}")) == { + "cost_breakdown": {"total_cost": None, "output_cost": 0.006}, + "costs": [None], + } + + +def test_a_set_holding_a_non_finite_float_sorts_its_finite_members_and_puts_null_last() -> None: + assert json.loads(safe_dumps({"costs": {2.0, float("inf"), 1.0}}, non_finite_as_null=True)) == { + "costs": [1.0, 2.0, None] + } + + +def test_non_finite_floats_keep_their_value_unless_the_caller_asks_for_null() -> None: + assert json.loads(safe_dumps({"timeout": float("inf"), "costs": [float("-inf")]})) == { + "timeout": float("inf"), + "costs": [float("-inf")], + } + + def test_clean_strings_are_not_run_through_replace(): """Regression for LIT-3910. diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 28587bf4749..bbf1897b3a7 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -910,6 +910,71 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" +def test_stream_chunk_builder_finds_upstream_provider_in_later_chunk_regression() -> None: + first: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + later: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" world", role=None), + ) + ], + ) + first_hidden_params: Final = {"origin": "first"} + later_hidden_params: Final = {"upstream_provider": "serving-site-a"} + chunks: Final = [ + {**first.model_dump(), "_hidden_params": first_hidden_params}, + {**later.model_dump(), "_hidden_params": later_hidden_params}, + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + assert first_hidden_params.get("upstream_provider") is None + assert later_hidden_params.get("upstream_provider") == "serving-site-a" + + +def test_stream_chunk_builder_preserves_source_provider_metadata_guard() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + hidden_params: Final = {"upstream_provider": "serving-site-a"} + chunk_dict: Final = {**chunk.model_dump(), "_hidden_params": hidden_params} + + response: Final = stream_chunk_builder(chunks=[chunk_dict]) + + assert response is not None + assert chunk_dict["_hidden_params"]["upstream_provider"] == "serving-site-a" + + def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): """Vertex AI metadata on streaming chunks must appear on assembled response.""" grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] @@ -1077,6 +1142,71 @@ def test_cost_field_in_usage_chunks(): assert usage.completion_tokens == 5 +def test_the_rebuilt_response_keeps_the_upstream_cost_a_stamped_usage_chunk_saved(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.5}) + text_chunk = ModelResponseStream( + id="chatcmpl-1", + model="gpt-4.1", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="Hi"))], + ) + usage_chunk = ModelResponseStream( + id="chatcmpl-1", + model="gpt-4.1", + choices=[], + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.0015), + ) + usage_chunk._hidden_params = {"response_cost": 0.0015, "upstream_usage_cost": 0.001} + + rebuilt = stream_chunk_builder([text_chunk, usage_chunk], messages=[{"role": "user", "content": "hi"}]) + + assert rebuilt is not None + assert rebuilt._hidden_params["upstream_usage_cost"] == 0.001 + assert litellm.completion_cost(rebuilt, custom_llm_provider="openai") == pytest.approx(0.0015), ( + "the upstream's 0.001 gets the 50% margin once, litellm's 0.0015 bill is not margined again" + ) + + + +def _perplexity_usage_chunk(cost: Mapping[str, float], finish_reason: str | None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-pplx", + created=1745513206, + model="sonar", + choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content="ok"))], + usage={"prompt_tokens": 26, "completion_tokens": 1, "total_tokens": 27, "cost": dict(cost)}, + ) + + +@pytest.mark.parametrize("poisoned_field", ["request_cost", "total_cost"]) +def test_a_non_finite_final_stream_cost_does_not_fall_back_to_an_earlier_chunk_cost(poisoned_field: str): + final_cost: Final = {"input_tokens_cost": 0.0012, "output_tokens_cost": 1e-05, "request_cost": 0.0025} + chunks: Final = [ + _perplexity_usage_chunk({"input_tokens_cost": 0, "output_tokens_cost": 0, "total_cost": 0}, None), + _perplexity_usage_chunk({**final_cost, "total_cost": 0.00371, poisoned_field: float("nan")}, "stop"), + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert getattr(response.usage, "cost", None) is None + assert response.usage.total_tokens == 27 + + +@pytest.mark.parametrize("provider_key", ["dropped_non_finite_cost", "_dropped_non_finite_cost"]) +def test_a_provider_usage_key_named_like_the_dropped_cost_flag_keeps_the_stream_cost(provider_key: str): + chunk: Final = ModelResponseStream( + id="chatcmpl-flag-key", + created=1745513206, + model="sonar", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="ok"))], + usage={"prompt_tokens": 26, "completion_tokens": 1, "total_tokens": 27, "cost": 0.00371, provider_key: True}, + ) + + response: Final = stream_chunk_builder(chunks=[chunk]) + + assert response is not None + assert getattr(response.usage, "cost", None) == 0.00371 + def test_stream_chunk_builder_tolerates_trailing_chunk_without_choices(): """Regression for https://github.com/BerriAI/litellm/issues/32051 diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b9c8feb223d..4d1699c7fe4 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -11,7 +11,7 @@ import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from types import CodeType, FrameType, MappingProxyType +from types import CodeType, FrameType, MappingProxyType, TracebackType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -6292,53 +6292,76 @@ def test_openai_chunk_creator_skips_full_chunk_dump_and_delta_rebuild( assert finish_output.choices[0].delta.content is None -def _openai_chunk(chunk_id: str, delta: Mapping[str, object]) -> ChatCompletionChunk: +def _openai_chunk( + chunk_id: str, + delta: Mapping[str, object], + choice_extras: Mapping[str, object] = MappingProxyType({}), + finish_reason: str | None = None, + chunk_extras: Mapping[str, object] = MappingProxyType({}), +) -> ChatCompletionChunk: return ChatCompletionChunk( id=chunk_id, created=1, model="gpt-4o", object="chat.completion.chunk", - choices=[{"index": 0, "delta": delta, "finish_reason": None}], + choices=[{"index": 0, "delta": delta, "finish_reason": finish_reason, **choice_extras}], + **chunk_extras, ) -def _missed_attributes(wrapper: CustomStreamWrapper, chunks: Sequence[ChatCompletionChunk]) -> tuple[str, ...]: - attribute_lookup: Final = BaseModel.__getattr__.__code__ - misses: Final[list[str]] = [] +def _raised_attribute_errors(wrapper: CustomStreamWrapper, chunks: Sequence[ChatCompletionChunk]) -> tuple[str, ...]: + raised: Final[list[str]] = [] - def record_miss(frame: FrameType, event: str, arg: object) -> Callable[[FrameType, str, object], object]: - if event == "exception" and isinstance(arg, tuple) and isinstance(arg[1], AttributeError): - misses.append(str(frame.f_locals["item"])) - return record_miss - - def trace_lookups(frame: FrameType, event: str, _arg: object) -> Callable[[FrameType, str, object], object] | None: - return record_miss if event == "call" and frame.f_code is attribute_lookup else None + def record_raise(frame: FrameType, event: str, arg: object) -> Callable[[FrameType, str, object], object]: + if event == "exception" and isinstance(arg, tuple): + _, error, traceback = arg + if isinstance(error, AttributeError) and isinstance(traceback, TracebackType) and traceback.tb_next is None: + raised.append(f"{frame.f_code.co_qualname}: {error}") + return record_raise previous_tracer: Final = sys.gettrace() - sys.settrace(trace_lookups) + sys.settrace(record_raise) try: tuple(wrapper.chunk_creator(chunk=chunk) for chunk in chunks) finally: sys.settrace(previous_tracer) - return tuple(misses) + return tuple(raised) -def test_openai_chunk_creator_builds_one_delta_and_misses_no_attribute_per_later_chunk( - logging_obj: Logging, +_AZURE_CHOICE_EXTRAS: Final = MappingProxyType( + {"content_filter_results": {"hate": {"filtered": False, "severity": "safe"}}} +) +_OPENAI_AND_AZURE_CHOICE_EXTRAS: Final = ( + pytest.param("openai", MappingProxyType({}), id="openai"), + pytest.param("azure", _AZURE_CHOICE_EXTRAS, id="azure"), +) + + +@pytest.mark.parametrize(("custom_llm_provider", "choice_extras"), _OPENAI_AND_AZURE_CHOICE_EXTRAS) +def test_openai_chunk_creator_builds_one_delta_and_raises_no_attribute_error_per_later_chunk( + logging_obj: Logging, custom_llm_provider: str, choice_extras: Mapping[str, object] ) -> None: wrapper: Final = CustomStreamWrapper( completion_stream=iter(()), model="gpt-4o", - custom_llm_provider="openai", + custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) - wrapper.chunk_creator(chunk=_openai_chunk("chunk-1", {"role": "assistant", "content": ""})) - plain_chunks: Final = ( - _openai_chunk("chunk-2", {"content": "hello"}), - _openai_chunk("chunk-3", {"content": " world"}), + + def later_chunk(chunk_id: str, delta: Mapping[str, object]) -> ChatCompletionChunk: + return _openai_chunk( + chunk_id, delta, choice_extras, chunk_extras={"service_tier": "default", "obfuscation": chunk_id} + ) + + nested_delta: Final = {"content": "!", "reasoning_details": [{"type": "reasoning.text", "text": "thought"}]} + first_chunk_of_each_shape: Final = ( + later_chunk("chunk-1", {"role": "assistant", "content": ""}), + later_chunk("chunk-2", nested_delta), ) - nested_chunk: Final = _openai_chunk( - "chunk-4", {"content": "!", "reasoning_details": [{"type": "reasoning.text", "text": "thought"}]} + tuple(wrapper.chunk_creator(chunk=chunk) for chunk in first_chunk_of_each_shape) + plain_chunks: Final = ( + later_chunk("chunk-3", {"content": "hello"}), + later_chunk("chunk-4", {"content": " world"}), ) delta_init: Final = Delta.__init__.__code__ streaming_choices_init: Final = StreamingChoices.__init__.__code__ @@ -6348,7 +6371,145 @@ def test_openai_chunk_creator_builds_one_delta_and_misses_no_attribute_per_later assert sum(call.code is delta_init for call in calls) == len(plain_chunks) assert sum(call.code is streaming_choices_init for call in calls) == len(plain_chunks) assert [output.choices[0].delta.content if output else None for output in outputs] == ["hello", " world"] - assert _missed_attributes(wrapper, (*plain_chunks, nested_chunk)) == () + assert _raised_attribute_errors(wrapper, (*plain_chunks, later_chunk("chunk-5", nested_delta))) == () + + +_OPENAI_STREAM_JSON: Final = ( + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":"","role":"assistant",' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o0"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":"one","role":null,' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o1"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":" two","role":null,' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o2"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":"stop","index":0,' + '"delta":{"provider_specific_fields":null,"content":null,"role":null,"function_call":null,' + '"tool_calls":null,"audio":null},"logprobs":null}],"provider_specific_fields":null,' + '"service_tier":"default","obfuscation":"o3"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[],"provider_specific_fields":null,' + '"usage":{"completion_tokens":3,"prompt_tokens":9,"total_tokens":12,"completion_tokens_details":null,' + '"prompt_tokens_details":null}}' + ), +) +_AZURE_STREAM_JSON: Final = ( + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":"","role":"assistant","function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o0"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":"one","role":null,"function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o1"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":" two","role":null,"function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o2"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":"stop","index":0,' + '"delta":{"provider_specific_fields":null,"content":null,"role":null,"function_call":null,' + '"tool_calls":null,"audio":null},"logprobs":null}],"provider_specific_fields":null,' + '"service_tier":"default","obfuscation":"o3"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[],"provider_specific_fields":null,' + '"usage":{"completion_tokens":3,"prompt_tokens":9,"total_tokens":12,"completion_tokens_details":null,' + '"prompt_tokens_details":null}}' + ), +) + + +@pytest.mark.parametrize( + ("custom_llm_provider", "choice_extras", "expected_json"), + [ + pytest.param("openai", MappingProxyType({}), _OPENAI_STREAM_JSON, id="openai"), + pytest.param("azure", _AZURE_CHOICE_EXTRAS, _AZURE_STREAM_JSON, id="azure"), + ], +) +def test_openai_chunk_creator_emits_each_chunk_as_the_same_json( + logging_obj: Logging, + custom_llm_provider: str, + choice_extras: Mapping[str, object], + expected_json: tuple[str | None, ...], +) -> None: + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + chunks: Final = ( + _openai_chunk( + "chatcmpl-1", + {"role": "assistant", "content": ""}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o0"}, + ), + _openai_chunk( + "chatcmpl-1", + {"content": "one"}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o1"}, + ), + _openai_chunk( + "chatcmpl-1", + {"content": " two"}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o2"}, + ), + _openai_chunk( + "chatcmpl-1", {}, choice_extras, "stop", chunk_extras={"service_tier": "default", "obfuscation": "o3"} + ), + ChatCompletionChunk( + id="chatcmpl-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[], + service_tier="default", + obfuscation="o4", + usage={"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}, + ), + ) + + outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in chunks) + + assert tuple( + output.model_copy(update={"created": 1}).model_dump_json() if output else None for output in outputs + ) == expected_json def test_chunk_creator_skips_full_dump_of_model_response_stream_chunks( @@ -6686,8 +6847,39 @@ def test_the_usage_chunk_carries_the_bill_the_builder_priced_for_the_spend_row() assert usage_chunk.usage is not None and usage_chunk.usage.cost == 0.00032, "the client sees what the row bills" +@pytest.mark.parametrize("upstream_cost", [None, 0.001]) +def test_rebuilding_a_priced_stream_bills_the_margin_once(monkeypatch: pytest.MonkeyPatch, upstream_cost: float | None) -> None: + from litellm.litellm_core_utils.core_helpers import stamp_litellm_bill_on_usage + from litellm.litellm_core_utils.streaming_handler import ( + _stamp_usage_cost, # pyright: ignore[reportPrivateUsage] # the helper under test + ) + + monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.5}) + model: Final = "anthropic/claude-sonnet-4-5" + shared_usage: Final = Usage(prompt_tokens=9, completion_tokens=16, total_tokens=25) + if upstream_cost is not None: + setattr(shared_usage, "cost", upstream_cost) + reference: Final = ModelResponse(model=model, usage=shared_usage.model_copy()) + billed_once: Final = litellm.completion_cost(reference, custom_llm_provider="anthropic") + complete: Final = ModelResponse(model=model, usage=shared_usage) + stamp_litellm_bill_on_usage(complete, shared_usage, billed_once) + text_chunk: Final = ModelResponseStream( + id="chunk-1", model=model, choices=[StreamingChoices(index=0, delta=Delta(content="hi"), finish_reason="stop")] + ) + usage_chunk: Final = ModelResponseStream(id="chunk-1", model=model, choices=[], usage=shared_usage) + + _stamp_usage_cost(usage_chunk, complete) + rebuilt: Final = litellm.stream_chunk_builder([text_chunk, usage_chunk], messages=[{"role": "user", "content": "hi"}]) + + assert rebuilt is not None and rebuilt.usage is not None # pyright: ignore[reportAttributeAccessIssue] # ModelResponse carries usage at runtime + assert usage_chunk.usage is not None and usage_chunk.usage.cost == pytest.approx(billed_once) + assert litellm.completion_cost(rebuilt, custom_llm_provider="anthropic") == pytest.approx(billed_once), ( + "pricing the rebuilt stream must not margin litellm's own bill again" + ) + + def test_chunk_creator_keeps_provider_model_private_across_stream(): - from litellm.router_utils.add_retry_fallback_headers import ( + from litellm.litellm_core_utils.add_retry_fallback_headers import ( get_hidden_params_dict, ) @@ -6748,7 +6940,7 @@ def test_chunk_creator_keeps_provider_model_private_across_stream(): def test_assembled_stream_uses_later_provider_model_for_cost( monkeypatch: pytest.MonkeyPatch, ): - from litellm.router_utils.add_retry_fallback_headers import ( + from litellm.litellm_core_utils.add_retry_fallback_headers import ( get_hidden_params_dict, ) diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index d815e1eedd3..c2ee58b67ed 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -6,6 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic """ import json +import logging import re from typing import Any, Literal, Optional, Final from unittest.mock import MagicMock, patch @@ -3391,3 +3392,18 @@ class TestAnthropicMessagesHandlerPostCallHookResponse: native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]} assert AnthropicMessagesHandler().post_call_hook_response(native) is native + + +def test_reading_a_dropped_budget_for_a_guardrail_logs_no_drop_warning(caplog: pytest.LogCaptureFixture) -> None: + data: Final = { + "model": "gemini-2.5-flash", + "max_tokens": 64, + "thinking": {"type": "enabled", "budget_tokens": "high"}, + "messages": [{"role": "user", "content": "hi"}], + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + messages: Final = AnthropicMessagesHandler().get_structured_messages(data) + + assert messages == [{"role": "user", "content": "hi"}] + assert not [record for record in caplog.records if "Dropping" in record.getMessage()] diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 4d5d6cc17a9..cbb60caa57b 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1,5 +1,6 @@ import copy import json +import logging from itertools import chain from typing import Final from unittest.mock import MagicMock, patch @@ -3189,6 +3190,47 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model( ) +def _reasoning_effort_drop_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + record.getMessage() for record in caplog.records if "Dropping reasoning_effort=" in record.getMessage() + ) + + +@pytest.mark.parametrize("reasoning_effort_value", [5, {"effort": 5, "summary": "concise"}]) +def test_reasoning_effort_that_is_not_a_level_name_is_dropped_with_one_warning( + reasoning_effort_value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = AnthropicConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort_value}, + optional_params={}, + model="claude-sonnet-4-6-20260219", + drop_params=False, + ) + + warnings: Final = _reasoning_effort_drop_warnings(caplog) + assert "thinking" not in result + assert "output_config" not in result + assert len(warnings) == 1 + assert f"reasoning_effort={reasoning_effort_value!r}" in warnings[0] + assert "claude-sonnet-4-6-20260219" in warnings[0] + + +@pytest.mark.parametrize("reasoning_effort_value", [{"summary": "concise"}, {"effort": "low", "summary": "concise"}]) +def test_reasoning_effort_with_a_level_name_or_no_level_logs_no_drop( + reasoning_effort_value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + AnthropicConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort_value}, + optional_params={}, + model="claude-sonnet-4-6-20260219", + drop_params=False, + ) + + assert _reasoning_effort_drop_warnings(caplog) == () + + @pytest.mark.parametrize( "model,budget_tokens,expected", [ diff --git a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py index c30644e7fdc..e0882653324 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py @@ -1,6 +1,7 @@ import base64 import functools import json +import logging from collections.abc import Mapping from types import MappingProxyType from typing import Any, Final, cast @@ -2871,14 +2872,21 @@ def test_translate_completion_input_params_rejects_a_non_numeric_budget_without_ @pytest.mark.parametrize("budget", ["2048", True]) def test_translate_completion_input_params_drops_a_non_numeric_budget_under_drop_params( - budget: object, monkeypatch: pytest.MonkeyPatch + budget: object, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ): monkeypatch.setattr(litellm, "drop_params", False) - translated: Final = AnthropicAdapter().translate_completion_input_params( - _gemini_request_with_budget(budget, drop_params=True) - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + translated: Final = AnthropicAdapter().translate_completion_input_params( + _gemini_request_with_budget(budget, drop_params=True) + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert translated is not None assert not {"reasoning_effort", "thinking"} & translated.keys() + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in (f"thinking.budget_tokens={budget!r}", "gemini-2.5-flash", "(drop_params=True)") + ) def test_translate_openai_content_to_anthropic_reasoning_content_without_thinking_blocks(): diff --git a/tests/unit/llms/anthropic/messages/adapters/test_handler_reasoning_effort_normalization.py b/tests/unit/llms/anthropic/messages/adapters/test_handler_reasoning_effort_normalization.py index e504513dd36..27ed1f5ebd6 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_handler_reasoning_effort_normalization.py @@ -8,6 +8,7 @@ leaves the adapter, in the shape the target expects. import functools import json +import logging from collections.abc import Mapping from typing import Final @@ -510,26 +511,33 @@ class TestADeclaredDeploymentRoundsTheMessagesTierOntoItsLevels: @pytest.mark.asyncio @pytest.mark.parametrize("budget", ["2048", True]) @pytest.mark.parametrize("litellm_params, model_info, upstream_url, upstream_body", _MALFORMED_BUDGET_TARGETS) - async def test_drop_params_drops_a_non_numeric_budget( + async def test_drop_params_drops_a_non_numeric_budget_with_one_warning( self, local_model_cost_map: None, respx_mock: respx.MockRouter, + caplog: pytest.LogCaptureFixture, litellm_params: Mapping[str, object], model_info: Mapping[str, object], upstream_url: str, upstream_body: Mapping[str, object], budget: object, ) -> None: - sent: Final = await _body_sent_upstream( - respx_mock, - {**litellm_params, "drop_params": True}, - model_info, - upstream_url, - upstream_body, - thinking={"type": "enabled", "budget_tokens": budget}, - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + sent: Final = await _body_sent_upstream( + respx_mock, + {**litellm_params, "drop_params": True}, + model_info, + upstream_url, + upstream_body, + thinking={"type": "enabled", "budget_tokens": budget}, + ) assert not {"reasoning_effort", "reasoning", "thinking"} & sent.keys() + dropped: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + assert len(dropped) == 1 + assert f"thinking.budget_tokens={budget!r}" in dropped[0] + assert "(drop_params=True)" in dropped[0] + assert str(litellm_params["model"]).split("/", 1)[1] in dropped[0] @pytest.mark.asyncio @pytest.mark.parametrize( diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index 5ce6d03f477..af3108e5072 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -159,7 +159,7 @@ async def test_openai_model_does_not_forward_stream_options_to_responses_api(): await litellm.anthropic.messages.acreate( max_tokens=100, messages=[{"role": "user", "content": "Hello, how are you?"}], - model="openai/gpt-5.5", + model="openai/gpt-6.1-sol", api_key="test-api-key", stream_options={"include_usage": True}, ) @@ -669,7 +669,9 @@ class TestThinkingSummaryPreservation: ) thinking = {"type": "enabled", "budget_tokens": 5000, "summary": "concise"} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( + thinking, model="openai/gpt-6.1-sol" + ) assert result == {"effort": "high", "summary": "concise"} def test_responses_adapter_no_summary_by_default(self): @@ -683,7 +685,9 @@ class TestThinkingSummaryPreservation: try: litellm.reasoning_auto_summary = False thinking = {"type": "enabled", "budget_tokens": 5000} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( + thinking, model="openai/gpt-6.1-sol" + ) assert result == {"effort": "high"} assert result is not None and "summary" not in result finally: diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py index 596cb968a0e..1d0d70d8e60 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import base64 import json +import logging import struct import time import uuid @@ -891,24 +892,32 @@ async def test_native_messages_send_no_effort_for_a_default_effort_beside_a_thin @pytest.mark.asyncio @pytest.mark.parametrize("model_info", [None, _DECLARING_LOW_AND_HIGH]) @pytest.mark.parametrize("drop_globally", [False, True]) -async def test_native_messages_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude( +async def test_native_messages_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude_with_one_warning( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, model_info: Mapping[str, object] | None, drop_globally: bool, ) -> None: monkeypatch.setattr(litellm, "drop_params", drop_globally) request_drop: Final = {} if drop_globally else {"drop_params": True} - dropped: Final = await _native_opus_body_sent( - respx_mock, model_info, thinking={"type": "enabled", "budget_tokens": "lots"}, **request_drop - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + dropped: Final = await _native_opus_body_sent( + respx_mock, model_info, thinking={"type": "enabled", "budget_tokens": "lots"}, **request_drop + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] adaptive_alone: Final = await _native_opus_body_sent( respx_mock, model_info, thinking={"type": "adaptive"}, **request_drop ) assert "output_config" not in adaptive_alone assert dropped == adaptive_alone + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in ("thinking.budget_tokens='lots'", "claude-opus-5-5", "(drop_params=True)") + ) _MIXED_GROUP: Final = "gpt-5.6-mixed" diff --git a/tests/unit/llms/anthropic/messages/native/test_reasoning_effort_translation.py b/tests/unit/llms/anthropic/messages/native/test_reasoning_effort_translation.py index 50dc0deabfd..28abb348371 100644 --- a/tests/unit/llms/anthropic/messages/native/test_reasoning_effort_translation.py +++ b/tests/unit/llms/anthropic/messages/native/test_reasoning_effort_translation.py @@ -1,5 +1,8 @@ """Tests for ``reasoning_effort`` translation on the Anthropic /v1/messages route.""" +import logging +from typing import Final + import pytest from litellm.constants import ( @@ -46,6 +49,49 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model( assert result.get("output_config") == {"effort": expected_effort} +def _reasoning_effort_drop_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + record.getMessage() for record in caplog.records if "Dropping reasoning_effort=" in record.getMessage() + ) + + +@pytest.mark.parametrize("reasoning_effort", [5, {"effort": "high"}]) +def test_reasoning_effort_that_is_not_a_string_is_dropped_with_one_warning( + reasoning_effort: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-opus-4-7", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 1024, "reasoning_effort": reasoning_effort}, + litellm_params={}, + headers={}, + ) + + warnings: Final = _reasoning_effort_drop_warnings(caplog) + assert "reasoning_effort" not in result + assert "output_config" not in result + assert len(warnings) == 1 + assert f"reasoning_effort={reasoning_effort!r}" in warnings[0] + assert "claude-opus-4-7" in warnings[0] + + +@pytest.mark.parametrize("reasoning_effort", ["high", {"summary": "concise"}]) +def test_reasoning_effort_that_is_a_string_or_names_no_level_logs_no_drop( + reasoning_effort: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-opus-4-7", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 1024, "reasoning_effort": reasoning_effort}, + litellm_params={}, + headers={}, + ) + + assert _reasoning_effort_drop_warnings(caplog) == () + + def test_reasoning_effort_none_clears_thinking_and_output_config(): config = AnthropicMessagesConfig() optional_params = { diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_transformation.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_transformation.py index f6fcb9e6f1e..ece0cf5d423 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_transformation.py @@ -44,6 +44,7 @@ def _make_request(**overrides) -> AnthropicMessagesRequest: _ADAPTER = LiteLLMAnthropicToResponsesAPIAdapter() +_MODEL: Final = "openai/gpt-6.1-sol" # --------------------------------------------------------------------------- @@ -957,13 +958,13 @@ class TestTranslateThinkingToReasoning: """Anthropic thinking param -> Responses API reasoning param.""" def test_budget_high_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}, model=_MODEL) # Default (reasoning_auto_summary=False): only effort, no summary assert result == {"effort": "high"} assert result is not None and "summary" not in result def test_budget_above_threshold_high_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 50000}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 50000}, model=_MODEL) assert result is not None assert result["effort"] == "high" assert "summary" not in result @@ -973,7 +974,8 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result == {"effort": "medium"} assert result is not None and "summary" not in result @@ -983,13 +985,14 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result == {"effort": "low"} assert result is not None and "summary" not in result def test_budget_minimal_effort(self): - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 500}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 500}, model=_MODEL) assert result == {"effort": "minimal"} assert result is not None and "summary" not in result @@ -998,7 +1001,8 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result_high is not None assert result_high["effort"] == "high" @@ -1006,7 +1010,8 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result_medium is not None assert result_medium["effort"] == "medium" @@ -1015,23 +1020,24 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result_low is not None assert result_low["effort"] == "low" assert "summary" not in result_low def test_disabled_type_returns_none(self): - result = _ADAPTER.translate_thinking_to_reasoning({"type": "disabled"}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "disabled"}, model=_MODEL) assert result is None def test_non_dict_returns_none(self): - result = _ADAPTER.translate_thinking_to_reasoning("enabled") # type: ignore + result = _ADAPTER.translate_thinking_to_reasoning("enabled", model=_MODEL) # type: ignore assert result is None def test_missing_budget_defaults_to_minimal(self): """Missing budget_tokens defaults to 0, below the low threshold -> minimal.""" - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled"}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled"}, model=_MODEL) assert result == {"effort": "minimal"} assert result is not None and "summary" not in result @@ -1042,7 +1048,7 @@ class TestTranslateThinkingToReasoning: original = litellm.reasoning_auto_summary try: litellm.reasoning_auto_summary = True - result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}) + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled", "budget_tokens": 10000}, model=_MODEL) assert result == {"effort": "high", "summary": "detailed"} finally: litellm.reasoning_auto_summary = original @@ -1059,7 +1065,8 @@ class TestTranslateThinkingToReasoning: { "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - } + }, + model=_MODEL, ) assert result == {"effort": "medium", "summary": "detailed"} finally: diff --git a/tests/unit/llms/anthropic/messages/test_reasoning_effort_fields.py b/tests/unit/llms/anthropic/messages/test_reasoning_effort_fields.py index 3432c9503de..7a26f403408 100644 --- a/tests/unit/llms/anthropic/messages/test_reasoning_effort_fields.py +++ b/tests/unit/llms/anthropic/messages/test_reasoning_effort_fields.py @@ -14,12 +14,12 @@ from typing import Any, Dict import pytest import litellm +from litellm.litellm_core_utils.reasoning_effort_capability import ( + resolve_supported_reasoning_efforts, +) from litellm.llms.anthropic.messages.utils import ( normalize_reasoning_effort_value, ) -from litellm.router_utils.reasoning_effort_capability import ( - resolve_supported_reasoning_efforts, -) from litellm.utils import get_model_info @@ -162,7 +162,8 @@ class TestAdapterAdaptiveThinking: adapter = LiteLLMAnthropicMessagesAdapter() result = adapter.translate_anthropic_thinking_to_reasoning_effort( - {"type": "adaptive"} + {"type": "adaptive"}, + model="openai/gpt-6.1-sol", ) assert result == "medium" @@ -198,6 +199,7 @@ class TestAdapterAdaptiveThinking: result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( thinking={"type": "adaptive"}, output_config={"effort": "xhigh"}, + model="openai/gpt-6.1-sol", ) assert result is not None assert result["effort"] == "xhigh" @@ -210,6 +212,7 @@ class TestAdapterAdaptiveThinking: result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( thinking={"type": "adaptive"}, + model="openai/gpt-6.1-sol", ) assert result is not None assert result["effort"] == "medium" diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 3ec72f29d34..d18ef28c68e 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -11,6 +11,7 @@ Verifies that: """ import json +import logging import os import sys import threading @@ -3543,12 +3544,11 @@ class TestWifProviderAllowlist: monkeypatch.setenv("WIF_TEST_JWT", "jwt-assertion-value") def test_vertex_anthropic_never_mints_or_sends_the_assertion(self, monkeypatch, wif_engine): + import litellm from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) - import litellm - poster, calls = wif_engine self._env_only_wif(monkeypatch) @@ -3774,7 +3774,7 @@ class TestModelDiscovery: ) monkeypatch.setattr("litellm.module_level_client", client) - with pytest.raises(Exception, match="invalid x-api-key") as exc_info: # noqa: B017, PT011 # the callee raises a bare Exception; match pins the sanitized text + with pytest.raises(Exception, match="invalid x-api-key") as exc_info: # noqa: B017 # the callee raises a bare Exception; match pins the sanitized text AnthropicModelInfo().get_models(api_base="https://api.anthropic.com") assert "invalid x-api-key" in str(exc_info.value) @@ -3946,10 +3946,10 @@ class TestWifDisabledOnClientRedirectedBase: """Setting the sentinel is only half of it. get_litellm_params rebuilds litellm_params from kwargs, so a field it does not carry is dropped on the way and the deployment federates for the caller-chosen base after all.""" - from litellm.litellm_core_utils.get_litellm_params import get_litellm_params - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( DISABLE_WORKLOAD_IDENTITY_PARAM, ) + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params funneled = get_litellm_params(**{DISABLE_WORKLOAD_IDENTITY_PARAM: True}) @@ -3958,7 +3958,7 @@ class TestWifDisabledOnClientRedirectedBase: def test_the_sentinel_is_not_client_settable(self): """It is server-owned in both directions: a caller must not be able to set it, and must not be able to clear it either.""" - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( DISABLE_WORKLOAD_IDENTITY_PARAM, ) from litellm.types.router import reject_server_owned_wif_params @@ -3969,11 +3969,11 @@ class TestWifDisabledOnClientRedirectedBase: def test_base_override_clears_wif_and_sets_the_sentinel(self): """A federation token minted for a client-chosen api_base would send the workload's assertion, and then the minted bearer, to that host.""" - from litellm.llms.anthropic.wif import resolve_anthropic_wif_params - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( DISABLE_WORKLOAD_IDENTITY_PARAM, get_dynamic_litellm_params, ) + from litellm.llms.anthropic.wif import resolve_anthropic_wif_params admin_deployment = { "model": "anthropic/claude-sonnet-5", @@ -3994,8 +3994,8 @@ class TestWifDisabledOnClientRedirectedBase: def test_sentinel_blocks_env_var_configured_federation(self, monkeypatch): """Environment-configured federation cannot be cleared out of a dict, so the sentinel is what stops it on a redirected deployment.""" + from litellm.litellm_core_utils.clientside_credential_handler import DISABLE_WORKLOAD_IDENTITY_PARAM from litellm.llms.anthropic.wif import resolve_anthropic_wif_params - from litellm.router_utils.clientside_credential_handler import DISABLE_WORKLOAD_IDENTITY_PARAM monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_env") monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-env") @@ -4009,11 +4009,11 @@ class TestWifDisabledOnClientRedirectedBase: """Same failure mode the legacy-path test above guards against, for the internal_issuer identity source: a signing_key_ref resolved for a client-chosen api_base would mint an assertion, and then a bearer token, for that host.""" - from litellm.llms.anthropic.wif import resolve_anthropic_wif_params - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( DISABLE_WORKLOAD_IDENTITY_PARAM, get_dynamic_litellm_params, ) + from litellm.llms.anthropic.wif import resolve_anthropic_wif_params admin_deployment = { "model": "anthropic/claude-sonnet-5", @@ -4038,11 +4038,11 @@ class TestWifDisabledOnClientRedirectedBase: def test_base_override_clears_keycloak_fields(self): """Same as the internal_issuer case above, for the keycloak identity source: a client_secret_ref resolved for a client-chosen api_base must not follow it there.""" - from litellm.llms.anthropic.wif import resolve_anthropic_wif_params - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( DISABLE_WORKLOAD_IDENTITY_PARAM, get_dynamic_litellm_params, ) + from litellm.llms.anthropic.wif import resolve_anthropic_wif_params admin_deployment = { "model": "anthropic/claude-sonnet-5", @@ -4237,15 +4237,18 @@ def test_thinking_display_beta_requires_active_thinking(thinking: object, expect @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize( - ("display", "expected_thinking"), + ("provider", "display", "expected_thinking"), ( - ("summarized", {"type": "adaptive", "display": "summarized"}), - ("omitted", {"type": "adaptive", "display": "omitted"}), - ("updates", {"type": "adaptive"}), + ("azure_ai", "summarized", {"type": "adaptive", "display": "summarized"}), + ("azure_ai", "omitted", {"type": "adaptive", "display": "omitted"}), + ("azure_ai", "updates", {"type": "adaptive", "display": "updates"}), + ("vertex_ai", "updates", {"type": "adaptive", "display": "updates"}), + ("anthropic", "updates", {"type": "adaptive", "display": "updates"}), + ("databricks", "updates", {"type": "adaptive"}), ), ) def test_shared_legacy_thinking_translation_preserves_supported_display( - display: str, expected_thinking: dict[str, str] + provider: str, display: str, expected_thinking: dict[str, str] ) -> None: from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -4256,7 +4259,7 @@ def test_shared_legacy_thinking_translation_preserves_supported_display( AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( model="claude-opus-5", optional_params=optional_params, - custom_llm_provider="azure_ai", + custom_llm_provider=provider, drop_params=False, ) @@ -4530,16 +4533,66 @@ def test_shared_legacy_thinking_translation_rejects_a_non_numeric_budget(budget_ @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize("budget_tokens", _NON_NUMERIC_BUDGETS) -def test_shared_legacy_thinking_translation_drops_a_non_numeric_budget_under_drop_params(budget_tokens: object) -> None: +def test_shared_legacy_thinking_translation_drops_a_non_numeric_budget_under_drop_params( + budget_tokens: object, caplog: pytest.LogCaptureFixture +) -> None: from litellm.llms.anthropic.common_utils import AnthropicModelInfo optional_params: Final = {"thinking": {"type": "enabled", "budget_tokens": budget_tokens}} - AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( - model="claude-opus-5", optional_params=optional_params, custom_llm_provider="azure_ai", drop_params=True - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model="claude-opus-5", optional_params=optional_params, custom_llm_provider="azure_ai", drop_params=True + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert optional_params == {"thinking": {"type": "adaptive"}} + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in (f"thinking.budget_tokens={budget_tokens!r}", "claude-opus-5", "(drop_params=True)") + ) + + +_SIGNED_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "summary so far", "signature": "sig"} + + +@pytest.mark.parametrize( + "provider,expected", + [("anthropic", True), ("vertex_ai", True), ("azure_ai", True), ("bedrock", False), ("bedrock_converse", False)], +) +@pytest.mark.parametrize( + "optional_params,messages", + [ + ({"compaction": {"type": "summarize"}}, [{"role": "user", "content": "hi"}]), + ({}, [{"role": "assistant", "content": [_SIGNED_COMPACTION_BLOCK]}, {"role": "user", "content": "go on"}]), + ], + ids=["compaction_param", "signed_compaction_block"], +) +def test_requires_native_compaction_beta_per_provider( + provider: str, expected: bool, optional_params: dict[str, object], messages: list[object] +) -> None: + from litellm.llms.anthropic.common_utils import requires_native_compaction_beta + + assert requires_native_compaction_beta(provider, optional_params, messages) is expected + + +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) +@pytest.mark.parametrize( + "optional_params,messages", + [ + ({"max_tokens": 64}, [{"role": "user", "content": "hi"}]), + ({}, [{"role": "assistant", "content": [{"type": "compaction", "content": "unsigned"}]}]), + ({"compaction": None}, [{"role": "user", "content": "hi"}]), + ], + ids=["no_signal", "unsigned_block", "compaction_none"], +) +def test_requires_native_compaction_beta_is_false_without_a_signal( + provider: str, optional_params: dict[str, object], messages: list[object] +) -> None: + from litellm.llms.anthropic.common_utils import requires_native_compaction_beta + + assert requires_native_compaction_beta(provider, optional_params, messages) is False @pytest.mark.parametrize( diff --git a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py index 4d0ebef4f2d..d95f831df3e 100644 --- a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py +++ b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Mapping from datetime import datetime from types import SimpleNamespace @@ -597,6 +598,45 @@ async def test_native_count_preserves_settings_and_requires_every_prefix( assert unexpected.call_count == 0 +@pytest.mark.usefixtures("local_model_cost_map") +def test_native_baseline_body_drops_a_non_numeric_budget_without_logging_a_drop( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + request: Final = { + "model": "claude-opus-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + "thinking": {"type": "enabled", "budget_tokens": "high"}, + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + body: Final = prepare_native_baseline_body(request, "anthropic/claude-opus-5") + + assert body is not None + assert body["thinking"] == {"type": "adaptive"} + assert [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] == [] + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_native_baseline_body_drops_a_non_string_reasoning_effort_without_logging_a_drop( + caplog: pytest.LogCaptureFixture, +) -> None: + request: Final = { + "model": "claude-opus-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + "reasoning_effort": 5, + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + body: Final = prepare_native_baseline_body(request, "anthropic/claude-opus-5") + + assert body is not None + assert "reasoning_effort" not in body + assert [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] == [] + + def test_native_baseline_body_matches_the_sent_body_when_drop_params_meets_an_auto_added_beta( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 4913d2b0d1d..7807975be24 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -215,3 +215,56 @@ async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest assert result is not None assert result.total_tokens == 6 assert route.called + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("assistant_content", "expected_beta"), + [ + ( + [{"type": "compaction", "content": "Summary of the chat so far.", "signature": "sig-abc"}], + "compact-2026-09-04,token-counting-2024-11-01", + ), + ("Use Recipe and Step tables.", "token-counting-2024-11-01"), + ], + ids=("signed_compaction_block", "plain_turn"), +) +async def test_count_sends_the_compact_beta_only_when_the_history_replays_a_compaction_block( + monkeypatch: pytest.MonkeyPatch, assistant_content: object, expected_beta: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 376})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-5-5", + messages=[ + {"role": "assistant", "content": assistant_content}, + {"role": "user", "content": "What tables did we pick?"}, + ], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-5-5", + ) + assert result is not None + assert result.total_tokens == 376 + assert route.calls.last.request.headers["anthropic-beta"] == expected_beta + + +@pytest.mark.asyncio +async def test_count_sends_the_tool_change_beta_when_the_history_adds_a_tool(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 541})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-5-5", + messages=[ + {"role": "user", "content": "Use the ping tool once."}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-5-5", + ) + assert result is not None + assert result.total_tokens == 541 + assert route.calls.last.request.headers["anthropic-beta"] == "mid-conversation-tool-changes-2026-07-01,token-counting-2024-11-01" diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 145d45317c3..9f68a4fb73f 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -1,4 +1,5 @@ import json +import logging import traceback from collections.abc import Callable from types import MappingProxyType @@ -758,6 +759,66 @@ def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): assert result["tools"][1]["function"]["name"] == "read_file" +def test_extra_parameters_error_drops_reasoning_effort_with_one_warning(caplog: pytest.LogCaptureFixture) -> None: + import httpx + + error_text: Final = json.dumps( + { + "error": { + "message": "Extra parameters ['reasoning_effort'] are not allowed when extra-parameters is not set " + "or set to be 'error'. Set extra-parameters to 'pass-through' to pass to the model." + } + } + ) + error: Final = httpx.HTTPStatusError( + message="422", + request=httpx.Request("POST", "https://example.services.ai.azure.com/models/chat/completions"), + response=httpx.Response(422, text=error_text), + ) + request_data: Final = { + "model": "Kimi-K2.6", + "messages": [{"role": "user", "content": "Say hi."}], + "reasoning_effort": "high", + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + retried: Final = AzureAIStudioConfig().transform_request_on_unprocessable_entity_error(error, request_data) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert retried == {"model": "Kimi-K2.6", "messages": [{"role": "user", "content": "Say hi."}]} + assert len(drop_warnings) == 1 + assert all(fragment in drop_warnings[0] for fragment in ("reasoning_effort='high'", "Kimi-K2.6")) + assert "drop_params" not in drop_warnings[0] + + +def test_extra_inputs_error_names_drop_params_in_the_reasoning_drop_warning(caplog: pytest.LogCaptureFixture) -> None: + import httpx + + detail: Final = [ + {"type": "extra_forbidden", "loc": ["body", "reasoning_effort"], "msg": "Extra inputs are not permitted"} + ] + error: Final = httpx.HTTPStatusError( + message="422", + request=httpx.Request("POST", "https://example.services.ai.azure.com/models/chat/completions"), + response=httpx.Response(422, text=json.dumps({"error": {"message": json.dumps({"detail": detail})}})), + ) + request_data: Final = { + "model": "Mistral-Large-3", + "messages": [{"role": "user", "content": "Say hi."}], + "reasoning_effort": "high", + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + retried: Final = AzureAIStudioConfig().transform_request_on_unprocessable_entity_error(error, request_data) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert retried == {"model": "Mistral-Large-3", "messages": [{"role": "user", "content": "Say hi."}]} + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] for fragment in ("reasoning_effort='high'", "Mistral-Large-3", "(drop_params=True)") + ) + + def _find_key_anywhere(obj, key: str) -> bool: if isinstance(obj, dict): if key in obj: diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 5e0395f84dc..1c4097ad7e0 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -607,3 +607,91 @@ async def test_outgoing_azure_request_keeps_an_object_output_format_under_drop_p body: Final = await _azure_request_body_for_output_format(output_format, drop_params=True) assert body["output_format"] == output_format + + +_SIGNED_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "Summary of the chat so far.", "signature": "sig"} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("messages", "params"), + [ + ([{"role": "user", "content": "Summarize the chat so far."}], {"compaction": {"type": "summarize"}}), + ( + [{"role": "assistant", "content": [_SIGNED_COMPACTION_BLOCK]}, {"role": "user", "content": "Go on"}], + {}, + ), + ], + ids=["compaction_param", "signed_compaction_block"], +) +async def test_outgoing_azure_compaction_request_carries_the_compact_2026_09_04_beta( + messages: list[object], params: dict[str, object] +) -> None: + """Foundry answers 400 "compaction: this parameter requires anthropic-beta: compact-2026-09-04" for the param + and 400 "Input tag 'compaction' ... does not match" for a replayed block without the header, and 200 with it + (live, claude-sonnet-5-5 on Azure AI Foundry, 2026-10-20).""" + received: Final[list[httpx.Request]] = [] + + def foundry_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _foundry_answer(request) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=256, + messages=messages, + model="azure_ai/claude-sonnet-5-5", + api_key="test-api-key", + api_base="https://test-resource.services.ai.azure.com", + client=upstream, + **params, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + assert received[0].headers["anthropic-beta"].split(",").count("compact-2026-09-04") == 1, dict(received[0].headers) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("feature", ("thinking", "output_config", "tool_addition", "tool_removal", "none")) +@pytest.mark.parametrize("explicit_beta", (False, True)) +def test_messages_feature_beta_headers(feature: str, explicit_beta: bool) -> None: + from typing import Final + + from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig + from litellm.types.llms.anthropic import ( + ANTHROPIC_BETA_HEADER_VALUES, + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, + ) + + beta: Final = ( + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + if feature == "thinking" + else ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value + if feature == "output_config" + else ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER + ) + message: Final = ( + {"role": "system", "content": [], "output_config": {"effort": "high"}} + if feature == "output_config" + else {"role": "system", "content": [{"type": feature, "tool": {"type": "tool_reference", "name": "ping"}}]} + if feature in ("tool_addition", "tool_removal") + else {"role": "user", "content": "Reply with OK"} + ) + headers, _ = AzureAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={"anthropic-beta": f"existing-beta,{beta}" if explicit_beta else "existing-beta"}, + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Hello"}, message], + optional_params={"thinking": {"type": "adaptive", "display": "updates"}} if feature == "thinking" else {}, + litellm_params={"vertex_ai_project": "test-project", "vertex_ai_location": "us-east5"}, + api_key="test-key", + api_base="https://example.com/anthropic", + ) + + betas: Final = headers["anthropic-beta"].split(",") + assert betas.count(beta) == int(feature != "none" or explicit_beta) + assert betas.count("existing-beta") == 1 diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py index 1179b0dd5da..6ff2a6973a3 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py @@ -5,10 +5,12 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) +from typing import Final from unittest.mock import MagicMock, patch import pytest +from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig from litellm.types.router import GenericLiteLLMParams @@ -484,3 +486,85 @@ def test_chat_flagged_model_keeps_mid_conversation_system_role_in_place(local_mo "role": "system", "content": [{"type": "text", "text": "Answer with exactly one word."}], } + + +@pytest.mark.parametrize( + ("optional_params", "messages", "expected"), + [ + ({"compaction": {"type": "summarize"}}, [{"role": "user", "content": "Hello"}], True), + ( + {}, + [ + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary.", "signature": "sig"}]}, + {"role": "user", "content": "Go on"}, + ], + True, + ), + ({}, [{"role": "user", "content": "Hello"}], False), + ], + ids=["compaction_param", "signed_compaction_block", "no_compaction"], +) +def test_azure_chat_compaction_sends_the_compact_2026_09_04_beta( + local_beta_headers_config: None, optional_params: dict[str, object], messages: list[object], expected: bool +) -> None: + headers: Final = AzureAnthropicConfig().update_headers_with_optional_anthropic_beta( + headers={}, optional_params=optional_params, messages=messages + ) + + filtered: Final = update_headers_with_filtered_beta(headers=headers, provider="azure_ai") + + assert ("compact-2026-09-04" in filtered.get("anthropic-beta", "").split(",")) is expected + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("thinking_type", ("adaptive", "enabled")) +@pytest.mark.parametrize("display", ("summarized", "omitted", "updates")) +def test_chat_preserves_thinking_display_and_beta(thinking_type: str, display: str) -> None: + from typing import Final + + from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + config: Final = AzureAnthropicConfig() + optional_params: Final = config.map_openai_params( + non_default_params={"thinking": {"type": thinking_type, "display": display, "budget_tokens": 2048}}, + optional_params={}, + model="claude-fable-5-1", + drop_params=False, + ) + headers: Final = config.validate_environment( + headers={"anthropic-beta": "existing-beta"}, + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Reply with OK"}], + optional_params=optional_params, + litellm_params={}, + api_key="test-key", + ) + result: Final = config.transform_request( + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Reply with OK"}], + optional_params=optional_params, + litellm_params={}, + headers=headers, + ) + + assert result["thinking"]["type"] == "adaptive" + assert result["thinking"]["display"] == display + assert (ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in headers["anthropic-beta"].split(",")) == (display == "updates") + assert "existing-beta" in headers["anthropic-beta"].split(",") + + +@pytest.mark.parametrize("thinking", ("enabled", {"type": "bogus"}, {"type": "adaptive", "display": "future-mode"})) +def test_beta_detection_preserves_unrecognized_thinking(thinking: object) -> None: + from typing import Final + + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + params: Final = {"thinking": thinking} + headers: Final = AzureAnthropicConfig().validate_environment( + headers={}, model="claude-fable-5-1", messages=[], optional_params=params, + litellm_params={}, api_key="test-key", + ) + + assert params["thinking"] == thinking + assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER not in headers.get("anthropic-beta", "").split(",") diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index 03459ba16c8..ec6f118837f 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -6,14 +6,19 @@ import httpx import pytest from pydantic import TypeAdapter, ValidationError +import litellm +from litellm.llms.azure_ai.decisions.transformation import AzureAIDecisionsConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( + BaseDecisionsConfig, + DecisionsInputLimits, ir_to_systemone_request, ir_to_systemone_response, parse_systemone_response, systemone_request_to_ir, ) from litellm.llms.openai.decisions.transformation import openai_request_to_ir +from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig from litellm.types.decisions import ( MAX_DECISION_QUESTIONS, DecisionsIRRefusal, @@ -49,6 +54,61 @@ _SYSTEMONE_REQUEST: Final[Mapping[str, object]] = { } +@pytest.mark.parametrize("declared", (4096, 0, -1, True, None)) +def test_decision_input_limits_use_only_positive_catalog_limits( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, declared: int | None +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "bespoke/test-input-limit", + {"litellm_provider": "bespoke", "mode": "evaluation", "max_input_tokens": declared}, + ) + + limits: Final = BaseDecisionsConfig().get_input_limits("test-input-limit", "bespoke", {}) + + assert limits == DecisionsInputLimits(max_request_tokens=4096 if declared == 4096 else None) + + +def test_unmapped_decision_model_limits_remain_unknown(local_model_cost_map: None) -> None: + assert BaseDecisionsConfig().get_input_limits("unmapped-input-limit", "databricks", {}) == DecisionsInputLimits() + + +@pytest.mark.parametrize( + ("headers", "long_context"), + ( + ({}, False), + ({"TypeSafe-Danger-Zone-Long-Context": "false"}, False), + ({"typesafe-danger-zone-long-context": " TRUE "}, True), + ({"TypeSafe-Danger-Zone-Long-Context": "true", "typesafe-danger-zone-long-context": "true"}, True), + ({"TypeSafe-Danger-Zone-Long-Context": "true", "typesafe-danger-zone-long-context": "false"}, False), + ({"typesafe-danger-zone-long-context": "true,false"}, False), + ), +) +def test_typesafe_limits_require_unambiguous_long_context_opt_in( + headers: Mapping[str, str], long_context: bool +) -> None: + """Provider contract from https://docs.typesafe.ai/models, verified 2026-10-30""" + limits: Final = TypeSafeDecisionsConfig().get_input_limits("jev-latest", "typesafe", headers) + + assert limits == ( + DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144) + if long_context + else DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768) + ) + + +@pytest.mark.parametrize("model", ("Microsoft-Decision-1", "unmapped-decision-deployment")) +def test_azure_decision_deployments_use_catalog_limits_without_typesafe_modes( + local_model_cost_map: None, model: str +) -> None: + catalog_limit: Final = litellm.get_model_info("Microsoft-Decision-1", "azure_ai")["max_input_tokens"] + limits: Final = AzureAIDecisionsConfig().get_input_limits( + model, "azure_ai", {"TypeSafe-Danger-Zone-Long-Context": "true"} + ) + + assert limits == DecisionsInputLimits(max_request_tokens=catalog_limit) + + def _openai_ir(raw: Mapping[str, object]) -> DecisionsIRRequest: return openai_request_to_ir(_OPENAI_BODY.validate_python(raw)) diff --git a/tests/unit/llms/baseten/chat/test_baseten_completions.py b/tests/unit/llms/baseten/chat/test_baseten_completions.py index 4f135a1a106..407b39319e9 100644 --- a/tests/unit/llms/baseten/chat/test_baseten_completions.py +++ b/tests/unit/llms/baseten/chat/test_baseten_completions.py @@ -7,8 +7,9 @@ import pytest import respx import litellm +from litellm.litellm_core_utils.model_retirement import is_past_deprecation_date from litellm.llms.baseten.chat import BasetenConfig -from litellm.utils import get_valid_models +from litellm.utils import _model_cache, get_valid_models class TestBasetenRouting: @@ -65,7 +66,7 @@ class TestBasetenModelDiscovery: servable_registry_models: Final = { key for key, value in litellm.model_cost.items() - if value.get("litellm_provider") == "baseten" and not litellm._is_past_deprecation_date(value) + if value.get("litellm_provider") == "baseten" and not is_past_deprecation_date(value) } assert servable_registry_models @@ -92,6 +93,20 @@ class TestBasetenModelDiscovery: litellm.baseten_models.discard(model) litellm.add_known_models({}) + def test_wildcard_drops_a_model_whose_retirement_date_passes_while_running(self, monkeypatch: pytest.MonkeyPatch): + model: Final = "baseten/acme/retires-while-running" + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "baseten", "deprecation_date": "2999-01-01"}) + try: + litellm.add_known_models({model: litellm.model_cost[model]}) + assert model in get_valid_models(custom_llm_provider="baseten") + + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "baseten", "deprecation_date": "2000-01-01"}) + + assert model not in get_valid_models(custom_llm_provider="baseten") + finally: + litellm.baseten_models.discard(model) + litellm.add_known_models({}) + @pytest.mark.parametrize( "value, expected", [ @@ -102,7 +117,7 @@ class TestBasetenModelDiscovery: ], ) def test_is_past_deprecation_date(self, value, expected): - assert litellm._is_past_deprecation_date(value) is expected + assert is_past_deprecation_date(value) is expected def test_get_models_calls_model_api(self, respx_mock: respx.MockRouter) -> None: route: Final = respx_mock.get("https://inference.baseten.co/v1/models").mock( @@ -156,6 +171,18 @@ class TestBasetenModelDiscovery: assert models == ["baseten/brand-new/model"] + @patch.dict(os.environ, {"BASETEN_API_KEY": "test-key"}) + def test_live_catalog_keeps_a_model_the_cost_map_dates_as_retired(self, respx_mock: respx.MockRouter) -> None: + assert is_past_deprecation_date(litellm.model_cost["baseten/zai-org/GLM-5"]) + _model_cache.flush_cache() + respx_mock.get("https://inference.baseten.co/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "zai-org/GLM-5"}]}) + ) + + models: Final = get_valid_models(check_provider_endpoint=True, custom_llm_provider="baseten") + + assert models == ["baseten/zai-org/GLM-5"] + class TestBasetenReasoningEffort: def test_reasoning_effort_passthrough_with_drop_params(self): diff --git a/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py b/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py index 3faae2e3289..f4dcff1969e 100644 --- a/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py +++ b/tests/unit/llms/bedrock/chat/chat_completions/test_bedrock_chat_completions_transformation.py @@ -1,6 +1,7 @@ """Bedrock Runtime Chat Completions: the default for GPT 5.6 and newer, ``bedrock/chat_completions/`` for the rest.""" import json +import logging import re import httpx @@ -706,18 +707,30 @@ def test_with_max_completion_tokens_leaves_other_params_alone(): @pytest.mark.parametrize( - "model", - ["us.xai.grok-4.6", "bedrock/us-gov-west-1/us.xai.grok-4.6"], + ("model", "logged_model"), + [ + ("us.xai.grok-4.6", "us.xai.grok-4.6"), + ("chat_completions/us.xai.grok-4.6", "us.xai.grok-4.6"), + ("bedrock/us-gov-west-1/us.xai.grok-4.6", "us-gov-west-1/us.xai.grok-4.6"), + ], ) -def test_map_openai_params_drops_reasoning_effort_none_for_grok(model): +def test_map_openai_params_drops_reasoning_effort_none_for_grok_with_one_warning( + model: str, logged_model: str, caplog: pytest.LogCaptureFixture +) -> None: cfg = AmazonBedrockRuntimeChatCompletionsConfig() - mapped = cfg.map_openai_params( - non_default_params={"reasoning_effort": "none", "max_tokens": 64}, - optional_params={}, - model=model, - drop_params=False, - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + mapped = cfg.map_openai_params( + non_default_params={"reasoning_effort": "none", "max_tokens": 64}, + optional_params={}, + model=model, + drop_params=False, + ) + drop_warnings = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoning_effort" not in mapped + assert len(drop_warnings) == 1 + assert "reasoning_effort='none'" in drop_warnings[0] + assert f"model={logged_model!r}" in drop_warnings[0] + assert "drop_params" not in drop_warnings[0] def test_map_openai_params_keeps_reasoning_effort_low_for_grok(): @@ -750,17 +763,24 @@ def test_map_openai_params_refuses_a_non_string_reasoning_effort_without_drop_pa @pytest.mark.parametrize("reasoning_effort", [["low"], {"effort": "low"}, 5], ids=["list", "object", "int"]) @pytest.mark.parametrize("drop_params_via", ["request", "litellm.drop_params"]) def test_map_openai_params_drops_a_non_string_reasoning_effort_under_drop_params( - monkeypatch, model, reasoning_effort, drop_params_via + monkeypatch, caplog, model, reasoning_effort, drop_params_via ): monkeypatch.setattr(litellm, "drop_params", drop_params_via == "litellm.drop_params") - mapped = AmazonBedrockRuntimeChatCompletionsConfig().map_openai_params( - non_default_params={"reasoning_effort": reasoning_effort, "max_tokens": 64}, - optional_params={}, - model=model, - drop_params=drop_params_via == "request", - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + mapped = AmazonBedrockRuntimeChatCompletionsConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort, "max_tokens": 64}, + optional_params={}, + model=model, + drop_params=drop_params_via == "request", + ) + drop_warnings = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoning_effort" not in mapped assert mapped["max_completion_tokens"] == 64 + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in (f"reasoning_effort={reasoning_effort!r}", repr(model), "(drop_params=True)") + ) def test_map_openai_params_keeps_reasoning_effort_none_for_gpt56(): diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index d82ba192b1f..524caaba7a1 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -13,6 +13,7 @@ import pytest import respx import litellm +from litellm._logging import verbose_logger from litellm import ModelResponse, Router, completion from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, @@ -1221,6 +1222,97 @@ def test_transform_request_drops_lookaround_regex_when_the_deployment_model_info assert properties == _LOOKAROUND_FREE_PROPERTIES +@pytest.mark.parametrize( + ("model", "litellm_params"), + [ + ("us.newvendor.model-1", None), + ("converse/us.newvendor.model-1", None), + ("us-east-1/us.newvendor.model-1", None), + ("us.newvendor.model-1", {"base_model": "bedrock/converse/us.newvendor.model-1"}), + ], +) +def test_transform_request_drops_lookaround_regex_for_a_model_the_cost_map_does_not_list(model, litellm_params): + assert "us.newvendor.model-1" not in litellm.model_cost + + properties: Final = _tool_schema_properties(model, _ARTIFACT_DATA_ANTHROPIC_TOOL, litellm_params) + + assert properties == _LOOKAROUND_FREE_PROPERTIES + + +@pytest.mark.parametrize( + "model", + [ + "arn:aws:bedrock:us-east-1:123456789012:custom-model-deployment/abc123", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + "us.anthropic.claude-unreleased-v1:0", + "eu.openai.gpt-unreleased", + "apac.newvendor.listed-1", + ], +) +def test_transform_request_keeps_lookaround_regex_for_an_unlisted_model_litellm_still_recognizes(model, monkeypatch): + monkeypatch.setitem(litellm.model_cost, "us.newvendor.listed-1", {"litellm_provider": "bedrock_converse"}) + + properties = _tool_schema_properties(model, _ARTIFACT_DATA_ANTHROPIC_TOOL) + + assert properties == _ARTIFACT_DATA_INPUT_SCHEMA["properties"] + + +def test_transform_request_keeps_lookaround_regex_for_an_unlisted_model_whose_deployment_opts_out(): + properties = _tool_schema_properties( + "us.newvendor.model-1", _ARTIFACT_DATA_ANTHROPIC_TOOL, {"model_info": {"supports_regex_lookaround": True}} + ) + + assert properties == _ARTIFACT_DATA_INPUT_SCHEMA["properties"] + + +@pytest.mark.parametrize( + "deployment_params", + [ + {"model": "bedrock/us.newvendor.model-1"}, + {"model": "bedrock/converse/us.newvendor.model-1", "allowed_openai_params": ["tools"]}, + ], +) +def test_a_router_deployment_of_a_model_the_cost_map_does_not_list_drops_lookaround_regex( + respx_mock, deployment_params +): + route: Final = respx_mock.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.newvendor.model-1/converse" + ).mock( + return_value=httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + ) + ) + router: Final = Router( + model_list=[ + { + "model_name": "unlisted", + "litellm_params": { + **deployment_params, + "aws_access_key_id": "fake-access-key", + "aws_secret_access_key": "fake-secret-key", + "aws_region_name": "us-east-1", + }, + "model_info": {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6}, + } + ] + ) + + router.completion( + model="unlisted", + messages=[{"role": "user", "content": "hi"}], + tools=[_ARTIFACT_DATA_OPENAI_TOOL], + client=HTTPHandler(), + ) + + tool: Final = json.loads(route.calls[0].request.content)["toolConfig"]["tools"][0] + assert tool["toolSpec"]["inputSchema"]["json"]["properties"] == _LOOKAROUND_FREE_PROPERTIES + + def test_transform_request_keeps_lookaround_regex_when_the_deployment_model_info_opts_out(): properties = _tool_schema_properties( "global.moonshotai.kimi-k3", _ARTIFACT_DATA_ANTHROPIC_TOOL, {"model_info": {"supports_regex_lookaround": True}} @@ -9189,18 +9281,223 @@ def test_bedrock_tools_transformation_valid_params(): assert "test" in result[0]["toolSpec"]["inputSchema"]["json"]["required"] -def test_not_found_error(): +def _invalid_model_identifier_reply(sent: list[httpx.Request]): + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response( + 400, + json={"message": "The provided model identifier is invalid."}, + headers={"x-amzn-ErrorType": "ValidationException", "x-amzn-RequestId": "req-123"}, + ) + + return respond + + +@pytest.mark.parametrize( + "model, path", + [ + ("bedrock/bad_model", b"/model/bad_model/converse"), + ("bedrock/us-east-1/bad_model", b"/model/bad_model/converse"), + ("bedrock/newvendor.listed-v1:0", b"/model/newvendor.listed-v1%3A0/converse"), + ("bedrock/us-east-1/newvendor.listed-v1:0", b"/model/newvendor.listed-v1%3A0/converse"), + ], +) +def test_not_found_error(model, path, monkeypatch): + monkeypatch.setitem(litellm.model_cost, "newvendor.listed-v1:0", {"litellm_provider": "bedrock", "mode": "chat"}) + sent: Final[list[httpx.Request]] = [] + + with pytest.raises(litellm.NotFoundError) as exc_info: + completion( + model=model, + messages=[{"role": "user", "content": "What is the meaning of life"}], + client=HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_invalid_model_identifier_reply(sent))) + ), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert exc_info.value.status_code == 404 + assert exc_info.value.response.status_code == 404 + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-123" + assert [request.url.raw_path for request in sent] == [path] + + +@pytest.mark.parametrize( + "model", ["bedrock/", "bedrock/us-east-1/", "bedrock/../../converse", "bedrock/xx-fake-1/us.xai.grok-4.7"] +) +def test_a_model_id_that_is_empty_or_a_path_fails_locally_without_calling_aws(model): + sent: Final[list[httpx.Request]] = [] + with pytest.raises(litellm.NotFoundError): completion( - model="bedrock/bad_model", - messages=[ - { - "role": "user", - "content": "What is the meaning of life", - } - ], + model=model, + messages=[{"role": "user", "content": "What is the meaning of life"}], + client=HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_invalid_model_identifier_reply(sent))) + ), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", ) + assert sent == [] + + +@pytest.mark.parametrize( + "model, path", + [ + ("bedrock/converse/bad_model", b"/model/bad_model/converse"), + ("bedrock/anthropic.claude-typo-v1:0", b"/model/anthropic.claude-typo-v1:0/invoke"), + ], +) +def test_invalid_model_identifier_stays_a_bad_request_for_a_model_litellm_recognizes(model, path): + sent: Final[list[httpx.Request]] = [] + + with pytest.raises(litellm.BadRequestError) as exc_info: + completion( + model=model, + messages=[{"role": "user", "content": "What is the meaning of life"}], + client=HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_invalid_model_identifier_reply(sent))) + ), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert exc_info.value.status_code == 400 + assert [request.url.raw_path for request in sent] == [path] + + +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +def _converse_tool_call_reply(sent: list[httpx.Request]): + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response( + 200, + json={ + "output": { + "message": { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "t1", "name": "get_weather", "input": {"city": "Paris"}}} + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + ) + + return respond + + +@pytest.mark.parametrize("model", ["bedrock/us.newvendor.model-1", "bedrock/us-east-1/us.newvendor.model-1"]) +def test_model_missing_from_cost_map_sends_tools_to_converse(model): + assert "us.newvendor.model-1" not in litellm.model_cost + sent: Final[list[httpx.Request]] = [] + + response: Final = completion( + model=model, + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_WEATHER_TOOL], + tool_choice="auto", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert response.choices[0].message.tool_calls[0].function.name == "get_weather" + assert [request.url.raw_path for request in sent] == [b"/model/us.newvendor.model-1/converse"] + tool_config: Final = json.loads(sent[0].content)["toolConfig"] + assert [tool["toolSpec"]["name"] for tool in tool_config["tools"]] == ["get_weather"] + assert tool_config["toolChoice"] == {"auto": {}} + + +def test_a_price_only_row_for_an_unlisted_model_still_sends_tools_to_converse(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, "us.newvendor.priced-1", {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6} + ) + sent: Final[list[httpx.Request]] = [] + + completion( + model="bedrock/us.newvendor.priced-1", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_WEATHER_TOOL], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + tool_config: Final = json.loads(sent[0].content)["toolConfig"] + assert [tool["toolSpec"]["name"] for tool in tool_config["tools"]] == ["get_weather"] + + +def test_explicit_converse_route_to_an_unlisted_model_still_drops_tools_under_drop_params(): + sent: Final[list[httpx.Request]] = [] + + completion( + model="bedrock/converse/us.newvendor.model-1", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_WEATHER_TOOL], + drop_params=True, + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert "toolConfig" not in json.loads(sent[0].content) + + +def test_unknown_model_gets_reasoning_effort_as_sent(): + sent: Final[list[httpx.Request]] = [] + + completion( + model="bedrock/us.newvendor.model-1", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + reasoning_effort="high", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert json.loads(sent[0].content)["additionalModelRequestFields"] == {"reasoning_effort": "high"} + + +def test_mapped_model_without_function_calling_still_rejects_tools(): + model: Final = "bedrock/converse/meta.llama3-8b-instruct-v1:0" + assert "meta.llama3-8b-instruct-v1:0" in litellm.model_cost + assert not litellm.supports_function_calling(model=model) + sent: Final[list[httpx.Request]] = [] + + with pytest.raises(litellm.UnsupportedParamsError): + completion( + model=model, + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_WEATHER_TOOL], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_converse_tool_call_reply(sent)))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert sent == [] + @pytest.mark.parametrize( "model, expected_base_model", @@ -10800,6 +11097,52 @@ def test_a_deployment_with_a_batch_control_plane_endpoint_keeps_it_out_of_chat_r assert "aws_bedrock_control_plane_endpoint" not in route.calls[0].request.content.decode() +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +_DEEPSEEK_R1_PROFILE_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.deepseek.r1-v1:0" + + +@pytest.mark.parametrize( + "param, value", (("thinking", {"type": "enabled", "budget_tokens": 1024}), ("reasoning_effort", "high")) +) +def test_bedrock_deepseek_arn_reasoning_drop_warns_once( + param: str, value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + optional_params: Final = AmazonConverseConfig().map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model=_DEEPSEEK_R1_PROFILE_ARN, + drop_params=False, + ) + + warnings: Final = _warnings_naming(caplog, param) + assert param not in optional_params + assert len(warnings) == 1 + assert _DEEPSEEK_R1_PROFILE_ARN in warnings[0] + + +@pytest.mark.parametrize( + "param, value", (("thinking", {"type": "enabled", "budget_tokens": 1024}), ("reasoning_effort", "high")) +) +def test_bedrock_deepseek_r1_drop_of_an_unadvertised_param_does_not_warn( + param: str, value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + AmazonConverseConfig().map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model="us.deepseek.r1-v1:0", + drop_params=False, + ) + + assert _warnings_naming(caplog, param) == () + + def _converse_bound_thinking_request( headers: dict[str, str] | None, body_betas: tuple[str, ...] = (), drop_params: bool = True ) -> dict[str, object]: diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index dd89d9c46b6..20c0e10cb81 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -214,24 +214,17 @@ def test_system_defined_inference_profile_arn_still_routes_to_converse(): assert route == "converse" -def test_other_opaque_arn_types_still_route_to_invoke(): - """ - Only application-inference-profile ARNs are auto-routed to converse. Other - opaque ARNs (provisioned-model, imported-model, custom-model-deployment) - also yield no invoke provider, but they are frequently invoke-only with - provider-specific body formats, so routing them to converse could break - them. Guards the deliberate scope against an over-broad "any opaque ARN -> - converse" generalization. - """ - for arn_segment in ( - "provisioned-model/abcdefgh1234", - "imported-model/abcdefgh1234", - "custom-model-deployment/abcdefgh1234", - ): - route = BedrockModelInfo.get_bedrock_route( - model=f"bedrock/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}" - ) - assert route == "invoke", f"{arn_segment} should stay on invoke route" +@pytest.mark.parametrize( + "arn_segment", + ["provisioned-model/abcdefgh1234", "imported-model/abcdefgh1234", "custom-model-deployment/abcdefgh1234"], +) +def test_bare_opaque_arn_with_no_invoke_provider_routes_to_converse(arn_segment: str): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + model: Final = f"bedrock/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}" + assert BaseAWSLLM.get_bedrock_invoke_provider(model) is None + assert BedrockModelInfo.get_bedrock_route(model) == "converse" + assert BedrockModelInfo.get_bedrock_route(f"bedrock/llama/arn:aws:bedrock:us-east-1:123412341234:{arn_segment}") == "invoke" def test_govcloud_cross_region_inference_prefix(): @@ -1098,6 +1091,84 @@ def test_unmapped_openai_family_model_routes_to_converse(): assert BedrockModelInfo.get_bedrock_route(imported) == "openai" +@pytest.mark.parametrize( + "unmapped", + [ + "bedrock/minimax.minimax-m99-unmapped", + "bedrock/us.nvidia.nemotron-99b-unmapped", + "bedrock/writer.palmyra-x99-unmapped", + "bedrock/zai.glm-99-unmapped", + "bedrock/google.gemma-99-unmapped", + "bedrock/global.xai.grok-99-unmapped", + ], +) +def test_unmapped_converse_only_family_model_routes_to_converse(unmapped: str): + import litellm + + assert BedrockModelInfo.get_base_model(unmapped) not in litellm.bedrock_converse_models + assert BedrockModelInfo.get_bedrock_route(unmapped) == "converse" + + +def test_unmapped_model_of_family_with_invoke_handler_still_routes_to_invoke(): + from typing import Final + + import litellm + + unmapped: Final = "bedrock/meta.llama99-unmapped-v1:0" + assert BedrockModelInfo.get_base_model(unmapped) not in litellm.bedrock_converse_models + assert BedrockModelInfo.get_bedrock_route(unmapped) == "invoke" + + +@pytest.mark.parametrize( + ("model", "model_path", "response_json"), + [ + ( + "bedrock/us.newvendor.model-1", + b"/model/us.newvendor.model-1/converse", + { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + ), + ( + "bedrock/mistral.mistral-99-unmapped-v1:0", + b"/model/mistral.mistral-99-unmapped-v1:0/invoke", + {"outputs": [{"text": "ok", "stop_reason": "stop"}]}, + ), + ], +) +def test_completion_on_unmapped_model_posts_converse_unless_its_id_names_an_invoke_provider( + model: str, model_path: bytes, response_json: dict[str, object] +): + import httpx + + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + assert BedrockModelInfo.get_base_model(model) not in litellm.bedrock_converse_models + sent: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response( + 200, + json=response_json, + headers={"x-amzn-bedrock-input-token-count": "1", "x-amzn-bedrock-output-token-count": "1"}, + ) + + response: Final = litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + aws_access_key_id="testing", + aws_secret_access_key="testing", + aws_region_name="us-west-2", + ) + + assert response.choices[0].message.content == "ok" + assert [request.url.raw_path for request in sent] == [model_path] + + @pytest.mark.parametrize( ("model", "expected"), [ diff --git a/tests/unit/llms/bedrock/test_nova_imported_models.py b/tests/unit/llms/bedrock/test_nova_imported_models.py index e3677aaf9e6..231f3097dcd 100644 --- a/tests/unit/llms/bedrock/test_nova_imported_models.py +++ b/tests/unit/llms/bedrock/test_nova_imported_models.py @@ -24,10 +24,8 @@ class TestGetBedrockRoute: def test_nova2_prefix_routes_to_converse(self): assert BedrockModelInfo.get_bedrock_route(NOVA2_MODEL) == "converse" - def test_plain_arn_routes_to_invoke(self): - # Without spec prefix, ARN doesn't match converse models - result = BedrockModelInfo.get_bedrock_route(f"bedrock/{NOVA_ARN}") - assert result == "invoke" + def test_plain_arn_without_an_invoke_provider_routes_to_converse(self): + assert BedrockModelInfo.get_bedrock_route(f"bedrock/{NOVA_ARN}") == "converse" class TestGetBedrockBaseModel: diff --git a/tests/unit/llms/clf_ai_gateway/test_clf_ai_gateway_chat_transformation.py b/tests/unit/llms/clf_ai_gateway/test_clf_ai_gateway_chat_transformation.py index ce1a1c1b870..bdaad6a8c2a 100644 --- a/tests/unit/llms/clf_ai_gateway/test_clf_ai_gateway_chat_transformation.py +++ b/tests/unit/llms/clf_ai_gateway/test_clf_ai_gateway_chat_transformation.py @@ -7,8 +7,8 @@ import respx import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts from litellm.llms.clf_ai_gateway.chat.transformation import ClfAiGatewayConfig -from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts DEFAULT_API_BASE: Final = "https://api.clfaigateway.dev/v1" CHAT_URL: Final = f"{DEFAULT_API_BASE}/chat/completions" diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index aed808223bd..4e4d5df9fb0 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -144,6 +144,49 @@ async def test_client_payload_error_mid_stream_raises_read_error(): assert mock_response.closed is True +@pytest.mark.asyncio +async def test_bytes_buffered_before_a_dropped_connection_are_yielded_before_the_read_error(): + loop: Final = asyncio.get_running_loop() + content: Final = aiohttp.StreamReader(aiohttp.base_protocol.BaseProtocol(loop), 2**16, loop=loop) + content.feed_data(b"event: message_start\n\n") + content.set_exception(aiohttp.ClientPayloadError("Response payload is not completed")) + mock_response: Final = MockAiohttpResponse() + mock_response.content = content + received: Final[list[bytes]] = [] + + async def _drain() -> None: + async for chunk in AiohttpResponseStream(mock_response): # pyright: ignore[reportArgumentType] # mock response + received.append(chunk) + + with pytest.raises(httpx.ReadError): + await _drain() + + assert received == [b"event: message_start\n\n"] + assert mock_response.closed is True + + +@pytest.mark.asyncio +async def test_bytes_already_read_before_a_dropped_connection_are_not_yielded_again(): + loop: Final = asyncio.get_running_loop() + content: Final = aiohttp.StreamReader(aiohttp.base_protocol.BaseProtocol(loop), 2**16, loop=loop) + content.feed_data(b"event: message_start\n\n") + content.feed_data(b"event: ping\n\n") + assert await content.read(7) == b"event: " + content.set_exception(aiohttp.ClientPayloadError("Response payload is not completed")) + mock_response: Final = MockAiohttpResponse() + mock_response.content = content + received: Final[list[bytes]] = [] + + async def _drain() -> None: + async for chunk in AiohttpResponseStream(mock_response): # pyright: ignore[reportArgumentType] # mock response + received.append(chunk) + + with pytest.raises(httpx.ReadError): + await _drain() + + assert received == [b"message_start\n\nevent: ping\n\n"] + + @pytest.mark.asyncio async def test_client_payload_error_before_first_chunk_raises_read_error(): """A connection reset before any body byte must surface, not yield an empty 200 body""" diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index 3fbb12881da..56af47aeebe 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Iterator from typing import Final @@ -6,6 +7,7 @@ import httpx import litellm import pytest import respx +from litellm._logging import verbose_logger from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.deepseek.chat.transformation import DeepSeekChatConfig @@ -951,3 +953,42 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) == "https://gateway.example.com/deepseek/chat/completions" ) + + +def _budget_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and "Dropping thinking.budget_tokens=" in r.getMessage() + ) + + +def test_enabled_thinking_budget_dropped_for_deepseek_warns_once(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = DeepSeekChatConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4000}}, + optional_params={}, + model="deepseek-reasoner", + drop_params=False, + ) + + warnings: Final = _budget_warnings(caplog) + assert result["thinking"] == {"type": "enabled"} + assert len(warnings) == 1 + assert "budget_tokens=4000" in warnings[0] + assert "deepseek-reasoner" in warnings[0] + + +@pytest.mark.parametrize( + "thinking", ({"type": "enabled"}, {"type": "disabled", "budget_tokens": 4000}), ids=("no-budget", "disabled") +) +def test_deepseek_thinking_without_a_requested_budget_does_not_warn( + thinking: dict[str, object], caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = DeepSeekChatConfig().map_openai_params( + non_default_params={"thinking": thinking}, optional_params={}, model="deepseek-reasoner", drop_params=False + ) + + assert result["thinking"] == {"type": thinking["type"]} + assert _budget_warnings(caplog) == () diff --git a/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py b/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py index b7a79f03f8d..f4a7455ff27 100644 --- a/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py +++ b/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py @@ -1,3 +1,6 @@ +import logging +from typing import Final + import httpx import pytest @@ -242,11 +245,33 @@ def test_map_openai_params_maps_reasoning_effort(effort, expected): assert mapped["reasoning"] is expected -def test_map_openai_params_drops_unknown_reasoning_effort_when_dropping(): - mapped = FalAIChatConfig().map_openai_params( - non_default_params={"reasoning_effort": "extreme"}, optional_params={}, model=MODEL, drop_params=True - ) +@pytest.mark.parametrize("drop_globally", [False, True]) +def test_map_openai_params_drops_unknown_reasoning_effort_with_one_warning( + drop_globally: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_globally) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + mapped: Final = FalAIChatConfig().map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model=MODEL, + drop_params=not drop_globally, + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoning" not in mapped + assert len(drop_warnings) == 1 + assert all(fragment in drop_warnings[0] for fragment in ("reasoning_effort='xhigh'", MODEL, "(drop_params=True)")) + + +def test_map_openai_params_rejects_unknown_reasoning_effort_without_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(FalAIError) as exc_info: + FalAIChatConfig().map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, optional_params={}, model=MODEL, drop_params=False + ) + assert exc_info.value.status_code == 400 @pytest.mark.parametrize("effort", [{"level": "low"}, ["low"], 1]) diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 28f4fdead9e..ed166d206c8 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Callable from typing import Final from unittest.mock import MagicMock, patch @@ -9,6 +10,7 @@ import respx from pydantic import ValidationError import litellm +from litellm._logging import verbose_logger from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -2361,6 +2363,39 @@ def test_global_disable_flag_with_transform_messages_helper( assert json.loads(route.calls.last.request.content)["messages"] == [{"role": "user", "content": content}] +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +def test_chat_template_kwargs_effort_dropped_for_non_reasoning_model_warns_once( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + FireworksAIConfig().map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"reasoning_budget": 512}}}, + _NON_REASONING_MODEL, + ) + + warnings: Final = _warnings_naming(caplog, "chat_template_kwargs") + assert len(warnings) == 1 + assert "'reasoning_budget': 512" in warnings[0] + assert "llama-v3-70b-instruct" in warnings[0] + + +def test_chat_template_kwargs_thinking_off_for_non_reasoning_model_does_not_warn( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + FireworksAIConfig().map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}}, + _NON_REASONING_MODEL, + ) + + assert _warnings_naming(caplog, "chat_template_kwargs") == () + + _CODEX_DEVELOPER_MESSAGE_AFTER_A_USER_TURN: Final = ( {"role": "system", "content": "You are Codex"}, {"role": "user", "content": "Hi there"}, diff --git a/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py b/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py index e5a77aa8d41..d35442d4095 100644 --- a/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py +++ b/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py @@ -1,7 +1,11 @@ +import logging +from typing import Final + import pytest import litellm +from litellm._logging import verbose_logger from litellm.llms.fireworks_ai.completion.transformation import ( @@ -204,3 +208,36 @@ def test_transform_text_completion_request_keeps_sdk_rejected_keys_in_extra_body assert "chat_template_kwargs" not in data["extra_body"] assert "best_of" not in data["extra_body"] assert "response_format" not in data + + +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +def test_chat_template_kwargs_effort_dropped_for_non_reasoning_model_warns_once( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + FireworksAITextCompletionConfig().map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"reasoning_budget": 512}}}, + _NON_REASONING_MODEL, + ) + + warnings: Final = _warnings_naming(caplog, "chat_template_kwargs") + assert len(warnings) == 1 + assert "'reasoning_budget': 512" in warnings[0] + assert "llama-v3-70b-instruct" in warnings[0] + + +def test_chat_template_kwargs_thinking_off_for_non_reasoning_model_does_not_warn( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + FireworksAITextCompletionConfig().map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}}, + _NON_REASONING_MODEL, + ) + + assert _warnings_naming(caplog, "chat_template_kwargs") == () diff --git a/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 9a62fcf6f0f..7461611e0fa 100644 --- a/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -148,6 +148,41 @@ class TestHostedVLLMRerankTransform: self.config._transform_response(response_dict) +@pytest.mark.parametrize( + "meta", + [ + { + "api_version": {"version": "2"}, + "tokens": {"input_tokens": 7, "output_tokens": 3}, + "billed_units": {"total_tokens": 12}, + }, + {"tokens": {"input_tokens": 0}, "billed_units": {"input_tokens": 0, "output_tokens": 0}}, + {"tokens": None, "billed_units": None}, + {}, + ], +) +@pytest.mark.parametrize("legacy_usage", [{}, {"usage": {"total_tokens": 99}}]) +def test_hosted_vllm_rerank_preserves_native_metadata_over_legacy_usage( + meta: dict[str, object], legacy_usage: dict[str, object] +) -> None: + response: Final = HostedVLLMRerankConfig()._transform_response( + {"id": "rerank-native", "results": [{"index": 0, "relevance_score": 0.9}], "meta": meta, **legacy_usage} + ) + assert response.meta == meta + assert response.model_dump() == { + "id": "rerank-native", + "results": [{"index": 0, "relevance_score": 0.9}], + "meta": meta, + } + + +def test_hosted_vllm_rerank_null_metadata_uses_legacy_usage() -> None: + response: Final = HostedVLLMRerankConfig()._transform_response( + {"results": [], "meta": None, "usage": {"total_tokens": 42}} + ) + assert response.meta == {"billed_units": {"total_tokens": 42}, "tokens": {"input_tokens": 42}} + + class TestGetOptionalRerankParamsInstruction: """`instruction` is threaded through get_optional_rerank_params only when set.""" diff --git a/tests/unit/llms/jina_ai/rerank/test_transformation.py b/tests/unit/llms/jina_ai/rerank/test_transformation.py index 9ec02a4548b..b752f7b4649 100644 --- a/tests/unit/llms/jina_ai/rerank/test_transformation.py +++ b/tests/unit/llms/jina_ai/rerank/test_transformation.py @@ -46,7 +46,10 @@ def test_transform_rerank_response_maps_results_and_usage(): ({"results": [], "usage": {}}, {"billed_units": {}, "tokens": {}}), ( {"results": [], "usage": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3, "unknown": 1}}, - {"billed_units": {"total_tokens": 12}, "tokens": {"input_tokens": 7, "output_tokens": 3}}, + { + "billed_units": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3}, + "tokens": {"input_tokens": 7, "output_tokens": 3}, + }, ), ], ) diff --git a/tests/unit/llms/microsoft_365_copilot/chat/test_handler.py b/tests/unit/llms/microsoft_365_copilot/chat/test_handler.py index 6c41a6dcad5..300e6fc58dd 100644 --- a/tests/unit/llms/microsoft_365_copilot/chat/test_handler.py +++ b/tests/unit/llms/microsoft_365_copilot/chat/test_handler.py @@ -324,7 +324,7 @@ def test_partial_token_exchange_configuration_fails_without_http_calls() -> None def test_api_base_override_clears_exchange_config_and_fails_closed_without_token_post() -> None: - from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + from litellm.litellm_core_utils.clientside_credential_handler import get_dynamic_litellm_params assertion: Final = "header-base-override.payload-base-override.signature-base-override" oauth_fields: Final = ( diff --git a/tests/unit/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py index c340f77c317..720a1b4f458 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -1,4 +1,5 @@ import json +import logging from collections.abc import Mapping from typing import Final, List, cast from unittest.mock import MagicMock, patch @@ -9,6 +10,7 @@ import pytest import respx from openai import AsyncOpenAI +from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY from litellm.llms.mistral.chat.transformation import ( MistralChatResponseIterator, @@ -1670,6 +1672,109 @@ def test_mistral_transform_request_drops_an_empty_tool_message_whose_call_id_is_ assert [m.get("role") for m in request["messages"]] == ["user", "assistant"] +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +@pytest.mark.parametrize( + "param, value", + ( + ("reasoning_effort", "none"), + ("reasoning_effort", "xhigh"), + ("thinking", {"type": "disabled"}), + ("thinking", {"type": ["enabled"]}), + ), +) +def test_magistral_reasoning_request_replaced_by_prompt_warns_once( + param: str, value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = MistralConfig().map_openai_params( + non_default_params={param: value}, + optional_params={}, + model="magistral-medium-latest", + drop_params=False, + ) + + warnings: Final = _warnings_naming(caplog, param) + assert result["_add_reasoning_prompt"] is True + assert len(warnings) == 1 + assert "magistral-medium-latest" in warnings[0] + + +@pytest.mark.parametrize( + "param, value", + ( + ("reasoning_effort", "low"), + ("reasoning_effort", "medium"), + ("reasoning_effort", "high"), + ("thinking", {"type": "enabled", "budget_tokens": 1024}), + ), +) +def test_magistral_reasoning_request_the_prompt_honors_does_not_warn( + param: str, value: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = MistralConfig().map_openai_params( + non_default_params={param: value}, + optional_params={}, + model="magistral-medium-latest", + drop_params=False, + ) + + assert result["_add_reasoning_prompt"] is True + assert _warnings_naming(caplog, param) == () + + +def test_non_magistral_reasoning_effort_is_sent_without_warning(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = MistralConfig().map_openai_params( + non_default_params={"reasoning_effort": "high"}, + optional_params={}, + model="mistral-medium-latest", + drop_params=False, + ) + + assert result["reasoning_effort"] == "high" + assert _warnings_naming(caplog, "reasoning_effort") == () + + +@pytest.mark.parametrize("reasoning_effort", [5, {"effort": 5}, {"effort": 5, "summary": "auto"}]) +def test_mistral_warns_once_when_drop_params_drops_a_level_it_cannot_send( + caplog: pytest.LogCaptureFixture, reasoning_effort: object +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = MistralConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort}, + optional_params={}, + model="mistral-medium-latest", + drop_params=True, + ) + + assert "reasoning_effort" not in result + warnings: Final = _warnings_naming(caplog, "reasoning_effort") + assert len(warnings) == 1 + assert "model='mistral-medium-latest'" in warnings[0] + assert "(drop_params=True)" in warnings[0] + + +@pytest.mark.parametrize("reasoning_effort", [{"summary": "auto"}, {"effort": "high", "summary": "auto"}]) +def test_mistral_stays_quiet_under_drop_params_when_no_level_is_lost( + caplog: pytest.LogCaptureFixture, reasoning_effort: object +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + MistralConfig().map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort}, + optional_params={}, + model="mistral-medium-latest", + drop_params=True, + ) + + assert _warnings_naming(caplog, "reasoning_effort") == () + + @pytest.mark.parametrize( ("model", "custom_llm_provider"), [ diff --git a/tests/unit/llms/oci/chat/test_oci_chat_transformation.py b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py index 0535279d177..f22b5e32593 100644 --- a/tests/unit/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py @@ -1,4 +1,5 @@ import datetime +import logging from typing import Final import httpx import pytest @@ -1103,19 +1104,25 @@ class TestOCIReasoningEffort: ) assert chat_request.get("reasoningEffort") == "HIGH" - def test_reasoning_effort_unsupported_on_cohere_dropped(self): - """drop_params=True → silently drop reasoning_effort for Cohere.""" + def test_reasoning_effort_unsupported_on_cohere_dropped_with_one_warning(self, caplog): from litellm.llms.oci.chat.transformation import OCIChatConfig config = OCIChatConfig() - result = config.map_openai_params( - non_default_params={"reasoning_effort": "low"}, - optional_params={}, - model="cohere.command-latest", - drop_params=True, - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result = config.map_openai_params( + non_default_params={"reasoning_effort": "low"}, + optional_params={}, + model="cohere.command-latest", + drop_params=True, + ) + drop_warnings = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoning_effort" not in result assert "reasoningEffort" not in result + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in ("reasoning_effort='low'", "cohere.command-latest", "(drop_params=True)") + ) def test_reasoning_effort_unsupported_on_cohere_raises(self): """drop_params=False → raise rather than ship a payload Cohere will reject.""" @@ -2110,16 +2117,20 @@ class TestReasoningEffortGating: assert "reasoning_effort" not in OCIChatConfig().get_supported_openai_params("meta.llama-3.3-70b-instruct") @pytest.mark.parametrize("model", SERVED_GROK_4_20_MODELS) - def test_reasoning_effort_dropped_under_drop_params(self, local_model_cost_map, model): - result = OCIChatConfig().map_openai_params( - non_default_params={"reasoning_effort": "low", "temperature": 0.1}, - optional_params={}, - model=model, - drop_params=True, - ) + def test_reasoning_effort_dropped_under_drop_params_with_one_warning(self, local_model_cost_map, model, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result = OCIChatConfig().map_openai_params( + non_default_params={"reasoning_effort": "low", "temperature": 0.1}, + optional_params={}, + model=model, + drop_params=True, + ) + drop_warnings = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] assert "reasoningEffort" not in result assert "reasoning_effort" not in result assert result["temperature"] == 0.1 + assert len(drop_warnings) == 1 + assert all(fragment in drop_warnings[0] for fragment in ("reasoning_effort='low'", model, "(drop_params=True)")) def test_reasoning_effort_raises_without_drop_params_for_gated_grok(self, local_model_cost_map, monkeypatch): monkeypatch.setattr(litellm, "drop_params", False) diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index df9f5fd9d14..6a80f29ac43 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1,5 +1,6 @@ import asyncio import inspect +import logging import os import sys from typing import Final, cast @@ -23,6 +24,7 @@ import json from unittest.mock import ANY, MagicMock, patch import litellm +from litellm._logging import verbose_logger from litellm.types.utils import Choices, EmbeddingResponse, Message, ModelResponse, ModelResponseStream from unittest import mock @@ -854,6 +856,99 @@ class TestOllamaReasoningContentStreaming: class TestOllamaToolCallTransformation: + @pytest.mark.parametrize("stream", [False, True]) + def test_transform_request_keeps_parallel_call_ids_paired_with_results(self, stream: bool): + config: Final = OllamaChatConfig() + messages: Final = cast( + list[AllMessageValues], + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_first", + "type": "function", + "function": { + "name": "read_fixture", + "arguments": '{"path": "first.txt"}', + }, + }, + { + "id": "call_second", + "type": "function", + "function": { + "name": "read_fixture", + "arguments": '{"path": "second.txt"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_second", + "content": "second result", + }, + { + "role": "tool", + "tool_call_id": "call_first", + "content": "first result", + }, + ], + ) + + result: Final = config.transform_request( + model="ollama_chat/test-model", + messages=messages, + optional_params={"stream": stream}, + litellm_params={}, + headers={}, + ) + + assert result["messages"][0]["tool_calls"] == [ + { + "id": "call_first", + "function": {"name": "read_fixture", "arguments": {"path": "first.txt"}}, + }, + { + "id": "call_second", + "function": {"name": "read_fixture", "arguments": {"path": "second.txt"}}, + }, + ] + assert [message.get("tool_call_id") for message in result["messages"][1:]] == [ + "call_second", + "call_first", + ] + + def test_transform_request_does_not_invent_an_id_for_an_unidentified_call(self): + config: Final = OllamaChatConfig() + messages: Final = cast( + list[AllMessageValues], + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": None, + "type": "function", + "function": {"name": "read_fixture", "arguments": "{}"}, + } + ], + } + ], + ) + + result: Final = config.transform_request( + model="ollama_chat/test-model", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["messages"][0]["tool_calls"] == [{"function": {"name": "read_fixture", "arguments": {}}}] + def test_transform_request_preserves_tool_calls(self): """ tool_calls on assistant messages must survive transform_request. @@ -1114,3 +1209,36 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) == "http://gateway.example.com/ollama/api/chat" ) + + +def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: + return tuple( + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() + ) + + +@pytest.mark.parametrize("effort", ("minimal", "xhigh")) +def test_reasoning_effort_mapped_to_think_off_warns_once(effort: str, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = OllamaChatConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model="qwen3", drop_params=False + ) + + warnings: Final = _warnings_naming(caplog, "reasoning_effort") + assert result["think"] is False + assert len(warnings) == 1 + assert f"reasoning_effort={effort!r}" in warnings[0] + assert "qwen3" in warnings[0] + + +@pytest.mark.parametrize("effort, think", (("none", False), ("low", True), ("high", True))) +def test_reasoning_effort_with_a_matching_think_value_does_not_warn( + effort: str, think: bool, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = OllamaChatConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model="qwen3", drop_params=False + ) + + assert result["think"] is think + assert _warnings_naming(caplog, "reasoning_effort") == () diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index 248e745ebff..edea39bb10c 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -1,6 +1,7 @@ import base64 import io import json +import logging import sys from typing import Final from unittest.mock import MagicMock, patch @@ -10,6 +11,7 @@ import pytest import respx import litellm +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.ollama.completion.transformation import ( @@ -990,3 +992,40 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) == "http://gateway.example.com/ollama/api/generate" ) + +def _reasoning_effort_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and "Dropping reasoning_effort=" in r.getMessage() + ) + + +@pytest.mark.parametrize("effort", ("minimal", "xhigh")) +def test_reasoning_effort_mapped_to_think_off_warns_once(effort: str, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = OllamaConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model="qwen3:0.6b", drop_params=False + ) + + warnings: Final = _reasoning_effort_warnings(caplog) + assert result["think"] is False + assert len(warnings) == 1 + assert f"reasoning_effort={effort!r}" in warnings[0] + assert "qwen3:0.6b" in warnings[0] + + +@pytest.mark.parametrize( + "model, effort, think", + (("qwen3:0.6b", "none", False), ("qwen3:0.6b", "high", True), ("gpt-oss:20b", "xhigh", "xhigh")), +) +def test_reasoning_effort_ollama_honors_does_not_warn( + model: str, effort: str, think: object, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = OllamaConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model=model, drop_params=False + ) + + assert result["think"] == think + assert _reasoning_effort_warnings(caplog) == () diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index 673e880a1c0..d732d922375 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -724,16 +724,23 @@ class TestToolResultItemsReachTextsGuardrails: assert blocked.value.status_code == 400 @pytest.mark.asyncio - @pytest.mark.parametrize("empty_output", ["", None], ids=["empty_string", "null"]) - async def test_empty_tool_output_sends_no_text_and_stays_as_sent(self, empty_output: str | None) -> None: + @pytest.mark.parametrize(("call_item", "output_type"), _TOOL_TURNS, ids=["function", "custom"]) + @pytest.mark.parametrize( + "empty_output", + ["", None, [], *[[{"type": kind, "text": ""}] for kind in ("input_text", "output_text", "text")]], + ids=["empty_string", "null", "empty_list", "input_text", "output_text", "text"], + ) + async def test_empty_tool_output_sends_no_text_and_stays_as_sent( + self, call_item: dict[str, str], output_type: str, empty_output: object + ) -> None: guardrail: Final = RecordingMaskingGuardrail(guardrail_name="test") - call_item, output_type = _TOOL_TURNS[0] result: Final = await OpenAIResponsesHandler().process_input_messages( _tool_result_request(call_item, output_type, empty_output), guardrail ) assert guardrail.seen_texts == [["Weather?"]] + assert guardrail.seen_message_contents == [["Weather?", None]] assert result["input"][2] == {"type": output_type, "call_id": "call_1", "output": empty_output} @pytest.mark.asyncio @@ -1103,7 +1110,6 @@ class TestToolResultItemsReachTextsGuardrails: ["visible"], id="non_dict_part", ), - pytest.param([], ["[]"], id="empty"), pytest.param({"result": "SECRET forecast"}, ['{"result": "SECRET forecast"}'], id="dict"), pytest.param( [{"type": "input_image", "file_id": "file-map"}], diff --git a/tests/unit/llms/openai/test_gpt_reasoning_transformation.py b/tests/unit/llms/openai/test_gpt_reasoning_transformation.py index 807a43c5d9d..663b443fbaf 100644 --- a/tests/unit/llms/openai/test_gpt_reasoning_transformation.py +++ b/tests/unit/llms/openai/test_gpt_reasoning_transformation.py @@ -1405,7 +1405,7 @@ def test_gpt5_6_never_advertises_reasoning_effort_max(model: str): """/v1/chat/completions answers max with "Unsupported value: 'reasoning_effort' does not support 'max' with this model. Supported values are: 'none', 'low', 'medium', 'high', and 'xhigh'", so no gpt-5.6 entry asserts supports_max_reasoning_effort and the advertised set stops at xhigh.""" - from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + from litellm.litellm_core_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts resolved = resolve_supported_reasoning_efforts(litellm.get_model_info(model), deployment_is_mapped=True) assert resolved is not None diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 3e1cbad9308..77a7c3b83fd 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -24,6 +24,7 @@ from litellm.llms.openai.openai import ( AsyncAssistantEventHandler, AsyncCursorPage, MessageData, + OpenAIBatchesAPI, OpenAIChatCompletion, OpenAIConfig, OpenAIMessage as Message, @@ -2320,3 +2321,43 @@ def test_openai_config_hoists_a_later_developer_message_only_for_a_model_off_ope model_info=None, ) assert [message["role"] for message in messages] == expected_roles + + +@pytest.mark.asyncio +async def test_sync_batch_list_given_an_async_client_returns_a_page_the_caller_can_await( + http_backend: ModuleType, +): + def respond(request: _Request) -> httpx.Response | SDKResponse: + assert str(request.url).startswith("http://vllm-router.test/v1/batches") + return http_backend.Response( + 200, + json={ + "object": "list", + "data": [ + { + "id": "batch_5f8dac", + "status": "pending", + "input_file_id": "file-b73f4c", + "created_at": 1793239207, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + } + ], + "first_id": "batch_5f8dac", + "last_id": "batch_5f8dac", + "has_more": False, + }, + ) + + async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client: + page: Final = OpenAIBatchesAPI().list_batches( + _is_async=False, + api_key="sk-test", + api_base="http://vllm-router.test/v1", + timeout=5.0, + max_retries=0, + organization=None, + client=AsyncOpenAI(api_key="sk-test", base_url="http://vllm-router.test/v1", http_client=http_client), + ) + + assert [(batch.id, batch.status) for batch in (await page).data] == [("batch_5f8dac", "pending")] diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 46c14fcfd8f..f0865cca7f9 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -1,4 +1,6 @@ import asyncio +import json +import logging from typing import Final from unittest.mock import MagicMock, patch @@ -9,7 +11,11 @@ import respx import litellm from litellm.litellm_core_utils.token_counter import token_counter -from litellm.llms.openai.common_utils import BaseOpenAILLM, is_openai_backed_api_base +from litellm.llms.openai.common_utils import ( + BaseOpenAILLM, + drop_params_from_unprocessable_entity_error, + is_openai_backed_api_base, +) # Test parameters for different API functions API_FUNCTION_PARAMS = [ @@ -319,6 +325,104 @@ def _completion_kwargs(provider: str, client, **overrides) -> dict: } +_REASONING_EFFORT_422: Final = { + "error": { + "message": json.dumps( + {"detail": [{"loc": ["body", "reasoning_effort"], "msg": "Extra inputs are not permitted"}]} + ) + } +} +_CHAT_COMPLETION: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "gpt-6.1-sol", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_CHAT_COMPLETION_STREAM: Final = ( + 'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,"model":"gpt-6.1-sol",' + '"choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}\n\n' + "data: [DONE]\n\n" +) + + +def _answer_422_while_reasoning_effort_is_sent(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + if "reasoning_effort" in body: + return httpx.Response(422, json=_REASONING_EFFORT_422) + if body.get("stream"): + return httpx.Response(200, text=_CHAT_COMPLETION_STREAM, headers={"content-type": "text/event-stream"}) + return httpx.Response(200, json=_CHAT_COMPLETION) + + +async def _complete_through_the_422_retry(call_shape: str) -> str: + transport: Final = httpx.MockTransport(_answer_422_while_reasoning_effort_is_sent) + kwargs: Final = { + "model": "gpt-6.1-sol", + "messages": [{"role": "user", "content": "hi"}], + "reasoning_effort": "high", + } + if call_shape == "sync": + sync_client: Final = openai.OpenAI(api_key="sk-not-a-real-key", http_client=httpx.Client(transport=transport)) + return litellm.completion(**kwargs, client=sync_client).choices[0].message.content + async_client: Final = openai.AsyncOpenAI( + api_key="sk-not-a-real-key", http_client=httpx.AsyncClient(transport=transport) + ) + if call_shape == "async": + return (await litellm.acompletion(**kwargs, client=async_client)).choices[0].message.content + stream: Final = await litellm.acompletion(**kwargs, client=async_client, stream=True) + return "".join([chunk.choices[0].delta.content or "" async for chunk in stream]) + + +@pytest.mark.parametrize("call_shape", ["sync", "async", "async_stream"]) +def test_a_422_naming_reasoning_effort_drops_it_with_one_warning( + call_shape: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + content: Final = asyncio.run(_complete_through_the_422_retry(call_shape)) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert content == "ok" + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in ("reasoning_effort='high'", "gpt-6.1-sol", "(drop_params=True)", "422") + ) + + +def test_a_422_with_an_unhashable_loc_still_drops_reasoning_effort_with_one_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + detail: Final = [ + {"loc": ["body", ["messages", 0]], "msg": "Invalid"}, + {"loc": ["body", "reasoning_effort"], "msg": "Extra inputs are not permitted"}, + ] + + def answer(request: httpx.Request) -> httpx.Response: + if "reasoning_effort" in json.loads(request.content): + return httpx.Response(422, json={"error": {"message": json.dumps({"detail": detail})}}) + return httpx.Response(200, json=_CHAT_COMPLETION) + + monkeypatch.setattr(litellm, "drop_params", True) + client: Final = openai.OpenAI( + api_key="sk-not-a-real-key", http_client=httpx.Client(transport=httpx.MockTransport(answer)) + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + response: Final = litellm.completion( + model="gpt-6.1-sol", + messages=[{"role": "user", "content": "hi"}], + reasoning_effort="high", + client=client, + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert response.choices[0].message.content == "ok" + assert len(drop_warnings) == 1 + assert "reasoning_effort='high'" in drop_warnings[0] + + @pytest.mark.parametrize("provider", ["openai", "azure"]) def test_sync_output_limit_400_maps_to_length_truncated_response(provider): response = litellm.completion( @@ -588,3 +692,24 @@ async def test_azure_gateway_and_image_variations_use_the_callers_async_session( ) assert response.data[0].url == "https://example.com/image.png" assert not session.is_closed + + + +def test_reasoning_drop_warning_caps_a_large_caller_value(caplog: pytest.LogCaptureFixture) -> None: + error: Final = httpx.HTTPStatusError( + message="422", + request=httpx.Request("POST", "https://example.invalid/v1/chat/completions"), + response=httpx.Response(422, json=_REASONING_EFFORT_422), + ) + data: Final[dict[str, object]] = {"model": "gpt-6.1-sol", "messages": [], "reasoning_effort": "x" * 5000} + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + retried: Final = drop_params_from_unprocessable_entity_error( + error, data, model="gpt-6.1-sol", drop_params_gated=True + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + + assert "reasoning_effort" not in retried + assert len(drop_warnings) == 1 + assert len(drop_warnings[0]) < 300 + assert all(fragment in drop_warnings[0] for fragment in ("reasoning_effort='xxx", "'gpt-6.1-sol'", "422")) diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index 1a1a4ca7047..b5810b7e3d1 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -974,7 +974,7 @@ class TestDiscoverModels: class TestClientsideBaseOverride: def test_client_api_base_override_clears_deployment_wif(self, deployment_wif: dict[str, str]) -> None: - from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + from litellm.litellm_core_utils.clientside_credential_handler import get_dynamic_litellm_params redirected: Final = get_dynamic_litellm_params( litellm_params={"model": "openai/gpt-4o-mini", **deployment_wif}, diff --git a/tests/unit/llms/openai_like/test_aiand_provider.py b/tests/unit/llms/openai_like/test_aiand_provider.py index 804f3a3b6eb..0274b58f19a 100644 --- a/tests/unit/llms/openai_like/test_aiand_provider.py +++ b/tests/unit/llms/openai_like/test_aiand_provider.py @@ -1,4 +1,6 @@ +import datetime import json +import time from typing import Final import httpx @@ -8,6 +10,8 @@ import respx import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.health_check_helpers import wildcard_health_check_models +from litellm.litellm_core_utils.model_retirement import is_past_deprecation_date from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.proxy.auth.model_checks import get_known_models_from_wildcard from litellm.utils import ProviderConfigManager @@ -89,11 +93,99 @@ def _aiand_catalog() -> frozenset[str]: return frozenset(name for name, info in litellm.model_cost.items() if info.get("litellm_provider") == "aiand") -def test_aiand_wildcard_lists_its_priced_models(): - catalog: Final = _aiand_catalog() +def _utc_today() -> datetime.date: + return datetime.datetime.now(datetime.timezone.utc).date() - assert catalog - assert frozenset(get_known_models_from_wildcard("aiand/*")) == catalog + +def test_aiand_wildcard_lists_only_models_not_yet_retired(): + catalog: Final = _aiand_catalog() + retired: Final = frozenset( + name + for name in catalog + if (retired_on := litellm.model_cost[name].get("deprecation_date")) is not None + and datetime.date.fromisoformat(retired_on) <= _utc_today() + ) + + assert catalog - retired + assert frozenset(get_known_models_from_wildcard("aiand/*")) == catalog - retired + + +def test_aiand_wildcard_drops_a_model_whose_retirement_date_passes_while_running(monkeypatch: pytest.MonkeyPatch): + model: Final = "aiand/acme/retires-while-running" + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "aiand", "deprecation_date": "2999-01-01"}) + try: + litellm.add_known_models({model: litellm.model_cost[model]}) + assert model in get_known_models_from_wildcard("aiand/*") + + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "aiand", "deprecation_date": "2000-01-01"}) + + assert model not in get_known_models_from_wildcard("aiand/*") + finally: + litellm.aiand_models.discard(model) + litellm.add_known_models({}) + + +def test_aiand_models_by_provider_leaves_out_models_retired_at_load(): + retired: Final = frozenset( + name + for name, row in litellm.model_cost.items() + if row.get("litellm_provider") == "aiand" and is_past_deprecation_date(row) + ) + + assert retired + assert retired.isdisjoint(litellm.models_by_provider["aiand"]) + + +def test_aiand_wildcard_health_check_skips_a_model_whose_retirement_date_passes_while_running( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "aiand/acme/free-and-retiring" + live_row: Final = { + "litellm_provider": "aiand", + "mode": "chat", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "deprecation_date": "2999-01-01", + } + monkeypatch.setitem(litellm.model_cost, model, live_row) + try: + litellm.add_known_models({model: live_row}) + assert model in wildcard_health_check_models(wildcard="aiand/*", custom_llm_provider="aiand") + + monkeypatch.setitem(litellm.model_cost, model, {**live_row, "deprecation_date": "2000-01-01"}) + + assert model not in wildcard_health_check_models(wildcard="aiand/*", custom_llm_provider="aiand") + finally: + litellm.aiand_models.discard(model) + litellm.add_known_models({}) + + +@pytest.mark.parametrize("host_zone", ["Etc/GMT+12", "Etc/GMT-14"]) +def test_aiand_models_retire_on_the_utc_date_whatever_the_host_zone(monkeypatch: pytest.MonkeyPatch, host_zone: str): + today: Final = _utc_today() + rows: Final = { + "aiand/acme/retires-today": {"litellm_provider": "aiand", "deprecation_date": today.isoformat()}, + "aiand/acme/retires-tomorrow": { + "litellm_provider": "aiand", + "deprecation_date": (today + datetime.timedelta(days=1)).isoformat(), + }, + } + for name, row in rows.items(): + monkeypatch.setitem(litellm.model_cost, name, row) + monkeypatch.setenv("TZ", host_zone) + time.tzset() + try: + litellm.add_known_models(rows) + listed: Final = frozenset(get_known_models_from_wildcard("aiand/*")) + + assert "aiand/acme/retires-today" not in listed + assert "aiand/acme/retires-tomorrow" in listed + finally: + monkeypatch.undo() + time.tzset() + for name in rows: + litellm.aiand_models.discard(name) + litellm.add_known_models({}) def test_aiand_calls_are_billed_from_their_cost_rows(): diff --git a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py index 59355dfabcb..3d659641f54 100644 --- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -1,6 +1,7 @@ import json from typing import Final +from unittest.mock import MagicMock import httpx import pytest @@ -464,6 +465,34 @@ def test_openrouter_cost_tracking_non_streaming(): ) + +@pytest.mark.parametrize("cost", [float("nan"), float("inf")]) +def test_openrouter_chat_leaves_a_non_finite_provider_cost_out_of_the_response_headers(cost: float) -> None: + body: Final = { + "id": "gen-123", + "object": "chat.completion", + "created": 1234567890, + "model": "openai/gpt-5.6-luna", + "choices": [{"message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop", "index": 0}], + "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11, "cost": cost}, + } + + result: Final = OpenrouterConfig().transform_response( + model="openrouter/openai/gpt-5.6-luna", + raw_response=httpx.Response( + 200, content=json.dumps(body).encode(), headers={"content-type": "application/json"} + ), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "Reply with ok"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert "llm_provider-x-litellm-response-cost" not in result._hidden_params.get("additional_headers", {}) + def test_openrouter_cost_tracking_streaming(): """ Test OpenRouter cost tracking for streaming completions. @@ -843,3 +872,120 @@ def test_openrouter_streaming_reasoning_is_preserved() -> None: assembled: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages) assert assembled is not None assert assembled.choices[0].message.reasoning_content == "Consider the evidence. Apply the rule." + + +def _transform_openrouter_response(body: dict[str, object]) -> litellm.ModelResponse: + return OpenrouterConfig().transform_response( + model="openrouter/openai/gpt-4o-mini", + raw_response=httpx.Response(200, json=body), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_openrouter_nonstream_provider_without_usage_is_preserved_regression() -> None: + response: Final = _transform_openrouter_response( + { + "id": "gen-provider-no-usage", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o-mini", + "provider": "serving-site-a", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + } + ) + + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + + +@pytest.mark.parametrize( + ("provider_fields", "case"), + [ + pytest.param({}, "missing", id="missing"), + pytest.param({"provider": None}, "null", id="null"), + pytest.param({"provider": ""}, "empty", id="empty"), + pytest.param({"provider": 123}, "non-string", id="non-string"), + ], +) +def test_openrouter_invalid_provider_keeps_reported_cost_guard( + provider_fields: dict[str, object], case: str +) -> None: + response: Final = _transform_openrouter_response( + { + "id": f"gen-invalid-provider-{case}", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o-mini", + **provider_fields, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + "cost": 0.00015, + }, + } + ) + + assert response._hidden_params.get("upstream_provider") is None + assert response._hidden_params["additional_headers"][OPENROUTER_COST_HEADER] == 0.00015 + + +def test_openrouter_streaming_chunk_preserves_provider_regression() -> None: + chunk: Final = { + "id": "gen-provider-stream", + "object": "chat.completion.chunk", + "created": 1, + "model": "openai/gpt-4o-mini", + "provider": "serving-site-a", + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + response: Final = OpenRouterChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ).chunk_parser(chunk) + + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + + +@pytest.mark.parametrize( + ("provider_fields", "case"), + [ + pytest.param({}, "missing", id="missing"), + pytest.param({"provider": None}, "null", id="null"), + pytest.param({"provider": ""}, "empty", id="empty"), + pytest.param({"provider": 123}, "non-string", id="non-string"), + ], +) +def test_openrouter_streaming_chunk_ignores_invalid_provider_guard(provider_fields: dict[str, object], case: str) -> None: + chunk: Final = { + "id": f"gen-invalid-provider-{case}", + "object": "chat.completion.chunk", + "created": 1, + "model": "openai/gpt-4o-mini", + **provider_fields, + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + response: Final = OpenRouterChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ).chunk_parser(chunk) + + assert response._hidden_params.get("upstream_provider") is None diff --git a/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py index 19517e58b31..78d7161c388 100644 --- a/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py +++ b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py @@ -9,6 +9,8 @@ reasoning.encrypted_content for multi-turn stateless workflows. Related issue: https://github.com/BerriAI/litellm/issues/22189 """ +import json +from typing import Final from unittest.mock import MagicMock import httpx @@ -157,6 +159,29 @@ def test_openrouter_responses_returns_provider_reported_cost(cost): assert get_response_cost_from_hidden_params(response.hidden_params) == cost +@pytest.mark.parametrize("cost", [float("nan"), float("inf")]) +def test_openrouter_responses_drops_a_non_finite_provider_cost(cost: float) -> None: + body: Final = { + "id": "resp_test", + "created_at": 1, + "model": "unmapped-model", + "object": "response", + "output": [], + "status": "completed", + "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5, "cost": cost}, + } + response: Final = OpenRouterResponsesAPIConfig().transform_response_api_response( + model="openrouter/unmapped-model", + raw_response=httpx.Response( + 200, content=json.dumps(body).encode(), headers={"content-type": "application/json"} + ), + logging_obj=MagicMock(), + ) + + assert get_response_cost_from_hidden_params(response.hidden_params) is None + assert json.loads(json.dumps(response.model_dump(), allow_nan=False, default=str))["usage"]["cost"] is None + + @pytest.mark.parametrize( "event_type", ["response.completed", "response.incomplete", "response.failed"], diff --git a/tests/unit/llms/openrouter/test_openrouter_common_utils.py b/tests/unit/llms/openrouter/test_openrouter_common_utils.py new file mode 100644 index 00000000000..087c05561b5 --- /dev/null +++ b/tests/unit/llms/openrouter/test_openrouter_common_utils.py @@ -0,0 +1,159 @@ +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.openrouter.common_utils import OpenRouterModelInfo + +DEFAULT_BASE: Final = "https://openrouter.ai/api/v1" +KEY_SOURCES: Final = ("api_key", "litellm.api_key", "litellm.openrouter_key", "OPENROUTER_API_KEY", "OR_API_KEY") +BASE_SOURCES: Final = ("api_base", "litellm.api_base", "OPENROUTER_API_BASE") + + +@pytest.fixture(autouse=True) +def _no_ambient_openrouter_settings(monkeypatch: pytest.MonkeyPatch): + for env_var in ("OPENROUTER_API_BASE", "OPENROUTER_API_KEY", "OR_API_KEY"): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setattr(litellm, "openrouter_key", None) + + +def _catalog(*model_ids: str) -> httpx.Response: + return httpx.Response(200, json={"data": [{"id": model_id, "name": model_id} for model_id in model_ids]}) + + +def _configure(monkeypatch: pytest.MonkeyPatch, source: str, value: str) -> None: + if source.startswith("litellm."): + monkeypatch.setattr(litellm, source.removeprefix("litellm."), value) + else: + monkeypatch.setenv(source, value) + + +def test_get_models_lists_only_the_chat_catalog_under_openrouter_keeping_the_org_namespace(): + with respx.mock() as respx_mock: + respx_mock.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-test") == [ + "openrouter/openai/gpt-5.4-mini", + "openrouter/anthropic/claude-haiku-5-5", + ] + assert [call.request.url.path for call in respx_mock.calls] == ["/api/v1/models/user"] + + +@respx.mock +def test_get_models_drops_batch_only_ids_and_keeps_every_other_variant(): + respx.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini:batch", "openai/gpt-5.4-mini", "qwen/qwen3-coder:free") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-test") == [ + "openrouter/openai/gpt-5.4-mini", + "openrouter/qwen/qwen3-coder:free", + ] + + +@pytest.mark.parametrize( + "api_base", + [ + "https://gateway.example/openrouter/api/v1/", + "https://gateway.example/openrouter/api/v1/chat/completions", + "https://gateway.example/openrouter/api/v1/chat/completions/", + ], +) +@respx.mock +def test_get_models_keeps_the_api_base_path_and_sends_the_deployment_key(api_base: str): + chat_route: Final = respx.get("https://gateway.example/openrouter/api/v1/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-deployment", api_base=api_base) == [ + "openrouter/openai/gpt-5.4-mini" + ] + assert chat_route.calls.last.request.headers["authorization"] == "Bearer sk-or-deployment" + + +@pytest.mark.parametrize("winner", range(len(KEY_SOURCES)), ids=KEY_SOURCES) +@respx.mock +def test_get_models_sends_the_first_configured_key_in_the_chat_paths_order( + monkeypatch: pytest.MonkeyPatch, winner: int +): + for source in KEY_SOURCES[winner + 1 :]: + _configure(monkeypatch, source, f"sk-or-from-{source}") + if winner > 0: + _configure(monkeypatch, KEY_SOURCES[winner], "sk-or-winner") + chat_route: Final = respx.get(f"{DEFAULT_BASE}/models/user").mock(return_value=_catalog("qwen/qwen3-coder")) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-winner" if winner == 0 else None) == [ + "openrouter/qwen/qwen3-coder" + ] + assert chat_route.calls.last.request.headers["authorization"] == "Bearer sk-or-winner" + + +@pytest.mark.parametrize("winner", range(len(BASE_SOURCES) + 1), ids=(*BASE_SOURCES, "default")) +@respx.mock +def test_get_models_calls_the_first_configured_base_in_the_chat_paths_order( + monkeypatch: pytest.MonkeyPatch, winner: int +): + for source in BASE_SOURCES[winner + 1 :]: + _configure(monkeypatch, source, f"https://{source.lower().replace('.', '-')}.example/api/v1") + if 0 < winner < len(BASE_SOURCES): + _configure(monkeypatch, BASE_SOURCES[winner], "https://winner.example/api/v1") + expected_base: Final = "https://winner.example/api/v1" if winner < len(BASE_SOURCES) else DEFAULT_BASE + respx.get(f"{expected_base}/models/user").mock(return_value=_catalog("qwen/qwen3-coder")) + + assert OpenRouterModelInfo().get_models( + api_key="sk-or-test", api_base="https://winner.example/api/v1" if winner == 0 else None + ) == ["openrouter/qwen/qwen3-coder"] + + +def test_get_models_without_any_key_raises_before_calling_openrouter(): + with respx.mock(assert_all_called=False) as respx_mock: + chat_route: Final = respx_mock.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini") + ) + + with pytest.raises(ValueError, match="OPENROUTER_API_KEY"): + OpenRouterModelInfo().get_models() + + assert not chat_route.called + + +@respx.mock +def test_get_models_raises_when_the_catalog_answers_an_error(): + respx.get(f"{DEFAULT_BASE}/models/user").mock(return_value=httpx.Response(401, json={"error": {"code": 401}})) + + with pytest.raises(httpx.HTTPStatusError): + OpenRouterModelInfo().get_models(api_key="sk-or-revoked") + + +def test_validate_environment_adds_the_resolved_key_and_keeps_the_callers_headers(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-env") + headers: Final = {"X-Title": "my-app"} + + assert OpenRouterModelInfo().validate_environment( + headers=headers, model="openrouter/openai/gpt-5.4-mini", messages=(), optional_params={}, litellm_params={} + ) == {"X-Title": "my-app", "Authorization": "Bearer sk-or-env"} + assert headers == {"X-Title": "my-app"} + + +def test_validate_environment_without_any_key_sends_no_authorization(): + assert OpenRouterModelInfo().validate_environment( + headers={"X-Title": "my-app"}, + model="openrouter/openai/gpt-5.4-mini", + messages=(), + optional_params={}, + litellm_params={}, + ) == {"X-Title": "my-app"} + + +@pytest.mark.parametrize( + ("model", "base_model"), + [("openrouter/openai/gpt-5.4-mini", "openai/gpt-5.4-mini"), ("openai/gpt-5.4-mini", "openai/gpt-5.4-mini")], +) +def test_get_base_model_drops_only_the_openrouter_prefix(model: str, base_model: str): + assert OpenRouterModelInfo.get_base_model(model) == base_model diff --git a/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py b/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py index b0f1064967a..caf2fcb4c4b 100644 --- a/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py +++ b/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py @@ -72,6 +72,72 @@ def test_search_sends_max_results_in_request_body(respx_mock: respx.MockRouter) assert [result.url for result in response.results] == ["https://example.com/ml"] +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_search_forwards_provider_filters_without_internal_params( + respx_mock: respx.MockRouter, httpx_transport: None, is_async: bool +) -> None: + route: Final = respx_mock.post(PERPLEXITY_SEARCH_URL).respond(json={"results": []}) + filters: Final = { + "search_recency_filter": "week", + "search_after_date_filter": "01/01/2026", + "search_before_date_filter": "10/01/2026", + "last_updated_after_filter": "01/01/2026", + "last_updated_before_filter": "10/01/2026", + "search_language_filter": ["en"], + "search_mode": "academic", + "max_tokens": 2048, + "search_domain_filter": ["europa.eu"], + } + params: Final = { + "query": "AI Act obligations for model providers", + "search_provider": "perplexity", + "metadata": {"private": "local-only"}, + "input_cost_per_query": 0.01, + "_litellm_private_marker": "local-only", + **filters, + } + + if is_async: + response = await litellm.asearch(**params) + else: + response = litellm.search(**params) + + assert response.object == "search" + assert json.loads(route.calls.last.request.content) == { + "query": "AI Act obligations for model providers", + **filters, + } + + +@pytest.mark.parametrize("query", ["original query", ["first query", "second query"]]) +def test_search_request_preserves_query_and_present_values_without_mutating_options(query: str | list[str]) -> None: + optional_params: Final[dict[str, object]] = { + "query": "replacement query", + "search_recency_filter": None, + "search_language_filter": [], + "max_tokens": 0, + "provider_option": {"enabled": False}, + } + original: Final = dict(optional_params) + + body: Final = PerplexitySearchConfig().transform_search_request( + query=query, + optional_params=optional_params, + api_key="transport-only", + api_base="https://transport.invalid", + headers={"Authorization": "transport-only"}, + ) + + assert body == { + "query": query, + "search_language_filter": [], + "max_tokens": 0, + "provider_option": {"enabled": False}, + } + assert optional_params == original + + def _transform(payload: object) -> SearchResponse: logging_obj = Logging( model="perplexity/search", @@ -213,11 +279,15 @@ async def test_router_search_with_search_tools( query="latest AI developments", search_tool_name="litellm-search", max_results=3, + search_recency_filter="week", + search_language_filter=["en"], ) assert json.loads(route.calls.last.request.content) == { "query": "latest AI developments", "max_results": 3, + "search_recency_filter": "week", + "search_language_filter": ["en"], } assert response.object == "search" assert [(result.title, result.url, result.snippet) for result in response.results] == [ diff --git a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py index 83c71479311..5e774b05207 100644 --- a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py @@ -6,9 +6,12 @@ search queries, and reasoning tokens. """ import json +import logging import math import os +from collections.abc import Mapping from datetime import datetime, timezone +from typing import Final from unittest.mock import patch import pytest @@ -117,7 +120,8 @@ class TestPerplexityCostCalculator: assert completion_cost == 0.008 assert prompt_cost + completion_cost == 0.008 - def test_uses_perplexity_provided_cost_when_normalized_to_float(self): + @pytest.mark.parametrize("reported_cost", (0.0, 0.008)) + def test_uses_perplexity_provided_cost_when_normalized_to_float(self, reported_cost: float): """ Regression: for Responses API / Agent API models, `ResponseAPIUsage.parse_cost` (litellm/types/llms/openai.py) already flattens Perplexity's @@ -133,12 +137,75 @@ class TestPerplexityCostCalculator: perplexity/openai/gpt-5.2 before rates existed). """ usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) - usage.cost = 0.008 + usage.cost = reported_cost prompt_cost, completion_cost = perplexity_cost_per_token(model="sonar-pro", usage=usage) assert prompt_cost == 0.0 - assert completion_cost == 0.008 + assert completion_cost == reported_cost + + @pytest.mark.parametrize( + ("legacy_pricing", "expected_search_rate"), + ( + pytest.param({}, 0.01, id="missing"), + pytest.param({"search_queries_cost_per_query": None}, 0.01, id="null"), + pytest.param({"search_queries_cost_per_query": 0.0}, 0.0, id="free"), + pytest.param({"search_queries_cost_per_query": 0.004}, 0.004, id="paid"), + ), + ) + def test_search_cost_preserves_zero_and_defaults_only_when_unset( + self, + monkeypatch: pytest.MonkeyPatch, + legacy_pricing: Mapping[str, float | None], + expected_search_rate: float, + ) -> None: + model: Final = "perplexity/search-pricing-fixture" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "litellm_provider": "perplexity", + "mode": "chat", + "input_cost_per_token": 2e-6, + "output_cost_per_token": 8e-6, + "output_cost_per_reasoning_token": 3e-6, + "citation_cost_per_token": 2e-6, + "search_context_cost_per_query": {"search_context_size_low": 0.01}, + **legacy_pricing, + }, + ) + usage: Final = Usage( + prompt_tokens=100, + completion_tokens=50, + prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=2), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=10), + citation_tokens=25, + ) + + prompt_cost, completion_cost = litellm.cost_per_token(model=model, usage_object=usage) + + assert prompt_cost == pytest.approx((100 + 25) * 2e-6) + assert completion_cost == pytest.approx(40 * 8e-6 + 10 * 3e-6 + 2 * expected_search_rate) + + @pytest.mark.parametrize( + "reported_cost", + [float("nan"), float("inf"), {"total_cost": float("nan")}], + ids=["nan", "inf", "nan-total-cost-dict"], + ) + def test_a_non_finite_perplexity_cost_is_priced_like_no_reported_cost_with_a_warning( + self, reported_cost: float | dict, caplog: pytest.LogCaptureFixture + ): + unreported: Final = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + reported: Final = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + reported.cost = reported_cost + + token_priced: Final = perplexity_cost_per_token(model="sonar-pro", usage=unreported) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + reported_priced: Final = perplexity_cost_per_token(model="sonar-pro", usage=reported) + + assert sum(token_priced) > 0 + assert reported_priced == token_priced + assert [record for record in caplog.records if "non-finite provider-reported cost" in record.getMessage()] OFF_PEAK_MODEL = "sonar-off-peak-test" OFF_PEAK_WINDOW = "14:00-00:00" diff --git a/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py b/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py index 4d6d252ae6e..1bf646fc2ad 100644 --- a/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py +++ b/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py @@ -1,5 +1,10 @@ +import logging +from typing import Final from unittest.mock import patch +import pytest + +from litellm._logging import verbose_logger from litellm.llms.tencent.chat.transformation import TencentChatConfig @@ -208,6 +213,27 @@ class TestAdaptiveThinkingCoercion: assert result["extra_body"]["thinking"] == {"type": "adaptive", "budget_tokens": 4096} + @pytest.mark.parametrize("budget, warned", (("high", 1), (4096, 0)), ids=("non-integer", "integer")) + def test_enabled_thinking_on_adaptive_only_model_warns_only_when_the_budget_is_dropped( + self, local_model_cost_map, budget: object, warned: int, caplog: pytest.LogCaptureFixture + ): + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = TencentChatConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget}}, + optional_params={}, + model="tencent/minimax-m3", + drop_params=False, + ) + + warnings: Final = tuple( + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and f"Dropping thinking.budget_tokens={budget!r}" in r.getMessage() + ) + assert result["extra_body"]["thinking"]["type"] == "adaptive" + assert result["extra_body"]["thinking"].get("budget_tokens") == (None if warned else budget) + assert len(warnings) == warned + def test_disabled_thinking_kept_for_adaptive_only_model(self, local_model_cost_map): config = TencentChatConfig() result = config.map_openai_params( diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 67c78b62151..98d1b8f1300 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1693,6 +1693,78 @@ def test_vertex_ai_map_thinking_param_without_budget_tokens_for_gemini_3(): assert result["thinkingConfig"] == {"includeThoughts": True} +def _budget_drop_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple( + r.getMessage() + for r in caplog.records + if r.levelno == logging.WARNING and "thinking.budget_tokens" in r.getMessage() + ) + + +@pytest.mark.parametrize("budget", ("high", 1024.5, {"tokens": 1024})) +def test_gemini_non_integer_thinking_budget_drop_warns_once(budget: object, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget}}, + optional_params={}, + model="gemini-2.5-flash", + drop_params=False, + ) + + warnings: Final = _budget_drop_warnings(caplog) + assert "thinkingBudget" not in result["thinkingConfig"] + assert len(warnings) == 1 + assert f"thinking.budget_tokens={budget!r}" in warnings[0] + assert "gemini-2.5-flash" in warnings[0] + + +def test_gemini_integer_thinking_budget_is_sent_without_warning(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 2048}}, + optional_params={}, + model="gemini-2.5-flash", + drop_params=False, + ) + + assert result["thinkingConfig"]["thinkingBudget"] == 2048 + assert _budget_drop_warnings(caplog) == () + + +def test_gemini_3_thinking_budget_drop_warns_once(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + result: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 2048}}, + optional_params={}, + model="gemini-3-pro-preview", + drop_params=False, + ) + + warnings: Final = _budget_drop_warnings(caplog) + assert "thinkingBudget" not in result["thinkingConfig"] + assert len(warnings) == 1 + assert "thinking.budget_tokens=2048" in warnings[0] + assert "gemini-3-pro-preview" in warnings[0] + + +@pytest.mark.parametrize( + "thinking", + ({"type": "enabled"}, {"type": "enabled", "budget_tokens": 0}, {"type": "disabled", "budget_tokens": 2048}), +) +def test_gemini_3_thinking_without_a_dropped_budget_does_not_warn( + thinking: dict[str, object], caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + VertexGeminiConfig().map_openai_params( + non_default_params={"thinking": thinking}, + optional_params={}, + model="gemini-3-pro-preview", + drop_params=False, + ) + + assert _budget_drop_warnings(caplog) == () + + def test_vertex_ai_map_tools(): v = VertexGeminiConfig() optional_params = {} diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 152ab2ffc13..3b222be6e5e 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1293,6 +1293,9 @@ def test_get_vertex_model_id_from_url(): model_id = get_vertex_model_id_from_url(url) assert model_id == "gemini-pro" + url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro?key=provider-secret#frag" + assert get_vertex_model_id_from_url(url) == "gemini-pro" + # Test with invalid URL url = "https://invalid-url.com" model_id = get_vertex_model_id_from_url(url) diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 7f16957ddc4..fc57dd13302 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -8,11 +8,13 @@ from threading import get_ident from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import requests import responses import litellm +from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.messages.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, @@ -115,7 +117,11 @@ def _vertex_client() -> MagicMock: async def _send_to_vertex(client: MagicMock, caller_headers: Mapping[str, str], drop_params: bool) -> None: with responses.RequestsMock(assert_all_requests_are_fired=False) as oauth: - oauth.add(responses.POST, "https://oauth2.googleapis.com/token", json={"access_token": "vertex-token", "expires_in": 3600}) + oauth.add( + responses.POST, + "https://oauth2.googleapis.com/token", + json={"access_token": "vertex-token", "expires_in": 3600}, + ) await litellm.anthropic.messages.acreate( model="vertex_ai/claude-sonnet-5", messages=[{"role": "user", "content": "hi"}], @@ -289,7 +295,12 @@ def test_no_safeguards_leaves_dangerous_tool_use_beta_header_out(): assert "dangerous-tool-use-2026-09-03" not in updated_headers.get("anthropic-beta", "") -def _validate_vertex_headers(client_headers, messages, thinking: Mapping[str, object] | None = None): +def _validate_vertex_headers( + client_headers: dict[str, str], + messages: list[object], + optional_params: dict[str, object] | None = None, + thinking: Mapping[str, object] | None = None, +) -> dict[str, str]: config = VertexAIPartnerModelsAnthropicMessagesConfig() litellm_params = { "vertex_ai_project": "test-project", @@ -305,7 +316,10 @@ def _validate_vertex_headers(client_headers, messages, thinking: Mapping[str, ob headers=client_headers, model="claude-opus-5-5", messages=messages, - optional_params={"max_tokens": 64, **({"thinking": thinking} if thinking is not None else {})}, + optional_params={ + **(optional_params or {"max_tokens": 64}), + **({"thinking": thinking} if thinking is not None else {}), + }, litellm_params=litellm_params, api_base=None, ) @@ -1040,3 +1054,178 @@ def test_vertex_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_f and info.get("supports_mid_conversation_system") is not True ] assert missing == [] + + +_SIGNED_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "summary so far", "signature": "sig"} + + +@pytest.mark.parametrize( + "client_headers", + [{"anthropic-beta": "compact-2026-09-04"}, {}], + ids=["client_sends_beta", "client_omits_beta"], +) +@pytest.mark.parametrize( + "optional_params,messages", + [ + ({"max_tokens": 64, "compaction": {"type": "summarize"}}, [{"role": "user", "content": "Hello"}]), + ( + {"max_tokens": 64}, + [{"role": "assistant", "content": [_SIGNED_COMPACTION_BLOCK]}, {"role": "user", "content": "Go on"}], + ), + ], + ids=["compaction_param", "signed_compaction_block"], +) +def test_native_compaction_reaches_vertex_with_compact_2026_09_04_beta( + local_beta_headers_config: None, + client_headers: dict[str, str], + optional_params: dict[str, object], + messages: list[object], +) -> None: + """Vertex serves on-demand compaction only behind compact-2026-09-04 (Google Cloud: beta on + https://platform.claude.com/docs/en/build-with-claude/compaction-on-demand, read 2026-10-05) + and answers 400 "compaction: Extra inputs are not permitted" without it, so the beta rides + along with the `compaction` param or a signed compaction block exactly once, whether or not + the client sent it, and survives the Vertex beta filter.""" + headers: Final = _validate_vertex_headers(client_headers, messages, optional_params) + + filtered: Final = update_headers_with_filtered_beta(headers=headers, provider="vertex_ai") + + assert filtered["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + + +def test_no_compaction_signal_leaves_compact_2026_09_04_beta_out() -> None: + headers: Final = _validate_vertex_headers({}, [{"role": "user", "content": "Hello"}]) + + assert "compact-2026-09-04" not in headers.get("anthropic-beta", "") + + +_VERTEX_MESSAGE_RESPONSE: Final = { + "id": "msg_vrtx_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "pong"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, +} + + +@pytest.mark.asyncio +async def test_vertex_messages_request_sends_compact_2026_09_04_beta_on_the_wire( + local_beta_headers_config: None, +) -> None: + """A /v1/messages compaction request to a Vertex Claude deployment reaches Vertex with the beta in the + anthropic-beta request header and never in the body, the shape Vertex accepts (live, 2026-10-05).""" + mock_response: Final = MagicMock() + mock_response.status_code = 200 + mock_response.text = json.dumps(_VERTEX_MESSAGE_RESPONSE) + mock_response.headers = httpx.Headers({}) + mock_response.json.return_value = _VERTEX_MESSAGE_RESPONSE + + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token_async", + AsyncMock(return_value=("token", "test-project")), + ): + await litellm.anthropic.messages.acreate( + model="vertex_ai/claude-sonnet-4-6", + max_tokens=64, + messages=[{"role": "user", "content": "Reply with the single word: pong"}], + compaction={"type": "summarize"}, + vertex_project="test-project", + vertex_location="us-east5", + client=client, + ) + + sent: Final = client.post.call_args.kwargs + sent_body: Final = sent["json"] if "json" in sent else json.loads(sent["data"]) + assert sent["headers"]["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + assert sent_body["compaction"] == {"type": "summarize"} + assert "anthropic_beta" not in sent_body + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize("feature", ("thinking", "output_config", "tool_addition", "tool_removal", "none")) +@pytest.mark.parametrize("explicit_beta", (False, True)) +def test_messages_feature_beta_headers(feature: str, explicit_beta: bool) -> None: + from typing import Final + from unittest.mock import patch + + from google.oauth2.credentials import Credentials + + from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.messages.transformation import ( + VertexAIPartnerModelsAnthropicMessagesConfig, + ) + from litellm.types.llms.anthropic import ( + ANTHROPIC_BETA_HEADER_VALUES, + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, + ) + + beta: Final = ( + ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + if feature == "thinking" + else ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value + if feature == "output_config" + else ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER + ) + message: Final = ( + {"role": "system", "content": [], "output_config": {"effort": "high"}} + if feature == "output_config" + else {"role": "system", "content": [{"type": feature, "tool": {"type": "tool_reference", "name": "ping"}}]} + if feature in ("tool_addition", "tool_removal") + else {"role": "user", "content": "Reply with OK"} + ) + with patch( + "google.auth.default", + return_value=(MagicMock(spec=Credentials, token="test-token", expired=False), "test-project"), + ): + headers, _ = VertexAIPartnerModelsAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={"anthropic-beta": f"existing-beta,{beta}" if explicit_beta else "existing-beta"}, + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Hello"}, message], + optional_params={"thinking": {"type": "adaptive", "display": "updates"}} if feature == "thinking" else {}, + litellm_params={"vertex_ai_project": "test-project", "vertex_ai_location": "us-east5"}, + api_key="test-key", + api_base="https://example.com/anthropic", + ) + + betas: Final = headers["anthropic-beta"].split(",") + assert betas.count(beta) == int(feature != "none" or explicit_beta) + assert betas.count("existing-beta") == 1 + + if feature in ("tool_addition", "tool_removal"): + filtered_headers: Final = update_headers_with_filtered_beta(headers=headers, provider="vertex_ai") + assert filtered_headers.get("anthropic-beta", "").split(",").count(beta) == 1 + + +@pytest.mark.parametrize("thinking", ("enabled", {"type": "bogus"}, {"type": "adaptive", "display": "future-mode"})) +def test_beta_detection_preserves_unrecognized_thinking(thinking: object) -> None: + from typing import Final + from unittest.mock import patch + + from google.oauth2.credentials import Credentials + + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + params: Final = {"thinking": thinking} + with patch( + "google.auth.default", + return_value=(MagicMock(spec=Credentials, token="test-token", expired=False), "test-project"), + ): + headers, _ = VertexAIPartnerModelsAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="claude-fable-5-1", + messages=[], + optional_params=params, + litellm_params={"vertex_ai_project": "test-project", "vertex_ai_location": "us-east5"}, + api_key="test-key", + api_base="https://example.com/anthropic", + ) + + assert params["thinking"] == thinking + assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER not in headers.get("anthropic-beta", "").split(",") diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 3ae208bb3a8..0a582c95749 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -4,7 +4,7 @@ import copy import json from collections.abc import Mapping from typing import Final -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import httpx import pytest @@ -12,10 +12,12 @@ import respx from google.oauth2.credentials import Credentials import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.anthropic_beta_headers_manager import ( update_headers_with_filtered_beta, update_request_with_filtered_beta, ) +from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig, ) @@ -475,8 +477,10 @@ def test_vertex_ai_anthropic_extra_headers_beta_propagation(): def test_vertex_ai_anthropic_inline_tools_beta_survives_the_chat_beta_filter(local_beta_headers_config: None) -> None: """The chat path runs the Vertex beta filter over both the header and the `anthropic_beta` body field right - before the request goes out, so inline-tools-2026-09-15 must survive both (source and date in - tests/unit/test_anthropic_beta_headers_filtering.py).""" + before the request goes out, so inline-tools-2026-09-15 must survive it as the header (source and date in + tests/unit/test_anthropic_beta_headers_filtering.py) and stay out of the body field, which Vertex rejects + with 400 "Unexpected value(s) `inline-tools-2026-09-15` for the `anthropic_beta` parameter" (live, global, + claude-sonnet-5-5, 2026-10-20).""" config: Final = VertexAIAnthropicConfig() headers: Final[dict[str, str]] = {} optional_params: Final[dict[str, object]] = { @@ -497,7 +501,7 @@ def test_vertex_ai_anthropic_inline_tools_beta_survives_the_chat_beta_filter(loc ) assert filtered_headers["anthropic-beta"].split(",").count("inline-tools-2026-09-15") == 1 - assert "inline-tools-2026-09-15" in filtered_request["anthropic_beta"] + assert "inline-tools-2026-09-15" not in filtered_request.get("anthropic_beta", []) def test_vertex_ai_anthropic_extra_headers_beta_merged_with_auto_betas(): @@ -1232,3 +1236,293 @@ def test_refusal_stop_details_kept_through_vertex_anthropic_transform(): provider_fields: Final = result.choices[0].message.provider_specific_fields assert provider_fields["stop_reason"] == "refusal" assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + + +_SIGNED_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "summary so far", "signature": "sig"} +_CLEAR_TOOL_USES: Final = {"edits": [{"type": "clear_tool_uses_20250919"}]} + + +def _vertex_request(optional_params: dict[str, object], headers: dict[str, str], litellm_params: dict[str, object] | None = None) -> dict: + return VertexAIAnthropicConfig().transform_request( + model="claude-opus-4-6", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"max_tokens": 100, "is_vertex_request": True, **optional_params}, + litellm_params=litellm_params or {}, + headers=headers, + ) + + +@pytest.mark.parametrize( + "optional_params,messages", + [ + ( + {"max_tokens": 100, "is_vertex_request": True, "compaction": {"type": "summarize"}}, + [{"role": "user", "content": "Hello"}], + ), + ( + {"max_tokens": 100, "is_vertex_request": True}, + [ + { + "role": "assistant", + "content": "", + "provider_specific_fields": {"compaction_blocks": [_SIGNED_COMPACTION_BLOCK]}, + }, + {"role": "user", "content": "Go on"}, + ], + ), + ], + ids=["compaction_param", "signed_compaction_block_sent_back"], +) +def test_vertex_ai_anthropic_native_compaction_adds_compact_2026_09_04_beta( + local_beta_headers_config: None, optional_params: dict[str, object], messages: list[dict[str, object]] +) -> None: + """The chat path skips the Anthropic beta injection for Vertex requests, so the Vertex config + adds compact-2026-09-04 itself (Google Cloud: beta on + https://platform.claude.com/docs/en/build-with-claude/compaction-on-demand, read 2026-10-05) + as the HTTP header only, since Vertex answered 400 "Unexpected value(s) `compact-2026-09-04` + for the `anthropic_beta` parameter" with the name in the body field (us-east5, + claude-sonnet-4-6, 2026-10-05), and the header must survive the Vertex beta filter.""" + headers: Final[dict[str, str]] = {} + + result: Final = VertexAIAnthropicConfig().transform_request( + model="claude-opus-4-6", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers=headers, + ) + filtered_headers, filtered_body = update_request_with_filtered_beta( + headers=headers, request_data=result, provider="vertex_ai" + ) + + assert filtered_headers["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + assert "compact-2026-09-04" not in filtered_body.get("anthropic_beta", []) + + +def test_vertex_ai_anthropic_no_compaction_signal_leaves_compact_2026_09_04_beta_out() -> None: + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request({}, headers) + + assert "compact-2026-09-04" not in result.get("anthropic_beta", []) + assert "compact-2026-09-04" not in headers.get("anthropic-beta", "") + + +def test_vertex_ai_anthropic_client_sent_compact_2026_09_04_stays_out_of_the_body_field( + local_beta_headers_config: None, +) -> None: + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request( + { + "compaction": {"type": "summarize"}, + "extra_headers": {"anthropic-beta": "compact-2026-09-04,interleaved-thinking-2025-05-14"}, + }, + headers, + ) + + assert headers["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + assert result["anthropic_beta"] == ["interleaved-thinking-2025-05-14"] + + +def test_vertex_ai_anthropic_compaction_keeps_the_betas_the_client_sent_as_headers( + local_beta_headers_config: None, +) -> None: + headers: Final = {"anthropic-beta": "context-1m-2025-08-07"} + + result: Final = _vertex_request({"compaction": {"type": "summarize"}}, headers) + + assert sorted(headers["anthropic-beta"].split(",")) == ["compact-2026-09-04", "context-1m-2025-08-07"] + assert "anthropic_beta" not in result + + +@pytest.mark.parametrize("header_only_beta", ["compact-2026-09-04", "inline-tools-2026-09-15"]) +def test_vertex_ai_anthropic_client_sent_body_header_only_beta_moves_to_the_header( + local_beta_headers_config: None, header_only_beta: str +) -> None: + """Vertex rejects these betas inside the body `anthropic_beta` field with 400 "Unexpected value(s)" and + accepts them as the header (live, global, claude-sonnet-5-5, 2026-10-20), so a client that opts in through that field still + gets the beta as a header and keeps its other betas in the body.""" + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request( + {"anthropic_beta": [header_only_beta, "interleaved-thinking-2025-05-14"]}, headers + ) + + assert headers["anthropic-beta"] == header_only_beta + assert result["anthropic_beta"] == ["interleaved-thinking-2025-05-14"] + + +@pytest.mark.parametrize( + "client_beta,body,header", + [ + ( + "context-1m-2025-08-07", + ["context-1m-2025-08-07", "context-management-2025-06-27"], + "context-management-2025-06-27", + ), + ( + "inline-tools-2026-09-15", + ["context-management-2025-06-27"], + "context-management-2025-06-27,inline-tools-2026-09-15", + ), + ], + ids=["body_beta", "header_only_beta"], +) +def test_vertex_ai_anthropic_litellm_betas_merge_with_the_client_body_betas( + local_beta_headers_config: None, client_beta: str, body: list[str], header: str +) -> None: + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request({"context_management": _CLEAR_TOOL_USES, "anthropic_beta": [client_beta]}, headers) + + assert result["anthropic_beta"] == body + assert headers["anthropic-beta"] == header + + +@pytest.mark.parametrize( + "anthropic_beta", + [["interleaved-thinking-2025-05-14"], ("interleaved-thinking-2025-05-14",)], + ids=["list", "sdk_tuple"], +) +def test_vertex_ai_anthropic_client_sent_body_betas_stay_out_of_the_header(anthropic_beta: object) -> None: + headers: Final[dict[str, str]] = {} + + result: Final = _vertex_request({"anthropic_beta": anthropic_beta}, headers) + + assert result["anthropic_beta"] == ["interleaved-thinking-2025-05-14"] + assert "anthropic-beta" not in headers + + +_MALFORMED_BETAS: Final = pytest.mark.parametrize( + "malformed", [5, [1, 2], {"beta": "x"}], ids=["int", "list_of_ints", "dict"] +) +_MALFORMED_BETA_SOURCES: Final = pytest.mark.parametrize( + "source,field", + [("anthropic_beta", "anthropic_beta"), ("extra_headers", "The anthropic-beta header")], + ids=["body_field", "extra_header"], +) + + +def _malformed_beta_params(source: str, malformed: object) -> dict[str, object]: + if source == "extra_headers": + return {"extra_headers": {"anthropic-beta": malformed}} + return {"anthropic_beta": malformed} + + +@_MALFORMED_BETAS +@_MALFORMED_BETA_SOURCES +def test_vertex_ai_anthropic_malformed_anthropic_beta_answers_400_naming_its_field( + malformed: object, source: str, field: str +) -> None: + with pytest.raises(AnthropicError) as excinfo: + _vertex_request(_malformed_beta_params(source, malformed), {}) + + assert excinfo.value.status_code == 400 + assert excinfo.value.message.startswith(f"{field} must be") + + +@_MALFORMED_BETAS +@_MALFORMED_BETA_SOURCES +@pytest.mark.parametrize("auto_beta", [False, True], ids=["alone", "next_to_litellm_beta"]) +def test_vertex_ai_anthropic_malformed_anthropic_beta_is_dropped_under_drop_params( + local_beta_headers_config: None, malformed: object, source: str, field: str, auto_beta: bool +) -> None: + headers: Final[dict[str, str]] = {} + params: Final = {**_malformed_beta_params(source, malformed), **({"context_management": _CLEAR_TOOL_USES} if auto_beta else {})} + + result: Final = _vertex_request(params, headers, litellm_params={"drop_params": True}) + + assert result.get("anthropic_beta") == (["context-management-2025-06-27"] if auto_beta else None) + assert headers.get("anthropic-beta") == ("context-management-2025-06-27" if auto_beta else None) + + +_VERTEX_MESSAGE_RESPONSE: Final = { + "id": "msg_vrtx_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "pong"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, +} + + +def test_vertex_ai_anthropic_completion_sends_compact_2026_09_04_beta_on_the_wire(local_beta_headers_config: None) -> None: + """A chat completion with `compaction` on a Vertex Claude deployment reaches Vertex with the beta in the + anthropic-beta request header and never in the body field, the shape Vertex accepts (live, 2026-10-05).""" + mock_response: Final = MagicMock() + mock_response.status_code = 200 + mock_response.text = json.dumps(_VERTEX_MESSAGE_RESPONSE) + mock_response.headers = {} + mock_response.json.return_value = _VERTEX_MESSAGE_RESPONSE + client: Final = MagicMock(spec=HTTPHandler) + client.post.return_value = mock_response + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token", + return_value=("token", "test-project"), + ): + litellm.completion( + model="vertex_ai/claude-sonnet-4-6", + messages=[{"role": "user", "content": "Reply with the single word: pong"}], + compaction={"type": "summarize"}, + vertex_project="test-project", + vertex_location="us-east5", + client=client, + ) + + sent: Final = client.post.call_args.kwargs + sent_body: Final = sent["json"] if "json" in sent else json.loads(sent["data"]) + assert sent["headers"]["anthropic-beta"].split(",").count("compact-2026-09-04") == 1 + assert sent_body["compaction"] == {"type": "summarize"} + assert "anthropic_beta" not in sent_body + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("thinking_type", ("adaptive", "enabled")) +@pytest.mark.parametrize("display", ("summarized", "omitted", "updates")) +def test_chat_preserves_thinking_display_and_beta(thinking_type: str, display: str) -> None: + from typing import Final + + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + config: Final = VertexAIAnthropicConfig() + optional_params: Final = config.map_openai_params( + non_default_params={"thinking": {"type": thinking_type, "display": display, "budget_tokens": 2048}}, + optional_params={}, + model="claude-fable-5-1", + drop_params=False, + ) + optional_params["extra_headers"] = {"anthropic-beta": "existing-beta"} + headers: Final = {} + result: Final = config.transform_request( + model="claude-fable-5-1", + messages=[{"role": "user", "content": "Reply with OK"}], + optional_params=optional_params, + litellm_params={}, + headers=headers, + ) + + assert result["thinking"]["type"] == "adaptive" + assert result["thinking"]["display"] == display + assert (ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in headers["anthropic-beta"].split(",")) == (display == "updates") + assert "existing-beta" in headers["anthropic-beta"].split(",") + + +@pytest.mark.parametrize("thinking", ("enabled", {"type": "bogus"}, {"type": "adaptive", "display": "future-mode"})) +def test_beta_detection_preserves_unrecognized_thinking(thinking: object) -> None: + from typing import Final + + from litellm.types.llms.anthropic import ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER + + headers: Final = {} + result: Final = VertexAIAnthropicConfig().transform_request( + model="claude-fable-5-1", messages=[{"role": "user", "content": "Hello"}], + optional_params={"thinking": thinking}, litellm_params={}, headers=headers, + ) + + assert result["thinking"] == thinking + assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER not in headers.get("anthropic-beta", "").split(",") diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py new file mode 100644 index 00000000000..c03f08ab610 --- /dev/null +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_handler.py @@ -0,0 +1,60 @@ +from typing import Final + +import pytest + +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_headers + +_HISTORY: Final = ({"role": "user", "content": "Hello"},) + + +@pytest.mark.parametrize( + "request_data", + [ + {"messages": [{"role": "assistant", "content": [{"type": "compaction", "content": "so far", "signature": "sig"}]}]}, + {"messages": list(_HISTORY), "compaction": {"type": "summarize"}}, + ], + ids=["signed_compaction_block", "compaction_param"], +) +def test_count_tokens_on_a_compaction_request_sends_the_compact_beta(request_data: dict[str, object]) -> None: + assert count_tokens_headers("token", request_data) == { + "Authorization": "Bearer token", + "anthropic-beta": "compact-2026-09-04", + } + + +@pytest.mark.parametrize( + "request_data", + [ + {"messages": list(_HISTORY)}, + {"messages": [{"role": "assistant", "content": [{"type": "compaction", "content": "so far", "signature": ""}]}]}, + ], + ids=["plain", "unsigned_compaction_block"], +) +def test_count_tokens_without_compaction_sends_no_beta(request_data: dict[str, object]) -> None: + assert count_tokens_headers("token", request_data) == {"Authorization": "Bearer token"} + + +_TOOL_ADDITION: Final = { + "role": "system", + "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}], +} + + +@pytest.mark.parametrize( + ("request_data", "expected_beta"), + [ + ({"messages": [*_HISTORY, _TOOL_ADDITION]}, "mid-conversation-tool-changes-2026-07-01"), + ( + {"messages": [*_HISTORY, _TOOL_ADDITION], "compaction": {"type": "summarize"}}, + "compact-2026-09-04,mid-conversation-tool-changes-2026-07-01", + ), + ], + ids=["tool_addition", "tool_addition_with_compaction"], +) +def test_count_tokens_on_a_tool_change_request_sends_the_tool_change_beta( + request_data: dict[str, object], expected_beta: str +) -> None: + assert count_tokens_headers("token", request_data) == { + "Authorization": "Bearer token", + "anthropic-beta": expected_beta, + } diff --git a/tests/unit/passthrough/test_main.py b/tests/unit/passthrough/test_main.py index 78fa83d126b..4997d736d0a 100644 --- a/tests/unit/passthrough/test_main.py +++ b/tests/unit/passthrough/test_main.py @@ -51,3 +51,178 @@ async def test_allm_passthrough_route_hosted_vllm_sends_normalized_model(respx_m outbound: Final = json.loads(route.calls[0].request.content) assert outbound["model"] == "my-deployment" assert outbound["messages"] == [{"role": "user", "content": "Hello"}] + + +@pytest.mark.asyncio +async def test_allm_passthrough_route_hosted_vllm_sends_deployment_api_key(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + api_base=_API_BASE, + api_key="sk-upstream", + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.calls[0].request.headers["authorization"] == "Bearer sk-upstream" + assert "x-api-key" not in route.calls[0].request.headers + + +@pytest.mark.parametrize("endpoint", ["v1/chat/completions", "chat/completions"]) +@pytest.mark.asyncio +async def test_allm_passthrough_route_hosted_vllm_reaches_v1_paths_under_an_api_base_ending_in_v1( + respx_mock: respx.MockRouter, endpoint: str +) -> None: + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint=endpoint, + model="hosted_vllm/my-deployment", + api_base=f"{_API_BASE}/v1", + api_key="sk-upstream", + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.call_count == 1 + + +@pytest.mark.asyncio +async def test_allm_passthrough_route_hosted_vllm_sends_no_auth_for_a_keyless_deployment( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("VLLM_API_KEY", "sk-vllm-route-only") + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + api_base=_API_BASE, + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert "authorization" not in route.calls[0].request.headers + assert "x-api-key" not in route.calls[0].request.headers + + +@pytest.mark.parametrize( + ("vllm_api_key", "expected_authorization"), + [("sk-vllm-env", "Bearer sk-vllm-env"), ("oidc/env/VLLM_OIDC_TOKEN", "Bearer sk-vllm-oidc")], +) +@pytest.mark.asyncio +async def test_allm_passthrough_route_keyless_hosted_vllm_without_api_base_sends_the_vllm_api_base_key( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, vllm_api_key: str, expected_authorization: str +) -> None: + monkeypatch.setenv("VLLM_API_BASE", _API_BASE) + monkeypatch.setenv("VLLM_API_KEY", vllm_api_key) + monkeypatch.setenv("VLLM_OIDC_TOKEN", "sk-vllm-oidc") + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.calls[0].request.headers["authorization"] == expected_authorization + + +@pytest.mark.parametrize( + ("model_api_base", "vllm_api_base"), [(f"{_API_BASE}/", _API_BASE), (_API_BASE, f"{_API_BASE}/")] +) +@pytest.mark.asyncio +async def test_allm_passthrough_route_keyless_hosted_vllm_matches_the_vllm_api_base_across_a_trailing_slash( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, model_api_base: str, vllm_api_base: str +) -> None: + monkeypatch.setenv("VLLM_API_BASE", vllm_api_base) + monkeypatch.setenv("VLLM_API_KEY", "sk-vllm-env") + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post(url__regex=rf"^{_API_BASE}/+v1/chat/completions$").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + api_base=model_api_base, + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.calls[0].request.headers["authorization"] == "Bearer sk-vllm-env" + + +@pytest.mark.parametrize("vllm_api_key", ["", "oidc/google/audience"]) +@pytest.mark.asyncio +async def test_allm_passthrough_route_unusable_vllm_api_key_sends_no_authorization_header( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, vllm_api_key: str +) -> None: + monkeypatch.setenv("VLLM_API_BASE", _API_BASE) + monkeypatch.setenv("VLLM_API_KEY", vllm_api_key) + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert "authorization" not in route.calls[0].request.headers + + +@pytest.mark.parametrize( + ("api_base", "api_key", "expected_authorization"), + [ + (None, "sk-upstream", "Bearer sk-upstream"), + (None, "sk-hosted-env", "Bearer sk-hosted-env"), + (_API_BASE, "sk-hosted-env", "Bearer sk-hosted-env"), + (None, None, "Bearer sk-hosted-env"), + ], +) +@pytest.mark.asyncio +async def test_allm_passthrough_route_hosted_vllm_on_the_vllm_api_base_sends_its_own_resolved_key( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + api_base: str | None, + api_key: str | None, + expected_authorization: str, +) -> None: + monkeypatch.setenv("VLLM_API_BASE", _API_BASE) + monkeypatch.setenv("VLLM_API_KEY", "sk-vllm-env") + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + monkeypatch.setenv("HOSTED_VLLM_API_KEY", "sk-hosted-env") + route: Final = respx_mock.post(f"{_API_BASE}/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await allm_passthrough_route( + method="POST", + endpoint="v1/chat/completions", + model="hosted_vllm/my-deployment", + api_base=api_base, + api_key=api_key, + json={"model": "anything", "messages": [{"role": "user", "content": "Hello"}]}, + client=AsyncHTTPHandler(), + ) + assert response.status_code == 200 + assert route.calls[0].request.headers["authorization"] == expected_authorization diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 7abe6ad7799..d6d6774d414 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -5,6 +5,7 @@ Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request """ import asyncio +import importlib.util import json import socket import sys @@ -17,8 +18,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from pydantic import BaseModel, ConfigDict import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.agents import AgentCaller @@ -44,6 +47,7 @@ def _sse_upstream_response(*lines: str, first_line_delay: float = 0.0) -> httpx. class CapturedAgentCall: request_id: object agent_extra_headers: dict[str, str] | None + agent_name: object @pytest.mark.asyncio @@ -627,7 +631,11 @@ async def _invoke_message_method( kwargs: Final = downstream.call_args.kwargs request_id: Final = kwargs["request"].__dict__["id"] if is_send else kwargs["request_id"] - return CapturedAgentCall(request_id=request_id, agent_extra_headers=kwargs.get("agent_extra_headers")) + return CapturedAgentCall( + request_id=request_id, + agent_extra_headers=kwargs.get("agent_extra_headers"), + agent_name=kwargs.get("agent_name"), + ) @pytest.mark.asyncio @@ -1684,6 +1692,76 @@ def test_build_message_send_params_proto_fallback_ignores_unknown_fields(): assert result.message.role.value == "user" +@pytest.mark.asyncio +async def test_message_stream_names_the_agent_by_its_card_name(): + agent: Final = _make_agent_mock() + agent.agent_card_params = {"url": "http://backend-agent:10001", "name": "trip-planner"} + + captured: Final = await _invoke_message_method( + "message/stream", + _make_request_mock("message/stream", _HELLO_MESSAGE_PARAMS), + UserAPIKeyAuth(api_key="sk-test", user_id="u1"), + agent=agent, + ) + + assert captured.agent_name == "trip-planner" + + +class _StreamedSpend(BaseModel): + model_config = ConfigDict(frozen=True) + + call_type: str + model: str + response_cost: float + + +class _StreamedSpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: Final = asyncio.Event() + self.spend: _StreamedSpend | None = None + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + spend: Final = _StreamedSpend.model_validate(kwargs["standard_logging_object"]) + if spend.call_type != "asend_message_streaming": + return + self.spend = spend + self.logged.set() + + +@pytest.mark.skipif(importlib.util.find_spec("a2a") is None, reason="needs the optional a2a-sdk") +@pytest.mark.asyncio +async def test_handle_stream_message_bills_a_bridge_agent_under_its_name(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message + + recorder: Final = _StreamedSpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + + response: Final = await _handle_stream_message( + api_base=None, + request_id="req-1", + params=dict(_HELLO_MESSAGE_PARAMS), + litellm_params={ + "custom_llm_provider": "langgraph", + "model": "agent", + "mock_response": "pong", + "cost_per_query": 0.25, + }, + agent_id="agent-1", + agent_name="trip-planner", + ) + _ = [chunk async for chunk in response.body_iterator] + + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + assert recorder.spend == _StreamedSpend( + call_type="asend_message_streaming", model="a2a_agent/trip-planner", response_cost=0.25 + ) + + @pytest.mark.asyncio async def test_handle_stream_message_rejects_invalid_params_with_32602(): from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e783f01769e..9d4fe67614e 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -46,7 +46,7 @@ from litellm.proxy.auth.auth_checks import ( _team_soft_budget_check, common_checks, ) -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import InternalUsageCache, ProxyLogging from litellm.proxy.utils import CallInfo from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER @@ -2425,3 +2425,51 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ assert refilled == [] assert redis.entries == {} + + +class _SignallingPublishRedisClient(_SlowPublishRedisClient): + def __init__(self) -> None: + super().__init__() + self.sent = asyncio.Event() + + async def _send(self, messages: Sequence[str]) -> None: + await super()._send(messages) + self.sent.set() + + +class _DeleteAfterBroadcastRedisCache(_AuthRedisCache): + def __init__(self, broadcast_sent: asyncio.Event) -> None: + super().__init__({}) + self._broadcast_sent = broadcast_sent + self.finished = False + + async def delete_cache_keys(self, keys: Sequence[str]) -> None: + await self._broadcast_sent.wait() + self.finished = True + + +@pytest.mark.asyncio +async def test_bulk_eviction_broadcasts_while_the_usage_cache_redis_is_wedged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16)) + monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) + client: Final = _SignallingPublishRedisClient() + usage_redis: Final = _DeleteAfterBroadcastRedisCache(broadcast_sent=client.sent) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.internal_usage_cache = InternalUsageCache( + dual_cache=DualCache(redis_cache=usage_redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer + ) + hashed_tokens: Final = ("token-a", "token-b") + + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_PubSubRedisCache(client), + ): + await delete_cache_key_objects( + hashed_tokens=hashed_tokens, user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=proxy_logging_obj + ) + await _drain_publishes() + + assert usage_redis.finished + assert {json.loads(message)["cache_key"] for message in client.published} == set(hashed_tokens) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 0bdc5f1910a..dc7bd8cb312 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import ( can_object_call_model, _can_object_call_vector_stores, _check_agent_access_group_model_access, + agent_model_grants, check_end_user_budget, _check_team_member_budget, personal_key_owner_teams, @@ -5883,7 +5884,7 @@ async def test_team_member_budget_check_reads_from_spend_counter(): class TestGuardrailModificationCheck: - """Defense-in-depth: `_guardrail_modification_check` must 403 when the + """Defense-in-depth: `guardrail_modification_check` must 403 when the caller's metadata attempts to modify any guardrail-related key and the team lacks the `modify_guardrails` permission. Checks both the historically-covered `guardrails` list and the bypass toggles that @@ -5891,11 +5892,11 @@ class TestGuardrailModificationCheck: """ def _call(self, request_body): - from litellm.proxy.auth.auth_checks import _guardrail_modification_check + from litellm.proxy.auth.auth_checks import guardrail_modification_check team_object = MagicMock() team_object.metadata = {} # no permission - return _guardrail_modification_check(request_body=request_body, team_object=team_object) + return guardrail_modification_check(request_body=request_body, team_object=team_object) def test_noop_when_no_guardrail_keys_present(self): # no-op — should return silently @@ -11294,6 +11295,15 @@ async def test_agent_access_groups_naming_no_model_deny_every_model(): assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied +@pytest.mark.asyncio +async def test_agent_model_grants_names_every_model_the_agents_access_groups_allow(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=[]) + resolve, asked = _agent_model_ceiling_resolver(frozenset({"gpt-5", "openai/gpt-5.6"})) + + assert await agent_model_grants(agent_key, resolve) == frozenset({"gpt-5", "openai/gpt-5.6"}) + assert asked == ["agent-1"] + + @pytest.mark.asyncio async def test_agent_without_access_groups_adds_no_model_ceiling(): agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 8612351eda6..928c73d4b04 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -540,7 +540,7 @@ def test_a_base_override_cannot_send_the_deployments_oidc_api_key( oidc_api_key_credential, deployment_params: dict, override: str ): from litellm.exceptions import BadRequestError - from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + from litellm.litellm_core_utils.clientside_credential_handler import get_dynamic_litellm_params with pytest.raises(BadRequestError, match="api_key is an oidc/ secret reference"): get_dynamic_litellm_params( @@ -550,7 +550,7 @@ def test_a_base_override_cannot_send_the_deployments_oidc_api_key( @pytest.mark.parametrize("deployment_params", _OIDC_DEPLOYMENT_PARAMS) def test_a_base_override_with_the_callers_own_api_key_is_allowed(oidc_api_key_credential, deployment_params: dict): - from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + from litellm.litellm_core_utils.clientside_credential_handler import get_dynamic_litellm_params out: Final = get_dynamic_litellm_params( litellm_params=dict(deployment_params), @@ -3009,7 +3009,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: """ def test_clears_admin_organization_and_extra_body_on_base_override(self): - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) @@ -3035,7 +3035,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert "api_version" not in out def test_clears_aws_and_vertex_secrets_on_base_override(self): - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) @@ -3058,7 +3058,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert "vertex_project" not in out def test_clears_nvcf_function_id_on_base_override(self): - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) @@ -3076,7 +3076,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert "nvcf_function_id" not in out def test_clears_use_ssl_on_base_override(self): - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) @@ -3094,7 +3094,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert "use_ssl" not in out def test_clears_oauth_token_exchange_fields_for_other_provider_base_overrides(self): - from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + from litellm.litellm_core_utils.clientside_credential_handler import get_dynamic_litellm_params oauth_values: Final = { "token_exchange_audience": "https://graph.microsoft.com", @@ -3116,7 +3116,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert all(field not in out for field in oauth_values) def test_clears_oauth_token_exchange_fields_for_microsoft_365_copilot_base_overrides(self): - from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + from litellm.litellm_core_utils.clientside_credential_handler import get_dynamic_litellm_params oauth_values: Final = { "token_exchange_audience": "https://graph.microsoft.com", @@ -3146,7 +3146,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: # the field name with any value (or empty string) to keep the # admin's value forwarded, which is the exfiltration vector this # test guards against. - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) @@ -3172,7 +3172,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: # deployment, which would flip the call to shared mode and send the # admin's Copilot token to the attacker's host. Keeping them is safe # because per-user mode ignores the caller api_base anyway. - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) from litellm.types.utils import CredentialItem @@ -3201,7 +3201,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert out["api_base"] == "https://attacker.example" def test_shared_oauth_credential_name_is_cleared_on_base_override(self): - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) from litellm.types.utils import CredentialItem @@ -3231,7 +3231,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: # Regression: a caller that echoes an admin-config field name with # an *empty* value (or any value) must not be able to keep the # admin's value in ``litellm_params``. - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) @@ -3253,7 +3253,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert "org-admin-secret" not in str(out) def test_no_clearing_when_only_api_key_overridden(self): - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) @@ -3272,7 +3272,7 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: assert out["api_base"] == "https://admin.upstream/v1" def test_client_api_key_used_when_supplied_with_base_override(self): - from litellm.router_utils.clientside_credential_handler import ( + from litellm.litellm_core_utils.clientside_credential_handler import ( get_dynamic_litellm_params, ) @@ -3971,7 +3971,7 @@ def test_admin_opt_in_proxy_wide_allows_riva_offline() -> None: def test_api_base_override_keeps_the_deployments_riva_offline_mode() -> None: - from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + from litellm.litellm_core_utils.clientside_credential_handler import get_dynamic_litellm_params out: Final = get_dynamic_litellm_params( litellm_params={"model": "nvidia_riva/nvidia/parakeet-0.6b-tdt", "riva_offline": True}, diff --git a/tests/unit/proxy/auth/test_fallback_budget.py b/tests/unit/proxy/auth/test_fallback_budget.py index 1117ca43657..e123fbc8177 100644 --- a/tests/unit/proxy/auth/test_fallback_budget.py +++ b/tests/unit/proxy/auth/test_fallback_budget.py @@ -164,7 +164,7 @@ async def test_router_fallback_budget_check_resolves_model_group_alias_before_me within_budget: Final = await budget_check( model="backup", request_kwargs={"metadata": {"user_api_key_auth": token}}, - llm_router=router, + ctx=router._core_ctx, ) assert within_budget is False @@ -195,7 +195,7 @@ async def test_router_fallback_budget_check_refuses_an_exhausted_alias_cap() -> within_budget: Final = await budget_check( model="backup", request_kwargs={"metadata": {"user_api_key_auth": token}}, - llm_router=router, + ctx=router._core_ctx, ) assert within_budget is False @@ -229,7 +229,7 @@ async def test_router_fallback_budget_check_allows_alias_when_both_names_are_wit within_budget: Final = await budget_check( model="backup", request_kwargs={"metadata": {"user_api_key_auth": token}}, - llm_router=router, + ctx=router._core_ctx, ) assert within_budget is True @@ -252,7 +252,7 @@ async def test_router_fallback_budget_check_deduplicates_non_alias_target() -> N within_budget: Final = await budget_check( model="paid-model", request_kwargs={"metadata": {"user_api_key_auth": token}}, - llm_router=router, + ctx=router._core_ctx, ) assert within_budget is True @@ -346,12 +346,12 @@ async def test_team_key_inherits_personal_budget_when_opted_in(monkeypatch): @pytest.mark.asyncio async def test_check_is_a_no_op_while_not_enforced(): request = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} - assert await NOT_ENFORCED(model="paid-model", request_kwargs=request, llm_router=_router()) is True + assert await NOT_ENFORCED(model="paid-model", request_kwargs=request, ctx=_router()._core_ctx) is True @pytest.mark.asyncio async def test_request_without_a_key_is_unrestricted(): - assert await ENFORCED(model="paid-model", request_kwargs={}, llm_router=_router()) is True + assert await ENFORCED(model="paid-model", request_kwargs={}, ctx=_router()._core_ctx) is True @pytest.mark.asyncio @@ -360,8 +360,8 @@ async def test_enforced_check_reads_the_key_from_request_metadata(metadata_field over = {metadata_field: {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} under = {metadata_field: {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} - assert await ENFORCED(model="paid-model", request_kwargs=over, llm_router=_router()) is False - assert await ENFORCED(model="paid-model", request_kwargs=under, llm_router=_router()) is True + assert await ENFORCED(model="paid-model", request_kwargs=over, ctx=_router()._core_ctx) is False + assert await ENFORCED(model="paid-model", request_kwargs=under, ctx=_router()._core_ctx) is True @pytest.mark.asyncio @@ -382,7 +382,7 @@ async def test_router_check_passes_member_limiter_from_request_metadata() -> Non within_budget: Final = await budget_check( model="paid-model", request_kwargs=request_kwargs, - llm_router=_router(), + ctx=_router()._core_ctx, ) assert within_budget is False @@ -424,7 +424,7 @@ async def test_check_fails_closed_when_the_spend_lookup_breaks(monkeypatch): monkeypatch.setattr(proxy_server, "get_current_spend", _boom, raising=False) request = {"metadata": {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} - assert await ENFORCED(model="paid-model", request_kwargs=request, llm_router=_router()) is False + assert await ENFORCED(model="paid-model", request_kwargs=request, ctx=_router()._core_ctx) is False @pytest.mark.asyncio @@ -446,12 +446,12 @@ async def test_router_skips_the_paid_fallback_target_when_over_budget(): under_attempted: Final = AttemptedFallbackTargets() assert ( - await _is_fallback_target_within_budget(router, "paid-model", "free-model", over, attempted=over_attempted) + await _is_fallback_target_within_budget(router._core_ctx, "paid-model", "free-model", over, attempted=over_attempted) is False ) assert ( await _is_fallback_target_within_budget( - router, "paid-model", "free-model", under, attempted=under_attempted + router._core_ctx, "paid-model", "free-model", under, attempted=under_attempted ) is True ) @@ -467,7 +467,7 @@ async def test_router_without_a_budget_check_attempts_every_fallback(): attempted: Final = AttemptedFallbackTargets() assert ( - await _is_fallback_target_within_budget(router, "paid-model", "free-model", over, attempted=attempted) + await _is_fallback_target_within_budget(router._core_ctx, "paid-model", "free-model", over, attempted=attempted) is True ) @@ -483,7 +483,7 @@ async def test_enforcement_is_on_by_default_and_opt_out_restores_the_leak(monkey over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) - assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, llm_router=_router()) is False + assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, ctx=_router()._core_ctx) is False monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": False}, raising=False) - assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, llm_router=_router()) is True + assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, ctx=_router()._core_ctx) is True diff --git a/tests/unit/proxy/auth/test_fallback_model_access.py b/tests/unit/proxy/auth/test_fallback_model_access.py index 39cb7d07b43..feeb2d26ff3 100644 --- a/tests/unit/proxy/auth/test_fallback_model_access.py +++ b/tests/unit/proxy/auth/test_fallback_model_access.py @@ -73,18 +73,18 @@ async def test_enforced_check_authorizes_the_key_carried_in_request_metadata(met router = _router() request_kwargs = _request_with_key(metadata_field) - assert await ENFORCED(model="open-model", request_kwargs=request_kwargs, llm_router=router) - assert not await ENFORCED(model="secret-model", request_kwargs=request_kwargs, llm_router=router) + assert await ENFORCED(model="open-model", request_kwargs=request_kwargs, ctx=router._core_ctx) + assert not await ENFORCED(model="secret-model", request_kwargs=request_kwargs, ctx=router._core_ctx) @pytest.mark.asyncio async def test_enforced_check_does_not_restrict_requests_without_a_key(): - assert await ENFORCED(model="secret-model", request_kwargs={"metadata": {}}, llm_router=_router()) + assert await ENFORCED(model="secret-model", request_kwargs={"metadata": {}}, ctx=_router()._core_ctx) @pytest.mark.asyncio async def test_check_allows_every_fallback_while_not_enforced(): - assert await NOT_ENFORCED(model="secret-model", request_kwargs=_request_with_key(), llm_router=_router()) + assert await NOT_ENFORCED(model="secret-model", request_kwargs=_request_with_key(), ctx=_router()._core_ctx) @pytest.mark.asyncio @@ -104,11 +104,11 @@ async def test_proxy_check_reads_enforce_fallback_model_access_from_general_sett monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) assert await router_fallback_access_check( - model="open-model", request_kwargs=_request_with_key(), llm_router=_router() + model="open-model", request_kwargs=_request_with_key(), ctx=_router()._core_ctx ) assert ( await router_fallback_access_check( - model="secret-model", request_kwargs=_request_with_key(), llm_router=_router() + model="secret-model", request_kwargs=_request_with_key(), ctx=_router()._core_ctx ) is expected ) @@ -140,8 +140,8 @@ async def test_search_tool_fallback_target_follows_the_key_search_tool_grant( }, } - assert await check(model="search-a", request_kwargs=request_kwargs, llm_router=router) - assert not await check(model="search-b", request_kwargs=request_kwargs, llm_router=router) + assert await check(model="search-a", request_kwargs=request_kwargs, ctx=router._core_ctx) + assert not await check(model="search-b", request_kwargs=request_kwargs, ctx=router._core_ctx) @pytest.mark.asyncio @@ -174,7 +174,7 @@ async def test_fallback_named_like_both_a_model_and_a_search_tool_follows_the_re allowed = await ENFORCED( model="shared-name", request_kwargs={**request_kwargs, "metadata": {"user_api_key_auth": key}}, - llm_router=router, + ctx=router._core_ctx, ) assert allowed is expected @@ -190,5 +190,5 @@ async def test_default_check_enforces_team_models_for_a_key_without_model_restri token: Final = UserAPIKeyAuth(api_key="hashed", models=[], team_id="qa-team", team_models=["open-model"]) request_kwargs: Final = {metadata_field: {"user_api_key_auth": token}} - assert await router_fallback_access_check(model="open-model", request_kwargs=request_kwargs, llm_router=router) - assert not await router_fallback_access_check(model="secret-model", request_kwargs=request_kwargs, llm_router=router) + assert await router_fallback_access_check(model="open-model", request_kwargs=request_kwargs, ctx=router._core_ctx) + assert not await router_fallback_access_check(model="secret-model", request_kwargs=request_kwargs, ctx=router._core_ctx) diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index c67610c24a1..d0156ff0195 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -7,6 +7,8 @@ import httpx import pytest import respx +from litellm import Router + def test_get_team_models_for_all_models_and_team_only_models(): from litellm.proxy.auth.model_checks import get_team_models @@ -1261,6 +1263,223 @@ def test_partial_bedrock_wildcard_lists_nothing_when_no_discovered_id_carries_it assert get_known_models_from_wildcard("bedrock/anthropic.*", deployment) == [] +@pytest.mark.parametrize( + ("wildcard", "deployment_model", "listed_ids"), + [ + ( + "openrouter/*", + "openrouter/*", + ["openrouter/openai/gpt-5.4-mini", "openrouter/anthropic/claude-haiku-5-5"], + ), + ("or-team/*", "openrouter/*", ["or-team/openai/gpt-5.4-mini", "or-team/anthropic/claude-haiku-5-5"]), + ("openrouter/openai/*", "openrouter/openai/*", ["openrouter/openai/gpt-5.4-mini"]), + ("*", "openrouter/*", ["openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5"]), + ], +) +@respx.mock +def test_openrouter_wildcard_lists_the_live_catalog_with_each_vendor_namespace( + monkeypatch: pytest.MonkeyPatch, wildcard: str, deployment_model: str, listed_ids: list[str] +) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + respx.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response( + 200, json={"data": [{"id": "openai/gpt-5.4-mini"}, {"id": "anthropic/claude-haiku-5-5"}]} + ) + ) + + assert ( + get_known_models_from_wildcard(wildcard, LiteLLM_Params(model=deployment_model, api_key=f"sk-or-{wildcard}")) + == listed_ids + ) + + +@respx.mock +def test_key_granted_openrouter_wildcard_without_an_openrouter_deployment_lists_no_environment_catalog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-environment-only") + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + ) + router: Final = Router( + model_list=[{"model_name": "or-team/*", "litellm_params": {"model": "openrouter/*", "api_key": "sk-or-team"}}] + ) + + listed: Final = get_complete_model_list( + key_models=["openrouter/*"], + team_models=[], + proxy_model_list=router.model_names, + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + + assert (listed, catalog.call_count) == ([], 0) + + +@contextmanager +def _openai_catalog_behind_a_catch_all( + monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str] +) -> Generator[Router, None, None]: + import litellm + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.get("https://api.openai.com/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "gpt-4o"}, {"id": "gpt-4o-mini"}]}) + ) + yield Router( + model_list=[ + {"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-openai-catch-all"}}, + { + "model_name": concrete_name, + "litellm_params": {"model": "openrouter/openai/gpt-4o", "api_key": "sk-or-gpt-4o"}, + "model_info": dict(concrete_model_info), + }, + ] + ) + + +@pytest.mark.parametrize( + ("concrete_name", "concrete_model_info", "team_id"), + [ + ("gpt-4o", {}, None), + ("model_name_team-x_gpt-4o", {"team_id": "team-x", "team_public_model_name": "gpt-4o"}, "team-x"), + ], +) +def test_catch_all_wildcard_lists_a_name_a_concrete_deployment_serves_once( + monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str], team_id: str | None +) -> None: + from litellm.proxy.auth.model_checks import get_complete_model_list + + with _openai_catalog_behind_a_catch_all(monkeypatch, concrete_name, concrete_model_info) as router: + listed: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=router.get_model_names(team_id=team_id), + user_model=None, + infer_model_from_keys=False, + llm_router=router, + team_id=team_id, + ) + + assert sorted(listed) == ["gpt-4o", "gpt-4o-mini"] + + +def test_catch_all_wildcard_model_info_keeps_only_the_concrete_row_for_a_deployed_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info + + with _openai_catalog_behind_a_catch_all(monkeypatch, "gpt-4o", {}) as router: + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + assert sorted((row["model_name"], row["litellm_params"]["model"]) for row in rows) == [ + ("gpt-4o", "openrouter/openai/gpt-4o"), + ("gpt-4o-mini", "openai/gpt-4o-mini"), + ] + + +@contextmanager +def _openrouter_catalog_beside_a_named_deployment( + monkeypatch: pytest.MonkeyPatch, + wildcard_name: str, + named: str, + catalog: tuple[str, ...] = ("openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5"), + wildcard_key: str | None = None, +) -> Generator[Router, None, None]: + import litellm + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": catalog_id} for catalog_id in catalog]}) + ) + yield Router( + model_list=[ + { + "model_name": wildcard_name, + "litellm_params": { + "model": "openrouter/*", + "api_key": wildcard_key or f"sk-or-wildcard-{wildcard_name}", + }, + }, + { + "model_name": named, + "litellm_params": {"model": "openrouter/openai/gpt-5.4-mini", "api_key": "sk-or-named"}, + }, + ] + ) + + +@pytest.mark.parametrize( + ("wildcard_name", "named"), + [("openrouter/*", "openrouter/openai/gpt-5.4-mini"), ("or-team/*", "or-team/openai/gpt-5.4-mini")], +) +def test_provider_wildcard_lists_a_name_a_named_deployment_serves_once( + monkeypatch: pytest.MonkeyPatch, wildcard_name: str, named: str +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + + with _openrouter_catalog_beside_a_named_deployment(monkeypatch, wildcard_name, named) as router: + listed: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=list(router.model_names), + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + vendor_prefix: Final = wildcard_name.removesuffix("*") + assert sorted(listed) == [f"{vendor_prefix}anthropic/claude-haiku-5-5", named] + assert sorted((row["model_name"], row["litellm_params"]["api_key"]) for row in rows) == [ + (f"{vendor_prefix}anthropic/claude-haiku-5-5", f"sk-or-wildcard-{wildcard_name}"), + (named, "sk-or-named"), + ] + + +@pytest.mark.parametrize( + ("wildcard_name", "named"), + [("openrouter/*", "openrouter/openai/gpt-5.4-mini"), ("or-team/*", "or-team/openai/gpt-5.4-mini")], +) +def test_provider_wildcard_whose_whole_catalog_is_deployed_by_name_shows_only_the_named_row( + monkeypatch: pytest.MonkeyPatch, wildcard_name: str, named: str +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info + + with _openrouter_catalog_beside_a_named_deployment( + monkeypatch, + wildcard_name, + named, + catalog=("openai/gpt-5.4-mini",), + wildcard_key=f"sk-or-one-model-{wildcard_name}", + ) as router: + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + assert [(row["model_name"], row["litellm_params"]["api_key"]) for row in rows] == [(named, "sk-or-named")] + + @respx.mock def test_custom_prefix_that_starts_a_discovered_id_still_prefixes_every_listed_model( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/auth/test_model_checks_fallbacks.py b/tests/unit/proxy/auth/test_model_checks_fallbacks.py index dc38cbbc879..071652cfc4d 100644 --- a/tests/unit/proxy/auth/test_model_checks_fallbacks.py +++ b/tests/unit/proxy/auth/test_model_checks_fallbacks.py @@ -50,7 +50,7 @@ def test_model_with_fallbacks_returns_complete_list(): router = create_mock_router(fallbacks=fallbacks_config) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: mock_get_fallback.return_value = ( ["bedrock-claude-sonnet-4", "google-claude-sonnet-4"], @@ -71,7 +71,7 @@ def test_model_without_fallbacks_returns_empty_list(): router = create_mock_router(fallbacks=fallbacks_config) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: mock_get_fallback.return_value = (None, None) @@ -87,7 +87,7 @@ def test_general_fallback_type(): router = create_mock_router(fallbacks=fallbacks_config) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: mock_get_fallback.return_value = (["bedrock-claude-sonnet-4"], None) @@ -110,7 +110,7 @@ def test_context_window_fallback_type(): router = create_mock_router(context_window_fallbacks=context_fallbacks_config) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: mock_get_fallback.return_value = (["gpt-3.5-turbo"], None) @@ -133,7 +133,7 @@ def test_content_policy_fallback_type(): router = create_mock_router(content_policy_fallbacks=content_fallbacks_config) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: mock_get_fallback.return_value = (["claude-3"], None) @@ -170,7 +170,7 @@ def test_exception_handling_returns_empty_list(): router = create_mock_router(fallbacks=[{"claude-4-sonnet": ["fallback"]}]) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: mock_get_fallback.side_effect = Exception("Test exception") @@ -195,7 +195,7 @@ def test_multiple_fallbacks_complete_list(): router = create_mock_router(fallbacks=fallbacks_config) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: mock_get_fallback.return_value = ( ["gpt-4-turbo", "gpt-3.5-turbo", "claude-3-haiku"], @@ -217,7 +217,7 @@ def test_wildcard_and_specific_fallbacks(): router = create_mock_router(fallbacks=fallbacks_config) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: # Test specific model fallbacks mock_get_fallback.return_value = ( @@ -241,7 +241,7 @@ def test_default_fallback_type_is_general(): router = create_mock_router(fallbacks=fallbacks_config) with patch( - "litellm.proxy.auth.model_checks.get_fallback_model_group" + "litellm.proxy.auth.model_checks.fallback_chain" ) as mock_get_fallback: mock_get_fallback.return_value = (["bedrock-claude-sonnet-4"], None) diff --git a/tests/unit/proxy/common_utils/test_cache_aware_routing.py b/tests/unit/proxy/common_utils/test_cache_aware_routing.py index 4eb22b2e34f..8063493927c 100644 --- a/tests/unit/proxy/common_utils/test_cache_aware_routing.py +++ b/tests/unit/proxy/common_utils/test_cache_aware_routing.py @@ -145,7 +145,7 @@ async def _select( counter_for_model: Callable[[str], TokenCounter], now: float, ) -> CacheAwareChoice | None: - complexity: Final = ComplexityRouter("smart", router, config.model_dump()) + complexity: Final = ComplexityRouter("smart", router._core_ctx, config.model_dump()) return await select_cached_model( router=router, config=config, @@ -233,10 +233,11 @@ async def test_disabled_setting_does_not_access_prediction_services() -> None: assert config.cache_aware_routing is False assert ( await choose_cached_model( - router=_router(), + ctx=_router()._core_ctx, config=config, params_for_model=partial( - ComplexityRouter("smart", _router(), config.model_dump())._litellm_params_for_model, request_kwargs={} + ComplexityRouter("smart", _router()._core_ctx, config.model_dump())._litellm_params_for_model, + request_kwargs={}, ), response=_response(), request_kwargs={}, diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 4861b2687de..5b37cee88c5 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -1771,6 +1771,53 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-ft", "8823717160934178816") + def _two_team_deployments(self): + """Two Vertex deployments sharing the base model, team A listed first.""" + return [ + { + "model_name": f"gemini-2.5-flash-team-{team}", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "custom_llm_provider": "vertex_ai"}, + "model_info": {"id": f"deploy-team-{team}"}, + } + for team in ("a", "b") + ] + + def test_flag_on_prefers_the_deployment_stored_on_the_row_over_the_first_base_model_match(self): + """Regression (LIT-10326): two deployments share gemini-2.5-flash with different + buckets. A batch created on team B's deployment stores deploy-team-b in file_object, + so the poller must read it back through team B, not the first match (team A).""" + router: Final = MagicMock() + router.registry.return_value.resolve.return_value = None + router.registry.return_value.deployments.return_value = self._two_team_deployments() + router.registry.return_value.deployment = MagicMock(return_value=self._vertex_deployment()) + instance: Final = self._instance(track_unmanaged=True, router=router) + job: Final = self._job( + file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-team-b"}) + ) + + with patch(_IS_B64, return_value=False): + result: Final = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-team-b", "8823717160934178816") + router.registry.return_value.deployment.assert_called_once_with("deploy-team-b") + + def test_flag_on_falls_back_to_base_model_match_when_the_stored_deployment_is_gone(self): + """The stored deployment was removed from the config after the batch was created: + the row still routes through today's base-model lookup instead of being skipped.""" + router: Final = MagicMock() + router.registry.return_value.resolve.return_value = None + router.registry.return_value.deployments.return_value = self._two_team_deployments() + router.registry.return_value.deployment = MagicMock(return_value=None) + instance: Final = self._instance(track_unmanaged=True, router=router) + job: Final = self._job( + file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-removed"}) + ) + + with patch(_IS_B64, return_value=False): + result: Final = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-team-a", "8823717160934178816") + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): """Flag on, but the only deployment for the model group is a non-vertex_ai provider: must not be selected, even though the model group name matches.""" diff --git a/tests/unit/proxy/common_utils/test_openai_error_payload.py b/tests/unit/proxy/common_utils/test_openai_error_payload.py index acefc4c1e2c..d4ab811d505 100644 --- a/tests/unit/proxy/common_utils/test_openai_error_payload.py +++ b/tests/unit/proxy/common_utils/test_openai_error_payload.py @@ -111,10 +111,13 @@ def test_the_default_status_stands_when_the_exception_carries_none(default: int) assert error_status_code(ValueError("boom"), default) == default -@pytest.mark.parametrize("carried_status", [True, False, "429", None, 429.0]) -def test_a_non_int_carried_status_falls_back_to_the_default(carried_status: object): - """True is an int in Python but not an HTTP status, and a stringified one would break - every caller that compares the code numerically.""" +@pytest.mark.parametrize( + "carried_status", + [True, False, "slow down", "", None, 429.0, 200, 1301, 0, "200", "1301", "0", "0429", " 429", "\u0664\u0662\u0669", "4" * 5000], +) +def test_a_carried_status_that_is_not_an_error_status_falls_back_to_the_default(carried_status: object): + """True is an int in Python but not an HTTP status, and answering a 200 or a 1301 either reads + as success or makes uvicorn drop the connection.""" class _Carrier(Exception): status_code = carried_status @@ -122,6 +125,22 @@ def test_a_non_int_carried_status_falls_back_to_the_default(carried_status: obje assert error_status_code(_Carrier("boom"), 500) == 500 +def test_a_status_carried_as_a_decimal_string_is_answered_as_that_status(): + """Provider stream-error chunks pass the vendor's code through as a string ``status_code``, + and the shared handler used to answer it as is, so dropping it turned a rate limit into a 500.""" + + class _Carrier(Exception): + status_code = "429" + + assert error_status_code(_Carrier("slow down"), 500) == 429 + assert proxy_exception_for(_Carrier("slow down"), default_status_code=500).to_dict() == { + "message": "slow down", + "type": "rate_limit_error", + "param": None, + "code": "429", + } + + def test_a_proxy_exception_keeps_the_status_it_was_raised_with(): """ProxyException stores its status as the string ``code`` rather than ``status_code``, so a route tail that rewraps one used to answer a 4xx rejection as a 500.""" @@ -130,8 +149,8 @@ def test_a_proxy_exception_keeps_the_status_it_was_raised_with(): assert error_status_code(rejection, 500) == 400 -@pytest.mark.parametrize("carried_code", [None, "None", "", "rate_limited", "4xx", 404]) -def test_a_code_that_is_not_a_decimal_string_falls_back_to_the_default(carried_code: object): +@pytest.mark.parametrize("carried_code", [None, "None", "", "rate_limited", "4xx", 404, "200", "1301"]) +def test_a_code_that_is_not_a_stringified_error_status_falls_back_to_the_default(carried_code: object): """Only ProxyException's stringified status is a status; ``code`` on anything else (OpenAI's ``invalid_api_key``, a stray int) says nothing about the HTTP answer.""" diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 47dd6d561b3..3d8e92988d9 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -3277,6 +3277,23 @@ class TestCredentialDisplayName: assert listed["openai-prod"]["credential_values"]["api_key"] != "sk-old" +def test_looking_up_a_missing_credential_is_labelled_a_client_error_not_a_gateway_failure( + restore_credential_list: None, +): + """Naming a credential that does not exist is the caller's own mistake, and the shared + handler used to label every one of these internal_server_error, so an SDK branching on + the type read a 404 it can never retry past as a gateway outage and backed off forever.""" + response: Final = _call_as("GET", "/credentials/by_name/definitely-not-there") + + assert response.status_code == 404, f"lookup answered {response.status_code}: {response.text}" + assert response.json()["error"] == { + "message": "Credential not found. Got credential name: definitely-not-there", + "type": "invalid_request_error", + "param": None, + "code": "404", + } + + def test_get_credential_by_model_masks_the_deployment_credentials(monkeypatch): import asyncio diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 24c12e06c9c..65da5630f33 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -1,10 +1,12 @@ import json +import logging from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm._logging import verbose_proxy_logger from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyTagSpendTransaction @@ -379,6 +381,104 @@ async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_fail assert restored_spend["org_member_list_transactions"] == {member_key: 1.5} +@pytest.mark.asyncio +async def test_a_single_parked_batch_with_a_nan_spend_drains_as_zero( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock +): + mock_redis_cache.async_lpop = AsyncMock( + return_value=json.dumps({"key_list_transactions": {"nan-key": float("nan"), "ok-key": 0.5}}, allow_nan=True) + ) + + drained: Final = await redis_update_buffer.get_all_update_transactions_from_redis_buffer() + + assert drained is not None + assert drained["key_list_transactions"] == {"nan-key": 0.0, "ok-key": 0.5} + + +PARKED_DAILY_KEY: Final = "user-1_2025-01-15_api-key-123_gpt-4o_openai_/chat/completions" + + +def _parked_daily_batch(spend: float) -> str: + return json.dumps( + { + PARKED_DAILY_KEY: { + "spend": spend, + "prompt_tokens": 12, + "completion_tokens": 7, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + }, + allow_nan=True, + ) + + +@pytest.mark.asyncio +async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_as_zero( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock +): + entity_batch: Final = json.dumps({"key_list_transactions": {"hashed-key": float("nan")}}, allow_nan=True) + window_batches: Final = tuple( + json.dumps( + [ + { + "entity_type": "key", + "entity_id": "hashed-key", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": spend, + "started_at": None, + } + ], + allow_nan=True, + ) + for spend in (float("nan"), 3.0) + ) + mock_redis_cache.async_lpop_pipeline = AsyncMock( + return_value=[ + [entity_batch], + [_parked_daily_batch(float("nan")), _parked_daily_batch(0.0042)], + None, + None, + None, + None, + list(window_batches), + ] + ) + + ( + db_spend, + daily_user, + *_rest, + window_spend, + ) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + + assert db_spend is not None + assert db_spend["key_list_transactions"] == {"hashed-key": 0} + assert daily_user is not None + assert daily_user[PARKED_DAILY_KEY]["spend"] == 0.0042 + assert daily_user[PARKED_DAILY_KEY]["api_requests"] == 2 + assert window_spend is not None + assert [transaction["spend"] for transaction in window_spend] == [3.0] + + +@pytest.mark.asyncio +async def test_a_nan_spend_parked_by_an_older_version_drains_from_a_daily_key_as_zero( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, caplog: pytest.LogCaptureFixture +): + mock_redis_cache.async_lpop = AsyncMock( + return_value=[_parked_daily_batch(float("inf")), _parked_daily_batch(0.0042)] + ) + + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + daily_tag: Final = await redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + + assert daily_tag is not None + assert daily_tag[PARKED_DAILY_KEY]["spend"] == 0.0042 + assert [record.getMessage() for record in caplog.records if "Infinity" in record.getMessage()] != [] + + @pytest.mark.asyncio async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): """When redis_cache is None, should return all Nones""" @@ -724,6 +824,10 @@ class _ListRedis: def __init__(self) -> None: self.rows: list[str] = [] + async def async_rpush(self, key: str, values: list[str]) -> int: + self.rows.extend(values) + return len(self.rows) + async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int: self.rows.extend(values) pushed_len = len(self.rows) @@ -752,6 +856,32 @@ async def test_store_spend_logs_in_redis_drops_oldest_rows_past_the_cap(): assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == () +@pytest.mark.asyncio +@pytest.mark.parametrize("spend", [float("nan"), float("inf")]) +async def test_a_non_finite_daily_spend_buffered_in_redis_drains_as_zero(spend: float): + buffer: Final = RedisUpdateBuffer(redis_cache=_ListRedis()) + queue: Final = DailySpendUpdateQueue() + await queue.add_update(json.loads(_parked_daily_batch(spend))) + await queue.add_update(json.loads(_parked_daily_batch(0.0042))) + + await buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) + drained: Final = await buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + + assert drained is not None + assert drained[PARKED_DAILY_KEY]["spend"] == 0.0 + assert drained[PARKED_DAILY_KEY]["api_requests"] == 2 + + +@pytest.mark.asyncio +async def test_a_spend_log_row_parked_with_a_nan_spend_is_read_back_with_zero_spend(): + buffer: Final = RedisUpdateBuffer(redis_cache=_ListRedis()) + buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True) + + assert await buffer.store_spend_logs_in_redis([{"request_id": "nan-row", "spend": float("nan")}]) is True + + assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == ({"request_id": "nan-row", "spend": 0.0},) + + @pytest.mark.asyncio async def test_store_spend_logs_in_redis_reports_failure_without_redis(): buffer = RedisUpdateBuffer(redis_cache=None) diff --git a/tests/unit/proxy/db/test_accounting_ledger.py b/tests/unit/proxy/db/test_accounting_ledger.py index 3916fe869e6..6ca719ed849 100644 --- a/tests/unit/proxy/db/test_accounting_ledger.py +++ b/tests/unit/proxy/db/test_accounting_ledger.py @@ -6,7 +6,7 @@ from typing import Final, Literal, NoReturn import pytest from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome -from litellm.proxy.db.accounting_ledger import AccountingLedger +from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingLedger, UsageWrite from litellm.types.utils import CallTypes @@ -101,3 +101,18 @@ async def test_unknown_delivery_target_is_rejected_before_database_access(event: ("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target envelope=envelope, ) + + +@pytest.mark.asyncio +async def test_empty_usage_batches_do_not_open_a_transaction() -> None: + ledger: Final = AccountingLedger(_no_transaction) + await ledger.insert_usage_batch(()) + assert await ledger.load_usage_batch(()) == () + + +@pytest.mark.asyncio +async def test_conflicting_usage_within_one_batch_is_rejected_before_database_access(event: InboundUsageEvent) -> None: + original: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")})) + changed: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json"), "pricing": {"rate": 2}})) + with pytest.raises(AccountingConflictError, match="conflicting copies"): + await AccountingLedger(_no_transaction).insert_usage_batch((original, changed)) diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index c3a93fd89be..298b153f621 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -2226,6 +2226,60 @@ async def test_shutdown_drain_finishes_the_rollups_update_database_started(): assert (team_row["api_requests"], team_row["spend"]) == (1, 0.1) +@pytest.mark.asyncio +@pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), float("-inf")]) +async def test_update_database_records_a_non_finite_response_cost_as_zero(response_cost: float) -> None: + db_writer: Final = DBSpendUpdateWriter() + prisma: Final = _tool_usage_prisma() + prisma.get_request_status = MagicMock(return_value="success") + response: Final = litellm.ModelResponse( + model="gpt-5.6-luna", + usage=litellm.Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={ + "model": "gpt-5.6-luna", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "req-non-finite-cost", + "response_cost": response_cost, + "litellm_params": {"metadata": {"user_api_key": "hashed-key", "user_api_key_team_id": "test-team"}}, + }, + completion_response=response, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=response_cost, + ) + await drain_logging_tasks() + + assert [row["spend"] for row in prisma.spend_log_transactions] == [0] + spend: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert spend["key_list_transactions"] == {"test-token": 0} + assert spend["team_list_transactions"] == {"test-team": 0} + daily_user: Final = ( + await db_writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_team: Final = ( + await db_writer.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + assert daily_user is not None and daily_team is not None + (user_row,) = daily_user.values() + (team_row,) = daily_team.values() + assert (user_row["api_requests"], user_row["spend"]) == (1, 0) + assert (team_row["api_requests"], team_row["spend"]) == (1, 0) + + @pytest.mark.asyncio async def test_shutdown_drain_waits_for_an_in_flight_rollup_enqueue(): """ diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index 75346ca08a6..d70d4403962 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -951,3 +951,33 @@ def test_is_lock_timeout_error_recognises_a_statement_cancelled_under_lock_timeo ) def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None: assert PrismaDBExceptionHandler.is_lock_timeout_error(error) is False + + +@pytest.mark.parametrize( + ("body", "retryable"), + [ + ('{"error_code":"P1017","message":"Server has closed the connection."}', True), + ('{"error_code":"P1000","message":"Authentication failed."}', False), + ('{"error_code":"P2021","message":"Table does not exist."}', False), + ('{"error_code":"P9999","message":"Unknown engine failure."}', False), + ("invalid engine response", False), + ], +) +def test_engine_connection_closure_is_recoverable_without_allowing_unverified_auth( + body: str, + retryable: bool, +) -> None: + from prisma.http import Response + + error: Final = EngineRequestError(Response(httpx.Response(500)), body) + assert PrismaDBExceptionHandler.is_database_transport_error(error) is retryable + assert PrismaDBExceptionHandler.is_permanent_database_fault(error) is not retryable + assert PrismaDBExceptionHandler.is_database_connection_error(error) is False + + +def test_connection_closure_code_survives_prisma_data_error_wrapping() -> None: + error: Final = DataError( + data={"user_facing_error": {"error_code": "P1017", "message": "Server has closed the connection."}} + ) + assert PrismaDBExceptionHandler.is_database_transport_error(error) is True + assert PrismaDBExceptionHandler.is_permanent_database_fault(error) is False diff --git a/tests/unit/proxy/db/test_routing_prisma_wrapper.py b/tests/unit/proxy/db/test_routing_prisma_wrapper.py index 911a6d42ed3..73225acd2c8 100644 --- a/tests/unit/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/unit/proxy/db/test_routing_prisma_wrapper.py @@ -254,6 +254,55 @@ def test_token_refresh_delegates_to_both_writer_and_reader(): reader.stop_token_refresh_task.assert_awaited_once() +def test_reader_token_refresh_stops_even_when_the_writer_task_failed() -> None: + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer: Final = MagicMock() + reader: Final = MagicMock() + writer.stop_token_refresh_task = AsyncMock(side_effect=RuntimeError("refresh failed")) + reader.stop_token_refresh_task = AsyncMock() + routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader) + + with pytest.raises(RuntimeError, match="refresh failed"): + asyncio.run(routing.stop_token_refresh_task()) + + reader.stop_token_refresh_task.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", (RuntimeError("refresh failed"), asyncio.CancelledError())) +async def test_reader_cleanup_finishes_before_writer_stop_failure_propagates(error: BaseException) -> None: + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + reader_stopping: Final = asyncio.Event() + release_reader: Final = asyncio.Event() + reader_stopped: Final = asyncio.Event() + + async def stop_reader() -> None: + reader_stopping.set() + await release_reader.wait() + reader_stopped.set() + + writer: Final = MagicMock() + reader: Final = MagicMock() + writer.stop_token_refresh_task = AsyncMock(side_effect=error) + reader.stop_token_refresh_task = AsyncMock(side_effect=stop_reader) + routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader) + stopping: Final = asyncio.create_task(routing.stop_token_refresh_task()) + try: + await reader_stopping.wait() + await asyncio.sleep(0) + assert not stopping.done() + release_reader.set() + with pytest.raises(type(error)): + await stopping + assert reader_stopped.is_set() + finally: + release_reader.set() + stopping.cancel() + await asyncio.gather(stopping, return_exceptions=True) + + def test_routed_actions_falls_back_to_writer_for_unknown_methods(): from litellm.proxy.db.routing_prisma_wrapper import _RoutedActions diff --git a/tests/unit/proxy/deps/__init__.py b/tests/unit/proxy/deps/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/deps/test_database.py b/tests/unit/proxy/deps/test_database.py new file mode 100644 index 00000000000..58502758987 --- /dev/null +++ b/tests/unit/proxy/deps/test_database.py @@ -0,0 +1,13 @@ +from typing import Final + +import pytest +from starlette.requests import Request + +from litellm.proxy.deps.database import DatabaseStateMissingError, get_database + + +@pytest.mark.asyncio +async def test_missing_lifespan_is_not_treated_as_a_disabled_database(): + request: Final = Request({"type": "http", "state": {}}) + with pytest.raises(DatabaseStateMissingError): + await get_database(request) diff --git a/tests/unit/proxy/deps/test_repositories.py b/tests/unit/proxy/deps/test_repositories.py new file mode 100644 index 00000000000..e802a4fda61 --- /dev/null +++ b/tests/unit/proxy/deps/test_repositories.py @@ -0,0 +1,112 @@ +from collections.abc import Mapping +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.models.deployment_warning import AdminIgnore +from litellm.models.gateway_activity import GatewayActivity +from litellm.models.user_banner import UserBanner +from litellm.proxy.deps.database import get_database +from litellm.proxy.deps.repositories import ( + DeploymentWarningIgnoreRepositoryDep, + GatewayActivityRepositoryDep, + UserBannerRepositoryDep, +) +from litellm.repositories import database as database_runtime +from tests._support.native_database import native_database + +NATIVE: Final = native_database() + + +def _app(database: database_runtime.Database | None) -> FastAPI: + @asynccontextmanager + async def lifespan(app: FastAPI): + yield {"database": database} + + app: Final = FastAPI(lifespan=lifespan) + + @app.get("/banner") + async def banner(banners: UserBannerRepositoryDep) -> UserBanner: + return UserBanner() if banners is None else await banners.get_banner() + + @app.get("/ignores") + async def ignores(ignores: DeploymentWarningIgnoreRepositoryDep) -> Mapping[str, AdminIgnore] | None: + return None if ignores is None else await ignores.get_ignores() + + @app.get("/activity") + async def activity(activity: GatewayActivityRepositoryDep) -> tuple[GatewayActivity, ...] | None: + return None if activity is None else await activity.get_activity("2026-01-01", "2026-01-31") + + return app + + +def test_lifespan_database_is_shared_until_shutdown(): + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + app: Final = _app(database_runtime.Database(prisma=prisma, native=NATIVE)) + with TestClient(app) as client: + assert client.get("/banner").json() == UserBanner().model_dump() + prisma.db.litellm_uisettings.find_unique.return_value = SimpleNamespace( + ui_settings={"enabled": True, "message": "published", "revision": "revision"} + ) + assert ( + client.get("/banner").json() + == UserBanner(enabled=True, message="published", revision="revision").model_dump() + ) + + +def test_both_repositories_use_the_lifespan_database(): + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock( + return_value=SimpleNamespace(ui_settings={"message": "stored banner"}) + ) + stored_ignores: Final = {"detailed_debug": {"ignored_by": "admin", "ignored_at": "2026-10-14T12:00:00Z"}} + prisma.writer_db.litellm_uisettings.find_unique = AsyncMock( + return_value=SimpleNamespace(ui_settings=stored_ignores) + ) + with TestClient(_app(database_runtime.Database(prisma=prisma, native=NATIVE))) as client: + assert client.get("/banner").json() == UserBanner(message="stored banner").model_dump() + assert client.get("/ignores").json() == { + "detailed_debug": {**stored_ignores["detailed_debug"], "source": "admin", "ignored_until": None} + } + + +def test_no_database_state_preserves_optional_repository(): + with TestClient(_app(None)) as client: + response: Final = client.get("/banner") + assert response.status_code == 200 + assert response.json() == UserBanner().model_dump() + + +def test_dependency_override_changes_next_request_without_changing_lifespan_owner(monkeypatch): + prisma: Final = MagicMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock( + return_value=SimpleNamespace(ui_settings={"message": "stored"}) + ) + app: Final = _app(database_runtime.Database(prisma=prisma, native=NATIVE)) + with TestClient(app) as client: + assert client.get("/banner").json() == UserBanner(message="stored").model_dump() + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + assert client.get("/banner").json() == UserBanner().model_dump() + + +def test_gateway_activity_reads_the_lifespan_reader_not_the_writer() -> None: + prisma: Final = MagicMock() + row: Final = {"kind": "status", "status_code": 429, "failed_requests": 3} + prisma.db.query_raw = AsyncMock(return_value=[row]) + prisma.writer_db.query_raw = AsyncMock(return_value=[{**row, "failed_requests": 99}]) + with TestClient(_app(database_runtime.Database(prisma=prisma, native=NATIVE))) as client: + response: Final = client.get("/activity") + assert response.status_code == 200 + assert response.json() == [row] + + +def test_gateway_activity_without_database_is_unavailable() -> None: + with TestClient(_app(None)) as client: + response: Final = client.get("/activity") + assert response.status_code == 200 + assert response.json() is None diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index 741fea5f2b9..0cdcb2e2d07 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -23,7 +23,6 @@ Tests cover: import datetime import json -import time from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -33,17 +32,22 @@ import respx from fastapi import HTTPException import litellm - from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( - HeadroomGuardrail, - has_headroom_retrieve_tool, HEADROOM_RETRIEVE_TOOL_NAME, + HeadroomGuardrail, + _session_affinity, + has_headroom_retrieve_tool, ) from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, ) from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, HEADROOM_STREAM_OPTIONS_KEY +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HEADROOM_REQUEST_CONTEXT_KEY, + HeadroomRequestContext, +) from litellm.types.utils import ( CallTypes, GenericGuardrailAPIInputs, @@ -148,6 +152,10 @@ def guardrail() -> HeadroomGuardrail: return _make_guardrail() +def _issued_context(*hashes: str, affinity: str | None = None) -> HeadroomRequestContext: + return HeadroomRequestContext(frozenset(hashes), affinity, float("inf")) + + def _recorded_guardrail_entries(request_data: dict) -> list: for container_key in ("metadata", "litellm_metadata"): container = request_data.get(container_key) @@ -175,20 +183,22 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages( structured_messages=ORIGINAL_MESSAGES, ) mock_response = _make_compress_response(COMPRESSED_MESSAGES) - request_data = {"model": "gpt-4o"} + auth: Final = UserAPIKeyAuth(api_key="tenant") + request_data: Final = {"session_id": "conversation", "metadata": {"user_api_key_auth": auth}} with patch.object( guardrail.async_handler, "post", new_callable=AsyncMock, return_value=mock_response, - ): + ) as mock_post: result = await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="request", ) + assert mock_post.call_args.kwargs["headers"]["x-headroom-affinity"] == _session_affinity(request_data) assert result.get("structured_messages") == EXPECTED_MESSAGES entries = _recorded_guardrail_entries(request_data) @@ -510,10 +520,6 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( tool_id="call_abc123", ) messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -530,11 +536,12 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac", affinity="original")}, ) assert plan.run_agentic_loop is True assert plan.request_patch is not None + assert HEADROOM_REQUEST_CONTEXT_KEY in plan.request_patch.kwargs follow_up = plan.request_patch.messages assert follow_up is not None @@ -545,6 +552,7 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( assert tool_result_message["tool_call_id"] == "call_abc123" mock_get.assert_called_once() + assert mock_get.call_args.kwargs["headers"]["x-headroom-affinity"] == "original" call_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0] assert "b573993006976af767214fac" in call_url @@ -576,10 +584,6 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( "content": "Retrieve more: hash=deadbeef000000000000dead", } ] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"deadbeef000000000000dead"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -596,7 +600,7 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("deadbeef000000000000dead")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -606,9 +610,13 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( @pytest.mark.asyncio +@pytest.mark.parametrize("context", [ + None, {"issued_hashes": ["deadbeef000000000000dead"]}, + HeadroomRequestContext(frozenset({"deadbeef000000000000dead"}), None, 0.0), +]) async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( - guardrail: HeadroomGuardrail, -): + guardrail: HeadroomGuardrail, context: object, +) -> None: """A hash-shaped string planted in message text must not be honored when this guardrail has no record of ever issuing it, even if it's echoed back in the current request's own messages (e.g. via prompt injection).""" @@ -625,8 +633,6 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( arguments={"hash": "deadbeef000000000000dead"}, tool_id="call_xyz", ) - assert not guardrail._issued_hashes_by_call_id - with patch.object( guardrail.async_handler, "get", @@ -641,7 +647,7 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-unknown"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: context}, ) mock_get.assert_not_called() @@ -653,17 +659,11 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( @pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_call( +async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_request( guardrail: HeadroomGuardrail, ): - """A hash issued for one request must not be retrievable by a different - request just because the second request echoes that hash-shaped string - back in its own messages -- retrieval must be scoped per litellm_call_id, - not derived by re-scanning attacker-controlled message text.""" - guardrail._issued_hashes_by_call_id["call-A"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) + issuer_data: Final[dict[str, object]] = {"model": "gpt-4o", "litellm_call_id": "call-A"} + await _apply(guardrail, ORIGINAL_MESSAGES, ["b573993006976af767214fac"], request_data=issuer_data) tool_calls = [ { @@ -693,7 +693,7 @@ async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_c anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-B"}, + kwargs={"litellm_call_id": "call-A"}, ) mock_get.assert_not_called() @@ -736,10 +736,6 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ } ] messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -756,7 +752,7 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -806,7 +802,6 @@ async def test_responses_retrieval_followup_reuses_the_model_the_caller_requeste "arguments": json.dumps({"hash": retrieve_hash}), } ] - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({retrieve_hash}), time.monotonic() + 999) with patch.object( guardrail.async_handler, @@ -832,7 +827,7 @@ async def test_responses_retrieval_followup_reuses_the_model_the_caller_requeste anthropic_messages_optional_request_params={}, logging_obj=logging_obj, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(retrieve_hash)}, ) assert plan.request_patch is not None @@ -876,10 +871,6 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag } ] messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -896,7 +887,7 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -929,16 +920,20 @@ HASH_SHAPED_HISTORY = [ ] -async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict: - request_data = {"model": "claude-sonnet-5"} +async def _apply( + guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None, + request_data: dict[str, object] | None = None, +) -> dict: + data: Final = request_data if request_data is not None else {"model": "claude-sonnet-5"} def _echo(**kwargs): + assert kwargs["json"].get("config") == ({"mode": "ccr"} if guardrail.ccr_retrieval else None) return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes) with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): return await guardrail.apply_guardrail( inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))), - request_data=request_data, + request_data=data, input_type="request", ) @@ -950,10 +945,11 @@ async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail tool must not appear and no hash may be registered as issued. Covers a service that omits ccr_hashes, returns it empty, or returns a non-list.""" for ccr_hashes in (None, [], "b573993006976af767214fac"): - result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes) + request_data: Final[dict[str, object]] = {} + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes, request_data=request_data) assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id + assert HEADROOM_REQUEST_CONTEXT_KEY not in request_data @pytest.mark.asyncio @@ -961,15 +957,18 @@ async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardr """Only the hashes the service reports in ccr_hashes are honored, in the service's own 12 to 24 hex grammar; anything else is dropped because each entry is interpolated into the /v1/retrieve URL.""" + request_data: Final[dict[str, object]] = {} result = await _apply( guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42], + request_data=request_data, ) assert has_headroom_retrieve_tool(result.get("tools") or []) - ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values() - assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) + context: Final = request_data[HEADROOM_REQUEST_CONTEXT_KEY] + assert isinstance(context, HeadroomRequestContext) + assert context.issued_hashes == frozenset({"98ca69107318", "b573993006976af767214fac"}) @pytest.mark.asyncio @@ -987,7 +986,6 @@ async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatc assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1] assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id @pytest.mark.asyncio @@ -1498,6 +1496,7 @@ async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_ ] guardrail = _make_guardrail(unreachable_fallback="fail_open") + request_data: Final = {HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("deadbeef000000000000dead")} inputs = GenericGuardrailAPIInputs( texts=["hello"], structured_messages=messages_with_fake_hash, @@ -1511,13 +1510,13 @@ async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_ ): result = await guardrail.apply_guardrail( inputs=inputs, - request_data={}, + request_data=request_data, input_type="request", ) assert result["structured_messages"] == messages_with_fake_hash assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id + assert HEADROOM_REQUEST_CONTEXT_KEY not in request_data @pytest.mark.asyncio @@ -2164,18 +2163,18 @@ def _openai_completion_payload(message: dict, finish_reason: str) -> dict: } -def _openai_tool_call_payload() -> dict: +def _openai_tool_call_payload(hash_value: str = CCR_HASH) -> dict: return _openai_completion_payload( { "role": "assistant", "content": None, "tool_calls": [ { - "id": "call_ccr", + "id": f"call_{hash_value}", "type": "function", "function": { "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), + "arguments": json.dumps({"hash": hash_value}), }, } ], @@ -2215,14 +2214,14 @@ def _openai_responses_payload(output_item: dict) -> dict: } -def _openai_responses_retrieve_call_payload() -> dict: +def _openai_responses_retrieve_call_payload(hash_value: str = CCR_HASH) -> dict: return _openai_responses_payload( { "type": "function_call", "id": "fc_ccr", - "call_id": "call_ccr", + "call_id": f"call_{hash_value}", "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), + "arguments": json.dumps({"hash": hash_value}), "status": "completed", } ) @@ -2381,10 +2380,6 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( to a client that never declared the tool.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "callbacks", [guardrail]) @@ -2392,6 +2387,7 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response(200, json=_openai_tool_call_payload()), + httpx.Response(200, json=_openai_tool_call_payload("a" * 24)), httpx.Response(200, json=_openai_text_payload(final_answer)), ] ) @@ -2408,20 +2404,22 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( tools=[_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH, "a" * 24, affinity="original-affinity")}, ) chunks = [chunk async for chunk in response] streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) assert streamed_text == final_answer assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert mock_get.call_count == 2 + assert mock_get.call_args.kwargs["url"].endswith("a" * 24) + assert all(call.kwargs["headers"]["x-headroom-affinity"] == "original-affinity" for call in mock_get.call_args_list) - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) + assert len(upstream.calls) == 3 + followup_body = json.loads(upstream.calls[2].request.content) assert not followup_body.get("stream") assert original_content in json.dumps(followup_body["messages"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) + assert not any(key.startswith("_headroom") for key in followup_body) @pytest.mark.asyncio @@ -2435,10 +2433,6 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( instead of streaming a headroom_retrieve function_call to the client.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "callbacks", [guardrail]) @@ -2446,6 +2440,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( side_effect=[ httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_retrieve_call_payload("a" * 24)), httpx.Response(200, json=_openai_responses_text_payload(final_answer)), ] ) @@ -2462,6 +2457,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( tools=[_responses_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH, "a" * 24, affinity="original-affinity")}, ) events = [event async for event in response] @@ -2471,14 +2467,15 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( assert streamed_text == final_answer assert not any("function_call" in str(getattr(event, "type", "")) for event in events) assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert mock_get.call_count == 2 + assert mock_get.call_args.kwargs["url"].endswith("a" * 24) + assert all(call.kwargs["headers"]["x-headroom-affinity"] == "original-affinity" for call in mock_get.call_args_list) - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) + assert len(upstream.calls) == 3 + followup_body = json.loads(upstream.calls[2].request.content) assert not followup_body.get("stream") assert original_content in json.dumps(followup_body["input"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) + assert not any(key.startswith("_headroom") for key in followup_body) def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( @@ -2491,10 +2488,6 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( completed response object.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -2518,6 +2511,7 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( tools=[_responses_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH)}, ) events = list(response) @@ -2917,7 +2911,6 @@ def _anthropic_response_with_text_and_tool_call() -> dict: async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({"h" * 24}), time.monotonic() + 60) logging_obj = MagicMock() logging_obj.litellm_call_id = "call-1" logging_obj.model_call_details = {} @@ -2938,7 +2931,7 @@ async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): anthropic_messages_optional_request_params={}, logging_obj=logging_obj, stream=False, - kwargs={}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("h" * 24)}, ) @@ -3033,7 +3026,7 @@ async def test_retrieve_call_passes_bounded_timeout(guardrail: HeadroomGuardrail new_callable=AsyncMock, return_value=_make_retrieve_response("original"), ) as mock_get: - result = await guardrail._call_retrieve("a" * 24) + result = await guardrail._call_retrieve("a" * 24, affinity=None) assert result == "original" timeout = _timeout_of(mock_get.call_args) @@ -3143,3 +3136,16 @@ def test_unusable_timeout_falls_back_to_the_default(configured: float): assert guardrail.timeout.read == 60.0 assert guardrail.timeout.connect == 5.0 + + +@pytest.mark.parametrize("session", [None, "conversation", "other"]) +@pytest.mark.parametrize("scope", [None, {}, {"api_key": "key"}, {"api_key": "other"}, {"user_id": "user"}]) +def test_affinity_requires_trusted_identity_and_caller_session( + scope: dict[str, str] | None, session: str | None +) -> None: + baseline: Final = {"session_id": "conversation", "metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key="key")}} + auth: Final = UserAPIKeyAuth(**scope) if scope else scope + candidate: Final = {"session_id": session, "metadata": {"user_api_key_auth": auth}} + affinity: Final = _session_affinity(candidate) + assert (affinity is None) is (session is None or not isinstance(auth, UserAPIKeyAuth)) + assert (affinity == _session_affinity(baseline)) is (candidate == baseline) diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index 60c3e4c22e1..e9e0a585949 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -26,19 +26,18 @@ import httpx import pytest import respx - import litellm from litellm.caching.caching import DualCache -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.litellm_core_utils.llm_cost_calc.finalization import finalize_cost -from litellm.types.utils import StandardLoggingPayload -from litellm.utils import dispatch_success_logging from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_cost_calc.finalization import finalize_cost from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import StandardLoggingPayload +from litellm.utils import dispatch_success_logging # --------------------------------------------------------------------------- # Helpers @@ -1520,12 +1519,12 @@ class TestArmDeferredStreamDispatch: sniffing the wrapper instead of the inner iterator armed the 1-arg native closure against the CSW's 2-arg stored shape and leaked a TypeError 500 frame into the stream.""" + from litellm.litellm_core_utils.add_retry_fallback_headers import ( + HiddenParamsAsyncIteratorWrapper, + ) from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) - from litellm.router_utils.add_retry_fallback_headers import ( - HiddenParamsAsyncIteratorWrapper, - ) logging_obj, recorded = self._dispatch_recording_logging_obj() wrapped = HiddenParamsAsyncIteratorWrapper(object.__new__(LiteLLMCompletionStreamingIterator)) diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index 6807a4fb569..7d912478877 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -13,7 +13,7 @@ import yaml from fastapi import HTTPException from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_endpoints import ( CreateGuardrailRequest, ListGuardrailsResponse, @@ -2198,6 +2198,39 @@ def _patch_apply_guardrail_env(mocker, guardrail_result): return mock_guardrail +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("vendor_status", "expected_type"), + [(401, "authentication_error"), (403, "permission_error")], +) +async def test_a_guardrail_vendors_rejection_is_labelled_the_way_chat_labels_it( + mocker, vendor_status: int, expected_type: str +): + """A vendor refusing the guardrail's own credentials is not the caller's key being rejected, + so it must not read as the proxy's auth_error the way a 401 the proxy raises does.""" + mock_guardrail: Final = _patch_apply_guardrail_env(mocker, None) + mock_guardrail.apply_guardrail.side_effect = HTTPException( + status_code=vendor_status, detail="vendor refused the request" + ) + mocker.patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", AsyncMock(return_value=None) + ) + + with pytest.raises(ProxyException) as raised: + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello"), + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert raised.value.to_dict() == { + "message": "vendor refused the request", + "type": expected_type, + "param": None, + "code": str(vendor_status), + } + + @pytest.mark.asyncio async def test_apply_guardrail_forwards_metadata_to_guardrail(mocker): """Client-supplied metadata must reach apply_guardrail via request_data so diff --git a/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py b/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py index 4583d8fd978..67aea8dfeb8 100644 --- a/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py @@ -1,10 +1,11 @@ import json import logging -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final +from unittest.mock import MagicMock import pytest from fastapi import FastAPI @@ -13,14 +14,15 @@ from pydantic import JsonValue from litellm._logging import verbose_logger from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.models.deployment_warning import AdminIgnore from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import get_deployment_warning_ignore_repository from litellm.proxy.health_endpoints.deployment_warning_endpoints import ( WarningCondition, config_audit_logger, current_warning_conditions, deployment_warning_conditions, - deployment_warning_ignore_store, license_expiry_condition, router, utc_now, @@ -28,7 +30,9 @@ from litellm.proxy.health_endpoints.deployment_warning_endpoints import ( from litellm.proxy.proxy_server import ( _dump_redacted_config, # pyright: ignore[reportPrivateUsage] # the exact redaction every config audit row goes through ) +from litellm.repositories.database import Database from litellm.repositories.deployment_warning_ignore_repository import DEPLOYMENT_WARNING_IGNORES_ROW_ID +from litellm.rust_bridge._native import NativeDatabase T0: Final = datetime(2026, 10, 14, 12, 0, tzinfo=timezone.utc) THIRTY_DAYS: Final = 30 * 24 * 60 * 60 @@ -36,13 +40,13 @@ THIRTY_DAYS: Final = 30 * 24 * 60 * 60 @dataclass class InMemoryIgnoreStore: - raw: object = None + saved: list[Mapping[str, AdminIgnore]] = field(default_factory=list) - async def get_raw_ignores(self) -> object: - return self.raw + async def get_ignores(self) -> Mapping[str, AdminIgnore]: + return self.saved[-1] if self.saved else {} - async def upsert_ignores(self, payload: str) -> None: - self.raw = json.loads(payload) + async def save_ignores(self, ignores: Mapping[str, AdminIgnore]) -> None: + self.saved.append(ignores) @dataclass @@ -87,7 +91,7 @@ class Harness: self.app.dependency_overrides[current_warning_conditions] = lambda: conditions def without_database(self) -> None: - self.app.dependency_overrides[deployment_warning_ignore_store] = lambda: None + self.app.dependency_overrides[get_deployment_warning_ignore_repository] = lambda: None def warnings(self) -> dict[str, dict[str, object]]: response = self.client.get("/health/deployment_warnings") @@ -102,7 +106,7 @@ def harness() -> Iterator[Harness]: store = InMemoryIgnoreStore() audit = RecordingAuditLogger() clock = [T0] - app.dependency_overrides[deployment_warning_ignore_store] = lambda: store + app.dependency_overrides[get_deployment_warning_ignore_repository] = lambda: store app.dependency_overrides[config_audit_logger] = lambda: audit app.dependency_overrides[utc_now] = lambda: clock[0] h = Harness(app=app, client=TestClient(app), store=store, audit=audit, clock=clock) @@ -233,28 +237,29 @@ class TestCurrentWarningConditions: class TestDeploymentWarningIgnoreStore: - def test_no_database_means_no_store(self, monkeypatch): - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) - assert deployment_warning_ignore_store() is None + @pytest.mark.asyncio + async def test_no_database_means_no_store(self): + assert await get_deployment_warning_ignore_repository(None) is None @pytest.mark.asyncio - async def test_reads_ignores_from_the_writer_not_a_lagging_replica(self, monkeypatch): + async def test_reads_ignores_from_the_writer_not_a_lagging_replica(self): def ui_settings_table(row: object) -> SimpleNamespace: async def find_unique(where: dict[str, str]) -> object: return row return SimpleNamespace(litellm_uisettings=SimpleNamespace(find_unique=find_unique)) - fresh = {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00Z"}} - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", - SimpleNamespace( - db=ui_settings_table(None), writer_db=ui_settings_table(SimpleNamespace(ui_settings=fresh)) - ), + fresh: Final = {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00Z"}} + prisma: Final = MagicMock( + db=ui_settings_table(None), writer_db=ui_settings_table(SimpleNamespace(ui_settings=fresh)) + ) + store: Final = await get_deployment_warning_ignore_repository( + Database(prisma=prisma, native=await NativeDatabase.connect()) ) - store = deployment_warning_ignore_store() assert store is not None - assert await store.get_raw_ignores() == fresh + assert await store.get_ignores() == { + "detailed_debug": AdminIgnore(ignored_by="a", ignored_at=datetime(2026, 10, 14, 12, tzinfo=timezone.utc)) + } class TestGetDeploymentWarnings: @@ -295,20 +300,6 @@ class TestGetDeploymentWarnings: ) assert harness.warnings()["no_redis"]["ignore"]["source"] == "env_var" - def test_stored_ignore_without_timezone_reads_as_unreadable(self, harness): - harness.store.raw = { - "detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00", "ignored_until": None} - } - warnings = harness.warnings() - assert set(warnings) == {"env_credential_login", "detailed_debug"} - assert warnings["detailed_debug"]["ignore"] is None - - def test_corrupt_stored_ignores_read_as_none(self, harness): - harness.store.raw = "{not json" - warnings = harness.warnings() - assert set(warnings) == {"env_credential_login", "detailed_debug"} - assert all(w["ignore"] is None for w in warnings.values()) - def test_without_database_ignoring_is_reported_unavailable(self, harness): harness.without_database() response = harness.client.get("/health/deployment_warnings") @@ -438,7 +429,7 @@ class TestUpdateDeploymentWarningIgnore: "/health/deployment_warnings/detailed_debug", json={"ignored": True, "snooze_seconds": snooze_seconds} ) assert response.status_code == 422 - assert harness.store.raw is None + assert harness.store.saved == [] def test_thirty_day_snooze_is_allowed(self, harness): response = harness.client.patch( @@ -457,14 +448,14 @@ class TestUpdateDeploymentWarningIgnore: def test_unknown_warning_is_rejected(self, harness, warning_id): response = harness.client.patch(f"/health/deployment_warnings/{warning_id}", json={"ignored": True}) assert response.status_code == 404 - assert harness.store.raw is None + assert harness.store.saved == [] @pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER]) def test_only_proxy_admins_can_ignore(self, harness, role): harness.as_role(role) response = harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True}) assert response.status_code == 403 - assert harness.store.raw is None + assert harness.store.saved == [] assert harness.audit.calls == [] def test_without_database_ignoring_is_unavailable(self, harness): diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 9847787efce..2258eb95fb8 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -1,8 +1,9 @@ import asyncio import copy +import itertools import json import time -from collections.abc import Awaitable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from datetime import datetime, timedelta from importlib import import_module @@ -22,7 +23,7 @@ import litellm import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64 from litellm.models.credentials import CredentialItem -from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, @@ -35,6 +36,7 @@ from litellm.proxy.health_endpoints._health_endpoints import ( test_model_connection as health_test_model_connection, ) from litellm.router import Router +from litellm.types.router import RouterModelGroupAliasItem from litellm.types.utils import oauth_token_exchange_litellm_params from litellm.types.workload_identity import ( ANTHROPIC_WIF_KWARGS_KEYS, @@ -925,7 +927,11 @@ async def test_test_model_connection_stored_audio_base_model_follows_the_stored_ deployment: Final = MappingProxyType( { "model_name": "custom-name-audio", - "litellm_params": {**deployment_params, "api_key": "fake-azure-key", "api_base": "https://x.openai.azure.com"}, + "litellm_params": { + **deployment_params, + "api_key": "fake-azure-key", + "api_base": "https://x.openai.azure.com", + }, "model_info": deployment_model_info, } ) @@ -1948,7 +1954,6 @@ def test_health_liveness_endpoint(proxy_client): assert duration_ms < 100, f"Health check took {duration_ms:.2f}ms, expected < 100ms for a simple endpoint" - def test_proxy_client_serves_full_proxy_routes_after_gateway_import(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy.proxy_server import app @@ -2472,6 +2477,7 @@ def _proxy_health_globals( llm_router: object, use_background_health_checks: bool = False, health_check_results: Mapping[str, object] | None = None, + prisma_client: object = None, ) -> Iterator[None]: with ( patch( # test-quality-ok: proxy module global, no injection seam @@ -2481,7 +2487,7 @@ def _proxy_health_globals( "litellm.proxy.proxy_server.llm_router", llm_router ), patch( # test-quality-ok: proxy module global, no injection seam - "litellm.proxy.proxy_server.prisma_client", None + "litellm.proxy.proxy_server.prisma_client", prisma_client ), patch( # test-quality-ok: proxy module global, no injection seam "litellm.proxy.proxy_server.use_background_health_checks", use_background_health_checks @@ -3812,7 +3818,11 @@ def test_clean_endpoint_data_never_displays_credential_fields(credential_field, async def _live_probed_model_ids( - model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth, model: str | None = None + model_list: Sequence[Mapping[str, object]], + user_api_key_dict: UserAPIKeyAuth, + model: str | None = None, + llm_router: Router | None = None, + prisma_client: object = None, ) -> set[str]: from fastapi import Response @@ -3825,7 +3835,7 @@ async def _live_probed_model_ids( return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} with ( - _proxy_health_globals(model_list, _router_for(model_list)), + _proxy_health_globals(model_list, llm_router or _router_for(model_list), prisma_client=prisma_client), patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", side_effect=fake_perform, @@ -4120,6 +4130,7 @@ async def _live_narrowed_model_ids( user_api_key_dict: UserAPIKeyAuth, model: str | None = None, model_id: str | None = None, + llm_router: Router | None = None, ) -> set[str]: from fastapi import Response @@ -4130,7 +4141,7 @@ async def _live_narrowed_model_ids( return probed, [], {} with ( - _proxy_health_globals(model_list, _router_for(model_list)), + _proxy_health_globals(model_list, llm_router or _router_for(model_list)), patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe ), @@ -4289,6 +4300,699 @@ async def test_health_endpoint_keeps_a_team_only_public_name_off_a_team_less_key probe.assert_not_awaited() +_PROVIDER_PREFIXED_MODEL_LIST: Final = [ + { + "model_name": "openai/gpt-5.4-nano", + "litellm_params": {"model": "openai/gpt-5.4-nano"}, + "model_info": {"id": "id-prefixed-nano"}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini"}, + }, + { + "model_name": "claude-haiku-5-5", + "litellm_params": {"model": "anthropic/claude-haiku-5-5"}, + "model_info": {"id": "id-haiku"}, + }, +] +_PROVIDER_PREFIXED_CACHED_RESULTS: Final = { + "healthy_endpoints": [ + {"model": "openai/gpt-5.4-nano", "model_id": "id-prefixed-nano"}, + {"model": "openai/gpt-5.4-mini", "model_id": "id-mini"}, + {"model": "anthropic/claude-haiku-5-5", "model_id": "id-haiku"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 3, + "unhealthy_count": 0, +} +_PROVIDER_WILDCARD_KEY: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["openai/*"]) +_RESTRICTED_TEAM_KEY: Final = UserAPIKeyAuth( + api_key="hashed-test-key", models=[], team_id="team-a", team_models=["gpt-5.4-mini"] +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller", "expected_ids"), + [(_PROVIDER_WILDCARD_KEY, {"id-prefixed-nano", "id-mini"}), (_RESTRICTED_TEAM_KEY, {"id-mini"})], +) +async def test_health_endpoint_probes_the_deployments_auth_lets_the_key_request( + caller: UserAPIKeyAuth, expected_ids: set[str] +): + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller) + + assert probed == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("caller", "expected_ids"), + [(_PROVIDER_WILDCARD_KEY, ["id-prefixed-nano", "id-mini"]), (_RESTRICTED_TEAM_KEY, ["id-mini"])], +) +async def test_health_endpoint_returns_the_deployments_auth_lets_the_key_request_on_background_cache_path( + caller: UserAPIKeyAuth, expected_ids: list[str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint(response=Response(), user_api_key_dict=caller, model=None, model_id=None) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_refuses_a_deployment_the_team_allowlist_forbids(use_background_health_checks: bool): + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), user_api_key_dict=_RESTRICTED_TEAM_KEY, model=None, model_id="id-prefixed-nano" + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +_KEY_ALIASED_AWAY_FROM_HAIKU: Final = UserAPIKeyAuth( + api_key="hashed-test-key", models=["gpt-5.4-mini"], aliases={"claude-haiku-5-5": "gpt-5.4-mini"} +) +_TEAM_ALIASED_AWAY_FROM_HAIKU: Final = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["claude-haiku-5-5"], + team_id="team-a", + team_model_aliases={"claude-haiku-5-5": "gpt-5.4-mini"}, +) +_MINI_ONLY_KEY: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini"]) +_ALIAS_SHADOW_CASES: Final = [ + pytest.param(_KEY_ALIASED_AWAY_FROM_HAIKU, {}, id="key-alias"), + pytest.param(_TEAM_ALIASED_AWAY_FROM_HAIKU, {}, id="team-alias"), + pytest.param(_MINI_ONLY_KEY, {"claude-haiku-5-5": "gpt-5.4-mini"}, id="router-alias"), +] + + +def _router_with_aliases(model_group_alias: Mapping[str, str]) -> Router: + return Router(model_list=copy.deepcopy(_PROVIDER_PREFIXED_MODEL_LIST), model_group_alias=dict(model_group_alias)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_probes_the_deployment_an_alias_sends_the_key_to_not_the_one_it_shadows( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + probed: Final = await _live_probed_model_ids( + _PROVIDER_PREFIXED_MODEL_LIST, caller, llm_router=_router_with_aliases(model_group_alias) + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_returns_the_deployment_an_alias_sends_the_key_to_on_background_cache_path( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_with_aliases(model_group_alias), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint(response=Response(), user_api_key_dict=caller, model=None, model_id=None) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_refuses_the_deployment_a_key_alias_shadows(use_background_health_checks: bool): + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_for(_PROVIDER_PREFIXED_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), user_api_key_dict=_KEY_ALIASED_AWAY_FROM_HAIKU, model=None, model_id="id-haiku" + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_targets_the_deployment_an_alias_name_sends_the_key_to( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + probed: Final = await _live_narrowed_model_ids( + _PROVIDER_PREFIXED_MODEL_LIST, + caller, + model="claude-haiku-5-5", + llm_router=_router_with_aliases(model_group_alias), + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_records_an_alias_check_under_the_name_the_caller_asked_for( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None): + return [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list], [], {} + + with ( + _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, _router_with_aliases(model_group_alias), prisma_client=MagicMock() + ), + patch( # test-quality-ok: the probe is the provider edge; the endpoint has no injection seam for it + "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe + ), + patch( # test-quality-ok: the history row is the assertion; the endpoint has no injection seam for its writer + "litellm.proxy.health_endpoints._health_endpoints._save_health_check_to_db", new_callable=AsyncMock + ) as save, + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=caller, model="claude-haiku-5-5", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + assert save.call_args.args[1] == "claude-haiku-5-5" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("caller", "model_group_alias"), _ALIAS_SHADOW_CASES) +async def test_health_endpoint_targets_the_deployment_an_alias_name_sends_the_key_to_on_background_cache_path( + caller: UserAPIKeyAuth, model_group_alias: Mapping[str, str] +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _PROVIDER_PREFIXED_MODEL_LIST, + _router_with_aliases(model_group_alias), + use_background_health_checks=True, + health_check_results=_PROVIDER_PREFIXED_CACHED_RESULTS, + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=caller, model="claude-haiku-5-5", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-mini"] + + +@pytest.mark.asyncio +async def test_health_endpoint_never_probes_a_deployment_without_a_model_name_for_a_restricted_key(): + nameless_deployment: Final = { + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-nameless"}, + } + + probed: Final = await _live_probed_model_ids( + [*_PROVIDER_PREFIXED_MODEL_LIST, nameless_deployment], + _MINI_ONLY_KEY, + llm_router=_router_for(_PROVIDER_PREFIXED_MODEL_LIST), + ) + + assert probed == {"id-mini"} + + +@pytest.mark.asyncio +async def test_health_endpoint_never_probes_the_deployment_a_keys_router_alias_sends_it_to_but_forbids(): + caller: Final = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["gpt-5.4-mini", "openai/gpt-5.4-nano"], + router_settings={"model_group_alias": {"gpt-5.4-mini": "claude-haiku-5-5"}}, + ) + + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller) + + assert probed == {"id-prefixed-nano"} + + +_WILDCARD_DEPLOYMENT: Final = { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild"}, +} +_CATCH_ALL_DEPLOYMENT: Final = {"model_name": "*", "litellm_params": {"model": "*"}, "model_info": {"id": "id-star"}} +_TEAM_WILDCARD_DEPLOYMENT: Final = { + "model_name": "openai-wild_team-a_7c1e", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-team-wild", "team_id": "team-a", "team_public_model_name": "openai/*"}, +} + + +def _team_key(team_models: list[str], team_model_aliases: Mapping[str, str] | None = None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed-test-key", + models=[], + team_id="team-a", + team_models=team_models, + team_model_aliases=dict(team_model_aliases) if team_model_aliases is not None else None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_list", "model_group_alias", "caller", "expected_ids"), + [ + pytest.param( + [_WILDCARD_DEPLOYMENT, _PROVIDER_PREFIXED_MODEL_LIST[2]], + {}, + _team_key(["openai/gpt-5.4-nano"]), + {"id-wild"}, + id="team-name-served-by-provider-wildcard", + ), + pytest.param( + [_TEAM_WILDCARD_DEPLOYMENT, _PROVIDER_PREFIXED_MODEL_LIST[2]], + {}, + _team_key(["openai/gpt-5.4-nano"]), + {"id-team-wild"}, + id="team-name-served-by-teams-own-wildcard", + ), + pytest.param( + [_CATCH_ALL_DEPLOYMENT], + {}, + _team_key(["gpt-5.4-mini"]), + {"id-star"}, + id="team-name-served-by-catch-all", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"gpt-4": "gpt-5.4-mini"}, + _team_key(["gpt-4"]), + {"id-mini"}, + id="team-allows-router-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {}, + _team_key(["my-alias"], {"my-alias": "gpt-5.4-mini"}), + {"id-mini"}, + id="team-allows-team-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"friendly-mini": RouterModelGroupAliasItem(model="gpt-5.4-mini", hidden=True)}, + UserAPIKeyAuth(api_key="hashed-test-key", models=["friendly-*"]), + {"id-mini"}, + id="wildcard-key-allows-hidden-router-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {}, + _team_key(["friendly-*"], {"friendly-mini": "gpt-5.4-mini"}), + {"id-mini"}, + id="wildcard-team-allows-team-alias", + ), + pytest.param( + _PROVIDER_PREFIXED_MODEL_LIST, + {"claude-haiku-5-5": "gpt-5.4-mini"}, + UserAPIKeyAuth(api_key="hashed-test-key", models=[]), + {"id-prefixed-nano", "id-mini", "id-haiku"}, + id="unrestricted-key-calls-shadowed-deployment-by-id", + ), + ], +) +async def test_health_endpoint_probes_every_deployment_the_callers_allowed_names_route_to( + model_list: Sequence[Mapping[str, object]], + model_group_alias: Mapping[str, str | RouterModelGroupAliasItem], + caller: UserAPIKeyAuth, + expected_ids: set[str], +): + llm_router: Final = Router( + model_list=copy.deepcopy([dict(deployment) for deployment in model_list]), + model_group_alias=dict(model_group_alias), + ) + + probed: Final = await _live_probed_model_ids(llm_router.get_model_list() or [], caller, llm_router=llm_router) + + assert probed == expected_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "caller", + [ + pytest.param( + UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini", "gpt-4-retired"]), id="key-grant" + ), + pytest.param(_team_key(["gpt-5.4-mini", "gpt-4-retired"]), id="team-grant"), + ], +) +async def test_health_endpoint_never_probes_the_default_fallback_for_a_granted_name_no_deployment_serves( + caller: UserAPIKeyAuth, +) -> None: + llm_router: Final = Router( + model_list=copy.deepcopy(_PROVIDER_PREFIXED_MODEL_LIST), default_fallbacks=["claude-haiku-5-5"] + ) + + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, caller, llm_router=llm_router) + + assert probed == {"id-mini"} + + +def _personal_key_whose_owner_may_call(models: list[str]) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import user_api_key_cache + + owner_id: Final = "health-scope-owner-" + "-".join(models) + user_api_key_cache.set_cache(key=owner_id, value=LiteLLM_UserTable(user_id=owner_id, models=models)) + return UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_id=owner_id) + + +def _managed_agent_key_granted(models: list[str]) -> UserAPIKeyAuth: + from litellm.types.agents import AgentResponse + + policy: Final = AgentResponse( + agent_id="managed", agent_name="Managed", agent_card_params={}, object_permission={"models": models} + ) + return UserAPIKeyAuth(api_key="hashed-test-key", models=[]).model_copy(update={"managed_agent_policy": policy}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("make_caller", "models", "expected_ids"), + [ + pytest.param(_personal_key_whose_owner_may_call, ["gpt-5.4-mini"], {"id-mini"}, id="key-owners-models"), + pytest.param(_managed_agent_key_granted, ["gpt-5.4-mini"], {"id-mini"}, id="managed-agents-models"), + pytest.param(_managed_agent_key_granted, [], set(), id="managed-agent-without-model-grants"), + ], +) +async def test_health_endpoint_probes_only_the_models_the_keys_user_or_agent_may_call( + make_caller: Callable[[list[str]], UserAPIKeyAuth], models: list[str], expected_ids: set[str] +): + probed: Final = await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, make_caller(models)) + + assert probed == expected_ids + + +@pytest.mark.asyncio +async def test_health_endpoint_refuses_a_model_id_the_keys_owner_may_not_call(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals(_PROVIDER_PREFIXED_MODEL_LIST, _router_for(_PROVIDER_PREFIXED_MODEL_LIST)), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=_personal_key_whose_owner_may_call(["gpt-5.4-mini"]), + model=None, + model_id="id-prefixed-nano", + ) + + assert refused.value.status_code == 403 + probe.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "make_caller", + [ + pytest.param(_personal_key_whose_owner_may_call, id="key-owner"), + pytest.param(_managed_agent_key_granted, id="managed-agent"), + ], +) +async def test_health_endpoint_keeps_the_wildcard_serving_a_name_the_keys_user_or_agent_may_call( + make_caller: Callable[[list[str]], UserAPIKeyAuth], +): + model_list: Final = [_WILDCARD_DEPLOYMENT, *_PROVIDER_PREFIXED_MODEL_LIST] + + probed: Final = await _live_probed_model_ids(model_list, make_caller(["openai/gpt-5.6"])) + + assert probed == {"id-wild"} + + +@pytest.mark.asyncio +async def test_health_endpoint_answers_an_agent_key_acting_for_an_unknown_user_the_way_chat_does(): + from litellm.types.agents import AgentCaller + + users_table: Final = MagicMock() + users_table.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + agent_key: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-mini"]).model_copy( + update={"agent_caller": AgentCaller(user_id="user-deleted-before-this-request")} + ) + + with pytest.raises(ProxyException) as failed: + await _live_probed_model_ids(_PROVIDER_PREFIXED_MODEL_LIST, agent_key, prisma_client=users_table) + + assert (failed.value.code, failed.value.type) == ("401", ProxyErrorTypes.auth_error) + assert "user-deleted-before-this-request" in failed.value.message + + +@pytest.mark.asyncio +async def test_health_endpoint_lets_other_requests_run_while_it_scopes_a_large_model_list(): + model_list: Final = [ + {"model_name": f"model-{n}", "litellm_params": {"model": f"openai/model-{n}"}, "model_info": {"id": f"id-{n}"}} + for n in range(50) + ] + other_request_turns: Final = itertools.count() + scoping_done: Final = asyncio.Event() + + async def scope() -> set[str]: + try: + return await _live_probed_model_ids( + model_list, UserAPIKeyAuth(api_key="hashed-test-key", models=["model-0"]) + ) + finally: + scoping_done.set() + + async def other_request() -> None: + while not scoping_done.is_set(): + next(other_request_turns) + await asyncio.sleep(0) + + probed, _ = await asyncio.gather(scope(), other_request()) + + assert probed == {"id-0"} + assert next(other_request_turns) >= len(model_list) + + +_WILDCARD_CACHED_RESULTS: Final = { + "healthy_endpoints": [ + {"model": "openai/gpt-5.6", "model_id": "id-wild"}, + {"model": "openai/gpt-5.6", "model_id": "id-team-wild"}, + {"model": "anthropic/claude-haiku-5-5", "model_id": "id-haiku"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 3, + "unhealthy_count": 0, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +@pytest.mark.parametrize( + ("wildcard_deployment", "expected_id"), + [ + pytest.param(_WILDCARD_DEPLOYMENT, "id-wild", id="proxy-wildcard"), + pytest.param(_TEAM_WILDCARD_DEPLOYMENT, "id-team-wild", id="teams-own-wildcard"), + ], +) +async def test_health_endpoint_model_param_reaches_the_wildcard_serving_a_team_allowed_name( + wildcard_deployment: Mapping[str, object], expected_id: str, use_background_health_checks: bool +): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + model_list: Final = [wildcard_deployment, _PROVIDER_PREFIXED_MODEL_LIST[2]] + + async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None): + return [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list], [], {} + + with ( + _proxy_health_globals( + model_list, + _router_for(model_list), + use_background_health_checks=use_background_health_checks, + health_check_results=_WILDCARD_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion + "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe + ), + ): + result: Final = await health_endpoint( + response=Response(), user_api_key_dict=_team_key(["openai/gpt-5.6"]), model="openai/gpt-5.6", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == [expected_id] + + +_ACCESS_GROUP_SPLIT_MODEL_LIST: Final = [ + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini-premium", "access_groups": ["premium"]}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-mini-standard"}, + }, +] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_access_groups_deployments_of_a_shared_name() -> None: + caller: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["premium"]) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_MODEL_LIST, caller) + + assert probed == {"id-mini-premium"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct_grant", ["gpt-5.4-mini", "*", "all-proxy-models"]) +async def test_health_endpoint_probes_every_deployment_of_a_name_the_key_also_holds_directly(direct_grant: str) -> None: + caller: Final = UserAPIKeyAuth(api_key="hashed-test-key", models=["premium", direct_grant]) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_MODEL_LIST, caller) + + assert probed == {"id-mini-premium", "id-mini-standard"} + + +_ACCESS_GROUP_SPLIT_WILDCARD_MODEL_LIST: Final = [ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild-premium", "access_groups": ["premium"]}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "id-wild-standard", "access_groups": ["standard"]}, + }, +] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_access_groups_wildcard_deployments_serving_a_name() -> None: + caller: Final = _personal_key_whose_owner_may_call(["openai/gpt-5.4-nano"]).model_copy( + update={"models": ["premium"]} + ) + + probed: Final = await _live_probed_model_ids(_ACCESS_GROUP_SPLIT_WILDCARD_MODEL_LIST, caller) + + assert probed == {"id-wild-premium"} + + +def _team_deployment( + internal_name: str, public_name: str, model: str, model_id: str, access_groups: list[str] +) -> Mapping[str, object]: + return { + "model_name": internal_name, + "litellm_params": {"model": model}, + "model_info": { + "id": model_id, + "team_id": "team-a", + "team_public_model_name": public_name, + "access_groups": access_groups, + }, + } + + +def _team_split_model_list(public_name: str, model: str) -> list[Mapping[str, object]]: + return [ + _team_deployment("model_name_team-a_7c1e4b2a", public_name, model, "id-team-premium", ["premium"]), + _team_deployment("model_name_team-a_9d2f6e01", public_name, model, "id-team-standard", ["standard"]), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("public_name", "model"), + [ + pytest.param("openai/*", "openai/*", id="team-wildcard"), + pytest.param("gpt-5.4-mini", "openai/gpt-5.4-mini", id="team-model"), + ], +) +async def test_health_endpoint_probes_only_the_access_groups_team_deployments_serving_a_name( + public_name: str, model: str +) -> None: + probed: Final = await _live_probed_model_ids(_team_split_model_list(public_name, model), _team_key(["premium"])) + + assert probed == {"id-team-premium"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("internal_name_grant", ["model_name_team-a_9d2f6e01", "model_name_team-a_*"]) +async def test_health_endpoint_probes_a_team_deployment_whose_internal_name_the_key_holds( + internal_name_grant: str, +) -> None: + caller: Final = _team_key(["premium", internal_name_grant]) + + probed: Final = await _live_probed_model_ids(_team_split_model_list("gpt-5.4-mini", "openai/gpt-5.4-mini"), caller) + + assert probed == {"id-team-premium", "id-team-standard"} + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_team_deployment_whose_internal_name_an_agent_holds() -> None: + from litellm.types.agents import AgentResponse + + policy: Final = AgentResponse( + agent_id="managed", + agent_name="Managed", + agent_card_params={}, + object_permission={"models": ["model_name_team-a_9d2f6e01"]}, + ) + caller: Final = _team_key(["gpt-5.4-mini"]).model_copy(update={"managed_agent_policy": policy}) + + probed: Final = await _live_probed_model_ids(_team_split_model_list("gpt-5.4-mini", "openai/gpt-5.4-mini"), caller) + + assert probed == {"id-team-standard"} + + def test_resolve_targeted_model_ids_lets_model_id_win_over_model(): resolve = _health_endpoints_module._resolve_targeted_model_ids diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 3aad2bda9bf..a50aa577a79 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -26,6 +26,7 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.litellm_core_utils import ptu_pricing +from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -49,7 +50,6 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.spend_tracking.input_tokens import models_counted_for from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token -from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.extraction import ExtractionResponse from litellm.types.llms.openai import ResponsesAPIResponse diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 8d4f8cd87ad..e482b6552a4 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -14,6 +14,7 @@ from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer @@ -36,6 +37,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payloa from litellm.proxy.utils import ( REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, + BilledFailedAttempt, ProxyUpdateSpend, ) from litellm.types.router import Deployment, LiteLLM_Params @@ -2676,6 +2678,229 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05 +@pytest.fixture() +def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> DualCache: + counters: Final = DualCache() + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", counters) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + return counters + + +def _failure_with_billed_attempts() -> Mapping[str, object]: + usage: Final = Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + return { + "combined_usage_object": usage, + "response_cost": 0.5, + "billed_failed_attempts": ( + BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25), + BilledFailedAttempt(attempted_retries=1, deployment_id="dep-b", usage=usage, response_cost=0.5), + ), + } + + +_LIVE_COUNTERS: Final = { + "spend:key:hashed-key": 0.2, + "spend:team:team-1": 0.3, + "spend:team_member:user-1:team-1": 0.35, + "spend:user:user-1": 0.4, + "spend:end_user:end-user-1": 0.45, + "spend:org:org-1": 0.5, + "spend:project:project-1": 0.55, + "spend:tag:batch-jobs": 0.6, + "spend:model_access_group:premium": 0.65, +} + + +async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None: + await logger.async_post_call_failure_hook( + request_data={ + "model": "gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"tags": ["batch-jobs"], MODEL_ACCESS_GROUP_METADATA_KEY: ["premium"]}, + "proxy_server_request": {"request_id": "rid"}, + **failure, + }, + original_exception=Exception("stream broke mid-flight"), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-key", + team_id="team-1", + user_id="user-1", + end_user_id="end-user-1", + org_id="org-1", + project_id="project-1", + ), + ) + + +async def _counter_values(counters: DualCache) -> dict[str, float]: + return {key: await counters.async_get_cache(key) for key in _LIVE_COUNTERS} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "charged"), + [ + ( + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + }, + 0.125, + ), + (_failure_with_billed_attempts(), 1.25), + ({}, 0.0), + ], +) +async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost( + live_spend_counters: DualCache, failure: Mapping[str, object], charged: float +): + """#35568: budget checks read the live counters, so a failure row's cost has to land there as well as in the + spend log, or an over-budget key keeps passing until the counter is next repaired from the database.""" + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request(ProxyDBLogger(), failure) + + assert await _counter_values(live_spend_counters) == pytest.approx( + {key: spend + charged for key, spend in _LIVE_COUNTERS.items()} + ) + + +@pytest.mark.asyncio +async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: DualCache): + """A failure whose cost was already finalized settles through the success path's write, which raises the + counters itself.""" + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request( + ProxyDBLogger(), {"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}} + ) + + assert await _counter_values(live_spend_counters) == pytest.approx( + {key: spend + 0.125 for key, spend in _LIVE_COUNTERS.items()} + ) + + +@pytest.mark.asyncio +async def test_failure_spend_rows_leave_the_live_spend_counters_alone_while_spend_updates_are_disabled( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_spend_updates": True}) + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request(ProxyDBLogger(), _failure_with_billed_attempts()) + + assert await _counter_values(live_spend_counters) == pytest.approx(_LIVE_COUNTERS) + + +@pytest.mark.asyncio +async def test_a_failing_spend_counter_store_is_still_sent_every_failure_spend_row(): + counters: Final = AsyncMock(side_effect=RuntimeError("redis unavailable")) + + await _fail_request(ProxyDBLogger(spend_counters=lambda: counters), _failure_with_billed_attempts()) + + assert [call.kwargs["response_cost"] for call in counters.call_args_list] == [0.5, 0.25, 0.5] + + +@pytest.mark.asyncio +@pytest.mark.timeout(30) +async def test_a_slow_spend_counter_store_lets_the_failure_out_and_raises_the_counters_once_it_answers(): + store_answers: Final = asyncio.Event() + raised_cost: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future() + + async def slow_counter_store(**counter_args: object) -> None: + await store_answers.wait() + raised_cost.set_result(counter_args["response_cost"]) + + logger: Final = ProxyDBLogger(spend_counters=lambda: slow_counter_store, failure_counter_wait_seconds=0) + + await _fail_request( + logger, + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + }, + ) + assert not raised_cost.done() + + store_answers.set() + assert await raised_cost == 0.125 + + +@pytest.mark.asyncio +async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_request_start( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + """The window row's one-time seed from the spend log stops at the request start, which keeps the failure row + it already logged from being counted twice.""" + from litellm.proxy.proxy_server import proxy_logging_obj + + window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue + await window_queue.flush_and_get_aggregated_window_spend_transactions() + key_cache: Final = DualCache() + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", key_cache) + request_start: Final = datetime(2026, 10, 20, 12, 0, tzinfo=timezone.utc) + reset_at: Final = request_start + timedelta(days=10) + await key_cache.async_set_cache( + "hashed-key", + UserAPIKeyAuth( + token="hashed-key", + budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}], + ), + ) + + await _fail_request( + ProxyDBLogger(), + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + "start_time": request_start, + }, + ) + + assert [ + (row["entity_type"], row["entity_id"], row["spend"], row["started_at"]) + for row in await window_queue.flush_and_get_aggregated_window_spend_transactions() + ] == [("key", "hashed-key", pytest.approx(0.125), "2026-10-20T12:00:00.000000")] + + +@pytest.mark.asyncio +async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_attributed_access_group( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + """Its spend row names the dropped deployment, which narrows nothing, so the counters must not narrow it to the + final attempt's groups either.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "sk-test"}, + "model_info": {"id": "dep-b", "access_groups": ["pool"]}, + } + ] + ), + ) + + await _fail_request( + ProxyDBLogger(), + { + **_failure_with_billed_attempts(), + "metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"], "model_info": {"id": "dep-b"}}, + }, + ) + + assert { + group: await live_spend_counters.async_get_cache(f"spend:model_access_group:{group}") + for group in ("premium", "pool") + } == pytest.approx({"premium": 0.25, "pool": 1.25}) + + @pytest.mark.asyncio async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): """MCP tool calls may only carry user_api_key; user/team rollups still need user_id.""" diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 8245c261a12..e43a72541c7 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -29,6 +29,7 @@ from PIL import Image from litellm import token_counter from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, estimate_image_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( @@ -46,7 +47,6 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, ) from litellm.proxy.utils import InternalUsageCache, hash_token -from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.llms.openai import ( InputTokensDetails, ResponseAPIUsage, diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index c6369a1f3d7..f808d07efd3 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -10,7 +10,6 @@ import pytest import litellm from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache -from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.proxy.hooks.model_max_budget_limiter import ( PROXY_VirtualKeyModelMaxBudgetLimiter, @@ -20,6 +19,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( model_budget_start_time_cache_key, resolve_model_budget, ) +from litellm.router_strategy.budget_limiter import budget_window_seconds async def _reset_at_of_open_window( @@ -30,7 +30,7 @@ async def _reset_at_of_open_window( entity_type=entity_type, entity_id=entity_id, budget_model=budget_model, budget_duration=period ) ) - return datetime.fromtimestamp(started_at + duration_in_seconds(period), tz=timezone.utc).isoformat() + return datetime.fromtimestamp(started_at + budget_window_seconds(started_at, period), tz=timezone.utc).isoformat() # Test class setup @@ -971,7 +971,7 @@ async def test_user_model_budget_window_resets_when_the_period_elapses(): # Age the window past its period. The next charge opens a new window rather # than adding to the exhausted one. - elapsed = duration_in_seconds("1mo") + 60 + elapsed = int(timedelta(days=31, minutes=1).total_seconds()) await dual_cache.async_set_cache( key=start_time_key, value=datetime.now(timezone.utc).timestamp() - elapsed, diff --git a/tests/unit/proxy/image_endpoints/test_endpoints.py b/tests/unit/proxy/image_endpoints/test_endpoints.py index f4aebecc11e..ea37b7b0213 100644 --- a/tests/unit/proxy/image_endpoints/test_endpoints.py +++ b/tests/unit/proxy/image_endpoints/test_endpoints.py @@ -2,8 +2,10 @@ import asyncio import copy import logging from collections.abc import Iterator, Mapping +from datetime import datetime +from queue import Queue from types import SimpleNamespace -from typing import Any, Dict +from typing import Any, Dict, Final, cast import orjson import pytest @@ -12,10 +14,20 @@ from fastapi.testclient import TestClient from starlette.requests import Request from starlette.responses import Response -from litellm._logging import verbose_proxy_logger +import litellm +from litellm._logging import session_id_var, trace_id_var, verbose_proxy_logger +from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.image_endpoints import endpoints +from litellm.proxy.proxy_server import PrismaClient, ProxyConfig +from litellm.proxy.utils import ProxyLogging +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_failures_for_current_minute +from litellm.types.utils import ImageObject, ImageResponse +from litellm.utils import custom_llm_setup @pytest.mark.asyncio @@ -410,3 +422,296 @@ async def test_failure_before_the_provider_call_bills_the_callers_litellm_call_i assert raised.value.headers["x-litellm-call-id"] == call_id assert [data["litellm_call_id"] for data in hook_request_data] == [call_id] + + +class _ImageRouterHandler(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[str] = [] + + async def aimage_generation( + self, + model: str, + prompt: str, + model_response: ImageResponse, + api_key: str | None, + api_base: str | None, + optional_params: Mapping[str, object], + logging_obj: object, + timeout: object | None = None, + client: object | None = None, + ) -> ImageResponse: + self.calls.append(model) + if model.endswith("failing"): + raise litellm.RateLimitError(message="primary failed", llm_provider="image-regression", model=model) + return ImageResponse(created=1, data=[ImageObject(url="https://images.example/generated.png")]) + + +def _image_generation_client( + monkeypatch: pytest.MonkeyPatch, + user_api_key_dict: UserAPIKeyAuth, + team_settings: Mapping[str, object] | None = None, + working_params: Mapping[str, object] | None = None, +) -> tuple[TestClient, _ImageRouterHandler]: + from litellm.proxy import proxy_server + + handler: Final = _ImageRouterHandler() + cache: Final = UserApiKeyCache() + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list]) + monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers]) + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "image-regression", "custom_handler": handler}]) + custom_llm_setup() + if team_settings is not None: + team_id: Final = user_api_key_dict.team_id + assert team_id is not None + cache.set_cache( + key=f"team_id:{team_id}", + value=LiteLLM_TeamTableCachedObj(team_id=team_id, router_settings=dict(team_settings)), + local_only=True, + model_type=LiteLLM_TeamTableCachedObj, + ) + monkeypatch.setattr(proxy_server, "prisma_client", cast(PrismaClient, object())) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "A", + "litellm_params": {"model": "image-regression/failing"}, + "model_info": {"id": "image-regression-failing"}, + }, + { + "model_name": "B", + "litellm_params": {"model": "image-regression/working", **(working_params or {})}, + "model_info": {"id": "image-regression-working"}, + }, + ], + num_retries=0, + ) + monkeypatch.setattr(litellm, "model_alias_map", {}) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "proxy_config", ProxyConfig()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "version", "test-version") + + app: Final = FastAPI() + app.include_router(endpoints.router) + app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler) + app.dependency_overrides[user_api_key_auth] = lambda: user_api_key_dict + return TestClient(app, raise_server_exceptions=False), handler + + +class _ImageOutcomeLogger(CustomLogger): + def __init__(self) -> None: + self.events: Queue[tuple[object, object, object]] = Queue() + self.messages: Queue[object] = Queue() + self.correlation: Queue[tuple[str, str]] = Queue() + + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + self.messages.put(kwargs.get("messages")) + self.correlation.put((trace_id_var.get(), session_id_var.get())) + self.events.put((kwargs.get("call_type"), kwargs.get("litellm_call_id"), kwargs.get("exception"))) + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + self.messages.put(kwargs.get("messages")) + self.events.put((kwargs.get("call_type"), kwargs.get("litellm_call_id"), response_obj)) + + +class _RewritingImageLogger(_ImageOutcomeLogger): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str + ) -> dict[str, object]: + return {**data, "messages": [{"role": "user", "content": "sanitized prompt"}]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path", ("/v1/images/generations", "/images/generations", "/openai/deployments/B/images/generations") +) +async def test_router_refused_image_generation_reaches_failure_callbacks( + monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + from litellm.proxy import proxy_server + + logger: Final = _ImageOutcomeLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + client, handler = _image_generation_client( + monkeypatch, UserAPIKeyAuth(), working_params={"max_budget": 1, "budget_duration": "1d"} + ) + router: Final = proxy_server.llm_router + assert router is not None and router.router_budget_logger is not None + router.router_budget_logger.dual_cache.set_cache("deployment_spend:image-regression-working:1d", 2) + + with client: + response: Final = client.post( + path, json={"model": "B", "prompt": "draw a tree"}, headers={"x-litellm-call-id": "refused-image"} + ) + assert response.status_code == 429, response.text + call_type, call_id, exception = logger.events.get(timeout=5) + + assert (call_type, call_id) == ("aimage_generation", "refused-image") + assert isinstance(exception, ValueError) and "crossed budget" in str(exception) + assert handler.calls == [] + assert logger.events.empty() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("supply_logging", (False, True)) +async def test_image_router_refusal_without_or_with_caller_logging( + monkeypatch: pytest.MonkeyPatch, supply_logging: bool +) -> None: + from litellm.proxy import proxy_server + + logger: Final = _ImageOutcomeLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + _image_generation_client(monkeypatch, UserAPIKeyAuth(), working_params={"max_budget": 1, "budget_duration": "1d"}) + router: Final = proxy_server.llm_router + assert router is not None and router.router_budget_logger is not None + router.router_budget_logger.dual_cache.set_cache("deployment_spend:image-regression-working:1d", 2) + kwargs: Final[dict[str, object]] = {"litellm_trace_id": "refused-trace", "litellm_session_id": "refused-session"} + if supply_logging: + logging_obj, _ = litellm.utils.function_setup( + "aimage_generation", litellm.utils.Rules(), datetime.now(), + model="B", prompt="draw a tree", litellm_call_id="supplied-image", **kwargs, + ) + kwargs.update(litellm_logging_obj=logging_obj, litellm_call_id="supplied-image") + + context_before: Final = (trace_id_var.get(), session_id_var.get()) + with pytest.raises(ValueError, match="crossed budget"): + await router.aimage_generation(model="B", prompt="draw a tree", **kwargs) + assert (trace_id_var.get(), session_id_var.get()) == context_before + call_type, call_id, exception = await asyncio.to_thread(logger.events.get, timeout=5) + assert call_type == "aimage_generation" and isinstance(exception, ValueError) + assert call_id == "supplied-image" if supply_logging else isinstance(call_id, str) and call_id + assert logger.events.empty() + assert logger.correlation.get_nowait() == ("refused-trace", "refused-session") + + +@pytest.mark.parametrize(("model", "status_code"), (("A", 429), ("B", 200))) +@pytest.mark.parametrize("rewrite_prompt", (False, True)) +def test_image_generation_provider_outcome_is_logged_once( + monkeypatch: pytest.MonkeyPatch, model: str, status_code: int, rewrite_prompt: bool +) -> None: + logger: Final = _RewritingImageLogger() if rewrite_prompt else _ImageOutcomeLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth()) + with client: + response: Final = client.post( + "/v1/images/generations", + json={"model": model, "prompt": "draw a tree"}, + headers={"x-litellm-call-id": "provider-image"}, + ) + assert response.status_code == status_code, response.text + call_type, call_id, outcome = logger.events.get(timeout=5) + assert (call_type, call_id) == ("aimage_generation", "provider-image") + assert isinstance(outcome, litellm.RateLimitError if model == "A" else ImageResponse) + assert logger.messages.get_nowait() == [ + {"role": "user", "content": "sanitized prompt" if rewrite_prompt else "draw a tree"} + ] + assert len(handler.calls) == 1 + assert logger.events.empty() + + +def test_image_generation_missing_prompt_keeps_its_validation_error(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth()) + with client: + response: Final = client.post("/v1/images/generations", json={"model": "B"}) + assert response.status_code == 400, response.text + assert response.json()["error"]["param"] == "prompt" + assert handler.calls == [] + + +def test_image_generation_uses_key_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"fallbacks": [{"A": ["B"]}]}), + ) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["failing", "working"], response.text + + +def test_image_generation_uses_team_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None: + team_id: Final = "image-regression-team" + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(team_id=team_id), + {"fallbacks": [{"A": ["B"]}]}, + ) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["failing", "working"], response.text + + +def test_image_generation_without_router_settings_does_not_use_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth()) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 429, response.text + assert handler.calls == ["failing"], response.text + + +def test_image_generation_uses_key_model_group_alias(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"model_group_alias": {"alias": "B"}}), + ) + + response: Final = client.post("/v1/images/generations", json={"model": "alias", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["working"], response.text + + +def test_image_generation_uses_key_num_retries(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + logger: Final = _ImageOutcomeLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"num_retries": 2}), + ) + + with client: + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + events: Final = [logger.events.get(timeout=5) for _ in range(3)] + + assert response.status_code == 429, response.text + assert handler.calls == ["failing", "failing", "failing"], response.text + assert all(isinstance(event[2], litellm.RateLimitError) for event in events) + assert logger.events.empty() + router: Final = proxy_server.llm_router + assert router is not None + assert get_deployment_failures_for_current_minute(router, "image-regression-failing") == 3 diff --git a/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index ceeac4919f2..65ebece406a 100644 --- a/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -20,6 +20,7 @@ from litellm.proxy._types import ( # Now we can safely import app from litellm.proxy.proxy_server import app +from litellm.router import Router from litellm.types.search import SearchToolInfoResponse from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs @@ -991,8 +992,7 @@ def _live_router_and_db(db_rows: list): proxy_config = ProxyConfig() proxy_config.update_config_state({}) - fake_router = MagicMock() - fake_router.search_tools = list(db_rows) + router = Router(model_list=[], search_tools=list(db_rows)) with contextlib.ExitStack() as stack: stack.enter_context( @@ -1002,7 +1002,7 @@ def _live_router_and_db(db_rows: list): patch("litellm.proxy.proxy_server.proxy_config", proxy_config) ) # test-quality-ok: proxy globals are the only seam; see the module note above stack.enter_context( - patch("litellm.proxy.proxy_server.llm_router", fake_router) + patch("litellm.proxy.proxy_server.llm_router", router) ) # test-quality-ok: proxy globals are the only seam; see the module note above stack.enter_context( patch( # test-quality-ok: proxy globals are the only seam; see the module note above @@ -1020,7 +1020,7 @@ def _live_router_and_db(db_rows: list): user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" ) try: - yield fake_router + yield router finally: app.dependency_overrides.pop(user_api_key_auth, None) @@ -1028,7 +1028,7 @@ def _live_router_and_db(db_rows: list): @pytest.mark.asyncio async def test_create_search_tool_reaches_the_router_before_the_response(): """A UI-created tool must be usable immediately, not only after the next config reload tick.""" - with _live_router_and_db([]) as fake_router: + with _live_router_and_db([]) as router: response = TestClient(app).post( "/search_tools", json={ @@ -1040,12 +1040,12 @@ async def test_create_search_tool_reaches_the_router_before_the_response(): ) assert response.status_code == 200 - assert [tool["search_tool_name"] for tool in fake_router.search_tools] == ["tavily-search"] + assert [tool["search_tool_name"] for tool in router.settings().search_tools] == ["tavily-search"] @pytest.mark.asyncio async def test_update_search_tool_reaches_the_router_before_the_response(): - with _live_router_and_db([_search_tool_row("tavily-search", provider="tavily")]) as fake_router: + with _live_router_and_db([_search_tool_row("tavily-search", provider="tavily")]) as router: response = TestClient(app).put( "/search_tools/tavily-search-id", json={ @@ -1057,17 +1057,17 @@ async def test_update_search_tool_reaches_the_router_before_the_response(): ) assert response.status_code == 200 - assert fake_router.search_tools[0]["litellm_params"]["search_provider"] == "exa_ai" + assert router.settings().search_tools[0]["litellm_params"]["search_provider"] == "exa_ai" @pytest.mark.asyncio async def test_delete_search_tool_removes_it_from_the_router(): """Deleting the last tool must clear the router; the old empty-list guard left it live.""" - with _live_router_and_db([_search_tool_row("tavily-search")]) as fake_router: + with _live_router_and_db([_search_tool_row("tavily-search")]) as router: response = TestClient(app).delete("/search_tools/tavily-search-id") assert response.status_code == 200 - assert fake_router.search_tools == [] + assert list(router.settings().search_tools) == [] @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py b/tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py index 7153fb909a8..c3bec658457 100644 --- a/tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py @@ -1,47 +1,25 @@ -import os from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Final -from unittest.mock import AsyncMock, MagicMock, call, patch - -from tests._master_key import MASTER_KEY - -# Patching ``litellm.proxy.proxy_server.prisma_client`` imports that module, whose -# module-level setup reads DATABASE_URL and LITELLM_MASTER_KEY. Tier-zero runners -# set neither, so pin throwaways first, as test_component_allowlists.py does. The -# prior values are restored below so a non-postgres URL cannot leak into sibling -# tests sharing the xdist worker and make them treat a phantom database as live. -_THROWAWAY_ENV = { - "DATABASE_URL": "sqlite:///:memory:", - "LITELLM_MASTER_KEY": MASTER_KEY, -} -_PRE_EXISTING_ENV = {key: os.environ.get(key) for key in _THROWAWAY_ENV} -for _key, _value in _THROWAWAY_ENV.items(): - os.environ.setdefault(_key, _value) +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import get_gateway_activity_repository from litellm.proxy.management_endpoints.gateway_request_endpoints import ( - _ACTIVITY_SQL, - _AggregateRow, _default_range, _fold_by_date, _fold_by_route, _fold_by_status_code, - _StatusCodeAggregateRow, get_gateway_daily_activity, router, ) - -for _key, _previous in _PRE_EXISTING_ENV.items(): - if _previous is None: - os.environ.pop(_key, None) - else: - os.environ[_key] = _previous +from litellm.repositories.gateway_activity_repository import GatewayActivityRepository # The handler stamps "today" from the wall clock, so any assertion that names a # date has to pin it. Recomputing the expected range in the assertion instead @@ -67,8 +45,8 @@ def _row( route: str = "/chat/completions", successful: int = 0, failed: int = 0, -) -> _AggregateRow: - return _AggregateRow( +) -> GatewayRouteActivity: + return GatewayRouteActivity( kind="route", date=date, category=category, @@ -86,9 +64,7 @@ def _prisma_returning(rows: Sequence[Mapping[str, str | int | None]] | None) -> client: Final = MagicMock() client.db = MagicMock() - async def query_raw(query: str, *_args: str) -> Sequence[Mapping[str, str | int | None]] | None: - if query != _ACTIVITY_SQL: - raise AssertionError(f"Unexpected gateway activity query: {query}") + async def query_raw(_query: str, *_args: str) -> Sequence[Mapping[str, str | int | None]] | None: return rows client.db.query_raw = AsyncMock(side_effect=query_raw) @@ -171,9 +147,9 @@ class TestFoldByStatusCode: def test_orders_by_failed_count_then_status_code(self): folded: Final = _fold_by_status_code( ( - _StatusCodeAggregateRow(kind="status", status_code=503, failed_requests=1), - _StatusCodeAggregateRow(kind="status", status_code=500, failed_requests=2), - _StatusCodeAggregateRow(kind="status", status_code=429, failed_requests=2), + GatewayStatusCodeActivity(kind="status", status_code=503, failed_requests=1), + GatewayStatusCodeActivity(kind="status", status_code=500, failed_requests=2), + GatewayStatusCodeActivity(kind="status", status_code=429, failed_requests=2), ) ) assert [(entry.status_code, entry.failed_requests) for entry in folded] == [(429, 2), (500, 2), (503, 1)] @@ -191,11 +167,11 @@ class TestGatewayDailyActivityEndpoint: ], ) async def test_refuses_every_non_admin_role(self, role): - with patch("litellm.proxy.proxy_server.prisma_client", _prisma_returning([])): - with pytest.raises(HTTPException) as exc: - await get_gateway_daily_activity( - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role), - ) + with pytest.raises(HTTPException) as exc: + await get_gateway_daily_activity( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role), + activity=None, + ) assert exc.value.status_code == 403 @pytest.mark.asyncio @@ -204,17 +180,18 @@ class TestGatewayDailyActivityEndpoint: [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], ) async def test_serves_both_admin_roles(self, role): - with patch("litellm.proxy.proxy_server.prisma_client", _prisma_returning([])): - response = await get_gateway_daily_activity( - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role), - ) + response = await get_gateway_daily_activity( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role), + activity=GatewayActivityRepository(_prisma_returning([])), + start_date="2026-01-01", + end_date="2026-01-31", + ) assert response.total_successful_requests == 0 @pytest.mark.asyncio async def test_reports_db_not_connected_rather_than_crashing(self): - with patch("litellm.proxy.proxy_server.prisma_client", None): - with pytest.raises(HTTPException) as exc: - await get_gateway_daily_activity(user_api_key_dict=_admin()) + with pytest.raises(HTTPException) as exc: + await get_gateway_daily_activity(user_api_key_dict=_admin(), activity=None) assert exc.value.status_code == 500 @pytest.mark.asyncio @@ -267,8 +244,12 @@ class TestGatewayDailyActivityEndpoint: }, ) prisma: Final = _prisma_returning(rows) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - response: Final = await get_gateway_daily_activity(user_api_key_dict=_admin()) + response: Final = await get_gateway_daily_activity( + user_api_key_dict=_admin(), + activity=GatewayActivityRepository(prisma), + start_date="2026-08-03", + end_date="2026-08-04", + ) assert response.total_successful_requests == 16 assert response.total_failed_requests == 4 @@ -277,13 +258,16 @@ class TestGatewayDailyActivityEndpoint: assert [(entry.status_code, entry.failed_requests) for entry in response.by_status_code] == [(429, 3), (500, 1)] assert [entry.date for entry in response.by_date] == ["2026-08-03", "2026-08-04"] assert [entry.route for entry in response.by_route] == ["/chat/completions", "/embeddings"] - assert prisma.db.query_raw.call_count == 1 @pytest.mark.asyncio async def test_a_null_result_set_is_not_an_error(self): client = _prisma_returning(None) - with patch("litellm.proxy.proxy_server.prisma_client", client): - response = await get_gateway_daily_activity(user_api_key_dict=_admin()) + response = await get_gateway_daily_activity( + user_api_key_dict=_admin(), + activity=GatewayActivityRepository(client), + start_date="2026-01-01", + end_date="2026-01-31", + ) assert response.total_successful_requests == 0 assert response.by_date == () assert response.by_route == () @@ -305,23 +289,23 @@ class TestGatewayDailyActivityRoute: app = FastAPI() app.include_router(router) app.dependency_overrides[user_api_key_auth] = _admin - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - response = TestClient(app).get( - "/gateway/daily/activity", - params={"start_date": "2026-01-01", "end_date": "2026-01-31"}, - ) + app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma) + response = TestClient(app).get( + "/gateway/daily/activity", + params={"start_date": "2026-01-01", "end_date": "2026-01-31"}, + ) assert response.status_code == 200 - assert prisma.db.query_raw.call_args_list == [call(_ACTIVITY_SQL, "2026-01-01", "2026-01-31")] + assert prisma.db.query_raw.call_args_list[0].args[1:] == ("2026-01-01", "2026-01-31") def test_omitted_dates_fall_back_to_the_default_window(self, frozen_clock): prisma = _prisma_returning([]) app = FastAPI() app.include_router(router) app.dependency_overrides[user_api_key_auth] = _admin - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - response = TestClient(app).get("/gateway/daily/activity") + app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma) + response = TestClient(app).get("/gateway/daily/activity") assert response.status_code == 200 - assert prisma.db.query_raw.call_args_list == [call(_ACTIVITY_SQL, *_FROZEN_RANGE)] + assert prisma.db.query_raw.call_args_list[0].args[1:] == _FROZEN_RANGE def test_serialized_response_carries_the_documented_shape(self): prisma: Final = _prisma_returning( @@ -349,8 +333,8 @@ class TestGatewayDailyActivityRoute: app: Final = FastAPI() app.include_router(router) app.dependency_overrides[user_api_key_auth] = _admin - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - body: Final = TestClient(app).get("/gateway/daily/activity").json() + app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma) + body: Final = TestClient(app).get("/gateway/daily/activity").json() assert body == { "total_successful_requests": 7, diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index a0c32f2235f..0198f582be7 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -7,7 +7,7 @@ from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch -from uuid import UUID +from uuid import UUID, uuid4 import pytest import yaml @@ -16,6 +16,7 @@ from fastapi.testclient import TestClient from pydantic import JsonValue, ValidationError import litellm +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.models.object_permission import LiteLLM_ObjectPermissionTable @@ -41,12 +42,15 @@ from litellm.proxy._types import ( RegenerateKeyRequest, ResetSpendRequest, UpdateKeyRequest, + hash_token, ) from litellm.proxy.auth.auth_checks import ( delete_cache_key_object, + get_key_object, jwt_key_mapping_cache_key, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -63,6 +67,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _enforce_upperbound_key_params, _execute_virtual_key_regeneration, _get_and_validate_existing_key, + _hash_token_if_needed, _is_team_key, _key_generation_required_param_check, _member_key_request, @@ -73,6 +78,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _team_admin_requests_member_key, _team_key_generation_check, _transform_verification_tokens_to_deleted_records, + _update_key_row_assigning_project, _validate_end_user_budget_id_change, _validate_max_budget, _validate_reset_spend_value, @@ -5755,6 +5761,231 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat ) +class _ChannelSubscription: + """Marks a message done when the subscriber comes back for the next one, so ``messages.join()`` + returns once every delivered message has been applied""" + + def __init__(self) -> None: + self.messages: asyncio.Queue[dict[str, object]] = asyncio.Queue() + self._holding_message = False + + async def subscribe(self, *channels: str) -> None: + return None + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> dict[str, object] | None: + if self._holding_message: + self._holding_message = False + self.messages.task_done() + try: + message: Final = await asyncio.wait_for(self.messages.get(), timeout) + except asyncio.TimeoutError: + return None + self._holding_message = True + return message + + async def aclose(self) -> None: + return None + + +class _ChannelPipeline: + def __init__(self, redis: "_SharedCoordinationRedis") -> None: + self._redis = redis + self._messages: tuple[str, ...] = () + + def publish(self, channel: str, message: str) -> "_ChannelPipeline": + self._messages = (*self._messages, message) + return self + + async def execute(self) -> None: + for message in self._messages: + self._redis.deliver(message) + + +class _SharedCoordinationRedis: + """The coordination Redis two workers share: whatever one worker publishes reaches the other's subscriber""" + + namespace = None + + def __init__(self) -> None: + self.published: tuple[str, ...] = () + self._subscriptions: tuple[_ChannelSubscription, ...] = () + + def init_pubsub_client(self) -> "_SharedCoordinationRedis": + return self + + def pubsub(self) -> _ChannelSubscription: + subscription: Final = _ChannelSubscription() + self._subscriptions = (*self._subscriptions, subscription) + return subscription + + def pipeline(self, transaction: bool = True) -> _ChannelPipeline: + return _ChannelPipeline(self) + + async def publish(self, channel: str, message: str) -> int: + self.deliver(message) + return 1 + + def deliver(self, message: str) -> None: + self.published = (*self.published, message) + for subscription in self._subscriptions: + subscription.messages.put_nowait({"type": "message", "data": message.encode()}) + + async def until_every_subscriber_applied_its_messages(self) -> None: + _ = await asyncio.gather(*(subscription.messages.join() for subscription in self._subscriptions)) + + +def _owned_key_row(hashed_token: str) -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken( + token=hashed_token, + user_id="user-123", + team_id=None, + key_alias="doomed-key", + spend=0.0, + max_budget=None, + models=[], + aliases={}, + config={}, + permissions={}, + metadata={}, + model_max_budget={}, + model_spend={}, + soft_budget_cooldown=False, + allowed_routes=[], + ) + + +def _prisma_with_key_row(key_row: LiteLLM_VerificationToken) -> AsyncMock: + prisma_client: Final = AsyncMock() + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[key_row]) + prisma_client.db.litellm_jwtkeymapping = _CascadingJWTMappingTable([]) + prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + prisma_client.delete_data = AsyncMock(return_value=[key_row.token]) + prisma_client.get_data = AsyncMock(return_value=None) + return prisma_client + + +@pytest.mark.asyncio +async def test_delete_verification_tokens_stops_the_key_on_a_peer_worker(monkeypatch: pytest.MonkeyPatch) -> None: + """Auth serves key objects cache-first, so a peer worker that is not told about the delete keeps + authenticating the deleted key from its own copy until user_api_key_cache_ttl expires""" + raw_key: Final = "sk-" + uuid4().hex + hashed_token: Final = hash_token(raw_key) + deleting_worker_cache: Final = UserApiKeyCache() + peer_worker_cache: Final = UserApiKeyCache() + for cache in (deleting_worker_cache, peer_worker_cache): + cache.set_cache(key=hashed_token, value=UserAPIKeyAuth(token=hashed_token, user_id="user-123")) + coordination_redis: Final = _SharedCoordinationRedis() + peer_subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=coordination_redis, user_api_key_cache=peer_worker_cache + ) + peer_subscriber.start() + prisma_client: Final = _prisma_with_key_row(_owned_key_row(hashed_token)) + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", coordination_redis) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + + try: + await delete_verification_tokens( + tokens=[raw_key], + user_api_key_cache=deleting_worker_cache, + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value + ), + ) + await coordination_redis.until_every_subscriber_applied_its_messages() + finally: + await peer_subscriber.stop() + + assert deleting_worker_cache.get_cache(key=hashed_token) is None + with pytest.raises(ProxyException) as denied: + await get_key_object( + hashed_token=hashed_token, prisma_client=prisma_client, user_api_key_cache=peer_worker_cache + ) + assert denied.value.code == "401" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deleter", + [ + UserAPIKeyAuth(user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value), + UserAPIKeyAuth(user_id="user-123", api_key="sk-owner", user_role=LitellmUserRoles.INTERNAL_USER.value), + ], + ids=["proxy_admin", "key_owner"], +) +async def test_delete_verification_tokens_evicts_a_requested_key_whose_row_is_already_gone( + monkeypatch: pytest.MonkeyPatch, deleter: UserAPIKeyAuth +) -> None: + """A key whose row was removed outside the API stays cached on every worker, so deleting it in a + batch beside a live key must still stop it everywhere instead of only evicting the rows found""" + live_key: Final = "sk-" + uuid4().hex + gone_key: Final = "sk-" + uuid4().hex + live_token: Final = hash_token(live_key) + gone_token: Final = hash_token(gone_key) + deleting_worker_cache: Final = UserApiKeyCache() + peer_worker_cache: Final = UserApiKeyCache() + for cache in (deleting_worker_cache, peer_worker_cache): + for token in (live_token, gone_token): + cache.set_cache(key=token, value=UserAPIKeyAuth(token=token, user_id="user-123")) + coordination_redis: Final = _SharedCoordinationRedis() + peer_subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=coordination_redis, user_api_key_cache=peer_worker_cache + ) + peer_subscriber.start() + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", coordination_redis) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_with_key_row(_owned_key_row(live_token))) + + try: + await delete_verification_tokens( + tokens=[live_key, gone_key], user_api_key_cache=deleting_worker_cache, user_api_key_dict=deleter + ) + await coordination_redis.until_every_subscriber_applied_its_messages() + finally: + await peer_subscriber.stop() + + assert [deleting_worker_cache.get_cache(key=token) for token in (live_token, gone_token)] == [None, None] + assert [peer_worker_cache.get_cache(key=token) for token in (live_token, gone_token)] == [None, None] + + +@pytest.mark.asyncio +async def test_delete_verification_tokens_leaves_unowned_cache_entries_alone_on_every_worker(monkeypatch: pytest.MonkeyPatch) -> None: + """/key/delete accepts any string, and a name that is not a key row passes through unhashed. A key + owner padding the request with another team's spend counter name must not wipe that counter from + every worker's memory: only the rows the caller was allowed to delete are evicted and broadcast""" + raw_key: Final = "sk-" + uuid4().hex + hashed_token: Final = hash_token(raw_key) + other_team_counter: Final = "spend:team:other-team" + deleting_worker_cache: Final = UserApiKeyCache() + peer_worker_cache: Final = UserApiKeyCache() + peer_spend_counters: Final = InMemoryCache() + peer_spend_counters.set_cache(other_team_counter, 42.0) + peer_worker_cache.set_cache(key=hashed_token, value=UserAPIKeyAuth(token=hashed_token, user_id="user-123")) + coordination_redis: Final = _SharedCoordinationRedis() + peer_subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=coordination_redis, + user_api_key_cache=peer_worker_cache, + additional_in_memory_caches=(peer_spend_counters,), + ) + peer_subscriber.start() + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", coordination_redis) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_with_key_row(_owned_key_row(hashed_token))) + + try: + await delete_verification_tokens( + tokens=[other_team_counter, raw_key], + user_api_key_cache=deleting_worker_cache, + user_api_key_dict=UserAPIKeyAuth( + user_id="user-123", api_key="sk-owner", user_role=LitellmUserRoles.INTERNAL_USER.value + ), + ) + await coordination_redis.until_every_subscriber_applied_its_messages() + finally: + await peer_subscriber.stop() + + assert peer_worker_cache.get_cache(key=hashed_token) is None + assert peer_spend_counters.get_cache(other_team_counter) == 42.0 + assert coordination_redis.published == (json.dumps({"cache_key": hashed_token}),) + + @pytest.mark.asyncio async def test_delete_key_fn_persists_deleted_keys(monkeypatch): from litellm.proxy._types import KeyRequest @@ -18750,7 +18981,7 @@ async def test_list_keys_rejects_invalid_expires(): expires="expred", ) - assert exc_info.value.code == "400" + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") assert "Invalid expires value" in str(exc_info.value.message) @@ -20746,12 +20977,20 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( async def _cache_with_project( - project_id: str, project_models: list[str], team_id: str | None = "team-lit-5823" + project_id: str, project_models: list[str], team_id: str | None = "team-lit-5823", + project_max_budget: float | None = None, ) -> UserApiKeyCache: user_api_key_cache = UserApiKeyCache() await user_api_key_cache.async_set_cache( key=project_cache_key(project_id), - value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id=team_id, models=project_models), + value=LiteLLM_ProjectTableCachedObj( + project_id=project_id, + team_id=team_id, + models=project_models, + litellm_budget_table=( + LiteLLM_BudgetTable(max_budget=project_max_budget) if project_max_budget is not None else None + ), + ), model_type=LiteLLM_ProjectTableCachedObj, ) return user_api_key_cache @@ -21184,6 +21423,487 @@ async def test_project_detachment_uses_effective_project_for_validation_on_unown assert expected in str(exc.value.detail) +def _project_assignment_prisma(stored_row: Mapping[str, object] | None = None) -> MagicMock: + database = MagicMock() + team_row: Final = LiteLLM_TeamTable(team_id="team-lit-5823", members=[]) + database.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + database.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + if stored_row is None: + return database + + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + + def _where_matches(where: Mapping[str, object], row: Mapping[str, object]) -> bool: + for key, clause in where.items(): + if key == "OR": + if not any(_where_matches(part, row) for part in clause): + return False + elif key == "models": + equals = clause["equals"] + if equals is None: + if row[key] is not None: + return False + elif row[key] is None or list(row[key]) != list(equals): + return False + elif row.get(key) != clause: + return False + return True + + async def _update_many(*, where, data): + if "object_permission" in data: + raise AssertionError("object_permission must not reach update_many") + return 1 if _where_matches(where, stored_row) else 0 + + tx.litellm_verificationtoken.update_many = AsyncMock(side_effect=_update_many) + written = MagicMock() + written.model_dump = MagicMock(return_value={**stored_row, "project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=written) + return database + + +@pytest.mark.asyncio +async def test_project_assignment_to_unassigned_key_on_same_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + data: Final = UpdateKeyRequest(key=existing.token, project_id="project-orbit") + + await _validate_update_key_data( + data, existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + + update: Final = await prepare_key_update_data(data=data, existing_key_row=existing) + assert update["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_project_on_another_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-other", models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_key_with_no_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id=None, models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_nonexistent_project(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _project_assignment_prisma() + database.db.litellm_projecttable.find_unique = AsyncMock(return_value=None) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-ghost"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, database, UserApiKeyCache(), + ) + assert exc.value.status_code == 404 + assert "Project not found" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_validates_existing_models_against_project(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-elsewhere"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "not in project's allowed models" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_with_explicit_null_max_budget_ignores_stored_budget(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=[], max_budget=100.0 + ) + cache: Final = await _cache_with_project("project-orbit", [], project_max_budget=10.0) + + request: Final = UpdateKeyRequest(key=existing.token, project_id="project-orbit", max_budget=None) + await _validate_update_key_data( + request, existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + + result: Final = await prepare_key_update_data(data=request, existing_key_row=existing) + assert result["max_budget"] is None + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_without_max_budget_uses_stored_budget(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=[], max_budget=100.0 + ) + cache: Final = await _cache_with_project("project-orbit", [], project_max_budget=10.0) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "exceeds project's max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_includes_new_project_assignment(): + existing: Final = LiteLLM_VerificationToken(token="project-assign-token", project_id=None) + + result: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing_key_row=existing, + ) + + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_requires_row_still_unassigned(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=1) + row = MagicMock() + row.model_dump = MagicMock(return_value={"project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["token"] == result["token"] + assert result["data"] == {"project_id": "project-orbit"} + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_row_gained_project(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=0) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + tx.litellm_verificationtoken.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_project_assignment_write_upserts_object_permission_in_same_tx(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + database.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=1) + permission_row = MagicMock() + permission_row.object_permission_id = "perm-lit-01" + tx.litellm_objectpermissiontable.upsert = AsyncMock(return_value=permission_row) + row = MagicMock() + row.model_dump = MagicMock(return_value={"project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823" + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit", "object_permission": {"agents": ["agent-a"]}}, + existing_key_row=existing, + ) + + tx.litellm_objectpermissiontable.upsert.assert_awaited_once() + update_data: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["data"] + assert update_data["object_permission_id"] == "perm-lit-01" + assert "object_permission" not in update_data + assert result["data"] == {"project_id": "project-orbit"} + + +@pytest.mark.asyncio +async def test_project_assignment_write_matches_null_models_column(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": None, + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=[] + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert result["data"]["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_strips_explicit_null_object_permission(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": None, + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=[] + ) + update_values: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit", object_permission=None), + existing_key_row=existing, + ) + assert "object_permission" in update_values + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values=update_values, + existing_key_row=existing, + ) + + assert result["data"]["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_max_budget_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": 50.0, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", + models=["model-orbit"], max_budget=100.0, + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_object_permission_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": "perm-new", + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_team_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-moved", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +def _assignment_project_update_prisma(stored_row: Mapping[str, object]) -> MagicMock: + database: Final = _project_assignment_prisma(stored_row=stored_row) + project_row: Final = LiteLLM_ProjectTable(project_id="project-orbit", team_id="team-lit-5823") + database.db.litellm_projecttable.find_unique = AsyncMock(return_value=project_row) + database.writer_db.litellm_projecttable.find_unique = AsyncMock(return_value=project_row) + return database + + +def _setup_assignment_update_mocks(monkeypatch: pytest.MonkeyPatch, database: MagicMock, cache: UserApiKeyCache) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", database) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + + +@pytest.mark.asyncio +async def test_update_key_fn_assigns_project_through_guarded_write(monkeypatch: pytest.MonkeyPatch): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _assignment_project_update_prisma(stored) + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + _setup_assignment_update_mocks(monkeypatch, database, cache) + + result: Final = await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many.assert_awaited_once() + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["project_id"] is None + tx.litellm_verificationtoken.update.assert_not_called() + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_update_key_fn_soft_budget_runs_guarded_assignment(monkeypatch: pytest.MonkeyPatch): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _assignment_project_update_prisma(stored) + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_budgettable.create = AsyncMock(return_value=SimpleNamespace(budget_id="budget-new")) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + _setup_assignment_update_mocks(monkeypatch, database, cache) + + result: Final = await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit", soft_budget=5.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + tx.litellm_budgettable.create.assert_awaited_once() + tx.litellm_verificationtoken.update_many.assert_awaited_once() + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["project_id"] is None + tx.litellm_verificationtoken.update.assert_not_called() + assert result["project_id"] == "project-orbit" + + @pytest.mark.asyncio async def test_key_creator_cannot_detach_project_without_admin_access(): existing: Final = LiteLLM_VerificationToken( @@ -21201,6 +21921,41 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) +@pytest.mark.asyncio +async def test_unauthorized_caller_gets_403_not_project_lookup_result(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", + user_id="user-orbit", created_by="user-orbit", + ) + database: Final = _project_assignment_prisma() + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-missing"), existing, + UserAPIKeyAuth(user_id="user-other", user_role=LitellmUserRoles.INTERNAL_USER), + None, False, database, UserApiKeyCache(), + ) + assert exc.value.status_code == 403 + assert "Project not found" not in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_unauthorized_caller_gets_403_not_project_team_leak(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", + user_id="user-orbit", created_by="user-orbit", + ) + cache: Final = await _cache_with_project("project-orbit", [], team_id="team-other") + database: Final = _project_assignment_prisma() + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_id="user-other", user_role=LitellmUserRoles.INTERNAL_USER), + None, False, database, cache, + ) + assert exc.value.status_code == 403 + assert "belongs to team" not in str(exc.value.detail) + + _OWNED_PROJECT: Final = "proj-owned-1" _OWNERSHIP_PROJECT_TEAM: Final = "ownership-project-team" _OWNERSHIP_KEY_TEAM: Final = "ownership-key-team" diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 872d1607351..ab6cfcfebfe 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -3,17 +3,20 @@ import contextlib import inspect import json import socket -from collections.abc import Iterator, Mapping -from types import SimpleNamespace -from typing import Dict, Final, Optional +from collections.abc import AsyncIterator, Iterator, Mapping +from datetime import datetime, time, timezone +from types import MappingProxyType, SimpleNamespace +from typing import Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException from fastapi.encoders import jsonable_encoder from fastapi.testclient import TestClient from litellm._uuid import uuid +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.models.credentials import CredentialItem from litellm.models.project import LiteLLM_ProjectTable @@ -44,7 +47,14 @@ from litellm.proxy.management_helpers.model_quota_policy import REQUIRE_BOTH_LIM from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.router_api import ApplyResult, ChangeSet -from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelAccessWindow, + ModelInfo, + updateDeployment, + updateLiteLLMParams, +) async def _passthrough_row(update_data): @@ -1639,11 +1649,350 @@ class TestDeletedAutoRouterAvailability: } +class _ContendedModelDb: + def __init__(self, row: LiteLLM_ProxyModelTable, other_writers: tuple[Mapping[str, object], ...]) -> None: + self.row = row + self.writes: tuple[Mapping[str, object], ...] = () + self._other_writers = other_writers + self._row_locked = False + self.db = self + self.litellm_proxymodeltable = self + self.litellm_teamtable = SimpleNamespace(find_unique=self._find_team) + + async def _find_team(self, *, where: Mapping[str, str]) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=where["team_id"]) + + async def find_unique(self, *, where: Mapping[str, object]) -> LiteLLM_ProxyModelTable: + read: Final = self.row + if not self._row_locked and self._other_writers: + self._commit_next_other_writer() + return read + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> LiteLLM_ProxyModelTable: + self.writes = (*self.writes, data) + self.row = self.row.model_copy( + update={k: json.loads(v) if k in ("litellm_params", "model_info") else v for k, v in data.items()} + ) + return self.row + + async def query_raw(self, query: str, *args: object) -> tuple[()]: + self._row_locked = self._row_locked or "FOR UPDATE" in query + return () + + @contextlib.asynccontextmanager + async def tx(self, timeout: object = None) -> AsyncIterator["_ContendedModelDb"]: + yield self + self._row_locked = False + while self._other_writers: + self._commit_next_other_writer() + + def _commit_next_other_writer(self) -> None: + self.row = self.row.model_copy( + update={"model_info": {**(self.row.model_info or {}), **self._other_writers[0]}} + ) + self._other_writers = self._other_writers[1:] + + +_AZURE_DEPLOYMENT_PARAMS: Final = MappingProxyType({"model": "azure/my-deployment", "api_key": "sk-existing"}) +_LOCAL_DEPLOYMENT_PARAMS: Final = MappingProxyType( + {"model": "hosted_vllm/local-model", "api_base": "https://local.test/v1", "api_key": "local-key"} +) + + +async def _discover_from_upstream(router: Router, model_card: Mapping[str, object]) -> None: + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"data": [dict(model_card)]})) + ) as client: + handler.client = client + await router.arefresh_model_info(client=handler) + + class TestUpdateModel: """ Tests for the update_model (POST /model/update) handler. """ + @staticmethod + async def _update_model( + stored_model_info: Mapping[str, object], + model_info: ModelInfo, + litellm_params: Mapping[str, object] = _AZURE_DEPLOYMENT_PARAMS, + upstream_model_card: Mapping[str, object] | None = None, + other_writers: tuple[Mapping[str, object], ...] = (), + ) -> "_ContendedModelDb": + db: Final = _ContendedModelDb( + LiteLLM_ProxyModelTable( + model_id="m1", + model_name="my-model", + litellm_params=dict(litellm_params), + model_info=dict(stored_model_info), + ), + other_writers, + ) + router: Final = Router( + model_list=[ + {"model_name": "my-model", "litellm_params": dict(litellm_params), "model_info": {"id": "m1"}} + ] + ) + if upstream_model_card is not None: + await _discover_from_upstream(router, upstream_model_card) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", db), + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( # test-quality-ok: [TQ008] the route reloads the proxy-global router with no seam to inject + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment(litellm_params=updateLiteLLMParams(), model_info=model_info), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + return db + + @pytest.mark.asyncio + async def test_update_model_persists_sent_model_info_over_stored_fields(self): + stored: Final = { + "id": "m1", + "db_model": True, + "base_model": "my-deployment", + "tier": "paid", + "team_id": "team-a", + } + + db: Final = await self._update_model( + stored, + ModelInfo(id="m1", base_model="test-provider/test-model", mode="chat", db_model=False, team_id="team-b"), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["mode"] == "chat" + assert {k: persisted[k] for k in ("db_model", "tier", "team_id")} == { + "db_model": True, + "tier": "paid", + "team_id": "team-a", + } + + @pytest.mark.asyncio + async def test_update_model_with_only_an_id_leaves_model_info_untouched(self): + db: Final = await self._update_model({"id": "m1", "base_model": "my-deployment"}, ModelInfo(id="m1")) + + assert tuple("model_info" in write for write in db.writes) == (False,) + + @pytest.mark.asyncio + async def test_update_model_resending_the_stored_model_info_does_not_rewrite_it(self): + db: Final = await self._update_model( + {"id": "m1", "db_model": True, "tier": "paid", "mode": "chat"}, + ModelInfo(id="m1", tier="paid", mode="chat"), + ) + + assert tuple("model_info" in write for write in db.writes) == (False,) + + @pytest.mark.asyncio + async def test_update_model_keeps_model_info_fields_a_concurrent_writer_stored(self): + db: Final = await self._update_model( + {"id": "m1", "db_model": True, "mode": "embedding"}, + ModelInfo(id="m1", mode="chat"), + other_writers=(MappingProxyType({"access_groups": ["group-a"]}), MappingProxyType({"tier": "paid"})), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in ("mode", "access_groups", "tier")} == { + "mode": "chat", + "access_groups": ["group-a"], + "tier": "paid", + } + + @pytest.mark.asyncio + async def test_update_model_does_not_pin_values_echoed_from_model_info(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + echoed: Final = ModelInfo( + id="m1", + base_model="azure/gpt-4o", + key=catalog_entry["key"], + max_tokens=catalog_entry["max_tokens"], + input_cost_per_token=catalog_entry["input_cost_per_token"], + ) + + db: Final = await self._update_model({"id": "m1"}, echoed) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "azure/gpt-4o" + assert not {"key", "max_tokens", "input_cost_per_token"} & persisted.keys() + + @pytest.mark.asyncio + async def test_update_model_saves_typed_values_that_equal_the_catalog(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + typed: Final = { + "max_tokens": catalog_entry["max_tokens"], + "max_input_tokens": catalog_entry["max_input_tokens"], + "supports_vision": catalog_entry["supports_vision"], + } + + db: Final = await self._update_model( + {"id": "m1", "max_input_tokens": catalog_entry["max_input_tokens"] // 2}, + ModelInfo(id="m1", base_model="azure/gpt-4o", **typed), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in typed} == typed + + @pytest.mark.asyncio + async def test_update_model_keeps_stored_value_when_a_field_is_sent_blank(self): + db: Final = await self._update_model( + {"id": "m1", "base_model": "my-deployment", "mode": "chat"}, + ModelInfo(id="m1", base_model="test-provider/test-model", mode=""), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["mode"] == "chat" + + @pytest.mark.asyncio + async def test_update_model_ignores_ptu_blocked_and_pricing_fields(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo( + id="m1", base_model="test-provider/test-model", ptu_count=5, blocked=True, input_cost_per_token=0.5 + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert not {"ptu_count", "blocked", "input_cost_per_token"} & persisted.keys() + + @pytest.mark.asyncio + @pytest.mark.parametrize("access_groups", [5, "group-a", ["group-a", 5], {"group": "group-a"}]) + async def test_update_model_rejects_access_groups_that_are_not_a_list_of_strings(self, access_groups: object): + with pytest.raises(ProxyException) as error: + await self._update_model({"id": "m1"}, ModelInfo.model_validate({"id": "m1", "access_groups": access_groups})) + + assert error.value.code == "400" + + @pytest.mark.asyncio + async def test_update_model_saves_access_groups_sent_as_a_list_of_strings(self): + db: Final = await self._update_model( + {"id": "m1"}, ModelInfo.model_validate({"id": "m1", "access_groups": ["group-a", "group-b"]}) + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["access_groups"] == ["group-a", "group-b"] + + def test_update_model_does_not_pin_the_previous_base_models_echoed_values(self): + from litellm.proxy.management_endpoints.model_management_endpoints import legacy_update_model_info_changes + + catalog: Final = MappingProxyType( + {"test-provider/previous": MappingProxyType({"supported_openai_params": ["previous-only-param"]})} + ) + stored: Final = Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params(model="test-provider/my-deployment"), + model_info=ModelInfo(id="m1", base_model="test-provider/previous"), + ) + + changes: Final = legacy_update_model_info_changes( + stored, + ModelInfo( + id="m1", + base_model="test-provider/next", + key="test-provider/previous", + supported_openai_params=["previous-only-param"], + ), + loaded_catalog=lambda: catalog, + ) + + assert dict(changes) == {"base_model": "test-provider/next"} + + @pytest.mark.asyncio + async def test_update_model_saves_a_value_the_previous_base_model_has_in_the_catalog(self): + import litellm + + previous: Final = litellm.get_model_info(model="azure/gpt-4o") + + db: Final = await self._update_model( + {"id": "m1", "base_model": "azure/gpt-4o", "max_input_tokens": previous["max_input_tokens"] // 2}, + ModelInfo(id="m1", base_model="azure/gpt-4.1", max_input_tokens=previous["max_input_tokens"]), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in ("base_model", "max_input_tokens")} == { + "base_model": "azure/gpt-4.1", + "max_input_tokens": previous["max_input_tokens"], + } + + @pytest.mark.asyncio + async def test_update_model_does_not_pin_discovered_limits_or_team_access_echoed_from_model_info(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo( + id="m1", + base_model="test-provider/test-model", + key="test-provider/test-model", + max_input_tokens=128000, + max_output_tokens=64000, + direct_access=True, + access_via_team_ids=["team-a"], + ), + litellm_params=_LOCAL_DEPLOYMENT_PARAMS, + upstream_model_card=MappingProxyType( + {"id": "local-model", "max_model_len": 128000, "max_output_tokens": 32000} + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["max_output_tokens"] == 64000 + assert not {"max_input_tokens", "direct_access", "access_via_team_ids"} & persisted.keys() + + @pytest.mark.asyncio + async def test_update_model_does_not_copy_litellm_params_pricing_into_model_info(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo(id="m1", tier="paid"), + litellm_params={**_AZURE_DEPLOYMENT_PARAMS, "input_cost_per_token": 0.000003}, + ) + + persisted: Final = db.row.model_info + assert persisted == {"id": "m1", "tier": "paid"} + + @pytest.mark.asyncio + async def test_update_model_resets_a_stored_override_echoed_back_at_its_catalog_value(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + stored_override: Final = catalog_entry["max_tokens"] // 2 + + db: Final = await self._update_model( + {"id": "m1", "base_model": "azure/gpt-4o", "max_tokens": stored_override}, + ModelInfo( + id="m1", base_model="azure/gpt-4o", key=catalog_entry["key"], max_tokens=catalog_entry["max_tokens"] + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert "max_tokens" not in persisted + @pytest.mark.asyncio async def test_update_model_clears_cache_after_db_write(self): """ @@ -3707,6 +4056,45 @@ class TestUpdateDBModelBlocked: assert "blocked" not in result +@pytest.mark.parametrize("window_patch", ["omit", "echo", "replace", "clear"]) +def test_model_edit_preserves_or_explicitly_changes_access_windows( + window_patch: Literal["omit", "echo", "replace", "clear"], +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + saved_window: Final = ModelAccessWindow( + start=time(13), end=time(17), timezone="Asia/Almaty", team_ids=("team-window", "team-shared") + ) + replacement_window: Final = ModelAccessWindow( + start=time(22, 0, 0, 123456), end=time(6, 0, 0, 654321), + timezone="America/New_York", team_ids=("team-other",), + ) + created_at: Final = datetime(2026, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc) + requested_windows: Final = { + "omit": None, + "echo": (saved_window,), + "replace": (replacement_window,), + "clear": (), + }[window_patch] + expected_windows: Final = (saved_window,) if requested_windows is None else requested_windows + result: Final = update_db_model( + db_model=Deployment( + model_name="window-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="window-deployment", access_windows=(saved_window,), created_at=created_at), + ), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(max_parallel_requests=7), + model_info=None if requested_windows is None else ModelInfo(access_windows=requested_windows), + ), + ) + + assert json.loads(result["litellm_params"])["max_parallel_requests"] == 7 + saved_info: Final = ModelInfo.model_validate_json(result["model_info"]) + assert saved_info.access_windows == expected_windows + assert saved_info.created_at == created_at + + class TestUpdateDBModelKeepsLegacyDropParams: def test_partial_patch_keeps_encrypted_string_drop_params(self, monkeypatch): from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 136cc86e67a..ba78fefd653 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -15,9 +15,15 @@ from fastapi import HTTPException, Request from pydantic import SecretStr import litellm +from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_UserTable, + NewUserResponse, + ProxyException, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.types import CustomOpenID @@ -287,6 +293,13 @@ def test_get_microsoft_callback_response(stubbed_graph_api): assert sso_assertion.issuer == "https://idp.example.com" + +def _jwt_handler_double() -> MagicMock: + handler: Final = MagicMock(spec=JWTHandler) + handler.litellm_jwtauth = LiteLLM_JWTAuth() + return handler + + GRAPH_ERROR_BODIES: Final = ( {"error": {"code": "Authentication_MissingOrMalformed", "message": "Access Token missing or malformed."}}, {"error": {"code": "InvalidAuthenticationToken", "message": "IDX14100: JWT is not well formed."}}, @@ -1158,7 +1171,7 @@ def test_generic_response_convertor_normalizes_email(): } # Mock JWT handler - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -2226,7 +2239,7 @@ async def test_get_generic_sso_response_with_additional_headers(): # Arrange mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] generic_client_id = "test_client_id" @@ -2299,7 +2312,7 @@ async def test_get_generic_sso_response_with_empty_headers(): # Arrange mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] generic_client_id = "test_client_id" @@ -2359,8 +2372,8 @@ async def test_get_generic_sso_response_includes_token_claims_when_enabled(monke from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) - mock_sso_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() + mock_sso_jwt_handler = _jwt_handler_double() mock_sso_jwt_handler.get_all_jwt_team_ids.return_value = ["team-from-userinfo"] mock_sso_jwt_handler.get_team_ids_from_jwt.return_value = [] @@ -2435,8 +2448,8 @@ async def test_get_generic_sso_response_does_not_include_token_claims_when_disab from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) - mock_sso_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() + mock_sso_jwt_handler = _jwt_handler_double() mock_sso_jwt_handler.get_all_jwt_team_ids.return_value = ["team-from-userinfo"] mock_sso_jwt_handler.get_team_ids_from_jwt.return_value = [] access_token = pyjwt.encode({"upn": "token-user@example.com"}, "test-secret", algorithm="HS256") @@ -2578,7 +2591,7 @@ async def test_get_generic_sso_response_pkce_merges_token_claims_and_excludes_cr ): result, received_response, _, _ = await get_generic_sso_response( request=request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="test-client", redirect_url="http://test.com/callback", sso_jwt_handler=None, @@ -2600,7 +2613,7 @@ async def test_get_generic_sso_response_ignores_opaque_and_empty_token_claims(mo from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() userinfo = { "preferred_username": "userinfo-user", "email": "userinfo@example.com", @@ -4961,7 +4974,7 @@ class TestProcessSSOJWTAccessToken: @pytest.fixture def mock_jwt_handler(self): """Create a mock JWT handler for testing""" - mock_handler = MagicMock(spec=JWTHandler) + mock_handler = _jwt_handler_double() mock_handler.get_team_ids_from_jwt.return_value = ["team1", "team2", "team3"] return mock_handler @@ -5348,7 +5361,7 @@ class TestGenericResponseConvertorNestedAttributes: from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor # Mock JWT handler - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] # Payload with nested attributes structure @@ -5419,7 +5432,7 @@ class TestGenericResponseConvertorUserRole: from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] sso_response = { @@ -5449,7 +5462,7 @@ class TestGenericResponseConvertorUserRole: """ from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] sso_response = { @@ -6813,7 +6826,9 @@ class TestAddMissingTeamMember: # Verify add_missing_team_member was called with correct args mock_add_member.assert_called_once_with( - user_info=new_user_info, sso_teams=["entra-group-1", "entra-group-2"] + user_info=new_user_info, + sso_teams=["entra-group-1", "entra-group-2"], + team_id_to_upsert=None, ) @pytest.mark.asyncio @@ -7897,7 +7912,7 @@ def test_generic_response_convertor_with_extra_attributes(monkeypatch): "custom_field3": {"nested": "data"}, } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7931,7 +7946,7 @@ def test_generic_response_convertor_without_extra_attributes(monkeypatch): "custom_field2": "value2", } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7959,7 +7974,7 @@ def test_generic_response_convertor_extra_attributes_with_nested_paths(monkeypat "org_info": {"department": "Engineering", "manager": "Jane Smith"}, } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7986,7 +8001,7 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch): "email": "user@example.com", } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -8865,7 +8880,7 @@ class TestPKCEStateCookieBinding: ): await get_generic_sso_response( request=mock_request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="cid", redirect_url="https://proxy.example.com/sso/callback", sso_jwt_handler=None, @@ -8917,7 +8932,7 @@ class TestPKCEStateCookieBinding: ): await get_generic_sso_response( request=mock_request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="cid", redirect_url="https://proxy.example.com/sso/callback", sso_jwt_handler=None, @@ -8980,7 +8995,7 @@ class TestPKCEStateCookieBinding: }, ), ): - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] result, _, _, _ = await get_generic_sso_response( request=mock_request, @@ -9087,7 +9102,7 @@ async def test_debug_sso_callback_renders_full_jwt_claims(): patch("litellm.proxy.proxy_server.general_settings", {}), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), ): # Microsoft / Google envs may leak in from other tests — ensure only # the generic path runs. @@ -9173,7 +9188,7 @@ async def test_debug_sso_callback_handles_missing_raw_response(): patch("litellm.proxy.proxy_server.general_settings", {}), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), ): for var in ("GENERIC_CLIENT_ID", "GOOGLE_CLIENT_ID"): os.environ.pop(var, None) @@ -9224,7 +9239,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: debug endpoint reads proxy globals patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: debug endpoint reads proxy DB patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: debug endpoint reads proxy globals - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), # test-quality-ok: debug endpoint reads proxy globals ] if force_inert: stack.append( @@ -9589,7 +9604,7 @@ async def test_pkce_arm_captures_sso_assertion(): }, ), ): - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] result, received_response, _, sso_assertion = await get_generic_sso_response( request=mock_request, @@ -9615,7 +9630,7 @@ async def test_verify_and_process_arm_captures_sso_assertion(): id_token = _ema_id_token() mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] mock_sso_instance = MagicMock() @@ -9680,7 +9695,7 @@ def _generic_sso_callback_request() -> Request: def _generic_sso_jwt_handler() -> MagicMock: - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] return jwt_handler @@ -11144,3 +11159,241 @@ async def test_cli_sso_login_never_serves_the_admin_login_form_on_a_ui_disabled_ body = response.body.decode() assert "Admin UI is Disabled" in body assert 'name="username"' not in body + + +def _upsert_jwt_handler(**jwtauth: object) -> JWTHandler: + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=LiteLLM_JWTAuth(**jwtauth) + ) + return handler + + +@pytest.mark.parametrize( + ("jwtauth", "userinfo", "expected_team_ids", "expected_upsert"), + [ + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": "finance"}, ["finance"], "finance"), + ({"team_id_jwt_field": "dept", "team_id_upsert": False}, {"dept": "finance"}, ["finance"], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {}, [], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ""}, [], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": " "}, [" "], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ["a", "b"]}, ["a", "b"], None), + ( + {"team_id_jwt_field": "dept", "team_id_upsert": True, "team_id_default": "everyone"}, + {}, + [], + None, + ), + ( + {"team_ids_jwt_field": "groups", "team_id_upsert": True}, + {"groups": ["g1", "g2"]}, + ["g1", "g2"], + None, + ), + ( + {"team_ids_jwt_field": "groups", "team_id_jwt_field": "dept", "team_id_upsert": True}, + {"groups": ["g1"], "dept": "finance"}, + ["g1", "finance"], + "finance", + ), + ], + ids=[ + "upsert-on", + "upsert-off", + "claim-missing", + "blank-claim-not-upserted", + "whitespace-claim-not-upserted", + "list-valued-claim-join-only", + "default-team-not-upserted", + "plural-claim-join-only", + "mixed", + ], +) +def test_generic_sso_login_upserts_only_the_team_id_jwt_field_claim( + jwtauth: dict[str, object], + userinfo: dict[str, object], + expected_team_ids: list[str], + expected_upsert: str | None, +) -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", **userinfo}, + jwt_handler=_upsert_jwt_handler(**jwtauth), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (expected_team_ids, expected_upsert) + + +def test_generic_sso_login_does_not_upsert_the_team_id_jwt_field_claim_when_jwt_auth_is_disabled() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "dept": "finance"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + jwt_auth_enabled=False, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (["finance"], None) + + +def test_generic_sso_login_marks_the_team_id_jwt_field_team_for_upsert_without_adding_it_to_team_ids() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "groups": ["g1"], "dept": "finance"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (["g1"], "finance") + + +def test_generic_sso_login_keeps_team_ids_empty_for_a_dept_only_user_so_the_access_token_group_fallback_runs() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "dept": "d1"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == ([], "d1") + + +def test_restricted_sso_group_still_rejects_a_user_whose_only_team_is_the_upserted_dept() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler, generic_response_convertor + + response: Final = {"preferred_username": "alice", "email": "alice@example.com", "dept": "X", "groups": []} + result: Final = generic_response_convertor( + response=response, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + with pytest.raises(ProxyException) as denied: + SSOAuthenticationHandler.verify_user_in_restricted_sso_group( + general_settings={"ui_access_mode": {"type": "restricted_sso_group", "restricted_sso_group": "X"}}, + result=result, + received_response=response, + ) + + assert denied.value.code == "403" + + +def _recording_team_calls( + events: list[tuple[str, str]], +) -> tuple[Callable[[str, NewUserResponse | LiteLLM_UserTable], Awaitable[None]], Callable[[str], Awaitable[None]]]: + async def add(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable) -> None: + events.append(("join", team_id)) + + async def create(team_id: str) -> None: + events.append(("create", team_id)) + + return add, create + + +async def _only_group_teams_exist(team_ids: Sequence[str]) -> frozenset[str] | None: + return frozenset(team_id for team_id in team_ids if team_id.startswith("g")) + + +@pytest.mark.asyncio +async def test_sso_login_joins_the_upsert_team_from_team_id_to_upsert_and_creates_it_before_the_join() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["g1", "g2"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + team_id_to_upsert="finance", + create_team=create, + ) + + assert events == [("join", "g1"), ("join", "g2"), ("create", "finance"), ("join", "finance")] + + +@pytest.mark.asyncio +async def test_sso_login_joins_only_the_claimed_groups_when_team_id_upsert_is_off() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["g1", "g2"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + create_team=create, + ) + + assert events == [("join", "g1"), ("join", "g2")] + + +@pytest.mark.asyncio +async def test_sso_login_does_not_create_the_team_id_jwt_field_team_when_it_already_exists() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + async def every_team_exists(team_ids: Sequence[str]) -> frozenset[str] | None: + return frozenset(team_ids) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["finance"], + add_team_member=add, + existing_team_ids=every_team_exists, + team_id_to_upsert="finance", + create_team=create, + ) + + assert events == [("join", "finance")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "create_error", + [ + HTTPException(status_code=400, detail={"error": "Team id = finance already exists"}), + HTTPException(status_code=403, detail={"error": "team limit reached for this license"}), + ], + ids=["lost-the-create-race", "create-rejected"], +) +async def test_sso_login_still_joins_the_team_and_warns_when_creating_it_fails( + caplog: pytest.LogCaptureFixture, create_error: HTTPException +) -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, _ = _recording_team_calls(events) + + async def failed_create(team_id: str) -> None: + raise create_error + + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["finance"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + team_id_to_upsert="finance", + create_team=failed_create, + ) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert events == [("join", "finance")] + assert len(warnings) == 1 and "finance" in warnings[0] and str(create_error.detail) in warnings[0] diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index ee3ef3f5323..1b95cbf3c94 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -21,6 +21,7 @@ from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable, + LitellmUserRoles, ProxyException, SpecialHeaders, SpecialMCPServerNames, @@ -13196,3 +13197,58 @@ async def test_personal_key_mcp_access_is_denied_when_owner_teams_cannot_be_read assert await MCPRequestHandler.get_allowed_mcp_servers(personal_key_owner_teams) == [] assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", personal_key_owner_teams) == [] assert await MCPRequestHandler.get_allowed_tools_for_server("srv-open", personal_key_owner_teams) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +@pytest.mark.parametrize( + "owner_role", + [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], + ids=["proxy_admin", "proxy_admin_viewer"], +) +@pytest.mark.parametrize( + ("key_team_id", "expected_servers"), + [("team-a", {"srv-a", "srv-open"}), ("team-without-mcp", {"srv-open"})], + ids=["team_grants_one_server", "team_grants_no_server"], +) +@pytest.mark.parametrize( + "credential_shape", + [{"is_session_token": False}, {"is_session_token": True}, {"api_key": None, "jwt_claims": {"sub": "admin"}}], + ids=["virtual_key", "cli_session", "jwt"], +) +async def test_admin_owned_team_key_reaches_only_its_team_mcp_servers( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, + owner_role: LitellmUserRoles, + key_team_id: str, + expected_servers: set[str], + credential_shape: dict[str, object], +) -> None: + from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return _personal_key_team(team_id) if team_id in _PERSONAL_KEY_TEAMS else LiteLLM_TeamTable(team_id=team_id) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + team_key: Final = personal_key_owner_teams.model_copy( + update={"team_id": key_team_id, "user_role": owner_role, **credential_shape} + ) + + assert set(await global_mcp_server_manager.get_allowed_mcp_servers(team_key)) == expected_servers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_admin_dashboard_session_keeps_the_whole_mcp_registry( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager + + dashboard_session: Final = personal_key_owner_teams.model_copy( + update={"team_id": UI_TEAM_ID, "user_role": LitellmUserRoles.PROXY_ADMIN} + ) + + assert set(await global_mcp_server_manager.get_allowed_mcp_servers(dashboard_session)) == set( + global_mcp_server_manager.get_registry() + ) diff --git a/tests/unit/proxy/mcp_server/test_mcp_logging.py b/tests/unit/proxy/mcp_server/test_mcp_logging.py index 9548135b530..b12cb916ecb 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_logging.py +++ b/tests/unit/proxy/mcp_server/test_mcp_logging.py @@ -1,27 +1,307 @@ import asyncio +import json import os import subprocess import sys +from collections.abc import Generator +from datetime import datetime from pathlib import Path -from unittest.mock import AsyncMock, patch +from queue import SimpleQueue +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from mcp.types import CallToolResult, TextContent +from fastapi import HTTPException +from mcp import Client +from mcp.server.auth.middleware.auth_context import auth_context_var +from mcp.server.context import ServerRequestContext +from mcp.server.session import ServerSession +from mcp.shared.exceptions import MCPError +from mcp.types import CallToolResult, EmptyResult, InitializeRequestParams, Request, RequestParams, TextContent from mcp.types import Tool as MCPTool +from pydantic import ConfigDict, ValidationError import litellm +from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth +from litellm.proxy._types import AlertType, LiteLLM_ObjectPermissionTable, UserAPIKeyAuth from litellm.proxy.mcp_server.mcp_server_manager import ( MCPServerManager, ) from litellm.proxy.mcp_server.server import ( + _log_sdk_rejection, mcp_server_tool_call, + server, set_auth_context, ) from litellm.types.mcp import MCPPostCallResponseObject +class SDKRejectionParams(RequestParams): + model_config = ConfigDict(frozen=True) + name: int | None = None + arguments: dict[str, str] | None = None + level: int | str | None = None + + +class SDKRejectionLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: list[dict[str, object]] = [] + self.pre_calls = 0 + self.post_failures = 0 + self.fail_post_hook = False + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str + ) -> dict[str, object]: + self.pre_calls += 1 + return data + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failures.append(kwargs) + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.post_failures += 1 + if self.fail_post_hook: + raise RuntimeError("callback unavailable") + + +@pytest.fixture +def sdk_rejection_logger() -> Generator[SDKRejectionLogger, None, None]: + logger: Final = SDKRejectionLogger() + token: Final = auth_context_var.set(None) + set_auth_context( + user_api_key_auth=UserAPIKeyAuth( + api_key="sk-sdk-rejection-unit-test", user_id="sdk-rejection-user", team_id="sdk-rejection-team" + ) + ) + try: + with patch("litellm.callbacks", [logger]): + yield logger + finally: + auth_context_var.reset(token) + litellm.logging_callback_manager._reset_all_callbacks() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method,params,code,message", + [ + ("tools/call-2", {"arguments": {"secret": "DO_NOT_LOG_THIS"}}, -32601, "Method not found"), + ("tools/call", {"name": 123, "arguments": {"secret": "DO_NOT_LOG_THIS"}}, -32602, "Invalid request parameters"), + ("tools/call", {"arguments": {"secret": "DO_NOT_LOG_THIS"}}, -32602, "Invalid request parameters"), + ("logging/setLevel", {"level": 123}, -32602, "Invalid request parameters"), + ("logging/setLevel", {"level": "info"}, -32601, "Method not found"), + ("completion/complete", {}, -32602, "Invalid request parameters"), + ("server/discover", {}, -32601, "Method not found"), + ("subscriptions/listen", {}, -32601, "Method not found"), + ], +) +async def test_sdk_rejection_emits_one_redacted_zero_cost_failure( + sdk_rejection_logger: SDKRejectionLogger, method: str, params: dict[str, object], code: int, message: str +) -> None: + async with Client(server, mode="legacy") as client: + with pytest.raises(MCPError) as rejected: + await client.session.send_request( + Request(method=method, params=SDKRejectionParams.model_validate(params)), EmptyResult + ) + assert rejected.value.code == code + assert rejected.value.message == message + await client.session.send_ping() + + assert sdk_rejection_logger.pre_calls == 0 + assert sdk_rejection_logger.post_failures == 1 + assert len(sdk_rejection_logger.failures) == 1 + payload: Final = sdk_rejection_logger.failures[0]["standard_logging_object"] + assert payload["status"] == "failure" + assert payload["response_cost"] == 0 + assert payload["prompt_tokens"] == 0 + assert payload["completion_tokens"] == 0 + assert payload["model"] == "MCP: request_rejected" + assert payload["metadata"]["user_api_key_user_id"] == "sdk-rejection-user" + assert payload["metadata"]["user_api_key_team_id"] == "sdk-rejection-team" + rejection: Final = payload["metadata"]["spend_logs_metadata"]["mcp_sdk_rejection"] + assert rejection["method"] == method + assert rejection["request_id"] + assert rejection["error_code"] == code + assert rejection["error_message"] == message + assert rejection["stage"] == "sdk_dispatch" + assert "DO_NOT_LOG_THIS" not in json.dumps(payload, default=str) + assert "sk-sdk-rejection-unit-test" not in json.dumps(payload, default=str) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method,params,code", + [ + ("tools/call-2", {}, -32601), + ("tools/call", {"name": 123}, -32602), + ], +) +async def test_sdk_caller_errors_do_not_trigger_llm_alerts( + sdk_rejection_logger: SDKRejectionLogger, method: str, params: dict[str, object], code: int +) -> None: + from litellm.proxy.proxy_server import proxy_logging_obj + + with ( + patch.object(proxy_logging_obj, "alert_types", [AlertType.llm_exceptions]), + patch.object(proxy_logging_obj, "alerting_handler", new_callable=AsyncMock) as alerts, + ): + async with Client(server, mode="legacy") as client: + with pytest.raises(MCPError) as rejected: + await client.session.send_request( + Request(method=method, params=SDKRejectionParams.model_validate(params)), EmptyResult + ) + assert rejected.value.code == code + await asyncio.sleep(0) + alerts.assert_not_awaited() + assert len(sdk_rejection_logger.failures) == 1 + assert sdk_rejection_logger.post_failures == 1 + assert sdk_rejection_logger.pre_calls == 0 + + await proxy_logging_obj.post_call_failure_hook( + request_data={}, + original_exception=HTTPException(status_code=500, detail="upstream unavailable"), + user_api_key_dict=UserAPIKeyAuth(), + ) + await asyncio.sleep(0) + alerts.assert_awaited_once() + assert alerts.await_args.kwargs["level"] == "High" + assert alerts.await_args.kwargs["alert_type"] == AlertType.llm_exceptions + + +@pytest.mark.asyncio +async def test_sdk_rejection_sync_logging_allows_concurrent_ping( + sdk_rejection_logger: SDKRejectionLogger, +) -> None: + loop: Final = asyncio.get_running_loop() + ping_responses: Final = SimpleQueue[EmptyResult]() + async with Client(server, mode="legacy") as client: + def sync_failure( + kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + ping_responses.put(asyncio.run_coroutine_threadsafe(client.session.send_ping(), loop).result(timeout=10)) + + with patch("litellm.failure_callback", [sync_failure]): + with pytest.raises(MCPError) as rejected: + await client.session.send_request(Request(method="unknown/method", params=RequestParams()), EmptyResult) + assert rejected.value.code == -32601 + + assert ping_responses.qsize() == 1 + assert isinstance(ping_responses.get_nowait(), EmptyResult) + assert len(sdk_rejection_logger.failures) == 1 + assert sdk_rejection_logger.post_failures == 1 + assert sdk_rejection_logger.pre_calls == 0 + + +@pytest.mark.asyncio +async def test_handler_method_not_found_does_not_emit_sdk_rejection( + sdk_rejection_logger: SDKRejectionLogger, +) -> None: + context: Final = ServerRequestContext( + session=MagicMock(spec=ServerSession), + lifespan_context=None, + protocol_version="2025-11-25", + method="tools/list", + params={}, + request_id="handler-error", + ) + handler_error: Final = MCPError(code=-32601, message="Handler rejected the request") + with pytest.raises(MCPError) as rejected: + await _log_sdk_rejection(context, AsyncMock(side_effect=handler_error)) + assert rejected.value is handler_error + assert sdk_rejection_logger.failures == [] + assert sdk_rejection_logger.post_failures == 0 + + +@pytest.mark.asyncio +async def test_sdk_initialization_errors_are_not_logged_as_tool_rejections( + sdk_rejection_logger: SDKRejectionLogger, +) -> None: + context: Final = ServerRequestContext( + session=MagicMock(spec=ServerSession), + lifespan_context=None, + protocol_version="2025-11-25", + method="initialize", + params={}, + request_id="bad-init", + ) + with pytest.raises(ValidationError) as sdk_error: + InitializeRequestParams.model_validate({}) + with pytest.raises(ValidationError) as rejected: + await _log_sdk_rejection(context, AsyncMock(side_effect=sdk_error.value)) + assert rejected.value is sdk_error.value + assert sdk_rejection_logger.pre_calls == 0 + assert sdk_rejection_logger.failures == [] + assert sdk_rejection_logger.post_failures == 0 + + +@pytest.mark.asyncio +async def test_sdk_response_survives_failure_callback_error(sdk_rejection_logger: SDKRejectionLogger) -> None: + sdk_rejection_logger.fail_post_hook = True + async with Client(server, mode="legacy") as client: + with pytest.raises(MCPError) as rejected: + await client.session.send_request(Request(method="tools/call-2", params=RequestParams()), EmptyResult) + assert rejected.value.code == -32601 + assert rejected.value.message == "Method not found" + assert sdk_rejection_logger.post_failures == 1 + assert len(sdk_rejection_logger.failures) == 1 + + +@pytest.mark.asyncio +async def test_sdk_rejection_without_admitted_identity_is_not_attributed( + sdk_rejection_logger: SDKRejectionLogger, +) -> None: + token: Final = auth_context_var.set(None) + try: + async with Client(server, mode="legacy") as client: + with pytest.raises(MCPError) as rejected: + await client.session.send_request(Request(method="tools/call-2", params=RequestParams()), EmptyResult) + assert rejected.value.code == -32601 + finally: + auth_context_var.reset(token) + assert sdk_rejection_logger.failures == [] + assert sdk_rejection_logger.post_failures == 0 + + +@pytest.mark.asyncio +async def test_sdk_rejection_keeps_concurrent_caller_identity(sdk_rejection_logger: SDKRejectionLogger) -> None: + async def reject_as(user: str) -> None: + token: Final = auth_context_var.set(None) + set_auth_context(user_api_key_auth=UserAPIKeyAuth(user_id=user, team_id=f"{user}-team")) + try: + async with Client(server, mode="legacy") as client: + with pytest.raises(MCPError) as rejected: + await client.session.send_request( + Request(method=f"unknown/{user}", params=RequestParams()), EmptyResult + ) + assert rejected.value.code == -32601 + finally: + auth_context_var.reset(token) + + await asyncio.gather(reject_as("caller-a"), reject_as("caller-b")) + assert len(sdk_rejection_logger.failures) == 2 + assert sdk_rejection_logger.post_failures == 2 + assert { + failure["standard_logging_object"]["metadata"]["user_api_key_user_id"] + for failure in sdk_rejection_logger.failures + } == {"caller-a", "caller-b"} + for failure in sdk_rejection_logger.failures: + metadata: Final = failure["standard_logging_object"]["metadata"] + identity: Final = metadata["user_api_key_user_id"] + assert metadata["user_api_key_team_id"] == f"{identity}-team" + assert metadata["spend_logs_metadata"]["mcp_sdk_rejection"]["method"] == f"unknown/{identity}" + + def _mcp_request_ctx(**overrides): from types import SimpleNamespace diff --git a/tests/unit/proxy/mcp_server/test_mcp_toolset_scope.py b/tests/unit/proxy/mcp_server/test_mcp_toolset_scope.py index f3cc38eacf9..d1337138cbc 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_toolset_scope.py +++ b/tests/unit/proxy/mcp_server/test_mcp_toolset_scope.py @@ -2,7 +2,7 @@ import asyncio from collections.abc import Awaitable, Callable -from typing import Dict, List, Optional +from typing import Dict, Final, List, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -112,6 +112,23 @@ class TestApplyToolsetScope: assert op.mcp_servers == ["server-a"] assert op.mcp_tool_permissions == toolset_perms + @pytest.mark.asyncio + @pytest.mark.parametrize( + "owner_role", + [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], + ids=["proxy_admin", "proxy_admin_viewer"], + ) + async def test_admin_owned_team_key_without_the_toolset_grant_is_denied(self, owner_role: LitellmUserRoles): + from litellm.proxy.mcp_server.server import apply_toolset_scope + + team_key: Final = UserAPIKeyAuth( + api_key="sk-test", team_id="team-a", user_role=owner_role, object_permission=None + ) + with pytest.raises(HTTPException) as exc_info: + await apply_toolset_scope(team_key, "toolset-123", granted=_granted_through_team("toolset-other")) + + assert exc_info.value.status_code == 403 + @pytest.mark.asyncio async def test_team_granted_toolset_is_served_to_a_key_without_its_own_grant(self): """A team key whose own row carries no toolset grant is admitted to the toolset its team diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index 82523572b28..238f1fa6495 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -3,10 +3,13 @@ from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +import respx from fastapi import HTTPException from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles +import litellm from litellm import Router from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth @@ -852,3 +855,112 @@ def test_a_team_keeps_its_first_config_file_openai_deployment_over_older_db_ones assert credentials is not None assert credentials["api_key"] == "deployment-z-key" + + +@pytest.mark.parametrize( + ("custom_llm_provider", "openrouter_key", "api_key"), + [("openai", "sk-or-files-openai", "openai-key"), ("openrouter", "sk-or-files-openrouter", "sk-or-files-openrouter")], +) +def test_a_provider_credential_lookup_never_lists_an_openrouter_wildcards_live_catalog( + monkeypatch: pytest.MonkeyPatch, custom_llm_provider: str, openrouter_key: str, api_key: str +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + router: Final = Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": openrouter_key, + "api_base": "https://openrouter.ai/api/v1", + }, + }, + {"model_name": "gpt-5.4-mini", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "openai-key"}}, + ] + ) + + with respx.mock(assert_all_called=False) as respx_mock: + catalog_route: Final = respx_mock.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + ) + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + custom_llm_provider=custom_llm_provider, + include_accessible_deployments=True, + ) + + assert credentials is not None + assert (credentials["custom_llm_provider"], credentials["api_key"]) == (custom_llm_provider, api_key) + assert catalog_route.call_count == 0 + + +@pytest.mark.parametrize( + ("key_models", "deployments", "custom_llm_provider", "api_key"), + [ + ( + ["anthropic/*"], + (("anthropic/claude-*", "anthropic/claude-*", "anthropic-narrow-key"),), + "anthropic", + "anthropic-narrow-key", + ), + (["anthropic/claude-3*"], (("anthropic/*", "anthropic/*", "anthropic-wide-key"),), "anthropic", "anthropic-wide-key"), + (["openrouter/*"], (("openrouter/openai/*", "openrouter/openai/*", "sk-or-narrow"),), "openrouter", "sk-or-narrow"), + ( + ["team/*"], + (("team/*", "openrouter/*", "sk-or-team"), ("team/openai/*", "openai/*", "openai-team-key")), + "openai", + "openai-team-key", + ), + ], +) +def test_a_key_wildcard_finds_the_deployment_it_grants_without_listing_any_provider_catalog( + monkeypatch: pytest.MonkeyPatch, + key_models: list[str], + deployments: tuple[tuple[str, str, str], ...], + custom_llm_provider: str, + api_key: str, +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + router: Final = Router( + model_list=[ + {"model_name": model_name, "litellm_params": {"model": model, "api_key": deployment_key}} + for model_name, model, deployment_key in deployments + ] + ) + + with respx.mock(assert_all_called=False) as respx_mock: + catalog_requests: Final = respx_mock.route().mock(return_value=httpx.Response(200, json={"data": []})) + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=key_models), + custom_llm_provider=custom_llm_provider, + include_accessible_deployments=True, + ) + + assert credentials is not None + assert (credentials["custom_llm_provider"], credentials["api_key"]) == (custom_llm_provider, api_key) + assert catalog_requests.call_count == 0 + + +@pytest.mark.parametrize( + ("deployed_name", "api_key"), [("gpt-4o-mini", None), ("openai/gpt-4o-mini", "openai-deployment-key")] +) +def test_a_provider_wildcard_grant_finds_only_deployments_named_under_it(deployed_name: str, api_key: str | None): + router: Final = Router( + model_list=[ + { + "model_name": deployed_name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "openai-deployment-key"}, + } + ] + ) + + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=["openai/*"]), + custom_llm_provider="openai", + include_accessible_deployments=True, + ) + + assert (None if credentials is None else credentials["api_key"]) == api_key diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py index 4011bc0294a..7f8aac983f5 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py @@ -323,6 +323,35 @@ def test_openrouter_decisions_response_uses_provider_cost_without_cost_map_row( assert response["kwargs"]["combined_usage_object"].total_tokens == 302 +@pytest.mark.parametrize("reported_cost", [float("nan"), float("inf")]) +def test_openrouter_non_finite_cost_logs_the_token_price(monkeypatch: pytest.MonkeyPatch, reported_cost: float) -> None: + monkeypatch.setitem( + litellm.model_cost, + "openrouter/typesafe/jev-1.13", + {"input_cost_per_token": 1.0, "output_cost_per_token": 3.0, "litellm_provider": "openrouter", "mode": "chat"}, + ) + logging_obj: Final = _logging_obj() + + result: Final = TypeSafePassthroughLoggingHandler.typesafe_passthrough_handler( + httpx_response=_response(), + response_body={ + "model": "typesafe/jev-1.13-20260917", + "usage": {"input_tokens": 2, "output_tokens": 1, "cost": reported_cost}, + }, + logging_obj=logging_obj, + url_route="https://openrouter.ai/api/alpha/decisions", + result='{"answers": {}}', + start_time=datetime(2000, 1, 1), + end_time=datetime(2000, 1, 1), + cache_hit=False, + request_body={"model": "typesafe/jev-1.13"}, + custom_llm_provider="openrouter", + ) + + assert result["kwargs"]["response_cost"] == pytest.approx(5.0) + assert logging_obj.model_call_details["response_cost"] == pytest.approx(5.0) + + def test_success_handler_dispatches_openrouter_to_the_shared_handler(): logging_obj = _logging_obj() normalized = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 806d88b2467..4a76aa77f14 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -25,12 +25,14 @@ from starlette.datastructures import FormData import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.prometheus import PrometheusLogger from tests._master_key import MASTER_KEY as SHARED_MASTER_KEY from litellm.caching.caching import DualCache -from litellm.types.utils import CallTypesLiteral +from litellm.types.utils import CallTypes, CallTypesLiteral from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException +from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id from litellm.proxy.utils import ProxyLogging from litellm.passthrough.main import raise_for_provider_status from litellm.llms.bedrock.common_utils import BedrockError @@ -2759,6 +2761,68 @@ class TestVLLMProxyRoute: mock_factory_route.assert_awaited_once() +@pytest.mark.parametrize( + ("vllm_api_base", "request_path"), + [ + ("http://vllm-upstream.test:8090", "/vllm/v1/chat/completions"), + ("http://vllm-upstream.test:8090/v1", "/vllm/v1/chat/completions"), + ("http://vllm-upstream.test:8090/v1", "/vllm/chat/completions"), + ], +) +def test_vllm_provider_route_reaches_v1_paths_with_its_key_as_a_bearer( + monkeypatch: pytest.MonkeyPatch, vllm_api_base: str, request_path: str +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setenv("VLLM_API_BASE", vllm_api_base) + monkeypatch.setenv("VLLM_API_KEY", "sk-vllm-upstream") + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem( + proxy_server.app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual") + ) + + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("http://vllm-upstream.test:8090/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"id": "chatcmpl-1"}) + ) + response: Final = TestClient(proxy_server.app).post( + request_path, + json={"model": "Qwen/Qwen3-0.6B", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert (response.status_code, response.json()) == (200, {"id": "chatcmpl-1"}) + assert route.calls.last.request.headers["authorization"] == "Bearer sk-vllm-upstream" + assert route.calls.last.request.headers["x-api-key"] == "sk-vllm-upstream" + + +def test_vllm_provider_route_with_an_empty_key_sends_no_credential_headers(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setenv("VLLM_API_BASE", "http://vllm-upstream.test:8090") + monkeypatch.setenv("VLLM_API_KEY", "") + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem( + proxy_server.app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual") + ) + + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("http://vllm-upstream.test:8090/v1/chat/completions").mock( + return_value=httpx.Response(200, json={"id": "chatcmpl-1"}) + ) + response: Final = TestClient(proxy_server.app).post( + "/vllm/v1/chat/completions", + json={"model": "Qwen/Qwen3-0.6B", "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert (response.status_code, response.json()) == (200, {"id": "chatcmpl-1"}) + assert "authorization" not in route.calls.last.request.headers + assert "x-api-key" not in route.calls.last.request.headers + + class TestGigachatProxyRoute: @pytest.mark.asyncio @patch( # test-quality-ok: patching litellm internal for unit test isolation @@ -6428,6 +6492,55 @@ class TestVertexAILiveWebsocketPassthrough: assert "default_vertex_config" in close_kwargs["reason"] assert len(close_kwargs["reason"].encode("utf-8")) <= 123 + @pytest.mark.asyncio + @pytest.mark.parametrize("model, expected_model", [("gemini-3.8-live", "gemini-3.8-live"), (None, "")]) + async def test_credential_failure_logs_the_named_model_under_vertex_ai( + self, monkeypatch: pytest.MonkeyPatch, model: str | None, expected_model: str + ): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + + class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.request_data: tuple[dict[str, object], ...] = () + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.request_data = (*self.request_data, request_data) + + recorder: Final = _FailureRecorder() + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()) + ) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None) + self._clear_vertex_env(monkeypatch) + ensure_token: Final = AsyncMock(side_effect=Exception("Unable to find your credentials")) + + with patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token): + await passthrough_module.vertex_ai_live_websocket_passthrough( + websocket=self._websocket(), + model=model, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert [ + ( + d["model"], + PrometheusLogger._extract_api_provider_from_request_data(d), + get_spend_logs_id(CallTypes.pass_through.value, {}, d) is not None, + ) + for d in recorder.request_data + ] == [(expected_model, "vertex_ai", True)] + class TestAnthropicProxyRoute: """The /anthropic passthrough route: custom auth headers must not clobber the @@ -6904,7 +7017,7 @@ class TestAzureRouterModelStreamingDispatch: logging_obj = MagicMock() logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() - from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment + from litellm.litellm_core_utils.add_retry_fallback_headers import prepare_response_for_header_attachment class StreamingRouter: async def allm_passthrough_route(self, **kwargs): diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 7b62853b0e9..b7e76ce0068 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -30,20 +30,26 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( + VertexPassthroughLoggingHandler, +) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, SafeRouteAdder, + _build_passthrough_failure_request_payload, _registered_pass_through_routes, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -61,13 +67,17 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id +from litellm.proxy.utils import ProxyLogging from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypes from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, EndpointType, PassthroughStandardLoggingPayload, + WebSocketSessionUsage, ) from tests._master_key import MASTER_KEY from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -140,9 +150,7 @@ def test_is_multipart(): ), ], ) -def test_get_endpoint_type_classifies_anthropic_messages_routes( - url: str, expected_endpoint_type: EndpointType -) -> None: +def test_get_endpoint_type_classifies_anthropic_messages_routes(url: str, expected_endpoint_type: EndpointType) -> None: assert HttpPassThroughEndpointHelpers.get_endpoint_type(url) == expected_endpoint_type @@ -430,53 +438,18 @@ async def test_pass_through_request_failure_handler(): Critical Test: When a users pass through endpoint request fails, we must log the failure code, exception in litellm spend logs. """ - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client") as mock_get_client: - with patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processing: - # Setup mock for post_call_failure_hook and pre_call_hook - mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.pre_call_hook = AsyncMock() + request_failure = httpx.HTTPError("Request failed") + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") - # Setup mock for httpx client - mock_client = MagicMock() - mock_client.client = MagicMock() - mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed")) - mock_get_client.return_value = mock_client + failure: Final = await _run_passthrough_failure( + "http://test.com", None, request_failure, user_api_key_dict=user_api_key_dict + ) - # Mock headers for custom headers - mock_processing.get_custom_headers.return_value = {} - - # Create mock request - mock_request = MagicMock(spec=Request) - mock_request.method = "POST" - mock_request.body = AsyncMock(return_value=b'{"test": "data"}') - mock_request.headers = Headers({}) - - # Create a simple empty QueryParams - mock_request.query_params = QueryParams({}) - - # Create mock user API key dict - mock_user_api_key_dict = MagicMock() - - # Call the function with a target that will trigger an HTTPError - with pytest.raises(ProxyException): - await pass_through_request( - request=mock_request, - target="http://test.com", - custom_headers={}, - user_api_key_dict=mock_user_api_key_dict, - ) - - # Assert post_call_failure_hook was called - mock_proxy_logging.post_call_failure_hook.assert_called_once() - - # Verify the arguments to post_call_failure_hook - call_args = mock_proxy_logging.post_call_failure_hook.call_args[1] - assert call_args["user_api_key_dict"] == mock_user_api_key_dict - assert isinstance(call_args["original_exception"], TypeError) # Now expecting TypeError - assert "traceback_str" in call_args + assert failure.user_api_key_dict is user_api_key_dict + assert failure.original_exception is request_failure + assert "Request failed" in (failure.traceback_str or "") + assert failure.request_data["contents"] == "not a list" + assert failure.request_data["model"] == "" @pytest.mark.asyncio @@ -3656,9 +3629,7 @@ async def test_pass_through_request_preserves_target_query_without_client_query( @pytest.mark.asyncio -async def test_config_pass_through_without_configured_headers_preserves_no_authorization( - tmp_path, monkeypatch -): +async def test_config_pass_through_without_configured_headers_preserves_no_authorization(tmp_path, monkeypatch): proxy = await _boot_db_backed_proxy( tmp_path, monkeypatch, @@ -3725,9 +3696,7 @@ async def test_pass_through_rpm_limit_is_sequential_and_scoped_per_key(tmp_path, proxy = await _boot_db_backed_proxy( tmp_path, monkeypatch, - config_pass_through_endpoints=[ - {"path": "/rpm-rerank", "target": "http://upstream.test/rerank", "auth": True} - ], + config_pass_through_endpoints=[{"path": "/rpm-rerank", "target": "http://upstream.test/rerank", "auth": True}], db_pass_through_endpoints=[], master_key=MASTER_KEY, ) @@ -4902,8 +4871,10 @@ from fastapi import HTTPException as _FastAPIHTTPException from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, + ModifyResponseException, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + _PT_MODULE = "litellm.proxy.pass_through_endpoints.pass_through_endpoints" @@ -7766,6 +7737,102 @@ async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text upstream_ws.close.assert_awaited_once() +@dataclass(frozen=True, slots=True) +class _LengthMeter: + rate: float + + def measure(self, client_text: str) -> int: + return len(client_text) + + def price(self, total: int) -> WebSocketSessionUsage: + return WebSocketSessionUsage(model="voice-1", custom_llm_provider="tts-provider", response_cost=total * self.rate) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_the_meters_price_for_the_client_text_it_delivered(): + """Text frames reach the upstream unchanged and are measured, binary frames are not, and the session is logged at + the model, provider and cost the meter prices for the total.""" + frames = ("hello", b"\x00\x01binary", "world!") + upstream_ws = RecordingUpstreamWebSocket() + websocket = _client_websocket( + AsyncMock( + side_effect=[ + *( + {"type": "websocket.receive", "bytes": frame} + if isinstance(frame, bytes) + else {"type": "websocket.receive", "text": frame} + for frame in frames + ), + {"type": "websocket.disconnect"}, + ] + ) + ) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.example/v1/stream", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/provider/v1/stream", + accept_websocket=False, + meter=_LengthMeter(rate=0.5), + ) + + assert [call.args[0] for call in upstream_ws.send.await_args_list] == list(frames) + logging_obj = success_handler.call_args.kwargs["logging_obj"] + assert logging_obj.model == "voice-1" + assert logging_obj.model_call_details["model"] == "voice-1" + assert logging_obj.model_call_details["custom_llm_provider"] == "tts-provider" + assert logging_obj.model_call_details["response_cost"] == (len("hello") + len("world!")) * 0.5 + + +@pytest.mark.asyncio +async def test_websocket_passthrough_meters_only_client_text_the_upstream_accepted(): + upstream_ws = RecordingUpstreamWebSocket() + upstream_ws.send = AsyncMock(side_effect=[None, ConnectionError("upstream went away")]) + websocket = _client_websocket( + AsyncMock( + side_effect=[ + {"type": "websocket.receive", "text": "hello"}, + {"type": "websocket.receive", "text": "never delivered"}, + {"type": "websocket.disconnect"}, + ] + ) + ) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.example/v1/stream", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/provider/v1/stream", + accept_websocket=False, + meter=_LengthMeter(rate=0.5), + ) + + logging_obj = success_handler.call_args.kwargs["logging_obj"] + assert [call.args[0] for call in upstream_ws.send.await_args_list] == ["hello", "never delivered"] + assert logging_obj.model_call_details["response_cost"] == len("hello") * 0.5 + + @pytest.mark.asyncio async def test_websocket_passthrough_vertex_live_setup_ack_names_the_model_but_is_not_billed_as_usage(): """Vertex Live keeps its special first frame: the setup acknowledgement is forwarded verbatim, read for the @@ -8182,8 +8249,10 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata } ) body: Final = { - "model": "upstream-only-model", metadata_slot: { - "model_group": "managed-model", "customer_label": "retained", + "model": "upstream-only-model", + metadata_slot: { + "model_group": "managed-model", + "customer_label": "retained", "user_api_key_team_model_max_budget": budget, }, } @@ -8191,48 +8260,84 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata assert await limiter.is_team_within_model_budget("shared-team", budget, None, "managed-model") start: Final = datetime.now() logging_obj: Final = LiteLLMLoggingObj( - model="upstream-only-model", messages=[], stream=False, call_type="pass_through_endpoint", - start_time=start, litellm_call_id="custom-budget", function_id="custom-budget", kwargs={}, + model="upstream-only-model", + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=start, + litellm_call_id="custom-budget", + function_id="custom-budget", + kwargs={}, dynamic_async_success_callbacks=[limiter], ) payload: Final = { - "url": "https://upstream.test/echo", "request_body": body, "request_method": "POST", "cost_per_request": 0.25, + "url": "https://upstream.test/echo", + "request_body": body, + "request_method": "POST", + "cost_per_request": 0.25, } kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( - request=request, user_api_key_dict=auth, passthrough_logging_payload=payload, logging_obj=logging_obj, - _parsed_body=body, litellm_call_id="custom-budget", + request=request, + user_api_key_dict=auth, + passthrough_logging_payload=payload, + logging_obj=logging_obj, + _parsed_body=body, + litellm_call_id="custom-budget", ) logging_obj.update_environment_variables( - model="upstream-only-model", user="unknown", optional_params={}, - litellm_params=kwargs["litellm_params"], call_type="pass_through_endpoint", + model="upstream-only-model", + user="unknown", + optional_params={}, + litellm_params=kwargs["litellm_params"], + call_type="pass_through_endpoint", ) response: Final = httpx.Response( - 200, request=httpx.Request("POST", "https://upstream.test/echo"), json={"ok": True}, + 200, + request=httpx.Request("POST", "https://upstream.test/echo"), + json={"ok": True}, ) await PassThroughEndpointLogging().pass_through_async_success_handler( - httpx_response=response, response_body={"ok": True}, request_body=body, logging_obj=logging_obj, - url_route="https://upstream.test/echo", result=response.text, start_time=start, end_time=datetime.now(), - cache_hit=False, **kwargs, + httpx_response=response, + response_body={"ok": True}, + request_body=body, + logging_obj=logging_obj, + url_route="https://upstream.test/echo", + result=response.text, + start_time=start, + end_time=datetime.now(), + cache_hit=False, + **kwargs, ) assert logging_obj.model_call_details["response_cost"] == 0.25 assert await limiter.is_team_within_model_budget("shared-team", budget, None, "managed-model") metadata: Final = kwargs["litellm_params"]["metadata"] assert (metadata["model_group"], metadata["customer_label"]) == ("managed-model", "retained") - assert metadata.keys().isdisjoint({ - "user_api_key_model_max_budget", "user_api_key_team_model_max_budget", - "user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget", - }) + assert metadata.keys().isdisjoint( + { + "user_api_key_model_max_budget", + "user_api_key_team_model_max_budget", + "user_api_key_user_model_max_budget", + "user_api_key_end_user_model_max_budget", + } + ) @pytest.mark.parametrize("metadata_slot", ["metadata", "litellm_metadata"]) def test_builtin_passthrough_pins_model_group_to_the_resolved_model(metadata_slot: str) -> None: - request: Final = Request({ - "type": "http", "method": "POST", "path": "/gemini/v1beta/models/gemini-2.5-flash:generateContent", - "headers": [], "query_string": b"", - }) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/gemini/v1beta/models/gemini-2.5-flash:generateContent", + "headers": [], + "query_string": b"", + } + ) kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( - request=request, user_api_key_dict=UserAPIKeyAuth(token="hash", user_id="u-1"), - passthrough_logging_payload=MagicMock(), logging_obj=MagicMock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(token="hash", user_id="u-1"), + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), _parsed_body={"contents": [], metadata_slot: {"model_group": "unbounded-client-choice"}}, ) assert kwargs["litellm_params"]["metadata"]["model_group"] == "gemini-2.5-flash" @@ -8591,6 +8696,581 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) +_GEMINI_FLASH_TARGET: Final = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash:generateContent" +_VERTEX_FLASH_TARGET: Final = ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/google/models/gemini-3.8-flash:generateContent" +) +_VERTEX_CLAUDE_RAW_PREDICT_TARGET: Final = ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/anthropic/models/claude-opus-5:rawPredict" +) +_VERTEX_SEARCH_TARGET: Final = ( + "https://discoveryengine.googleapis.com/v1/projects/p/locations/global" + "/dataStores/default/servingConfigs/default:search" +) +_RELAYED_GEMINI_COUNT_TOKENS_TARGET: Final = "http://127.0.0.1:30732/v1beta/models/gemini-3.8-flash:countTokens" +_TRANSLATE_TARGET: Final = "https://translate.googleapis.com/v3/projects/p/locations/global:translateText" +_NON_GOOGLE_VERTEX_SHAPED_TARGET: Final = "https://llm.example.com/v1/projects/a/locations/b/models/my-model:predict" +_VERTEX_LIVE_TARGET: Final = ( + "wss://us-central1-aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" +) +_VEO_OPERATION_SUFFIX: Final = "/models/veo-3.0-generate-001/operations/op-8f2c1" +_RELAYED_VERTEX_COUNT_TOKENS_TARGET: Final = ( + "http://127.0.0.1:30732/v1/projects/p/locations/us-east5/publishers/google/models/gemini-3.8-flash:countTokens" +) + + +@dataclass(frozen=True, slots=True) +class _RecordedFailure: + request_data: dict[str, object] + original_exception: Exception + user_api_key_dict: UserAPIKeyAuth + traceback_str: str | None + + +class _FailureHookRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: tuple[_RecordedFailure, ...] = () + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.failures = ( + *self.failures, + _RecordedFailure(request_data, original_exception, user_api_key_dict, traceback_str), + ) + + +@dataclass(frozen=True, slots=True) +class _PassthroughRun: + failure: _RecordedFailure + response: Response | None + + +async def _drained_pass_through_response( + request: Request, + target: str, + custom_llm_provider: str | None, + user_api_key_dict: UserAPIKeyAuth, + stream: bool, +) -> Response | None: + try: + received: Final = await pass_through_request( + request=request, + target=target, + custom_headers={}, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + stream=stream, + ) + if isinstance(received, StreamingResponse): + async for _ in received.body_iterator: + pass + except (ProxyException, httpx.HTTPError): + return None + return received if isinstance(received, Response) else None + + +async def _run_passthrough_failure( + target: str, + custom_llm_provider: str | None, + upstream: httpx.Response | Exception, + stream: bool = False, + user_api_key_dict: UserAPIKeyAuth | None = None, + dispatched_endpoint: Callable[..., object] | None = None, + guardrail: CustomGuardrail | None = None, + body: bytes = b'{"contents": "not a list"}', +) -> _RecordedFailure: + run: Final = await _run_passthrough( + target, custom_llm_provider, upstream, stream, user_api_key_dict, dispatched_endpoint, guardrail, body + ) + return run.failure + + +async def _run_passthrough( + target: str, + custom_llm_provider: str | None, + upstream: httpx.Response | Exception, + stream: bool = False, + user_api_key_dict: UserAPIKeyAuth | None = None, + dispatched_endpoint: Callable[..., object] | None = None, + guardrail: CustomGuardrail | None = None, + body: bytes = b'{"contents": "not a list"}', +) -> _PassthroughRun: + recorder: Final = _FailureHookRecorder() + with ( + pytest.MonkeyPatch.context() as litellm_patch, + patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request; there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()) + ), + respx.mock(assert_all_called=False) as upstream_mock, + ): + litellm_patch.setattr(litellm, "disable_aiohttp_transport", True) + litellm_patch.setattr(litellm, "callbacks", [recorder] if guardrail is None else [guardrail, recorder]) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream_route: Final = upstream_mock.route(method="POST") + if isinstance(upstream, Exception): + upstream_route.mock(side_effect=upstream) + else: + upstream_route.mock(return_value=upstream) + mock_request: Final = MagicMock(spec=Request) + mock_request.method = "POST" + proxy_url: Final = httpx.URL(target.replace("https://", "http://test-proxy.com/")) + mock_request.url = proxy_url + mock_request.scope = { + "type": "http", + "method": "POST", + "path": proxy_url.path, + **({"endpoint": dispatched_endpoint} if dispatched_endpoint is not None else {}), + } + mock_request.body = AsyncMock(return_value=body) + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + response: Final = await _drained_pass_through_response( + mock_request, + target, + custom_llm_provider, + user_api_key_dict or UserAPIKeyAuth(api_key="sk-test"), + stream, + ) + await asyncio.sleep(0) + litellm.in_memory_llm_clients_cache.flush_cache() + assert len(recorder.failures) == 1 + return _PassthroughRun(recorder.failures[0], response) + + +def _upstream_400(target: str) -> httpx.Response: + return httpx.Response( + status_code=400, + headers={"content-type": "application/json"}, + content=json.dumps({"error": {"code": 400, "status": "INVALID_ARGUMENT"}}).encode(), + request=httpx.Request("POST", target), + ) + + +def _upstream_400_as_event_stream(target: str) -> httpx.Response: + return httpx.Response( + status_code=400, + headers={"content-type": "text/event-stream"}, + content=b'data: {"error": {"code": 400, "status": "INVALID_ARGUMENT"}}\n\n', + request=httpx.Request("POST", target), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target, custom_llm_provider, stream, expected_provider", + [ + (_GEMINI_FLASH_TARGET, "gemini", False, "gemini"), + (_GEMINI_FLASH_TARGET, "gemini", True, "gemini"), + (_VERTEX_FLASH_TARGET, None, False, "vertex_ai"), + (_VERTEX_FLASH_TARGET, None, True, "vertex_ai"), + (_GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), "gemini", False, "gemini"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), None, False, "vertex_ai"), + (_RELAYED_GEMINI_COUNT_TOKENS_TARGET, "gemini", False, "gemini"), + ], +) +async def test_passthrough_upstream_error_logs_model_from_url( + target: str, custom_llm_provider: str | None, stream: bool, expected_provider: str +): + failure: Final = await _run_passthrough_failure(target, custom_llm_provider, _upstream_400(target), stream=stream) + assert failure.request_data["model"] == "gemini-3.8-flash" + assert failure.request_data["custom_llm_provider"] == expected_provider + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target", + [ + _RELAYED_VERTEX_COUNT_TOKENS_TARGET, + _TRANSLATE_TARGET, + _NON_GOOGLE_VERTEX_SHAPED_TARGET, + ], +) +async def test_passthrough_upstream_error_to_a_host_outside_gemini_and_vertex_gets_no_url_attribution(target: str): + failure: Final = await _run_passthrough_failure(target, None, _upstream_400(target)) + assert failure.request_data["model"] == "" + assert failure.request_data.get("custom_llm_provider") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target", [_GEMINI_FLASH_TARGET, _VERTEX_FLASH_TARGET]) +async def test_passthrough_upstream_error_on_a_user_defined_route_to_google_gets_no_url_attribution(target: str): + failure: Final = await _run_passthrough_failure( + target, None, _upstream_400(target), dispatched_endpoint=_marked_pass_through_endpoint() + ) + assert failure.request_data["model"] == "" + assert failure.request_data.get("custom_llm_provider") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target, custom_llm_provider, expected_provider", + [ + (_GEMINI_FLASH_TARGET, "gemini", "gemini"), + (_VERTEX_FLASH_TARGET, None, "vertex_ai"), + ], +) +async def test_passthrough_detected_stream_upstream_error_logs_model_from_url( + target: str, custom_llm_provider: str | None, expected_provider: str +): + failure: Final = await _run_passthrough_failure( + target, custom_llm_provider, _upstream_400_as_event_stream(target), stream=False + ) + assert failure.request_data["model"] == "gemini-3.8-flash" + assert failure.request_data["custom_llm_provider"] == expected_provider + + +@pytest.mark.asyncio +async def test_passthrough_internal_failure_logs_model_from_url(): + failure: Final = await _run_passthrough_failure( + _GEMINI_FLASH_TARGET, "gemini", httpx.ReadTimeout("upstream timed out") + ) + assert failure.request_data["model"] == "gemini-3.8-flash" + assert failure.request_data["custom_llm_provider"] == "gemini" + + +def _upstream_dropping_stream(target: str) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "text/event-stream"}, + stream=_UpstreamDroppingMidStream(), + request=httpx.Request("POST", target), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize( + "target, custom_llm_provider, expected_provider", + [ + (_GEMINI_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), "gemini", "gemini"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), None, "vertex_ai"), + ], +) +async def test_passthrough_mid_stream_drop_logs_model_from_url( + target: str, custom_llm_provider: str | None, expected_provider: str, stream: bool +): + failure: Final = await _run_passthrough_failure( + target, custom_llm_provider, _upstream_dropping_stream(target), stream=stream + ) + assert isinstance(failure.original_exception, httpx.ReadError) + assert failure.request_data["model"] == "gemini-3.8-flash" + assert failure.request_data["custom_llm_provider"] == expected_provider + + +class _ModifyResponseOnPreCallGuardrail(CustomGuardrail): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: str, + ) -> None: + self.raise_passthrough_exception(violation_message="blocked by policy", request_data=data, detection_info=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target, custom_llm_provider, body, expected_model, expected_provider", + [ + (_GEMINI_FLASH_TARGET, "gemini", b'{"contents": "not a list"}', "gemini-3.8-flash", "gemini"), + (_VERTEX_FLASH_TARGET, None, b'{"contents": "not a list"}', "gemini-3.8-flash", "vertex_ai"), + (_VERTEX_FLASH_TARGET, None, b'{"custom_llm_provider": "gemini"}', "gemini-3.8-flash", "vertex_ai"), + ], +) +async def test_passthrough_guardrail_block_logs_model_from_url( + target: str, custom_llm_provider: str | None, body: bytes, expected_model: str, expected_provider: str +): + guardrail: Final = _ModifyResponseOnPreCallGuardrail( + guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + failure: Final = await _run_passthrough_failure( + target, custom_llm_provider, _upstream_400(target), guardrail=guardrail, body=body + ) + assert isinstance(failure.original_exception, ModifyResponseException) + assert failure.request_data["model"] == expected_model + assert failure.request_data["custom_llm_provider"] == expected_provider + assert PrometheusLogger._extract_api_provider_from_request_data(failure.request_data) == expected_provider + assert get_spend_logs_id(CallTypes.pass_through.value, {}, failure.request_data) is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [b'{"contents": "not a list"}', b'{"litellm_call_id": "client-chosen-id"}']) +async def test_passthrough_guardrail_block_response_names_the_call_id_of_its_spend_row(body: bytes): + guardrail: Final = _ModifyResponseOnPreCallGuardrail( + guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + run: Final = await _run_passthrough( + _GEMINI_FLASH_TARGET, "gemini", _upstream_400(_GEMINI_FLASH_TARGET), guardrail=guardrail, body=body + ) + assert isinstance(run.failure.original_exception, ModifyResponseException) + assert run.response is not None + assert run.response.headers["x-litellm-call-id"] == get_spend_logs_id( + CallTypes.pass_through.value, {}, run.failure.request_data + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [b'{"contents": "not a list"}', b'{"custom_llm_provider": "vertex_ai"}']) +async def test_passthrough_guardrail_block_on_a_user_defined_route_gets_no_url_attribution(body: bytes): + guardrail: Final = _ModifyResponseOnPreCallGuardrail( + guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + failure: Final = await _run_passthrough_failure( + _GEMINI_FLASH_TARGET, + None, + _upstream_400(_GEMINI_FLASH_TARGET), + dispatched_endpoint=_marked_pass_through_endpoint(), + guardrail=guardrail, + body=body, + ) + assert isinstance(failure.original_exception, ModifyResponseException) + assert failure.request_data["model"] == "" + assert "custom_llm_provider" not in failure.request_data + + +def _websocket_connect_failure(kind: str) -> Exception: + from websockets.datastructures import Headers as WebSocketHeaders + from websockets.exceptions import InvalidStatus + from websockets.http11 import Response as WebSocketHandshakeResponse + + if kind == "rejected_handshake": + return InvalidStatus(WebSocketHandshakeResponse(403, "Forbidden", WebSocketHeaders())) + return OSError("connection refused") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_kind", ["rejected_handshake", "unreachable_host"]) +@pytest.mark.parametrize( + "target, builtin_provider_route, expected_provider", + [ + (_VERTEX_LIVE_TARGET, True, "vertex_ai"), + (_VERTEX_LIVE_TARGET, False, None), + ("wss://api.openai.com/v1/realtime?model=gpt-realtime", True, None), + ], +) +async def test_websocket_passthrough_connect_failure_logs_the_provider_of_the_upstream_url( + failure_kind: str, target: str, builtin_provider_route: bool, expected_provider: str | None +): + failure: Final = await _run_websocket_connect_failure(failure_kind, target, builtin_provider_route) + assert failure.request_data.get("custom_llm_provider") == expected_provider + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_kind", ["rejected_handshake", "unreachable_host"]) +@pytest.mark.parametrize("model, expected_model", [("gemini-3.8-live", "gemini-3.8-live"), (None, "")]) +async def test_websocket_passthrough_connect_failure_logs_the_model_the_caller_named( + failure_kind: str, model: str | None, expected_model: str +): + failure: Final = await _run_websocket_connect_failure(failure_kind, _VERTEX_LIVE_TARGET, True, model=model) + assert failure.request_data["model"] == expected_model + + +async def _run_websocket_connect_failure( + failure_kind: str, target: str, builtin_provider_route: bool, model: str | None = None +) -> _RecordedFailure: + recorder: Final = _FailureHookRecorder() + with ( + pytest.MonkeyPatch.context() as litellm_patch, + patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside websocket_passthrough_request; there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()) + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + side_effect=_websocket_connect_failure(failure_kind), + ), + ): + litellm_patch.setattr(litellm, "callbacks", [recorder]) + await websocket_passthrough_request( + websocket=_client_websocket(_pending_receive), + target=target, + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + endpoint="/vertex_ai/live", + accept_websocket=False, + builtin_provider_route=builtin_provider_route, + model=model, + ) + assert len(recorder.failures) == 1 + return recorder.failures[0] + + +@pytest.mark.parametrize( + "parsed_body, url, custom_llm_provider, expected_model, expected_provider", + [ + ({"model": "from-body"}, _GEMINI_FLASH_TARGET, "gemini", "from-body", "gemini"), + ({"contents": []}, _GEMINI_FLASH_TARGET, "gemini", "gemini-3.8-flash", "gemini"), + ({"contents": []}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "vertex_ai"), + ({"contents": []}, _VERTEX_FLASH_TARGET, "anthropic", "", "anthropic"), + ({}, "https://generativelanguage.googleapis.com/v1beta/openai/models/gemini-3.8-flash", "openai", "", "openai"), + ({"custom_llm_provider": "from-body"}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "vertex_ai"), + ({"custom_llm_provider": "gemini"}, _VERTEX_FLASH_TARGET, "anthropic", "", "anthropic"), + ({"contents": []}, _GEMINI_FLASH_TARGET, None, "gemini-3.8-flash", "gemini"), + ( + {}, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash?key=provider-secret", + "gemini", + "gemini-3.8-flash", + "gemini", + ), + ( + {}, + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/google/models/gemini-3.8-flash?key=provider-secret", + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ({}, _VERTEX_CLAUDE_RAW_PREDICT_TARGET, None, "vertex_ai/claude-opus-5", "vertex_ai"), + ( + {}, + _VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict?alt=sse"), + None, + "claude-opus-5", + "vertex_ai", + ), + ({}, "https://generativelanguage.googleapis.com/v1beta/models?pageToken=/models/x", "gemini", "", "gemini"), + ( + {"contents": []}, + "https://evilaiplatform.googleapis.com/v1/projects/p/locations/l" + "/publishers/google/models/gemini-3.8-flash:generateContent", + None, + "", + None, + ), + ({}, _VERTEX_SEARCH_TARGET, None, "vertex_ai/search_api", "vertex_ai"), + ({}, f"{_VERTEX_FLASH_TARGET}?tag=rawPredict", None, "gemini-3.8-flash", "vertex_ai"), + ( + {}, + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs" + "?filter=:search", + None, + "", + "vertex_ai", + ), + ( + {}, + f"https://generativelanguage.googleapis.com/v1beta{_VEO_OPERATION_SUFFIX}", + "gemini", + "veo-3.0-generate-001", + "gemini", + ), + ( + {}, + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1" + f"/publishers/google{_VEO_OPERATION_SUFFIX}", + None, + "veo-3.0-generate-001", + "vertex_ai", + ), + ( + {"contents": []}, + "https://aiplatform.us.rep.googleapis.com/v1/projects/p/locations/us" + "/publishers/google/models/gemini-3.8-flash:generateContent", + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ( + {}, + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs", + None, + "", + "vertex_ai", + ), + ( + {"contents": []}, + _GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), + "gemini", + "gemini-3.8-flash", + "gemini", + ), + ( + {"contents": []}, + _VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ( + {"contents": []}, + "https://aiplatform.googleapis.com/v1/projects/p/locations/global" + "/publishers/google/models/gemini-3.8-flash:countTokens", + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ( + {}, + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/cachedContents", + None, + "", + "vertex_ai", + ), + ({"contents": []}, _RELAYED_GEMINI_COUNT_TOKENS_TARGET, "gemini", "gemini-3.8-flash", "gemini"), + ({"contents": []}, _RELAYED_VERTEX_COUNT_TOKENS_TARGET, None, "", None), + ({"contents": []}, _TRANSLATE_TARGET, None, "", None), + ({"contents": []}, _NON_GOOGLE_VERTEX_SHAPED_TARGET, None, "", None), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:countTokens", None, "", None), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:predict", None, "", None), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, "", None), + (None, None, None, "", None), + ], +) +def test_build_passthrough_failure_request_payload_model_and_provider_precedence( + parsed_body: dict | None, + url: str | None, + custom_llm_provider: str | None, + expected_model: str, + expected_provider: str | None, +): + request_payload: Final = _build_passthrough_failure_request_payload( + parsed_body=parsed_body, + kwargs=None, + logging_obj=None, + custom_llm_provider=custom_llm_provider, + url=httpx.URL(url) if url else None, + builtin_provider_route=True, + ) + assert request_payload["model"] == expected_model + assert request_payload.get("custom_llm_provider") == expected_provider + + +@pytest.mark.parametrize( + "url_route, expected_model", + [ + (_VERTEX_FLASH_TARGET, "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":predict"), "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":embedContent"), "gemini-3.8-flash"), + (_VERTEX_CLAUDE_RAW_PREDICT_TARGET, "vertex_ai/claude-opus-5"), + (_VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict"), "vertex_ai/claude-opus-5"), + ( + _VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict?alt=sse"), + "claude-opus-5", + ), + (_VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":rawPredict?alt=sse"), "claude-opus-5"), + (f"{_VERTEX_FLASH_TARGET}?tag=rawPredict", "gemini-3.8-flash"), + (_VERTEX_SEARCH_TARGET, "vertex_ai/search_api"), + ("https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs", None), + ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs" + "?filter=:search", + None, + ), + ], +) +def test_vertex_failure_model_name_matches_success_handler_naming(url_route: str, expected_model: str | None): + assert VertexPassthroughLoggingHandler.model_from_url_route(url_route) == expected_model + + @pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"]) def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_body(client_metadata_key: str): """A provider route that resolved a router deployment stashes its model_info on request.state. That @@ -8598,7 +9278,9 @@ def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_bo the call to (LIT-1761: passthrough successes carried model_id="").""" mock_request = MagicMock(spec=Request) mock_request.method = "POST" - mock_request.url = httpx.URL("http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent") + mock_request.url = httpx.URL( + "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent" + ) mock_request.headers = Headers({}) mock_request.scope = {} mock_request.state = SimpleNamespace( @@ -9668,12 +10350,14 @@ async def _drain_logging_worker(): await GLOBAL_LOGGING_WORKER.stop() yield + @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate") def test_update_pass_through_route_updates_registry(): """ @@ -9763,9 +10447,7 @@ def mock_request(): return self._dict.values() class MockRequest: - def __init__( - self, headers=None, method="POST", request_body: Optional[dict] = None - ): + def __init__(self, headers=None, method="POST", request_body: Optional[dict] = None): self.headers = headers or {} self.query_params = QueryParams() self.method = method @@ -9873,9 +10555,9 @@ def test_pass_through_routes_support_all_methods(): path = route.path methods = set(route.methods) allowed = PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES.get(path, expected_methods) - assert ( - methods == allowed - ), f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}" + assert methods == allowed, ( + f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}" + ) check_router_methods(llm_router) @@ -9905,19 +10587,11 @@ def test_is_bedrock_agent_runtime_route(): # Test agent runtime endpoints (should return True) assert _is_bedrock_agent_runtime_route("/knowledgebases/kb-123/retrieve") is True - assert ( - _is_bedrock_agent_runtime_route("/agents/knowledgebases/kb-123/retrieve") - is True - ) + assert _is_bedrock_agent_runtime_route("/agents/knowledgebases/kb-123/retrieve") is True # Test regular bedrock runtime endpoints (should return False) - assert ( - _is_bedrock_agent_runtime_route("/guardrail/test-id/version/1/apply") is False - ) - assert ( - _is_bedrock_agent_runtime_route("/model/cohere.command-r-v1:0/converse") - is False - ) + assert _is_bedrock_agent_runtime_route("/guardrail/test-id/version/1/apply") is False + assert _is_bedrock_agent_runtime_route("/model/cohere.command-r-v1:0/converse") is False assert _is_bedrock_agent_runtime_route("/some/random/endpoint") is False @@ -9984,9 +10658,7 @@ def test_custom_pricing_used_in_cost_calculation(): # Verify custom cost is different from standard cost (unless prices happen to match) # This confirms custom pricing is actually being applied - assert ( - custom_cost != standard_cost - ), "Custom pricing should produce different cost than standard pricing" + assert custom_cost != standard_cost, "Custom pricing should produce different cost than standard pricing" # Test 3: Custom pricing with cache_read_input_token_cost and input_cost_per_token_batches # This specifically tests the parameters that were causing the original issue @@ -10562,6 +11234,7 @@ def test_init_kwargs_client_metadata_cannot_mark_a_pass_through_call_as_a_silent assert "is_silent_experiment" not in json.loads(row["metadata"]) assert result["litellm_params"]["metadata"]["client_key"] == "kept" + @pytest.mark.asyncio async def test_direct_passthrough_keeps_mapped_baggage_tags_outside_upstream_body(): from litellm.proxy.common_utils.http_parsing_utils import read_request_body diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py index 40403de846f..ef19482b30a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py @@ -73,6 +73,25 @@ class TestPassthroughGuardrailHandlerGetGuardrailNames: assert set(result) == {"pii-detection", "content-moderation"} +class TestPassthroughGuardrailHandlerCollectGuardrails: + def test_inherits_a_key_guardrail_with_its_settings(self, monkeypatch: pytest.MonkeyPatch): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + key = UserAPIKeyAuth( + api_key="hashed", metadata={"guardrails": [{"pii": {"extra_body": {"threshold": 1}}}, "toxic"]} + ) + + guardrails = PassthroughGuardrailHandler.collect_guardrails(key, ["passthrough-guard", "toxic"]) + + assert guardrails == ["passthrough-guard", "toxic", {"pii": {"extra_body": {"threshold": 1}}}] + request_data = {"metadata": {"guardrails": guardrails}} + assert CustomGuardrail(guardrail_name="pii").get_guardrail_dynamic_request_body_params(request_data) == { + "threshold": 1 + } + + class TestPassthroughGuardrailHandlerNormalizeConfig: """Tests for PassthroughGuardrailHandler.normalize_config method.""" diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index ae95c27950c..8127ffe558b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -41,12 +41,13 @@ def _make_streaming_response(chunks): return mock -def _unarmed_logging_obj(): +def _unarmed_logging_obj(model: str = "unknown"): """Real Logging objects only carry _on_deferred_stream_complete when the proxy arms deferred dispatch; a bare MagicMock's auto-attribute is truthy and would spuriously trigger the deferral branch.""" obj = MagicMock() obj._on_deferred_stream_complete = None + obj.model = model return obj @@ -409,7 +410,7 @@ async def _collect_openai_passthrough_chunks(chunks, endpoint_type): async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "gpt-4o-mini", "stream": True}, - litellm_logging_obj=_unarmed_logging_obj(), + litellm_logging_obj=_unarmed_logging_obj(model="gpt-4o-mini"), endpoint_type=endpoint_type, start_time=datetime(2026, 1, 1), passthrough_success_handler_obj=MagicMock(), diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index c721e808ed5..cf00545f331 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -346,6 +346,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_enabled(): # Setup logging object with model info litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.model = "claude-sonnet-4@20250514" litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None @@ -544,6 +545,7 @@ async def test_vertex_ai_anthropic_streaming_model_extraction(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.model = "claude-sonnet-4@20250514" litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {} litellm_logging_obj.completion_start_time = None diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index da855753b38..27e8dfb7a2c 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -18,6 +18,7 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite from litellm.proxy._types import UserAPIKeyAuth @@ -34,7 +35,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) -from litellm.types.utils import CallTypesLiteral +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs try: from fastapi.exceptions import HTTPException @@ -1161,6 +1162,7 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch): data=data, user_api_key_dict=MagicMock(), call_type="completion", + metadata_key="metadata", ) assert outcome == "pass" @@ -1274,7 +1276,7 @@ def _assert_passed_with_discard_warning(result, caplog, reason): assert len(discard_warnings) == 1 assert reason in discard_warnings[0] assert "text rewrites included" in discard_warnings[0] - assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] @pytest.mark.asyncio @@ -1437,6 +1439,60 @@ async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_he assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] +class _SelfRecordingGuardrail(_TextReturningGuardrail): + records_own_guardrail_information = True + + +@pytest.mark.asyncio +async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", [_SelfRecordingGuardrail(["hello world"])]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + + +class _LitellmMetadataOpeningGuardrail(CustomGuardrail): + """Opens ``litellm_metadata`` on the request while scanning, as the A2A translation does.""" + + def __init__(self) -> None: + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + request_data["litellm_metadata"] = {"user_api_key_hash": "hashed"} + return inputs + + +@pytest.mark.asyncio +async def test_streaming_step_is_named_in_the_metadata_the_request_started_with( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_LitellmMetadataOpeningGuardrail()]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert "applied_guardrails" not in result.modified_data["litellm_metadata"] + + +@pytest.mark.asyncio +async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert [(step.outcome, step.action_taken) for step in result.step_results] == [("fail", "next")] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + + @pytest.mark.asyncio async def test_streaming_step_records_guardrail_information_once_on_block(monkeypatch): monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) diff --git a/tests/unit/proxy/proxy_server/test_background_health.py b/tests/unit/proxy/proxy_server/test_background_health.py index ebc33b1c788..4c1c3d4ece9 100644 --- a/tests/unit/proxy/proxy_server/test_background_health.py +++ b/tests/unit/proxy/proxy_server/test_background_health.py @@ -432,18 +432,6 @@ def test_write_health_state_to_router_cache_sets_states(monkeypatch): monkeypatch.setattr(hc, "build_deployment_health_states", lambda **_kw: fake_states) - import litellm.router_utils.cooldown_handlers as cd - - monkeypatch.setattr(cd, "set_cooldown_deployments", lambda **_kw: None) - - import litellm.router_utils.router_callbacks.track_deployment_metrics as tdm - - monkeypatch.setattr( - tdm, - "increment_deployment_failures_for_current_minute", - lambda **_kw: None, - ) - healthy = [{"model_id": "m1"}] unhealthy = [{"model_id": "m2"}] exceptions = {"m2": SimpleNamespace(status_code=500)} @@ -451,6 +439,7 @@ def test_write_health_state_to_router_cache_sets_states(monkeypatch): _write_health_state_to_router_cache(healthy, unhealthy, exceptions) health_checks.set_deployment_health_states.assert_called_once_with(fake_states) + fake_router.health.return_value.record_failure.assert_called_once_with("m2", exceptions["m2"], 500) call_args = health_checks.set_deployment_health_states.call_args[0][0] assert normalize( @@ -503,26 +492,6 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp monkeypatch.setattr(hc, "build_deployment_health_states", lambda **_kw: fake_states) - cooldowns: list[str] = [] - - import litellm.router_utils.cooldown_handlers as cd - - monkeypatch.setattr( - cd, - "set_cooldown_deployments", - lambda **kw: cooldowns.append(kw.get("deployment")), - ) - - failures: list[str] = [] - - import litellm.router_utils.router_callbacks.track_deployment_metrics as tdm - - monkeypatch.setattr( - tdm, - "increment_deployment_failures_for_current_minute", - lambda **kw: failures.append(kw.get("deployment_id")), - ) - _write_health_state_to_router_cache( [{"model_id": "m1"}], [{"model_id": "m2"}], @@ -530,8 +499,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp ) fake_router.health.return_value.health_checks.set_deployment_health_states.assert_called_once_with(fake_states) - assert cooldowns == [] - assert failures == [] + fake_router.health.return_value.record_failure.assert_not_called() def test_write_health_state_to_router_cache_swallows_internal_failures(monkeypatch): diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index a8706c5bccd..db5be07fb6f 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -23,7 +23,8 @@ import json import logging import os import subprocess -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -35,7 +36,6 @@ from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( - _proxy_startup_event, ProxyStartupEvent, _resolve_pydantic_type, _resolve_typed_dict_type, @@ -51,9 +51,19 @@ from litellm.proxy.proxy_server import ( run_ollama_serve, save_worker_config, ) +from litellm.proxy.utils import PrismaClient from .conftest import normalize + +@asynccontextmanager +async def _closing_prisma(prisma: MagicMock) -> AsyncIterator[None]: + try: + yield + finally: + await PrismaClient.aclose(prisma) + + # --------------------------------------------------------------------------- # cleanup_router_config_variables # --------------------------------------------------------------------------- @@ -101,9 +111,10 @@ def test_cleanup_router_config_variables_fails_on_unknown_attr_raises(): @pytest.mark.asyncio -async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): +async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeypatch): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock() + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) monkeypatch.setattr(ps, "master_key", "sk-x", raising=False) @@ -117,7 +128,9 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with _closing_prisma(fake_prisma): + await proxy_shutdown_event() + fake_prisma.disconnect.assert_not_awaited() observed = { "disconnect_called": fake_prisma.disconnect.await_count == 1, @@ -204,6 +217,7 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) async def _record_flush(client, accumulator): @@ -222,7 +236,8 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with _closing_prisma(fake_prisma): + await proxy_shutdown_event() assert calls == ["flush", "disconnect"] @@ -239,6 +254,7 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) async def _record_flush(client, accumulator): @@ -262,7 +278,8 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor monkeypatch.setattr(litellm, "success_callback", [], raising=False) track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts())) - await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + async with _closing_prisma(fake_prisma): + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) assert calls == ["flush", "audit", "disconnect"] @@ -318,9 +335,10 @@ async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_t @pytest.mark.asyncio -async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): +async def test_database_lifespan_logs_disconnect_error_after_shutdown_instead_of_raising(monkeypatch, caplog): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone")) + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) fake_jwt = MagicMock() @@ -332,8 +350,12 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - with pytest.raises(RuntimeError, match="db gone"): - await proxy_shutdown_event() + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + async with _closing_prisma(fake_prisma): + await proxy_shutdown_event() + + fake_prisma.disconnect.assert_awaited_once() + assert "db gone" in caplog.text # --------------------------------------------------------------------------- @@ -1008,7 +1030,9 @@ async def test_proxy_shutdown_commits_spend_before_waiting_on_a_slow_logging_tas task_may_finish: Final = asyncio.Event() task_finished: Final = asyncio.Event() - flushes: Final[list[tuple[str, bool]]] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion + flushes: Final[ + list[tuple[str, bool]] + ] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion async def slow_logging_task() -> None: await task_may_finish.wait() @@ -1060,7 +1084,7 @@ def test_otel_global_provider_published_after_callback_init(): span was orphaned. The publish (``publish_global_otel_v2_provider``) must therefore appear after ``_initialize_startup_logging`` in the lifespan source. """ - wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(") publish_pos = source.find("publish_global_otel_v2_provider(") @@ -1146,7 +1170,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): The call must follow Prisma setup so DB-backed deployments do not false-positive. Direct ``_warn_budget_without_db`` tests cover the warning behavior itself. """ - wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:") warn_pos = source.find("warn_budget_without_db(") diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 49f9e99cf48..d9ff7fef9f9 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -2263,7 +2263,6 @@ def test_ProxyConfig_merge_config_and_db_search_tools_prefers_db_duplicate(): @pytest.mark.asyncio async def test_ProxyConfig__init_search_tools_in_db_loads_merged_tools(monkeypatch): from litellm.proxy import proxy_server - from litellm.router_utils.search_api_router import SearchAPIRouter pc = ProxyConfig() pc.update_config_state( @@ -2289,28 +2288,24 @@ async def test_ProxyConfig__init_search_tools_in_db_loads_merged_tools(monkeypat }, } ] - fake_router = MagicMock() + router = litellm.Router() mock_get_db_tools = AsyncMock(return_value=db_tools) - mock_update_router = AsyncMock() - monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr( "litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db", mock_get_db_tools, ) - monkeypatch.setattr(SearchAPIRouter, "update_router_search_tools", mock_update_router) await pc._init_search_tools_in_db(prisma_client=MagicMock()) mock_get_db_tools.assert_awaited_once() - mock_update_router.assert_awaited_once() - update_kwargs = mock_update_router.await_args.kwargs - assert update_kwargs["router_instance"] is fake_router - assert [tool["search_tool_name"] for tool in update_kwargs["search_tools"]] == [ + loaded_tools = router.settings().search_tools + assert [tool["search_tool_name"] for tool in loaded_tools] == [ "config-only", "shared-search", ] - assert update_kwargs["search_tools"][1]["litellm_params"]["api_key"] == "fake-db-key" + assert loaded_tools[1]["litellm_params"]["api_key"] == "fake-db-key" @pytest.mark.asyncio @@ -2320,11 +2315,10 @@ async def test_ProxyConfig__init_search_tools_in_db_clears_router_when_last_tool pc = ProxyConfig() pc.update_config_state({}) - fake_router = MagicMock() - fake_router.search_tools = [{"search_tool_name": "deleted-search", "litellm_params": {}}] + router = litellm.Router(search_tools=[{"search_tool_name": "deleted-search", "litellm_params": {}}]) mock_get_db_tools = AsyncMock(return_value=[]) - monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr( "litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db", mock_get_db_tools, @@ -2333,7 +2327,7 @@ async def test_ProxyConfig__init_search_tools_in_db_clears_router_when_last_tool await pc._init_search_tools_in_db(prisma_client=MagicMock()) mock_get_db_tools.assert_awaited_once() - assert fake_router.search_tools == [] + assert router.settings().search_tools == [] @pytest.mark.asyncio @@ -2347,8 +2341,7 @@ async def test_ProxyConfig__init_search_tools_in_db_keeps_loaded_tools_whose_par "search_tool_name": "rotated-search", "litellm_params": {"search_provider": "perplexity", "api_key": "pplx-loaded"}, } - fake_router = MagicMock() - fake_router.settings.return_value = RouterSettings( + router = litellm.Router( search_tools=[ loaded_tool, { @@ -2378,7 +2371,7 @@ async def test_ProxyConfig__init_search_tools_in_db_keeps_loaded_tools_whose_par "litellm_params": {"search_provider": "Tavily", "api_key": "tvly-edited"}, }, ] - monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr( "litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db", AsyncMock(return_value=db_tools), @@ -2386,7 +2379,7 @@ async def test_ProxyConfig__init_search_tools_in_db_keeps_loaded_tools_whose_par await pc._init_search_tools_in_db(prisma_client=MagicMock()) - assert [tool["litellm_params"] for tool in fake_router.search_tools] == [ + assert [tool["litellm_params"] for tool in router.settings().search_tools] == [ {"search_provider": "perplexity", "api_key": "pplx-loaded"}, {"search_provider": "tavily", "api_key": "tvly-fresh"}, {"search_provider": "Tavily", "api_key": "tvly-edited"}, diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index 86364d93ed0..bffa4612964 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -550,15 +550,12 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke deployment = { "model_name": "*", - "litellm_params": {"model": "openai/*"}, + "litellm_params": {"model": "openai/*", "api_key": "sk-star"}, } - router = route_entry_points_to_stubs(MagicMock()) - router.get_model_access_groups = MagicMock(return_value={}) - router.get_model_names = MagicMock(return_value=["*"]) - router.get_model_list = MagicMock(return_value=[deployment]) + router = litellm.Router(model_list=[deployment]) monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) - expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment]) + expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info(router.get_model_list(), router) allowed_model_names = await proxy_server._get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=UserAPIKeyAuth( api_key="sk-test", @@ -575,7 +572,8 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke allowed_model_names=allowed_model_names, ) - assert [model["model_name"] for model in result] == ["openai/gpt-4o"] + assert [(model["model_name"], model["litellm_params"]["model"]) for model in result] == [("gpt-4o", "openai/gpt-4o")] + assert router.registry().names("gpt-4o", kind="wildcard") == frozenset({"openai/gpt-4o"}) # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index ffd23fa3a6f..874f471f3e7 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -9,10 +9,16 @@ Pins (PR2): from __future__ import annotations +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field from typing import Final from unittest.mock import MagicMock +import httpx import pytest +import respx import litellm from litellm import Router @@ -541,6 +547,73 @@ def test_codex_service_tiers_resolved_via_internal_team_key(client, auth_as, pat assert (entry["slug"], [tier["id"] for tier in entry["service_tiers"]]) == ("gpt-4-team", ["ultrafast"]) +OPENROUTER_BASE: Final = "https://openrouter.ai/api/v1" + + +@dataclass(frozen=True, slots=True) +class _CatalogHeldUntilTheLoopRuns: + requested: threading.Event = field(default_factory=threading.Event) + released: threading.Event = field(default_factory=threading.Event) + + def answer(self, _request: httpx.Request) -> httpx.Response: + self.requested.set() + if not self.released.wait(timeout=30): + return httpx.Response(504) + return httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + + async def release_from_the_loop(self) -> None: + assert await asyncio.to_thread(self.requested.wait, 30) + self.released.set() + + +@pytest.mark.parametrize("scope", [None, "expand"]) +def test_v1_models_fetches_a_live_openrouter_catalog_without_holding_the_event_loop_or_its_default_thread_pool( + monkeypatch: pytest.MonkeyPatch, scope: str | None +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": f"sk-or-v1-models-off-loop-{scope}", + "api_base": OPENROUTER_BASE, + }, + } + ] + ), + ) + catalog: Final = _CatalogHeldUntilTheLoopRuns() + + async def listed_ids_on_a_one_thread_default_pool() -> list[str]: + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + listing: Final = asyncio.ensure_future( + proxy_server.model_list( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ), + scope=scope, + ) + ) + await catalog.release_from_the_loop() + response: Final = await listing + return [str(model["id"]) for model in response["data"]] + + with respx.mock() as respx_mock: + catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer) + listed_ids: Final = asyncio.run(listed_ids_on_a_one_thread_default_pool()) + + assert catalog_route.calls.last.response.status_code == 200 + assert listed_ids == ["openrouter/openai/gpt-5.4-mini"] + + def _priced_jev_models() -> frozenset[str]: return frozenset(name.removeprefix("typesafe/") for name in litellm.model_cost if name.startswith("typesafe/")) diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 48672c7a305..72755633376 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -28,6 +28,28 @@ from litellm.types.proxy.public_endpoints.public_endpoints import ProviderCreate from litellm.types.utils import LlmProviders +def test_meta_model_api_fields_are_separate_from_meta_llama() -> None: + app: Final = FastAPI() + app.include_router(router) + response: Final = TestClient(app).get("/public/providers/fields") + assert response.status_code == 200 + providers: Final = response.json() + meta: Final = next((p for p in providers if p["litellm_provider"] == "meta"), None) + llama: Final = next(p for p in providers if p["litellm_provider"] == "meta_llama") + assert meta is not None + assert meta["provider"] == "META" + assert meta["provider_display_name"] == "Meta Model API" + assert meta["default_model_placeholder"] == "meta/muse-spark-1.1" + fields: Final = {field["key"]: field for field in meta["credential_fields"]} + assert set(fields) == {"api_base", "api_key"} + assert fields["api_base"]["field_type"] == "text" + assert fields["api_key"]["field_type"] == "password" + assert fields["api_base"]["required"] is False + assert fields["api_key"]["required"] is False + assert llama["provider"] == "LLAMA" + assert llama["provider_display_name"] == "Meta Llama" + + def test_bigmodel_fields_require_private_key_and_share_default_base_with_adapter() -> None: from litellm.llms.bigmodel.common_utils import BIGMODEL_API_BASE @@ -386,6 +408,34 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_zai_provider_fields(): + from litellm.llms.zai.chat.transformation import ZAI_API_BASE + + app_instance: Final = FastAPI() + app_instance.include_router(router) + test_client: Final = TestClient(app_instance) + + response: Final = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers: Final = response.json() + + zai: Final = next((p for p in providers if p["provider"] == "ZAI"), None) + assert zai is not None, "Z.AI provider entry not found" + + assert zai["provider_display_name"] == "Z.AI (Zhipu AI)" + assert zai["litellm_provider"] == "zai" + assert zai["default_model_placeholder"].startswith("zai/") + + fields_by_key: Final = {f["key"]: f for f in zai["credential_fields"]} + + assert fields_by_key["api_key"]["required"] is True + assert fields_by_key["api_key"]["field_type"] == "password" + + assert fields_by_key["api_base"]["required"] is False + assert fields_by_key["api_base"]["field_type"] == "text" + assert fields_by_key["api_base"]["placeholder"] == ZAI_API_BASE + + def test_zerogpu_add_model_form_defaults_resolve_to_a_priced_zerogpu_model(): import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -568,7 +618,6 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( "libertai", "litellm_agent", "manus", - "meta", "modelscope", "mongodb", "nano-gpt", @@ -590,7 +639,6 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( "transcribe", "valkey", "xiaomi_mimo", - "zai", } ) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 1e4dded801e..07f85cb9628 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -7,6 +7,7 @@ import contextlib import dataclasses import datetime import json +import re import unittest from collections.abc import Callable, Mapping from typing import Any, Final, Literal, cast @@ -22,12 +23,17 @@ from fastapi.testclient import TestClient from httpx import Response import litellm +from litellm.caching.caching import DualCache from litellm.constants import EMPTY_MAPPING -from litellm.proxy._types import ProxyException +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.proxy_server import app +from litellm.proxy.utils import ProxyLogging from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors +from litellm.types.utils import CallTypesLiteral from tests._master_key import MASTER_KEY @@ -2213,6 +2219,59 @@ class TestResponsesWSFirstFrameValidation: ws.close.assert_awaited_once_with(code=1011, reason="Internal server error") +class _FrameGuardrail(CustomGuardrail): + def __init__( + self, + name: str, + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + rewrites: bool = False, + blocks: bool = False, + ): + super().__init__(guardrail_name=name, event_hook=event_hook, default_on=False) + self.inputs: list[object] = [] + self.configs: list[object] = [] + self.rewrites = rewrites + self.blocks = blocks + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + return self._scan(data) + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + data["input"] = self._scan(data)["input"] + + def _scan(self, data: dict[str, object]) -> dict[str, object]: + self.inputs.append(data["input"]) + self.configs.append(cast(dict[str, object], data["litellm_metadata"]).get("guardrail_config")) + if self.blocks: + raise HTTPException( + status_code=400, + detail={"error": f"Content blocked by {self.guardrail_name}", "guardrail_name": self.guardrail_name}, + ) + return {**data, "input": f"{data['input']} [{self.guardrail_name}]"} if self.rewrites else data + + +@contextlib.contextmanager +def _guardrails_on(*guardrails: CustomGuardrail): + with patch.object(litellm, "callbacks", list(guardrails)): + ProxyLogging._callback_capabilities_cache.clear() + try: + yield ProxyLogging(user_api_key_cache=DualCache()) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + +def _key() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="hashed") + + class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio async def test_endpoint_enforces_auth_after_model_from_first_frame(self): @@ -2344,7 +2403,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_name": "turn"}] return {**data, "input": data["input"] + " [guarded]"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) admitted = { "model": "gpt-4o-mini", @@ -2375,7 +2434,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] raise ValueError("blocked") - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock() spend_logger = MagicMock(spec=ProxyDBLogger) @@ -2413,7 +2472,7 @@ class TestResponsesWSFirstFrameModelAuth: async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): return {**data, "input": data["input"] + " [guarded]", "tools": data["tools"][:1]} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2432,10 +2491,10 @@ class TestResponsesWSFirstFrameModelAuth: async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( - {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() ) metadata = { "user_api_key_metadata": {"disable_global_guardrails": True}, @@ -2454,17 +2513,18 @@ class TestResponsesWSFirstFrameModelAuth: [ {"type": "response.create", "input": "forbidden", "guardrails": []}, {"type": "response.create", "input": "forbidden", "metadata": {"guardrails": []}}, + {"type": "response.create", "input": "forbidden", "litellm_metadata": {"guardrails": []}}, {"type": "response.create", "response": {"input": "forbidden", "guardrails": [], "policies": []}}, ], ) async def test_frame_guard_keeps_the_admitted_guardrail_list_when_a_frame_tries_to_replace_it(self, frame): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, - MagicMock(), + _key(), proxy_logging_obj, MagicMock(), ) @@ -2476,13 +2536,258 @@ class TestResponsesWSFirstFrameModelAuth: assert "guardrails" not in (scanned.get("metadata") or {}) assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "selected", + [ + {"guardrails": ["frame-pick"]}, + {"guardrails": [{"frame-pick": {"extra_body": {"threshold": 0.8}}}]}, + {"metadata": {"guardrails": ["frame-pick"], "tag": "kept"}}, + {"litellm_metadata": {"guardrails": ["frame-pick"]}}, + {"response": {"guardrails": ["frame-pick"]}}, + ], + ids=["top-level", "with-extra-body", "metadata", "litellm-metadata", "nested"], + ) + async def test_frame_guard_runs_the_guardrails_a_frame_selects_on_top_of_the_admitted_ones(self, selected: dict[str, object]): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick", rewrites=True) + unselected = _FrameGuardrail("unselected") + admitted = {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}} + nested = "response" in selected + body = {"input": "hi", "store": False, **(selected["response"] if nested else selected)} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with _guardrails_on(key_required, frame_pick, unselected) as proxy_logging_obj: + guard = _responses_ws_frame_guard(admitted, _key(), proxy_logging_obj, MagicMock()) + selected_turn = await guard(json.dumps(frame)) + plain_turn = await guard(json.dumps({"type": "response.create", "input": "again"})) + + kept_metadata = {"metadata": {"tag": "kept"}} if "metadata" in selected else {} + forwarded = {"input": "hi [frame-pick]", "store": False, **kept_metadata} + assert json.loads(selected_turn.frame) == ( + {"type": "response.create", "response": forwarded} if nested else {"type": "response.create", **forwarded} + ) + assert json.loads(plain_turn.frame) == {"type": "response.create", "input": "again"} + assert key_required.inputs == ["hi", "again"] + assert frame_pick.inputs == ["hi"] + assert unselected.inputs == [] + assert admitted["litellm_metadata"] == {"guardrails": ["key-required"]} + + @pytest.mark.asyncio + async def test_frame_guard_forwards_string_metadata_named_like_a_pick_without_running_it(self): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick") + metadata = {"guardrails": "frame-pick", "guardrail_config": "strict", "trace": "t1"} + + async def locked_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": False}}) + + with _guardrails_on(key_required, frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + team_object=locked_team, + ) + guarded = await guard(json.dumps({"type": "response.create", "input": "hi", "metadata": metadata})) + + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": metadata} + assert key_required.inputs == ["hi"] + assert frame_pick.inputs == [] + + @pytest.mark.asyncio + async def test_frame_guard_hands_a_selected_guardrail_its_config_with_the_top_level_winning(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + frame = { + "type": "response.create", + "input": "hi", + "metadata": {"guardrail_config": {"frame-pick": {"threshold": 1}, "metadata-only": {"threshold": 3}}}, + "guardrails": ["frame-pick"], + "guardrail_config": {"frame-pick": {"threshold": 2}}, + } + + with _guardrails_on(frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrail_config": {"key-required": {"threshold": 0}}}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard(json.dumps(frame)) + + assert frame_pick.configs == [ + {"key-required": {"threshold": 0}, "frame-pick": {"threshold": 2}, "metadata-only": {"threshold": 3}} + ] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": {}} + + @pytest.mark.asyncio + async def test_frame_guard_runs_streaming_scoped_guardrails_and_skips_non_streaming_ones_on_a_socket_turn(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + streaming_only = _FrameGuardrail("streaming-only") + streaming_only.apply_stream_scope("streaming") + streaming_during_call = _FrameGuardrail("streaming-during-call", event_hook=GuardrailEventHooks.during_call) + streaming_during_call.apply_stream_scope("streaming") + non_streaming_only = _FrameGuardrail("non-streaming-only") + non_streaming_only.apply_stream_scope("non_streaming") + frame = { + "type": "response.create", + "input": "hi", + "guardrails": ["streaming-only", "streaming-during-call", "non-streaming-only"], + } + + with _guardrails_on(streaming_only, streaming_during_call, non_streaming_only) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + guarded = await guard(json.dumps(frame)) + + assert streaming_only.inputs == ["hi"] + assert streaming_during_call.inputs == ["hi"] + assert non_streaming_only.inputs == [] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + + @pytest.mark.asyncio + async def test_frame_guard_blocks_a_turn_with_the_message_a_selected_during_call_guardrail_raises(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + during_pick = _FrameGuardrail("during-pick", event_hook=GuardrailEventHooks.during_call, blocks=True) + + with _guardrails_on(during_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + unselected_turn = await guard(json.dumps({"type": "response.create", "input": "plain"})) + with pytest.raises(ValueError, match=r"^Content blocked by during-pick$"): + await guard(json.dumps({"type": "response.create", "input": "hi", "guardrails": ["during-pick"]})) + + assert json.loads(unselected_turn.frame) == {"type": "response.create", "input": "plain"} + assert during_pick.inputs == ["hi"] + + @pytest.mark.asyncio + async def test_frame_guard_runs_only_the_moderation_hooks_that_enforce_request_content(self): + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + class _ModerationHook(CustomLogger): + def __init__(self, enforces_request_content: bool) -> None: + super().__init__() + self.enforces_request_content = enforces_request_content + self.inputs: list[object] = [] + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + self.inputs.append(data.get("input")) + + vendor_moderation = _ModerationHook(enforces_request_content=False) + content_enforcer = _ModerationHook(enforces_request_content=True) + + with _guardrails_on(vendor_moderation, content_enforcer) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + assert vendor_moderation.inputs == [] + assert content_enforcer.inputs == ["hi"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("selection", "message"), + [ + ( + {"guardrails": "frame-pick"}, + "response.create guardrails: 'str' instances are not allowed as a Sequence value", + ), + ({"guardrails": [123]}, "response.create guardrails.0.str: Input should be a valid string"), + ({"guardrail_config": []}, "response.create guardrail_config: Input should be a valid dictionary"), + ( + {"policies": "frame-policy"}, + "response.create policies: 'str' instances are not allowed as a Sequence value", + ), + ({"policies": [1]}, "response.create policies.0: Input should be a valid string"), + ( + {"litellm_metadata": {"guardrails": "frame-pick"}}, + "response.create litellm_metadata.guardrails: 'str' instances are not allowed as a Sequence value", + ), + ( + {"litellm_metadata": {"guardrail_config": []}}, + "response.create litellm_metadata.guardrail_config: Input should be a valid dictionary", + ), + ], + ) + async def test_frame_guard_refuses_and_bills_a_frame_whose_guardrail_selection_is_malformed( + self, selection: dict[str, object], message: str + ): + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + + with _guardrails_on(key_required, spend_logger) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + await guard(json.dumps({"type": "response.create", "input": "hi", **selection})) + + assert key_required.inputs == [] + spend_logger.async_post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_frame_guard_runs_the_guardrails_of_the_policies_a_frame_selects(self): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + key_required = _FrameGuardrail("key-required") + policy_pick = _FrameGuardrail("policy-pick") + policy_registry = get_policy_registry() + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + try: + policy_registry.add_policy("frame-policy", Policy(guardrails=PolicyGuardrails(add=["policy-pick"]))) + with _guardrails_on(key_required, policy_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard( + json.dumps({"type": "response.create", "input": "hi", "policies": ["frame-policy"]}) + ) + finally: + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + + assert key_required.inputs == ["hi"] + assert policy_pick.inputs == ["hi"] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + @pytest.mark.asyncio async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self): import threading from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) span = threading.RLock() guard = _responses_ws_frame_guard( @@ -2510,7 +2815,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] return {**data, "model": "onprem-llama"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock() @@ -2543,7 +2848,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] raise ValueError("blocked") - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock() @@ -2571,7 +2876,7 @@ class TestResponsesWSFirstFrameModelAuth: async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): return {key: value for key, value in data.items() if key != "instructions"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2596,7 +2901,7 @@ class TestResponsesWSFirstFrameModelAuth: data["metadata"]["pii_tokens"] = {"": "alice@example.com"} return {**data, "input": "email please"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2617,19 +2922,279 @@ class TestResponsesWSFirstFrameModelAuth: assert dict(guarded.pii_tokens) == {"": "alice@example.com"} @pytest.mark.asyncio - async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self): + @pytest.mark.parametrize("field", ["input", "instructions", "tools"]) + async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field: str): + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + ) + message = f"response.create sends {field} at the top level next to response; send it inside response" + + with ( + patch.object(litellm, "callbacks", [spend_logger]), + pytest.raises(ValueError, match=f"^{re.escape(message)}$"), + ): + await guard(json.dumps({"type": "response.create", field: "unguarded", "response": {"input": "hi"}})) + + proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert spend_logger.async_post_call_failure_hook.await_args.kwargs["request_data"]["input"] == "hi" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "top_level", + [ + {"guardrails": ["frame-pick"]}, + {"metadata": {"guardrails": ["frame-pick"], "tag": "kept"}}, + {"litellm_metadata": {"guardrails": ["frame-pick"]}}, + ], + ids=["top-level", "metadata", "litellm-metadata"], + ) + async def test_frame_guard_runs_the_guardrails_a_nested_frame_selects_beside_its_response(self, top_level: dict[str, object]): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick", rewrites=True) + unselected = _FrameGuardrail("unselected") + + with _guardrails_on(key_required, frame_pick, unselected) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard(json.dumps({"type": "response.create", **top_level, "response": {"input": "hi"}})) + + kept_metadata = {"metadata": {"tag": "kept"}} if "metadata" in top_level else {} + assert json.loads(guarded.frame) == { + "type": "response.create", + **kept_metadata, + "response": {"input": "hi [frame-pick]"}, + } + assert key_required.inputs == ["hi"] + assert frame_pick.inputs == ["hi"] + assert unselected.inputs == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}, + {"type": "response.create", "input": "hi", "metadata": {"guardrails": ["frame-pick"]}}, + {"type": "response.create", "input": "hi", "metadata": {"disable_global_guardrails": True}}, + {"type": "response.create", "input": "hi", "metadata": {"disable_global_guardrails": "true"}}, + {"type": "response.create", "response": {"input": "hi", "guardrails": ["frame-pick"]}}, + {"type": "response.create", "litellm_metadata": {"guardrails": ["frame-pick"]}, "response": {"input": "hi"}}, + {"type": "response.create", "input": "hi", "policies": ["frame-policy"]}, + {"type": "response.create", "response": {"input": "hi", "policies": ["frame-policy"]}}, + ], + ids=[ + "top-level", + "metadata", + "global-toggle", + "string-global-toggle", + "nested", + "beside-nested", + "policies", + "nested-policies", + ], + ) + async def test_frame_guard_refuses_guardrail_and_policy_picks_from_a_team_without_permission(self, frame: dict[str, object]): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick") + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + team_lookups: list[UserAPIKeyAuth] = [] + key = _key() + + async def locked_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + team_lookups.append(user_api_key_dict) + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": False}}) + + with _guardrails_on(key_required, frame_pick, spend_logger) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + key, + proxy_logging_obj, + MagicMock(), + team_object=locked_team, + ) + plain_turn = await guard(json.dumps({"type": "response.create", "input": "plain"})) + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(json.dumps(frame)) + + assert json.loads(plain_turn.frame) == {"type": "response.create", "input": "plain"} + assert key_required.inputs == ["plain"] + assert frame_pick.inputs == [] + assert team_lookups == [key] + spend_logger.async_post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_frame_guard_runs_a_frame_s_picks_for_a_team_allowed_to_modify_guardrails(self): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + + async def open_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": True}}) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, + _key(), + proxy_logging_obj, + MagicMock(), + team_object=open_team, + ) + await guard(json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]})) + + assert frame_pick.inputs == ["hi"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("modify_guardrails", "picked_inputs"), + [(True, ["hi"]), (False, [])], + ids=["open", "locked"], + ) + @pytest.mark.parametrize( + "team_lookup", + [ + lambda: patch( + "litellm.proxy.response_api_endpoints.endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist"})), + ), + lambda: patch("litellm.proxy.proxy_server.prisma_client", None), + ], + ids=["team-row-gone", "no-database"], + ) + async def test_frame_guard_reads_the_key_s_team_metadata_when_the_team_cannot_be_loaded( + self, + modify_guardrails: bool, + picked_inputs: list[str], + team_lookup: Callable[[], contextlib.AbstractContextManager[object]], + ): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + key = UserAPIKeyAuth( + api_key="hashed", + team_id="deleted-team", + team_metadata={"guardrails": {"modify_guardrails": modify_guardrails}}, + ) + frame = json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + + with _guardrails_on(frame_pick) as proxy_logging_obj, team_lookup(): + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, key, proxy_logging_obj, MagicMock() + ) + if modify_guardrails: + await guard(frame) + else: + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(frame) + + assert frame_pick.inputs == picked_inputs + + @pytest.mark.asyncio + async def test_frame_guard_forwards_the_pre_call_output_and_not_a_during_call_rewrite_like_http(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + pre_pick = _FrameGuardrail("pre-pick", rewrites=True) + during_pick = _FrameGuardrail("during-pick", event_hook=GuardrailEventHooks.during_call, rewrites=True) + + with _guardrails_on(pre_pick, during_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + guarded = await guard( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["pre-pick", "during-pick"]}) + ) + + assert during_pick.inputs == ["hi [pre-pick]"] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi [pre-pick]"} + + @pytest.mark.asyncio + async def test_frame_guard_forwards_a_nested_turn_without_top_level_fields_that_pick_no_guardrail(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() ) - with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): - await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + with patch.object(litellm, "callbacks", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "include_guardrail_response": True, + "guardrail_config": {"frame-pick": {"threshold": 1}}, + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + ) + ) - proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert json.loads(guarded.frame) == { + "type": "response.create", + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + + @pytest.mark.asyncio + async def test_a_blocked_socket_turn_answers_with_the_guardrail_fields_http_returns(self): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + frame_pick = _FrameGuardrail("frame-pick", blocks=True) + websocket = MagicMock(send_text=AsyncMock()) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + handler = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-haiku-5-5", + logging_obj=Logging( + model="anthropic/claude-haiku-5-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + frame_guard=_responses_ws_frame_guard( + {"model": "claude-managed", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ), + ) + await handler._process_response_create( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + ) + + assert json.loads(websocket.send_text.await_args.args[0]) == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "Content blocked by frame-pick", + "provider_specific_fields": { + "error": "Content blocked by frame-pick", + "guardrail_name": "frame-pick", + "guardrail_mode": "pre_call", + }, + }, + } @pytest.mark.asyncio async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self): @@ -2637,7 +3202,7 @@ class TestResponsesWSFirstFrameModelAuth: from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=ValueError("blocked")) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock(side_effect=RuntimeError("db down")) @@ -2679,7 +3244,7 @@ class TestResponsesWSFirstFrameModelAuth: async def fake_llm_call(): return failure if provider_rejected else None - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() @@ -2768,7 +3333,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() row_written = asyncio.Event() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=lambda **kwargs: row_written.set()) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) @@ -2858,7 +3423,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) @@ -2948,7 +3513,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) @@ -3022,7 +3587,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: { **data, @@ -3083,7 +3648,7 @@ class TestResponsesWSFirstFrameModelAuth: rejection = litellm.RateLimitError( message="origin deployment is cooling down", model="gpt-4o-mini", llm_provider="openai" ) - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 41acd487b67..6912049d225 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -75,7 +75,7 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: if kind == "deployment_names": return router.deployment_names if kind == "team_public": - return router.team_public_model_names + return router.team_public_model_names if team_id is None else router.team_public_model_names_of(team_id) if kind == "litellm_models": return router.resolved_litellm_models(model_name, **_given(team_id=team_id)) if kind == "route_candidates": @@ -143,6 +143,10 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: registry.deployment_dict.side_effect = lambda id: router.get_model_info(id=id) registry.has_deployment.side_effect = lambda id: router.has_model_id(id) registry.discovered_model_info.side_effect = lambda id: router.get_discovered_model_info(id) + registry.routing_group.side_effect = lambda name: router.get_routing_group(name) + registry.client.side_effect = lambda deployment, *, client_type=None: router._get_client( + deployment=deployment, kwargs={}, client_type=client_type + ) def _credentials(id, *, form="params", **team): # Old call sites passed team_id (even None) exactly when the new ones do, so pass it on as given. @@ -159,6 +163,9 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: registry.plugins.side_effect = lambda kind: getattr(router, f"{kind}_routers") health.fully_blocked_model_names.side_effect = lambda: router.get_fully_blocked_model_names() health.split_priority.side_effect = lambda model, kwargs: router.split_request_priority(model=model, kwargs=kwargs) + health.priority_header_fields.side_effect = lambda model, kwargs, header: router.priority_header_update( + model=model, kwargs=kwargs, header=header + ) registry.index.side_effect = lambda name: getattr(router, f"{name}_to_deployment_indices") registry.patterns.side_effect = lambda team=False: router.team_pattern_routers if team else router.pattern_router health.is_blocked.side_effect = lambda deployment: router.is_deployment_blocked(deployment) diff --git a/tests/unit/proxy/spend_tracking/test_savings.py b/tests/unit/proxy/spend_tracking/test_savings.py index f9566fd1ba0..8503a124b01 100644 --- a/tests/unit/proxy/spend_tracking/test_savings.py +++ b/tests/unit/proxy/spend_tracking/test_savings.py @@ -137,6 +137,34 @@ def test_compression_savings_priced_at_input_rate(): assert result.prompt_caching == 0.0 +def test_an_overflowing_deployment_price_records_finite_savings() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "overflowing-sonnet", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "input_cost_per_token": 1e308, + "cache_read_input_token_cost": 1e-07, + }, + }, + ] + ) + deployment_id: Final = router.get_model_list(model_name="overflowing-sonnet")[0]["model_info"]["id"] + + result: Final = compute_savings_spend( + model="claude-sonnet-5-5", + custom_llm_provider="anthropic", + compression_saved_tokens=4389, + gateway_injected_cache=True, + usage_object=_caching_usage(read=1000, written=0), + model_id=deployment_id, + llm_router=lambda: router, + ) + + assert result == (0.0, 0.0, 0.0, 0.0) + + def test_prompt_caching_savings_priced_at_input_minus_cache_read(): input_cost, cache_read_cost = _anthropic_costs("claude-sonnet-5") # A model that supports prompt caching must charge less to read from cache; diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 1ea445fda0d..a91cdd512fe 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2394,6 +2394,65 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): ) +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_carries_upstream_provider_from_hidden_params_regression() -> None: + standard_logging_payload: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + payload_with_provider: Final = { + **dict(standard_logging_payload), + "hidden_params": { + **dict(standard_logging_payload["hidden_params"]), + "upstream_provider": "serving-site-a", + }, + } + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": payload_with_provider, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + }, + }, + start_time=datetime.datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc), + ) + + assert json.loads(payload["metadata"]).get("upstream_provider") == "serving-site-a" + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_omits_upstream_provider_when_hidden_param_absent_guard() -> None: + standard_logging_payload: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": standard_logging_payload, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + }, + }, + start_time=datetime.datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc), + ) + + assert "upstream_provider" not in json.loads(payload["metadata"]) + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_strips_null_bytes_from_request_tags(): @@ -6911,7 +6970,10 @@ def test_spend_logs_payload(model_id: str | None): ) assert len(payload["request_id"]) > 0 - expected_metadata_keys: Final = SpendLogsMetadata.__annotations__.keys() - {"is_silent_experiment"} + expected_metadata_keys: Final = SpendLogsMetadata.__annotations__.keys() - { + "is_silent_experiment", + "upstream_provider", + } assert "metadata" in payload assert isinstance(payload["metadata"], str) diff --git a/tests/unit/proxy/spend_tracking/test_usage_meter.py b/tests/unit/proxy/spend_tracking/test_usage_meter.py new file mode 100644 index 00000000000..de7d7c6591d --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_usage_meter.py @@ -0,0 +1,389 @@ +import asyncio +import json +from datetime import datetime, timezone +from decimal import Decimal +from typing import Final + +import pytest + +from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome +from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite +from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings +from litellm.types.utils import CallTypes + + +def _write(identity: str) -> UsageWrite: + event: Final = InboundUsageEvent( + event_id=identity, + request_id=identity, + call_type=CallTypes.acompletion, + started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + outcome=Outcome.SUCCEEDED, + dimensions=(), + metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), + ) + return UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")})) + + +class CommitBarrier: + def __init__(self, *, conflict: str | None = None, failure: Exception | None = None) -> None: + self.entered: Final = asyncio.Event() + self.release: Final = asyncio.Event() + self.committed: tuple[str, ...] = () + self.transactions = 0 + self.conflict: Final = conflict + self.failure: Final = failure + + async def write(self, writes: tuple[UsageWrite, ...]) -> None: + self.transactions += 1 + self.entered.set() + await self.release.wait() + if self.failure is not None: + raise self.failure + identities: Final = tuple(write.event.event_id for write in writes) + if self.conflict in identities: + raise AccountingConflictError("Conflicting test event") + self.committed += identities + + +async def _manual_flush(wakeup: asyncio.Event, timeout: float) -> None: + await wakeup.wait() + + +@pytest.mark.asyncio +async def test_receipts_wait_for_one_batched_commit_and_keep_inflight_capacity_reserved() -> None: + database: Final = CommitBarrier() + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) + first: Final = meter.submit((_write("one"),)) + second: Final = meter.submit((_write("two"),)) + await database.entered.wait() + assert not first.done() and not second.done() + assert meter.pending_records == 2 + with pytest.raises(UsageCaptureError, match="capacity"): + meter.submit((_write("three"),)) + database.release.set() + await asyncio.gather(first, second) + await meter.close() + assert database.committed == ("one", "two") + assert database.transactions == 1 + assert (meter.pending_bytes, meter.pending_records) == (0, 0) + + +@pytest.mark.asyncio +async def test_byte_limit_rejects_new_work_without_evicting_admitted_work() -> None: + database: Final = CommitBarrier() + write: Final = _write("large" * 300) + size: Final = len(write.envelope.encode()) + meter: Final = DurableUsageMeter( + database.write, UsageCaptureSettings(max_pending_bytes=size), wait_for_batch=_manual_flush + ) + receipt: Final = meter.submit((write,)) + with pytest.raises(UsageCaptureError, match="capacity"): + meter.submit((_write("extra"),)) + assert meter.pending_bytes == size + database.release.set() + await meter.close() + await receipt + assert database.committed == (write.event.event_id,) + with pytest.raises(UsageCaptureError, match="closed"): + meter.submit((_write("after-close"),)) + + +@pytest.mark.asyncio +async def test_conflicting_group_rolls_back_and_valid_neighbor_still_commits() -> None: + database: Final = CommitBarrier(conflict="conflict") + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3)) + rejected: Final = meter.submit((_write("caller"), _write("conflict"))) + accepted: Final = meter.submit((_write("valid"),)) + await database.entered.wait() + database.release.set() + with pytest.raises(UsageCaptureError, match="identity"): + await rejected + await accepted + await meter.close() + assert database.committed == ("valid",) + assert meter.pending_records == 0 + + +@pytest.mark.asyncio +async def test_permanent_persistence_failure_reaches_each_receipt_and_releases_capacity() -> None: + failure: Final = ValueError("invalid envelope") + database: Final = CommitBarrier(failure=failure) + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) + first: Final = meter.submit((_write("one"),)) + second: Final = meter.submit((_write("two"),)) + await database.entered.wait() + database.release.set() + results: Final = await asyncio.gather(first, second, return_exceptions=True) + await meter.close() + assert results == [failure, failure] + assert meter.pending_bytes == 0 + + +@pytest.mark.asyncio +async def test_cancelling_record_waiter_does_not_cancel_persistence() -> None: + database: Final = CommitBarrier() + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) + waiting: Final = asyncio.create_task(meter.record(_write("caller").event)) + await database.entered.wait() + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + database.release.set() + await meter.close() + assert database.committed == ("caller",) + assert meter.pending_records == 0 + + +@pytest.mark.asyncio +async def test_shutdown_deadline_rejects_unacknowledged_work_and_releases_capacity() -> None: + database: Final = CommitBarrier() + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) + first: Final = meter.submit((_write("in-flight"), _write("in-flight-two"))) + await database.entered.wait() + queued: Final = meter.submit((_write("queued"),)) + await meter.stop(0) + results: Final = await asyncio.gather(first, queued, return_exceptions=True) + assert all(isinstance(result, UsageCaptureError) for result in results) + assert meter.pending_records == 0 + assert meter.pending_bytes == 0 + assert database.committed == () + + +@pytest.mark.asyncio +async def test_shutdown_during_database_backoff_releases_reserved_capacity() -> None: + retrying: Final = asyncio.Event() + blocked: Final = asyncio.Event() + + async def unavailable(writes: tuple[UsageWrite, ...]) -> None: + raise ConnectionError("database unavailable") + + async def backoff(delay: float) -> None: + retrying.set() + await blocked.wait() + + meter: Final = DurableUsageMeter(unavailable, UsageCaptureSettings(batch_size=2), retry_delay=backoff) + receipt: Final = meter.submit((_write("one"), _write("two"))) + await retrying.wait() + await meter.stop(0) + with pytest.raises(UsageCaptureError, match="interrupted"): + await receipt + assert (meter.pending_records, meter.pending_bytes) == (0, 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wrapped", [False, True]) +async def test_closed_pooled_connection_retains_batch_until_automatic_recovery(wrapped: bool) -> None: + import httpx + from prisma.engine.errors import EngineRequestError + from prisma.http import Response + + database: Final = CommitBarrier() + attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]() + retrying: Final = asyncio.Event() + recovered: Final = asyncio.Event() + closure: Final = EngineRequestError( + Response(httpx.Response(500)), '{"error_code":"P1017","message":"Server has closed the connection."}' + ) + failure: Final = RuntimeError("ledger transaction failed") if wrapped else closure + if wrapped: + failure.__cause__ = closure + + async def persist(writes: tuple[UsageWrite, ...]) -> None: + attempts.put_nowait(writes) + if attempts.qsize() == 1: + raise failure + await database.write(writes) + + async def backoff(delay: float) -> None: + retrying.set() + await recovered.wait() + + meter: Final = DurableUsageMeter( + persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), retry_delay=backoff + ) + writes: Final = (_write("caller"), _write("attempt")) + receipt: Final = meter.submit(writes) + retry_signal: Final = asyncio.create_task(retrying.wait()) + reserved: Final = meter.pending_bytes + try: + await asyncio.wait((receipt, retry_signal), return_when=asyncio.FIRST_COMPLETED) + assert not receipt.done(), "A dropped pooled connection must retain admitted usage for retry" + assert (meter.pending_records, meter.pending_bytes) == (2, reserved) + with pytest.raises(UsageCaptureError, match="capacity"): + meter.submit((_write("overflow"),)) + database.release.set() + recovered.set() + await receipt + assert attempts.get_nowait() == attempts.get_nowait() == writes + assert database.committed == ("caller", "attempt") + assert (meter.pending_records, meter.pending_bytes) == (0, 0) + finally: + retry_signal.cancel() + await asyncio.gather(retry_signal, return_exceptions=True) + await meter.stop(0) + await asyncio.gather(receipt, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("first_commits", [False, True]) +async def test_shutdown_during_conflict_isolation_completes_every_receipt(first_commits: bool) -> None: + isolating: Final = asyncio.Event() + blocked: Final = asyncio.Event() + + async def persist(writes: tuple[UsageWrite, ...]) -> None: + if len(writes) > 1: + raise AccountingConflictError("one producer conflicts") + if first_commits and writes[0].event.event_id == "one": + return + isolating.set() + await blocked.wait() + + meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=2)) + first: Final = meter.submit((_write("one"),)) + second: Final = meter.submit((_write("two"),)) + await isolating.wait() + await meter.stop(0) + results: Final = await asyncio.gather(first, second, return_exceptions=True) + assert results[0] is None if first_commits else isinstance(results[0], UsageCaptureError) + assert isinstance(results[1], UsageCaptureError) + assert (meter.pending_records, meter.pending_bytes) == (0, 0) + + +@pytest.mark.asyncio +async def test_elapsed_flush_deadline_commits_a_partial_batch() -> None: + from litellm.proxy.spend_tracking.usage_meter import _wait_for_batch + + database: Final = CommitBarrier() + database.release.set() + + async def expired_flush(wakeup: asyncio.Event, timeout: float) -> None: + assert timeout == 5.0 + await _wait_for_batch(wakeup, 0) + + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=10), wait_for_batch=expired_flush) + receipt: Final = meter.submit((_write("partial-batch"),)) + try: + await receipt + assert database.committed == ("partial-batch",) + assert database.transactions == 1 + assert (meter.pending_records, meter.pending_bytes) == (0, 0) + finally: + await meter.stop(0) + + +@pytest.mark.asyncio +async def test_burst_drains_full_batches_without_waiting_for_partial_batch_deadline() -> None: + database: Final = CommitBarrier() + database.release.set() + waits: Final = asyncio.Queue[bool]() + + async def observe_wait(wakeup: asyncio.Event, timeout: float) -> None: + waits.put_nowait(wakeup.is_set()) + await wakeup.wait() + + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2), wait_for_batch=observe_wait) + receipts: Final = tuple(meter.submit((_write(str(index)),)) for index in range(5)) + try: + assert await waits.get() is True + await asyncio.gather(*receipts[:2]) + assert await waits.get() is True, "A full queued batch must not wait for the flush deadline" + await asyncio.gather(*receipts[2:4]) + assert await waits.get() is False + assert database.committed == ("0", "1", "2", "3") + assert meter.pending_records == 1 + assert not receipts[-1].done() + await meter.close() + await receipts[-1] + assert database.committed == ("0", "1", "2", "3", "4") + assert database.transactions == 3 + assert (meter.pending_records, meter.pending_bytes) == (0, 0) + finally: + await meter.stop(0) + await asyncio.gather(*receipts, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sqlstate", ["22P05", "22003", "42P01"]) +async def test_invalid_submission_isolated_without_retrying_systemic_faults(sqlstate: str) -> None: + from prisma.errors import DataError + + database: Final = CommitBarrier() + database.release.set() + attempts: Final = asyncio.Queue[tuple[str, ...]]() + + async def persist(writes: tuple[UsageWrite, ...]) -> None: + identities: Final = tuple(write.event.event_id for write in writes) + attempts.put_nowait(identities) + if "invalid" in identities or sqlstate == "42P01": + raise DataError( + data={"user_facing_error": {"message": "database rejected input", "meta": {"code": sqlstate}}} + ) + await database.write(writes) + + meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=3)) + rejected: Final = meter.submit((_write("caller"), _write("invalid"))) + neighbor: Final = meter.submit((_write("neighbor"),)) + try: + with pytest.raises(DataError): + await rejected + if sqlstate == "42P01": + with pytest.raises(DataError): + await neighbor + assert attempts.qsize() == 1 + else: + await neighbor + assert database.committed == ("neighbor",) + assert attempts.qsize() == 3 + finally: + await meter.close() + assert (meter.pending_records, meter.pending_bytes) == (0, 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("dependency_outcome", ["commit", "failure", "cancel"]) +async def test_commit_dependencies_preserve_capacity_and_isolate_failed_dependents(dependency_outcome: str) -> None: + database: Final = CommitBarrier() + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3, max_pending_records=4)) + prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() + dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,)) + neighbor: Final = meter.submit((_write("neighbor"),)) + assert meter.pending_records == 3 + assert not dependent.done() + with pytest.raises(UsageCaptureError, match="capacity"): + meter.submit((_write("excess-one"), _write("excess-two"))) + if dependency_outcome == "commit": + prerequisite.set_result(None) + elif dependency_outcome == "failure": + prerequisite.set_exception(ValueError("dependency failed")) + else: + prerequisite.cancel() + database.release.set() + await meter.close() + await neighbor + if dependency_outcome == "commit": + await dependent + assert database.committed == ("parent", "final", "neighbor") + else: + with pytest.raises(UsageCaptureError, match="dependency"): + await dependent + assert database.committed == ("neighbor",) + assert meter.pending_records == 0 + assert meter.pending_bytes == 0 + + +@pytest.mark.asyncio +async def test_shutdown_does_not_cancel_dependency_receipt_or_leak_reserved_capacity() -> None: + database: Final = CommitBarrier() + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) + prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() + dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,)) + await meter.stop(0) + with pytest.raises(UsageCaptureError, match="acknowledgement"): + await dependent + assert not prerequisite.done() + assert database.committed == () + assert meter.pending_records == 0 + assert meter.pending_bytes == 0 + prerequisite.set_result(None) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 2d08609c93b..4d1b67f9212 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -74,6 +74,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck @@ -83,10 +84,11 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( team_member_budget_entity_id, ) from litellm.proxy.utils import ProxyLogging +from litellm.responses.mcp.parked_stream_dispatch import IN_FLIGHT_ROUND from litellm.router import Router from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.utils import ModelResponseListIterator -from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment +from litellm.litellm_core_utils.add_retry_fallback_headers import prepare_response_for_header_attachment def test_attach_guardrail_information_copies_recorded_entries_onto_model_response(): @@ -9051,6 +9053,33 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + @pytest.mark.asyncio + async def test_disconnect_whose_in_flight_round_billing_raises_still_closes_the_stream_and_frees_the_slot(self): + import types + + response = await self._start_partial_stream() + closed: Final = AsyncMock() + + async def cannot_price_the_in_flight_round() -> None: + raise RuntimeError("the in-flight round could not be priced") + + setattr(response.logging_obj, IN_FLIGHT_ROUND, cannot_price_the_in_flight_round) + stream: Final = types.SimpleNamespace(chunks=[], messages=None, aclose=closed) + proxy_logging_obj: Final = types.SimpleNamespace(arelease_max_parallel_requests_on_disconnect=AsyncMock()) + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=stream, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=proxy_logging_obj, + ) + + closed.assert_awaited_once() + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + async def _bill_and_collect_success_event(self, prepare=None, request_data=None): recorder = _RecordingSuccessLogger() original_callbacks = litellm.callbacks @@ -10193,6 +10222,8 @@ class TestInjectCostIntoUsageDict: discount streamed a ``usage.cost`` that disagreed with the callback's ``response_cost``.""" class _StubLoggingObj: + model = "claude-haiku-4-5" + def __init__(self, cost): self._cost = cost self.captured_result = None @@ -10229,6 +10260,8 @@ class TestInjectCostIntoUsageDict: def test_message_delta_falls_back_to_model_pricing_when_the_logging_obj_returns_no_cost(self): class _StubLoggingObj: + model = "claude-haiku-4-5" + def response_cost_calculator(self, result): return None @@ -10254,6 +10287,8 @@ class TestInjectCostIntoUsageDict: model-name pricing rather than propagating into the response body.""" class _StubLoggingObj: + model = "claude-haiku-4-5" + def response_cost_calculator(self, result): raise ValueError("no pricing for this deployment") @@ -10335,6 +10370,8 @@ class TestInjectCostIntoUsageDict: streaming /v1/chat/completions gets its negotiated price instead of sticker.""" class _StubLoggingObj: + model = "gpt-4o-mini" + def __init__(self, cost): self._cost = cost self.captured_result = None @@ -10363,6 +10400,8 @@ class TestInjectCostIntoUsageDict: def test_openai_chunk_falls_back_to_model_pricing_when_the_logging_obj_returns_no_cost(self): class _StubLoggingObj: + model = "gpt-4o-mini" + def response_cost_calculator(self, result): return None @@ -10378,6 +10417,60 @@ class TestInjectCostIntoUsageDict: assert result is not None assert result["usage"]["cost"] == pytest.approx(self._expected_cost("gpt-4o-mini", 11, 4)) + @pytest.mark.parametrize( + ("deployment_model", "served_model", "custom_llm_provider"), + [ + ("vertex_ai/claude-sonnet-5-5", "claude-sonnet-5-5", "vertex_ai"), + ("anthropic/claude-sonnet-5-5", "claude-sonnet-5-5", "anthropic"), + ("azure_ai/claude-sonnet-5-5", "claude-sonnet-5-5", "azure_ai"), + ], + ) + def test_message_delta_bills_search_like_the_logged_response_on_a_token_priced_claude_deployment( + self, local_model_cost_map: None, deployment_model: str, served_model: str, custom_llm_provider: str + ) -> None: + token_rates: Final = {"input_cost_per_token": 3e-6, "output_cost_per_token": 1.5e-5} + router: Final = Router( + model_list=[ + { + "model_name": "priced-claude", + "litellm_params": { + "model": deployment_model, + "api_key": "k", + "vertex_project": "p", + "vertex_location": "global", + }, + "model_info": token_rates, + } + ] + ) + deployment_id: Final = router.model_list[0]["model_info"]["id"] + logging_obj: Final = LiteLLMLoggingObj( + model=deployment_model, + messages=[{"role": "user", "content": "test"}], + stream=True, + call_type="anthropic_messages", + start_time=None, + litellm_call_id="token-priced-claude-search", + function_id="token-priced-claude-search", + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {"model_info": {"id": deployment_id, **token_rates}}}, optional_params={} + ) + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + wire_usage: Final = {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}} + event: Final = {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": wire_usage} + + result: Final = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "priced-claude", logging_obj) + + stream_usage: Final = AnthropicConfig().calculate_usage(usage_object=wire_usage, reasoning_content=None) + logged_cost: Final = logging_obj.response_cost_calculator( + result=litellm.ModelResponse(model=served_model, usage=stream_usage) + ) + token_cost: Final = 1840 * token_rates["input_cost_per_token"] + 412 * token_rates["output_cost_per_token"] + assert logged_cost is not None and logged_cost > token_cost + assert result is not None + assert result["usage"]["cost"] == pytest.approx(logged_cost) + class TestProcessChunkWithCostInjection: @pytest.mark.parametrize("model", ["claude-opus-5-5", "vertex_ai/claude-opus-5-5"]) @@ -10444,6 +10537,8 @@ class TestProcessChunkWithCostInjection: monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) class _StubLoggingObj: + model = "claude-haiku-4-5" + def response_cost_calculator(self, result): return 0.00042 @@ -12468,6 +12563,55 @@ class TestResponseCostFromLoggingObjForCacheHits: ) assert {name: headers.get(name) for name in cost_headers} == dict.fromkeys(cost_headers, expected_header) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("response_cost", "expected_header"), + ((float("nan"), None), (float("inf"), None), (0.0042, "0.0042"), (2, "2")), + ids=("nan", "inf", "finite", "int"), + ) + async def test_a_non_finite_response_cost_sends_no_cost_header( + self, response_cost: float, expected_header: str | None + ): + response: Final = litellm.ModelResponse() + response._hidden_params = {"response_cost": response_cost} + + headers: Final = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data={"model": "gpt-5.4-mini"}, + request=Request({"type": "http", "method": "POST", "headers": []}), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + logging_obj=self._logging_obj(), + version=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + assert headers.get("x-litellm-response-cost") == expected_header + + @pytest.mark.asyncio + async def test_an_overflowing_cost_breakdown_sends_no_infinite_cost_header(self): + response: Final = litellm.ModelResponse() + response._hidden_params = {"response_cost": 0.0} + logging_obj: Final = self._logging_obj() + logging_obj.cost_breakdown = { + "original_cost": float("inf"), + "input_cost": float("inf"), + "output_cost": 0.0012, + "total_cost": float("inf"), + } + + headers: Final = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data={"model": "gpt-5.4-mini"}, + request=Request({"type": "http", "method": "POST", "headers": []}), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + logging_obj=logging_obj, + version=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + cost_headers: Final = {name: value for name, value in headers.items() if name.startswith("x-litellm-response-cost")} + assert cost_headers == {"x-litellm-response-cost-output": "0.0012"} + def test_a_fresh_zero_cost_v1_messages_call_still_sends_no_cost_header(self): logging_obj: Final = self._logging_obj() logging_obj.model_call_details["response_cost"] = 0.0 diff --git a/tests/unit/proxy/test_health_check_functions.py b/tests/unit/proxy/test_health_check_functions.py index 2c0198cd538..79eec44958f 100644 --- a/tests/unit/proxy/test_health_check_functions.py +++ b/tests/unit/proxy/test_health_check_functions.py @@ -741,6 +741,7 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c result = await _perform_health_check_and_save( model_list=model_list, target_model=None, + recorded_model_name=None, cli_model=None, details=True, prisma_client=None, @@ -777,6 +778,7 @@ async def test_perform_health_check_and_save_forwards_skip_disabled_background_f await _perform_health_check_and_save( model_list=model_list, target_model=None, + recorded_model_name=None, cli_model=None, details=True, prisma_client=None, diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c5136ff6e4a..91c87c63067 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -4323,6 +4323,186 @@ def test_add_litellm_metadata_from_request_headers_traceparent_and_baggage_toget assert data["litellm_session_id"] == "baggage-session-42" +def _bedrock_passthrough_data( + client_body: dict[str, object], custom_llm_provider: str = "bedrock" +) -> dict[str, object]: + return { + "method": "POST", + "endpoint": "model/us.anthropic.claude-haiku-4-5-20251001-v1:0/converse", + "custom_llm_provider": custom_llm_provider, + "data": client_body, + "litellm_metadata": {}, + } + + +@pytest.mark.parametrize( + ("headers", "client_body"), + [ + pytest.param({}, {"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}, id="converse_body"), + pytest.param( + {"X-Amzn-Bedrock-Request-Metadata": '{"threadId": "thread-abc-123", "team": "search"}'}, + {"messages": []}, + id="invoke_header", + ), + ], +) +def test_bedrock_request_metadata_thread_id_becomes_session_id( + headers: dict[str, str], client_body: dict[str, object] +): + sent_body: Final = copy.deepcopy(client_body) + data: Final = _bedrock_passthrough_data(client_body) + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert data["litellm_session_id"] == "thread-abc-123" + assert data["litellm_metadata"] == {"session_id": "thread-abc-123"} + assert "litellm_trace_id" not in data, "a conversation tag must not merge the turns into one trace" + assert data["data"] == sent_body, "the tags must still reach Bedrock unchanged" + + +@pytest.mark.parametrize( + ("tags", "expected"), + [ + ({"threadId": "thread-abc-123", "sessionId": "session-xyz-789"}, "session-xyz-789"), + ({"conversationId": "conversation-42", "threadId": "thread-abc-123"}, "thread-abc-123"), + ({"conversationId": "conversation-42", "user": "shrey-1234"}, "conversation-42"), + ({"sessionId": "", "threadId": "thread-abc-123"}, "thread-abc-123"), + ({"sessionId": "conv-42"}, "conv-42"), + ({"sessionId": "user:123 / thread.7"}, "user:123 / thread.7"), + ], +) +def test_bedrock_request_metadata_session_key_priority(tags: dict[str, str], expected: str): + data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": tags}) + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert data["litellm_session_id"] == expected + + +@pytest.mark.parametrize( + ("headers", "data"), + [ + pytest.param( + {"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"'}, + _bedrock_passthrough_data({"messages": []}), + id="malformed_header_json", + ), + pytest.param( + {"x-amzn-bedrock-request-metadata": '["thread-abc-123"]'}, + _bedrock_passthrough_data({"messages": []}), + id="header_json_not_an_object", + ), + pytest.param( + {}, _bedrock_passthrough_data({"requestMetadata": {"threadId": 1234567890}}), id="non_string_value" + ), + pytest.param({}, _bedrock_passthrough_data({"requestMetadata": {"threadId": ""}}), id="empty"), + pytest.param({}, _bedrock_passthrough_data({"requestMetadata": {"user": "shrey-1234"}}), id="no_session_key"), + pytest.param( + {"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"}'}, + {"model": "claude", "custom_llm_provider": "bedrock", "litellm_metadata": {}}, + id="not_a_passthrough_body", + ), + pytest.param( + {}, + _bedrock_passthrough_data( + {"requestMetadata": {"threadId": "thread-abc-123"}}, custom_llm_provider="vertex_ai" + ), + id="not_bedrock", + ), + ], +) +def test_bedrock_request_metadata_without_usable_session_id_sets_no_session( + headers: dict[str, str], data: dict[str, object] +): + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert "litellm_session_id" not in data + + +@pytest.mark.parametrize( + ("body_tags", "header_tags", "expected"), + [ + pytest.param({"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key"), + pytest.param({"sessionId": "body-session"}, {"sessionId": "header-session"}, "body-session", id="body_wins"), + ], +) +def test_bedrock_request_metadata_reads_body_tags_then_header_tags( + body_tags: dict[str, str], header_tags: dict[str, str], expected: str +): + data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": body_tags}) + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-amzn-bedrock-request-metadata": json.dumps(header_tags)}, + data=data, + _metadata_variable_name="litellm_metadata", + ) + + assert data["litellm_session_id"] == expected + + +@pytest.mark.parametrize("metadata_bucket", ["litellm_metadata", "metadata"]) +def test_bedrock_request_metadata_never_replaces_a_caller_metadata_session(metadata_bucket: str): + data: Final = { + **_bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-B"}}), + metadata_bucket: {"session_id": "conversation-A"}, + } + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert "litellm_session_id" not in data + assert data[metadata_bucket]["session_id"] == "conversation-A" + + +def test_bedrock_request_metadata_stamps_a_tag_equal_to_the_caller_metadata_session(): + data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}) + data["litellm_metadata"] = {"session_id": "thread-abc-123"} + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert data["litellm_session_id"] == "thread-abc-123" + + +def test_bedrock_request_metadata_header_leaves_non_passthrough_caller_session_alone(): + data: Final = {"model": "claude", "custom_llm_provider": "bedrock", "metadata": {"session_id": "caller-session-1"}} + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"}'}, + data=data, + _metadata_variable_name="metadata", + ) + + assert data["metadata"] == {"session_id": "caller-session-1"} + assert "litellm_session_id" not in data + + +@pytest.mark.parametrize( + ("headers", "expected"), + [ + pytest.param({"x-litellm-session-id": "explicit-session-1"}, "explicit-session-1", id="explicit_header_wins"), + pytest.param({"baggage": "session.id=baggage-session-42"}, "thread-abc-123", id="beats_baggage"), + ], +) +def test_bedrock_request_metadata_session_precedence(headers: dict[str, str], expected: str): + data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}) + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert data["litellm_session_id"] == expected + assert data["litellm_metadata"]["session_id"] == expected + + def test_add_litellm_metadata_from_request_headers_explicit_trace_id_beats_traceparent(): """x-litellm-trace-id must win over a traceparent header carrying a different trace-id - explicit litellm headers are always highest priority.""" @@ -5360,6 +5540,88 @@ async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_ assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"] +@pytest.mark.asyncio +async def test_add_guardrails_from_policy_engine_keeps_configured_guardrails_and_skips_ones_already_selected(): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"guardrails": [configured, "audit"]}, + } + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [PolicyAttachment(policy="baseline", scope="*")] + attachment_registry._initialized = True + + try: + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + +def test_add_guardrails_from_auth_metadata_lists_each_key_and_team_guardrail_once_keeping_the_key_s_settings(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {}} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + metadata={"guardrails": [configured, "audit"]}, + team_metadata={"guardrails": ["audit", "pii_blocker", "toxic"]}, + ) + + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, data=data, metadata_variable_name="metadata" + ) + + assert data["metadata"]["guardrails"] == [configured, "audit", "toxic"] + + +def test_add_guardrails_from_auth_metadata_keeps_a_selected_guardrail_s_settings_when_a_key_policy_adds_it(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {"guardrails": [configured, "audit"]}} + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + + try: + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", metadata={"policies": ["baseline"]}), + data=data, + metadata_variable_name="metadata", + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_applies_inherited_parent_guardrail_when_child_condition_misses(): from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry @@ -9516,6 +9778,19 @@ async def test_missing_session_id_reject_accepts_body_metadata_session_id(): assert updated["metadata"]["session_id"] == "body-session-1" +@pytest.mark.asyncio +async def test_missing_session_id_reject_accepts_bedrock_request_metadata_session_id(): + updated: Final = await add_litellm_data_to_request( + data=_bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}), + request=_request_for("/bedrock/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/converse"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert _spend_log_session_id(updated, "litellm_metadata") == "thread-abc-123" + + @pytest.mark.asyncio async def test_missing_session_id_reject_returns_400_without_session_id(): with pytest.raises(ProxyException) as exc_info: diff --git a/tests/unit/proxy/test_mcp_asgi_response.py b/tests/unit/proxy/test_mcp_asgi_response.py index d030f65af4b..a9e884d680b 100644 --- a/tests/unit/proxy/test_mcp_asgi_response.py +++ b/tests/unit/proxy/test_mcp_asgi_response.py @@ -1,11 +1,70 @@ import asyncio +from typing import Final import pytest from fastapi import HTTPException +from starlette.types import Message, Receive, Scope, Send from litellm.proxy.proxy_server import _stream_mcp_asgi_response +@pytest.mark.asyncio +async def test_completed_response_finishes_handler_cleanup(): + body_consumed: Final = asyncio.Event() + completed: Final = asyncio.Event() + + async def handle_fn(_scope: Scope, _receive: Receive, send: Send) -> None: + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"initialized", "more_body": False}) + await body_consumed.wait() + completed.set() + + async def receive() -> Message: + return {"type": "http.request", "body": b"", "more_body": False} + + response: Final = await _stream_mcp_asgi_response( + handle_fn, {"type": "http", "method": "POST", "path": "/mcp", "headers": []}, receive + ) + body: Final = aiter(response.body_iterator) + assert await anext(body) == b"initialized" + body_consumed.set() + assert await anext(body, None) is None + assert completed.is_set() + + +@pytest.mark.asyncio +async def test_disconnected_response_cancels_handler(): + body_received: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + hold_open: Final = asyncio.Event() + + async def handle_fn(_scope: Scope, _receive: Receive, send: Send) -> None: + try: + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"event", "more_body": True}) + await hold_open.wait() + finally: + cancelled.set() + + async def receive() -> Message: + return {"type": "http.request", "body": b"", "more_body": False} + + response: Final = await _stream_mcp_asgi_response( + handle_fn, {"type": "http", "method": "POST", "path": "/mcp", "headers": []}, receive + ) + + async def consume() -> None: + async for _ in response.body_iterator: + body_received.set() + + consumer: Final = asyncio.create_task(consume()) + await body_received.wait() + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + assert cancelled.is_set() + + @pytest.mark.asyncio async def test_stream_mcp_asgi_response_propagates_pre_header_http_exception(): async def handle_fn(_scope, _receive, _send): diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 139cfbbe58c..b3725f01512 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -116,7 +116,7 @@ async def test_admin_mcp_failure_still_closes_proxy_resources( message: Final = "LITELLM_LICENSE" if failure_phase == "license" else f"{failure_phase} failed" async def run_lifespan() -> None: async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state: - assert state == {"tracing_receiver": None} + assert state == {"tracing_receiver": None, "database": None} if failure_phase == "serving": raise RuntimeError("serving failed") if failure_phase == "cancelled": @@ -168,7 +168,7 @@ async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector( transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000" ) as client: async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state: - assert state == {"tracing_receiver": None} + assert state == {"tracing_receiver": None, "database": None} request: Final = asyncio.create_task(client.post( "/admin/mcp", headers={"Authorization": "Bearer admin", "Accept": "application/json, text/event-stream"}, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index bc3a228321a..3647146eb4e 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -1654,6 +1654,15 @@ class MockPrisma: async def disconnect(self): pass + async def stop_background_tasks(self): + pass + + async def aclose(self): + try: + await self.stop_background_tasks() + finally: + await self.disconnect() + mock_prisma = MockPrisma() @@ -1764,7 +1773,7 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_even_when_the_database from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError from litellm.proxy.proxy_server import proxy_startup_event - async def unreachable(): + async def unreachable(resources): raise ConnectionError("database is down") _, announced = _boot_with_general_settings(monkeypatch, tmp_path, general_settings) @@ -1809,7 +1818,7 @@ async def test_proxy_startup_asks_to_migrate_only_when_the_database_holds_values database = _DatabaseWithOneStoredCredential(encrypt_value_helper("sk-provider", new_encryption_key=encrypted_with)) - async def connected(): + async def connected(resources): return database _, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": PUBLICLY_KNOWN_KEY}) @@ -1844,6 +1853,21 @@ async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_delet class _PrismaClientWhoseDatabaseRejectsQueries: + def __init__(self) -> None: + self.disconnect_calls = 0 + + async def disconnect(self) -> None: + self.disconnect_calls += 1 + + async def stop_background_tasks(self) -> None: + pass + + async def aclose(self) -> None: + try: + await self.stop_background_tasks() + finally: + await self.disconnect() + class _Database: async def query_raw(self, query, *args): raise RuntimeError("permission denied for table LiteLLM_CredentialsTable") @@ -1858,7 +1882,8 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc from litellm.proxy.proxy_server import proxy_startup_event _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + prisma: Final = _PrismaClientWhoseDatabaseRejectsQueries() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY) with ( @@ -1871,6 +1896,25 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message] assert len(notices) == 1 assert "Could not migrate stored values" in notices[0] + assert prisma.disconnect_calls == 1 + + +@pytest.mark.asyncio +async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_client(monkeypatch, tmp_path): + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY) + with pytest.raises(RuntimeError, match="permission denied"): + async with proxy_startup_event(FastAPI()): + pass + + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "") + async with proxy_startup_event(FastAPI()) as state: + assert state["database"] is None @pytest.mark.asyncio @@ -1917,11 +1961,11 @@ async def test_proxy_startup_boots_an_unsafe_master_key_under_the_override(monke class _ShutdownAwarePrisma(MockPrisma): def __init__(self): super().__init__() - self.stop_view_setup_task = AsyncMock() + self.stop_background_tasks = AsyncMock() @pytest.mark.asyncio -async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): +async def test_proxy_shutdown_stops_database_background_tasks(monkeypatch, tmp_path): import yaml from fastapi import FastAPI @@ -1936,11 +1980,11 @@ async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): monkeypatch.setattr(proxy_server_module, "store_model_in_db", False) async with proxy_startup_event(FastAPI()): - stopped_while_serving = fake_prisma.stop_view_setup_task.await_count + stopped_while_serving = fake_prisma.stop_background_tasks.await_count actual = { "stopped_while_serving": stopped_while_serving, - "stopped_after_shutdown": fake_prisma.stop_view_setup_task.await_count, + "stopped_after_shutdown": fake_prisma.stop_background_tasks.await_count, } assert actual == {"stopped_while_serving": 0, "stopped_after_shutdown": 1} @@ -14231,12 +14275,22 @@ async def test_team_window_spend_row_carries_the_request_start_time(): def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client = MagicMock() client.connect = AsyncMock(side_effect=connect_error) + client.disconnect = AsyncMock() + client.stop_background_tasks = AsyncMock() client.db.start_token_refresh_task = AsyncMock() client.check_view_exists = AsyncMock() client._set_spend_logs_row_count_in_proxy_state = AsyncMock() client.start_view_setup_task = MagicMock() client.start_db_health_watchdog_task = AsyncMock() client.health_check = AsyncMock(side_effect=health_check_error) + + async def close_client() -> None: + try: + await client.stop_background_tasks() + finally: + await client.disconnect() + + client.aclose = AsyncMock(side_effect=close_client) return client @@ -14343,6 +14397,8 @@ async def test_setup_prisma_client_raises_when_db_unavailable_is_not_allowed(mon mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out")) with pytest.raises(httpx.ReadTimeout): await _run_setup_prisma_client(mock_client) + mock_client.stop_background_tasks.assert_awaited_once() + mock_client.disconnect.assert_awaited_once() @pytest.mark.asyncio @@ -14363,6 +14419,31 @@ async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkey assert result is None assert mock_client.start_db_health_watchdog_task.await_count == 0 assert mock_client.health_check.await_count == 0 + mock_client.stop_background_tasks.assert_awaited_once() + mock_client.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_cancelling_prisma_setup_closes_the_client_before_publishing_it(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False") + checking: Final = asyncio.Event() + client: Final = _mock_startup_prisma_client() + + async def health_check() -> None: + checking.set() + await asyncio.Event().wait() + + client.health_check = health_check + setup: Final = asyncio.create_task(_run_setup_prisma_client(client)) + try: + await asyncio.wait_for(checking.wait(), timeout=2) + finally: + setup.cancel() + with pytest.raises(asyncio.CancelledError): + await setup + + client.stop_background_tasks.assert_awaited_once() + client.disconnect.assert_awaited_once() async def _run_scheduled_background_jobs(): @@ -15317,12 +15398,12 @@ async def test_load_config_router_budget_checks_fallback_targets_against_the_cal # on by default: an over-budget caller is refused the paid fallback with no config at all monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) - assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is False - assert await router.fallback_budget_check(model="m", request_kwargs=under_budget, llm_router=router) is True + assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, ctx=router._core_ctx) is False + assert await router.fallback_budget_check(model="m", request_kwargs=under_budget, ctx=router._core_ctx) is True # explicit opt-out restores the unguarded behaviour monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": False}, raising=False) - assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is True + assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, ctx=router._core_ctx) is True @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index c182cd1ec75..68eb05e60dd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -179,6 +179,39 @@ async def test_token_counting_model_not_in_model_list(): assert response.model_used == "special-alias" +@pytest.mark.asyncio +async def test_asking_the_provider_endpoint_leaves_an_openrouter_deployments_count_unchanged( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "or-gpt", + "litellm_params": { + "model": "openai/gpt-4o", + "custom_llm_provider": "openrouter", + "api_key": "sk-or-token-count", + }, + } + ] + ), + ) + request: Final = TokenCountRequest( + model="or-gpt", messages=[{"role": "user", "content": "Count these tokens, ünïcödé 你好 🙂"}] + ) + + via_provider_endpoint: Final = await token_counter(request=request, call_endpoint=True) + counted_locally: Final = await token_counter(request=request, call_endpoint=False) + + assert (via_provider_endpoint.total_tokens, via_provider_endpoint.model_used) == ( + counted_locally.total_tokens, + counted_locally.model_used, + ) + + @pytest.mark.asyncio async def test_gpt_token_counting(): """ diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 3cdd2edb2fb..0a483840336 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1844,6 +1844,26 @@ class TestPostCallFailureHookLLMExceptionAlerting: alerting_handler = await self._alerted(HTTPException(status_code=429, detail="rate limited")) assert alerting_handler.called is False + @pytest.mark.asyncio + @pytest.mark.parametrize("status_code, should_alert", ((400, False), (422, False), (429, False), (500, True))) + async def test_budget_exceeded_alerts_only_for_server_status( + self, status_code: int, should_alert: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "budget_exceeded_status_code", status_code) + exc: Final = litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5) + assert (await self._alerted(exc)).called is should_alert + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "exc", + ( + litellm.AuthenticationError(message="invalid provider key", model="test-model", llm_provider="openai"), + litellm.RateLimitError(message="provider quota exceeded", model="test-model", llm_provider="openai"), + ), + ) + async def test_provider_client_status_still_alerts(self, exc: Exception) -> None: + assert (await self._alerted(exc)).called is True + class TestPostCallFailureHookProxyExceptionLogging: """A guardrail block raises a ProxyException; on an LLM route it must still @@ -4109,3 +4129,32 @@ def test_handle_exception_on_proxy_logs_bug_report_only_for_unmapped_500(caplog) assert provider_result.code == internal_result.code == "500" assert ISSUE_URL_BASE in caplog.text assert ISSUE_URL_BASE not in internal_result.message + + +def test_is_client_error_exception_budget_exceeded_is_client_error() -> None: + from litellm.exceptions import BudgetExceededError + from litellm.proxy.utils import _is_client_error_exception + + exc: Final = BudgetExceededError(current_cost=100.56, max_budget=100.0) + assert _is_client_error_exception(exc) is True + + +def test_is_client_error_exception_budget_exceeded_5xx_is_not_client_error() -> None: + from litellm.exceptions import BudgetExceededError + from litellm.proxy.utils import _is_client_error_exception + + exc: Final = BudgetExceededError(current_cost=100.56, max_budget=100.0) + exc.status_code = 500 + assert _is_client_error_exception(exc) is False + + +def test_is_client_error_exception_existing_behaviour_unchanged() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.utils import _is_client_error_exception + + assert _is_client_error_exception(HTTPException(status_code=400, detail="x")) is True + assert _is_client_error_exception(HTTPException(status_code=500, detail="x")) is False + assert _is_client_error_exception( + ProxyException(message="x", type="x", param=None, code="429") + ) is True + assert _is_client_error_exception(ValueError("x")) is False diff --git a/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py index c014ffd04c6..5200b67598d 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py @@ -1,5 +1,7 @@ +import asyncio import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -7,8 +9,14 @@ from fastapi.testclient import TestClient from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.database import get_database +from litellm.proxy.deps.repositories import get_user_banner_repository from litellm.proxy.proxy_server import app +from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import get_user_banner_audit +from litellm.repositories.database import Database +from litellm.repositories.user_banner_repository import UserBannerRepository from tests._master_key import MASTER_KEY +from tests._support.native_database import native_database client = TestClient(app) @@ -21,6 +29,8 @@ PUBLISHED_BANNER = { PUBLISH_BODY = {k: v for k, v in PUBLISHED_BANNER.items() if k != "revision"} DISABLED_BANNER = {"enabled": False, "message": "", "severity": "info", "revision": ""} +NATIVE: Final = native_database() + def _auth_override(role: LitellmUserRoles): async def override() -> UserAPIKeyAuth: @@ -46,7 +56,7 @@ def internal_user_auth(): @pytest.fixture def mock_audit_log(monkeypatch): audit_mock = AsyncMock(return_value=None) - monkeypatch.setattr("litellm.proxy.proxy_server.create_config_audit_log", audit_mock) + monkeypatch.setitem(app.dependency_overrides, get_user_banner_audit, lambda: audit_mock) return audit_mock @@ -54,7 +64,7 @@ def _mock_prisma(monkeypatch, record=None): mock_prisma = MagicMock() mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record) mock_prisma.db.litellm_uisettings.upsert = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: Database(prisma=mock_prisma, native=NATIVE)) return mock_prisma @@ -94,21 +104,77 @@ class TestGetUserBanner: assert response.json() == DISABLED_BANNER def test_no_database_returns_disabled_banner(self, admin_auth, monkeypatch): - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) response = client.get("/get/user_banner") assert response.status_code == 200 assert response.json() == DISABLED_BANNER class TestUpdateUserBanner: + def test_storage_finishes_before_audit_runs(self, admin_auth, monkeypatch, mock_audit_log): + events = [] + prisma = _mock_prisma(monkeypatch) + + async def save(**kwargs): + await asyncio.sleep(0) + events.append("saved") + + async def audit(before, after, user): + events.append("audited") + + prisma.db.litellm_uisettings.upsert.side_effect = save + mock_audit_log.side_effect = audit + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == 200 + assert events == ["saved", "audited"] + + @pytest.mark.parametrize( + ("role", "status", "detail"), + [ + (LitellmUserRoles.INTERNAL_USER, 403, "Only proxy admins can update the user banner."), + (LitellmUserRoles.PROXY_ADMIN, 500, "Database not connected. Please connect a database."), + ], + ) + def test_no_database_preserves_authorization_precedence(self, monkeypatch, role, status, detail): + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override(role)) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == status + assert response.json() == {"detail": detail} + + def test_failed_write_does_not_schedule_success_audit(self, admin_auth, monkeypatch, mock_audit_log): + prisma = _mock_prisma(monkeypatch) + prisma.db.litellm_uisettings.upsert.side_effect = RuntimeError("storage unavailable") + with pytest.raises(RuntimeError, match="storage unavailable"): + client.patch("/update/user_banner", json=PUBLISH_BODY) + mock_audit_log.assert_not_called() + + def test_audit_receives_previous_and_saved_banner(self, admin_auth, monkeypatch, mock_audit_log): + prisma = _mock_prisma(monkeypatch, record=SimpleNamespace(ui_settings=json.dumps(PUBLISHED_BANNER))) + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == 200 + before, after, user = mock_audit_log.call_args.args + assert before.model_dump() == PUBLISHED_BANNER + assert after.model_dump() == response.json()["banner"] + assert user.user_id == "test-user" + assert ( + json.loads(prisma.db.litellm_uisettings.upsert.await_args.kwargs["data"]["create"]["ui_settings"]) + == after.model_dump() + ) + def test_rejects_non_admin(self, internal_user_auth, monkeypatch): mock_prisma = _mock_prisma(monkeypatch) response = client.patch("/update/user_banner", json=PUBLISH_BODY) assert response.status_code == 403 mock_prisma.db.litellm_uisettings.upsert.assert_not_awaited() - def test_persists_and_round_trips(self, admin_auth, monkeypatch, mock_audit_log): - mock_prisma = _mock_prisma(monkeypatch, record=None) + @pytest.mark.parametrize("override_repository", [False, True]) + def test_persists_and_round_trips(self, admin_auth, monkeypatch, mock_audit_log, override_repository: bool): + mock_prisma: Final = _mock_prisma(monkeypatch, record=None) + if override_repository: + banners: Final = UserBannerRepository(mock_prisma) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + monkeypatch.setitem(app.dependency_overrides, get_user_banner_repository, lambda: banners) response = client.patch("/update/user_banner", json=PUBLISH_BODY) assert response.status_code == 200 diff --git a/tests/unit/proxy/utils/helpers/test_error_helpers.py b/tests/unit/proxy/utils/helpers/test_error_helpers.py index df399c8b1d2..b20857dbdaf 100644 --- a/tests/unit/proxy/utils/helpers/test_error_helpers.py +++ b/tests/unit/proxy/utils/helpers/test_error_helpers.py @@ -1,13 +1,20 @@ import asyncio import json +from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx +import openai import pytest from fastapi import HTTPException from prisma.errors import DataError +import litellm from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.common_request_processing import proxy_exception_from_route_error +from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.utils import get_error_message_str, handle_exception_on_proxy +from litellm.types.router import RouterErrors def normalize(value): @@ -96,7 +103,7 @@ def test_handle_exception_on_proxy_happy_path_http_exception(): assert snapshot == { "is_proxy_exception": True, "message": "forbidden", - "type": ProxyErrorTypes.internal_server_error.value, + "type": "auth_error", "code": "403", } @@ -156,7 +163,7 @@ def test_handle_exception_on_proxy_uses_attached_status_code_when_present(): assert snapshot == { "code": "418", "message": "teapot", - "type": ProxyErrorTypes.internal_server_error.value, + "type": "invalid_request_error", } @@ -176,6 +183,176 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500(): } +@pytest.mark.parametrize( + "status_code, expected_type", + [ + (400, "invalid_request_error"), + (401, "auth_error"), + (403, "auth_error"), + (404, "invalid_request_error"), + (408, "invalid_request_error"), + (409, "invalid_request_error"), + (422, "invalid_request_error"), + (429, "rate_limit_error"), + (500, ProxyErrorTypes.internal_server_error.value), + (503, ProxyErrorTypes.internal_server_error.value), + ], +) +def test_handle_exception_on_proxy_labels_an_http_exception_by_the_status_it_answers( + status_code: int, expected_type: str +): + """A management route that rejects a caller's own argument answers 400, and calling that + internal_server_error tells the caller's retry loop the gateway is down when nothing but + their request is wrong.""" + result: Final = handle_exception_on_proxy(HTTPException(status_code=status_code, detail="boom")) + + assert (result.code, result.type) == (str(status_code), expected_type) + + +def _exception_carrying_status(carried: int | None, message: str = "boom") -> Exception: + class _CarriedStatus(Exception): + status_code = carried + + return _CarriedStatus(message) + + +@pytest.mark.parametrize( + "status_code, expected_type", + [ + (400, "invalid_request_error"), + (404, "invalid_request_error"), + (408, "invalid_request_error"), + (429, "rate_limit_error"), + (500, ProxyErrorTypes.internal_server_error.value), + ], +) +def test_handle_exception_on_proxy_labels_a_carried_status_the_same_way(status_code: int, expected_type: str): + """Routes also reach this handler with plain exceptions carrying their own status, and + those must be named by that status rather than all landing on the 500 label.""" + result: Final = handle_exception_on_proxy(_exception_carrying_status(status_code)) + + assert (result.code, result.type) == (str(status_code), expected_type) + + +def test_handle_exception_on_proxy_keeps_the_type_the_exception_already_names(): + """A rate limit litellm raised already names itself, and overwriting that with the status + map's guess would drop what the client branches on.""" + result: Final = handle_exception_on_proxy( + litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-4o") + ) + + assert (result.code, result.type) == ("429", "throttling_error") + + +@pytest.mark.parametrize( + "exc", + [ + openai.InternalServerError( + message="upstream down", + response=httpx.Response(500, request=httpx.Request("GET", "https://api.openai.com/v1/batches/b")), + body={"type": "server_error"}, + ), + litellm.Timeout(message="took too long", model="gpt-6.1-sol", llm_provider="openai"), + litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-6.1-sol"), + litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-6.1-sol"), + litellm.BadRequestError(message="bad input", model="gpt-6.1-sol", llm_provider="openai"), + HTTPException(status_code=404, detail="no such batch"), + ], + ids=["provider_5xx_naming_its_own_type", "timeout", "rate_limit", "bad_key", "bad_request", "http_404"], +) +def test_handle_exception_on_proxy_labels_a_failure_the_way_the_chat_route_does(exc: Exception): + """Batch, file, and management routes answer through this handler while chat answers through the route + error path, and one client calling both must read the same type for the same failure.""" + shared: Final = handle_exception_on_proxy(exc) + chat: Final = proxy_exception_from_route_error(exc) + + assert (shared.code, shared.type) == (chat.code, chat.type) + + +@pytest.mark.parametrize( + "exc", + [ + HTTPException(status_code=400, detail="bad budget_duration"), + HTTPException(status_code=401, detail="not allowed"), + HTTPException(status_code=403, detail="admin only"), + HTTPException(status_code=404, detail="no such team"), + litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-6.1-sol"), + ValueError(f"{RouterErrors.no_deployments_with_tag_routing.value}=['paid']"), + _exception_carrying_status(409, "duplicate credential"), + ], + ids=["http_400", "http_401", "http_403", "http_404", "provider_401", "tag_denial", "carried_409"], +) +def test_handle_exception_on_proxy_labels_a_failure_the_way_route_catch_alls_do(exc: Exception): + """Management routes answer some failures through their own catch-all and the rest through this handler, + so a client reading /key/list and /key/generate must see the same type and status for the same failure.""" + shared: Final = handle_exception_on_proxy(exc) + catch_all: Final = proxy_exception_for(exc, default_status_code=500) + + assert (shared.code, shared.type, shared.message) == (catch_all.code, catch_all.type, catch_all.message) + + +@pytest.mark.parametrize( + "message, expected", + [ + (f"{RouterErrors.no_deployments_with_provider_budget_routing.value}, try later", ("429", "rate_limit_error")), + (f"{RouterErrors.no_deployments_with_tag_routing.value}=['paid']", ("401", "auth_error")), + ], + ids=["no_deployment_left", "no_deployment_for_tags"], +) +def test_handle_exception_on_proxy_names_a_router_failure_by_the_status_it_answers( + message: str, expected: tuple[str, str] +): + """A router failure that carries no status answers the 429 or 401 its message maps to, and a 429 labelled + internal_server_error reads as an outage to a client that branches on the type.""" + result: Final = handle_exception_on_proxy(ValueError(message)) + + assert (result.code, result.type) == expected + + +def test_handle_exception_on_proxy_keeps_the_path_a_rejected_request_names(): + """A provider rejecting a batch input names the caller's own bucket path, and redacting the /tmp part of it + leaves the caller nothing to fix their request with.""" + result: Final = handle_exception_on_proxy( + litellm.BadRequestError( + message="Invalid input_file_id gs://my-bucket/tmp/batch_input.jsonl", + model="gemini-3.1-pro-preview", + llm_provider="vertex_ai", + ) + ) + + assert result.code == "400" + assert "gs://my-bucket/tmp/batch_input.jsonl" in result.message + + +def test_handle_exception_on_proxy_keeps_the_message_an_exception_with_a_status_carries(): + """Main sends a status-bearing failure's message verbatim, and relabeling its type must not change that text.""" + result: Final = handle_exception_on_proxy(_exception_carrying_status(503, "Cannot connect to host 10.12.4.7:8443")) + + assert (result.code, result.message) == ("503", "Cannot connect to host 10.12.4.7:8443") + + +@pytest.mark.parametrize( + "message", + [ + "boom", + "Cannot connect to host 10.12.4.7:8443", + "Cannot connect to host guardrail.internal.svc.cluster.local:443", + "Cannot connect to host [fd00::12]:8443", + ], + ids=["plain", "ipv4", "cluster_dns", "ipv6"], +) +def test_handle_exception_on_proxy_answers_a_generic_500_for_an_exception_whose_status_is_none(message: str): + """A guardrail whose upstream refuses the connection raises with status_code=None, which crashed the handler + into a 503 blaming the auth database, and its text names private hosts no virtual key was ever sent.""" + result: Final = handle_exception_on_proxy(_exception_carrying_status(None, message)) + + assert (result.code, result.type, result.message) == ( + "500", + ProxyErrorTypes.internal_server_error.value, + "Internal server error", + ) + + @pytest.mark.parametrize( "exc", [ diff --git a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py index e2c1a7c2c09..dd933c508cb 100644 --- a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py +++ b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py @@ -362,7 +362,17 @@ def test_merge_guardrails_with_existing_creates_metadata_when_missing(): } -def test_merge_guardrails_with_existing_raises_on_unhashable_guardrail(): - data = {"metadata": {"guardrails": [{"unhashable": True}]}} - with pytest.raises(TypeError): - _merge_guardrails_with_existing(data, ["g1"]) +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_over_the_model_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + result = _merge_guardrails_with_existing({"metadata": {"guardrails": [configured]}}, ["pii-redact", "toxic-filter"]) + + assert result["metadata"]["guardrails"] == [configured, "toxic-filter"] + + +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_listed_after_the_key_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + result = _merge_guardrails_with_existing({"metadata": {"guardrails": ["pii-redact", configured]}}, ["toxic-filter"]) + + assert result["metadata"]["guardrails"] == [configured, "toxic-filter"] diff --git a/tests/unit/proxy/utils/helpers/test_model_access.py b/tests/unit/proxy/utils/helpers/test_model_access.py index 324360cc2ee..30d24e9eddc 100644 --- a/tests/unit/proxy/utils/helpers/test_model_access.py +++ b/tests/unit/proxy/utils/helpers/test_model_access.py @@ -1,10 +1,17 @@ +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Final from unittest.mock import MagicMock +import httpx import pytest +import respx from fastapi import HTTPException import litellm -from litellm import ModelResponse +from litellm import ModelResponse, Router from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ( create_model_info_response, @@ -650,3 +657,61 @@ async def test_listing_is_unchanged_without_an_agent_or_without_attached_groups( resolve_agent_ceiling=_agent_ceiling(None), ) assert (plain_key, agent_without_groups) == (["model-a", "model-b"], ["model-a", "model-b"]) + + +OPENROUTER_BASE: Final = "https://openrouter.ai/api/v1" + + +@dataclass(frozen=True, slots=True) +class _CatalogHeldUntilTheLoopRuns: + requested: threading.Event = field(default_factory=threading.Event) + released: threading.Event = field(default_factory=threading.Event) + + def answer(self, _request: httpx.Request) -> httpx.Response: + self.requested.set() + if not self.released.wait(timeout=30): + return httpx.Response(504) + return httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + + async def release_from_the_loop(self) -> None: + assert await asyncio.to_thread(self.requested.wait, 30) + self.released.set() + + +def test_get_available_models_for_user_fetches_a_live_catalog_without_holding_the_event_loop_or_its_default_pool( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + catalog: Final = _CatalogHeldUntilTheLoopRuns() + router: Final = Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": "sk-or-available-models-off-loop", + "api_base": OPENROUTER_BASE, + }, + } + ] + ) + + async def list_models_on_a_one_thread_default_pool() -> list[str]: + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + listing: Final = asyncio.ensure_future( + get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-key", user_id="user-1"), + llm_router=router, + general_settings={}, + user_model=None, + ) + ) + await catalog.release_from_the_loop() + return await listing + + with respx.mock() as respx_mock: + catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer) + models: Final = asyncio.run(list_models_on_a_one_thread_default_pool()) + + assert catalog_route.calls.last.response.status_code == 200 + assert models == ["openrouter/openai/gpt-5.4-mini"] diff --git a/tests/unit/proxy/utils/prisma_and_spend/conftest.py b/tests/unit/proxy/utils/prisma_and_spend/conftest.py index aefb95516af..afc282a14d4 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/unit/proxy/utils/prisma_and_spend/conftest.py @@ -139,6 +139,7 @@ def mock_prisma_client() -> MagicMock: client.db.is_connected = MagicMock(return_value=False) client.db.connect = AsyncMock() client.db.disconnect = AsyncMock() + client.db.stop_token_refresh_task = AsyncMock() client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}]) client.db.execute_raw = AsyncMock() client.db.tx = MagicMock() diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index 9aa57c7a19b..35241c4c664 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -14,12 +14,161 @@ from __future__ import annotations import asyncio import logging -from typing import Any -from unittest.mock import AsyncMock, MagicMock +import threading +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import _SHUTDOWN_RECONNECT_WAIT_SECONDS, PrismaClient + + +@pytest.mark.asyncio +@pytest.mark.parametrize("engine_dead", (False, True)) +async def test_shutdown_waits_for_hung_request_reconnect_to_time_out( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, engine_dead: bool +) -> None: + monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") + reconnecting: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + + async def hung_recreate(*args: object, **kwargs: object) -> None: + reconnecting.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + prisma_client.db.recreate_prisma_client = AsyncMock(side_effect=hung_recreate) + prisma_client._engine_confirmed_dead = engine_dead + with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock: + reconnect: Final = asyncio.create_task( + prisma_client.attempt_db_reconnect(reason="request", force=True, force_recreate=True) + ) + await reconnecting.wait() + closing: Final = asyncio.create_task(prisma_client.aclose()) + try: + await asyncio.sleep(0) + assert not closing.done() + prisma_client.db.disconnect.assert_not_awaited() + + clock.return_value = 1000.0 + prisma_client._db_watchdog_reconnect_timeout_seconds + 1 + assert await reconnect is False + await closing + finally: + reconnect.cancel() + closing.cancel() + await asyncio.gather(reconnect, closing, return_exceptions=True) + + assert cancelled.is_set() + assert not prisma_client._db_reconnect_lock.locked() + prisma_client.db.disconnect.assert_awaited_once() + assert await prisma_client.attempt_db_reconnect(reason="request", force=True) is False + prisma_client.db.recreate_prisma_client.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_shutdown_stops_waiting_for_a_held_reconnect_lock_after_the_fixed_limit( + prisma_client: PrismaClient, +) -> None: + assert _SHUTDOWN_RECONNECT_WAIT_SECONDS < prisma_client._db_watchdog_reconnect_timeout_seconds + await prisma_client._db_reconnect_lock.acquire() + with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock: + closing: Final = asyncio.create_task(prisma_client.aclose()) + try: + await asyncio.sleep(0) + assert not closing.done() + prisma_client.db.disconnect.assert_not_awaited() + + clock.return_value = 1000.0 + _SHUTDOWN_RECONNECT_WAIT_SECONDS + 1 + await closing + finally: + closing.cancel() + await asyncio.gather(closing, return_exceptions=True) + + prisma_client.db.disconnect.assert_awaited_once() + prisma_client._db_reconnect_lock.release() + await asyncio.sleep(0) + assert not prisma_client._db_reconnect_lock.locked() + + +@pytest.mark.asyncio +async def test_shutdown_rejects_a_reconnect_already_waiting_for_the_lock(prisma_client: PrismaClient) -> None: + await prisma_client._db_reconnect_lock.acquire() + reconnect: Final = asyncio.create_task( + prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True) + ) + await asyncio.sleep(0) + closing: Final = asyncio.create_task(prisma_client.stop_background_tasks()) + await asyncio.sleep(0) + prisma_client._db_reconnect_lock.release() + + await asyncio.wait_for(closing, timeout=2) + assert await reconnect is False + assert await prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True) is False + prisma_client.db.connect.assert_not_awaited() + prisma_client.db.disconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_spawned_engine_work_and_prevents_rearming(prisma_client: PrismaClient) -> None: + running: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + released: Final = asyncio.Event() + finished: Final = asyncio.Event() + + async def engine_work() -> None: + running.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + await released.wait() + finished.set() + + prisma_client._spawn_engine_task(engine_work()) + await asyncio.wait_for(running.wait(), timeout=2) + closing: Final = asyncio.create_task(prisma_client.stop_background_tasks()) + try: + await asyncio.wait_for(cancelled.wait(), timeout=2) + assert not finished.is_set() + assert not closing.done() + finally: + released.set() + await asyncio.wait_for(closing, timeout=2) + + assert finished.is_set() + prisma_client._handle_writer_engine_replaced() + await asyncio.sleep(0) + assert not prisma_client._engine_tasks + + +@pytest.mark.asyncio +async def test_shutdown_disconnect_returns_when_an_engine_wait_thread_never_exits( + patched_prisma_import: MagicMock, mock_prisma_client: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False) + client: Final = PrismaClient( + database_url="postgresql://test:test@localhost:5432/test", + proxy_logging_obj=MagicMock(failure_handler=AsyncMock()), + engine_exit_wait_seconds=0.05, + ) + client.db = mock_prisma_client.db + release: Final = threading.Event() + stuck: Final = threading.Thread(target=release.wait, daemon=True) + stuck.start() + client._engine_wait_threads = (stuck,) + client._closing = True + try: + await asyncio.wait_for(client.disconnect(), timeout=2) + finally: + release.set() + stuck.join(timeout=2) + + client.db.disconnect.assert_awaited_once() + assert client._engine_wait_threads == () + _PROBE_SQL = "SELECT to_regclass($1) IS NOT NULL AS present" @@ -448,3 +597,32 @@ async def test_start_view_setup_task_twice_keeps_the_first_task(prisma_client: P "second_is_first": second is first, } assert actual == {"first_is_task": True, "second_is_first": True} + + +@pytest.mark.asyncio +async def test_aclose_disconnects_and_logs_when_stopping_background_tasks_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + prisma: Final = MagicMock() + prisma.stop_background_tasks = AsyncMock(side_effect=RuntimeError("watchdog failed")) + prisma.disconnect = AsyncMock() + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + await PrismaClient.aclose(prisma) + + prisma.disconnect.assert_awaited_once() + assert "watchdog failed" in caplog.text + + +@pytest.mark.asyncio +async def test_aclose_returns_and_logs_when_disconnect_fails(caplog: pytest.LogCaptureFixture) -> None: + prisma: Final = MagicMock() + prisma.stop_background_tasks = AsyncMock() + prisma.disconnect = AsyncMock(side_effect=RuntimeError("engine hung")) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + await PrismaClient.aclose(prisma) + + prisma.stop_background_tasks.assert_awaited_once() + prisma.disconnect.assert_awaited_once() + assert "engine hung" in caplog.text diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 5640d7d94ec..5fd40696224 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -12,6 +12,7 @@ from __future__ import annotations import asyncio import json from collections.abc import AsyncGenerator, Iterator +from contextlib import AbstractContextManager, nullcontext from copy import deepcopy import logging from typing import Any, Callable, Dict, Final, List @@ -30,14 +31,20 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + get_logging_caching_headers, +) from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler from litellm.proxy.policy_engine.pipeline_executor import recorded_guardrail_information from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.caching.dual_cache import DualCache +from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, @@ -1163,6 +1170,831 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe assert seen["count"] == 1 +class _PipelineStepGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, event_hook: str | list[str]) -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=False) + self.calls = 0 + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: + self.calls += 1 + return data + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: + self.calls += 1 + + +def _two_step_pipeline_data( + mode: str, first_on_pass: str, first_on_fail: str = "block", second_on_fail: str = "block" +) -> dict[str, object]: + pipeline: Final = GuardrailPipeline( + mode=mode, + steps=[ + PipelineStep(guardrail="first", on_pass=first_on_pass, on_fail=first_on_fail), + PipelineStep(guardrail="second", on_pass="next", on_fail=second_on_fail), + ], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second"], + "_guardrail_pipelines": [("governance", pipeline)], + "_pipeline_managed_guardrails": {"first", "second"}, + }, + } + + +def _logging_obj(*, log_parked: bool, call_type: str = "acompletion") -> LiteLLMLoggingObj: + logging_obj: Final = LiteLLMLoggingObj( + model="m", + messages=[], + stream=False, + call_type=call_type, + start_time=None, + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.defer_async_logging = True + logging_obj.enqueue_deferred_logging = (lambda: None) if log_parked else None + return logging_obj + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first_on_pass", "expected_calls", "expected_header"), + [("allow", (1, 0), ["first"]), ("next", (1, 1), ["first", "second"])], +) +async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + first_on_pass: str, + expected_calls: tuple[int, int], + expected_header: list[str], +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", first_on_pass), + call_type="completion", + ) + + assert (first.calls, second.calls) == expected_calls + assert result["metadata"]["applied_guardrails"] == expected_header + + +@pytest.mark.asyncio +async def test_pre_call_hook_names_a_skipped_step_whose_guardrail_also_runs_outside_the_pipeline( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "allow"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 0) + assert result["metadata"]["applied_guardrails"] == ["first", "second"] + + +class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail): + def __init__(self, guardrail_name: str, error: Exception) -> None: + super().__init__(guardrail_name, "pre_call") + self.error = error + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: + self.calls += 1 + raise self.error + + +_PIPELINE_STEP_ERRORS: Final = ( + pytest.param(HTTPException(status_code=400, detail="flagged"), id="flagged"), + pytest.param(RuntimeError("vendor unreachable"), id="errored"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) +async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_on( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + error: Exception, +) -> None: + first: Final = _RaisingPipelineStepGuardrail("first", error) + second: Final = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 1) + assert result["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) +async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + error: Exception, +) -> None: + first: Final = _RaisingPipelineStepGuardrail("first", error) + second: Final = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), + call_type="completion", + ) + + with pytest.raises(type(error)): + await first.async_pre_call_deployment_hook( + kwargs={**result, "guardrails": ["first"]}, call_type=CallTypes.completion + ) + await second.async_pre_call_deployment_hook( + kwargs={**result, "guardrails": ["second"]}, call_type=CallTypes.completion + ) + + assert (first.calls, second.calls) == (2, 1) + + +class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: + self.calls += 1 + data["litellm_metadata"] = {} + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + return data + + +@pytest.mark.asyncio +async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_later_step_opened( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _OpensLitellmMetadataGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 1) + assert "applied_guardrails" not in result["metadata"] + assert result["litellm_metadata"]["applied_guardrails"] == ["second", "first"] + assert set(get_logging_caching_headers(result)["x-litellm-applied-guardrails"].split(",")) == {"first", "second"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("log_parked", "call_type", "bucket", "expected_after_response"), + [ + pytest.param(True, "acompletion", "metadata", ["first"], id="success-log-waits-for-the-response-hooks"), + pytest.param(True, "acompletion", "litellm_metadata", ["first"], id="success-log-waits-litellm-metadata"), + pytest.param(False, "acompletion", "metadata", ["first", "second"], id="success-log-already-built"), + pytest.param( + True, "allm_passthrough_route", "metadata", ["first", "second"], id="pass-through-headers-already-out" + ), + pytest.param( + True, "agenerate_content", "metadata", ["first", "second"], id="generate-content-adapter-logged-its-call" + ), + ], +) +async def test_post_call_pipeline_drops_the_steps_that_did_not_run_only_while_nothing_has_read_them( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + log_parked: bool, + call_type: str, + bucket: str, + expected_after_response: list[str], +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + request: Final = _two_step_pipeline_data("post_call", "allow") + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + "model": request["model"], + "messages": request["messages"], + bucket: request["metadata"], + "litellm_logging_obj": _logging_obj(log_parked=log_parked, call_type=call_type), + }, + call_type="completion", + ) + assert data[bucket]["applied_guardrails"] == ["first", "second"] + logged_metadata: Final = data[bucket].copy() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 0) + assert data[bucket]["applied_guardrails"] == expected_after_response + assert logged_metadata["applied_guardrails"] == expected_after_response + + +def _three_step_two_pipeline_data() -> dict[str, object]: + runs_second: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="first", on_pass="next", on_fail="block"), + PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + ], + ) + stops_before_second: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="third", on_pass="allow", on_fail="block"), + PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + ], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second", "third"], + "_guardrail_pipelines": [("runs-second", runs_second), ("stops-before-second", stops_before_second)], + "_pipeline_managed_guardrails": {"first", "second", "third"}, + }, + } + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")] + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + **_three_step_two_pipeline_data(), + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1] + assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"] + + +@pytest.mark.asyncio +async def test_held_stream_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")] + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + logging_obj: Final = _logging_obj(log_parked=False) + + async def stream_log(*args: object) -> None: + return None + + logging_obj._on_deferred_stream_complete = stream_log + auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions") + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, + data={ + **_three_step_two_pipeline_data(), + "stream": True, + "litellm_logging_obj": logging_obj, + }, + call_type="completion", + ) + + delivered: Final[list[object]] = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data + ) + ] + + assert len(delivered) == 2 + assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1] + assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"] + + +def _ends_before_a_later_pipeline_data(ending: str) -> dict[str, object]: + ends: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass=ending)]) + never_reached: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="second")]) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second"], + "_guardrail_pipelines": [("ends", ends), ("never-reached", never_reached)], + "_pipeline_managed_guardrails": {"first", "second"}, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("ending", "raised"), [("block", HTTPException), ("modify_response", ModifyResponseException)]) +async def test_post_call_pipeline_that_ends_the_request_drops_the_steps_of_a_later_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + ending: str, + raised: type[Exception], +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_ends_before_a_later_pipeline_data(ending), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + + with pytest.raises(raised): + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 0) + assert data["metadata"]["applied_guardrails"] == ["first"] + + +def _blocks_before_a_later_request_pipeline_data() -> dict[str, object]: + blocks: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass="block")]) + request: Final = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="gate", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "gate", "second"], + "_guardrail_pipelines": [("blocks", blocks), ("request", request)], + "_pipeline_managed_guardrails": {"first", "gate", "second"}, + }, + } + + +@pytest.mark.asyncio +async def test_post_call_pipeline_that_blocks_keeps_a_later_request_step_whose_guardrail_runs_outside_pipelines( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + gate: Final = _PipelineStepGuardrail("gate", "pre_call") + second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + monkeypatch.setattr(litellm, "callbacks", [first, gate, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_blocks_before_a_later_request_pipeline_data(), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert sorted(data["metadata"]["applied_guardrails"]) == ["first", "gate", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("ending", "stream_end"), [("block", pytest.raises(HTTPException)), ("modify_response", nullcontext())] +) +async def test_held_stream_pipeline_that_ends_the_stream_drops_the_steps_of_a_later_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + ending: str, + stream_end: AbstractContextManager[object], +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + logging_obj: Final = _logging_obj(log_parked=False) + + async def stream_log(*args: object) -> None: + return None + + logging_obj._on_deferred_stream_complete = stream_log + auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions") + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, + data={**_ends_before_a_later_pipeline_data(ending), "stream": True, "litellm_logging_obj": logging_obj}, + call_type="completion", + ) + + with stream_end: + async for _item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data + ): + pass + + assert (first.calls, second.calls) == (1, 0) + assert data["metadata"]["applied_guardrails"] == ["first"] + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_that_never_ran_whose_guardrail_ran_outside_the_pipeline( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("steps", "expected"), + [(("first", "second"), ["second"]), (("first",), None)], + ids=["skipped-then-ran", "only-step-skipped"], +) +async def test_post_call_pipeline_drops_a_step_its_guardrail_skips_on_a_whole_response( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + steps: tuple[str, ...], + expected: list[str] | None, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in steps] + guardrails[0].apply_stream_scope("streaming") + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline: Final = GuardrailPipeline( + mode="post_call", steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="block") for name in steps] + ) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": list(steps), + "_guardrail_pipelines": [("governance", pipeline)], + "_pipeline_managed_guardrails": set(steps), + }, + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert [guardrail.calls for guardrail in guardrails] == [0, 1][: len(steps)] + assert data["metadata"].get("applied_guardrails") == expected + + +class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: + self.calls += 1 + data["litellm_metadata"] = {"user_api_key_hash": "hashed"} + + +@pytest.mark.asyncio +async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_step_opens_litellm_metadata( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_two_step_pipeline_data("post_call", "next"), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_a_later_pipeline_runs_after_an_earlier_one_opens_litellm_metadata( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + stops_before_second: Final = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="first", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")], + ) + stops_before_first: Final = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="second", on_pass="allow"), PipelineStep(guardrail="first", on_pass="next")], + ) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second"], + "_guardrail_pipelines": [ + ("stops-before-second", stops_before_second), + ("stops-before-first", stops_before_first), + ], + "_pipeline_managed_guardrails": {"first", "second"}, + }, + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_names_every_pipeline_step_when_the_pipelines_run_on_another_request( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "allow"), + call_type="completion", + skip_guardrails=True, + ) + + assert (first.calls, second.calls) == (0, 0) + assert result["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_deployment_hook_names_a_skipped_pre_call_step_it_ran( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "allow"), + call_type="completion", + ) + + await second.async_pre_call_deployment_hook( + kwargs={**data, "guardrails": ["second"]}, call_type=CallTypes.completion + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + assert second.deployment_hook_ran(data) + assert not first.deployment_hook_ran(data) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) +async def test_deployment_hook_names_a_skipped_pre_call_step_that_raised( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + error: Exception, +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _RaisingPipelineStepGuardrail("second", error) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "allow"), + call_type="completion", + ) + + with pytest.raises(type(error)): + await second.async_pre_call_deployment_hook( + kwargs={**data, "guardrails": ["second"]}, call_type=CallTypes.completion + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_skipped_step_a_deployment_hook_ran( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + + await second.async_post_call_success_deployment_hook( + request_data={**data, "guardrails": ["second"]}, + response=litellm.ModelResponse(), + call_type=CallTypes.acompletion, + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_deployment_hook_leaves_a_request_without_policy_pipelines_untouched() -> None: + guardrail: Final = _PipelineStepGuardrail("second", "pre_call") + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": {"guardrails": ["second"]}, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.calls == 1 + assert kwargs["metadata"] == {"guardrails": ["second"]} + + +@pytest.mark.asyncio +async def test_deployment_hook_names_a_pipeline_step_it_ran_before_any_guardrail_was_named() -> None: + guardrail: Final = _PipelineStepGuardrail("second", "pre_call") + metadata: Final[dict[str, object]] = { + "guardrails": ["second"], + "_pipeline_managed_guardrails": frozenset({"second"}), + } + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": metadata, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.calls == 1 + assert metadata["applied_guardrails"] == ["second"] + + +@pytest.mark.asyncio +async def test_deployment_hook_names_a_pipeline_step_requested_only_in_the_other_metadata_bucket() -> None: + guardrail: Final = _PipelineStepGuardrail("second", "pre_call") + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": { + "guardrails": ["first"], + "applied_guardrails": ["first"], + "_pipeline_managed_guardrails": frozenset({"first", "second"}), + }, + "litellm_metadata": {"guardrails": ["second"]}, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.calls == 1 + assert sorted(get_logging_caching_headers(kwargs)["x-litellm-applied-guardrails"].split(",")) == ["first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("named_bucket", ["metadata", "litellm_metadata"]) +async def test_deployment_hook_names_a_pipeline_step_next_to_the_steps_already_named_when_both_buckets_exist( + named_bucket: str, +) -> None: + guardrail: Final = _PipelineStepGuardrail("second", "pre_call") + metadata: Final[dict[str, object]] = { + "guardrails": ["first", "second"], + "_pipeline_managed_guardrails": frozenset({"first", "second"}), + } + litellm_metadata: Final[dict[str, object]] = {} + {"metadata": metadata, "litellm_metadata": litellm_metadata}[named_bucket]["applied_guardrails"] = ["first"] + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": metadata, + "litellm_metadata": litellm_metadata, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.calls == 1 + assert sorted(get_logging_caching_headers(kwargs)["x-litellm-applied-guardrails"].split(",")) == ["first", "second"] + + +@pytest.mark.asyncio +async def test_deployment_hook_records_a_pipeline_step_it_ran_without_naming_one_the_request_did_not_list() -> None: + guardrail: Final = _PipelineStepGuardrail("second", "pre_call") + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": {"_pipeline_managed_guardrails": frozenset({"second"})}, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.deployment_hook_ran(kwargs) + assert not _PipelineStepGuardrail("first", "pre_call").deployment_hook_ran(kwargs) + assert "applied_guardrails" not in kwargs["metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("log_waits", "expected_after_stream"), + [(True, ["first"]), (False, ["first", "second"])], + ids=["stream-log-waits-for-the-pipeline", "stream-log-not-deferred"], +) +async def test_held_stream_pipeline_drops_the_steps_that_did_not_run_only_while_its_log_waits( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + log_waits: bool, + expected_after_stream: list[str], +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + logging_obj: Final = _logging_obj(log_parked=False) + + async def stream_log(*args: object) -> None: + return None + + logging_obj._on_deferred_stream_complete = stream_log if log_waits else None + auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions") + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, + data={**_two_step_pipeline_data("post_call", "allow"), "stream": True, "litellm_logging_obj": logging_obj}, + call_type="completion", + ) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + logged_metadata: Final = data["metadata"].copy() + + delivered: Final[list[object]] = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data + ) + ] + + assert len(delivered) == 2 + assert (first.calls, second.calls) == (1, 0) + assert data["metadata"]["applied_guardrails"] == expected_after_stream + assert logged_metadata["applied_guardrails"] == expected_after_stream + + @pytest.mark.asyncio async def test_post_call_pipeline_replacement_response_reaches_caller( proxy_logging, make_user_api_key_auth, monkeypatch diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index b7fc8c6dbeb..5013ecd98a2 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -170,8 +170,9 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox proxy_logging.slack_alerting_instance = MagicMock(alerting=None) invoked = {} - def fake_process(d): + def fake_process(d: object, *, skip_guardrails: bool) -> None: invoked["data"] = d + invoked["skip_guardrails"] = skip_guardrails proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment] out = await proxy_logging.pre_call_hook( @@ -180,7 +181,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox call_type="completion", ) assert out is data - assert invoked["data"] is data + assert invoked == {"data": data, "skip_guardrails": False} @pytest.mark.asyncio diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py new file mode 100644 index 00000000000..d9343a796d2 --- /dev/null +++ b/tests/unit/repositories/test_database.py @@ -0,0 +1,18 @@ +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.repositories.database import Database + + +@pytest.mark.asyncio +async def test_closing_the_database_closes_prisma() -> None: + prisma: Final = MagicMock() + prisma.aclose = AsyncMock() + database: Final = await Database.open(prisma) + + prisma.aclose.assert_not_awaited() + await database.aclose() + + prisma.aclose.assert_awaited_once() diff --git a/tests/unit/repositories/test_gateway_activity_repository.py b/tests/unit/repositories/test_gateway_activity_repository.py new file mode 100644 index 00000000000..e88702c5bdc --- /dev/null +++ b/tests/unit/repositories/test_gateway_activity_repository.py @@ -0,0 +1,47 @@ +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import ValidationError + +from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity +from litellm.repositories.gateway_activity_repository import GatewayActivityRepository + + +@pytest.mark.asyncio +async def test_gateway_activity_decodes_both_row_kinds_with_nullable_union_columns() -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[ + { + "kind": "route", "date": "2026-01-02", "category": "llm", "route": "/chat/completions", + "successful_requests": 7, "failed_requests": 2, "status_code": None, + }, + { + "kind": "status", "date": None, "category": None, "route": None, + "successful_requests": 0, "failed_requests": 2, "status_code": 429, + }, + ]) + result: Final = await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31") + assert result == ( + GatewayRouteActivity( + kind="route", date="2026-01-02", category="llm", route="/chat/completions", + successful_requests=7, failed_requests=2, + ), + GatewayStatusCodeActivity(kind="status", status_code=429, failed_requests=2), + ) + + +@pytest.mark.asyncio +async def test_gateway_activity_rejects_malformed_database_rows() -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[{"kind": "status", "status_code": None, "failed_requests": 2}]) + with pytest.raises(ValidationError): + await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31") + + +@pytest.mark.asyncio +async def test_gateway_activity_does_not_hide_database_failure() -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(side_effect=RuntimeError("reader unavailable")) + with pytest.raises(RuntimeError, match="reader unavailable"): + await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31") diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index 4e410e22b7f..56843640888 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -3,7 +3,7 @@ Tests for gateway repository layer. """ import json -from datetime import datetime +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any, Dict, Final, List, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -23,7 +23,11 @@ from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.credentials_repository import CredentialsRepository -from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository +from litellm.models.deployment_warning import AdminIgnore +from litellm.repositories.deployment_warning_ignore_repository import ( + DeploymentWarningIgnoreRepository, + parse_ignores, +) from litellm.repositories.model_repository import ModelRepository from litellm.repositories.object_permission_repository import ( ObjectPermissionRepository, @@ -1618,15 +1622,27 @@ class TestDeploymentWarningIgnoreRepository: } writer_table._records.update({row_id: dict(row) for row_id, row in other_rows.items()}) repo = DeploymentWarningIgnoreRepository(client, use_writer=True) + ignored_at = datetime(2026, 10, 14, 12, tzinfo=timezone.utc) + snoozed = AdminIgnore(ignored_by="admin", ignored_at=ignored_at, ignored_until=ignored_at + timedelta(hours=1)) - assert await repo.get_raw_ignores() is None - await repo.upsert_ignores('{"no_redis": {}}') - await repo.upsert_ignores('{"detailed_debug": {}}') + assert await repo.get_ignores() == {} + await repo.save_ignores({"no_redis": AdminIgnore(ignored_by="admin", ignored_at=ignored_at)}) + await repo.save_ignores({"detailed_debug": snoozed}) - assert await repo.get_raw_ignores() == '{"detailed_debug": {}}' + assert await repo.get_ignores() == {"detailed_debug": snoozed} assert {row_id: writer_table._records[row_id] for row_id in other_rows} == other_rows assert len(writer_table._records) == len(other_rows) + 1 + @pytest.mark.parametrize( + "stored", + [ + "{not json", + {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00", "ignored_until": None}}, + ], + ) + def test_unreadable_rows_read_as_no_ignores(self, stored): + assert parse_ignores(stored) == {} + class TestVerificationTokenRepositoryExtended: @pytest.fixture diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index ef198b487b9..5627167432d 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -23,6 +23,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( ) from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import Message, ModelResponse +from tests.unit.proxy.conftest import httpx_transport def test_response_api_handler_merges_metadata_and_service_tier_without_error(): @@ -593,6 +594,31 @@ def _anthropic_message(*content: Mapping[str, object]) -> httpx.Response: ) +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.parametrize("output", ["", [], [{"type": "input_text", "text": ""}]]) +def test_empty_responses_tool_output_reaches_anthropic_without_json_text( + respx_mock: respx.MockRouter, output: str | list[dict[str, str]] +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "ok"}) + ) + response: Final = litellm.responses( + model=_MODEL, + api_key="mock_api_key", + input=[ + {"role": "user", "content": "Check the weather"}, + {"type": "function_call", "call_id": "call_empty", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_empty", "output": output}, + ], + ) + sent: Final = json.loads(route.calls.last.request.content) + assert response.output_text == "ok" + assert sent["messages"][-1] == { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "call_empty", "content": ""}], + } + + def _replay(output: Sequence[BaseModel]) -> list[object]: return [*_FIRST_INPUT, *(item.model_dump(exclude_none=True) for item in output), *_NEXT_INPUT] @@ -1569,3 +1595,57 @@ async def test_previous_response_id_keeps_the_tools_an_earlier_codex_tool_search [("tool_result", "toolu_standup"), ("text", None)], ] assert prisma_client.db.queries == 1 + + +async def _codex_follow_up_body(**kwargs: object) -> Mapping[str, object]: + anthropic: Final = _RecordingRawAnthropicMessages() + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_SEARCH_TURN, _CALENDAR_TURN])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-calendar-turn", + input=[{"type": "function_call_output", "call_id": "toolu_standup", "output": "Created Standup."}], + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + **kwargs, + ) + return json.loads(anthropic.body) + + +@pytest.mark.asyncio +async def test_codex_tool_search_follow_up_rejects_text_verbosity_anthropic_does_not_support( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + anthropic: Final = _RecordingRawAnthropicMessages() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_SEARCH_TURN, _CALENDAR_TURN])), + pytest.raises(litellm.UnsupportedParamsError, match="verbosity"), + ): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-calendar-turn", + input=[{"type": "function_call_output", "call_id": "toolu_standup", "output": "Created Standup."}], + tools=[_CODEX_TOOL_SEARCH, _SHELL_TOOL], + text={"verbosity": "low"}, + api_key="sk-ant-fake", + num_retries=0, + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + assert anthropic.body == "" + + +@pytest.mark.asyncio +async def test_codex_tool_search_follow_up_text_verbosity_follows_provider_param_validation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + without_text: Final = await _codex_follow_up_body() + + dropped: Final = await _codex_follow_up_body(text={"verbosity": "low"}, drop_params=True) + allowed: Final = await _codex_follow_up_body(text={"verbosity": "low"}, allowed_openai_params=["verbosity"]) + + assert dropped == without_text + assert allowed == {**without_text, "verbosity": "low"} diff --git a/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py b/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py index 65800e42124..fce06598d29 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py +++ b/tests/unit/responses/litellm_completion_transformation/test_function_call_output_normalization.py @@ -50,9 +50,10 @@ def test_image_detail_mime_and_adjacent_text_survive_normalization(kind: str, de else "data:application/octet-stream;base64,AAAA", **({} if nested_url else metadata), } - parts: Final = [{"type": "input_text", "text": "Screenshot follows"}, image] + parts: Final = [{"type": "input_text", "text": ""}, {"type": "input_text", "text": "Screenshot follows"}, image] original: Final = deepcopy(parts) assert _normalize_media_content(kind, parts) == [ + *([{"type": "text", "text": ""}] if kind == "direct" else []), {"type": "text", "text": "Screenshot follows"}, {"type": "image_url", "image_url": {"url": "data:application/octet-stream;base64,AAAA", **metadata}}, ] @@ -97,6 +98,12 @@ def test_media_normalization_preserves_formats_references_and_metadata(kind: str ("output", "expected"), [ (None, ""), + ("", ""), + ([], ""), + ([{"type": "input_text", "text": ""}], ""), + ([{"type": kind, "text": ""} for kind in ("input_text", "output_text", "text")], ""), + ("[]", "[]"), + ([{"type": "input_text", "text": " \n"}], " \n"), ('{"ok":true}', '{"ok":true}'), ([{"type": "input_text", "text": "hello"}, {"type": "input_text", "text": " world"}], "hello world"), ({"type": "file", "status": "ready"}, '{"type": "file", "status": "ready"}'), @@ -106,6 +113,14 @@ def test_media_normalization_preserves_formats_references_and_metadata(kind: str ), ({"type": "input_image", "format": 17}, '{"type": "input_image", "format": 17}'), ([{"type": "unknown", "value": 17}], '[{"type": "unknown", "value": 17}]'), + ([{"type": "input_text", "text": 17}], '[{"type": "input_text", "text": 17}]'), + ([{"type": "input_text"}], '[{"type": "input_text"}]'), + ([{"type": "input_text", "text": None}], '[{"type": "input_text", "text": null}]'), + ( + [{"type": "input_text", "text": ""}, {"type": "unknown", "value": 17}], + '[{"type": "input_text", "text": ""}, {"type": "unknown", "value": 17}]', + ), + ([{"type": "input_text", "text": ""}, 17], '[{"type": "input_text", "text": ""}, 17]'), ( [{"type": "input_text", "text": "keep"}, {"type": "input_image", "file_id": "file-only"}], "keep", diff --git a/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_wire.py b/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_wire.py index f63088a606f..c7acc70ba95 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_wire.py +++ b/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_wire.py @@ -84,6 +84,34 @@ async def test_web_search_preview_is_sent_as_google_search() -> None: assert request_body["tools"] == [{"googleSearch": {}}] +@pytest.mark.asyncio +@pytest.mark.parametrize("output", ["", [], [{"type": "input_text", "text": ""}]]) +async def test_empty_responses_tool_output_reaches_gemini_without_json_text( + output: str | list[dict[str, str]], +) -> None: + with respx.mock() as mock_router: + route: Final = mock_router.post(url__regex=_GEMINI_URL).mock( + return_value=httpx.Response(status_code=200, json=_function_call_response("signature")) + ) + response: Final = await litellm.aresponses( + model="gemini/gemini-2.5-flash", + api_key="test-key", + input=[ + {"role": "user", "content": "Check the weather"}, + {"type": "function_call", "call_id": "call_empty", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_empty", "output": output}, + ], + ) + request_body: Final = _JSON_OBJECT.validate_json(route.calls.last.request.content) + + assert isinstance(response, ResponsesAPIResponse) + assert request_body["contents"] == [ + {"role": "user", "parts": [{"text": "Check the weather"}]}, + {"role": "model", "parts": [{"function_call": {"name": "get_weather", "args": {}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "get_weather", "response": {"content": ""}}}]}, + ] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("assistant_content", "assistant_parts"), diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ec5968c2a1a..c0d32742e8a 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1,17 +1,25 @@ import asyncio +import gc import json +import queue import sys +import threading import types -from collections.abc import Mapping, Sequence +import warnings +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager from datetime import datetime from types import MappingProxyType from typing import Final, cast -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from httpx import Response from mcp.types import Tool as MCPTool +from starlette.responses import StreamingResponse +from starlette.types import Message import litellm from litellm.integrations.custom_logger import CustomLogger @@ -25,7 +33,8 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.utils import ModelResponse, StandardLoggingPayload +from litellm.types.utils import ModelResponse, ModelResponseStream, StandardLoggingPayload, Usage +from litellm.utils import CustomStreamWrapper @pytest.mark.asyncio @@ -1556,6 +1565,7 @@ _CHAT_ROUND_TWO: Final = MappingProxyType( ), } ) +_NEVER_LOGGED_SECONDS: Final = 60 _CONTEXT_WINDOW_ERROR: Final = MappingProxyType( { "error": MappingProxyType( @@ -1575,16 +1585,31 @@ class _ChatRoundRecorder(CustomLogger): super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () + self._recorded: Final = asyncio.Condition() async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + await self._announce() async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + await self._announce() + + async def _announce(self) -> None: + async with self._recorded: + self._recorded.notify_all() + + async def settled(self, logged: int = 0, failed: int = 0) -> None: + async with self._recorded: + await asyncio.wait_for( + self._recorded.wait_for(lambda: len(self.logged) >= logged and len(self.failed) >= failed), + timeout=_NEVER_LOGGED_SECONDS, + ) + await _settled() def _chat_round_cost(body: Mapping[str, object]) -> float: @@ -1600,18 +1625,20 @@ async def _settled() -> None: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) -async def _auto_execute_chat_echo( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], call_id: str -) -> tuple[_ChatRoundRecorder, ModelResponse | Exception]: +def _serve_chat_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response] +) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() respx_mock.post("https://mcp-chat-rounds.example.com/v1/chat/completions").mock(side_effect=list(rounds)) - recorder: Final = _ChatRoundRecorder() - logging_obj: Final = Logging( + + +def _chat_rounds_logging_obj(recorder: _ChatRoundRecorder, call_id: str, stream: bool) -> Logging: + return Logging( model="openai/gpt-5.6", messages=[{"role": "user", "content": "Call echo with PING"}], - stream=False, + stream=stream, call_type="acompletion", start_time=datetime(2026, 1, 1), litellm_call_id=call_id, @@ -1619,6 +1646,10 @@ async def _auto_execute_chat_echo( dynamic_async_success_callbacks=[recorder], dynamic_async_failure_callbacks=[recorder], ) + + +@contextmanager +def _echo_tool_served() -> Iterator[None]: echo_tool: Final = MCPTool.model_validate( {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, by_name=False, @@ -1635,6 +1666,16 @@ async def _auto_execute_chat_echo( AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}]), ), ): + yield + + +async def _auto_execute_chat_echo( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], call_id: str +) -> tuple[_ChatRoundRecorder, ModelResponse | Exception]: + _serve_chat_rounds(respx_mock, monkeypatch, rounds) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, call_id, stream=False) + with _echo_tool_served(): try: response: Final = await litellm.acompletion( model="openai/gpt-5.6", @@ -1712,3 +1753,994 @@ async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follo assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19) assert recorder.failed[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) assert recorder.failed[0]["response_cost"] > 0 + + +_FIXED_MARGIN: Final = 0.001 + + +def _chat_stream_events(body: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + completed: Final = json.loads(json.dumps(body, default=dict)) + choice: Final = completed["choices"][0] + message: Final = choice["message"] + calls: Final = message.get("tool_calls") or [] + delta: Final = ( + {**message, "tool_calls": [{**call, "index": index} for index, call in enumerate(calls)]} if calls else message + ) + head: Final = {"id": completed["id"], "object": "chat.completion.chunk", "created": 1, "model": completed["model"]} + return ( + {**head, "choices": [{"index": 0, "delta": delta, "finish_reason": None}]}, + {**head, "choices": [{"index": 0, "delta": {}, "finish_reason": choice["finish_reason"]}]}, + {**head, "choices": [], "usage": completed["usage"]}, + ) + + +def _sse(events: Sequence[Mapping[str, object]]) -> bytes: + return b"".join(f"data: {json.dumps(event)}\n\n".encode() for event in events) + + +def _streamed(body: Mapping[str, object]) -> Response: + return Response( + 200, + content=_sse(_chat_stream_events(body)) + b"data: [DONE]\n\n", + headers={"content-type": "text/event-stream"}, + ) + + +class _CutOnceTheAnswerStarts(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + raise httpx.ReadError("connection reset mid-stream") + + +def _cut_once_the_answer_starts() -> Response: + return Response(200, stream=_CutOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}) + + +async def _open_chat_echo_stream( + logging_obj: Logging, metadata: dict[str, object] | None = None +) -> CustomStreamWrapper: + stream: Final = await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + api_key="sk-mcp-chat-rounds", + api_base="https://mcp-chat-rounds.example.com/v1", + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + metadata=metadata, + ) + assert isinstance(stream, CustomStreamWrapper) + return stream + + +async def _read_all(stream: CustomStreamWrapper) -> tuple[ModelResponseStream, ...]: + return tuple([chunk async for chunk in stream]) + + +async def _read_until(stream: CustomStreamWrapper, content: str) -> None: + async for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == content: + return + pytest.fail(f"the stream never sent {content!r}") + + +def _streamed_usage(chunks: Sequence[ModelResponseStream]) -> Usage: + usages: Final = tuple(usage for chunk in chunks if isinstance(usage := getattr(chunk, "usage", None), Usage)) + assert len(usages) == 1, usages + return usages[0] + + +def _assert_one_stream_row_bills_both_rounds_once( + logged: Sequence[StandardLoggingPayload], chunks: Sequence[ModelResponseStream] +) -> None: + charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN + assert len(logged) == 1 + row: Final = logged[0] + assert row["id"] == chunks[0].id + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert row["metadata"]["usage_object"]["cost"] == pytest.approx(row["response_cost"]) + usage: Final = _streamed_usage(chunks) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71 + 149, 19 + 5, 90 + 154) + assert usage.cost == pytest.approx(row["response_cost"]) + + +def _prompt_tokens_the_follow_up_sent(respx_mock: respx.MockRouter) -> int: + follow_up_request: Final = json.loads(respx_mock.calls[1].request.content) + return litellm.token_counter(model="gpt-5.6", messages=follow_up_request["messages"]) + + +def _assert_one_row_bills_round_one(logged: Sequence[StandardLoggingPayload]) -> None: + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (71, 19) + assert logged[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_read_to_the_end_logs_one_row_whose_usage_cost_counts_the_fixed_margin_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-read-to-the-end", stream=True) + + with _echo_tool_served(): + chunks: Final = await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_when_the_proxy_releases( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rounds-deferred", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_all(stream) + await _settled() + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await recorder.settled(logged=1) + + assert logged_before_release == () + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_rounds_on_the_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.exceptions import GuardrailRaisedException + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rounds-blocked", stream=True) + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_all(stream) + assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await recorder.settled(failed=1) + + charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71 + 149, 19 + 5) + assert recorder.failed[0]["response_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_rejected_by_the_proxy_during_the_follow_up_bills_round_one_with_its_own_breakdown( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from fastapi import HTTPException + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-rejected-mid-follow-up", stream=True) + rejected: Final = HTTPException(status_code=400, detail="rejected by a stream hook") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_until(stream, "PING") + assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, rejected) + await logging_obj.async_failure_handler(rejected, "rejected by a stream hook") + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + breakdown: Final = failed["cost_breakdown"] + assert (failed["prompt_tokens"], failed["completion_tokens"]) == (71, 19) + assert failed["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE)) + assert breakdown is not None + assert breakdown["total_cost"] == pytest.approx(failed["response_cost"]) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round_once_the_answer_finishes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_TWO),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-one-round", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + async for chunk in stream: + if chunk.choices and chunk.choices[0].finish_reason is not None: + break + await recorder.settled(logged=1) + logged_once_the_answer_finished: Final = recorder.logged + after_the_answer: Final = await _read_all(stream) + await _settled() + + assert len(logged_once_the_answer_finished) == 1 + assert recorder.logged == logged_once_the_answer_finished + assert recorder.logged[0]["id"] == _CHAT_ROUND_TWO["id"] + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (149, 5) + assert after_the_answer[-1:] and isinstance(after_the_answer[-1].usage, Usage) + usage: Final = _streamed_usage(after_the_answer) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + recorder.logged[0]["prompt_tokens"], + recorder.logged[0]["completion_tokens"], + recorder.logged[0]["total_tokens"], + ) + assert recorder.logged[0]["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_TWO)) + assert recorder.failed == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("closes", (True, False), ids=("aclose", "dropped")) +async def test_stream_auto_execute_bills_the_finished_first_round_when_the_caller_stops_reading_the_follow_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, closes: bool +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-caller-stops", stream=True) + + async def stop_reading_once_the_follow_up_answers() -> None: + stream: Final = await _open_chat_echo_stream(logging_obj) + await _read_until(stream, "PING") + if closes: + await stream.aclose() + + with _echo_tool_served(): + await stop_reading_once_the_follow_up_answers() + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_first_round_breaks_mid_answer_bills_only_its_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_cut_once_the_answer_starts(),)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-round-breaks", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + with pytest.raises(litellm.ServiceUnavailableError): + await _read_all(stream) + await recorder.settled(failed=1) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert recorder.failed[0]["completion_tokens"] > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging, _failure_fields_to_lift + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks", stream=True) + request_data: Final = MappingProxyType({"litellm_logging_obj": logging_obj}) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + await _read_all(stream) + if not ProxyLogging._discard_deferred_stream_logging_for_failure(dict(request_data), broken.value): + ProxyLogging.fire_deferred_stream_logging(dict(request_data)) + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + unmargined: Final = ( + litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=failed["prompt_tokens"], + completion_tokens=failed["completion_tokens"], + total_tokens=failed["total_tokens"], + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + - _FIXED_MARGIN + ) + assert failed["prompt_tokens"] == _prompt_tokens_the_follow_up_sent(respx_mock) + assert failed["completion_tokens"] > 0 + assert failed["response_cost"] == pytest.approx(unmargined) + request_failure: Final = _failure_fields_to_lift(request_data) + assert cast(StandardLoggingPayload, request_failure["standard_logging_object"])["id"] != recorder.logged[0]["id"] + assert request_failure["response_cost"] == pytest.approx(unmargined) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_whose_follow_up_breaks_reports_the_request_guardrail_result_on_the_failure_row_only( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks-guarded", stream=True) + metadata: Final[dict[str, object]] = {} + request_data: Final = {"metadata": metadata, "litellm_logging_obj": logging_obj} + content_filter: Final = CustomGuardrail(guardrail_name="request-content-filter", event_hook="post_call") + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj, metadata) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError) as broken: + await _read_all(stream) + content_filter.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=str(broken.value), + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + ) + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, broken.value): + ProxyLogging.fire_deferred_stream_logging(request_data) + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert recorder.logged[0]["guardrail_information"] is None + assert [entry["guardrail_name"] for entry in recorder.failed[0]["guardrail_information"] or ()] == [ + "request-content-filter" + ] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_served_through_the_proxy_stream_hook_logs_both_rounds_once_when_it_ends( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-proxy-stream-hook", stream=True) + + with _echo_tool_served(): + served: Final = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=await _open_chat_echo_stream(logging_obj), + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj}, + ) + chunks: Final = tuple([chunk async for chunk in served]) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_the_first_round_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + nothing_streamed: Final = Response(200, content=b"data: [DONE]\n\n", headers={"content-type": "text/event-stream"}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), nothing_streamed)) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-empty-follow-up", stream=True) + + with _echo_tool_served(): + await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +_ANSWER_WITH_UNREADABLE_AUDIO: Final = MappingProxyType( + { + **_CHAT_ROUND_TWO, + "choices": ( + MappingProxyType( + { + "index": 0, + "finish_reason": "stop", + "message": MappingProxyType( + { + "role": "assistant", + "content": "PING", + "audio": MappingProxyType({"id": "audio_echo", "data": "not base64", "transcript": "PING"}), + } + ), + } + ), + ), + } +) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with_the_usage_of_the_round_it_billed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_ANSWER_WITH_UNREADABLE_AUDIO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-unassembled", stream=True) + + with _echo_tool_served(): + chunks: Final = await _read_all(await _open_chat_echo_stream(logging_obj)) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert any(chunk.choices and chunk.choices[0].delta.content == "PING" for chunk in chunks) + usage: Final = _streamed_usage(chunks) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71, 19, 90) + assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +@pytest.mark.parametrize("answer", [_CHAT_ROUND_TWO, _ANSWER_WITH_UNREADABLE_AUDIO], ids=("assembled", "unassembled")) +async def test_stream_auto_execute_sends_its_usage_chunk_under_the_id_of_the_row_it_billed( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + deferred_for_a_post_call_guardrail: bool, + answer: Mapping[str, object], +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(answer))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-usage-chunk-id", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_all(stream) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await recorder.settled(logged=1) + + usage_chunks: Final = tuple(chunk for chunk in chunks if isinstance(getattr(chunk, "usage", None), Usage)) + answer_chunks: Final = tuple(chunk for chunk in chunks if chunk.choices and chunk.choices[0].delta.content == "PING") + assert len(recorder.logged) == 1 + assert [chunk.id for chunk in answer_chunks] == [_CHAT_ROUND_TWO["id"]] + assert len(usage_chunks) == 1 + assert usage_chunks[0].id == chunks[0].id == recorder.logged[0]["id"] == _CHAT_ROUND_ONE["id"] + + +def _assert_one_row_bills_round_one_and_the_partial_answer( + recorder: _ChatRoundRecorder, respx_mock: respx.MockRouter +) -> None: + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert row["id"] == _CHAT_ROUND_ONE["id"] + partial_prompt, partial_completion = (row["prompt_tokens"] - 71, row["completion_tokens"] - 19) + assert partial_prompt == _prompt_tokens_the_follow_up_sent(respx_mock) + assert partial_completion > 0 + partial_answer: Final = litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=partial_prompt, + completion_tokens=partial_completion, + total_tokens=partial_prompt + partial_completion, + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + assert row["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE) + partial_answer - _FIXED_MARGIN) + assert recorder.failed == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_and_the_partial_answer_on_one_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-mid-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _read_until(stream, "PING") + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +def _chat_rounds_router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "mcp-chat-rounds", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-mcp-chat-rounds", + "api_base": "https://mcp-chat-rounds.example.com/v1", + }, + } + ] + ) + + +class _StallsOnceTheAnswerStarts(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + await asyncio.Event().wait() + + +def _streams_content(body: object, content: str) -> bool: + lines: Final = body.decode().splitlines() if isinstance(body, bytes) else () + events: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: {")) + return any(event["choices"] and event["choices"][0]["delta"].get("content") == content for event in events) + + +async def _serve_until_the_client_hangs_up_after(response: StreamingResponse, content: str) -> None: + read: Final = asyncio.Event() + + async def receive() -> Message: + await read.wait() + return {"type": "http.disconnect"} + + async def send(message: Message) -> None: + if message["type"] == "http.response.body" and _streams_content(message.get("body"), content): + read.set() + + uvicorn_scope: Final = {"type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, "headers": []} + await response(uvicorn_scope, receive, send) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_answer_when_the_client_hangs_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, create_response + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds( + respx_mock, + monkeypatch, + ( + _streamed(_CHAT_ROUND_ONE), + Response(200, stream=_StallsOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"}), + ), + ) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-hang-up", stream=True) + + with _echo_tool_served(): + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + served: Final = await create_response( + proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ), + "text/event-stream", + {}, + ) + assert isinstance(served, StreamingResponse) + await _serve_until_the_client_hangs_up_after(served, "PING") + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +async def _stop_reading_once_the_proxy_sent(served: AsyncIterator[object], content: str) -> None: + async for frame in served: + if _streams_content(frame, content): + return + pytest.fail(f"the proxy never sent {content!r}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_routed_through_the_proxy_bills_the_partial_answer_when_hung_up_between_reads( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-hang-up-between-reads", stream=True) + + with _echo_tool_served(): + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + served: Final = proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ) + await _stop_reading_once_the_proxy_sent(served, "PING") + await served.aclose() + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +_LONGEST_WAIT_FOR_THE_LOOP_TO_MOVE_SECONDS: Final = 10 + + +class _CountsTheToolResultWhileTheLoopRuns: + def __init__(self, loop: asyncio.AbstractEventLoop, count: Callable[..., int]) -> None: + self._loop: Final = loop + self._count: Final = count + self._counting: Final = asyncio.Event() + self._loop_moved: Final = threading.Event() + self.loop_moved_while_counting: tuple[bool, ...] = () + + def __call__(self, model: str | None = None, messages: list[object] | None = None, **kwargs: object) -> int: + if any(isinstance(message, Mapping) and message.get("role") == "tool" for message in messages or ()): + self._loop.call_soon_threadsafe(self._counting.set) + moved: Final = self._loop_moved.wait(timeout=_LONGEST_WAIT_FOR_THE_LOOP_TO_MOVE_SECONDS) + self.loop_moved_while_counting = (*self.loop_moved_while_counting, moved) + return self._count(model=model, messages=messages, **kwargs) + + async def keep_the_loop_moving(self) -> None: + await self._counting.wait() + self._loop_moved.set() + + +@contextmanager +def _partial_count_watched(monkeypatch: pytest.MonkeyPatch) -> Iterator[_CountsTheToolResultWhileTheLoopRuns]: + from litellm.litellm_core_utils import streaming_chunk_builder_utils + + counter: Final = _CountsTheToolResultWhileTheLoopRuns( + asyncio.get_running_loop(), streaming_chunk_builder_utils.token_counter + ) + monkeypatch.setattr(streaming_chunk_builder_utils, "token_counter", counter) + other_task: Final = asyncio.create_task(counter.keep_the_loop_moving()) + try: + yield counter + finally: + other_task.cancel() + + +class _StallsAfterTheAnswerStarts(httpx.AsyncByteStream): + def __init__(self) -> None: + self.stalled: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield _sse(_chat_stream_events(_CHAT_ROUND_TWO)[:1]) + self.stalled.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +async def test_stream_auto_execute_cancelled_mid_answer_counts_the_partial_off_the_loop_and_bills_it_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + answer: Final = _StallsAfterTheAnswerStarts() + _serve_chat_rounds( + respx_mock, + monkeypatch, + (_streamed(_CHAT_ROUND_ONE), Response(200, stream=answer, headers={"content-type": "text/event-stream"})), + ) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-cancelled-off-the-loop", stream=True) + + with _echo_tool_served(), _partial_count_watched(monkeypatch) as counter: + stream: Final = await _open_chat_echo_stream(logging_obj) + reading: Final = asyncio.create_task(_read_all(stream)) + await asyncio.wait_for(answer.stalled.wait(), timeout=10) + reading.cancel() + with pytest.raises(asyncio.CancelledError): + await reading + await stream.aclose() + await recorder.settled(logged=1) + + assert counter.loop_moved_while_counting == (True,) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_routed_through_the_proxy_and_closed_between_reads_counts_the_partial_off_the_loop( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + router: Final = _chat_rounds_router() + monkeypatch.setattr(proxy_server, "llm_router", router) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-routed-closed-off-the-loop", stream=True) + + with _echo_tool_served(), _partial_count_watched(monkeypatch) as counter: + stream: Final = await router.acompletion( + model="mcp-chat-rounds", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + stream=True, + stream_options={"include_usage": True}, + litellm_logging_obj=logging_obj, + ) + served: Final = proxy_server.async_data_generator( + response=stream, + user_api_key_dict=UserAPIKeyAuth(), + request_data={"litellm_logging_obj": logging_obj, "model": "mcp-chat-rounds", "metadata": {}}, + ) + await _stop_reading_once_the_proxy_sent(served, "PING") + await served.aclose() + await recorder.settled(logged=1) + + assert counter.loop_moved_while_counting == (True,) + _assert_one_row_bills_round_one_and_the_partial_answer(recorder, respx_mock) + + +class _StallsBeforeTheAnswerStarts(httpx.AsyncByteStream): + def __init__(self) -> None: + self.opened: Final = asyncio.Event() + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b": the answer has not started\n\n" + self.opened.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_before_the_answer_starts_bills_round_one_alone_on_one_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + answer: Final = _StallsBeforeTheAnswerStarts() + _serve_chat_rounds( + respx_mock, + monkeypatch, + (_streamed(_CHAT_ROUND_ONE), Response(200, stream=answer, headers={"content-type": "text/event-stream"})), + ) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-before-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + reading: Final = asyncio.create_task(_read_all(stream)) + await asyncio.wait_for(answer.opened.wait(), timeout=10) + reading.cancel() + with pytest.raises(asyncio.CancelledError): + await reading + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_row_bills_round_one(recorder.logged) + assert recorder.failed == () + + +async def _read_through_the_usage_chunk( + chunks: AsyncIterator[ModelResponseStream], read: tuple[ModelResponseStream, ...] = () +) -> tuple[ModelResponseStream, ...]: + chunk: Final = await anext(chunks) + if isinstance(getattr(chunk, "usage", None), Usage): + return (*read, chunk) + return await _read_through_the_usage_chunk(chunks, (*read, chunk)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True]) +async def test_stream_auto_execute_hung_up_after_its_whole_answer_bills_both_rounds_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-after-answer", stream=True) + + with _echo_tool_served(): + stream: Final = await _open_chat_echo_stream(logging_obj) + if deferred_for_a_post_call_guardrail: + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + chunks: Final = await _read_through_the_usage_chunk(aiter(stream)) + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True + ) + await recorder.settled(logged=1) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks) + assert recorder.failed == () + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bridged_to_messages_bills_round_one_when_its_follow_up_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts())) + recorder: Final = _ChatRoundRecorder() + logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-bridged-to-messages", stream=True) + + async def bridged(stream: CustomStreamWrapper) -> AsyncIterator[ModelResponseStream]: + async for chunk in stream: + yield chunk + + with _echo_tool_served(): + relayed: Final = bridged(await _open_chat_echo_stream(logging_obj)) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=relayed, route_type="anthropic_messages", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + with pytest.raises(litellm.ServiceUnavailableError): + _ = [chunk async for chunk in relayed] + await recorder.settled(logged=1, failed=1) + + _assert_one_row_bills_round_one(recorder.logged) + + +def _read_synchronously_until_the_follow_up_answers( + logging_obj: Logging, loop: asyncio.AbstractEventLoop, before_dropping: Callable[[], None] = lambda: None +) -> None: + stream: Final = loop.run_until_complete(_open_chat_echo_stream(logging_obj)) + for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content == "PING": + before_dropping() + return + pytest.fail("the stream never sent 'PING'") + + +_CHAT_STREAM_BILLING_NAMES: Final = ("MCPStreamingIterator", "ChatStreamRoundBilling", "_log_success") + + +def _names_chat_stream_billing(text: str) -> bool: + return any(name in text for name in _CHAT_STREAM_BILLING_NAMES) + + +class _UnraisableRecorder: + def __init__(self) -> None: + self.raised: tuple[str, ...] = () + + def __call__(self, unraisable: "sys.UnraisableHookArgs") -> None: + self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") + + +def test_stream_auto_execute_read_synchronously_and_dropped_mid_answer_bills_the_finished_first_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + loop: Final = asyncio.new_event_loop() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) + + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped", stream=True), loop + ) + gc.collect() + loop.run_until_complete(recorder.settled(logged=1)) + finally: + asyncio.set_event_loop(None) + loop.close() + + _assert_one_row_bills_round_one(recorder.logged) + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] + + +def test_stream_auto_execute_read_synchronously_and_dropped_after_its_loop_closed_drops_without_errors( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO))) + recorder: Final = _ChatRoundRecorder() + unraisable: Final = _UnraisableRecorder() + loop: Final = asyncio.new_event_loop() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + asyncio.set_event_loop(loop) + + try: + with _echo_tool_served(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _read_synchronously_until_the_follow_up_answers( + _chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped-after-close", stream=True), + loop, + before_dropping=loop.close, + ) + gc.collect() + finally: + asyncio.set_event_loop(None) + loop.close() + + assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index 2f27622d77c..a16661c0174 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -4,6 +4,8 @@ import subprocess import sys import textwrap import types +from types import MappingProxyType +from collections.abc import Mapping, Sequence from contextlib import nullcontext from typing import Any, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -18,10 +20,14 @@ import litellm from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.mcp_server import operations as mcp_operations +from litellm.proxy.common_utils.user_api_key_cache import object_permission_cache_key from litellm.proxy.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.proxy.mcp_server.mcp_server_manager import ListedToolsCaller, MCPServerManager +from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTableCachedObj +from litellm.proxy.mcp_server.ui_session_utils import GrantedToolsetIds, granted_toolset_ids +from litellm.types.mcp_server.mcp_toolset import MCPToolset from litellm.proxy.utils import ProxyLogging from litellm.responses import main as responses_main from litellm.responses.mcp import litellm_proxy_mcp_handler as mcp_handler_module @@ -1575,61 +1581,309 @@ async def test_concurrent_bridge_calls_use_their_own_served_metadata(monkeypatch ProxyLogging._callback_capabilities_cache.clear() -def _toolset_gateway_manager(toolset_id: str, server_id: str) -> types.SimpleNamespace: - return types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - get_allowed_mcp_servers=AsyncMock(return_value=[]), - get_mcp_servers_from_ids=MagicMock(return_value=[]), - get_mcp_server_by_name=MagicMock(return_value=None), - get_toolset_by_name_cached=AsyncMock(return_value=types.SimpleNamespace(toolset_id=toolset_id)), - resolve_toolset_tool_permissions=AsyncMock(return_value={server_id: ["add"]}), - ) +TOOLSET_ROWS: Final = ( + MCPToolset( + toolset_id="ts-granted", toolset_name="team-toolset", tools=[{"server_id": "srv-1", "tool_name": "add"}] + ), + MCPToolset( + toolset_id="ts-other", toolset_name="other-toolset", tools=[{"server_id": "srv-2", "tool_name": "lookup"}] + ), + MCPToolset(toolset_id="ts-shouty", toolset_name="CALC", tools=[{"server_id": "srv-2", "tool_name": "lookup"}]), + MCPToolset(toolset_id="ts-group", toolset_name="reference", tools=[{"server_id": "srv-1", "tool_name": "add"}]), +) +GATEWAY_SERVERS: Final = ( + MCPServer(server_id="srv-1", name="calc", alias="calc", transport=MCPTransport.http, access_groups=["wiki"]), + MCPServer(server_id="srv-2", name="docs", alias="docs", transport=MCPTransport.http, access_groups=["reference"]), +) +UPSTREAM_TOOLS: Final = MappingProxyType( + { + "calc": ( + MCPTool(name="add", inputSchema={"type": "object"}), + MCPTool(name="subtract", inputSchema={"type": "object"}), + ), + "docs": (MCPTool(name="lookup", inputSchema={"type": "object"}),), + } +) -async def _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id: str) -> dict[str, object]: - from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.mcp_server.ui_session_utils import granted_toolset_ids +def _toolset_table(lookup_error: Exception | None = None) -> types.SimpleNamespace: + async def find_first(*, where: Mapping[str, str]) -> MCPToolset | None: + if lookup_error is not None: + raise lookup_error + return next((row for row in TOOLSET_ROWS if row.toolset_name == where["toolset_name"]), None) - mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=[], outcomes={})) - monkeypatch.setattr("litellm.proxy.mcp_server.server._get_tools_from_mcp_servers", mock_get_tools) - monkeypatch.setattr( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", - _toolset_gateway_manager("ts-granted", "srv-1"), - ) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + async def find_many(*, where: Mapping[str, Mapping[str, Sequence[str]]]) -> list[MCPToolset]: + return [row for row in TOOLSET_ROWS if row.toolset_id in where["toolset_id"]["in"]] + return types.SimpleNamespace(find_first=find_first, find_many=find_many) + + +def _team_grants(toolset_id: str) -> GrantedToolsetIds: async def team_permission(context: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable: - return LiteLLM_ObjectPermissionTable(object_permission_id="op-team", mcp_toolsets=[team_toolset_id]) + return LiteLLM_ObjectPermissionTable(object_permission_id="op-team", mcp_toolsets=[toolset_id]) async def granted_through_team(context: UserAPIKeyAuth) -> frozenset[str]: return await granted_toolset_ids(context, team_object_permission=team_permission, require_key_access=False) - team_key: Final = UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER) - await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( - user_api_key_auth=team_key, - mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/team-toolset"}], - granted_toolsets=granted_through_team, + return granted_through_team + + +async def _no_toolset_grants(context: UserAPIKeyAuth) -> frozenset[str]: + return frozenset() + + +async def _tools_listed_through_the_gateway( + monkeypatch: pytest.MonkeyPatch, + user_api_key_auth: UserAPIKeyAuth, + server_urls: tuple[str, ...], + granted_toolsets: GrantedToolsetIds, + toolset_table: types.SimpleNamespace | None = None, + agent_permission: LiteLLM_ObjectPermissionTable | None = None, +) -> list[str]: + async def upstream_listing(_client: object, server_name: str) -> list[MCPTool]: + return list(UPSTREAM_TOOLS[server_name]) + + manager: Final = MCPServerManager() + manager.registry = {server.server_id: server for server in GATEWAY_SERVERS} + manager.config_mcp_servers = dict(manager.registry) + manager.create_mcp_client = AsyncMock(return_value=object()) + manager._fetch_tools_with_timeout = AsyncMock(side_effect=upstream_listing) + tables: Final = types.SimpleNamespace(litellm_mcptoolsettable=toolset_table or _toolset_table()) + monkeypatch.setattr("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_operations, "global_mcp_server_manager", manager) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", types.SimpleNamespace(db=tables, writer_db=tables)) + cache: Final = DualCache() + await cache.async_set_cache(key="team_id:team-1", value=LiteLLM_TeamTableCachedObj(team_id="team-1")) + if agent_permission is not None: + await cache.async_set_cache( + key=f"agent_object_permission_id:{user_api_key_auth.agent_id}", + value=agent_permission.object_permission_id, + ) + await cache.async_set_cache( + key=object_permission_cache_key(agent_permission.object_permission_id), value=agent_permission + ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"supported_db_objects": []}) + try: + tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=user_api_key_auth, + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": url} for url in server_urls], + granted_toolsets=granted_toolsets, + ) + finally: + for server in GATEWAY_SERVERS: + manager._drop_listed_tools(server.server_id) + return sorted(tool.name for tool in tools) + + +def _key_holding(*server_ids: str, role: LitellmUserRoles = LitellmUserRoles.INTERNAL_USER) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-team", + team_id="team-1", + user_role=role, + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=list(server_ids)), ) - assert mock_get_tools.await_args is not None - return mock_get_tools.await_args.kwargs @pytest.mark.asyncio -async def test_toolset_gateway_url_scopes_a_team_granted_toolset_for_a_key_without_its_own_grant(monkeypatch): - kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-granted") - scoped = kwargs["user_api_key_auth"].object_permission - assert scoped is not None - assert scoped.mcp_servers == ["srv-1"] - assert scoped.mcp_tool_permissions == {"srv-1": ["add"]} - assert kwargs["mcp_servers"] is None +async def test_toolset_gateway_url_scopes_a_team_granted_toolset_for_a_key_without_its_own_grant( + monkeypatch: pytest.MonkeyPatch, +): + listed: Final = await _tools_listed_through_the_gateway( + monkeypatch, + UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER), + ("litellm_proxy/mcp/team-toolset",), + _team_grants("ts-granted"), + ) + assert listed == ["calc-add"] @pytest.mark.asyncio -async def test_toolset_gateway_url_skips_a_toolset_the_team_does_not_grant(monkeypatch): - kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-other") - assert kwargs["user_api_key_auth"].object_permission is None - assert kwargs["mcp_servers"] is None +@pytest.mark.parametrize( + "owner_role", + [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], + ids=["internal_user", "proxy_admin", "proxy_admin_viewer"], +) +async def test_toolset_gateway_url_answers_403_for_a_toolset_the_team_does_not_grant( + monkeypatch: pytest.MonkeyPatch, owner_role: LitellmUserRoles +): + with pytest.raises(HTTPException) as denial: + await _tools_listed_through_the_gateway( + monkeypatch, + UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=owner_role), + ("litellm_proxy/mcp/team-toolset",), + _team_grants("ts-other"), + ) + assert denial.value.status_code == 403 + assert "team-toolset" in str(denial.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_toolsets", + [("team-toolset", "other-toolset"), ("other-toolset", "team-toolset")], + ids=["granted_first", "denied_first"], +) +async def test_a_denied_toolset_beside_a_granted_one_adds_nothing_to_the_scoped_key( + monkeypatch: pytest.MonkeyPatch, requested_toolsets: tuple[str, ...] +): + listed: Final = await _tools_listed_through_the_gateway( + monkeypatch, + UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.PROXY_ADMIN), + tuple(f"litellm_proxy/mcp/{name}" for name in requested_toolsets), + _team_grants("ts-granted"), + ) + assert listed == ["calc-add"] + + +@pytest.mark.asyncio +async def test_a_whole_gateway_entry_beside_an_ungranted_toolset_still_lists_the_keys_own_servers( + monkeypatch: pytest.MonkeyPatch, +): + listed: Final = await _tools_listed_through_the_gateway( + monkeypatch, + _key_holding("srv-1", role=LitellmUserRoles.PROXY_ADMIN), + ("litellm_proxy", "litellm_proxy/mcp/other-toolset"), + _team_grants("ts-granted"), + ) + assert listed == ["calc-add", "calc-subtract"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("team_id", [None, UI_TEAM_ID], ids=["no_team", "dashboard_session"]) +async def test_an_admin_credential_without_a_team_keeps_a_toolset_no_grant_names( + monkeypatch: pytest.MonkeyPatch, team_id: str | None +): + listed: Final = await _tools_listed_through_the_gateway( + monkeypatch, + UserAPIKeyAuth(api_key="sk-admin", team_id=team_id, user_role=LitellmUserRoles.PROXY_ADMIN), + ("litellm_proxy/mcp/team-toolset",), + _team_grants("ts-other"), + ) + assert listed == ["calc-add"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lookup_error", [None, ConnectionError("database unreachable")], ids=["no_such_toolset", "toolset_lookup_fails"] +) +async def test_a_name_that_does_not_resolve_to_a_toolset_is_listed_as_a_server_name( + monkeypatch: pytest.MonkeyPatch, lookup_error: Exception | None +): + listed: Final = await _tools_listed_through_the_gateway( + monkeypatch, + _key_holding("srv-1", "srv-2"), + ("litellm_proxy", "litellm_proxy/mcp/wiki"), + _no_toolset_grants, + _toolset_table(lookup_error), + ) + assert listed == ["calc-add", "calc-subtract"] + + +@pytest.mark.asyncio +async def test_an_ungranted_toolset_named_like_a_server_the_key_holds_answers_403_once( + monkeypatch: pytest.MonkeyPatch, +): + with pytest.raises(HTTPException) as denial: + await _tools_listed_through_the_gateway( + monkeypatch, + _key_holding("srv-1", "srv-2"), + ("litellm_proxy/mcp/CALC", "litellm_proxy/mcp/CALC"), + _no_toolset_grants, + ) + assert denial.value.status_code == 403 + assert str(denial.value.detail).count("CALC") == 1 + + +@pytest.mark.asyncio +async def test_an_ungranted_toolset_named_like_a_server_adds_no_server_beside_a_named_one( + monkeypatch: pytest.MonkeyPatch, +): + listed: Final = await _tools_listed_through_the_gateway( + monkeypatch, + _key_holding("srv-1", "srv-2"), + ("litellm_proxy/mcp/docs", "litellm_proxy/mcp/CALC"), + _no_toolset_grants, + ) + assert listed == ["docs-lookup"] + + +@pytest.mark.asyncio +async def test_ungranted_toolsets_are_named_in_request_order_in_the_403(monkeypatch: pytest.MonkeyPatch): + with pytest.raises(HTTPException) as denial: + await _tools_listed_through_the_gateway( + monkeypatch, + _key_holding("srv-1", "srv-2"), + ("litellm_proxy/mcp/other-toolset", "litellm_proxy/mcp/CALC"), + _no_toolset_grants, + ) + assert denial.value.status_code == 403 + assert str(denial.value.detail).endswith("other-toolset, CALC'}") + + +@pytest.mark.asyncio +async def test_an_agent_bound_key_naming_an_ungranted_toolset_gets_the_generic_403(monkeypatch: pytest.MonkeyPatch): + with pytest.raises(HTTPException) as denial: + await _tools_listed_through_the_gateway( + monkeypatch, + _key_holding("srv-1", "srv-2").model_copy(update={"agent_id": "agent-1"}), + ("litellm_proxy/mcp/CALC",), + _no_toolset_grants, + ) + assert denial.value.status_code == 403 + assert denial.value.detail == {"error": "The key is not allowed to access the requested MCP servers: CALC"} + + +@pytest.mark.asyncio +async def test_an_access_group_the_key_holds_outranks_an_ungranted_toolset_of_the_same_name( + monkeypatch: pytest.MonkeyPatch, +): + listed: Final = await _tools_listed_through_the_gateway( + monkeypatch, + _key_holding("srv-1", "srv-2"), + ("litellm_proxy/mcp/reference",), + _no_toolset_grants, + ) + assert listed == ["docs-lookup"] + + +@pytest.mark.asyncio +async def test_an_access_group_the_keys_agent_does_not_grant_leaves_the_same_named_toolset_denied( + monkeypatch: pytest.MonkeyPatch, +): + key_holding_the_group: Final = UserAPIKeyAuth( + api_key="sk-team", + team_id="team-1", + user_role=LitellmUserRoles.INTERNAL_USER, + agent_id="agent-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", mcp_servers=["srv-1"], mcp_access_groups=["reference"] + ), + ) + with pytest.raises(HTTPException) as denial: + await _tools_listed_through_the_gateway( + monkeypatch, + key_holding_the_group, + ("litellm_proxy/mcp/reference",), + _no_toolset_grants, + agent_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-agent", mcp_servers=["srv-1"]), + ) + assert denial.value.status_code == 403 + assert denial.value.detail == {"error": "The key is not allowed to access the requested MCP servers: reference"} + + +@pytest.mark.asyncio +async def test_a_granted_toolset_outranks_an_access_group_of_the_same_name(monkeypatch: pytest.MonkeyPatch): + async def granted_reference_toolset(context: UserAPIKeyAuth) -> frozenset[str]: + return frozenset({"ts-group"}) + + listed: Final = await _tools_listed_through_the_gateway( + monkeypatch, + _key_holding("srv-1", "srv-2"), + ("litellm_proxy/mcp/reference",), + granted_reference_toolset, + ) + assert listed == ["calc-add"] @pytest.mark.asyncio diff --git a/tests/unit/responses/mcp/test_parked_stream_dispatch.py b/tests/unit/responses/mcp/test_parked_stream_dispatch.py new file mode 100644 index 00000000000..dd380c582c8 --- /dev/null +++ b/tests/unit/responses/mcp/test_parked_stream_dispatch.py @@ -0,0 +1,17 @@ +from typing import Final + +import pytest + +from litellm.responses.mcp.parked_stream_dispatch import parked_success +from litellm.types.utils import ModelResponse + +_ASSEMBLED: Final = ModelResponse(id="chatcmpl-parked", model="gpt-5.6") + + +@pytest.mark.parametrize( + "parked", + [(), (_ASSEMBLED,), (_ASSEMBLED, False, "not a logger"), (_ASSEMBLED, False, None, None)], + ids=("empty", "assembled-alone", "third-part-is-not-a-logger", "four-parts"), +) +def test_a_parked_dispatch_of_an_unknown_shape_is_never_dispatched(parked: tuple[object, ...]): + assert parked_success(parked) is None diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index de96e995dd4..210488f0888 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -554,7 +554,6 @@ class TestUseResponsesApiBridgeFlag: truncation="auto", metadata={"k": "v"}, background=False, - text={"verbosity": "low"}, api_base="https://example.databricks.test/serving-endpoints", api_key="fake-databricks-api-key", num_retries=0, @@ -1103,3 +1102,281 @@ async def test_chatgpt_completion_routes_to_responses_without_a_provider_catalog def test_missing_catalog_entry_does_not_force_other_providers_to_responses(provider: str) -> None: info, model = responses_api_bridge_check(model="test-chatgpt-routing-control", custom_llm_provider=provider) assert (info.get("mode"), model) == (None, "test-chatgpt-routing-control") + + +_GPT_5_6_CHUNKS: Final = ( + { + "id": "chatcmpl-123", + "object": "chat.completion.chunk", + "created": 1677652288, + "model": "gpt-5.6", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Answer"}, "finish_reason": None}], + }, + { + "id": "chatcmpl-123", + "object": "chat.completion.chunk", + "created": 1677652288, + "model": "gpt-5.6", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, +) + + +def _openai_chat_upstream(respx_mock: respx.MockRouter, stream: bool) -> respx.Route: + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions") + if stream: + return route.respond( + 200, + headers={"content-type": "text/event-stream"}, + text="".join(f"data: {json.dumps(chunk)}\n\n" for chunk in _GPT_5_6_CHUNKS) + "data: [DONE]\n\n", + ) + return route.respond( + 200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + + +async def _bridged_responses_call(asynchronous: bool, stream: bool, **kwargs: object) -> None: + if not asynchronous: + response = litellm.responses( + input="Hello", stream=stream, api_key="fake-provider-api-key", num_retries=0, **kwargs + ) + _ = list(response) if stream else response + return + response = await litellm.aresponses( + input="Hello", stream=stream, api_key="fake-provider-api-key", num_retries=0, **kwargs + ) + _ = [event async for event in response] if stream else response + + +@pytest.mark.parametrize("verbosity", ["low", "high"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_bridge_sends_text_verbosity_as_chat_verbosity( + verbosity: str, asynchronous: bool, stream: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = _openai_chat_upstream(respx_mock, stream) + + await _bridged_responses_call( + asynchronous, stream, model="openai/gpt-5.6", use_chat_completions_api=True, text={"verbosity": verbosity} + ) + + assert upstream.call_count == 1 + assert json.loads(upstream.calls[0].request.read()) == { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "Hello"}], + "verbosity": verbosity, + **({"stream": True, "stream_options": {"include_usage": True}} if stream else {}), + } + + +_JSON_SCHEMA_TEXT_FORMAT: Final = { + "type": "json_schema", + "name": "answer", + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + "strict": True, +} +_JSON_SCHEMA_RESPONSE_FORMAT: Final = { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + "strict": True, + }, +} + + +@pytest.mark.parametrize( + ("text", "expected_fields"), + [ + pytest.param( + {"verbosity": "low", "format": _JSON_SCHEMA_TEXT_FORMAT}, + {"response_format": _JSON_SCHEMA_RESPONSE_FORMAT, "verbosity": "low"}, + id="verbosity-and-format", + ), + pytest.param( + {"format": _JSON_SCHEMA_TEXT_FORMAT}, {"response_format": _JSON_SCHEMA_RESPONSE_FORMAT}, id="format-only" + ), + pytest.param(None, {}, id="no-text"), + ], +) +def test_bridge_maps_text_format_alongside_verbosity( + text: dict[str, object] | None, + expected_fields: dict[str, object], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + upstream: Final = _openai_chat_upstream(respx_mock, stream=False) + + litellm.responses( + model="openai/gpt-5.6", + input="Hello", + use_chat_completions_api=True, + api_key="fake-provider-api-key", + num_retries=0, + **({"text": text} if text is not None else {}), + ) + + assert upstream.call_count == 1 + assert json.loads(upstream.calls[0].request.read()) == { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "Hello"}], + **expected_fields, + } + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_bridge_rejects_text_verbosity_the_provider_does_not_support_when_drop_params_is_off( + asynchronous: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = _databricks_chat_upstream(respx_mock) + + with pytest.raises(litellm.UnsupportedParamsError) as raised: + await _bridged_responses_call( + asynchronous, + False, + model="databricks/my-custom-model", + text={"verbosity": "low"}, + api_base="https://example.databricks.test/serving-endpoints", + ) + + assert upstream.call_count == 0, "the provider must not be called with a verbosity it does not support" + assert raised.value.status_code == 400 + assert "verbosity" in str(raised.value) + + +@pytest.mark.parametrize( + ("call_kwargs", "global_drop_params", "expected_fields"), + [ + pytest.param({"drop_params": True}, False, {}, id="request-drop-params"), + pytest.param({}, True, {}, id="global-drop-params"), + pytest.param({"allowed_openai_params": ["verbosity"]}, False, {"verbosity": "low"}, id="allowed-openai-params"), + ], +) +def test_bridge_text_verbosity_follows_provider_param_validation( + call_kwargs: dict[str, object], + global_drop_params: bool, + expected_fields: dict[str, object], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + upstream: Final = _databricks_chat_upstream(respx_mock) + + response: Final = litellm.responses( + model="databricks/my-custom-model", + input="Hello", + text={"verbosity": "low"}, + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + **call_kwargs, + ) + + assert upstream.call_count == 1 + assert json.loads(upstream.calls[0].request.read()) == { + "model": "my-custom-model", + "messages": [{"role": "user", "content": "Hello"}], + **expected_fields, + } + assert response.output[0].content[0].text == "Answer" + + +def test_native_responses_route_keeps_text_verbosity_nested( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").respond( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + + litellm.responses( + model="openai/gpt-5.6", + input="Hello", + text={"verbosity": "low"}, + api_key="fake-provider-api-key", + num_retries=0, + ) + + assert upstream.call_count == 1 + assert json.loads(upstream.calls[0].request.read()) == { + "model": "gpt-5.6", + "input": "Hello", + "text": {"verbosity": "low"}, + } + + +_CLIENT_TOOL_SEARCH: Final = { + "type": "tool_search", + "execution": "client", + "description": "Search the deferred app tools.", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}, +} + + +@pytest.mark.parametrize("verbosity", ["low", "high"]) +@pytest.mark.asyncio +async def test_bridge_sends_text_verbosity_on_a_client_tool_search_follow_up_turn( + verbosity: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = _openai_chat_upstream(respx_mock, stream=False) + + await _bridged_responses_call( + True, + False, + model="openai/gpt-5.6", + use_chat_completions_api=True, + previous_response_id="chatcmpl-earlier-turn", + tools=[_CLIENT_TOOL_SEARCH], + text={"verbosity": verbosity}, + ) + + assert upstream.call_count == 1 + assert json.loads(upstream.calls[0].request.read()) == { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "Hello"}], + "tools": [ + { + "type": "function", + "function": { + "name": "tool_search", + "description": _CLIENT_TOOL_SEARCH["description"], + "parameters": _CLIENT_TOOL_SEARCH["parameters"], + }, + } + ], + "verbosity": verbosity, + } diff --git a/tests/unit/responses/test_responses_router_cooldown.py b/tests/unit/responses/test_responses_router_cooldown.py index e4bad083c87..078890c2395 100644 --- a/tests/unit/responses/test_responses_router_cooldown.py +++ b/tests/unit/responses/test_responses_router_cooldown.py @@ -77,7 +77,7 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown(): ) cooldown_ids = await async_get_cooldown_deployments( - litellm_router_instance=router, parent_otel_span=None + ctx=router._core_ctx, parent_otel_span=None ) assert failing_deployment_id in cooldown_ids, ( f"Responses API failure callback did not register cooldown for " diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index ab52e6d29c3..d0e26798090 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -8,15 +8,21 @@ Tests that: """ import asyncio +import itertools import json -from collections.abc import Mapping -from typing import Final -from unittest.mock import MagicMock +from collections.abc import AsyncIterator, Mapping +from datetime import datetime, timezone +from functools import partial +from typing import TYPE_CHECKING, Final +from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +import respx from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig +from litellm.llms.custom_llm import CustomLLM from litellm.llms.databricks.responses.transformation import ( DatabricksResponsesAPIConfig, ) @@ -42,6 +48,13 @@ from litellm.llms.volcengine.responses.transformation import ( VolcEngineResponsesAPIConfig, ) from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.utils import GenericStreamingChunk + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck class TestResponsesAPIWebSocketSupport: @@ -1591,7 +1604,7 @@ def _admission_refusing_the_next_turn(): from litellm.exceptions import RateLimitError - return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), None]) + return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), {}]) class _ClientSocket: @@ -1630,7 +1643,7 @@ class _StalledClientSocket(_ClientSocket): def _admission_crashing_the_next_turn(): from unittest.mock import AsyncMock - return AsyncMock(side_effect=[ValueError("custom check failed"), None]) + return AsyncMock(side_effect=[ValueError("custom check failed"), {}]) def _guard_blocking(word: str): @@ -1655,6 +1668,229 @@ def _guard_recording(bill_refusal: MagicMock): return guard +_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-34" +_TURNED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-35" +_OPENED_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) +_TURNED_AT: Final = datetime(2026, 1, 1, 12, 35, 5, tzinfo=timezone.utc) +_WS_DEPLOYMENT: Final = { + "litellm_params": {"model": "openai/gpt-5-mini", "otpm": 4096}, + "model_info": {"id": "ws-dep"}, +} +_OTHER_USAGE: Final = 10 +_OUTPUT_TOKENS_PER_TURN: Final = 2 +_UPSTREAM_ERROR: Final = json.dumps({"type": "error", "error": {"type": "server_error", "message": "upstream failed"}}) + + +def _turn_completed() -> str: + return json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_turn", + "status": "completed", + "output": [], + "usage": {"input_tokens": 5, "output_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + }, + } + ) + + +async def _socket_opened_on(dual_cache: "DualCache") -> Mapping[str, object]: + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + + await dual_cache.async_increment_cache(key=_OPENED_ON, value=_OTHER_USAGE, ttl=60) + await dual_cache.async_increment_cache(key=_TURNED_ON, value=_OTHER_USAGE, ttl=60) + return await async_admit_io_token_turn(dual_cache, _WS_DEPLOYMENT, _OPENED_AT) + + +def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: finished_at + ) + + +def _later_turns_admitted_through(dual_cache: "DualCache") -> "ResponsesWsTurnAdmission": + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + + return partial(async_admit_io_token_turn, dual_cache, _WS_DEPLOYMENT, _TURNED_AT) + + +_ITPM_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:itpm:12-34" +_ITPM_DEPLOYMENT: Final = { + "litellm_params": {"model": "openai/gpt-5-mini", "itpm": 40_000}, + "model_info": {"id": "ws-dep"}, +} +_TURN_INPUT_TOKENS: Final = 1_000 + + +def _turn_completed_reading_cache(cached_tokens: int) -> str: + return json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_cached_turn", + "status": "completed", + "output": [], + "usage": { + "input_tokens": _TURN_INPUT_TOKENS, + "input_tokens_details": {"cached_tokens": cached_tokens}, + "output_tokens": _OUTPUT_TOKENS_PER_TURN, + "total_tokens": _TURN_INPUT_TOKENS + _OUTPUT_TOKENS_PER_TURN, + }, + }, + } + ) + + +def _connection_logged_through(rate_limits: "ModelRateLimitingCheck", opened_with: Mapping[str, object]) -> "Logging": + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CallTypes + + connection: Final = Logging( + model="gpt-5-mini", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="connection-call", + function_id="connection-call", + dynamic_async_success_callbacks=[rate_limits], + dynamic_async_failure_callbacks=[rate_limits], + ) + connection.update_environment_variables( + litellm_params={"metadata": dict(opened_with), CallTypes.aresponses_websocket.value: True}, + optional_params={}, + ) + return connection + + +def _provider_turn_stream(_request: httpx.Request) -> httpx.Response: + completed: Final = { + "id": "resp_turn", + "object": "response", + "created_at": 1767271505, + "status": "completed", + "model": "gpt-5-mini", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 5, "output_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + } + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join( + f"event: {kind}\ndata: {json.dumps({'type': kind, 'response': completed})}\n\n" + for kind in ("response.created", "response.completed") + ), + ) + + +class _ProviderResponsesSocket: + """A provider's native Responses socket that answers each ``response.create`` with a completed turn.""" + + def __init__(self) -> None: + self.sent: Final[list[dict[str, object]]] = [] + self._replies: Final[asyncio.Queue[str]] = asyncio.Queue() + + def connect(self, _url: str, **_kwargs: object) -> "_ProviderResponsesSocket": + return self + + async def __aenter__(self) -> "_ProviderResponsesSocket": + return self + + async def __aexit__(self, *_exc_info: object) -> None: + return None + + async def send(self, message: str) -> None: + self.sent.append(json.loads(message)) + await self._replies.put(_turn_completed()) + + async def recv(self, decode: bool = False) -> str: + return await self._replies.get() + + async def close(self) -> None: + return None + + +class _TwoOutputTokenTurns(CustomLLM): + def __init__(self, failing_turns: frozenset[int] = frozenset()) -> None: + super().__init__() + self._failing_turns: Final = failing_turns + self._turns: Final = itertools.count(1) + + async def astreaming(self, *args: object, **kwargs: object) -> AsyncIterator[GenericStreamingChunk]: + if next(self._turns) in self._failing_turns: + raise RuntimeError("upstream failed") + yield GenericStreamingChunk( + text="hi", + is_finished=True, + finish_reason="stop", + index=0, + tool_use=None, + usage={"prompt_tokens": 5, "completion_tokens": _OUTPUT_TOKENS_PER_TURN, "total_tokens": 7}, + ) + + +async def _run_managed_relay( + monkeypatch: pytest.MonkeyPatch, + provider: CustomLLM, + rate_limits: "ModelRateLimitingCheck", + frames: tuple[str, ...], + turn_admission: "ResponsesWsTurnAdmission | None", + opened_with: Mapping[str, object], +) -> _ClientSocket: + import litellm + + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "ws-turn-llm", "custom_handler": provider}]) + return await _relay_turns_to(monkeypatch, "ws-turn-llm/turns", rate_limits, frames, turn_admission, opened_with) + + +async def _relay_turns_to( + monkeypatch: pytest.MonkeyPatch, + model: str, + rate_limits: "ModelRateLimitingCheck", + frames: tuple[str, ...], + turn_admission: "ResponsesWsTurnAdmission | None", + opened_with: Mapping[str, object], + **connection: object, +) -> _ClientSocket: + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler, drain_logging_tasks + + monkeypatch.setattr(litellm, "callbacks", [rate_limits]) + first, *later = frames + websocket: Final = _ClientSocket(*later) + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model=model, + logging_obj=Logging( + model=model, + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + litellm_metadata={"user_api_key_hash": "k1", **opened_with}, + first_message=first, + turn_admission=turn_admission, + **connection, + ) + + await handler.run() + await drain_logging_tasks() + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + return websocket + + def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]: return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"] @@ -1722,6 +1958,195 @@ class TestWebSocketTurnAdmission: assert isinstance(failure, RateLimitError) assert "crossed budget" in str(failure) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("first_turn_end", "first_turn_tokens"), + [ + pytest.param(_turn_completed(), _OUTPUT_TOKENS_PER_TURN, id="first_turn_completed"), + pytest.param(_UPSTREAM_ERROR, 0, id="first_turn_failed"), + ], + ) + async def test_native_relay_counts_every_turn_on_the_minute_it_finishes( + self, first_turn_end: str, first_turn_tokens: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.caching.dual_cache import DualCache + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _TURNED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await _socket_opened_on(dual_cache) + first, second, third = _turn_frames() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[first_turn_end, _turn_completed(), _turn_completed(), Exception("stop")] + ) + handler: Final = _make_streaming( + websocket=_ClientSocket(second, third), + backend_ws=backend_ws, + request_data={}, + logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + first_message=first, + turn_admission=_later_turns_admitted_through(dual_cache), + ) + + await handler.client_to_backend() + await handler.backend_to_client() + await drain_logging_tasks() + + assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert ( + await dual_cache.async_get_cache(key=_TURNED_ON) + == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("cached_tokens", "counted_input_tokens"), + [ + pytest.param(800, _TURN_INPUT_TOKENS - 800, id="cached_input"), + pytest.param(0, _TURN_INPUT_TOKENS, id="uncached_input"), + ], + ) + async def test_native_relay_leaves_a_turn_s_cached_input_out_of_the_deployment_s_itpm( + self, cached_tokens: int, counted_input_tokens: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.caching.dual_cache import DualCache + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _OPENED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await async_admit_io_token_turn(dual_cache, _ITPM_DEPLOYMENT, _OPENED_AT) + first, _, _ = _turn_frames() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=[_turn_completed_reading_cache(cached_tokens), Exception("stop")]) + handler: Final = _make_streaming( + websocket=_ClientSocket(), + backend_ws=backend_ws, + request_data={}, + logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + first_message=first, + ) + + await handler.client_to_backend() + await handler.backend_to_client() + await drain_logging_tasks() + + assert await dual_cache.async_get_cache(key=_ITPM_OPENED_ON) == counted_input_tokens + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("failing_turns", "first_turn_tokens"), + [ + pytest.param(frozenset(), _OUTPUT_TOKENS_PER_TURN, id="first_turn_completed"), + pytest.param(frozenset({1}), 0, id="first_turn_failed"), + ], + ) + async def test_managed_relay_counts_every_turn_on_the_minute_it_finishes( + self, monkeypatch: pytest.MonkeyPatch, failing_turns: frozenset[int], first_turn_tokens: int + ) -> None: + from litellm.caching.dual_cache import DualCache + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + _turns_finish_at(monkeypatch, _TURNED_AT) + dual_cache: Final = DualCache() + opened_with: Final = await _socket_opened_on(dual_cache) + + await _run_managed_relay( + monkeypatch, + _TwoOutputTokenTurns(failing_turns), + ModelRateLimitingCheck(dual_cache), + _turn_frames(), + _later_turns_admitted_through(dual_cache), + opened_with, + ) + + assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert ( + await dual_cache.async_get_cache(key=_TURNED_ON) + == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN + ) + + @pytest.mark.asyncio + async def test_managed_relay_sends_the_provider_only_the_metadata_each_turn_carries( + self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter + ) -> None: + import litellm + from litellm.caching.dual_cache import DualCache + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + _turns_finish_at(monkeypatch, _TURNED_AT) + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=_provider_turn_stream) + client_metadata: Final = {"session": "abc", "team": "search"} + dual_cache: Final = DualCache() + + await _relay_turns_to( + monkeypatch, + "openai/gpt-5-mini", + ModelRateLimitingCheck(dual_cache), + tuple( + json.dumps({"type": "response.create", "input": word, "metadata": client_metadata}) + for word in ("one", "two") + ), + _later_turns_admitted_through(dual_cache), + await _socket_opened_on(dual_cache), + api_key="sk-fake", + ) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [ + client_metadata, + client_metadata, + ] + assert await dual_cache.async_get_cache(key=_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN + + @pytest.mark.asyncio + async def test_native_relay_sends_the_provider_only_the_metadata_the_socket_opened_with( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from unittest.mock import patch + + from litellm import Router + + monkeypatch.setenv("LITELLM_RUST", "false") + provider: Final = _ProviderResponsesSocket() + client_metadata: Final = {"session": "abc", "team": "search"} + first, second, _ = _turn_frames() + router: Final = Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": { + "model": "openai/gpt-5-mini", + "api_key": "sk-fake", + "itpm": 40_000, + "otpm": 4096, + }, + "model_info": {"id": "ws-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + with patch("websockets.connect", provider.connect): + await router._aresponses_websocket( + model="ws-io", + websocket=_ClientSocket(second), + first_message=first, + metadata=dict(client_metadata), + ) + + assert [frame.get("metadata") for frame in provider.sent] == [client_metadata, client_metadata] + @pytest.mark.asyncio async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch: pytest.MonkeyPatch): from unittest.mock import AsyncMock @@ -1875,7 +2300,7 @@ class TestWebSocketTurnAdmission: request_data={}, first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) close_time_failure: Final = await handler.bidirectional_forward() @@ -1904,7 +2329,7 @@ class TestWebSocketTurnAdmission: request_data={}, first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) relay: Final = asyncio.ensure_future(handler.bidirectional_forward()) await asyncio.wait_for(client.send_started.wait(), timeout=1) @@ -1955,7 +2380,7 @@ class TestWebSocketTurnAdmission: ), first_message=first, frame_guard=_guard_recording(bill_refusal), - turn_admission=AsyncMock(side_effect=[refusal, None]), + turn_admission=AsyncMock(side_effect=[refusal, {}]), ) close_time_failure: Final = await handler.run() diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 7445070bd3a..14ecb185f06 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -166,21 +166,24 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): assert not isinstance(exc_info.value, MidStreamFallbackError) -def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(): +@pytest.mark.parametrize( + "message", + ( + "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", + "Your input exceeds the context window of this model. Please adjust your input and try again.", + ), +) +def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(message: str) -> None: """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming path and, being a non-retriable client error, is raised directly rather than wrapped for mid-stream fallback, preserving the direct-SDK 400 contract from issue #15785.""" - iterator = _make_iterator() - chunk = _make_error_chunk( - "invalid_request_error", - "context_length_exceeded", - "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", - ) + iterator: Final = _make_iterator() + chunk: Final = _make_error_chunk("invalid_request_error", "context_length_exceeded", message) with pytest.raises(litellm.ContextWindowExceededError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) - assert "maximum context length" in str(exc_info.value) + assert message in str(exc_info.value) CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." diff --git a/tests/unit/router_strategy/complexity_router/test_context_compaction.py b/tests/unit/router_strategy/complexity_router/test_context_compaction.py index c6242f775c8..850accf574b 100644 --- a/tests/unit/router_strategy/complexity_router/test_context_compaction.py +++ b/tests/unit/router_strategy/complexity_router/test_context_compaction.py @@ -11,6 +11,7 @@ import respx import litellm from litellm.llms import compaction as native +from litellm.router_api import RoutingContext from litellm.router_strategy.complexity_router.config import ContextCompactionConfig from litellm.router_strategy.complexity_router.context_compaction import ( CompactionState, @@ -184,7 +185,7 @@ async def test_all_surfaces_compact_and_keep_selected_answerer( protocol: native.CompactionProtocol, request: Mapping[str, object], parent_model: str | None = None ) -> Mapping[str, object]: assert parent_model == "auto" - result: Final = await native.dispatch(router, protocol, request) + result: Final = await native.dispatch(RoutingContext(router), protocol, request) captured.put_nowait(result) return result @@ -385,9 +386,9 @@ async def test_client_owned_history_bypasses_compaction( counted: Final = router._count_pre_call_check_tokens(None, request["input"], request) if arm_first and counted > 512: with pytest.raises(litellm.ContextWindowExceededError): - await compact_to_fit(router, deployment.model_dump(), request, "responses") + await compact_to_fit(RoutingContext(router), deployment.model_dump(), request, "responses") else: - result: Final = await compact_to_fit(router, deployment.model_dump(), request, "responses") + result: Final = await compact_to_fit(RoutingContext(router), deployment.model_dump(), request, "responses") assert result is request assert {key: value for key, value in request.items() if key != "_context_compaction_state"} == original assert all(route.call_count == 0 for route in wire) @@ -454,15 +455,16 @@ async def test_unusable_output_budget_still_enforces_known_input_window( request: Final = {**payload, "model": "small", "max_tokens": output, "_context_compaction_state": state} if overflow: with pytest.raises(litellm.BadRequestError, match="known input window"): - await compact_to_fit(router, deployment.model_dump(), request, "chat") + await compact_to_fit(RoutingContext(router), deployment.model_dump(), request, "chat") else: - assert await compact_to_fit(router, deployment.model_dump(), request, "chat") is request + assert await compact_to_fit(RoutingContext(router), deployment.model_dump(), request, "chat") is request assert all(route.call_count == 0 for route in wire) @pytest.mark.parametrize("outcome", ["success", "timeout", "cancel"]) async def test_retry_reuses_summary_or_terminal_cancellation(outcome: Literal["success", "timeout", "cancel"]) -> None: router: Final = make_router() + ctx: Final = RoutingContext(router) deployment: Final = router.get_deployment(model_id="pinned-answer") assert deployment is not None state: Final = CompactionState(config=ContextCompactionConfig(model="large", max_tokens=512, timeout_seconds=0.02)) @@ -483,21 +485,21 @@ async def test_retry_reuses_summary_or_terminal_cancellation(outcome: Literal["s token: Final = compaction_executor.set(execute) try: - first: Final = asyncio.create_task(compact_to_fit(router, deployment.model_dump(), request, "messages")) + first: Final = asyncio.create_task(compact_to_fit(ctx, deployment.model_dump(), request, "messages")) await asyncio.wait_for(started.wait(), timeout=2) if outcome == "cancel": first.cancel() if outcome == "success": - assert await first == await compact_to_fit(router, deployment.model_dump(), request, "messages") + assert await first == await compact_to_fit(ctx, deployment.model_dump(), request, "messages") changed: Final = {**request, "messages": [{"role": "user", "content": "new history"}, *request["messages"]]} with pytest.raises(litellm.BadRequestError, match="History changed"): - await compact_to_fit(router, deployment.model_dump(), changed, "messages") + await compact_to_fit(ctx, deployment.model_dump(), changed, "messages") else: error: Final = asyncio.CancelledError if outcome == "cancel" else asyncio.TimeoutError with pytest.raises(error): await first with pytest.raises(error): - await compact_to_fit(router, deployment.model_dump(), request, "messages") + await compact_to_fit(ctx, deployment.model_dump(), request, "messages") assert calls.qsize() == 1 and stopped.is_set() finally: compaction_executor.reset(token) diff --git a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py index 77d465526a4..abbb1ed80d0 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -9,6 +9,7 @@ from unittest.mock import create_autospec import httpx import pytest import respx +import tiktoken import litellm from litellm._logging import verbose_router_logger @@ -16,12 +17,15 @@ from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig +from litellm.router_api import RoutingContext +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter, _ClassifierCircuitBreaker from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, JevClassifierConfig from litellm.router_strategy.complexity_router.jev_classifier import ( DEFAULT_JEV_INSTRUCTIONS, HttpJevClassifierClient, JevChoiceAnswer, + JevSystemOneRequest, JevSystemOneResponse, JevUsage, build_jev_request, @@ -82,7 +86,7 @@ async def test_jev_logging_failure_preserves_verdict_and_keeps_circuit_closed( handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) router: Final = ComplexityRouter( "jev-logging-failure", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), {"classifier_type": "jev", "jev_classifier_config": {}, "tiers": {"SIMPLE": "cheap"}}, jev_client=HttpJevClassifierClient("test", "https://typesafe.test", handler), derive_savings_baseline=False, @@ -198,7 +202,7 @@ async def test_jev_accounts_once_with_parent_identity_even_when_the_verdict_fail provider: Final = HttpJevClassifierClient("test", "https://typesafe.test", handler) router: Final = ComplexityRouter( "jev-router", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), { "classifier_type": "jev" if legacy else "oss_classifier", "jev_classifier_config" if legacy else "opensource_classifier_config": { @@ -280,7 +284,7 @@ async def test_jev_uses_bounded_history_and_separates_operator_instructions(incl handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) router: Final = ComplexityRouter( "jev-context", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), { "classifier_type": "jev", "jev_classifier_config": {"instructions": "operator-only rubric"}, @@ -344,7 +348,7 @@ async def test_jev_encrypted_task_skips_provider_without_disabling_plaintext_cla handler.client = httpx.AsyncClient(transport=transport) router: Final = ComplexityRouter( "jev-encrypted", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), { "classifier_type": "jev", "jev_classifier_config": {}, @@ -412,7 +416,7 @@ async def test_jev_cancellation_propagates_without_opening_timeout_breaker() -> handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) router: Final = ComplexityRouter( "jev-cancellation", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), {"classifier_type": "jev", "jev_classifier_config": {}, "tiers": {"SIMPLE": "cheap"}}, jev_client=HttpJevClassifierClient("test", "https://typesafe.test", handler), derive_savings_baseline=False, @@ -521,7 +525,7 @@ async def test_oss_routes_with_its_own_credentials_and_accounts_the_checkpoint( monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) router: Final = ComplexityRouter( f"{provider}-route", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), { "classifier_type": "jev" if legacy else "oss_classifier", "jev_classifier_config" if legacy else "opensource_classifier_config": { @@ -588,7 +592,7 @@ async def test_databricks_routes_through_the_serving_endpoint_and_accounts_the_e monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) router: Final = ComplexityRouter( "databricks-route", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), { "classifier_type": "jev" if legacy else "oss_classifier", "jev_classifier_config" if legacy else "opensource_classifier_config": { @@ -645,7 +649,7 @@ async def test_databricks_routes_through_ai_decide_without_a_model_field_and_acc monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) router: Final = ComplexityRouter( "databricks-route", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), { "classifier_type": "jev" if legacy else "oss_classifier", "jev_classifier_config" if legacy else "opensource_classifier_config": { @@ -685,7 +689,7 @@ def test_databricks_without_a_base_or_key_fails_at_client_build(monkeypatch: pyt with pytest.raises(ValueError, match="DATABRICKS_API_BASE"): ComplexityRouter( "databricks-route", - litellm.Router(model_list=[]), + RoutingContext(litellm.Router(model_list=[])), { "classifier_type": "oss_classifier", "opensource_classifier_config": {"provider": "databricks", "model": "ai_decide"}, @@ -770,7 +774,9 @@ async def test_jev_connection_resolves_only_explicit_environment_references( expected_key: Final = ( "default-key" if connection == "defaults" - else "resolved-key" if connection in ("key-reference", "references") else "literal-key" + else "resolved-key" + if connection in ("key-reference", "references") + else "literal-key" ) assert outcome.cause == "jev_classifier" assert route.calls[0].request.headers["authorization"] == f"Bearer {expected_key}" @@ -1055,3 +1061,258 @@ def test_switching_to_deployed_classifier_does_not_copy_legacy_credentials() -> ) assert write.effective is not None assert write.effective["opensource_classifier_config"] == {"deployment_name": "registered"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("large_field", ["user", "system"]) +@pytest.mark.parametrize("disable_counter", [False, True]) +async def test_jev_bounds_complete_input_without_changing_forwarded_messages( + large_field: str, disable_counter: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_token_counter", disable_counter) + limit: Final = TypeSafeDecisionsConfig().get_input_limits("jev-latest", "typesafe", {}).max_question_tokens + assert limit is not None + rubric: Final = "Operator routing rubric " + "document the task category " * 500 + criterion: Final = ("Simple tasks " + "the classification criteria " * 15).rstrip() + large_text: Final = "BEGIN_SOURCE " + "budget 预算 function(x) return x; " * 18000 + " END_SOURCE" + ask: Final = large_text if large_field == "user" else "CURRENT_ASK Explain the example" + constraints: Final = large_text if large_field == "system" else "Use concise language" + messages: Final = [ + {"role": "system", "content": constraints}, + {"role": "user", "content": "PRIOR_TURN Explain the budget"}, + {"role": "user", "content": ask}, + ] + original: Final = deepcopy(messages) + handler: Final = AsyncHTTPHandler() + handler.client = httpx.AsyncClient() + classifier: Final = ComplexityRouter( + "bounded-classifier", + litellm.Router(model_list=[]), + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"instructions": rubric}, + "tier_definitions": [{"name": "SIMPLE", "description": criterion}, {"name": "REASONING"}], + "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, + "fallback_tier": "SIMPLE", + "session_affinity": False, + "deployment_affinity": False, + }, + jev_client=HttpJevClassifierClient("test", "https://budget.test", handler), + derive_savings_baseline=False, + ) + try: + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("https://budget.test/v1/systemone").respond( + json={"answers": {"tier": _answer().model_dump()}} + ) + result: Final = await classifier.async_pre_routing_hook( + model="bounded-classifier", request_kwargs={"messages": messages}, messages=messages + ) + sent: Final = JevSystemOneRequest.model_validate_json(route.calls[0].request.content) + complete_input: Final = sent.model_dump_json(exclude={"model"}) + measured: Final = len(tiktoken.get_encoding("o200k_base").encode(complete_input)) + assert measured < limit + assert sent.questions["tier"].instructions == rubric + assert sent.questions["tier"].criteria["SIMPLE"] == criterion + assert sent.state.count("BEGIN_SOURCE") == sent.state.count("END_SOURCE") == 1 + assert sent.state.count("PRIOR_TURN") == 1 + assert result is not None and result.routing_decision is not None + decision: Final = result.routing_decision + assert decision["cause"] == "jev_classifier" + assert decision["classifier_input_budget_status"] == "truncated" + assert decision["classifier_input_truncated"] is True + assert measured <= decision["classifier_input_tokens"] <= decision["classifier_input_budget_tokens"] < limit + assert messages == original + assert result.messages is None or result.messages == original + if disable_counter: + assert len(complete_input.encode()) < limit + assert decision["classifier_input_count_method"] == "utf8_bytes" + finally: + await GLOBAL_LOGGING_WORKER.flush() + await handler.client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mode", "default_header", "explicit_cap", "other_member", "expected_status"), + [ + ("false", None, None, None, "truncated"), + ("true", None, None, None, "within_budget"), + ("true", "false", None, None, "truncated"), + ("true", None, 8192, None, "truncated"), + ("true", None, 8192, "default", "truncated"), + ("true", None, None, "default", "truncated"), + ("false", None, None, "unknown", "unknown_limits"), + ], +) +async def test_jev_budget_uses_effective_headers_and_every_deployment( + mode: str, + default_header: str | None, + explicit_cap: int | None, + other_member: str | None, + expected_status: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(litellm.model_cost, "typesafe/jev-latest", deepcopy(litellm.model_cost["typesafe/jev-latest"])) + header: Final = "TypeSafe-Danger-Zone-Long-Context" + deployment: Final = { + "model_name": "budget-model", + "litellm_params": { + "model": "typesafe/jev-latest", + "api_key": "test", + "api_base": "https://budget.test", + "extra_headers": {header: mode}, + }, + "model_info": {"max_input_tokens": explicit_cap} if explicit_cap is not None else {}, + } + sibling: Final = { + "model_name": "budget-model", + "litellm_params": { + "model": "databricks/unknown-budget-endpoint" if other_member == "unknown" else "typesafe/jev-latest", + "api_key": "test", + "api_base": "https://budget.test", + }, + } + router: Final = litellm.Router( + model_list=[deployment, sibling] if other_member else [deployment], + default_litellm_params={"extra_headers": {header: default_header}} if default_header else {}, + num_retries=0, + ) + classifier: Final = ComplexityRouter( + "bounded-classifier", + router, + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"deployment_name": "budget-model"}, + "tiers": {"SIMPLE": "cheap"}, + }, + derive_savings_baseline=False, + ) + ask: Final = "Explain this code " + "word001 word002 word003 " * 12000 + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.route(method="POST").respond(json={"answers": {"tier": _answer().model_dump()}}) + outcome: Final = await classifier.aclassify(ask) + assert outcome.cause == "jev_classifier" + telemetry: Final = outcome.classifier_input_telemetry + assert telemetry is not None and telemetry["classifier_input_budget_status"] == expected_status + assert telemetry["classifier_input_tokens"] <= telemetry["classifier_input_budget_tokens"] + sent: Final = JevSystemOneRequest.model_validate_json(route.calls[0].request.content) + if other_member is None: + assert route.calls[0].request.headers[header] == (default_header or mode) + if explicit_cap is not None: + assert telemetry["classifier_input_budget_tokens"] < explicit_cap + assert (ask in sent.state) is (expected_status == "within_budget") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("oversized_field", "fallback"), + [("instructions", "heuristic"), ("instructions", "default_model"), ("criteria", "custom")], +) +async def test_jev_untrimmable_rubric_uses_configured_fallback_without_a_provider_call( + oversized_field: str, fallback: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(litellm.model_cost, "typesafe/jev-latest", deepcopy(litellm.model_cost["typesafe/jev-latest"])) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "budget-model", + "litellm_params": { + "model": "typesafe/jev-latest", + "api_key": "test", + "api_base": "https://budget.test", + }, + "model_info": {"max_input_tokens": 128 if oversized_field == "criteria" else 1024}, + } + ] + ) + rubric: Final = "word001 word002 word003 " * 2000 + classifier: Final = ComplexityRouter( + "bounded-classifier", + router, + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": { + "deployment_name": "budget-model", + "instructions": rubric if oversized_field == "instructions" else "Choose", + }, + "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, + **( + { + "tier_definitions": [ + {"name": name, "description": "word001 " * 50} for name in ("SIMPLE", "REASONING") + ], + "fallback_tier": "SIMPLE", + } + if fallback == "custom" + else {"classifier_fallback": fallback, "default_model": "default"} + ), + }, + derive_savings_baseline=False, + ) + with respx.mock(assert_all_called=False) as upstream: + route: Final = upstream.route(method="POST").respond(json={"answers": {"tier": _answer().model_dump()}}) + outcome: Final = await classifier.aclassify("What is two plus two?") + assert route.call_count == 0 + assert ( + outcome.cause + == { + "heuristic": "heuristic_scorer", + "default_model": "default_model_fallback", + "custom": "classifier_fallback", + }[fallback] + ) + assert outcome.classifier_failure_reason == "unsupported_input" + assert outcome.classifier_input_telemetry is not None + assert outcome.classifier_input_telemetry["classifier_input_budget_status"] == "untrimmable_input" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback", ["heuristic", "default_model"]) +async def test_jev_open_breaker_preserves_input_budget_in_fallback_decision(fallback: str) -> None: + def clock() -> float: + return 100.0 + + transport: Final = create_autospec(httpx.AsyncBaseTransport, instance=True) + transport.handle_async_request.side_effect = TimeoutError("controlled classifier timeout") + handler: Final = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=transport) + classifier: Final = ComplexityRouter( + "breaker-budget", + litellm.Router(model_list=[]), + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {}, + "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, + "classifier_fallback": fallback, + "default_model": "default", + "session_affinity": False, + "deployment_affinity": False, + }, + jev_client=HttpJevClassifierClient("test", "https://budget.test", handler), + derive_savings_baseline=False, + ) + classifier._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=clock) + try: + first: Final = await classifier.async_pre_routing_hook( + model="breaker-budget", request_kwargs={}, messages=[{"role": "user", "content": "Explain budgets"}] + ) + second: Final = await classifier.async_pre_routing_hook( + model="breaker-budget", + request_kwargs={}, + messages=[{"role": "user", "content": "Explain this code " + "word001 word002 word003 " * 12000}], + ) + assert first is not None and first.routing_decision is not None + assert first.routing_decision["classifier_failure_reason"] == "timeout" + assert second is not None and second.routing_decision is not None + decision: Final = second.routing_decision + assert decision["classifier_failure_reason"] == "circuit_open" + assert decision["cause"] == ("heuristic_scorer" if fallback == "heuristic" else "default_model_fallback") + assert decision.get("classifier_input_budget_status") == "truncated" + assert decision["classifier_input_truncated"] is True + assert 0 < decision["classifier_input_tokens"] <= decision["classifier_input_budget_tokens"] + transport.handle_async_request.assert_awaited_once() + finally: + await handler.client.aclose() diff --git a/tests/unit/router_strategy/test_auto_router.py b/tests/unit/router_strategy/test_auto_router.py index de775b9246f..5d88ebc5aea 100644 --- a/tests/unit/router_strategy/test_auto_router.py +++ b/tests/unit/router_strategy/test_auto_router.py @@ -7,6 +7,7 @@ import pytest from litellm.exceptions import Timeout as LiteLLMTimeout +from litellm.router_api import RoutingContext from litellm.router_strategy.auto_router.auto_router import AutoRouter pytestmark_skip_beta = pytest.mark.skip( @@ -200,7 +201,7 @@ class TestAutoRouter: auto_router_config_path=router_config_path, default_model=default_model, embedding_model=embedding_model, - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), ) # Assert @@ -208,7 +209,7 @@ class TestAutoRouter: assert auto_router.auto_sync_value == AutoRouter.DEFAULT_AUTO_SYNC_VALUE assert auto_router.default_model == default_model assert auto_router.embedding_model == embedding_model - assert auto_router.litellm_router_instance == mock_router_instance + assert auto_router.ctx.belongs_to(mock_router_instance) assert auto_router.routelayer is None mock_semantic_router_class.from_json.assert_called_once_with(router_config_path) @@ -237,7 +238,7 @@ class TestAutoRouter: auto_router_config_path="test/path/router.json", default_model="gpt-4o-mini", embedding_model="text-embedding-model", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), ) messages = [{"role": "user", "content": "test message"}] @@ -278,7 +279,7 @@ class TestAutoRouter: auto_router_config_path="test/path/router.json", default_model="gpt-4o-mini", embedding_model="text-embedding-model", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), ) messages = [{"role": "user", "content": "test message"}] @@ -303,7 +304,7 @@ class TestAutoRouter: auto_router_config_path="test/path/router.json", default_model="gpt-4o-mini", embedding_model="text-embedding-model", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), ) # Act @@ -385,13 +386,13 @@ class FailingEmbeddingRouter(StubEmbeddingRouter): raise self.error -def _auto_router(routelayer: Any, litellm_router_instance: Any = None, **kwargs: Any) -> AutoRouter: +def _auto_router(routelayer: Any, router: Any = None, **kwargs: Any) -> AutoRouter: auto_router: Final = AutoRouter( model_name="my-auto-router", auto_router_config=ROUTER_CONFIG, default_model="fallback-model", embedding_model="text-embedding-3-small", - litellm_router_instance=litellm_router_instance or StubEmbeddingRouter(), + ctx=RoutingContext(router or StubEmbeddingRouter()), **kwargs, ) auto_router.routelayer = routelayer @@ -405,7 +406,7 @@ def test_should_load_no_routes_from_an_empty_non_list_routes_value(routes: objec auto_router_config=json.dumps({"routes": routes}), default_model="fallback-model", embedding_model="text-embedding-3-small", - litellm_router_instance=StubEmbeddingRouter(), + ctx=RoutingContext(StubEmbeddingRouter()), ) assert auto_router.loaded_routes == [] @@ -417,7 +418,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: @pytest.mark.asyncio @pytest.mark.parametrize("request_model", ["my-auto-router", "public-router-alias"]) async def test_should_fall_back_to_default_model_when_the_embedding_call_fails(self, request_model: str) -> None: - auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=FailingEmbeddingRouter()) + auto_router: Final = _auto_router(FixedRouteLayer(None), router=FailingEmbeddingRouter()) result: Final = await auto_router.async_pre_routing_hook( model=request_model, @@ -495,7 +496,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: from semantic_router.schema import RouteChoice router: Final = StubEmbeddingRouter() - auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), router=router) result: Final = await auto_router.async_pre_routing_hook( model=request_model, @@ -524,7 +525,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: async def test_should_record_timeout_without_exception_message(self, error: Exception, wrapped: bool) -> None: auto_router: Final = _auto_router( FixedRouteLayer(None if wrapped else error), - litellm_router_instance=FailingEmbeddingRouter(error) if wrapped else StubEmbeddingRouter(), + router=FailingEmbeddingRouter(error) if wrapped else StubEmbeddingRouter(), ) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", request_kwargs={}, messages=[{"role": "user", "content": "route this"}] @@ -555,7 +556,7 @@ class TestAutoRouterEmbeddingInputCap: @pytest.mark.asyncio async def test_should_build_its_encoder_with_the_configured_cap(self): - auto_router: Final = _auto_router(None, litellm_router_instance=StubEmbeddingRouter(), max_input_chars=777) + auto_router: Final = _auto_router(None, router=StubEmbeddingRouter(), max_input_chars=777) await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -575,7 +576,7 @@ class TestAutoRouterRoutesResponsesApiInput: from semantic_router.schema import RouteChoice router: Final = StubEmbeddingRouter() - auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), router=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -596,7 +597,7 @@ class TestAutoRouterRoutesResponsesApiInput: from semantic_router.schema import RouteChoice router: Final = StubEmbeddingRouter() - auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), router=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -622,7 +623,7 @@ class TestAutoRouterRoutesResponsesApiInput: @pytest.mark.asyncio async def test_should_skip_routing_when_neither_messages_nor_input_is_present(self): router: Final = StubEmbeddingRouter() - auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router) + auto_router: Final = _auto_router(FixedRouteLayer(None), router=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -636,7 +637,7 @@ class TestAutoRouterRoutesResponsesApiInput: @pytest.mark.asyncio async def test_should_keep_routing_an_empty_messages_list_to_the_default_model(self): router: Final = StubEmbeddingRouter() - auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router) + auto_router: Final = _auto_router(FixedRouteLayer(None), router=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -657,7 +658,7 @@ class TestAutoRouterAttributesItsEmbeddingSpend: from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY router: Final = StubEmbeddingRouter() - auto_router: Final = _auto_router(None, litellm_router_instance=router) + auto_router: Final = _auto_router(None, router=router) request_kwargs: Final = { "metadata": { "user_api_key": "hashed-key", @@ -710,7 +711,7 @@ class TestAutoRouterColdStartDoesNotBlockTheEventLoop: import threading embedding_router: Final = ThreadTrackingEmbeddingRouter() - auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + auto_router: Final = _auto_router(None, router=embedding_router) event_loop_thread: Final = threading.get_ident() result: Final = await auto_router.async_pre_routing_hook( @@ -727,7 +728,7 @@ class TestAutoRouterColdStartDoesNotBlockTheEventLoop: @pytest.mark.asyncio async def test_should_build_the_routelayer_exactly_once_under_concurrent_cold_start_requests(self): embedding_router: Final = ThreadTrackingEmbeddingRouter() - auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + auto_router: Final = _auto_router(None, router=embedding_router) results: Final = await asyncio.gather( *( @@ -760,7 +761,7 @@ class TestAutoRouterColdStartDoesNotBlockTheEventLoop: return super().embedding(input, model, **kwargs) embedding_router: Final = BlockingEmbeddingRouter() - auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + auto_router: Final = _auto_router(None, router=embedding_router) first_call: Final = asyncio.ensure_future( auto_router.async_pre_routing_hook( @@ -812,7 +813,7 @@ class TestAutoRouterColdStartDoesNotBlockTheEventLoop: return super().embedding(input, model, **kwargs) embedding_router: Final = FailsOnFirstAttemptEmbeddingRouter() - auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + auto_router: Final = _auto_router(None, router=embedding_router) first_call: Final = asyncio.ensure_future( auto_router.async_pre_routing_hook( diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index ab168b2a950..ef7b9e485d4 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -15,7 +15,7 @@ from copy import deepcopy from functools import partial from types import MappingProxyType from typing import Dict, Final, List, Literal -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import DEFAULT, AsyncMock, MagicMock, patch import httpx import pytest @@ -37,6 +37,7 @@ from litellm.constants import ( from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import as_output_cap, stamped_client_ceiling +from litellm.router_api import RoutingContext from litellm.router_strategy.auto_router.auto_router_model_naming import ( CUSTOMIZATION_CAPABILITY, GATED_AUTO_ROUTER_CAPABILITIES, @@ -113,11 +114,16 @@ def _heuristic_v2_artifact() -> TrainedTierArtifact: ) -def _answer_deployments_named(router: MagicMock) -> None: - """Make a mock router answer `deployments_named` from its own `model_list`, the way Router answers from its index.""" - router.deployments_named.side_effect = lambda name: ( - tuple(deployment for deployment in router.model_list if deployment["model_name"] == name) or None - ) +def _answer_deployments_named(router: MagicMock, model_list: Sequence[Mapping[str, object]]) -> None: + """Make a mock router's registry answer `deployments(name, source="named")` from `model_list`, + the way Router answers from its index.""" + + def deployments(model_name: str | None = None, *, source: str = "group", **_: object) -> object: + if source != "named": + return DEFAULT + return tuple(deployment for deployment in model_list if deployment["model_name"] == model_name) or None + + router.registry.return_value.deployments.side_effect = deployments @pytest.fixture @@ -150,7 +156,7 @@ def complexity_router(mock_router_instance, basic_config): """Create a ComplexityRouter instance with basic config.""" return ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) @@ -208,7 +214,7 @@ class TestComplexityRouterInit: """Test initialization with configuration.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) assert router.model_name == "test-router" @@ -230,7 +236,7 @@ class TestComplexityRouterInit: ask = "Derive the amortized complexity of a splay tree access" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "reminder_markers": [ @@ -255,7 +261,7 @@ class TestComplexityRouterInit: """A config that never mentions reminder_markers keeps stripping .""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) @@ -272,7 +278,7 @@ class TestComplexityRouterInit: """Test initialization without configuration uses defaults.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), ) assert router.model_name == "test-router" # Should have equivalent default values but NOT be the same instance @@ -283,7 +289,7 @@ class TestComplexityRouterInit: """Test initialization with default_model override.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, default_model="fallback-model", ) @@ -297,7 +303,7 @@ class TestComplexityRouterInit: config = {**basic_config, "return_raw_model_name": return_raw_model_name} router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) request_kwargs = {} @@ -330,7 +336,7 @@ class TestComplexityRouterInit: ) router = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "classifier_type": "jev", "jev_classifier_config": {"api_key": "test", "timeout_ms": 100}, @@ -378,7 +384,7 @@ class TestComplexityRouterInit: ) router = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "classifier_type": "jev", "jev_classifier_config": {"api_key": "test", "timeout_ms": 100}, @@ -417,7 +423,7 @@ class TestComplexityRouterInit: ) router = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "classifier_type": "jev", "jev_classifier_config": {"api_key": "test"}, @@ -456,7 +462,7 @@ class TestComplexityRouterInit: ) router = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "classifier_type": "jev", "jev_classifier_config": {"api_key": "test"}, @@ -477,7 +483,7 @@ class TestComplexityRouterInit: client = _TimeoutJevClient() router = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "classifier_type": "jev", "jev_classifier_config": {"api_key": "test", "timeout_ms": 1}, @@ -518,7 +524,7 @@ class TestComplexityRouterInit: client = _StaticJevClient(response) router = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "classifier_type": "jev", "jev_classifier_config": {"api_key": "test"}, @@ -616,7 +622,7 @@ class TestReasoningMarkerScoring: """0 is a real floor, not an absent one, so the markers alone promote again.""" router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "reasoning_override_min_score": 0.0}, ) tier, score, _ = router.classify("hi, step by step, pros and cons") @@ -630,12 +636,12 @@ class TestReasoningMarkerScoring: ) low = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "tier_boundaries": {"simple_medium": 0.20}}, ) high = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "tier_boundaries": {"simple_medium": 0.30}}, ) assert low._effective_reasoning_override_min_score() == 0.20 @@ -650,7 +656,7 @@ class TestReasoningMarkerScoring: ) router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "tier_boundaries": {"simple_medium": 0.10}, @@ -669,7 +675,7 @@ class TestReasoningMarkerScoring: ) router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "reasoning_override_min_score": 0.25}, ) tier, score, _ = router.classify(prompt) @@ -795,7 +801,7 @@ class TestModelSelection: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) model = router.get_model_for_tier(ComplexityTier.SIMPLE) @@ -804,7 +810,7 @@ class TestModelSelection: def test_get_model_for_tier_list_random_choice(self, mock_router_instance): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": ["cheap", "premium"], "MEDIUM": "mid"}, "default_model": "mid", @@ -822,7 +828,7 @@ class TestModelSelection: def test_get_model_for_tier_empty_pool_raises(self, mock_router_instance): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": []}, "default_model": "mid", @@ -950,7 +956,7 @@ class TestConfigOverrides: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) # With very low thresholds, even neutral prompts should be COMPLEX or higher @@ -974,7 +980,7 @@ class TestConfigOverrides: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) # A longer prompt (~150 tokens) should be considered "long" with these thresholds @@ -991,7 +997,7 @@ class TestCustomTechnicalKeywords: """Custom keywords should be appended to the default technical keywords.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={"custom_technical_keywords": ["udp", "kafka"]}, ) assert router.technical_keywords == DEFAULT_TECHNICAL_KEYWORDS + ["udp", "kafka"] @@ -1000,7 +1006,7 @@ class TestCustomTechnicalKeywords: """Custom keywords should be appended to a technical_keywords override.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "technical_keywords": ["quantum", "photonics"], "custom_technical_keywords": ["udp"], @@ -1012,7 +1018,7 @@ class TestCustomTechnicalKeywords: """Duplicates against the base list and within the custom list should be dropped.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={"custom_technical_keywords": ["TCP", "udp", "UDP", "kafka"]}, ) lowered = [kw.lower() for kw in router.technical_keywords] @@ -1025,12 +1031,12 @@ class TestCustomTechnicalKeywords: """Absent or None custom_technical_keywords should leave the keyword list identical.""" router_absent = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={"tiers": {"MEDIUM": "gpt-4o"}}, ) router_none = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={"custom_technical_keywords": None}, ) assert router_absent.technical_keywords == DEFAULT_TECHNICAL_KEYWORDS @@ -1041,12 +1047,12 @@ class TestCustomTechnicalKeywords: prompt = "Configure udp multicast between kafka brokers" baseline_router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) custom_router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "custom_technical_keywords": ["UDP", "Kafka"], @@ -1083,10 +1089,10 @@ class TestCustomDimensions: def test_custom_dimension_changes_only_matching_requests( self, mock_router_instance: MagicMock, matchers: dict[str, object], prompt: str ) -> None: - baseline: Final = ComplexityRouter("test-router", mock_router_instance) + baseline: Final = ComplexityRouter("test-router", RoutingContext(mock_router_instance)) configured: Final = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), {"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, **matchers}]}, ) baseline_tier, baseline_score, baseline_signals = baseline.classify(prompt) @@ -1193,7 +1199,7 @@ class TestCustomDimensions: ) -> None: router: Final = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}, "dimension_weights": {}, @@ -1238,7 +1244,7 @@ class TestCustomDimensions: ) -> None: router: Final = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "custom_dimensions": [ { @@ -1250,7 +1256,7 @@ class TestCustomDimensions: ] }, ) - baseline: Final = ComplexityRouter("test-router", mock_router_instance) + baseline: Final = ComplexityRouter("test-router", RoutingContext(mock_router_instance)) assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2] assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2] second_hit_past_the_bound: Final = "yyy " + "a" * 2044 + " zzz" @@ -1280,11 +1286,13 @@ class TestCustomDimensions: "keywords": ["orbitmesh", "ORBITMESH", "fluxgate"], "patterns": [r"\bcreate\s{1,4}table\b", r"\bcreate\s{1,4}table\b", r"\balter\s{1,4}table\b"], } - baseline: Final = ComplexityRouter("test-router", mock_router_instance) - binary: Final = ComplexityRouter("test-router", mock_router_instance, {"custom_dimensions": [dimension]}) + baseline: Final = ComplexityRouter("test-router", RoutingContext(mock_router_instance)) + binary: Final = ComplexityRouter( + "test-router", RoutingContext(mock_router_instance), {"custom_dimensions": [dimension]} + ) graded: Final = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), {"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]}, ) _, baseline_score, baseline_signals = baseline.classify(prompt) @@ -1329,7 +1337,7 @@ class TestCustomDimensions: assert custom_pattern_work(pattern) == work router: Final = ComplexityRouter( "test-router", - mock_router_instance, + RoutingContext(mock_router_instance), { "custom_dimensions": [ {"name": "bounded", "weight": 0.7, "patterns": [pattern]}, @@ -1508,7 +1516,7 @@ class TestSingletonMutation: # Create router with empty config and custom default_model router1 = ComplexityRouter( model_name="test-router-1", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=None, default_model="custom-fallback", ) @@ -1516,7 +1524,7 @@ class TestSingletonMutation: # Create another router without config router2 = ComplexityRouter( model_name="test-router-2", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=None, ) @@ -2804,7 +2812,7 @@ def llm_complexity_router(mock_router_instance, llm_classifier_config): """ComplexityRouter configured to classify via an LLM call.""" return ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=llm_classifier_config, ) @@ -3053,7 +3061,7 @@ class TestCapabilityClassifier: def _router(mock_router_instance, **overrides): return ComplexityRouter( model_name="capability-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=_capability_router_config(**overrides), ) @@ -3185,11 +3193,13 @@ class TestCapabilityClassifier: from litellm.types.router import RequestType mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply)) - mock_router_instance.model_list = [ - {"model_name": "efficient-model", "litellm_params": {"input_cost_per_token": 0.000001}}, - {"model_name": "capable-model", "litellm_params": {"input_cost_per_token": 0.00001}}, - ] - _answer_deployments_named(mock_router_instance) + _answer_deployments_named( + mock_router_instance, + [ + {"model_name": "efficient-model", "litellm_params": {"input_cost_per_token": 0.000001}}, + {"model_name": "capable-model", "litellm_params": {"input_cost_per_token": 0.00001}}, + ], + ) router: Final = self._router( mock_router_instance, adaptive=True, @@ -3289,7 +3299,7 @@ class TestCapabilityClassifier: ) router = ComplexityRouter( model_name="capability-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) outcome = await router.aclassify("do the task") @@ -3389,7 +3399,7 @@ class TestCapabilityClassifier: bypass: Literal["literal_keyword_match", "session_affinity_pin", "housekeeping"], ) -> None: mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8))) - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router: Final = self._router( mock_router_instance, session_affinity=bypass == "session_affinity_pin", @@ -3519,12 +3529,12 @@ class TestTierLabels: """ renamed = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "tier_labels": CUSTOM_TIER_LABELS}, ) canonical = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) @@ -3590,16 +3600,24 @@ def _native_classifier_router( else AsyncMock(return_value=_native_classifier_response(output), side_effect=failure) ), acompletion=AsyncMock(return_value=_llm_response('{"tier":"SIMPLE"}')), - get_model_list=( - native_router.get_model_list - if native_router is not None - else MagicMock(return_value=[{"litellm_params": {"model": deployment_model}}]) + registry=MagicMock( + return_value=MagicMock( + deployments=( + native_router.registry().deployments + if native_router is not None + else MagicMock( + side_effect=lambda model_name=None, *, source="group", **_: ( + [{"litellm_params": {"model": deployment_model}}] if source == "group" else DEFAULT + ) + ) + ) + ) ), ) return ( ComplexityRouter( model_name="encrypted-router", - litellm_router_instance=dependency, + ctx=RoutingContext(dependency), complexity_router_config={ "tiers": {"SIMPLE": "cheap-model", "REASONING": "deep-model"}, "classifier_type": classifier_type, @@ -3954,7 +3972,7 @@ class TestLLMClassifier: artifact: Final = _heuristic_v2_artifact() router: Final = ComplexityRouter( model_name="tier-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "classifier_type": "heuristic_v2", "heuristic_v2_artifact": artifact, @@ -4156,7 +4174,7 @@ class TestLLMClassifier: ) router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=real_router, + ctx=RoutingContext(real_router), complexity_router_config=llm_classifier_config, ) outcome = await router.aclassify("hi") @@ -4193,7 +4211,7 @@ class TestLLMClassifier: } router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=real_router, + ctx=RoutingContext(real_router), complexity_router_config=config, ) @@ -4224,7 +4242,7 @@ class TestLLMClassifier: } router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -4241,7 +4259,7 @@ class TestLLMClassifier: mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out")) router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=llm_classifier_config, ) @@ -4313,7 +4331,7 @@ class TestLLMClassifier: ) router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=llm_classifier_config, ) router._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now) @@ -4334,7 +4352,7 @@ class TestLLMClassifier: mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out")) router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "classifier_llm_config": { @@ -4451,7 +4469,7 @@ class TestLLMClassifier: ) router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=real_router, + ctx=RoutingContext(real_router), complexity_router_config=llm_classifier_config, ) @@ -4532,7 +4550,7 @@ class TestLLMClassifier: } router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "classifier_llm_config": classifier_llm_config}, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) @@ -4690,7 +4708,7 @@ class TestLLMClassifier: """ router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "tier_labels": CUSTOM_TIER_LABELS}, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "Deep"}')) @@ -4730,7 +4748,7 @@ class TestLLMClassifier: ): router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "tier_labels": CUSTOM_TIER_LABELS}, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "%s"}' % verdict)) @@ -4751,7 +4769,7 @@ class TestLLMClassifier: """ router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "tier_labels": CUSTOM_TIER_LABELS}, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "Expensive"}')) @@ -5660,25 +5678,27 @@ class TestAdaptiveSoftFloors: @pytest.fixture def adaptive_router_instance(self): router = MagicMock() - router.model_list = [ - { - "model_name": "cheap", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "input_cost_per_token": 0.00000015, + _answer_deployments_named( + router, + [ + { + "model_name": "cheap", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + "model_info": {"adaptive_router_preferences": {"quality_tier": 1, "strengths": []}}, }, - "model_info": {"adaptive_router_preferences": {"quality_tier": 1, "strengths": []}}, - }, - { - "model_name": "premium", - "litellm_params": { - "model": "openai/gpt-4o", - "input_cost_per_token": 0.000005, + { + "model_name": "premium", + "litellm_params": { + "model": "openai/gpt-4o", + "input_cost_per_token": 0.000005, + }, + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, }, - "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, - }, - ] - _answer_deployments_named(router) + ], + ) return router @pytest.fixture @@ -5703,7 +5723,7 @@ class TestAdaptiveSoftFloors: def test_cold_start_randomly_samples_unobserved_classified_tier_models(self, adaptive_router_instance): cr = ComplexityRouter( model_name="hybrid", - litellm_router_instance=adaptive_router_instance, + ctx=RoutingContext(adaptive_router_instance), complexity_router_config={ "adaptive": True, "tiers": { @@ -5732,7 +5752,7 @@ class TestAdaptiveSoftFloors: def test_get_model_for_tier_list_without_adaptive_random_choice(self, mock_router_instance): router = ComplexityRouter( model_name="test", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "adaptive": False, "tiers": {"SIMPLE": ["cheap", "premium"], "MEDIUM": "mid"}, @@ -5754,7 +5774,7 @@ class TestAdaptiveSoftFloors: cr = ComplexityRouter( model_name="hybrid", - litellm_router_instance=adaptive_router_instance, + ctx=RoutingContext(adaptive_router_instance), complexity_router_config=hybrid_config, ) adaptive = cr._ensure_adaptive_router() @@ -5776,7 +5796,7 @@ class TestAdaptiveSoftFloors: cr = ComplexityRouter( model_name="hybrid", - litellm_router_instance=adaptive_router_instance, + ctx=RoutingContext(adaptive_router_instance), complexity_router_config=hybrid_config, ) adaptive = cr._ensure_adaptive_router() @@ -5797,7 +5817,7 @@ class TestAdaptiveSoftFloors: cr = ComplexityRouter( model_name="hybrid", - litellm_router_instance=adaptive_router_instance, + ctx=RoutingContext(adaptive_router_instance), complexity_router_config={ "adaptive": True, "tiers": { @@ -5829,7 +5849,7 @@ class TestAdaptiveSoftFloors: async def test_pre_routing_hook_adaptive_stashes_chosen_model(self, adaptive_router_instance, hybrid_config): cr = ComplexityRouter( model_name="hybrid", - litellm_router_instance=adaptive_router_instance, + ctx=RoutingContext(adaptive_router_instance), complexity_router_config=hybrid_config, ) request_kwargs: Dict = {"metadata": {}} @@ -5867,7 +5887,7 @@ class TestLexicalKeywordTierRules: """A prompt hitting a rule keyword routes to that tier, not the scored tier.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=rule_config, ) prompt = "please deploy to k8s now" @@ -5895,7 +5915,7 @@ class TestLexicalKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) result = await router.async_pre_routing_hook( @@ -5919,7 +5939,7 @@ class TestLexicalKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) result = await router.async_pre_routing_hook( @@ -5941,7 +5961,7 @@ class TestLexicalKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) assert router._lexical_tier_override("hi there, please advise") == KeywordOverride( @@ -5963,7 +5983,7 @@ class TestLexicalKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) result = await router.async_pre_routing_hook( @@ -5982,7 +6002,7 @@ class TestLexicalKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) assert router._lexical_tier_override("running my k8s cluster") == KeywordOverride( @@ -5997,7 +6017,7 @@ class TestCjkKeywordTierRules: def _router(self, mock_router_instance, basic_config, keywords: List[str]) -> ComplexityRouter: return ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "keyword_tier_rules": [{"keywords": keywords, "tier": "REASONING"}], @@ -6121,6 +6141,12 @@ class FakeEmbeddingRouter: self.async_embedding_kwargs.append(kwargs) return _make_embedding_response(self._vectors(docs)) + def registry(self) -> "FakeEmbeddingRouter": + return self + + def deployments(self, model_name: str | None = None, **_: object) -> List[Dict]: + return [] + def utterance_embedding_count(self, utterance: str) -> int: """How many times the given route utterance was embedded == number of route-index builds.""" return sum(1 for batch in self.embedded_batches if utterance in batch) @@ -6146,7 +6172,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) result = await router.async_pre_routing_hook( @@ -6179,7 +6205,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) # Only "kubernetes deployment" is close to this query (cos 1.0); "thanks" and @@ -6210,7 +6236,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) caller_metadata = {"user_api_key_hash": "hash-abc", "user_api_key_team_id": "team-1"} @@ -6247,7 +6273,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) await router.async_pre_routing_hook( @@ -6280,7 +6306,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) await router.async_pre_routing_hook( @@ -6312,7 +6338,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) caller_metadata = { @@ -6363,7 +6389,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) loop_thread_id = threading.get_ident() @@ -6394,7 +6420,7 @@ class TestSemanticKeywordTierRules: def _make_router(fake): return ComplexityRouter( model_name="test-router", - litellm_router_instance=fake, + ctx=RoutingContext(fake), complexity_router_config=config, ) @@ -6429,7 +6455,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) # "hello there friend" embeds orthogonal to the REASONING route (cos 0 < 0.9). @@ -6457,7 +6483,7 @@ class TestSemanticKeywordTierRules: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) assert router._semantic_routelayer is None @@ -6644,7 +6670,7 @@ class TestKeywordOverrideEdgeCases: } return ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -6652,7 +6678,7 @@ class TestKeywordOverrideEdgeCases: """No keyword_tier_rules configured -> lexical override is a no-op.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) assert router._lexical_tier_override("deploy to k8s and reason step by step") is None @@ -6663,7 +6689,7 @@ class TestKeywordOverrideEdgeCases: config = {**basic_config, "keyword_tier_rules": [{"keywords": ["k8s"], "tier": "REASONING"}]} router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) assert router.config.embedding_model is None @@ -6743,7 +6769,7 @@ class TestRoutingDecisionCauseLogging: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) await router.async_pre_routing_hook( @@ -6769,7 +6795,7 @@ class TestRoutingDecisionCauseLogging: } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=fake_router, + ctx=RoutingContext(fake_router), complexity_router_config=config, ) await router.async_pre_routing_hook( @@ -6787,7 +6813,7 @@ class TestRoutingDecisionCauseLogging: # No keyword rules -> the scorer decides, and its line must be tagged as such. router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) await router.async_pre_routing_hook( @@ -6835,12 +6861,11 @@ class TestTierModelAffinity: deployment_affinity: bool = True, plugins: bool = False, ) -> ComplexityRouter: - mock_router_instance.cache = DualCache() - mock_router_instance.model_list = [] - _answer_deployments_named(mock_router_instance) + mock_router_instance.health.return_value.store = DualCache() + _answer_deployments_named(mock_router_instance, []) return ComplexityRouter( model_name="affinity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { tier: [ @@ -6934,7 +6959,7 @@ class TestTierModelAffinity: async def test_replaces_pin_outside_the_context_candidate_domain(self, adaptive: bool) -> None: router: Final = ComplexityRouter( model_name="affinity-router", - litellm_router_instance=_windowed_router(_SMALL, _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, _BIG)), complexity_router_config={ "tiers": {"SIMPLE": ["small-model", "big-model"]}, "enable_context_window_escalation": True, @@ -6961,11 +6986,14 @@ class TestTierModelAffinity: async def test_temporary_replay_gate_keeps_the_held_tiers_model_preference( self, mock_router_instance: MagicMock, session_affinity: bool, gate: Literal["image", "health"] ) -> None: - async def get_healthy_deployments( + async def select( model: str, request_kwargs: Mapping[str, object], + *, + kind: Literal["healthy"], messages: Sequence[Mapping[str, object]] | None = None, input: object = None, + specific_deployment: bool | None = False, parent_otel_span: object = None, health_check_probe: bool = False, ) -> list[dict[str, object]]: @@ -6979,8 +7007,8 @@ class TestTierModelAffinity: return [] if unavailable else [{"model_name": model, "model_info": {"id": f"deployment-{model}"}}] cache: Final = DualCache() - mock_router_instance.cache = cache - mock_router_instance.async_get_healthy_deployments = get_healthy_deployments + mock_router_instance.health.return_value.store = cache + mock_router_instance.select = select router: Final = TestModalityRouting._router( mock_router_instance, { @@ -7028,10 +7056,10 @@ class TestTierModelAffinity: self, mock_router_instance: MagicMock ) -> None: clock: Final = MagicMock(return_value=100.0) - mock_router_instance.cache = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + mock_router_instance.health.return_value.store = DualCache(in_memory_cache=InMemoryCache(clock=clock)) router: Final = ComplexityRouter( model_name="affinity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": ["model-a", "model-b"]}, "classification_mode": "user_turn", @@ -7065,10 +7093,10 @@ class TestTierModelAffinity: self, mock_router_instance: MagicMock ) -> None: cache: Final = DualCache() - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache router: Final = ComplexityRouter( model_name="affinity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { "SIMPLE": "base", @@ -7118,10 +7146,10 @@ class TestTierModelAffinity: expected_tier: str, ) -> None: cache: Final = DualCache() - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache router: Final = ComplexityRouter( model_name="affinity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "model-a", "MEDIUM": "model-b"}, "session_affinity": True, @@ -7143,10 +7171,10 @@ class TestTierModelAffinity: async def test_custom_tier_keeps_its_own_model( self, mock_router_instance: MagicMock, classification_mode: Literal["every_request", "user_turn"] ) -> None: - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router: Final = ComplexityRouter( model_name="affinity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=_custom_tier_config( tiers={ "SIMPLE": ["model-a", "model-b"], @@ -7212,10 +7240,10 @@ class TestSessionAffinity: """The hook response's session_affinity_ttl_seconds is what the Router stamps as the deployment-affinity marker, so both the classify path (turn 1) and the session-pin path (turn 2) must carry the configured TTL.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**session_affinity_config, "session_affinity_ttl_seconds": 321}, ) request_kwargs = self._request_kwargs("marker-session") @@ -7252,10 +7280,10 @@ class TestSessionAffinity: ): """Deployment affinity retains a model per tier while classification continues. Session affinity keeps the first tier too; plugins suppress both affinity policies.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "session_affinity": session_affinity, @@ -7279,15 +7307,15 @@ class TestSessionAffinity: async def test_hook_response_has_no_session_affinity_ttl_when_disabled_or_plugins( self, mock_router_instance, basic_config, session_affinity_config ): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() disabled_router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "deployment_affinity": False}, ) plugin_router = ComplexityRouter( model_name="test-router-plugins", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**session_affinity_config, "plugins": [_DummyPlugin()]}, ) disabled = await disabled_router.async_pre_routing_hook( @@ -7303,10 +7331,10 @@ class TestSessionAffinity: async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): """With session_affinity off, a shared session can move from REASONING to SIMPLE.""" assert "session_affinity" not in basic_config - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) request_kwargs = self._request_kwargs("session-1") @@ -7323,10 +7351,10 @@ class TestSessionAffinity: async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config): """A session id the proxy generated for a request that had none is per request, so it must not create a pin even with session_affinity enabled.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=session_affinity_config, ) request_kwargs = {"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}} @@ -7343,10 +7371,10 @@ class TestSessionAffinity: async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config): """Regression: session_affinity=True is the opt-in, so a shared session_id reuses the first turn's model instead of reclassifying.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=session_affinity_config, ) request_kwargs = self._request_kwargs("session-1") @@ -7361,10 +7389,10 @@ class TestSessionAffinity: @pytest.mark.asyncio async def test_pins_model_after_first_turn(self, mock_router_instance, session_affinity_config): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=session_affinity_config, ) request_kwargs = self._request_kwargs("session-1") @@ -7387,13 +7415,13 @@ class TestSessionAffinity: TTL, so a session whose only turn landed on the cooldown fallback must classify again once the breaker closes instead of holding that fallback's model.""" now = 100.0 - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() mock_router_instance.acompletion = AsyncMock( side_effect=[TimeoutError("classifier timed out"), _llm_response('{"tier": "REASONING"}')] ) router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **session_affinity_config, "classifier_type": "llm", @@ -7422,10 +7450,10 @@ class TestSessionAffinity: @pytest.mark.asyncio async def test_a_pinned_turn_reports_the_tier_that_serves_it(self, mock_router_instance, session_affinity_config): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=session_affinity_config, ) request_kwargs = self._request_kwargs("session-1") @@ -7439,10 +7467,10 @@ class TestSessionAffinity: @pytest.mark.asyncio async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=session_affinity_config, ) reasoning = await router.async_pre_routing_hook( @@ -7461,10 +7489,10 @@ class TestSessionAffinity: ) -> None: clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache router: Final = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "session_affinity": True, @@ -7565,10 +7593,10 @@ class TestSessionAffinity: litellm.in_memory_llm_clients_cache.set_cache(key=client_key, value=client) try: cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock), redis_cache=redis_cache) - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache router: Final = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "classification_mode": "user_turn", @@ -7623,10 +7651,10 @@ class TestSessionAffinity: async def test_different_api_keys_do_not_share_pin(self, mock_router_instance, session_affinity_config): """A session_id is client-supplied and unauthenticated; two different callers (API keys) reusing the same session_id must not poison each other's pin.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=session_affinity_config, ) caller_a_kwargs = {"metadata": {"session_id": "shared-session", "user_api_key_hash": "key-a"}} @@ -7647,10 +7675,10 @@ class TestSessionAffinity: @pytest.mark.asyncio async def test_no_session_id_falls_back_to_reclassify(self, mock_router_instance, session_affinity_config): cache = AsyncMock() - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=session_affinity_config, ) result = await router.async_pre_routing_hook( @@ -7665,18 +7693,20 @@ class TestSessionAffinity: """Regression: skipping classification on a pinned turn must not break the adaptive bandit's reward-feedback loop, which only records a turn's outcome when ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY is present in the request metadata.""" - mock_router_instance.cache = DualCache() - mock_router_instance.model_list = [ - { - "model_name": "cheap", - "litellm_params": {"model": "openai/gpt-4o-mini", "input_cost_per_token": 0.0}, - "model_info": {}, - }, - ] - _answer_deployments_named(mock_router_instance) + mock_router_instance.health.return_value.store = DualCache() + _answer_deployments_named( + mock_router_instance, + [ + { + "model_name": "cheap", + "litellm_params": {"model": "openai/gpt-4o-mini", "input_cost_per_token": 0.0}, + "model_info": {}, + }, + ], + ) router = ComplexityRouter( model_name="hybrid", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "adaptive": True, "session_affinity": True, @@ -7915,10 +7945,10 @@ class TestClassificationMode: return {"metadata": {"session_id": session_id}} def _router(self, mock_router_instance, config: dict) -> ComplexityRouter: - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() return ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -7936,7 +7966,7 @@ class TestClassificationMode: with pytest.raises(ValidationError): ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "classification_mode": "sometimes"}, ) @@ -8252,7 +8282,7 @@ class TestRoutingPlugins: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": ["gpt-4o-mini", "gpt-4o-nano"]}, "plugins": [ExcludeGpt4oMini()], @@ -8280,7 +8310,7 @@ class TestRoutingPlugins: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini"}, "default_model": "gpt-4o-fallback", @@ -8303,7 +8333,7 @@ class TestRoutingPlugins: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini"}, "plugins": [BlockEverything()], @@ -8327,7 +8357,7 @@ class TestRoutingPlugins: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini"}, "plugins": [CaptureMetadata()], @@ -8351,7 +8381,7 @@ class TestRoutingPlugins: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": ["gpt-4o-mini", "gpt-4o-nano"]}, "keyword_tier_rules": [{"keywords": ["hello"], "tier": "SIMPLE"}], @@ -8380,7 +8410,7 @@ class TestRoutingPlugins: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"MEDIUM": ["gpt-4o-default", "gpt-4o-nano"]}, "default_model": "gpt-4o-default", @@ -8407,7 +8437,7 @@ class TestRoutingPlugins: PR #33251 after the plugin-bypass fix changed this priority unconditionally.""" router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"MEDIUM": ["gpt-4o-medium-tier"]}, "default_model": "gpt-4o-configured-default", @@ -8449,7 +8479,7 @@ class TestRoutingPlugins: (e.g. a budget cap crossed mid-session) was only ever enforced on a session's first turn. With plugins configured, every turn must go through _classify_and_route (and therefore the plugin pipeline) again.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() class AllowAll: async def run(self, context): @@ -8457,7 +8487,7 @@ class TestRoutingPlugins: router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": ["gpt-4o-mini"]}, "session_affinity": True, @@ -8521,7 +8551,7 @@ def _plugin_router(mock_router_instance, plugin, **config_overrides): } return ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -8686,7 +8716,7 @@ class TestClassifierPlugin: """A built-in tier the operator gave no models is a decline, not a later routing error.""" router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini"}, "classifier_type": "custom", @@ -8826,7 +8856,7 @@ class TestEscalationKeywords: def test_escalate_tier_skips_unconfigured_intermediate(self, mock_router_instance): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "o1-preview"}}, ) assert router._escalate_tier(ComplexityTier.SIMPLE) == ComplexityTier.REASONING @@ -8834,7 +8864,7 @@ class TestEscalationKeywords: def test_tier_for_model_returns_most_severe(self, mock_router_instance): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={"tiers": {"SIMPLE": "shared", "COMPLEX": "shared", "REASONING": "top"}}, ) assert router._tier_for_model("shared") == ComplexityTier.COMPLEX @@ -8845,7 +8875,7 @@ class TestEscalationKeywords: async def test_escalation_bumps_classified_tier(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) # Baseline: this prompt classifies SIMPLE. @@ -8865,7 +8895,7 @@ class TestEscalationKeywords: async def test_lowercase_keyword_does_not_escalate(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) result = await router.async_pre_routing_hook( @@ -8879,7 +8909,7 @@ class TestEscalationKeywords: async def test_custom_escalation_keyword(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "escalation_keywords": ["MAKE IT BETTER"]}, ) # The default keyword no longer triggers once a custom list is supplied. @@ -8901,7 +8931,7 @@ class TestEscalationKeywords: async def test_empty_keyword_list_disables_escalation(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "escalation_keywords": []}, ) result = await router.async_pre_routing_hook( @@ -8915,7 +8945,7 @@ class TestEscalationKeywords: async def test_escalation_caps_at_highest_tier(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) result = await router.async_pre_routing_hook( @@ -8934,7 +8964,7 @@ class TestEscalationKeywords: async def test_escalation_bumps_keyword_tier_override(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "keyword_tier_rules": [{"keywords": ["billing"], "tier": "SIMPLE"}], @@ -8956,10 +8986,10 @@ class TestEscalationKeywords: async def test_escalation_overrides_session_pin_and_persists(self, mock_router_instance, basic_config): """Mid-session escalation bumps relative to the pinned model (never below it) and the bumped model persists for later turns.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "session_affinity": True}, ) request_kwargs = self._request_kwargs("session-1") @@ -9019,10 +9049,10 @@ class TestEscalationKeywords: survives across the plumbing turns after it, so reading escalation off it re-fires per turn and, with the pin persisted, walks the session to the top tier. Escalation reads the newest turn's ask. """ - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "session_affinity": True}, ) request_kwargs = self._request_kwargs("session-plumbing") @@ -9057,7 +9087,7 @@ class TestEscalationKeywords: """ router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) @@ -9096,7 +9126,7 @@ class TestEscalationKeywords: async def test_blank_escalation_keyword_does_not_escalate_everything(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "escalation_keywords": [""]}, ) assert router.escalation_keywords == () @@ -9112,7 +9142,7 @@ class TestEscalationKeywords: when that tier's pool has peers `get_model_for_tier` could randomly pick instead.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}}, ) for pinned in ("o1-a", "o1-b", "o1-c"): @@ -9121,17 +9151,17 @@ class TestEscalationKeywords: @pytest.mark.asyncio async def test_session_escalation_at_ceiling_keeps_multi_model_pin(self, mock_router_instance): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}, "session_affinity": True, }, ) cache_key = router._get_session_affinity_cache_key("session-top", {}) - await mock_router_instance.cache.async_set_cache(key=cache_key, value="o1-b") + await mock_router_instance.health().store.async_set_cache(key=cache_key, value="o1-b") result = await router.async_pre_routing_hook( model="test-model", request_kwargs=self._request_kwargs("session-top"), @@ -9168,7 +9198,7 @@ class TestStallEscalation: async def test_repeated_tool_calls_escalate_the_classified_tier(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "stall_escalation_enabled": True}, ) messages = [*_stalled_tool_history(), {"role": "user", "content": "Hello there!"}] @@ -9179,7 +9209,7 @@ class TestStallEscalation: async def test_varied_tool_calls_do_not_escalate(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "stall_escalation_enabled": True}, ) messages = [ @@ -9200,7 +9230,7 @@ class TestStallEscalation: async def test_disabled_by_default_ignores_repeated_tool_calls(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=basic_config, ) messages = [*_stalled_tool_history(), {"role": "user", "content": "Hello there!"}] @@ -9211,7 +9241,7 @@ class TestStallEscalation: async def test_signals_record_stall_escalation(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "stall_escalation_enabled": True}, ) messages = [*_stalled_tool_history(), {"role": "user", "content": "Hello there!"}] @@ -9222,7 +9252,7 @@ class TestStallEscalation: async def test_stall_escalation_caps_at_highest_tier(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "stall_escalation_enabled": True}, ) messages = [ @@ -9236,7 +9266,7 @@ class TestStallEscalation: async def test_stall_escalation_stacks_with_keyword_escalation(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "stall_escalation_enabled": True}, ) messages = [*_stalled_tool_history(), {"role": "user", "content": "LITELLM ESCALATE Hello there!"}] @@ -9250,7 +9280,7 @@ class TestStallEscalation: would be the one path a stall could never lift.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **basic_config, "stall_escalation_enabled": True, @@ -9276,7 +9306,7 @@ class TestStallEscalation: before it: escalation still fires on the turn carrying that follow-up.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "stall_escalation_enabled": True}, ) messages = [*_stalled_tool_history(), {"role": "user", "content": "try again"}] @@ -9403,7 +9433,7 @@ class TestRoutingDecisionContents: } router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) response = await router.async_pre_routing_hook( @@ -9437,7 +9467,7 @@ class TestRoutingDecisionContents: async def test_a_default_model_fallback_claims_no_tier(self, mock_router_instance, basic_config): router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "default_model": "gpt-4o"}, ) response = await router.async_pre_routing_hook( @@ -9452,15 +9482,15 @@ class TestRoutingDecisionContents: @pytest.mark.asyncio async def test_session_pin_decision(self, mock_router_instance, basic_config): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "session_affinity": True}, ) request_kwargs = {"metadata": {"session_id": "session-decision"}} cache_key = router._get_session_affinity_cache_key("session-decision", request_kwargs) - await mock_router_instance.cache.async_set_cache(key=cache_key, value="gpt-4o") + await mock_router_instance.health().store.async_set_cache(key=cache_key, value="gpt-4o") response = await router.async_pre_routing_hook( model="test-complexity-router", request_kwargs=request_kwargs, @@ -9514,7 +9544,7 @@ class TestRoutingDecisionContents: """ router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "tier_labels": CUSTOM_TIER_LABELS}, ) response = await router.async_pre_routing_hook( @@ -9534,7 +9564,7 @@ class TestRoutingDecisionContents: """A partial map must not stamp a redundant label on the tiers it left alone.""" router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "tier_labels": {"REASONING": "Deep"}}, ) simple = await router.async_pre_routing_hook( @@ -9841,7 +9871,7 @@ class TestEscalationIsRecordedConsistently: async def test_scorer_path_at_ceiling_keeps_the_keyword_and_reports_no_bump(self, mock_router_instance): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **self.CEILING_CONFIG, "tier_boundaries": {"simple_medium": -99, "medium_complex": -99, "complex_reasoning": -99}, @@ -9870,15 +9900,15 @@ class TestEscalationIsRecordedConsistently: @pytest.mark.asyncio async def test_session_pin_at_ceiling_still_records_the_ask(self, mock_router_instance): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**self.CEILING_CONFIG, "session_affinity": True}, ) request_kwargs = {"metadata": {"session_id": "session-ceiling"}} cache_key = router._get_session_affinity_cache_key("session-ceiling", request_kwargs) - await mock_router_instance.cache.async_set_cache(key=cache_key, value="o1-preview") + await mock_router_instance.health().store.async_set_cache(key=cache_key, value="o1-preview") response = await router.async_pre_routing_hook( model="test-router", @@ -9895,15 +9925,15 @@ class TestEscalationIsRecordedConsistently: @pytest.mark.asyncio async def test_session_pin_below_ceiling_reports_the_bump(self, mock_router_instance): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**self.CEILING_CONFIG, "session_affinity": True}, ) request_kwargs = {"metadata": {"session_id": "session-below"}} cache_key = router._get_session_affinity_cache_key("session-below", request_kwargs) - await mock_router_instance.cache.async_set_cache(key=cache_key, value="gpt-4o-mini") + await mock_router_instance.health().store.async_set_cache(key=cache_key, value="gpt-4o-mini") response = await router.async_pre_routing_hook( model="test-router", @@ -10192,7 +10222,7 @@ class TestContextAwareClassifier: dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}'))) router: Final = ComplexityRouter( "test-complexity-router", - dependency, + RoutingContext(dependency), { **llm_classifier_config, "classifier_context_include_assistant_turns": True, @@ -10253,7 +10283,7 @@ class TestContextAwareClassifier: dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}'))) router: Final = ComplexityRouter( "test-complexity-router", - dependency, + RoutingContext(dependency), {**llm_classifier_config, "classifier_context_window_size": 0}, ) messages: Final = [ @@ -10317,7 +10347,7 @@ class TestContextAwareClassifier: ) router: Final = ComplexityRouter( model_name="router", - litellm_router_instance=MagicMock(acompletion=completion, cache=DualCache()), + ctx=RoutingContext(MagicMock(acompletion=completion, **{"health.return_value.store": DualCache()})), complexity_router_config={ "tiers": {"SIMPLE": "simple-model", "COMPLEX": "task-model"}, "classifier_type": "llm", @@ -10382,7 +10412,7 @@ class TestContextAwareClassifier: completion: Final = AsyncMock(return_value=_llm_response('{"tier":"COMPLEX"}')) router: Final = ComplexityRouter( model_name="router", - litellm_router_instance=MagicMock(acompletion=completion), + ctx=RoutingContext(MagicMock(acompletion=completion)), complexity_router_config={ "tiers": {"COMPLEX": "task-model"}, "default_model": "fallback-model", @@ -10450,7 +10480,7 @@ class TestContextAwareClassifier: completion: Final = AsyncMock(return_value=_llm_response('{"tier":"COMPLEX"}')) router: Final = ComplexityRouter( model_name="codex-router", - litellm_router_instance=MagicMock(acompletion=completion), + ctx=RoutingContext(MagicMock(acompletion=completion)), complexity_router_config={ "tiers": {"COMPLEX": "task-model", "REASONING": "escalated-model"}, "classifier_type": "llm", @@ -10498,7 +10528,7 @@ class TestContextAwareClassifier: completion: Final = AsyncMock(return_value=_llm_response('{"tier":"COMPLEX"}')) router: Final = ComplexityRouter( model_name="router", - litellm_router_instance=MagicMock(acompletion=completion), + ctx=RoutingContext(MagicMock(acompletion=completion)), complexity_router_config={ "tiers": {"COMPLEX": "task-model"}, "classifier_type": "llm", @@ -11031,7 +11061,7 @@ class TestContextAwareClassifier: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=llm_classifier_config, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) @@ -11058,7 +11088,7 @@ class TestContextAwareClassifier: """ router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "classifier_context_budget_chars": 1}, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) @@ -11344,7 +11374,7 @@ class TestContextAwareClassifier: """ router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini", "COMPLEX": "claude-sonnet-4-20250514"}, "classifier_type": "llm", @@ -11383,7 +11413,7 @@ class TestContextAwareClassifier: """ router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "classifier_context_include_assistant_turns": include_assistant, @@ -11425,7 +11455,7 @@ class TestContextAwareClassifier: """ router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "classifier_context_include_assistant_turns": include_assistant, @@ -11466,7 +11496,7 @@ class TestContextAwareClassifier: """ router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "classifier_context_include_assistant_turns": True, @@ -11516,7 +11546,7 @@ class TestClientHousekeepingCalls: """ router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=llm_classifier_config, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) @@ -11542,7 +11572,7 @@ class TestClientHousekeepingCalls: """ router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=llm_classifier_config, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) @@ -11566,7 +11596,7 @@ class TestClientHousekeepingCalls: """A caller who explicitly escalated asked for something; the cap must not silently undo it.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=llm_classifier_config, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) @@ -11585,7 +11615,7 @@ class TestClientHousekeepingCalls: """keyword_tier_rules are the operator's own instruction, decided before this ever runs.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "keyword_tier_rules": [{"keywords": ["livelocks under contention"], "tier": "REASONING"}], @@ -11606,7 +11636,7 @@ class TestClientHousekeepingCalls: """The floor is an operator guarantee about what plan-mode turns may run on, so it wins.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "plan_mode_min_tier": "COMPLEX"}, ) @@ -11629,7 +11659,7 @@ class TestClientHousekeepingCalls: """An operator who wants these classified keeps the old behaviour, classifier call included.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "route_housekeeping_to_cheapest_tier": False}, ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) @@ -11649,7 +11679,7 @@ class TestClientHousekeepingCalls: """Client wording drifts with releases, so coverage has to be extensible without a code change.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "housekeeping_patterns": ["Summarize this thread for the sidebar"], @@ -11678,7 +11708,7 @@ class TestClientHousekeepingCalls: """A tier can be declared with no pool, and routing to an empty pool is a different bug.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"COMPLEX": "claude-sonnet-4-20250514", "REASONING": "o1-preview"}, "default_model": "gpt-4o-mini", @@ -11712,7 +11742,7 @@ class TestClientHousekeepingCalls: router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "o1-preview"}, "classifier_type": "custom", @@ -11731,22 +11761,24 @@ class TestClientHousekeepingCalls: def _adaptive_router(self, tier_distance_penalty: float, plan_mode_min_tier: str | None = None) -> ComplexityRouter: adaptive_instance = MagicMock() - adaptive_instance.model_list = [ - { - "model_name": "cheap", - "litellm_params": {"model": "openai/gpt-4o-mini", "input_cost_per_token": 0.00000015}, - "model_info": {"adaptive_router_preferences": {"quality_tier": 1, "strengths": []}}, - }, - { - "model_name": "premium", - "litellm_params": {"model": "openai/gpt-4o", "input_cost_per_token": 0.000005}, - "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, - }, - ] - _answer_deployments_named(adaptive_instance) + _answer_deployments_named( + adaptive_instance, + [ + { + "model_name": "cheap", + "litellm_params": {"model": "openai/gpt-4o-mini", "input_cost_per_token": 0.00000015}, + "model_info": {"adaptive_router_preferences": {"quality_tier": 1, "strengths": []}}, + }, + { + "model_name": "premium", + "litellm_params": {"model": "openai/gpt-4o", "input_cost_per_token": 0.000005}, + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, + }, + ], + ) router = ComplexityRouter( model_name="hybrid", - litellm_router_instance=adaptive_instance, + ctx=RoutingContext(adaptive_instance), complexity_router_config={ "adaptive": True, "adaptive_eligible": "all", @@ -11806,11 +11838,11 @@ class TestClientHousekeepingCalls: An agent names the conversation on its first turn, so the cheapest tier would be the pin every session starts with and the real work that follows would run there for the whole TTL. """ - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { "SIMPLE": "gpt-4o-mini", @@ -11848,7 +11880,7 @@ class TestClientHousekeepingCalls: """ router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=llm_classifier_config, ) @@ -11918,7 +11950,7 @@ class TestClassifierTrustBoundary: router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "o1-preview"}, "classifier_type": "llm", @@ -11980,7 +12012,7 @@ class TestClassifierTrustBoundary: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "classifier_context_include_assistant_turns": include_assistant, @@ -12030,7 +12062,7 @@ class TestConversationShapeDiscriminator: def _router(mock_router_instance, basic_config) -> ComplexityRouter: return ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**basic_config, "session_affinity": False}, ) @@ -12038,7 +12070,7 @@ class TestConversationShapeDiscriminator: async def test_a_single_ask_is_a_first_turn(self, mock_router_instance, basic_config): """Nothing is cached for any model yet, so the baseline would have paid the same cache write and the saving is the plain rate difference.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() result = await self._router(mock_router_instance, basic_config).async_pre_routing_hook( model="test-model", request_kwargs={"metadata": {}}, @@ -12050,7 +12082,7 @@ class TestConversationShapeDiscriminator: async def test_a_second_ask_means_the_baseline_was_already_warm(self, mock_router_instance, basic_config): """An earlier turn was served, so a single-model deployment wrote the prompt then and would only read it now; this request's write is what switching cost.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() result = await self._router(mock_router_instance, basic_config).async_pre_routing_hook( model="test-model", request_kwargs={"metadata": {}}, @@ -12066,7 +12098,7 @@ class TestConversationShapeDiscriminator: async def test_it_needs_no_session_id(self, mock_router_instance, basic_config): """The whole point of reading the conversation rather than remembering it: a caller that sends no session header is still classified correctly.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = self._router(mock_router_instance, basic_config) first = await router.async_pre_routing_hook( model="test-model", request_kwargs={}, messages=[{"role": "user", "content": "Hello!"}] @@ -12089,7 +12121,7 @@ class TestConversationShapeDiscriminator: round-trip, and with it a cache failure that would read as a first turn.""" cache = AsyncMock() cache.async_get_cache = AsyncMock(return_value=None) - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache result = await self._router(mock_router_instance, basic_config).async_pre_routing_hook( model="test-model", request_kwargs={"metadata": {}}, @@ -12243,7 +12275,7 @@ class TestCustomClassifierSystemPrompt: **config_patch, ) return ComplexityRouter( - model_name="test-complexity-router", litellm_router_instance=MagicMock(), complexity_router_config=config + model_name="test-complexity-router", ctx=RoutingContext(MagicMock()), complexity_router_config=config ) def test_custom_instructions_keep_the_rubric_criteria_and_examples(self): @@ -12290,7 +12322,7 @@ class TestCustomClassifierSystemPrompt: classification_prompt="Grade the request.", ) router = ComplexityRouter( - model_name="test-complexity-router", litellm_router_instance=MagicMock(), complexity_router_config=config + model_name="test-complexity-router", ctx=RoutingContext(MagicMock()), complexity_router_config=config ) prompt = router._classifier_system_prompt assert prompt is not None @@ -12364,7 +12396,7 @@ class TestCustomClassifierSystemPrompt: ) router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "classifier_llm_config": { @@ -12391,7 +12423,7 @@ class TestCustomClassifierSystemPrompt: rather than escaping as an exception to the caller's request.""" router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "classifier_llm_config": { @@ -12425,7 +12457,7 @@ class TestClassifierFallbackChoice: def default_model_fallback_router(self, mock_router_instance, llm_classifier_config): return ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ **llm_classifier_config, "classifier_fallback": "default_model", @@ -12442,7 +12474,7 @@ class TestClassifierFallbackChoice: with pytest.raises(ValueError, match="requires a default model"): ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "classifier_fallback": "default_model"}, ) @@ -12453,7 +12485,7 @@ class TestClassifierFallbackChoice: validator would have rejected this valid deployment.""" router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={**llm_classifier_config, "classifier_fallback": "default_model"}, default_model="gpt-4o", ) @@ -12508,7 +12540,7 @@ class TestClassifierFallbackChoice: disabled here so the next turn isolates and verifies the affinity contract.""" router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { "SIMPLE": "gpt-4o-mini", @@ -12527,7 +12559,7 @@ class TestClassifierFallbackChoice: "session_affinity": True, }, ) - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() request_kwargs: Dict = {"metadata": {"session_id": "session-flaky"}} mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out")) @@ -12556,7 +12588,7 @@ class TestClassifierFallbackChoice: turn on a default_model-fallback router must still pin exactly as it did before.""" router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "o1-preview"}, "classifier_type": "llm", @@ -12566,7 +12598,7 @@ class TestClassifierFallbackChoice: "session_affinity": True, }, ) - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() request_kwargs: Dict = {"metadata": {"session_id": "session-steady"}} mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) @@ -12600,7 +12632,7 @@ class TestClassifierFallbackChoice: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"MEDIUM": ["gpt-4o-default", "gpt-4o-nano"]}, "classifier_type": "llm", @@ -12640,7 +12672,7 @@ class TestClassifierFallbackChoice: router = ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"COMPLEX": ["o1-preview"]}, "classifier_type": "llm", @@ -12690,7 +12722,7 @@ class TestSavingsBaselineOnDecision: ) return ComplexityRouter( model_name="savings-router", - litellm_router_instance=parent, + ctx=RoutingContext(parent), complexity_router_config={"tiers": tiers}, **kwargs, ) @@ -12716,7 +12748,7 @@ class TestSavingsBaselineOnDecision: """The deployment id is what lets the spend writer price a baseline whose deployment carries a configured rate instead of the public one.""" router = self._router_with_tiers({"SIMPLE": "cheap", "REASONING": "top"}) - expected_id = router.litellm_router_instance.get_model_list(model_name="top")[0]["model_info"]["id"] + expected_id = router.ctx.registry().deployments(model_name="top")[0]["model_info"]["id"] decision = router._build_routing_decision(routed_model="cheap", cause="heuristic_scorer") assert decision["savings_baseline_model"] == "anthropic/claude-fable-5" assert decision["savings_baseline_deployment_id"] == expected_id @@ -12760,7 +12792,7 @@ class TestSavingsBaselinePinnedPerInstance: ) router = ComplexityRouter( model_name="savings-router", - litellm_router_instance=parent, + ctx=RoutingContext(parent), complexity_router_config={"tiers": {"SIMPLE": "cheap", "REASONING": ["cheap", "top"]}}, ) return router, parent @@ -12779,7 +12811,7 @@ class TestSavingsBaselinePinnedPerInstance: parent.model_name_to_deployment_indices.clear() rebuilt = ComplexityRouter( model_name="savings-router", - litellm_router_instance=parent, + ctx=RoutingContext(parent), complexity_router_config={"tiers": {"SIMPLE": "cheap", "REASONING": ["cheap", "top"]}}, ) assert rebuilt.savings_baseline is None @@ -13071,7 +13103,7 @@ class TestTierDefinitions: def custom_tier_router(self, mock_router_instance): return ComplexityRouter( model_name="custom-tier-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=_custom_tier_config(), ) @@ -13186,7 +13218,7 @@ class TestTierDefinitions: the injection-defense paragraph would let a caller ask for a tier and get it.""" router = ComplexityRouter( model_name="custom-tier-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=_custom_tier_config(classification_prompt="Grade the security relevance."), ) mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) @@ -13205,7 +13237,7 @@ class TestTierDefinitions: after the defined tiers, and still above the injection guard.""" router = ComplexityRouter( model_name="custom-tier-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=_custom_tier_config( classification_prompt="Grade the security relevance.", classification_examples='- "audit this login handler" -> SECURITY_REVIEW', @@ -13257,7 +13289,7 @@ class TestTierDefinitions: tier_definitions wins, mirroring the built-in severity tie-break.""" router = ComplexityRouter( model_name="custom-tier-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=_custom_tier_config( keyword_tier_rules=[ {"keywords": ["audit"], "tier": "SECURITY_REVIEW"}, @@ -13517,7 +13549,7 @@ class TestPlanModeTierFloor: def _router(self, mock_router_instance, config: dict) -> ComplexityRouter: return ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -13618,7 +13650,7 @@ class TestPlanModeTierFloor: auto-routes back to it instead of staying premium.""" from litellm.caching.dual_cache import DualCache - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() config = {**basic_config, "plan_mode_min_tier": "COMPLEX", "session_affinity": True} router = self._router(mock_router_instance, config) session_kwargs = {"metadata": {"session_id": "plan-session"}} @@ -13659,7 +13691,7 @@ class TestPlanModeTierFloor: the first ordinary turn classifies and pins as if plan mode had never happened.""" from litellm.caching.dual_cache import DualCache - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() config = {**basic_config, "plan_mode_min_tier": "COMPLEX", "session_affinity": True} router = self._router(mock_router_instance, config) session_kwargs = {"metadata": {"session_id": "plan-first-session"}} @@ -13683,7 +13715,7 @@ class TestPlanModeTierFloor: async def test_pinned_session_at_or_above_floor_keeps_pin_cause(self, mock_router_instance, basic_config): from litellm.caching.dual_cache import DualCache - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() config = {**basic_config, "plan_mode_min_tier": "MEDIUM", "session_affinity": True} router = self._router(mock_router_instance, config) session_kwargs = {"metadata": {"session_id": "premium-session"}} @@ -13822,22 +13854,24 @@ class TestPlanModeTierFloor: from litellm.types.router import RequestType adaptive_instance = MagicMock() - adaptive_instance.model_list = [ - { - "model_name": "cheap", - "litellm_params": {"model": "openai/gpt-4o-mini", "input_cost_per_token": 0.00000015}, - "model_info": {"adaptive_router_preferences": {"quality_tier": 1, "strengths": []}}, - }, - { - "model_name": "premium", - "litellm_params": {"model": "openai/gpt-4o", "input_cost_per_token": 0.000005}, - "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, - }, - ] - _answer_deployments_named(adaptive_instance) + _answer_deployments_named( + adaptive_instance, + [ + { + "model_name": "cheap", + "litellm_params": {"model": "openai/gpt-4o-mini", "input_cost_per_token": 0.00000015}, + "model_info": {"adaptive_router_preferences": {"quality_tier": 1, "strengths": []}}, + }, + { + "model_name": "premium", + "litellm_params": {"model": "openai/gpt-4o", "input_cost_per_token": 0.000005}, + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, + }, + ], + ) router = ComplexityRouter( model_name="hybrid", - litellm_router_instance=adaptive_instance, + ctx=RoutingContext(adaptive_instance), complexity_router_config={ "adaptive": True, "tiers": {"SIMPLE": ["cheap"], "MEDIUM": ["cheap"], "COMPLEX": ["premium"]}, @@ -13864,7 +13898,7 @@ class TestPlanModeTierFloor: turn's pin would carry plan mode past its exit.""" from litellm.caching.dual_cache import DualCache - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() config = { **basic_config, "plan_mode_min_tier": "MEDIUM", @@ -14011,7 +14045,7 @@ def test_adaptive_empty_tier_pool_is_rejected(): def test_tier_model_params_are_used_by_pools_and_savings_baseline(mock_router_instance): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { "SIMPLE": "mini", @@ -14029,7 +14063,7 @@ def test_tier_model_params_are_used_by_pools_and_savings_baseline(mock_router_in async def test_tier_model_params_reach_the_hook_response_and_override_client_values(mock_router_instance): router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { "REASONING": { @@ -14065,13 +14099,13 @@ async def test_tier_params_mask_credentials_in_routing_decision(route, mock_rout } router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) request_kwargs = {"metadata": {"session_id": "masked-params-session"}} if route == "session": - mock_router_instance.cache = DualCache() - await mock_router_instance.cache.async_set_cache( + mock_router_instance.health.return_value.store = DualCache() + await mock_router_instance.health().store.async_set_cache( key=router._get_session_affinity_cache_key("masked-params-session", request_kwargs), value={"model": "opus", "tier": "REASONING"}, ) @@ -14094,10 +14128,10 @@ async def test_tier_params_mask_credentials_in_routing_decision(route, mock_rout @pytest.mark.asyncio async def test_session_pin_outside_tiers_does_not_inherit_medium_params(mock_router_instance): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { "SIMPLE": "mini", @@ -14108,7 +14142,7 @@ async def test_session_pin_outside_tiers_does_not_inherit_medium_params(mock_rou }, ) request_kwargs = {"metadata": {"session_id": "orphan-session"}} - await mock_router_instance.cache.async_set_cache( + await mock_router_instance.health().store.async_set_cache( key=router._get_session_affinity_cache_key("orphan-session", request_kwargs), value="orphan", ) @@ -14126,10 +14160,10 @@ async def test_session_pin_outside_tiers_does_not_inherit_medium_params(mock_rou @pytest.mark.asyncio async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mock_router_instance): - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { "SIMPLE": {"model_name": "shared", "litellm_params": {"reasoning_effort": "low"}}, @@ -14139,7 +14173,7 @@ async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mo }, ) request_kwargs = {"metadata": {"session_id": "shared-session"}} - await mock_router_instance.cache.async_set_cache( + await mock_router_instance.health().store.async_set_cache( key=router._get_session_affinity_cache_key("shared-session", request_kwargs), value={"model": "shared", "tier": "SIMPLE"}, ) @@ -14159,10 +14193,10 @@ async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mo @pytest.mark.asyncio async def test_session_pin_survives_json_list_round_trip(mock_router_instance): cache: Final = DualCache() - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache router = ComplexityRouter( model_name="test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": { "SIMPLE": {"model_name": "shared", "litellm_params": {"reasoning_effort": "low"}}, @@ -14223,7 +14257,7 @@ def _heuristic_first_router(mock_router_instance, **config_overrides): } return ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -14413,7 +14447,7 @@ def _hybrid_router(mock_router_instance, **config_overrides): } return ComplexityRouter( model_name="test-complexity-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -14645,7 +14679,7 @@ async def test_selected_v2_chain_routes_using_success_probabilities( mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}')) router: Final = ComplexityRouter( model_name="v2-chain", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "classifier_type": classifier_type, "local_heuristic": "heuristic_v2", @@ -14693,7 +14727,7 @@ async def test_v2_chain_judge_failure_uses_selected_fallback( dependency.acompletion = AsyncMock(side_effect=RuntimeError("judge failed")) router: Final = ComplexityRouter( model_name="v2-chain", - litellm_router_instance=dependency, + ctx=RoutingContext(dependency), complexity_router_config={ "classifier_type": classifier_type, "local_heuristic": "heuristic_v2", @@ -14778,7 +14812,7 @@ class TestContextWindowEscalation: """ router = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(_SMALL, _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, _BIG)), complexity_router_config=_tier_config(), ) @@ -14796,7 +14830,7 @@ class TestContextWindowEscalation: """The gate must be invisible for normal traffic: same model, no escalation facts.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(_SMALL, _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, _BIG)), complexity_router_config=_tier_config(), ) @@ -14815,7 +14849,7 @@ class TestContextWindowEscalation: that fits, which is cheaper than escalating and preserves the classifier's decision.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG)), complexity_router_config=_tier_config(tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}), ) @@ -14834,24 +14868,26 @@ class TestContextWindowEscalation: smallest resolvable window and escalates past it.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=Router( - model_list=[ - { - "model_name": "mixed-pool", - "litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"}, - "model_info": {"max_input_tokens": 16385}, - }, - { - "model_name": "mixed-pool", - "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"max_input_tokens": 200000}, - }, - { - "model_name": "big-model", - "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"max_input_tokens": 200000}, - }, - ] + ctx=RoutingContext( + Router( + model_list=[ + { + "model_name": "mixed-pool", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 16385}, + }, + { + "model_name": "mixed-pool", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + { + "model_name": "big-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + ] + ) ), complexity_router_config=_tier_config(tiers={"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}), ) @@ -14869,7 +14905,7 @@ class TestContextWindowEscalation: length, which the token count can never exceed.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(_SMALL, _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, _BIG)), complexity_router_config=_tier_config(), ) @@ -14904,7 +14940,7 @@ class TestContextWindowEscalation: when nothing provably fits, the classified tier stands and the client owns overflow.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(*deployments), + ctx=RoutingContext(_windowed_router(*deployments)), complexity_router_config=_tier_config(tiers=tiers), ) @@ -14934,7 +14970,7 @@ class TestContextWindowEscalation: ) router: Final = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(_SMALL, _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, _BIG)), complexity_router_config=config.model_dump(exclude_unset=not serialized, exclude_none=True), ) @@ -14954,7 +14990,7 @@ class TestContextWindowEscalation: dispatches a provably oversized request and the provider 400s anyway.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(_SMALL, _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, _BIG)), complexity_router_config=_tier_config(), ) @@ -14983,7 +15019,7 @@ class TestContextWindowEscalation: fresh instead of inheriting a pin.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(_SMALL, _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, _BIG)), complexity_router_config=_tier_config(session_affinity=True), ) @@ -15008,7 +15044,7 @@ class TestContextWindowEscalation: session's own model so the first turn that fits again routes exactly as pinned.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=_windowed_router(_SMALL, _BIG), + ctx=RoutingContext(_windowed_router(_SMALL, _BIG)), complexity_router_config=_tier_config(session_affinity=True), ) @@ -15039,19 +15075,21 @@ class TestContextWindowEscalation: unobserved, cold start samples only among models whose window holds the prompt.""" router = ComplexityRouter( model_name="test-router", - litellm_router_instance=Router( - model_list=[ - { - "model_name": "small-model", - "litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"}, - "model_info": {"max_input_tokens": 16385}, - }, - { - "model_name": "mid-model", - "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"max_input_tokens": 200000}, - }, - ] + ctx=RoutingContext( + Router( + model_list=[ + { + "model_name": "small-model", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 16385}, + }, + { + "model_name": "mid-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + ] + ) ), complexity_router_config=_tier_config(adaptive=True, tiers={"SIMPLE": ["small-model", "mid-model"]}), ) @@ -15073,15 +15111,17 @@ class TestContextWindowEscalation: (tmp_path / "api-key.json").write_text(json.dumps({"token": "tid=test", "expires_at": int(time.time()) + 3600})) router = ComplexityRouter( model_name="test-router", - litellm_router_instance=Router( - model_list=[ - {"model_name": "cop-pool", "litellm_params": {"model": "github_copilot/gpt-4o"}}, - { - "model_name": "big-model", - "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"max_input_tokens": 200000}, - }, - ] + ctx=RoutingContext( + Router( + model_list=[ + {"model_name": "cop-pool", "litellm_params": {"model": "github_copilot/gpt-4o"}}, + { + "model_name": "big-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + ] + ) ), complexity_router_config=_tier_config(tiers={"SIMPLE": "cop-pool", "COMPLEX": "big-model"}), ) @@ -15187,7 +15227,11 @@ class TestModalityRouting: def _router(mock_router_instance, config, vision_by_model): """vision_by_model: model name -> True/False (deployment model_info) or None (undeclared).""" - def get_model_list(model_name=None): + named: Final = mock_router_instance.registry.return_value.deployments.side_effect + + def deployments(model_name=None, *, source="group", **_): + if source == "named": + return DEFAULT if named is None else named(model_name, source=source) if model_name not in vision_by_model: return [] declared = vision_by_model[model_name] @@ -15199,10 +15243,10 @@ class TestModalityRouting: } ] - mock_router_instance.get_model_list = get_model_list + mock_router_instance.registry.return_value.deployments.side_effect = deployments return ComplexityRouter( model_name="modality-test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -15314,7 +15358,7 @@ class TestModalityRouting: ] elif path.startswith(("pin_kept", "pin_replacement", "pin_override")): cache: Final = DualCache() - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache config["session_affinity"] = True request_kwargs = {"metadata": {"session_id": "s1"}} if path == "pin_replacement_gated": @@ -15332,11 +15376,10 @@ class TestModalityRouting: config["modality_pin_override"] = path.startswith("pin_override") elif path == "adaptive_pick_rewritten": config["adaptive"] = True - mock_router_instance.model_list = [] - _answer_deployments_named(mock_router_instance) + _answer_deployments_named(mock_router_instance, []) router = self._router(mock_router_instance, config, vision) if path.startswith(("pin_kept", "pin_replacement", "pin_override")): - mock_router_instance.cache.in_memory_cache.set_cache( + mock_router_instance.health().store.in_memory_cache.set_cache( router._get_session_affinity_cache_key("s1", request_kwargs), {"model": "text-cheap", "tier": "SIMPLE"}, ) @@ -15409,7 +15452,9 @@ class TestModalityRouting: @pytest.mark.asyncio async def test_mixed_deployment_group_is_treated_text_only(self, mock_router_instance): - def get_model_list(model_name=None): + def deployments(model_name=None, *, source="group", **_): + if source != "group": + return DEFAULT declared = {"mixed-group": [True, False], "vision-big": [True]}.get(model_name) if declared is None: return [] @@ -15422,10 +15467,10 @@ class TestModalityRouting: for i, accepts in enumerate(declared) ] - mock_router_instance.get_model_list = get_model_list + mock_router_instance.registry.return_value.deployments.side_effect = deployments router = ComplexityRouter( model_name="modality-test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {"SIMPLE": "mixed-group", "COMPLEX": "vision-big"}, "modality_routing": True, @@ -15439,7 +15484,7 @@ class TestModalityRouting: async def test_continuation_turn_screenshot_escalates_past_the_held_model(self, mock_router_instance): """classification_mode user_turn replays the held model on continuation turns; a continuation carrying a screenshot must still be re-placed when that model is text-only.""" - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() config = { "tiers": dict(self.BASE_TIERS), "classification_mode": "user_turn", @@ -15508,7 +15553,7 @@ class TestModalityRouting: async def test_pin_override_serves_the_image_turn_without_repinning(self, mock_router_instance): """The override is for one request: the session keeps the model it was pinned to.""" cache: Final = DualCache() - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache router = self._router( mock_router_instance, { @@ -15542,7 +15587,7 @@ class TestModalityRouting: async def test_pin_override_with_no_capable_model_rejects_and_keeps_the_pin(self, mock_router_instance): """The clear 400 replaces the provider's, and a rejected turn must not cost the session its pin.""" cache: Final = DualCache() - mock_router_instance.cache = cache + mock_router_instance.health.return_value.store = cache router = self._router( mock_router_instance, { @@ -16290,7 +16335,7 @@ class TestTierHealthFailover: ): """ids_by_model: model group -> deployment ids the router knows. - The fake mirrors the real async_get_healthy_deployments contract, including how it says + The fake mirrors the real select(kind="healthy") contract, including how it says no: BadRequestError for a group with no deployment at all, RouterRateLimitError when every deployment is filtered out (cooling, admin-paused, or excluded by a request-scoped policy such as tags, team scoping or access groups), a per-model exception via raises_for (the @@ -16303,8 +16348,16 @@ class TestTierHealthFailover: probed_kwargs = [] probed_prompts = [] - async def get_healthy_deployments( - model, request_kwargs, messages=None, input=None, parent_otel_span=None, health_check_probe=False + async def select( + model, + request_kwargs, + *, + kind, + messages=None, + input=None, + specific_deployment=False, + parent_otel_span=None, + health_check_probe=False, ): probed_kwargs.append(request_kwargs) probed_prompts.append((messages, input)) @@ -16326,13 +16379,13 @@ class TestTierHealthFailover: ) return healthy - mock_router_instance.async_get_healthy_deployments = get_healthy_deployments + mock_router_instance.select = select mock_router_instance.probed_kwargs = probed_kwargs mock_router_instance.probed_prompts = probed_prompts - mock_router_instance.cache = DualCache() + mock_router_instance.health.return_value.store = DualCache() return ComplexityRouter( model_name="health-test-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -16357,7 +16410,7 @@ class TestTierHealthFailover: ) # Seed the pin onto the dead group directly so the replay path is exercised. key = router._get_session_affinity_cache_key("sess-dead", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) result = await router.async_pre_routing_hook( @@ -16422,7 +16475,7 @@ class TestTierHealthFailover: ) pinned = "only" if tiers else "dead-a" key = router._get_session_affinity_cache_key("sess-open", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": pinned, "tier": "SIMPLE"}, ttl=600 ) result = await router.async_pre_routing_hook( @@ -16452,7 +16505,7 @@ class TestTierHealthFailover: request_kwargs={"metadata": {"session_id": session_id}}, messages=self.SIMPLE_MESSAGE, ) - return await router.litellm_router_instance.cache.async_get_cache( + return await router.ctx.health().store.async_get_cache( key=router._get_session_affinity_cache_key(session_id, {}) ) @@ -16477,7 +16530,7 @@ class TestTierHealthFailover: messages=[{"role": "user", "content": TITLE_ASK}], ) assert result.routing_decision["cause"] in ("housekeeping", "health_failover") - stored = await router.litellm_router_instance.cache.async_get_cache( + stored = await router.ctx.health().store.async_get_cache( key=router._get_session_affinity_cache_key(session_id, {}) ) assert stored is None @@ -16502,7 +16555,7 @@ class TestTierHealthFailover: blocked=("id-b1",), ) key = router._get_session_affinity_cache_key("sess-paused", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) results = [ @@ -16544,7 +16597,7 @@ class TestTierHealthFailover: cooling=("id-a1",), ) key = router._get_session_affinity_cache_key("sess-adaptive", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) request_kwargs = {"metadata": {"session_id": "sess-adaptive", "adaptive_router_chosen_model": "dead-a"}} @@ -16560,7 +16613,9 @@ class TestTierHealthFailover: modality gate chose: serving a cooling vision model beats a hard 400.""" vision_by_model = {"dead-vision": True, "live-text": False} - def get_model_list(model_name=None): + def deployments(model_name=None, *, source="group", **_): + if source != "group": + return DEFAULT if model_name not in vision_by_model: return [] return [ @@ -16571,7 +16626,7 @@ class TestTierHealthFailover: } ] - mock_router_instance.get_model_list = get_model_list + mock_router_instance.registry.return_value.deployments.side_effect = deployments router = self._router( mock_router_instance, { @@ -16588,7 +16643,7 @@ class TestTierHealthFailover: cooling=("id-v1",), ) key = router._get_session_affinity_cache_key("sess-image", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-vision", "tier": "SIMPLE"}, ttl=600 ) image_message = [ @@ -16626,7 +16681,7 @@ class TestTierHealthFailover: router_instance.cache = DualCache() router = ComplexityRouter( model_name="health-window-router", - litellm_router_instance=router_instance, + ctx=RoutingContext(router_instance), complexity_router_config={ "tiers": {name: list(pool) for name in ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING")}, "session_affinity": True, @@ -16634,7 +16689,7 @@ class TestTierHealthFailover: }, ) key = router._get_session_affinity_cache_key("sess-window", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-big", "tier": "SIMPLE"}, ttl=600 ) results = [ @@ -16680,7 +16735,7 @@ class TestTierHealthFailover: {"live-b": ["id-b1"]}, ) key = router._get_session_affinity_cache_key("sess-unknown", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) result = await router.async_pre_routing_hook( @@ -16709,7 +16764,7 @@ class TestTierHealthFailover: excluded=("id-b1",), ) key = router._get_session_affinity_cache_key("sess-tagged", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) results = [ @@ -16731,7 +16786,7 @@ class TestTierHealthFailover: cooling=("id-a1",), ) key = router._get_session_affinity_cache_key("sess-kwargs", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) request_kwargs = { @@ -16745,7 +16800,7 @@ class TestTierHealthFailover: assert result.model == "live-b" assert request_kwargs["_target_order"] == 1 assert request_kwargs["_excluded_deployment_ids"] == ["id-x"] - assert all(probed is not request_kwargs for probed in router.litellm_router_instance.probed_kwargs) + assert all(probed is not request_kwargs for probed in mock_router_instance.probed_kwargs) @pytest.mark.asyncio async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(self, mock_router_instance): @@ -16769,7 +16824,7 @@ class TestTierHealthFailover: raises_for={"rpm-full-b": RouterRateLimitErrorBasic(model="rpm-full-b")}, ) key = router._get_session_affinity_cache_key("sess-rpm", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) results = [ @@ -16792,7 +16847,7 @@ class TestTierHealthFailover: cooling=("id-a1",), ) key = router._get_session_affinity_cache_key("sess-input", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) result = await router.async_pre_routing_hook( @@ -16803,7 +16858,7 @@ class TestTierHealthFailover: assert result.model == "live-b" assert any( probed_input == "summarize this document for me" - for _, probed_input in router.litellm_router_instance.probed_prompts + for _, probed_input in mock_router_instance.probed_prompts ), "the eligibility probe must forward `input` to the owner" @pytest.mark.asyncio @@ -16842,7 +16897,7 @@ class TestTierHealthFailover: sessions: Final = tuple(f"sess-exhausted-{sample}" for sample in range(20)) await asyncio.gather( *( - router.litellm_router_instance.cache.async_set_cache( + router.ctx.health().store.async_set_cache( key=router._get_session_affinity_cache_key(session_id, {}), value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600, @@ -16889,7 +16944,7 @@ class TestTierHealthFailover: cooling=("id-a1",), ) key = router._get_session_affinity_cache_key("sess-missing", {}) - await router.litellm_router_instance.cache.async_set_cache( + await router.ctx.health().store.async_set_cache( key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 ) results = [ @@ -16912,7 +16967,9 @@ class TestClassifierVision: @staticmethod def _router(mock_router_instance, *, vision, classifier_declares_vision=True, classifier_type="llm", **extra): - def get_model_list(model_name=None): + def deployments(model_name=None, *, source="group", **_): + if source != "group": + return DEFAULT if model_name != "clf": return [{"model_name": model_name, "litellm_params": {"model": "openai/gpt-4o"}}] declared = classifier_declares_vision @@ -16924,11 +16981,11 @@ class TestClassifierVision: } ] - mock_router_instance.get_model_list = get_model_list + mock_router_instance.registry.return_value.deployments.side_effect = deployments classifier_llm_config = {"model": "clf", "circuit_breaker_enabled": False} return ComplexityRouter( model_name="vision-classifier-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "classifier_type": classifier_type, "classifier_llm_config": ( @@ -17569,7 +17626,7 @@ class TestMaxTokensFromTierModel: async def test_a_default_model_landing_with_no_tier_still_gets_its_ceiling(self): strategy = ComplexityRouter( model_name="smart-router", - litellm_router_instance=self._router(), + ctx=RoutingContext(self._router()), complexity_router_config={"tiers": {"SIMPLE": "small"}, "default_model": "big"}, ) @@ -17783,7 +17840,7 @@ class TestNonReasoningTier: } return ComplexityRouter( model_name="test-non-reasoning-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config=config, ) @@ -17923,7 +17980,7 @@ class TestNonReasoningTier: """The four-class artifact's 1-based index must keep mapping onto SIMPLE..REASONING.""" router: Final = ComplexityRouter( model_name="v2-router", - litellm_router_instance=mock_router_instance, + ctx=RoutingContext(mock_router_instance), complexity_router_config={ "tiers": {k: v for k, v in NON_REASONING_TIERS.items() if k != "NON_REASONING"}, "classifier_type": "heuristic_v2", @@ -17946,7 +18003,7 @@ class TestNonReasoningTier: async def test_custom_classifier_infers_tiers_and_routes_fallback(verdict: str | None, expected: str) -> None: router: Final = ComplexityRouter( model_name="auto", - litellm_router_instance=litellm.Router(model_list=[]), + ctx=RoutingContext(litellm.Router(model_list=[])), default_model="backup", complexity_router_config={ "classifier_type": "custom", @@ -17966,7 +18023,7 @@ async def test_custom_classifier_infers_tiers_and_routes_fallback(verdict: str | async def test_custom_classifier_can_extend_builtin_tiers() -> None: router: Final = ComplexityRouter( model_name="auto", - litellm_router_instance=litellm.Router(model_list=[]), + ctx=RoutingContext(litellm.Router(model_list=[])), default_model="backup", complexity_router_config={ "classifier_type": "custom", @@ -18000,7 +18057,7 @@ def test_inferred_custom_tiers_reject_unusable_routes(tiers: dict[str, str | lis async def test_defined_custom_tiers_allow_default_model_fallback(explicit_fallback: str | None) -> None: router: Final = ComplexityRouter( model_name="auto", - litellm_router_instance=litellm.Router(model_list=[]), + ctx=RoutingContext(litellm.Router(model_list=[])), default_model="backup", complexity_router_config={ "classifier_type": "custom", @@ -18050,7 +18107,7 @@ async def test_custom_classifier_failure_preserves_plan_floor_and_direct_default ) router: Final = ComplexityRouter( model_name="auto", - litellm_router_instance=litellm.Router(model_list=[]), + ctx=RoutingContext(litellm.Router(model_list=[])), default_model=default_model, complexity_router_config={ "classifier_type": "custom", diff --git a/tests/unit/router_strategy/test_lar1_routing.py b/tests/unit/router_strategy/test_lar1_routing.py index 4710bbb6896..64cd84fd4ff 100644 --- a/tests/unit/router_strategy/test_lar1_routing.py +++ b/tests/unit/router_strategy/test_lar1_routing.py @@ -56,7 +56,7 @@ def _create_test_router(): @pytest.mark.asyncio async def test_low_confidence_routes_to_cloud_smart(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -68,7 +68,7 @@ async def test_low_confidence_routes_to_cloud_smart(): @pytest.mark.asyncio async def test_high_confidence_routes_to_deep(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -80,7 +80,7 @@ async def test_high_confidence_routes_to_deep(): @pytest.mark.asyncio async def test_unverified_evidence_fallback(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -97,7 +97,7 @@ async def test_unverified_evidence_fallback(): async def test_custom_thresholds(): router = _create_test_router() custom_strategy = LAR1RoutingStrategy( - router, + router._core_ctx, thresholds={"low": 0.1, "medium": 0.3, "high": 0.9}, ) @@ -111,7 +111,7 @@ async def test_custom_thresholds(): @pytest.mark.asyncio async def test_mem_time_routes_to_cloud_fast(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -127,7 +127,7 @@ async def test_mem_time_routes_to_cloud_fast(): @pytest.mark.asyncio async def test_invalid_confidence_falls_back_to_local(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -141,7 +141,7 @@ async def test_invalid_confidence_falls_back_to_local(): @pytest.mark.asyncio async def test_no_metadata_defaults_to_local(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -185,7 +185,7 @@ async def test_router_init_lar1_default_thresholds(): @pytest.mark.asyncio async def test_mid_confidence_routes_to_cloud_fast(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -205,7 +205,7 @@ async def test_no_router_returns_none(): @pytest.mark.asyncio async def test_request_kwargs_none_uses_defaults(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -217,7 +217,7 @@ async def test_request_kwargs_none_uses_defaults(): @pytest.mark.asyncio async def test_invalid_lar1_metadata_type_uses_defaults(caplog): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) with caplog.at_level("WARNING"): result = await strategy.async_get_available_deployment( @@ -232,7 +232,7 @@ async def test_invalid_lar1_metadata_type_uses_defaults(caplog): @pytest.mark.asyncio async def test_confirmed_evidence_routes_by_confidence(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) result = await strategy.async_get_available_deployment( model="agent-router", @@ -248,8 +248,8 @@ async def test_confirmed_evidence_routes_by_confidence(): @pytest.mark.asyncio async def test_empty_healthy_deployments_returns_none(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) - router.async_get_healthy_deployments = AsyncMock(return_value=[]) + strategy = LAR1RoutingStrategy(router._core_ctx) + router.select = AsyncMock(return_value=[]) result = await strategy.async_get_available_deployment( model="agent-router", @@ -265,8 +265,8 @@ async def test_specific_deployment_dict_short_circuit(): "litellm_params": {"model": "ollama/mythos"}, } router = _create_test_router() - strategy = LAR1RoutingStrategy(router) - router.async_get_healthy_deployments = AsyncMock(return_value=deployment) + strategy = LAR1RoutingStrategy(router._core_ctx) + router.select = AsyncMock(return_value=deployment) result = await strategy.async_get_available_deployment( model="agent-router", @@ -279,8 +279,8 @@ async def test_specific_deployment_dict_short_circuit(): @pytest.mark.asyncio async def test_non_dict_healthy_deployment_returns_none(): router = _create_test_router() - strategy = LAR1RoutingStrategy(router) - router.async_get_healthy_deployments = AsyncMock(return_value=[None]) + strategy = LAR1RoutingStrategy(router._core_ctx) + router.select = AsyncMock(return_value=[None]) result = await strategy.async_get_available_deployment( model="agent-router", @@ -398,7 +398,7 @@ async def test_update_settings_routing_strategy_args_relinks_lar1(): def test_apply_lar1_routing_strategy_wires_custom_selector(): router = Router(model_list=_model_list(), routing_strategy="simple-shuffle") - apply_lar1_routing_strategy(router, {"confidence_threshold_high": 0.9}) + apply_lar1_routing_strategy(router._core_ctx, {"confidence_threshold_high": 0.9}) assert router.routing_strategy == "lar1" with pytest.raises(NotImplementedError, match="async routing"): router.get_available_deployment(model="agent-router") @@ -409,7 +409,7 @@ def test_apply_lar1_invalid_thresholds_leaves_router_unchanged(): with pytest.raises(ValueError, match="LAR-1 thresholds must satisfy"): apply_lar1_routing_strategy( - router, + router._core_ctx, { "confidence_threshold_low": 0.9, "confidence_threshold_medium": 0.5, @@ -448,7 +448,7 @@ async def test_missing_target_type_falls_back_with_warning(caplog): } ] ) - strategy = LAR1RoutingStrategy(router) + strategy = LAR1RoutingStrategy(router._core_ctx) with caplog.at_level("WARNING"): result = await strategy.async_get_available_deployment( diff --git a/tests/unit/router_strategy/test_litellm_encoder.py b/tests/unit/router_strategy/test_litellm_encoder.py index 46187f52adb..d950025afdb 100644 --- a/tests/unit/router_strategy/test_litellm_encoder.py +++ b/tests/unit/router_strategy/test_litellm_encoder.py @@ -8,6 +8,7 @@ import pytest import litellm from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS +from litellm.router_api import RoutingContext if sys.version_info >= (3, 14): pytest.skip("The semantic-router extra excludes Python 3.14", allow_module_level=True) @@ -39,7 +40,7 @@ class RecordingRouter: def _encoder(router: RecordingRouter, **kwargs: Any) -> LiteLLMRouterEncoder: return LiteLLMRouterEncoder( - litellm_router_instance=router, # pyright: ignore[reportArgumentType] # test double stands in for Router + ctx=RoutingContext(router), # pyright: ignore[reportArgumentType] # test double stands in for Router model_name="text-embedding-3-small", **kwargs, ) diff --git a/tests/unit/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py index 8c1e9e48661..993da777c66 100644 --- a/tests/unit/router_strategy/test_llm_v2.py +++ b/tests/unit/router_strategy/test_llm_v2.py @@ -9,6 +9,7 @@ from pydantic import ValidationError import litellm from litellm import ModelResponse, Router from litellm.caching.dual_cache import DualCache +from litellm.router_api import RoutingContext from litellm.router_strategy.auto_router.auto_router_model_naming import strategy_router_dependencies from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier @@ -90,7 +91,7 @@ def _router(content: str, config: ComplexityRouterConfig | None = None) -> tuple client.acompletion = AsyncMock(return_value=_response(content)) router: Final = ComplexityRouter( model_name="v2-router", - litellm_router_instance=client, + ctx=RoutingContext(client), complexity_router_config=(config or _config()).model_dump(), derive_savings_baseline=False, ) @@ -570,7 +571,7 @@ def test_response_schema_requires_both_model_forecasts() -> None: @pytest.mark.asyncio async def test_user_turn_mode_reuses_forecast_until_a_new_user_requirement() -> None: router, client = _router(_verdict().model_dump_json(), _config(classification_mode="user_turn")) - client.cache = DualCache() + client.health.return_value.store = DualCache() initial: Final = [{"role": "user", "content": "Fix nested behavior"}] first: Final = await router.async_pre_routing_hook( model="v2-router", messages=initial, request_kwargs={"metadata": {"session_id": "v2-task"}} @@ -670,7 +671,7 @@ def test_v2_judge_is_a_declared_dependency_for_authorization() -> None: async def test_v2_forwards_inline_images_only_when_vision_is_enabled(vision_enabled: bool) -> None: config: Final = _config(classifier_llm_config={"model": "judge", "vision": {"enabled": vision_enabled}}) router, client = _router(_verdict().model_dump_json(), config) - client.get_model_list.return_value = [ + client.registry.return_value.deployments.return_value = [ {"model_name": "judge", "litellm_params": {"model": "judge"}, "model_info": {"supports_vision": True}} ] image: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}} diff --git a/tests/unit/router_strategy/test_quality_router.py b/tests/unit/router_strategy/test_quality_router.py index 4e87652f8b3..249e5254412 100644 --- a/tests/unit/router_strategy/test_quality_router.py +++ b/tests/unit/router_strategy/test_quality_router.py @@ -15,6 +15,7 @@ from unittest.mock import MagicMock import pytest +from litellm.router_api import RoutingContext from litellm.router_strategy.quality_router.config import ( DEFAULT_COMPLEXITY_TO_QUALITY, ) @@ -74,7 +75,7 @@ def four_tier_model_list() -> List[Dict[str, Any]]: def mock_router(four_tier_model_list): """A MagicMock router preloaded with the four-tier model list.""" router = MagicMock() - router.model_list = four_tier_model_list + router.registry.return_value.all.return_value = four_tier_model_list return router @@ -87,7 +88,7 @@ def quality_router(mock_router) -> QualityRouter: } return QualityRouter( model_name="quality-router-test", - litellm_router_instance=mock_router, + ctx=RoutingContext(mock_router), default_model="haiku", quality_router_config=config, ) @@ -109,11 +110,11 @@ class TestTierIndex: # Add a model the config doesn't list — it should be ignored. extra = _make_model_list([{"model_name": "ghost", "quality_tier": 5}]) router = MagicMock() - router.model_list = four_tier_model_list + extra + router.registry.return_value.all.return_value = four_tier_model_list + extra qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="haiku", quality_router_config={ "available_models": ["haiku", "sonnet", "opus", "opus-next"] @@ -132,14 +133,14 @@ class TestTierIndex: ] ) router = MagicMock() - router.model_list = ml + router.registry.return_value.all.return_value = ml # Construction succeeds (tier index is lazy); the error surfaces on # first use so the router entry doesn't have to appear after all of # its referenced models in config.yaml. qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="haiku", quality_router_config={"available_models": ["haiku", "sonnet"]}, ) @@ -163,11 +164,11 @@ class TestResolveModelForQualityTier: {"model_name": "opus-next", "quality_tier": 4}, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="haiku", quality_router_config={"available_models": ["haiku", "opus", "opus-next"]}, ) @@ -179,11 +180,11 @@ class TestResolveModelForQualityTier: # rounds DOWN to the closest lower tier — tier 1. spec = [{"model_name": "haiku", "quality_tier": 1}] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="emergency-default", quality_router_config={"available_models": ["haiku"]}, ) @@ -197,11 +198,11 @@ class TestResolveModelForQualityTier: {"model_name": "sonnet", "quality_tier": 2}, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="emergency-default", quality_router_config={"available_models": ["haiku", "sonnet"]}, ) @@ -215,11 +216,11 @@ class TestResolveModelForQualityTier: {"model_name": "opus", "quality_tier": 3}, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="emergency-default", quality_router_config={"available_models": ["haiku", "opus"]}, ) @@ -244,11 +245,11 @@ class TestRoutingPreferencesValidation: } ] router = MagicMock() - router.model_list = ml + router.registry.return_value.all.return_value = ml qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="haiku", quality_router_config={"available_models": ["haiku"]}, ) @@ -266,18 +267,18 @@ class TestConfigOrderingIndependence: # appears before its target deployments). The tier index must not be # built eagerly — it's deferred until first use. router = MagicMock() - router.model_list = [] # <- targets haven't been added yet + router.registry.return_value.all.return_value = [] # <- targets haven't been added yet qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="haiku", quality_router_config={"available_models": ["haiku", "sonnet", "opus"]}, ) # Now the targets come online. This mirrors the incremental add by # `Router._create_deployment`. - router.model_list = _make_model_list( + router.registry.return_value.all.return_value = _make_model_list( [ {"model_name": "haiku", "quality_tier": 1}, {"model_name": "sonnet", "quality_tier": 2}, @@ -489,10 +490,10 @@ def keyword_router(): }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) return QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="default-haiku", quality_router_config={ "available_models": [ @@ -547,10 +548,10 @@ class TestKeywordOverride: }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="cheap", quality_router_config={"available_models": ["expensive", "cheap"]}, ) @@ -574,10 +575,10 @@ class TestKeywordOverride: }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="no-price", quality_router_config={"available_models": ["no-price", "with-price"]}, ) @@ -633,10 +634,10 @@ class TestKeywordOverride: }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="ordered-tier2", quality_router_config={ "available_models": ["ordered-tier2", "implicit-tier3"] @@ -663,10 +664,10 @@ class TestKeywordOverride: }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="default-tier3", quality_router_config={"available_models": ["preferred", "default-tier3"]}, ) @@ -691,10 +692,10 @@ class TestKeywordOverride: }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="cheap-default", quality_router_config={ "available_models": ["expensive-but-preferred", "cheap-default"] @@ -719,10 +720,10 @@ class TestKeywordOverride: }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="first", quality_router_config={"available_models": ["first", "second"]}, ) @@ -756,10 +757,10 @@ class TestKeywordOverride: }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="low-tier", quality_router_config={ "available_models": [ @@ -787,10 +788,10 @@ class TestKeywordOverride: }, ] router = MagicMock() - router.model_list = _make_model_list(spec) + router.registry.return_value.all.return_value = _make_model_list(spec) qr = QualityRouter( model_name="qr", - litellm_router_instance=router, + ctx=RoutingContext(router), default_model="default-pick", quality_router_config={ "available_models": ["default-pick", "preferred-pick"] diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index ba35b64b1a4..9aba894eb36 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -20,6 +20,7 @@ import litellm from litellm import Router from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger +from litellm.router_api import RoutingContext from litellm.router_strategy.least_busy import picks_reserve from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy from litellm.utils import Rules, function_setup @@ -1604,7 +1605,7 @@ async def _call_and_get_cooldowns(router, model): messages=[{"role": "user", "content": "hi"}], mock_response="litellm.RateLimitError", ) - return await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + return await async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) @pytest.mark.asyncio @@ -2097,7 +2098,7 @@ async def test_priority_group_preserves_responses_continuity_on_a_backup() -> No def _recording_fallback_gate(allowed: frozenset[str], checked: list[str]) -> FallbackAccessCheck: - async def check(*, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool: + async def check(*, model: str, request_kwargs: Mapping[str, object], ctx: RoutingContext) -> bool: checked.append(model) return model in allowed diff --git a/tests/unit/router_strategy/test_router_tag_regex_routing.py b/tests/unit/router_strategy/test_router_tag_regex_routing.py index 46e79a1b8c8..4f80cc54d3f 100644 --- a/tests/unit/router_strategy/test_router_tag_regex_routing.py +++ b/tests/unit/router_strategy/test_router_tag_regex_routing.py @@ -9,9 +9,8 @@ existing tag-based routing. import pytest - -from unittest.mock import MagicMock - +from litellm import Router +from litellm.router_api import RoutingContext from litellm.router_strategy import tag_based_routing from litellm.router_strategy.tag_based_routing import get_deployments_for_tag @@ -105,20 +104,17 @@ REGULAR_DEPLOYMENT = { ALL_DEPLOYMENTS = [CLAUDE_CODE_DEPLOYMENT, REGULAR_DEPLOYMENT] -def _make_router_mock(enable_tag_filtering=True, match_any=True): - mock = MagicMock() - mock.enable_tag_filtering = enable_tag_filtering - mock.tag_filtering_match_any = match_any - mock.tag_routing_prefix = "" - return mock +def _make_routing_ctx(enable_tag_filtering=True, match_any=True) -> RoutingContext: + """The candidates come in as healthy_deployments, so the router itself configures none.""" + return Router(model_list=[], enable_tag_filtering=enable_tag_filtering, tag_filtering_match_any=match_any)._core_ctx @pytest.mark.asyncio async def test_claude_code_ua_routes_to_cc_deployment(): """claude-code/x.y.z UA → claude-code-deployment via tag_regex.""" - router = _make_router_mock() + ctx = _make_routing_ctx() result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=ALL_DEPLOYMENTS, request_deployments=ALL_DEPLOYMENTS, @@ -131,9 +127,9 @@ async def test_claude_code_ua_routes_to_cc_deployment(): @pytest.mark.asyncio async def test_regular_ua_routes_to_default_deployment(): """Mozilla UA → regular-deployment via default tag fallback.""" - router = _make_router_mock() + ctx = _make_routing_ctx() result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=ALL_DEPLOYMENTS, request_deployments=ALL_DEPLOYMENTS, @@ -146,9 +142,9 @@ async def test_regular_ua_routes_to_default_deployment(): @pytest.mark.asyncio async def test_no_ua_routes_to_default_deployment(): """No User-Agent → default deployment.""" - router = _make_router_mock() + ctx = _make_routing_ctx() result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=ALL_DEPLOYMENTS, request_deployments=ALL_DEPLOYMENTS, @@ -161,10 +157,10 @@ async def test_no_ua_routes_to_default_deployment(): @pytest.mark.asyncio async def test_tag_routing_metadata_written_for_regex_match(): """tag_routing metadata block is populated when regex matches.""" - router = _make_router_mock() + ctx = _make_routing_ctx() metadata: dict = {"user_agent": "claude-code/2.0.0-beta.1"} await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=ALL_DEPLOYMENTS, request_deployments=ALL_DEPLOYMENTS, @@ -180,9 +176,9 @@ async def test_tag_routing_metadata_written_for_regex_match(): @pytest.mark.asyncio async def test_tag_filtering_disabled_returns_all_deployments(): """When enable_tag_filtering is False, all deployments returned regardless of UA.""" - router = _make_router_mock(enable_tag_filtering=False) + ctx = _make_routing_ctx(enable_tag_filtering=False) result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=ALL_DEPLOYMENTS, request_deployments=ALL_DEPLOYMENTS, @@ -199,9 +195,9 @@ async def test_empty_healthy_deployments_with_request_tags_returns_empty_list(): Tag-based filtering runs only against non-empty healthy_deployments; an empty list is returned unchanged for the router's standard handling. """ - router = _make_router_mock() + ctx = _make_routing_ctx() result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="gpt-5.2", healthy_deployments=[], request_deployments=[], @@ -225,13 +221,13 @@ async def test_explicit_tag_match_takes_precedence_over_regex(): }, "model_info": {"id": "both-deployment"}, } - router = _make_router_mock() + ctx = _make_routing_ctx() metadata: dict = { "tags": ["premium"], "user_agent": "claude-code/1.0", } result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=[deployment_with_both], request_deployments=[deployment_with_both], @@ -269,11 +265,11 @@ async def test_user_agent_present_no_tag_regex_deployments_does_not_raise(): "model_info": {"id": "free-deployment"}, }, ] - router = _make_router_mock() + ctx = _make_routing_ctx() # The request has a User-Agent (as all proxy requests do) but NO tags and # neither deployment has tag_regex — must not raise, must return all. result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="gpt-4", healthy_deployments=plain_tag_only_deployments, request_deployments=plain_tag_only_deployments, @@ -307,10 +303,10 @@ async def test_tag_routing_metadata_not_overwritten_for_multiple_matches(): }, "model_info": {"id": "cc-deployment-b"}, } - router = _make_router_mock() + ctx = _make_routing_ctx() metadata: dict = {"user_agent": "claude-code/1.0"} result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=[deployment_a, deployment_b], request_deployments=[deployment_a, deployment_b], @@ -357,13 +353,13 @@ async def test_match_any_false_strict_tag_check_blocks_regex_fallback(): # match_any=False: req_set must be a subset of dep_set. # Request has "vip" which is NOT in ["premium"], so tag check fails. # Even though UA matches tag_regex, the deployment must NOT be selected. - router = _make_router_mock(enable_tag_filtering=True, match_any=False) + ctx = _make_routing_ctx(enable_tag_filtering=True, match_any=False) metadata: dict = { "tags": ["vip"], # "vip" not in deployment tags → strict check fails "user_agent": "claude-code/1.0", } result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=[deployment_strict, default_deployment], request_deployments=[deployment_strict, default_deployment], @@ -392,9 +388,9 @@ async def test_match_any_false_regex_only_deployment_still_matches(): }, "model_info": {"id": "regex-only-deployment"}, } - router = _make_router_mock(enable_tag_filtering=True, match_any=False) + ctx = _make_routing_ctx(enable_tag_filtering=True, match_any=False) result = await get_deployments_for_tag( - llm_router_instance=router, + ctx=ctx, model="claude-sonnet", healthy_deployments=[regex_only_deployment], request_deployments=[regex_only_deployment], diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index 26d04be7e25..28192704bd1 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -1168,31 +1168,29 @@ async def test_request_level_enable_tag_filtering_false_cannot_disable_global(): # --- model_info.enable_tag_filtering per-chain override --- -class _FakeRouterForChainOverride: - def __init__(self, all_deployments): - self._all_deployments = all_deployments - - def get_all_deployments(self, model_name): - return self._all_deployments +def _chain_override_router(*model_infos: dict) -> litellm.Router: + return litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}, "model_info": {"id": f"gpt-4-{i}", **info}} + for i, info in enumerate(model_infos) + ] + ) def test_chain_tag_filtering_override_reads_any_member(): from litellm.router_strategy.tag_based_routing import _chain_tag_filtering_override - deployments = [ - {"model_info": {}}, - {"model_info": {"enable_tag_filtering": False}}, - ] - router = _FakeRouterForChainOverride(deployments) - assert _chain_tag_filtering_override(router, "gpt-4", deployments) is False + router = _chain_override_router({}, {"enable_tag_filtering": False}) + deployments = router.registry().deployments("gpt-4", source="all") + assert _chain_tag_filtering_override(router._core_ctx, "gpt-4", deployments) is False def test_chain_tag_filtering_override_none_when_unset_anywhere(): from litellm.router_strategy.tag_based_routing import _chain_tag_filtering_override - deployments = [{"model_info": {}}, {}] - router = _FakeRouterForChainOverride(deployments) - assert _chain_tag_filtering_override(router, "gpt-4", deployments) is None + router = _chain_override_router({}, {}) + deployments = router.registry().deployments("gpt-4", source="all") + assert _chain_tag_filtering_override(router._core_ctx, "gpt-4", deployments) is None def test_chain_tag_filtering_override_survives_the_overriding_member_going_unhealthy(): @@ -1203,27 +1201,28 @@ def test_chain_tag_filtering_override_survives_the_overriding_member_going_unhea # the one deployment that carries the group's only explicit override. from litellm.router_strategy.tag_based_routing import _chain_tag_filtering_override - all_deployments = [ - {"model_info": {"enable_tag_filtering": True}}, - {"model_info": {}}, - ] - router = _FakeRouterForChainOverride(all_deployments) + router = _chain_override_router({"enable_tag_filtering": True}, {}) + all_deployments = router.registry().deployments("gpt-4", source="all") # The overriding deployment (index 0) is cooled down and absent from # healthy_deployments -- the override must still be found via the full-group # lookup, not silently lost. healthy_deployments = [all_deployments[1]] - assert _chain_tag_filtering_override(router, "gpt-4", healthy_deployments) is True + assert _chain_tag_filtering_override(router._core_ctx, "gpt-4", healthy_deployments) is True def test_chain_tag_filtering_override_falls_back_to_healthy_deployments_on_lookup_error(): from litellm.router_strategy.tag_based_routing import _chain_tag_filtering_override - class _BrokenRouter: - def get_all_deployments(self, model_name): + class _BrokenRegistry: + def deployments(self, model_name, *, source): raise RuntimeError("model group not found") + class _BrokenCtx: + def registry(self): + return _BrokenRegistry() + healthy_deployments = [{"model_info": {"enable_tag_filtering": False}}] - assert _chain_tag_filtering_override(_BrokenRouter(), "gpt-4", healthy_deployments) is False + assert _chain_tag_filtering_override(_BrokenCtx(), "gpt-4", healthy_deployments) is False @pytest.mark.asyncio() @@ -2754,14 +2753,14 @@ def test_tag_known_to_group_true_for_real_tag(): from litellm.router_strategy.tag_based_routing import _tag_known_to_group router = _quality_high_cost_low_router() - assert _tag_known_to_group(router, "gpt-4", ["quality:high"], frozenset()) is True + assert _tag_known_to_group(router._core_ctx, "gpt-4", ["quality:high"], frozenset()) is True def test_tag_known_to_group_false_for_foreign_tag(): from litellm.router_strategy.tag_based_routing import _tag_known_to_group router = _quality_high_cost_low_router() - assert _tag_known_to_group(router, "gpt-4", ["llm-preference-include:unrelated"], frozenset()) is False + assert _tag_known_to_group(router._core_ctx, "gpt-4", ["llm-preference-include:unrelated"], frozenset()) is False def test_inherited_constraint_sets_none_when_inherited_tags_absent(): @@ -3456,7 +3455,7 @@ async def _routed_to(router: litellm.Router, metadata: dict[str, object]) -> str async def _pool(router: litellm.Router, metadata: dict[str, object]) -> frozenset[str]: request_deployments: Final = router.get_model_list(model_name="chat") deployments: Final = await get_deployments_for_tag( - llm_router_instance=router, + ctx=router._core_ctx, model="chat", healthy_deployments=request_deployments, request_deployments=request_deployments, diff --git a/tests/unit/router_strategy/test_savings_baseline.py b/tests/unit/router_strategy/test_savings_baseline.py index 4241a541202..a6c99ba061c 100644 --- a/tests/unit/router_strategy/test_savings_baseline.py +++ b/tests/unit/router_strategy/test_savings_baseline.py @@ -1,6 +1,7 @@ import pytest from litellm.router import Router +from litellm.router_api import RoutingContext from litellm.router_strategy.savings_baseline import ( Baseline, canonical_model, @@ -11,7 +12,7 @@ from litellm.router_strategy.savings_baseline import ( @pytest.fixture -def parent() -> Router: +def parent() -> RoutingContext: return Router( model_list=[ {"model_name": "cheap", "litellm_params": {"model": "anthropic/claude-haiku-4-5"}}, @@ -19,7 +20,7 @@ def parent() -> Router: {"model_name": "pool", "litellm_params": {"model": "anthropic/claude-haiku-4-5"}}, {"model_name": "pool", "litellm_params": {"model": "anthropic/claude-opus-5"}}, ] - ) + )._core_ctx class TestCanonicalModel: @@ -107,7 +108,7 @@ class TestResolveBaseline: served; a dashboard counterfactual must not be able to take routing down.""" class Exploding: - def deployments_named(self, model_name): + def registry(self): raise RuntimeError("router is mid-reload") assert resolve_baseline(Exploding(), ["anything"]) is None @@ -128,8 +129,8 @@ class TestDeploymentsPricedByBaseModel: """ @staticmethod - def _router(*deployments: dict) -> Router: - return Router(model_list=list(deployments)) + def _router(*deployments: dict) -> RoutingContext: + return Router(model_list=list(deployments))._core_ctx def test_model_info_base_model_is_preferred_over_the_deployment_name(self): router = self._router( @@ -192,7 +193,7 @@ class TestDeploymentPricingOverrides: {"model_name": "top", "litellm_params": {"model": "anthropic/claude-opus-5"}}, ] ) - assert resolve_baseline(router, ["cheap", "top"]).model == "anthropic/claude-opus-5" + assert resolve_baseline(router._core_ctx, ["cheap", "top"]).model == "anthropic/claude-opus-5" overridden = Router( model_list=[ @@ -207,4 +208,4 @@ class TestDeploymentPricingOverrides: {"model_name": "top", "litellm_params": {"model": "anthropic/claude-opus-5"}}, ] ) - assert resolve_baseline(overridden, ["cheap", "top"]).model == "anthropic/claude-haiku-4-5" + assert resolve_baseline(overridden._core_ctx, ["cheap", "top"]).model == "anthropic/claude-haiku-4-5" diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index a3745c7b23f..d184bf84035 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -24,6 +24,7 @@ import pytest import litellm from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_api import RoutingContext from litellm.types.llms.openai import ResponsesAPIResponse @@ -1016,13 +1017,13 @@ async def test_affinity_falls_back_to_same_boundary_on_alias_switch(): def test_boundary_fallback_no_router_ref_returns_empty(): """ Standalone use (no router wired in) -> the boundary lookup short-circuits - to ``[]`` instead of crashing on ``None.get_deployment``. + to ``[]`` instead of crashing on ``None.registry()``. """ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) - check = EncryptedContentAffinityCheck(router=None) + check = EncryptedContentAffinityCheck(ctx=None) healthy = [ { "model_info": {"id": "dep-1"}, @@ -1040,7 +1041,7 @@ def test_boundary_fallback_no_router_ref_returns_empty(): def test_boundary_fallback_originating_deployment_removed_returns_empty(): """ If the originating deployment has been removed from the router (e.g. via - /model/delete), ``router.get_deployment`` returns None and we return [] so + /model/delete), ``ctx.registry().deployment`` returns None and we return [] so the caller falls back to the full healthy_deployments list. """ from unittest.mock import MagicMock @@ -1049,10 +1050,10 @@ def test_boundary_fallback_originating_deployment_removed_returns_empty(): EncryptedContentAffinityCheck, ) - mock_router = MagicMock() - mock_router.get_deployment.return_value = None + mock_ctx = MagicMock() + mock_ctx.registry.return_value.deployment.return_value = None - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) healthy = [ { "model_info": {"id": "dep-1"}, @@ -1065,7 +1066,7 @@ def test_boundary_fallback_originating_deployment_removed_returns_empty(): ) assert matches == [] assert originating is None - mock_router.get_deployment.assert_called_once_with(model_id="dep-removed") + mock_ctx.registry().deployment.assert_called_once_with("dep-removed") def test_boundary_key_accepts_pydantic_litellm_params_instance(): @@ -1241,7 +1242,7 @@ def test_boundary_fallback_matches_deployments_with_same_named_credential_values ], num_retries=0, ) - check = EncryptedContentAffinityCheck(router=router) + check = EncryptedContentAffinityCheck(ctx=RoutingContext(router)) healthy_deployments = [ { "model_info": {"id": "peer-same-boundary"}, @@ -1308,19 +1309,21 @@ def _make_router_mock_with_cooldown( routed_group_model_ids: list[str] | None = None, ): """ - Build a MagicMock router whose ``cooldown_cache.async_get_active_cooldowns`` + Build a MagicMock routing context whose ``health().cooldowns.async_get_active_cooldowns`` returns ``cooldown_entries`` (defaulting to ``[]`` — no active cooldown), and - whose ``get_candidate_model_ids_for_route`` returns ``routed_group_model_ids`` + whose ``registry().names(model, kind="route_candidates")`` returns ``routed_group_model_ids`` (the deployment ids the router resolves for the routed model; defaulting to ``[]`` — origin absent from the routed group, i.e. a tier change). """ from unittest.mock import AsyncMock, MagicMock - mock_router = MagicMock() - mock_router.get_deployment.return_value = originating - mock_router.cooldown_cache.async_get_active_cooldowns = AsyncMock(return_value=list(cooldown_entries or [])) - mock_router.get_candidate_model_ids_for_route.return_value = frozenset(routed_group_model_ids or []) - return mock_router + mock_ctx = MagicMock() + mock_ctx.registry.return_value.deployment.return_value = originating + mock_ctx.health.return_value.cooldowns.async_get_active_cooldowns = AsyncMock( + return_value=list(cooldown_entries or []) + ) + mock_ctx.registry.return_value.names.return_value = frozenset(routed_group_model_ids or []) + return mock_ctx @pytest.mark.asyncio @@ -1336,7 +1339,7 @@ async def test_affinity_strips_and_dispatches_when_origin_cooled_for_non_429(): ) originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[ ( @@ -1352,7 +1355,7 @@ async def test_affinity_strips_and_dispatches_when_origin_cooled_for_non_429(): routed_group_model_ids=["deployment-a-cooled", "deployment-b"], ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-cooled", "rs_test") healthy_only_b = [ { @@ -1402,7 +1405,7 @@ async def test_affinity_strips_and_dispatches_when_origin_cooled_for_429(): originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") cooldown_started = time.time() - 5.0 - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[ ( @@ -1418,7 +1421,7 @@ async def test_affinity_strips_and_dispatches_when_origin_cooled_for_429(): routed_group_model_ids=["deployment-a-cooled-429", "deployment-b"], ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-cooled-429", "rs_test") healthy_only_b = [ { @@ -1457,11 +1460,11 @@ async def test_affinity_strips_and_dispatches_when_origin_filtered_without_coold ) originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[], routed_group_model_ids=["deployment-a-filtered", "deployment-b"] ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-filtered", "rs_test") healthy_only_b = [ { @@ -1501,10 +1504,10 @@ async def test_affinity_serves_sibling_when_candidate_origin_has_no_boundary_pee ) originating = _make_originating_mock("https://region-a.example.com/v1", "shared-key") - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[], routed_group_model_ids=["region-a", "region-b", "region-c"] ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", "region-a") siblings = [ { @@ -1563,10 +1566,10 @@ async def test_affinity_strips_and_dispatches_when_origin_is_unknown_or_removed( EncryptedContentAffinityCheck, ) - mock_router = MagicMock() - mock_router.get_deployment.return_value = None + mock_ctx = MagicMock() + mock_ctx.registry.return_value.deployment.return_value = None - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-removed") routed_pool = [ { @@ -1601,7 +1604,7 @@ async def test_affinity_strips_and_dispatches_when_origin_is_unknown_or_removed( assert result is routed_pool assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in request_kwargs["input"]) - mock_router.get_candidate_model_ids_for_route.assert_not_called() + mock_ctx.registry().names.assert_not_called() @pytest.mark.asyncio @@ -1622,10 +1625,10 @@ async def test_affinity_does_not_raise_when_boundary_peer_available(): "api_base": "https://account-a.openai.azure.com/", "api_key": "key-a", } - mock_router = MagicMock() - mock_router.get_deployment.return_value = originating + mock_ctx = MagicMock() + mock_ctx.registry.return_value.deployment.return_value = originating - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a", "rs_test") peer = { "model_info": {"id": "deployment-a-peer"}, @@ -1881,10 +1884,10 @@ async def test_encrypted_content_affinity_strips_bridge_reasoning_from_messages_ same way they reject a foreign one. """ originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[], routed_group_model_ids=["openai-org-b"] ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) routed_pool = [{"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5-nano"}}] messages = _bridge_replayed_anthropic_messages(minted_by="openai-org-a") assistant_content = messages[1]["content"] @@ -2134,7 +2137,7 @@ async def test_affinity_keeps_mixed_origins_on_the_same_encryption_boundary(): shared_api_base = "https://account-a.openai.azure.com/" shared_api_key = "shared-key" origin_d2 = _make_originating_mock(shared_api_base, shared_api_key) - mock_router = _make_router_mock_with_cooldown(origin_d2, cooldown_entries=[], routed_group_model_ids=["d1", "d2"]) + mock_ctx = _make_router_mock_with_cooldown(origin_d2, cooldown_entries=[], routed_group_model_ids=["d1", "d2"]) deployment_d1 = { "model_info": {"id": "d1"}, "litellm_params": {"api_base": shared_api_base, "api_key": shared_api_key}, @@ -2158,8 +2161,8 @@ async def test_affinity_keeps_mixed_origins_on_the_same_encryption_boundary(): d2_item.copy(), ] } - mock_router.get_deployment.side_effect = lambda model_id: origin_d2 if model_id == "d2" else None - check = EncryptedContentAffinityCheck(router=mock_router) + mock_ctx.registry.return_value.deployment.side_effect = lambda model_id: origin_d2 if model_id == "d2" else None + check = EncryptedContentAffinityCheck(ctx=mock_ctx) result = await check.async_filter_deployments( model="gpt-5.4", @@ -2180,8 +2183,11 @@ async def test_boundary_pin_strips_reasoning_from_a_different_origin(): origin_a = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") origin_b = _make_originating_mock("https://account-b.openai.azure.com/", "key-b") - mock_router = _make_router_mock_with_cooldown(origin_a, cooldown_entries=[], routed_group_model_ids=["peer-a"]) - mock_router.get_deployment.side_effect = lambda model_id: {"origin-a": origin_a, "origin-b": origin_b}.get(model_id) + mock_ctx = _make_router_mock_with_cooldown(origin_a, cooldown_entries=[], routed_group_model_ids=["peer-a"]) + mock_ctx.registry.return_value.deployment.side_effect = lambda model_id: { + "origin-a": origin_a, + "origin-b": origin_b, + }.get(model_id) peer_a = { "model_info": {"id": "peer-a"}, "litellm_params": {"api_base": "https://account-a.openai.azure.com/", "api_key": "key-a"}, @@ -2204,7 +2210,7 @@ async def test_boundary_pin_strips_reasoning_from_a_different_origin(): }, ] } - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) result = await check.async_filter_deployments( model="gpt-5.4", @@ -2234,8 +2240,11 @@ async def test_affinity_keeps_only_anthropic_reasoning_from_the_pinned_origin(): origin_a = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") origin_b = _make_originating_mock("https://account-b.openai.azure.com/", "key-b") - mock_router = _make_router_mock_with_cooldown(origin_b, cooldown_entries=[], routed_group_model_ids=["origin-a"]) - mock_router.get_deployment.side_effect = lambda model_id: {"origin-a": origin_a, "origin-b": origin_b}.get(model_id) + mock_ctx = _make_router_mock_with_cooldown(origin_b, cooldown_entries=[], routed_group_model_ids=["origin-a"]) + mock_ctx.registry.return_value.deployment.side_effect = lambda model_id: { + "origin-a": origin_a, + "origin-b": origin_b, + }.get(model_id) deployment_a = { "model_info": {"id": "origin-a"}, "litellm_params": {"api_base": "https://account-a.openai.azure.com/", "api_key": "key-a"}, @@ -2248,7 +2257,7 @@ async def test_affinity_keeps_only_anthropic_reasoning_from_the_pinned_origin(): foreign_messages = _bridge_replayed_anthropic_messages(minted_by="origin-b") assistant_content = messages[1]["content"] assistant_content.insert(3, foreign_messages[1]["content"][1]) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) result = await check.async_filter_deployments( model="gpt-5.4", @@ -2288,8 +2297,8 @@ async def test_affinity_strips_unknown_origins_but_leaves_unmarked_encrypted_con EncryptedContentAffinityCheck, ) - mock_router = MagicMock() - mock_router.get_deployment.return_value = None + mock_ctx = MagicMock() + mock_ctx.registry.return_value.deployment.return_value = None deployment_a = { "model_info": {"id": "origin-a"}, "litellm_params": {"api_base": "https://account-a.openai.azure.com/", "api_key": "key-a"}, @@ -2316,7 +2325,7 @@ async def test_affinity_strips_unknown_origins_but_leaves_unmarked_encrypted_con }, ] } - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) result = await check.async_filter_deployments( model="gpt-5.4", @@ -2354,11 +2363,13 @@ def _router_without_the_origin(): ) -@pytest.mark.parametrize("router", [None, _router_without_the_origin()], ids=["no router", "origin removed"]) +@pytest.mark.parametrize( + "ctx", [None, RoutingContext(_router_without_the_origin())], ids=["no router", "origin removed"] +) @pytest.mark.parametrize( "unmarked_origin", ["origin-removed", None], ids=["failed deployment named", "failed deployment unknown"] ) -def test_hop_strip_drops_unmarked_reasoning_whose_origin_cannot_be_resolved(router, unmarked_origin): +def test_hop_strip_drops_unmarked_reasoning_whose_origin_cannot_be_resolved(ctx, unmarked_origin): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -2378,7 +2389,7 @@ def test_hop_strip_drops_unmarked_reasoning_whose_origin_cannot_be_resolved(rout } EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( - router, request_input, None, (target,), unmarked_origin=unmarked_origin + ctx, request_input, None, (target,), unmarked_origin=unmarked_origin ) assert request_input == [ @@ -2417,10 +2428,10 @@ async def test_affinity_strips_encrypted_reasoning_when_routed_to_another_model_ ) originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[], routed_group_model_ids=["deployment-b"] ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) routed_pool = [ { "model_info": {"id": "deployment-b"}, @@ -2475,10 +2486,10 @@ async def test_affinity_degrades_within_the_origins_own_group(): originating = _make_originating_mock( "https://account-a.openai.azure.com/", "key-a", model_name="gpt-reasoning-tier" ) - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[], routed_group_model_ids=["deployment-a", "deployment-a2"] ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) sibling_pool = [ { "model_info": {"id": "deployment-a2"}, @@ -2517,10 +2528,10 @@ async def test_affinity_strips_when_group_is_spelled_differently_but_same_by_id( ) originating = _make_originating_mock(None, "key-a", model_name="openai/gpt-5.4-mini") - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[], routed_group_model_ids=["deployment-mini-a", "deployment-mini-b"] ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-mini-a") sibling_pool = [ { @@ -2569,10 +2580,10 @@ async def test_affinity_strips_for_team_and_pattern_routes(): ) originating = _make_originating_mock(None, "key-a", model_name="model_name_teamA_uuid") - mock_router = _make_router_mock_with_cooldown( + mock_ctx = _make_router_mock_with_cooldown( originating, cooldown_entries=[], routed_group_model_ids=["deployment-team-a", "deployment-team-b"] ) - check = EncryptedContentAffinityCheck(router=mock_router) + check = EncryptedContentAffinityCheck(ctx=mock_ctx) wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-team-a") sibling_pool = [ { @@ -2657,7 +2668,7 @@ async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_re group: Final = router.model_list reasoning_input: Final = _reasoning_input_minted_by("region-a") - result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments( + result: Final = await EncryptedContentAffinityCheck(ctx=RoutingContext(router)).async_filter_deployments( model="gpt-5.4", healthy_deployments=group, messages=None, @@ -2681,7 +2692,9 @@ async def test_affinity_keeps_origin_and_reasoning_when_previous_response_id_bin reasoning_input: Final = _reasoning_input_minted_by("region-a") original_reasoning: Final = dict(reasoning_input[0]) - assert await EncryptedContentAffinityCheck(router=litellm.Router(model_list=list(pool))).async_filter_deployments( + assert await EncryptedContentAffinityCheck( + ctx=RoutingContext(litellm.Router(model_list=list(pool))) + ).async_filter_deployments( model="gpt-5.4", healthy_deployments=list(pool), messages=None, @@ -2706,7 +2719,7 @@ async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_th reasoning_input: Final = _reasoning_input_minted_by("region-a") original_reasoning: Final = dict(reasoning_input[0]) - result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments( + result: Final = await EncryptedContentAffinityCheck(ctx=RoutingContext(router)).async_filter_deployments( model="gpt-5.4", healthy_deployments=router.model_list, messages=None, diff --git a/tests/unit/router_utils/test_access_windows.py b/tests/unit/router_utils/test_access_windows.py index a797a788b4e..a283b2d57c1 100644 --- a/tests/unit/router_utils/test_access_windows.py +++ b/tests/unit/router_utils/test_access_windows.py @@ -2,7 +2,7 @@ from datetime import datetime, time, timezone from typing import Final -from litellm.router_utils.access_windows import ( +from litellm.litellm_core_utils.access_windows import ( access_windows_config_error, filter_reserved_deployments, is_window_active, diff --git a/tests/unit/router_utils/test_add_retry_fallback_headers.py b/tests/unit/router_utils/test_add_retry_fallback_headers.py index 1f4ac3be6ac..0f65e0ecf48 100644 --- a/tests/unit/router_utils/test_add_retry_fallback_headers.py +++ b/tests/unit/router_utils/test_add_retry_fallback_headers.py @@ -5,7 +5,7 @@ from typing import Final, Literal import pytest from pydantic import BaseModel -from litellm.router_utils.add_retry_fallback_headers import ( +from litellm.litellm_core_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, add_retry_headers_to_response, complexity_router_decision_headers, diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 2e1e3acaabb..7deeb255e6c 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -20,6 +20,7 @@ from litellm._internal_context import current_service_target from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials +from litellm.router_api import RoutingContext from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback from litellm.router_utils.cooldown_handlers import ( @@ -57,6 +58,25 @@ from litellm.types.utils import CooldownDeploymentInfo, CredentialItem from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +def _ctx_double( + *, + deployment_dict: Mapping[str, object] | None = None, + cooldown_time: float = 60.0, + allowed_fails: int = 0, + store: DualCache | MagicMock | None = None, +) -> MagicMock: + """A RoutingContext double answering the registry, settings and health reads the cooldown helpers make.""" + ctx: Final = MagicMock() + ctx.registry.return_value.deployment.return_value = None + ctx.registry.return_value.deployment_dict.return_value = deployment_dict + ctx.settings.return_value.cooldown_time = cooldown_time + ctx.settings.return_value.allowed_fails = allowed_fails + ctx.settings.return_value.allowed_fails_policy = None + ctx.settings.return_value.allowed_fails_for.return_value = None + ctx.health.return_value.store = store if store is not None else DualCache(in_memory_cache=InMemoryCache()) + return ctx + + @pytest.mark.asyncio @pytest.mark.parametrize( ("exception_type", "expected_cooldown"), @@ -80,25 +100,20 @@ async def test_unsupported_endpoint_preserves_cooldown_exclusion_with_bad_reques ) exception: Final = exception_type(message="Request failed", model="claude-sonnet-5-5", llm_provider="anthropic") set_cooldown_deployments( - litellm_router_instance=router, + ctx=RoutingContext(router), original_exception=exception, exception_status=exception.status_code, deployment="deployment", time_to_cooldown=60, ) - assert await async_get_cooldown_deployments(router, parent_otel_span=None) == ( + assert await async_get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == ( ["deployment"] if expected_cooldown else [] ) class TestGetDeploymentCooldownPolicy: - def _make_router(self, deployment_id: str, model_info: dict | None = None): - router = MagicMock() - if model_info is None: - router.get_model_info.return_value = None - else: - router.get_model_info.return_value = {"model_info": model_info} - return router + def _make_router(self, deployment_id: str, model_info: dict | None = None) -> MagicMock: + return _ctx_double(deployment_dict=None if model_info is None else {"model_info": model_info}) def test_deployment_not_found_returns_none_none(self): router = self._make_router("dep-1") @@ -107,8 +122,7 @@ class TestGetDeploymentCooldownPolicy: assert allowed is None def test_no_model_info_returns_none_none(self): - router = MagicMock() - router.get_model_info.return_value = {"model_info": {}} + router = _ctx_double(deployment_dict={"model_info": {}}) policy, allowed = _get_deployment_cooldown_policy(router, "dep-1") assert policy is None assert allowed is None @@ -186,13 +200,8 @@ class TestResolveAllowedFailsFromPolicy: class TestShouldCooldownBasedOnDeploymentPolicy: - def _make_router(self, model_info: dict | None = None): - router = MagicMock() - if model_info is None: - router.get_model_info.return_value = None - else: - router.get_model_info.return_value = model_info - return router + def _make_router(self, model_info: dict | None = None) -> MagicMock: + return _ctx_double(deployment_dict=model_info) def test_policy_match_uses_exception_type_as_cache_key_suffix(self): policy = {"RateLimitErrorAllowedFails": 0} @@ -351,31 +360,25 @@ class TestShouldCooldownBasedOnDeploymentPolicy: class TestShouldCooldownBasedOnAllowedFailsPolicy: - def _make_router(self, cooldown_time: float = 60.0, cache: DualCache | None = None) -> MagicMock: - router = MagicMock() - router.cooldown_time = cooldown_time - router.allowed_fails = 0 - router.allowed_fails_policy = None - router.get_allowed_fails_from_policy.return_value = None - router.cache = cache if cache is not None else DualCache(in_memory_cache=InMemoryCache()) - return router + def _make_router(self, cooldown_time: float = 60.0, cache: DualCache | MagicMock | None = None) -> MagicMock: + return _ctx_double(cooldown_time=cooldown_time, allowed_fails=0, store=cache) def test_cooldown_time_override_zero_is_not_falsy(self): """cooldown_time_override=0 must be honored; it must not fall through to the router-level value.""" - router = self._make_router(cooldown_time=60.0) - router.cache = MagicMock() - router.cache.increment_cache.return_value = 1 + store = MagicMock() + store.increment_cache.return_value = 1 + ctx = self._make_router(cooldown_time=60.0, cache=store) exc = litellm.RateLimitError("429", "openai", "gpt-4") should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5, cooldown_time_override=0.0, ) - increment_call = router.cache.increment_cache.call_args + increment_call = store.increment_cache.call_args assert increment_call is not None assert increment_call[1]["ttl"] == 0.0, ( "cooldown_time_override=0 should be used as TTL, not the router-level 60.0" @@ -389,7 +392,7 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy: results = [ should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=workers[i % 2], + ctx=workers[i % 2], deployment="dep-1", original_exception=exc, allowed_fails_override=5, @@ -408,7 +411,7 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy: exc = litellm.AuthenticationError("401", "openai", "gpt-4") result = should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=router, deployment="dep-1", original_exception=exc, allowed_fails_override=5, @@ -426,7 +429,7 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy: results = [ should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=router, deployment="dep-1", original_exception=exc, allowed_fails_override=5, @@ -472,7 +475,7 @@ class TestRoutingGroupCooldownAlternatives: ) assert ( _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="cg-deploy-1", exception_status=429, original_exception=Exception("rate limited"), @@ -495,7 +498,7 @@ class TestRoutingGroupCooldownAlternatives: ) assert ( _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="cg-deploy-1", exception_status=429, original_exception=Exception("rate limited"), @@ -510,7 +513,7 @@ class TestRoutingGroupCooldownAlternatives: router = self._router(routing_groups=None) assert ( _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="cg-deploy-1", exception_status=429, original_exception=Exception("rate limited"), @@ -543,7 +546,7 @@ class TestTeamModelCooldownAlternatives: router = self._router(team_deployments=2) assert ( _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="team-deploy-0", exception_status=429, original_exception=Exception("rate limited"), @@ -558,7 +561,7 @@ class TestTeamModelCooldownAlternatives: router = self._router(team_deployments=1) assert ( _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="team-deploy-0", exception_status=429, original_exception=Exception("rate limited"), @@ -573,7 +576,7 @@ class TestTeamModelCooldownAlternatives: router = self._router(team_deployments=2, blocked_ids=frozenset({"team-deploy-1"})) assert ( _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="team-deploy-0", exception_status=429, original_exception=Exception("rate limited"), @@ -675,7 +678,7 @@ class TestDeploymentLevelAllowedFails: _exception = litellm.RateLimitError("Rate limit", "openai", "gpt-4") should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="primary", exception_status=429, original_exception=_exception, @@ -709,7 +712,7 @@ class TestDeploymentLevelAllowedFails: _exception = litellm.RateLimitError("Rate limit", "openai", "gpt-4") should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="secondary", exception_status=429, original_exception=_exception, @@ -731,7 +734,7 @@ def _primary_cooldown_decisions_on_repeated_503s(router: Router) -> tuple[bool, exception: Final = litellm.ServiceUnavailableError("down", "openai", "gpt-4") return tuple( _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="primary", exception_status=503, original_exception=exception, @@ -743,7 +746,9 @@ def _primary_cooldown_decisions_on_repeated_503s(router: Router) -> tuple[bool, @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestRouterLevelAllowedFailsEqualToDefault: def test_allowed_fails_equal_to_the_default_cools_down_after_that_many_fails(self): - router: Final = _make_router(model_list=list(_gpt_4_group_with_a_sibling()), allowed_fails=litellm.allowed_fails) + router: Final = _make_router( + model_list=list(_gpt_4_group_with_a_sibling()), allowed_fails=litellm.allowed_fails + ) assert _primary_cooldown_decisions_on_repeated_503s(router) == (False,) * litellm.allowed_fails + (True,) @@ -785,7 +790,7 @@ class TestDeploymentLevelAllowedFailsPolicyByExceptionType: rate_limit_exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="primary", exception_status=429, original_exception=rate_limit_exc, @@ -818,7 +823,7 @@ class TestDeploymentLevelAllowedFailsPolicyByExceptionType: for _ in range(5): should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="primary", exception_status=500, original_exception=ise, @@ -826,7 +831,7 @@ class TestDeploymentLevelAllowedFailsPolicyByExceptionType: assert should_cooldown is False, "Should not cooldown within the allowed_fails threshold" should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="primary", exception_status=500, original_exception=ise, @@ -857,7 +862,7 @@ class TestExceptionTypeCountersTrackedIndependently: for _ in range(3): should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="primary", original_exception=rate_limit_exc, allowed_fails_override=5, @@ -871,7 +876,7 @@ class TestExceptionTypeCountersTrackedIndependently: assert generic_counter == 0, "generic counter must be untouched by RateLimitError increments" should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="primary", original_exception=ise, allowed_fails_override=5, @@ -994,16 +999,14 @@ class TestFallbackDeploymentCooldown: _trigger_cooldown_for_failed_deployment must call set_cooldown_deployments with the deployment ID stamped on the exception. """ - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _ctx_double(cooldown_time=60.0) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs={}, exception=exc, ) @@ -1018,11 +1021,11 @@ class TestFallbackDeploymentCooldown: _trigger_cooldown_for_failed_deployment must not raise and must skip set_cooldown_deployments when the exception has no failed_deployment_id. """ - mock_router = MagicMock() + ctx = _ctx_double() with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs={}, exception=RuntimeError("no stamped deployment id"), ) @@ -1036,9 +1039,7 @@ class TestFallbackDeploymentCooldown: set metadata must not be able to get an arbitrary deployment cooled down by forging a deployment_model_name marker. """ - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _ctx_double(cooldown_time=60.0) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") kwargs = { @@ -1050,7 +1051,7 @@ class TestFallbackDeploymentCooldown: with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs=kwargs, exception=exc, ) @@ -1063,9 +1064,7 @@ class TestFallbackDeploymentCooldown: primary path uses, or repeated fallback failures never accumulate toward the default percent-fail-rate cooldown threshold. """ - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _ctx_double(cooldown_time=60.0) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" @@ -1076,11 +1075,9 @@ class TestFallbackDeploymentCooldown: "litellm.router_utils.fallback_event_handlers.increment_deployment_failures_for_current_minute" ) as mock_increment, ): - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) - mock_increment.assert_called_once_with( - litellm_router_instance=mock_router, deployment_id="fallback-deployment" - ) + mock_increment.assert_called_once_with(ctx=ctx, deployment_id="fallback-deployment") mock_set_cooldown.assert_called_once() def test_trigger_cooldown_uses_deployment_cooldown_time_override(self): @@ -1088,16 +1085,14 @@ class TestFallbackDeploymentCooldown: When the deployment has a model_info.cooldown_time, that value must be passed as time_to_cooldown rather than the router-level cooldown_time. """ - mock_router = MagicMock() - mock_router.cooldown_time = 300.0 - mock_router.get_model_info.return_value = {"model_info": {"cooldown_time": 30.0}} + ctx = _ctx_double(cooldown_time=300.0, deployment_dict={"model_info": {"cooldown_time": 30.0}}) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs={}, exception=exc, ) @@ -1113,9 +1108,7 @@ class TestFallbackDeploymentCooldown: deployment) must not cool down the fallback deployment, matching the same guard already applied in Router.deployment_callback_on_failure. """ - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _ctx_double(cooldown_time=60.0) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" @@ -1123,7 +1116,7 @@ class TestFallbackDeploymentCooldown: with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs={}, exception=exc, ) @@ -1137,16 +1130,14 @@ class TestFallbackDeploymentCooldown: honored as a fallback when model_info doesn't set it, unlike the new allowed_fails/allowed_fails_policy fields which are model_info-only. """ - mock_router = MagicMock() - mock_router.cooldown_time = 300.0 - mock_router.get_model_info.return_value = {"litellm_params": {"cooldown_time": 30.0}} + ctx = _ctx_double(cooldown_time=300.0, deployment_dict={"litellm_params": {"cooldown_time": 30.0}}) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs={}, exception=exc, ) @@ -1157,19 +1148,20 @@ class TestFallbackDeploymentCooldown: ) def test_trigger_cooldown_prefers_model_info_cooldown_time_over_litellm_params(self): - mock_router = MagicMock() - mock_router.cooldown_time = 300.0 - mock_router.get_model_info.return_value = { - "model_info": {"cooldown_time": 15.0}, - "litellm_params": {"cooldown_time": 30.0}, - } + ctx = _ctx_double( + cooldown_time=300.0, + deployment_dict={ + "model_info": {"cooldown_time": 15.0}, + "litellm_params": {"cooldown_time": 30.0}, + }, + ) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs={}, exception=exc, ) @@ -1200,7 +1192,7 @@ class TestSingleDeploymentModelGroupProtection: exc = Exception("Internal error") for _ in range(2): should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="solo", exception_status=500, original_exception=exc, @@ -1230,7 +1222,7 @@ class TestSingleDeploymentModelGroupProtection: exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="solo", exception_status=429, original_exception=exc, @@ -1260,7 +1252,7 @@ class TestShouldCooldownBasedOnAllowedFailsPolicyFalsyZero: exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") should_cooldown = should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="primary", original_exception=exc, ) @@ -1425,8 +1417,8 @@ class TestCallerScopedOAuthAuthFailureCooldown: result: Final = self._callback(router, model_id, self._auth_exception(status)) assert result is False - assert get_deployment_failures_for_current_minute(router, model_id) == 0 - assert get_cooldown_deployments(router, parent_otel_span=None) == [] + assert get_deployment_failures_for_current_minute(RoutingContext(router), model_id) == 0 + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [] @pytest.mark.asyncio @pytest.mark.parametrize("status", (401, 403)) @@ -1456,8 +1448,8 @@ class TestCallerScopedOAuthAuthFailureCooldown: result: Final = self._callback(router, model_id, self._auth_exception(status)) assert result is False - assert get_deployment_failures_for_current_minute(router, model_id) == 0 - assert get_cooldown_deployments(router, parent_otel_span=None) == [] + assert get_deployment_failures_for_current_minute(RoutingContext(router), model_id) == 0 + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [] @pytest.mark.asyncio async def test_api_key_deployment_still_cools_down_on_401(self) -> None: @@ -1470,8 +1462,8 @@ class TestCallerScopedOAuthAuthFailureCooldown: litellm.AuthenticationError("upstream rejected the API key", "openai", "gpt-4o-mini"), ) - assert get_deployment_failures_for_current_minute(router, model_id) == 1 - assert get_cooldown_deployments(router, parent_otel_span=None) == [model_id] + assert get_deployment_failures_for_current_minute(RoutingContext(router), model_id) == 1 + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [model_id] @pytest.mark.asyncio @pytest.mark.parametrize("status", (429, 500)) @@ -1497,8 +1489,8 @@ class TestCallerScopedOAuthAuthFailureCooldown: self._callback(router, model_id, exception) - assert get_deployment_failures_for_current_minute(router, model_id) == 1 - assert get_cooldown_deployments(router, parent_otel_span=None) == [model_id] + assert get_deployment_failures_for_current_minute(RoutingContext(router), model_id) == 1 + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [model_id] @pytest.mark.parametrize("status", (401, 403)) def test_fallback_oauth_caller_auth_failure_does_not_cooldown(self, status: int) -> None: @@ -1514,10 +1506,10 @@ class TestCallerScopedOAuthAuthFailureCooldown: exception: Final = self._auth_exception(status) exception.failed_deployment_id = model_id - _trigger_cooldown_for_failed_deployment(litellm_router=router, kwargs={}, exception=exception) + _trigger_cooldown_for_failed_deployment(ctx=RoutingContext(router), kwargs={}, exception=exception) - assert get_deployment_failures_for_current_minute(router, model_id) == 0 - assert get_cooldown_deployments(router, parent_otel_span=None) == [] + assert get_deployment_failures_for_current_minute(RoutingContext(router), model_id) == 0 + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [] @pytest.mark.asyncio async def test_fallback_api_key_auth_failure_still_cools_down(self) -> None: @@ -1526,10 +1518,10 @@ class TestCallerScopedOAuthAuthFailureCooldown: exception: Final = litellm.AuthenticationError("API key rejected", "openai", "gpt-4o-mini") exception.failed_deployment_id = model_id - _trigger_cooldown_for_failed_deployment(litellm_router=router, kwargs={}, exception=exception) + _trigger_cooldown_for_failed_deployment(ctx=RoutingContext(router), kwargs={}, exception=exception) - assert get_deployment_failures_for_current_minute(router, model_id) == 1 - assert get_cooldown_deployments(router, parent_otel_span=None) == [model_id] + assert get_deployment_failures_for_current_minute(RoutingContext(router), model_id) == 1 + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [model_id] @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -1595,7 +1587,7 @@ class TestPerUserCopilotAuthFailureCooldown: ) result: Final = self._callback(router, model_id, exception) assert result is False - assert get_cooldown_deployments(router, parent_otel_span=None) == [] + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [] @pytest.mark.asyncio async def test_shared_mode_copilot_auth_failure_still_cools_down(self) -> None: @@ -1612,8 +1604,8 @@ class TestPerUserCopilotAuthFailureCooldown: model_id, litellm.AuthenticationError("upstream rejected the shared token", "github_copilot", "gpt-4o"), ) - assert get_deployment_failures_for_current_minute(router, model_id) == 1 - assert get_cooldown_deployments(router, parent_otel_span=None) == [model_id] + assert get_deployment_failures_for_current_minute(RoutingContext(router), model_id) == 1 + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [model_id] @pytest.mark.asyncio async def test_fallback_per_user_caller_auth_failure_does_not_cooldown( @@ -1640,9 +1632,9 @@ class TestPerUserCopilotAuthFailureCooldown: ) exception.failed_deployment_id = model_id - _trigger_cooldown_for_failed_deployment(litellm_router=router, kwargs={}, exception=exception) + _trigger_cooldown_for_failed_deployment(ctx=RoutingContext(router), kwargs={}, exception=exception) - assert get_cooldown_deployments(router, parent_otel_span=None) == [] + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [] @pytest.mark.asyncio async def test_fallback_shared_mode_auth_failure_still_cools_down(self) -> None: @@ -1659,9 +1651,9 @@ class TestPerUserCopilotAuthFailureCooldown: ) exception.failed_deployment_id = model_id - _trigger_cooldown_for_failed_deployment(litellm_router=router, kwargs={}, exception=exception) + _trigger_cooldown_for_failed_deployment(ctx=RoutingContext(router), kwargs={}, exception=exception) - assert get_cooldown_deployments(router, parent_otel_span=None) == [model_id] + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [model_id] @pytest.mark.asyncio async def test_fallback_rate_limit_on_per_user_deployment_does_not_cooldown( @@ -1690,9 +1682,9 @@ class TestPerUserCopilotAuthFailureCooldown: exception: Final = litellm.RateLimitError("copilot 429", "github_copilot", "gpt-4o") exception.failed_deployment_id = model_id - _trigger_cooldown_for_failed_deployment(litellm_router=router, kwargs={}, exception=exception) + _trigger_cooldown_for_failed_deployment(ctx=RoutingContext(router), kwargs={}, exception=exception) - assert get_cooldown_deployments(router, parent_otel_span=None) == [] + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [] @pytest.mark.asyncio async def test_fallback_rate_limit_on_shared_deployment_still_cools_down(self) -> None: @@ -1711,9 +1703,9 @@ class TestPerUserCopilotAuthFailureCooldown: exception: Final = litellm.RateLimitError("copilot 429", "github_copilot", "gpt-4o") exception.failed_deployment_id = model_id - _trigger_cooldown_for_failed_deployment(litellm_router=router, kwargs={}, exception=exception) + _trigger_cooldown_for_failed_deployment(ctx=RoutingContext(router), kwargs={}, exception=exception) - assert get_cooldown_deployments(router, parent_otel_span=None) == [model_id] + assert get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) == [model_id] @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -1814,19 +1806,18 @@ async def test_router_cooldown_event_callback_no_deployment(): In this scenario it should do nothing """ - # Mock Router instance - mock_router = MagicMock() - mock_router.get_deployment.return_value = None + # Routing context whose registry has no such deployment + ctx = _ctx_double() await router_cooldown_event_callback( - litellm_router_instance=mock_router, + ctx=ctx, deployment_id="test-deployment", exception_status="429", cooldown_time=60.0, ) - # Assert that the router's get_deployment method was called - mock_router.get_deployment.assert_called_once_with(model_id="test-deployment") + # Assert that the registry was asked for the deployment + ctx.registry().deployment.assert_called_once_with("test-deployment") @pytest.fixture @@ -1856,15 +1847,21 @@ def testing_litellm_router(): def test_should_run_cooldown_logic(testing_litellm_router): testing_litellm_router.disable_cooldowns = True # don't run cooldown logic if disable_cooldowns is True - assert _should_run_cooldown_logic(testing_litellm_router, "test_deployment", 500, Exception("Test")) is False + assert ( + _should_run_cooldown_logic(RoutingContext(testing_litellm_router), "test_deployment", 500, Exception("Test")) + is False + ) # don't cooldown if deployment is None testing_litellm_router.disable_cooldowns = False - assert _should_run_cooldown_logic(testing_litellm_router, None, 500, Exception("Test")) is False + assert _should_run_cooldown_logic(RoutingContext(testing_litellm_router), None, 500, Exception("Test")) is False # don't cooldown if it's a provider default deployment testing_litellm_router.provider_default_deployment_ids = ["test_deployment"] - assert _should_run_cooldown_logic(testing_litellm_router, "test_deployment", 500, Exception("Test")) is False + assert ( + _should_run_cooldown_logic(RoutingContext(testing_litellm_router), "test_deployment", 500, Exception("Test")) + is False + ) @pytest.fixture @@ -1893,7 +1890,7 @@ def test_should_run_cooldown_logic_generic_bad_request_excluded_by_default( configured for that exception type. This is the pre-existing, intentional default: a client error is usually not the deployment's fault.""" exc = litellm.BadRequestError("bad request", "openai", "gpt-5-mini") - assert _should_run_cooldown_logic(single_deployment_router, "dep-1", 400, exc) is False + assert _should_run_cooldown_logic(RoutingContext(single_deployment_router), "dep-1", 400, exc) is False @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -1907,7 +1904,7 @@ def test_should_run_cooldown_logic_router_level_policy_does_not_override_bad_req see test_should_run_cooldown_logic_explicit_deployment_level_policy_overrides_content_policy_exclusion.""" exc = litellm.BadRequestError("bad request", "openai", "gpt-5-mini") single_deployment_router.allowed_fails_policy = AllowedFailsPolicy(BadRequestErrorAllowedFails=5) - assert _should_run_cooldown_logic(single_deployment_router, "dep-1", 400, exc) is False + assert _should_run_cooldown_logic(RoutingContext(single_deployment_router), "dep-1", 400, exc) is False @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -1919,14 +1916,17 @@ def test_should_run_cooldown_logic_explicit_deployment_level_policy_overrides_co exc = litellm.ContentPolicyViolationError("flagged content", "openai", "gpt-5-mini") deployment_dict = single_deployment_router.get_model_info(id="dep-1") deployment_dict["model_info"]["allowed_fails_policy"] = {"ContentPolicyViolationErrorAllowedFails": 0} - assert _should_run_cooldown_logic(single_deployment_router, "dep-1", 400, exc) is True + assert _should_run_cooldown_logic(RoutingContext(single_deployment_router), "dep-1", 400, exc) is True @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestHasExplicitAllowedFailsPolicyForException: def test_no_policy_anywhere_returns_false(self, single_deployment_router): exc = litellm.BadRequestError("bad request", "openai", "gpt-5-mini") - assert _has_explicit_allowed_fails_policy_for_exception(single_deployment_router, "dep-1", exc) is False + assert ( + _has_explicit_allowed_fails_policy_for_exception(RoutingContext(single_deployment_router), "dep-1", exc) + is False + ) def test_router_level_policy_for_matching_exception_returns_false(self, single_deployment_router): """Deliberately scoped to deployment-level only: a router-level policy @@ -1934,23 +1934,35 @@ class TestHasExplicitAllowedFailsPolicyForException: opt-in for cooldown-gate purposes.""" exc = litellm.RateLimitError("rate limited", "openai", "gpt-5-mini") single_deployment_router.allowed_fails_policy = AllowedFailsPolicy(RateLimitErrorAllowedFails=3) - assert _has_explicit_allowed_fails_policy_for_exception(single_deployment_router, "dep-1", exc) is False + assert ( + _has_explicit_allowed_fails_policy_for_exception(RoutingContext(single_deployment_router), "dep-1", exc) + is False + ) def test_router_level_policy_for_different_exception_returns_false(self, single_deployment_router): exc = litellm.BadRequestError("bad request", "openai", "gpt-5-mini") single_deployment_router.allowed_fails_policy = AllowedFailsPolicy(RateLimitErrorAllowedFails=3) - assert _has_explicit_allowed_fails_policy_for_exception(single_deployment_router, "dep-1", exc) is False + assert ( + _has_explicit_allowed_fails_policy_for_exception(RoutingContext(single_deployment_router), "dep-1", exc) + is False + ) def test_deployment_level_policy_for_matching_exception_returns_true(self, single_deployment_router): exc = litellm.ContentPolicyViolationError("flagged", "openai", "gpt-5-mini") deployment_dict = single_deployment_router.get_model_info(id="dep-1") deployment_dict["model_info"]["allowed_fails_policy"] = {"ContentPolicyViolationErrorAllowedFails": 0} - assert _has_explicit_allowed_fails_policy_for_exception(single_deployment_router, "dep-1", exc) is True + assert ( + _has_explicit_allowed_fails_policy_for_exception(RoutingContext(single_deployment_router), "dep-1", exc) + is True + ) def test_none_deployment_returns_false(self, single_deployment_router): exc = litellm.RateLimitError("rate limited", "openai", "gpt-5-mini") single_deployment_router.allowed_fails_policy = AllowedFailsPolicy(RateLimitErrorAllowedFails=3) - assert _has_explicit_allowed_fails_policy_for_exception(single_deployment_router, None, exc) is False + assert ( + _has_explicit_allowed_fails_policy_for_exception(RoutingContext(single_deployment_router), None, exc) + is False + ) @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -1960,7 +1972,9 @@ def test_should_cooldown_deployment_rate_limit_error(testing_litellm_router): """ # Test 429 error (rate limit) -> always cooldown a deployment returning 429s _exception = litellm.exceptions.RateLimitError("Rate limit", "openai", "gpt-5-mini") - assert _should_cooldown_deployment(testing_litellm_router, "test_deployment", 429, _exception) is True + assert ( + _should_cooldown_deployment(RoutingContext(testing_litellm_router), "test_deployment", 429, _exception) is True + ) @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -1970,7 +1984,9 @@ def test_should_cooldown_deployment_auth_limit_error(testing_litellm_router): """ # Test 401 error (auth limit) -> always cooldown a deployment returning 401s _exception = litellm.exceptions.AuthenticationError("Unauthorized", "openai", "gpt-5-mini") - assert _should_cooldown_deployment(testing_litellm_router, "test_deployment", 401, _exception) is True + assert ( + _should_cooldown_deployment(RoutingContext(testing_litellm_router), "test_deployment", 401, _exception) is True + ) @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -1978,7 +1994,7 @@ def test_should_cooldown_deployment_auth_limit_error(testing_litellm_router): def test_is_cooldown_required_for_account_errors(testing_litellm_router, exception_status): assert ( _is_cooldown_required( - litellm_router_instance=testing_litellm_router, + ctx=RoutingContext(testing_litellm_router), model_id="test_deployment", exception_status=exception_status, ) @@ -1993,15 +2009,17 @@ def test_single_deployment_402_does_not_cooldown( ) -> None: assert ( _should_cooldown_deployment( - Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini"}, - "model_info": {"id": "dep-1"}, - }, - ], - allowed_fails=allowed_fails, + RoutingContext( + Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini"}, + "model_info": {"id": "dep-1"}, + }, + ], + allowed_fails=allowed_fails, + ) ), "dep-1", 402, @@ -2019,15 +2037,17 @@ def test_single_deployment_402_does_not_cooldown( def test_single_deployment_402_respects_router_allowed_fails_policy() -> None: assert ( _should_cooldown_deployment( - Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini"}, - "model_info": {"id": "dep-1"}, - }, - ], - allowed_fails_policy=AllowedFailsPolicy(BadRequestErrorAllowedFails=0), + RoutingContext( + Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini"}, + "model_info": {"id": "dep-1"}, + }, + ], + allowed_fails_policy=AllowedFailsPolicy(BadRequestErrorAllowedFails=0), + ) ), "dep-1", 402, @@ -2045,17 +2065,19 @@ def test_single_deployment_402_respects_router_allowed_fails_policy() -> None: def test_single_deployment_402_respects_deployment_allowed_fails_policy() -> None: assert ( _should_cooldown_deployment( - Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini"}, - "model_info": { - "id": "dep-1", - "allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}, + RoutingContext( + Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini"}, + "model_info": { + "id": "dep-1", + "allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}, + }, }, - }, - ], + ], + ) ), "dep-1", 402, @@ -2073,7 +2095,7 @@ def test_single_deployment_402_respects_deployment_allowed_fails_policy() -> Non def test_multi_deployment_402_cools_down(testing_litellm_router: Router) -> None: assert ( _should_cooldown_deployment( - testing_litellm_router, + RoutingContext(testing_litellm_router), "test_deployment", 402, litellm.PaymentRequiredError( @@ -2100,7 +2122,9 @@ async def test_should_cooldown_deployment(testing_litellm_router): # Test 429 error (rate limit) -> always cooldown a deployment returning 429s _exception = litellm.exceptions.RateLimitError("Rate limit", "openai", "gpt-5-mini") - assert _should_cooldown_deployment(testing_litellm_router, "test_deployment", 429, _exception) is True + assert ( + _should_cooldown_deployment(RoutingContext(testing_litellm_router), "test_deployment", 429, _exception) is True + ) available_deployment = testing_litellm_router.get_available_deployment(model="test_deployment") print("available_deployment", available_deployment) @@ -2112,7 +2136,7 @@ async def test_should_cooldown_deployment(testing_litellm_router): # set current success for deployment to 40 for _ in range(40): increment_deployment_successes_for_current_minute( - litellm_router_instance=testing_litellm_router, deployment_id=deployment_id + ctx=RoutingContext(testing_litellm_router), deployment_id=deployment_id ) # now we fail 40 requests in a row @@ -2134,7 +2158,10 @@ async def test_should_cooldown_deployment(testing_litellm_router): await asyncio.sleep(1) # expect this to fail since it's now 51% of requests are failing - assert _should_cooldown_deployment(testing_litellm_router, deployment_id, 500, Exception("Test")) is True + assert ( + _should_cooldown_deployment(RoutingContext(testing_litellm_router), deployment_id, 500, Exception("Test")) + is True + ) @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -2159,9 +2186,9 @@ async def test_should_cooldown_deployment_allowed_fails_set_on_router(): # should not cooldown when fails are below the allowed limit for _ in range(100): - assert _should_cooldown_deployment(router, "test_deployment", 500, Exception("Test")) is False + assert _should_cooldown_deployment(RoutingContext(router), "test_deployment", 500, Exception("Test")) is False - assert _should_cooldown_deployment(router, "test_deployment", 500, Exception("Test")) is True + assert _should_cooldown_deployment(RoutingContext(router), "test_deployment", 500, Exception("Test")) is True @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @@ -2187,11 +2214,11 @@ def test_increment_deployment_successes_for_current_minute_does_not_write_to_red # Call the function we're testing increment_deployment_successes_for_current_minute( - litellm_router_instance=testing_litellm_router, deployment_id="test_deployment" + ctx=RoutingContext(testing_litellm_router), deployment_id="test_deployment" ) increment_deployment_failures_for_current_minute( - litellm_router_instance=testing_litellm_router, deployment_id="test_deployment" + ctx=RoutingContext(testing_litellm_router), deployment_id="test_deployment" ) time.sleep(1) @@ -2239,7 +2266,7 @@ def test_should_cooldown_high_traffic_all_fails(mock_failures, mock_successes, r mock_successes.return_value = 0 should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="gpt-4--0", exception_status=500, original_exception=Exception("Test error"), @@ -2257,7 +2284,7 @@ def test_no_cooldown_low_traffic(mock_failures, mock_successes, router): mock_successes.return_value = 0 should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="gpt-4--0", exception_status=500, original_exception=Exception("Test error"), @@ -2277,7 +2304,7 @@ def test_cooldown_rate_limit(mock_failures, mock_successes, router): mock_successes.return_value = 0 should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="gpt-4--0", exception_status=429, # Rate limit error original_exception=Exception("Rate limit exceeded"), @@ -2295,7 +2322,7 @@ def test_mixed_success_failure(mock_failures, mock_successes, router): mock_successes.return_value = 7 should_cooldown = _should_cooldown_deployment( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="gpt-4--0", exception_status=500, original_exception=Exception("Test error"), @@ -2310,7 +2337,7 @@ def test_is_cooldown_required_empty_string_exception_status(testing_litellm_rout Test that _is_cooldown_required returns False when exception_status is an empty string """ result = _is_cooldown_required( - litellm_router_instance=testing_litellm_router, + ctx=RoutingContext(testing_litellm_router), model_id="test_deployment", exception_status="", ) @@ -2342,13 +2369,15 @@ def test_should_cooldown_deployment_minimum_request_threshold(testing_litellm_ro # Simulate only 1 failure (below minimum threshold) # This should NOT trigger cooldown even though 100% > 50% increment_deployment_failures_for_current_minute( - litellm_router_instance=testing_litellm_router, deployment_id=deployment_id + ctx=RoutingContext(testing_litellm_router), deployment_id=deployment_id ) _exception = litellm.exceptions.InternalServerError("Internal error", "openai", "gpt-5-mini") # With only 1 request, should NOT cooldown (below minimum threshold) - should_cooldown = _should_cooldown_deployment(testing_litellm_router, deployment_id, 500, _exception) + should_cooldown = _should_cooldown_deployment( + RoutingContext(testing_litellm_router), deployment_id, 500, _exception + ) assert should_cooldown is False, ( f"Should NOT cooldown with only 1 failed request (below minimum threshold of {DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS})" ) @@ -2356,11 +2385,13 @@ def test_should_cooldown_deployment_minimum_request_threshold(testing_litellm_ro # Now add more failures to reach the minimum threshold for _ in range(DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS - 1): increment_deployment_failures_for_current_minute( - litellm_router_instance=testing_litellm_router, deployment_id=deployment_id + ctx=RoutingContext(testing_litellm_router), deployment_id=deployment_id ) # Now with enough requests (all failures), it SHOULD trigger cooldown - should_cooldown = _should_cooldown_deployment(testing_litellm_router, deployment_id, 500, _exception) + should_cooldown = _should_cooldown_deployment( + RoutingContext(testing_litellm_router), deployment_id, 500, _exception + ) assert should_cooldown is True, ( f"Should cooldown when we have {DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS} failed requests (100% failure rate)" ) @@ -2446,7 +2477,7 @@ async def test_cooldown_time_zero_uses_zero_not_default(): mock_add_cooldown.assert_not_called() - cooldown_list = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + cooldown_list = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) assert len(cooldown_list) == 0 healthy_deployments, _ = await router._async_get_healthy_deployments(model="gpt-3.5-turbo", parent_otel_span=None) @@ -2473,7 +2504,7 @@ def test_should_run_cooldown_logic_early_exit_on_zero_cooldown(): ) result = _should_run_cooldown_logic( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="test-deployment-id", exception_status=429, original_exception=litellm.RateLimitError("test error", "openai", "gpt-3.5-turbo"), @@ -2482,7 +2513,7 @@ def test_should_run_cooldown_logic_early_exit_on_zero_cooldown(): assert result is False, "Should not run cooldown logic when time_to_cooldown is 0" result = _should_run_cooldown_logic( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="test-deployment-id", exception_status=429, original_exception=litellm.RateLimitError("test error", "openai", "gpt-3.5-turbo"), @@ -2491,7 +2522,7 @@ def test_should_run_cooldown_logic_early_exit_on_zero_cooldown(): assert result is False, "Should not run cooldown logic when time_to_cooldown is effectively 0" result = _should_run_cooldown_logic( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="test-deployment-id", exception_status=429, original_exception=litellm.RateLimitError("test error", "openai", "gpt-3.5-turbo"), @@ -2500,7 +2531,7 @@ def test_should_run_cooldown_logic_early_exit_on_zero_cooldown(): assert result is True, "Should run cooldown logic when time_to_cooldown is None" result = _should_run_cooldown_logic( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="test-deployment-id", exception_status=429, original_exception=litellm.RateLimitError("test error", "openai", "gpt-3.5-turbo"), @@ -2675,7 +2706,7 @@ async def test_high_traffic_cooldowns_all_healthy_deployments(): raise e print("model_stats: ", model_stats) - cooldown_list = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + cooldown_list = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) assert len(cooldown_list) == 0 @@ -2769,7 +2800,7 @@ async def test_high_traffic_cooldowns_one_bad_deployment(): raise e print("model_stats: ", model_stats) - cooldown_list = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + cooldown_list = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) assert len(cooldown_list) == 1 @@ -2866,7 +2897,7 @@ async def test_high_traffic_cooldowns_one_rate_limited_deployment(): raise e print("model_stats: ", model_stats) - cooldown_list = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + cooldown_list = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) assert len(cooldown_list) == 1 @@ -2932,7 +2963,7 @@ async def test_router_fallbacks_with_cooldowns_and_dynamic_credentials(): model="gpt-3.5-turbo", messages=messages, api_key="my-bad-key-1", mock_response="litellm.RateLimitError" ) await asyncio.sleep(1) - assert await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) == [] + assert await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) == [] response = await router.acompletion( model="gpt-3.5-turbo", messages=messages, api_key="my-good-key-2", mock_response="served with credential 2" @@ -2942,7 +2973,7 @@ async def test_router_fallbacks_with_cooldowns_and_dynamic_credentials(): with pytest.raises(litellm.RateLimitError): await router.acompletion(model="gpt-3.5-turbo", messages=messages, mock_response="litellm.RateLimitError") await asyncio.sleep(1) - cooled_down = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + cooled_down = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) assert len(cooled_down) == 1 and cooled_down[0] in {"123", "456"} @@ -2959,7 +2990,7 @@ async def test_router_fallbacks_with_cooldowns_and_dynamic_credentials(): def test_caller_credential_errors_never_cool_down_the_shared_deployment(single_deployment_router, exception, status): """A per-user credential failure is scoped to one caller's stored token; cooling down the shared deployment would punish every other user on the group.""" - assert _should_run_cooldown_logic(single_deployment_router, "dep-1", status, exception) is False + assert _should_run_cooldown_logic(RoutingContext(single_deployment_router), "dep-1", status, exception) is False @pytest.mark.asyncio @@ -3000,7 +3031,7 @@ async def test_bad_request_error_does_not_cool_down_a_deployment( ) cooldown_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=router, + ctx=RoutingContext(router), parent_otel_span=None, ) @@ -3065,7 +3096,7 @@ async def test_fallback_deployment_stays_up_while_the_primary_is_in_cooldown( served_by: Final = [await _complete_like_the_proxy(router, "primary") for _ in range(20)] cooldown_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=router, + ctx=RoutingContext(router), parent_otel_span=None, ) @@ -3098,7 +3129,7 @@ async def test_fallback_group_stays_up_after_a_bad_key_cools_the_primary_under_t served_by: Final = [await _complete_like_the_proxy(router, "primary") for _ in range(12)] cooldown_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=router, + ctx=RoutingContext(router), parent_otel_span=None, ) @@ -3206,7 +3237,7 @@ def test_has_fallback_group_reads_string_targets_and_skips_unparseable_entries( router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT]) router.fallbacks = list(fallbacks) - assert _has_fallback_group(router, "test-model") is expected + assert _has_fallback_group(RoutingContext(router), "test-model") is expected async def _group_call_outcome(router: Router) -> str: @@ -3239,7 +3270,7 @@ async def test_router_allowed_fails_still_benches_a_deployment_that_has_a_siblin while len(timed_out.calls) < 2: await _group_call_outcome(router) cooldown_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=router, parent_otel_span=None + ctx=RoutingContext(router), parent_otel_span=None ) served: Final = tuple([await _group_call_outcome(router) for _ in range(5)]) @@ -3264,7 +3295,9 @@ def test_per_user_session_upstream_errors_never_cool_down_the_shared_deployment( (litellm.InternalServerError("copilot 500", "github_copilot", "gpt-4o"), 500), ): assert ( - _should_run_cooldown_logic(single_deployment_router, "dep-1", status, exc, request_kwargs=session_kwargs) + _should_run_cooldown_logic( + RoutingContext(single_deployment_router), "dep-1", status, exc, request_kwargs=session_kwargs + ) is False ) @@ -3275,11 +3308,15 @@ def test_shared_mode_upstream_429_still_cools_down_the_deployment(single_deploym exc = litellm.RateLimitError("copilot 429", "github_copilot", "gpt-4o") assert ( _should_run_cooldown_logic( - single_deployment_router, "dep-1", 429, exc, request_kwargs={"model": "github_copilot/gpt-4o"} + RoutingContext(single_deployment_router), + "dep-1", + 429, + exc, + request_kwargs={"model": "github_copilot/gpt-4o"}, ) is True ) - assert _should_run_cooldown_logic(single_deployment_router, "dep-1", 429, exc) is True + assert _should_run_cooldown_logic(RoutingContext(single_deployment_router), "dep-1", 429, exc) is True REFUSAL_MESSAGES: Final = [{"role": "user", "content": "cooldown"}] diff --git a/tests/unit/router_utils/test_deployment_pricing.py b/tests/unit/router_utils/test_deployment_pricing.py index 45cc211a89c..96a9436fcd1 100644 --- a/tests/unit/router_utils/test_deployment_pricing.py +++ b/tests/unit/router_utils/test_deployment_pricing.py @@ -1,6 +1,6 @@ import pytest -from litellm.router_utils.deployment_pricing import deployment_pricing_models, pricing_models +from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models, pricing_models @pytest.mark.parametrize( diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index c5a476c450b..26dbb96c60a 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -20,6 +20,7 @@ from litellm.llms.anthropic.messages.native.fake_stream_iterator import FakeAnth from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.router_api import RouterSettings, RoutingContext from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments, mark_advisor_orchestration_failure from litellm.types.utils import Choices, Message, ModelResponse, StandardLoggingPayload, Usage from litellm.router_utils.fallback_event_handlers import ( @@ -58,24 +59,28 @@ class StreamingWrapper: self._hidden_params = {"additional_headers": {}} -class FakeRouter: +class RouterDouble: + """The router members RoutingContext forwards to during a fallback: the fallback checks, the retry log and + the fallback loop (async_function_with_fallbacks, defined by each double).""" + fallback_access_check = None fallback_budget_check = None - def log_retry(self, kwargs, e): + def settings(self) -> RouterSettings: + return RouterSettings( + fallback_access_check=self.fallback_access_check, fallback_budget_check=self.fallback_budget_check + ) + + def _core_log_retry(self, kwargs, e): return kwargs + +class FakeRouter(RouterDouble): async def async_function_with_fallbacks(self, *args, **kwargs): return StreamingWrapper() -class AlwaysFailRouter: - fallback_access_check = None - fallback_budget_check = None - - def log_retry(self, kwargs, e): - return kwargs - +class AlwaysFailRouter(RouterDouble): async def async_function_with_fallbacks(self, *args, **kwargs): raise RuntimeError("fallback model also failed") @@ -83,7 +88,7 @@ class AlwaysFailRouter: @pytest.mark.asyncio async def test_run_async_fallback_adds_errors_when_opted_in(): response = await run_async_fallback( - litellm_router=FakeRouter(), + litellm_routing_ctx=RoutingContext(FakeRouter()), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("upstream limited request"), @@ -107,7 +112,7 @@ async def test_run_async_fallback_adds_errors_when_opted_in(): @pytest.mark.asyncio async def test_run_async_fallback_omits_errors_without_opt_in(): response = await run_async_fallback( - litellm_router=FakeRouter(), + litellm_routing_ctx=RoutingContext(FakeRouter()), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("upstream limited request"), @@ -124,7 +129,7 @@ async def test_run_async_fallback_omits_errors_without_opt_in(): 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(), + litellm_routing_ctx=RoutingContext(AlwaysFailRouter()), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("original request failed"), @@ -134,16 +139,10 @@ async def test_run_async_fallback_raises_when_all_fallbacks_fail(): ) -class RecordingRouter: - fallback_access_check = None - fallback_budget_check = None - +class RecordingRouter(RouterDouble): 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() @@ -211,7 +210,7 @@ async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_ logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) await run_async_fallback( - litellm_router=HopLoggingItsOwnSuccessRouter(), + litellm_routing_ctx=RoutingContext(HopLoggingItsOwnSuccessRouter()), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("stream dropped before content"), @@ -234,7 +233,7 @@ async def test_run_async_fallback_preserves_superseded_usage_when_hop_fails_befo with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( - litellm_router=AlwaysFailRouter(), + litellm_routing_ctx=RoutingContext(AlwaysFailRouter()), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("stream dropped before content"), @@ -325,7 +324,7 @@ async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_l with pytest.raises(_RateLimitedHop): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("stream dropped before content"), @@ -336,7 +335,7 @@ async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_l metadata={}, ) - cooled_down: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + cooled_down: Final = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) assert cooled_down == (["fallback-dep-1"] if failure_already_logged else []) @@ -348,7 +347,7 @@ async def test_run_async_fallback_cools_down_the_final_deployment_after_retry_fa with pytest.raises(_RateLimitedHop): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("stream dropped before content"), @@ -359,7 +358,7 @@ async def test_run_async_fallback_cools_down_the_final_deployment_after_retry_fa ) assert logging_obj.model_call_details.get("has_logged_async_failure", False) is True - cooled_down: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + cooled_down: Final = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) assert cooled_down == ["fallback-dep-2"] @@ -369,7 +368,7 @@ async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_cal flag has to reach the nested async_function_with_fallbacks call.""" router = RecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("upstream limited request"), @@ -385,7 +384,7 @@ async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_cal async def test_run_async_fallback_does_not_forward_flag_without_opt_in(): router = RecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("upstream limited request"), @@ -399,7 +398,7 @@ async def test_run_async_fallback_does_not_forward_flag_without_opt_in(): @pytest.mark.asyncio async def test_run_async_fallback_skips_original_model_group(): response = await run_async_fallback( - litellm_router=FakeRouter(), + litellm_routing_ctx=RoutingContext(FakeRouter()), fallback_model_group=["primary-model", "fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -410,17 +409,11 @@ async def test_run_async_fallback_skips_original_model_group(): assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 -class AttemptRecordingRouter: - fallback_access_check = None - fallback_budget_check = None - +class AttemptRecordingRouter(RouterDouble): def __init__(self): self.attempted_model_groups = [] self.received_kwargs = None - def log_retry(self, kwargs, e): - return kwargs - async def async_function_with_fallbacks(self, *args, **kwargs): self.attempted_model_groups.append(kwargs.get("model")) self.received_kwargs = kwargs @@ -468,7 +461,7 @@ async def test_run_async_fallback_keeps_uploaded_file_requests_in_their_model_gr with pytest.raises(RuntimeError, match="openai connection error"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["azure-group"], original_model_group="openai-group", original_exception=owning_provider_error, @@ -488,7 +481,7 @@ async def test_run_async_fallback_keeps_fine_tuning_requests_in_their_model_grou with pytest.raises(RuntimeError, match="openai connection error"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["azure-group"], original_model_group="openai-group", original_exception=RuntimeError("openai connection error"), @@ -509,7 +502,7 @@ async def test_run_async_fallback_allows_same_model_group_retry_for_uploaded_fil router = AttemptRecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[{"model": "openai-group", "_target_order": 2}], original_model_group="openai-group", original_exception=RuntimeError("first deployment failed"), @@ -532,7 +525,7 @@ async def test_run_async_fallback_keeps_file_creation_in_its_model_group(): with pytest.raises(RuntimeError, match="azure connection error"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["openai-group"], original_model_group="azure-group", original_exception=RuntimeError("azure connection error"), @@ -550,7 +543,7 @@ async def test_run_async_fallback_allows_same_model_group_retry_for_file_creatio router = AttemptRecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[{"model": "azure-group", "_target_order": 2}], original_model_group="azure-group", original_exception=RuntimeError("first deployment failed"), @@ -568,7 +561,7 @@ async def test_run_async_fallback_still_crosses_model_groups_without_an_uploaded router = AttemptRecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["azure-group"], original_model_group="openai-group", original_exception=RuntimeError("openai connection error"), @@ -612,7 +605,7 @@ async def test_run_async_fallback_keeps_provider_scoped_ids_in_their_model_group with pytest.raises(RuntimeError, match="openai connection error"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["azure-group"], original_model_group="openai-group", original_exception=RuntimeError("openai connection error"), @@ -635,7 +628,7 @@ async def test_run_async_fallback_ignores_stray_resource_ids_on_completion_calls router = AttemptRecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["azure-group"], original_model_group="openai-group", original_exception=RuntimeError("openai connection error"), @@ -654,7 +647,7 @@ async def test_run_async_fallback_allows_same_model_group_retry_for_batch_cancel router = AttemptRecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[{"model": "openai-group", "_target_order": 2}], original_model_group="openai-group", original_exception=RuntimeError("first deployment failed"), @@ -675,7 +668,7 @@ async def test_run_async_fallback_handles_explicitly_none_metadata(): router = AttemptRecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["azure-group"], original_model_group="openai-group", original_exception=RuntimeError("openai connection error"), @@ -699,7 +692,7 @@ async def test_run_async_fallback_records_batch_model_group_outside_provider_met router = AttemptRecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[{"model": "openai-group", "_target_order": 2}], original_model_group="openai-group", original_exception=RuntimeError("first deployment failed"), @@ -724,8 +717,8 @@ class AccessCheckedRouter(AttemptRecordingRouter): fallback_budget_check = None - async def fallback_access_check(self, *, model, request_kwargs, llm_router): - self.access_checks.append((model, request_kwargs["metadata"]["user_api_key"], llm_router is self)) + async def fallback_access_check(self, *, model, request_kwargs, ctx): + self.access_checks.append((model, request_kwargs["metadata"]["user_api_key"], ctx.belongs_to(self))) return model in self.allowed_models @@ -734,7 +727,7 @@ async def test_run_async_fallback_skips_targets_the_access_check_rejects(): router = AccessCheckedRouter(allowed_models=frozenset({"allowed-model"})) await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[ {"model": "secret-model", "messages": [{"role": "user", "content": "hi"}]}, "allowed-model", @@ -754,13 +747,39 @@ async def test_run_async_fallback_skips_targets_the_access_check_rejects(): ] +class LegacyAccessCheckedRouter(AccessCheckedRouter): + async def fallback_access_check(self, *, model, request_kwargs, llm_router): + self.access_checks.append((model, request_kwargs["metadata"]["user_api_key"], llm_router is self)) + return model in self.allowed_models + + +@pytest.mark.asyncio +async def test_run_async_fallback_still_runs_a_deprecated_llm_router_access_check(): + router = LegacyAccessCheckedRouter(allowed_models=frozenset({"allowed-model"})) + + with pytest.warns(DeprecationWarning, match="llm_router"): + await run_async_fallback( + litellm_routing_ctx=RoutingContext(router), + fallback_model_group=["secret-model", "allowed-model"], + original_model_group="primary-model", + original_exception=RuntimeError("primary failed"), + max_fallbacks=3, + fallback_depth=0, + model="primary-model", + metadata={"user_api_key": "hashed"}, + ) + + assert router.attempted_model_groups == ["allowed-model"] + assert router.access_checks == [("secret-model", "hashed", True), ("allowed-model", "hashed", True)] + + @pytest.mark.asyncio async def test_run_async_fallback_raises_original_error_when_no_target_is_authorized(): router = AccessCheckedRouter(allowed_models=frozenset()) with pytest.raises(RuntimeError, match="primary failed"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["secret-model", "other-secret-model"], original_model_group="primary-model", original_exception=RuntimeError("primary failed"), @@ -779,7 +798,7 @@ async def test_run_async_fallback_does_not_consult_access_check_for_same_model_g router = AccessCheckedRouter(allowed_models=frozenset()) await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[{"model": "primary-model", "_target_order": 2}], original_model_group="primary-model", original_exception=RuntimeError("first order level failed"), @@ -793,16 +812,10 @@ async def test_run_async_fallback_does_not_consult_access_check_for_same_model_g assert router.access_checks == [] -class RecordingFailRouter: - fallback_access_check = None - fallback_budget_check = None - +class RecordingFailRouter(RouterDouble): def __init__(self): self.attempted_models = [] - def log_retry(self, kwargs, e): - return kwargs - async def async_function_with_fallbacks(self, *args, **kwargs): self.attempted_models.append(kwargs.get("model")) raise RuntimeError("fallback model also failed") @@ -818,7 +831,7 @@ async def test_run_async_fallback_skips_model_group_already_attempted(): with pytest.raises(RuntimeError, match="original failed"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["already-attempted"], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -836,7 +849,7 @@ async def test_run_async_fallback_attempts_a_repeated_target_once(): with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["fallback-model", "fallback-model", "other-model"], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -855,7 +868,7 @@ async def test_run_async_fallback_forwards_attempted_model_groups_to_nested_call router = RecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -876,7 +889,7 @@ async def test_run_async_fallback_can_target_the_requested_group_when_a_pre_rout router = RecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["requested-model"], original_model_group="requested-model", original_exception=RuntimeError("selected tier failed"), @@ -905,7 +918,7 @@ async def test_run_async_fallback_still_retargets_the_same_group_via_dict_entry( router = RecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[entry], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -926,7 +939,7 @@ async def test_run_async_fallback_skips_a_repeated_dict_target(): with pytest.raises(RuntimeError, match="original failed"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[{"model": "already-attempted"}], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -945,7 +958,7 @@ async def test_run_async_fallback_attempts_a_repeated_dict_target_once(): with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[entry, entry, {"model": "other-model"}], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -965,7 +978,7 @@ async def test_run_async_fallback_keeps_a_request_override_distinct_from_the_bar with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=[{"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]}], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -1028,17 +1041,24 @@ def test_get_fallback_model_group_does_not_mutate_fallbacks(): assert fallbacks == [{"gpt-3.5-turbo": ["claude-3-haiku"]}, "gpt-4o-mini"] +def _cooldown_ctx(*, cooldown_time: float = 60.0, deployment_dict: Mapping[str, object] | None = None) -> MagicMock: + """A RoutingContext double for the fallback cooldown trigger: no deployment, this model info, this cooldown time.""" + ctx: Final = MagicMock() + ctx.registry.return_value.deployment.return_value = None + ctx.registry.return_value.deployment_dict.return_value = deployment_dict + ctx.settings.return_value.cooldown_time = cooldown_time + return ctx + + class TestTriggerCooldownForFailedDeployment: def test_calls_set_cooldown_deployments_with_stamped_deployment_id(self): - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) mock_set_cooldown.assert_called_once() call_kwargs = mock_set_cooldown.call_args[1] @@ -1050,9 +1070,7 @@ class TestTriggerCooldownForFailedDeployment: one without knowing this call's function_name, so a client with permission to set metadata must not be able to get an arbitrary deployment cooled down by forging a deployment_model_name marker.""" - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") kwargs = { @@ -1063,7 +1081,7 @@ class TestTriggerCooldownForFailedDeployment: } with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs=kwargs, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs=kwargs, exception=exc) mock_set_cooldown.assert_not_called() @@ -1071,9 +1089,7 @@ class TestTriggerCooldownForFailedDeployment: """The fallback path must feed the same per-minute failure counter the primary path uses, or repeated fallback failures never accumulate toward the default percent-fail-rate cooldown threshold.""" - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" @@ -1084,47 +1100,39 @@ class TestTriggerCooldownForFailedDeployment: "litellm.router_utils.fallback_event_handlers.increment_deployment_failures_for_current_minute" ) as mock_increment, ): - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) - mock_increment.assert_called_once_with( - litellm_router_instance=mock_router, deployment_id="fallback-deployment" - ) + mock_increment.assert_called_once_with(ctx=ctx, deployment_id="fallback-deployment") mock_set_cooldown.assert_called_once() def test_no_op_when_deployment_id_missing(self): - mock_router = MagicMock() + ctx = _cooldown_ctx() with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: - _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, kwargs={}, exception=RuntimeError("no metadata") - ) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=RuntimeError("no metadata")) mock_set_cooldown.assert_not_called() def test_skipped_for_advisor_orchestration_failure(self): - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" mark_advisor_orchestration_failure(exc) with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) mock_set_cooldown.assert_not_called() def test_uses_deployment_litellm_params_cooldown_time_override(self): - mock_router = MagicMock() - mock_router.cooldown_time = 300.0 - mock_router.get_model_info.return_value = {"litellm_params": {"cooldown_time": 30.0}} + ctx = _cooldown_ctx(cooldown_time=300.0, deployment_dict={"litellm_params": {"cooldown_time": 30.0}}) exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) call_kwargs = mock_set_cooldown.call_args[1] assert call_kwargs["time_to_cooldown"] == 30.0 @@ -1133,24 +1141,20 @@ class TestTriggerCooldownForFailedDeployment: """Precedence must match Router.deployment_callback_on_failure's primary path: deployment config, then the response's Retry-After header, then the router default.""" - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = {"litellm_params": {}} + ctx = _cooldown_ctx(cooldown_time=60.0, deployment_dict={"litellm_params": {}}) exc = RuntimeError("upstream error") exc.failed_deployment_id = "fallback-deployment" exc.litellm_response_headers = httpx.Headers({"retry-after": "45"}) with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) call_kwargs = mock_set_cooldown.call_args[1] assert call_kwargs["time_to_cooldown"] == 45 def test_silently_catches_exceptions(self): - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = RuntimeError("upstream error") exc.failed_deployment_id = "fallback-deployment" @@ -1159,16 +1163,14 @@ class TestTriggerCooldownForFailedDeployment: "litellm.router_utils.fallback_event_handlers.set_cooldown_deployments", side_effect=RuntimeError("cooldown error"), ): - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) def test_skips_request_scoped_404_on_generic_api_call(self): """A generic API call (files/batches/threads/rerank/...) forwards a caller-supplied resource id, so a 404 there means "that id doesn't exist", not "this deployment is unhealthy". Without this guard, a single bad id would 404 every deployment in the fallback chain and cool all of them down from one request.""" - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = litellm.NotFoundError("not found", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" @@ -1180,7 +1182,7 @@ class TestTriggerCooldownForFailedDeployment: ) as mock_increment, ): _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs={"original_generic_function": MagicMock()}, exception=exc, ) @@ -1192,15 +1194,13 @@ class TestTriggerCooldownForFailedDeployment: """The request-scoped-404 guard is scoped to generic API calls only: a 404 on a regular completion fallback (no original_generic_function in kwargs) must still cool down the deployment as before.""" - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = litellm.NotFoundError("not found", "openai", "gpt-4") exc.failed_deployment_id = "fallback-deployment" with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) mock_set_cooldown.assert_called_once() @@ -1213,9 +1213,7 @@ class TestTriggerCooldownForFailedDeployment: The failure logger never stamps end_time for a fallback hop (has_logged_async_failure is already set), so model_call_details still carries the previous hop's end_time, which predates this hop's api_call_start_time. The guard must not trust it.""" - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = litellm.Timeout(message="timeout", model="gpt-4", llm_provider="openai") exc.failed_deployment_id = "fallback-deployment" @@ -1227,7 +1225,7 @@ class TestTriggerCooldownForFailedDeployment: ) as mock_increment, ): _trigger_cooldown_for_failed_deployment( - litellm_router=mock_router, + ctx=ctx, kwargs={"client_side_timeout": True}, exception=exc, model_call_details={ @@ -1271,7 +1269,7 @@ class TestTriggerCooldownForFailedDeployment: started = datetime.now() _trigger_cooldown_for_failed_deployment( - litellm_router=router, + ctx=RoutingContext(router), kwargs={"client_side_timeout": True}, exception=exc, model_call_details={ @@ -1282,9 +1280,7 @@ class TestTriggerCooldownForFailedDeployment: ) assert ( - get_deployment_failures_for_current_minute( - litellm_router_instance=router, deployment_id="fallback-deployment" - ) + get_deployment_failures_for_current_minute(ctx=RoutingContext(router), deployment_id="fallback-deployment") == 1 ) active = router.cooldown_cache.get_active_cooldowns(model_ids=["fallback-deployment"], parent_otel_span=None) @@ -1294,15 +1290,13 @@ class TestTriggerCooldownForFailedDeployment: """The client-side-timeout guard is scoped to caller-supplied timeouts only: a 408 that did not come from x-litellm-timeout (no client_side_timeout in kwargs) must still cool down the deployment as before.""" - mock_router = MagicMock() - mock_router.cooldown_time = 60.0 - mock_router.get_model_info.return_value = None + ctx = _cooldown_ctx(cooldown_time=60.0) exc = litellm.Timeout(message="timeout", model="gpt-4", llm_provider="openai") exc.failed_deployment_id = "fallback-deployment" with patch("litellm.router_utils.fallback_event_handlers.set_cooldown_deployments") as mock_set_cooldown: - _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + _trigger_cooldown_for_failed_deployment(ctx=ctx, kwargs={}, exception=exc) mock_set_cooldown.assert_called_once() @@ -1368,26 +1362,14 @@ async def test_a_fallback_deployment_whose_stream_stalled_stays_out_of_cooldown( ) error.failed_deployment_id = "fallback-deployment" - _trigger_cooldown_for_failed_deployment(litellm_router=router, kwargs={}, exception=error) + _trigger_cooldown_for_failed_deployment(ctx=RoutingContext(router), kwargs={}, exception=error) active: Final = router.cooldown_cache.get_active_cooldowns(model_ids=["fallback-deployment"], parent_otel_span=None) assert tuple(entry[0] for entry in active) == cooled class TestRunAsyncFallbackTriggersCooldown: - class RouterWithLoggingKwarg: - fallback_access_check = None - fallback_budget_check = None - - def __init__(self): - self.cooldown_time = 60.0 - - def log_retry(self, kwargs, e): - return kwargs - - def get_model_info(self, id): - return None - + class RouterWithLoggingKwarg(RouterDouble): async def async_function_with_fallbacks(self, *args, **kwargs): logging_obj: Final = kwargs.get("litellm_logging_obj") if isinstance(logging_obj, Logging): @@ -1410,7 +1392,7 @@ class TestRunAsyncFallbackTriggersCooldown: ) as mock_trigger: with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( - litellm_router=self.RouterWithLoggingKwarg(), + litellm_routing_ctx=RoutingContext(self.RouterWithLoggingKwarg()), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("original request failed"), @@ -1432,7 +1414,7 @@ class TestRunAsyncFallbackTriggersCooldown: ) as mock_trigger: with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( - litellm_router=self.RouterWithLoggingKwarg(), + litellm_routing_ctx=RoutingContext(self.RouterWithLoggingKwarg()), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("original request failed"), @@ -1450,7 +1432,7 @@ class TestRunAsyncFallbackTriggersCooldown: ) as mock_trigger: with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( - litellm_router=self.RouterWithLoggingKwarg(), + litellm_routing_ctx=RoutingContext(self.RouterWithLoggingKwarg()), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("original request failed"), @@ -1472,7 +1454,7 @@ async def test_a_stored_fallback_target_cannot_carry_a_federation_field(): match="server-owned workload identity federation or OAuth token exchange parameter", ): await run_async_fallback( - litellm_router=FakeRouter(), + litellm_routing_ctx=RoutingContext(FakeRouter()), fallback_model_group=[{"model": "anthropic-backup", "anthropic_federation_workspace_id": "wrkspc_other"}], original_model_group="primary-model", original_exception=RuntimeError("upstream limited request"), @@ -1487,7 +1469,7 @@ async def test_a_stored_fallback_target_cannot_carry_an_openai_federation_field( naming a token file would pick which workload assertion is exchanged for the bearer.""" with pytest.raises(ValueError, match="openai_identity_token_file"): await run_async_fallback( - litellm_router=FakeRouter(), + litellm_routing_ctx=RoutingContext(FakeRouter()), fallback_model_group=[ {"model": "openai-backup", "openai_identity_token_file": "/var/run/secrets/tokens/other"} ], @@ -1526,7 +1508,7 @@ async def test_a_stored_fallback_target_cannot_make_the_proxy_read_an_oidc_api_k with pytest.raises(ValueError, match="oidc/"): await run_async_fallback( - litellm_router=FakeRouter(), + litellm_routing_ctx=RoutingContext(FakeRouter()), fallback_model_group=[target], original_model_group="primary-model", original_exception=RuntimeError("upstream limited request"), @@ -1541,7 +1523,7 @@ async def test_the_refusal_is_not_swallowed_as_a_fallback_error(): that target's failure and the run would quietly continue to the next one.""" with pytest.raises(ValueError, match="anthropic_issuer_signing_key_ref"): await run_async_fallback( - litellm_router=FakeRouter(), + litellm_routing_ctx=RoutingContext(FakeRouter()), fallback_model_group=[ {"model": "anthropic-backup", "anthropic_issuer_signing_key_ref": "os.environ/ADMIN_KEY"}, "a-perfectly-fine-model", @@ -1560,7 +1542,7 @@ async def test_run_async_fallback_stamps_fallback_info_into_metadata(): router = RecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["fallback-model"], original_model_group="primary-model", original_exception=RuntimeError("original failed"), @@ -1581,7 +1563,7 @@ async def test_run_async_fallback_preserves_original_model_group_on_nested_fallb router = RecordingRouter() await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), fallback_model_group=["second-fallback"], original_model_group="first-fallback", original_exception=RuntimeError("first fallback failed"), @@ -2592,7 +2574,7 @@ async def test_run_async_fallback(function_name): request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}] result = await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), original_function=original_function, num_retries=1, fallback_model_group=fallback_model_group, @@ -2702,7 +2684,7 @@ async def test_failed_fallbacks_raise_most_recent_exception(function_name): with pytest.raises(litellm.exceptions.RateLimitError): await run_async_fallback( - litellm_router=router, + litellm_routing_ctx=RoutingContext(router), original_function=original_function, num_retries=1, fallback_model_group=fallback_model_group, @@ -2854,7 +2836,7 @@ async def test_multiple_fallbacks(function_name, respx_mock: respx.MockRouter, m request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}] result = await run_async_fallback( - litellm_router=router_2, + litellm_routing_ctx=RoutingContext(router_2), original_function=original_function, num_retries=1, fallback_model_group=fallback_model_group, diff --git a/tests/unit/router_utils/test_get_retry_from_policy.py b/tests/unit/router_utils/test_get_retry_from_policy.py index 4ff869849dd..8847329d82c 100644 --- a/tests/unit/router_utils/test_get_retry_from_policy.py +++ b/tests/unit/router_utils/test_get_retry_from_policy.py @@ -5,7 +5,7 @@ import httpx import pytest import litellm -from litellm.router_utils.get_retry_from_policy import get_num_retries_from_retry_policy +from litellm.litellm_core_utils.get_retry_from_policy import get_num_retries_from_retry_policy from litellm.types.router import RetryPolicy _EXCEPTION_FOR_FIELD: Final = MappingProxyType( diff --git a/tests/unit/router_utils/test_handle_error.py b/tests/unit/router_utils/test_handle_error.py index 4b0690b2bd4..151fb0bfc68 100644 --- a/tests/unit/router_utils/test_handle_error.py +++ b/tests/unit/router_utils/test_handle_error.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest import litellm +from litellm.router_api import RouterSettings, RoutingObservers from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -13,9 +14,10 @@ async def test_send_llm_exception_alert_success(): """ Test that the function sends an alert when the router.slack_alerting_logger is set. """ - # Create a mock LitellmRouter instance - mock_router = MagicMock() - mock_router.slack_alerting_logger = AsyncMock() + # A routing context whose observers carry a slack_alerting_logger + slack_alerting_logger = AsyncMock() + ctx = MagicMock() + ctx.observers.return_value = RoutingObservers(slack_alerting_logger=slack_alerting_logger) # Create a mock exception mock_exception = Exception("Test exception") @@ -32,10 +34,10 @@ async def test_send_llm_exception_alert_success(): # Call the function from litellm.router_utils.handle_error import send_llm_exception_alert - await send_llm_exception_alert(mock_router, request_kwargs, error_traceback, mock_exception) + await send_llm_exception_alert(ctx, request_kwargs, error_traceback, mock_exception) # Assert that the slack_alerting_logger's send_alert method was called - mock_router.slack_alerting_logger.send_alert.assert_called_once() + slack_alerting_logger.send_alert.assert_called_once() @pytest.mark.asyncio @@ -43,9 +45,9 @@ async def test_send_llm_exception_alert_no_logger(): """ Test that the function does error out when no slack_alerting_logger is set """ - # Create a mock LitellmRouter instance without a slack_alerting_logger - mock_router = MagicMock() - mock_router.slack_alerting_logger = None + # A routing context whose observers have no slack_alerting_logger + ctx = MagicMock() + ctx.observers.return_value = RoutingObservers(slack_alerting_logger=None) # Create a mock exception mock_exception = Exception("Test exception") @@ -62,7 +64,7 @@ async def test_send_llm_exception_alert_no_logger(): # Call the function from litellm.router_utils.handle_error import send_llm_exception_alert - await send_llm_exception_alert(mock_router, request_kwargs, error_traceback, mock_exception) + await send_llm_exception_alert(ctx, request_kwargs, error_traceback, mock_exception) @pytest.mark.asyncio @@ -70,9 +72,10 @@ async def test_send_llm_exception_alert_when_proxy_server_request_in_kwargs(): """ Test that the function does not send an alert when the request kwargs contains a proxy_server_request key. """ - # Create a mock LitellmRouter instance with a slack_alerting_logger - mock_router = MagicMock() - mock_router.slack_alerting_logger = AsyncMock() + # A routing context whose observers carry a slack_alerting_logger + slack_alerting_logger = AsyncMock() + ctx = MagicMock() + ctx.observers.return_value = RoutingObservers(slack_alerting_logger=slack_alerting_logger) # Create a mock exception mock_exception = Exception("Test exception") @@ -90,11 +93,11 @@ async def test_send_llm_exception_alert_when_proxy_server_request_in_kwargs(): # Call the function from litellm.router_utils.handle_error import send_llm_exception_alert - await send_llm_exception_alert(mock_router, request_kwargs, error_traceback, mock_exception) + await send_llm_exception_alert(ctx, request_kwargs, error_traceback, mock_exception) # Assert that no exception was raised and the function completed successfully - mock_router.slack_alerting_logger.send_alert.assert_not_called() + slack_alerting_logger.send_alert.assert_not_called() @pytest.mark.asyncio @@ -108,11 +111,11 @@ async def test_async_raise_no_deployment_exception(): from litellm.router_utils.handle_error import async_raise_no_deployment_exception from litellm.types.router import RouterRateLimitError - # Create a mock LitellmRouter instance - mock_router = MagicMock() - mock_router.get_model_ids.return_value = ["deployment-1", "deployment-2"] - mock_router.cooldown_cache.get_min_cooldown.return_value = 30.0 - mock_router.enable_pre_call_checks = True + # A routing context: the group's deployment ids, their cooldown store and the router settings + ctx = MagicMock() + ctx.registry.return_value.names.return_value = ["deployment-1", "deployment-2"] + ctx.health.return_value.cooldowns.get_min_cooldown.return_value = 30.0 + ctx.settings.return_value = RouterSettings(enable_pre_call_checks=True) # Mock the _async_get_cooldown_deployments_with_debug_info function # It should return a list of tuples where each tuple contains (model_id, debug_info) @@ -128,7 +131,7 @@ async def test_async_raise_no_deployment_exception(): ): # Call the function result = await async_raise_no_deployment_exception( - litellm_router_instance=mock_router, + ctx=ctx, model="gpt-5-mini", parent_otel_span=None, ) @@ -150,8 +153,8 @@ async def test_async_raise_no_deployment_exception(): assert isinstance(item, str), f"Expected string ID, got {type(item)}: {item}" # Verify mock calls - mock_router.get_model_ids.assert_called_once_with(model_name="gpt-5-mini") - mock_router.cooldown_cache.get_min_cooldown.assert_called_once_with( + ctx.registry().names.assert_called_once_with("gpt-5-mini", kind="ids") + ctx.health().cooldowns.get_min_cooldown.assert_called_once_with( model_ids=["deployment-1", "deployment-2"], parent_otel_span=None ) @@ -166,11 +169,11 @@ async def test_async_raise_no_deployment_exception_empty_cooldown_list(): from litellm.router_utils.handle_error import async_raise_no_deployment_exception from litellm.types.router import RouterRateLimitError - # Create a mock LitellmRouter instance - mock_router = MagicMock() - mock_router.get_model_ids.return_value = ["deployment-1", "deployment-2"] - mock_router.cooldown_cache.get_min_cooldown.return_value = 15.0 - mock_router.enable_pre_call_checks = False + # A routing context: the group's deployment ids, their cooldown store and the router settings + ctx = MagicMock() + ctx.registry.return_value.names.return_value = ["deployment-1", "deployment-2"] + ctx.health.return_value.cooldowns.get_min_cooldown.return_value = 15.0 + ctx.settings.return_value = RouterSettings(enable_pre_call_checks=False) # Mock empty cooldown list mock_cooldown_list: List = [] @@ -181,7 +184,7 @@ async def test_async_raise_no_deployment_exception_empty_cooldown_list(): ): # Call the function result = await async_raise_no_deployment_exception( - litellm_router_instance=mock_router, + ctx=ctx, model="claude-3-sonnet", parent_otel_span=None, ) @@ -211,11 +214,11 @@ async def test_async_raise_no_deployment_exception_none_cooldown_list(): from litellm.router_utils.handle_error import async_raise_no_deployment_exception from litellm.types.router import RouterRateLimitError - # Create a mock LitellmRouter instance - mock_router = MagicMock() - mock_router.get_model_ids.return_value = [] - mock_router.cooldown_cache.get_min_cooldown.return_value = 45.0 - mock_router.enable_pre_call_checks = True + # A routing context: the group's deployment ids, their cooldown store and the router settings + ctx = MagicMock() + ctx.registry.return_value.names.return_value = [] + ctx.health.return_value.cooldowns.get_min_cooldown.return_value = 45.0 + ctx.settings.return_value = RouterSettings(enable_pre_call_checks=True) # Mock None cooldown list (though this shouldn't happen in practice) mock_cooldown_list = None @@ -226,7 +229,7 @@ async def test_async_raise_no_deployment_exception_none_cooldown_list(): ): # After the defensive fix, this should handle None gracefully and return empty list result = await async_raise_no_deployment_exception( - litellm_router_instance=mock_router, + ctx=ctx, model="gpt-5.5", parent_otel_span=None, ) diff --git a/tests/unit/router_utils/test_health_check_allowed_fails_integration.py b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py index d200b304c4e..74b3a8825b3 100644 --- a/tests/unit/router_utils/test_health_check_allowed_fails_integration.py +++ b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py @@ -6,13 +6,14 @@ should increment the same counters used by allowed_fails_policy, using the actual exception type from the health check error. """ -from unittest.mock import patch +from unittest.mock import ANY, patch import pytest import litellm from litellm.proxy.health_check import run_with_timeout from litellm.router import Router +from litellm.router_api import RoutingContext from litellm.types.router import AllowedFailsPolicy @@ -160,14 +161,14 @@ class TestHealthCheckCooldownIntegration: # First call: should not cooldown (1 <= 3) result = should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="deploy-1", original_exception=timeout_exc, ) assert result is False # Check counter was incremented - current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails") + current_fails = router.health().store.get_cache(key="deployment:deploy-1:allowed_fails") assert current_fails == 1 def test_health_check_failure_triggers_cooldown_at_threshold(self): @@ -186,7 +187,7 @@ class TestHealthCheckCooldownIntegration: # Fails 1 and 2: should not cooldown for _ in range(2): result = should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="deploy-1", original_exception=auth_exc, ) @@ -194,7 +195,7 @@ class TestHealthCheckCooldownIntegration: # Fail 3: should trigger cooldown (3 > 2) result = should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="deploy-1", original_exception=auth_exc, ) @@ -219,7 +220,7 @@ class TestHealthCheckCooldownIntegration: # Fail 1: should not cooldown (1 <= 1) result = should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="deploy-1", original_exception=generic_exc, ) @@ -227,7 +228,7 @@ class TestHealthCheckCooldownIntegration: # Fail 2: should trigger cooldown (2 > 1) result = should_cooldown_based_on_allowed_fails_policy( - litellm_router_instance=router, + ctx=RoutingContext(router), deployment="deploy-1", original_exception=generic_exc, ) @@ -266,11 +267,11 @@ class TestHealthCheckCooldownIntegration: timeout_exc = litellm.Timeout(message="Health check timeout", model="gpt-4", llm_provider="openai") result = set_cooldown_deployments( - litellm_router_instance=router, + ctx=RoutingContext(router), original_exception=timeout_exc, exception_status=500, deployment="deploy-1", - time_to_cooldown=router.cooldown_time, + time_to_cooldown=router.settings().cooldown_time, ) assert result is False @@ -306,12 +307,13 @@ class TestWriteHealthStateIntegration: exceptions_by_model_id={"deploy-1": timeout_exc}, ) mock_cooldown.assert_called_once_with( - litellm_router_instance=router, + ANY, original_exception=timeout_exc, exception_status=408, # Timeout has status_code 408 deployment="deploy-1", - time_to_cooldown=router.cooldown_time, + time_to_cooldown=router.settings().cooldown_time, ) + assert mock_cooldown.call_args.args[0].belongs_to(router) def test_unhealthy_endpoint_without_exception_skips_cooldown(self): """Unhealthy endpoints without an exception key should not trigger cooldown.""" @@ -364,10 +366,8 @@ class TestWriteHealthStateIntegration: unhealthy_endpoints=unhealthy_endpoints, exceptions_by_model_id={"deploy-1": rate_exc}, ) - mock_increment.assert_called_once_with( - litellm_router_instance=router, - deployment_id="deploy-1", - ) + mock_increment.assert_called_once_with(ANY, "deploy-1") + assert mock_increment.call_args.args[0].belongs_to(router) class TestHealthCheckFilterBypassWithPolicy: diff --git a/tests/unit/router_utils/test_pattern_match_deployments.py b/tests/unit/router_utils/test_pattern_match_deployments.py index 0bb33064c68..b5145ce9c71 100644 --- a/tests/unit/router_utils/test_pattern_match_deployments.py +++ b/tests/unit/router_utils/test_pattern_match_deployments.py @@ -1,4 +1,4 @@ -"""Behavior pins for ``litellm/router_utils/pattern_match_deployments.py``.""" +"""Behavior pins for ``litellm/litellm_core_utils/pattern_match_deployments.py``.""" from __future__ import annotations @@ -13,7 +13,7 @@ import litellm from litellm import Router from litellm.litellm_core_utils import get_llm_provider_logic from litellm.router import Deployment, LiteLLM_Params -from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils from litellm.types.router import ModelInfo import json from unittest.mock import MagicMock @@ -406,7 +406,7 @@ def test_sorted_patterns(): """ Tests that the pattern specificity is calculated correctly """ - from litellm.router_utils.pattern_match_deployments import PatternUtils + from litellm.litellm_core_utils.pattern_match_deployments import PatternUtils sorted_patterns = PatternUtils.sorted_patterns( { @@ -418,7 +418,7 @@ def test_sorted_patterns(): def test_calculate_pattern_specificity(): - from litellm.router_utils.pattern_match_deployments import PatternUtils + from litellm.litellm_core_utils.pattern_match_deployments import PatternUtils assert PatternUtils.calculate_pattern_specificity("llmengine/*") == (11, 1) assert PatternUtils.calculate_pattern_specificity("*") == (1, 1) diff --git a/tests/unit/router_utils/test_ptu_shares.py b/tests/unit/router_utils/test_ptu_shares.py index bc929f6455a..6a843f5f821 100644 --- a/tests/unit/router_utils/test_ptu_shares.py +++ b/tests/unit/router_utils/test_ptu_shares.py @@ -6,7 +6,7 @@ from typing import Final import pytest from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY -from litellm.router_utils.ptu_shares import ( +from litellm.litellm_core_utils.ptu_shares import ( filter_ptu_shared_deployments, model_group_deployments, model_group_ptu_capacity, diff --git a/tests/unit/router_utils/test_reasoning_effort_capability.py b/tests/unit/router_utils/test_reasoning_effort_capability.py index ce4cb84addf..c0ff141ffab 100644 --- a/tests/unit/router_utils/test_reasoning_effort_capability.py +++ b/tests/unit/router_utils/test_reasoning_effort_capability.py @@ -7,7 +7,7 @@ import pytest import litellm from litellm._logging import JsonFormatter, verbose_logger -from litellm.router_utils.reasoning_effort_capability import ( +from litellm.litellm_core_utils.reasoning_effort_capability import ( RequestedReasoningEffort, deployment_is_catalog_mapped, intersect_supported_reasoning_efforts, diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 8011ca37cfb..6ec72868b6e 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -10,7 +10,7 @@ from litellm import Router from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy._types import UserAPIKeyAuth -from litellm.router_utils.common_utils import ( +from litellm.litellm_core_utils.routing_common_utils import ( _deployment_supports_web_search, add_model_file_id_mappings, filter_team_based_models, diff --git a/tests/unit/rust_bridge/test_runtime.py b/tests/unit/rust_bridge/test_runtime.py index c0c416a62c9..75cdfde087b 100644 --- a/tests/unit/rust_bridge/test_runtime.py +++ b/tests/unit/rust_bridge/test_runtime.py @@ -6,8 +6,8 @@ from typing import Final, Protocol import pytest from litellm.exceptions import APIError +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import bindings, configuration, runtime from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule diff --git a/tests/unit/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py index deba7fe9587..c8fef95cf98 100644 --- a/tests/unit/test_anthropic_beta_headers_filtering.py +++ b/tests/unit/test_anthropic_beta_headers_filtering.py @@ -85,7 +85,8 @@ class TestAnthropicBetaHeadersFiltering: filtered = filter_and_transform_beta_headers( beta_headers=all_headers, provider=provider ) - assert ("compact-2026-09-04" in filtered) is (provider == "anthropic") + assert ("compact-2026-09-04" in filtered) is (provider in ("anthropic", "azure_ai", "vertex_ai")) + assert ("mid-conversation-tool-changes-2026-07-01" in filtered) is (provider != "bedrock_converse") for header in unsupported_headers: assert ( @@ -406,15 +407,49 @@ class TestAnthropicBetaHeadersFiltering: test_case["expected"] in filtered ), f"Header '{test_case['input']}' should be mapped to '{test_case['expected']}' for {test_case['provider']}, but got: {filtered}" - def test_filter_and_transform_beta_headers_vertex_ai_keeps_compact(self): - """Vertex AI supports compact context edits, so the compact beta header - must be forwarded instead of stripped (it was previously mapped to null, - which broke compact_20260112 context edits over /v1/messages).""" - filtered = filter_and_transform_beta_headers( - beta_headers=["compact-2026-01-12"], provider="vertex_ai" + @pytest.mark.parametrize("beta", ["compact-2026-09-04", "inline-tools-2026-09-15"]) + def test_vertex_ai_header_only_beta_is_forwarded_as_the_header_and_dropped_from_the_body( + self, beta: str + ) -> None: + headers, body = update_request_with_filtered_beta( + headers={"anthropic-beta": f"{beta},context-1m-2025-08-07"}, + request_data={"anthropic_beta": [beta, "context-1m-2025-08-07"]}, + provider="vertex_ai", ) - assert filtered == ["compact-2026-01-12"] + assert set(headers["anthropic-beta"].split(",")) == {beta, "context-1m-2025-08-07"} + assert body["anthropic_beta"] == ["context-1m-2025-08-07"] + + def test_header_only_betas_stay_out_of_the_provider_columns_older_proxies_read(self) -> None: + """Proxies released before header-only columns existed fetch this map from main, read only the + provider column, and copy its betas into the body field too, where Vertex answers 400. So a + header-only beta may sit in its provider column only if older proxies already forwarded it.""" + header_only_columns: Final = { + column.removesuffix("_header_only"): betas + for column, betas in self.config.items() + if column.endswith("_header_only") + } + forwarded_by_older_proxies: Final = { + provider: {beta for beta in betas if self.config[provider].get(beta) is not None} + for provider, betas in header_only_columns.items() + } + + assert "vertex_ai" in header_only_columns + assert {provider: betas for provider, betas in forwarded_by_older_proxies.items() if betas} == { + "vertex_ai": {"inline-tools-2026-09-15"} + } + + def test_filter_and_transform_beta_headers_vertex_ai_keeps_compact(self): + """Vertex AI serves both compaction betas (threshold compaction behind compact-2026-01-12 + and on-demand compaction behind compact-2026-09-04, Google Cloud: beta on + https://platform.claude.com/docs/en/build-with-claude/compaction-on-demand, read 2026-10-05), + so the Vertex filter forwards both instead of stripping them: a dropped header makes + Vertex answer 400 "compaction: Extra inputs are not permitted".""" + filtered = filter_and_transform_beta_headers( + beta_headers=["compact-2026-01-12", "compact-2026-09-04"], provider="vertex_ai" + ) + + assert filtered == ["compact-2026-01-12", "compact-2026-09-04"] @pytest.mark.parametrize("provider", ["bedrock_converse", "bedrock"]) def test_fine_grained_tool_streaming_forwarded_for_bedrock(self, provider): diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index f10baccd1cb..c536522e736 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -37,7 +37,7 @@ from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import CacheCreationTokenDetails, CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject -from litellm.utils import _invalidate_model_cost_lowercase_map, supports_prompt_caching +from litellm.utils import _invalidate_model_cost_lowercase_map, cached_get_model_info_helper, supports_prompt_caching import os import threading from unittest.mock import MagicMock, patch @@ -4598,23 +4598,40 @@ def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate( assert completion_cost_value == pytest.approx(200 * 15e-6 / 2) -def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate(): +@pytest.mark.parametrize("input_rate", (3e-6, 0.0)) +@pytest.mark.parametrize( + ("cache_pricing", "expected_write_rate"), + [ + pytest.param({}, None, id="missing"), + pytest.param({"cache_creation_input_token_cost": None}, None, id="null"), + pytest.param({"cache_creation_input_token_cost": 0.0}, 0.0, id="free"), + pytest.param({"cache_creation_input_token_cost": 3.75e-6}, 3.75e-6, id="paid"), + ], +) +def test_batch_cost_calculator_preserves_cache_write_prices_and_defaults_only_when_unset( + cache_pricing: ModelInfo, + expected_write_rate: float | None, + input_rate: float, +) -> None: from litellm.cost_calculator import batch_cost_calculator - model_info: ModelInfo = { + model_info: Final[ModelInfo] = { "supported_openai_params": [], - "input_cost_per_token": 3e-6, + "input_cost_per_token": input_rate, "output_cost_per_token": 15e-6, "cache_read_input_token_cost": 3e-7, + **cache_pricing, } - prompt_cost, _ = batch_cost_calculator( + prompt_cost, completion_cost_value = batch_cost_calculator( usage=_batch_cache_usage(), - model="claude-sonnet-4-5-20250929", + model="batch-pricing-fixture", custom_llm_provider="anthropic", model_info=model_info, ) - assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2) + write_rate: Final = input_rate if expected_write_rate is None else expected_write_rate + assert prompt_cost == pytest.approx((1000 * input_rate + 8000 * 3e-7 + 2000 * write_rate) / 2) + assert completion_cost_value == pytest.approx(200 * 15e-6 / 2) def test_completion_cost_bills_interactions_api_response(): @@ -7817,20 +7834,42 @@ def test_vertex_ai_mistral_predict_cost(usage): assert predictive_cost > 0 +@pytest.mark.usefixtures("_local_model_cost_map") +@pytest.mark.parametrize("call_type", ("speech", "aspeech")) +@pytest.mark.parametrize("model", ("speech-fixture", "openai/speech-fixture", "azure/speech-fixture")) @pytest.mark.parametrize( - "model", ["openai/tts-1", "azure/tts-1", "openai/gpt-4o-mini-tts"] + ("pricing", "expected_cost"), + ( + pytest.param({}, None, id="missing"), + pytest.param({"input_cost_per_character": None, "input_cost_per_token": None}, None, id="null"), + pytest.param({"input_cost_per_token": 0.0}, 0.0, id="free-tokens"), + pytest.param({"input_cost_per_character": 0.0, "input_cost_per_token": 2e-6}, 0.0, id="free-characters"), + pytest.param({"input_cost_per_token": 2e-6}, 100 * 2e-6, id="missing-character-rate"), + pytest.param({"input_cost_per_character": None, "input_cost_per_token": 2e-6}, 100 * 2e-6, id="null-character-rate"), + pytest.param({"input_cost_per_character": 3e-6, "input_cost_per_token": 2e-6}, 1000 * 3e-6, id="paid-characters"), + ), ) -def test_completion_cost_tts(model): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - cost = completion_cost( - model=model, - prompt="the quick brown fox jumped over the lazy dogs", - call_type="speech", +def test_registered_speech_prices_distinguish_unset_from_zero( + call_type: str, model: str, pricing: ModelInfo, expected_cost: float | None +) -> None: + provider: Final = "azure" if model.startswith("azure/") else "openai" + litellm.register_model( + {model: {"litellm_provider": provider, "mode": "audio_speech", **pricing}}, + persist_across_reloads=False, ) + compatibility_info: Final = cached_get_model_info_helper(model, provider) + assert compatibility_info["input_cost_per_token"] == (pricing.get("input_cost_per_token") or 0.0) + litellm.get_model_info(model) + response: Final = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=0, total_tokens=100)) - assert cost > 0 + if expected_cost is None: + with pytest.raises(ValueError, match="does not have 'input_cost_per_character' or 'input_cost_per_token'"): + completion_cost(completion_response=response, model=model, prompt="a" * 1000, call_type=call_type) + litellm.register_model({model: {"input_cost_per_token": 0.0}}, persist_across_reloads=False) + + assert completion_cost( + completion_response=response, model=model, prompt="a" * 1000, call_type=call_type + ) == pytest.approx(expected_cost or 0.0) def test_completion_cost_anthropic(): diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index b0bf1cb879e..9f2d2a9bbc5 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -5900,6 +5900,28 @@ def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): assert response._hidden_params["response_cost"] == pytest.approx(0.5) +@pytest.mark.parametrize("reported_cost", [float("nan"), float("inf")]) +def test_stream_chunk_builder_bills_the_token_price_when_the_provider_cost_is_not_finite(reported_cost: float): + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=reported_cost) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() + ) + + assert response is not None + prompt_cost, completion_cost = litellm.cost_per_token(model="gpt-4o", usage_object=response.usage) + token_price: Final = prompt_cost + completion_cost + assert token_price > 0 + assert getattr(response.usage, "cost", None) == pytest.approx(token_price) + assert response._hidden_params["response_cost"] == pytest.approx(token_price) + + def test_completion_cost_on_a_stamped_stream_still_bills_the_reported_charge(): usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) @@ -10240,9 +10262,10 @@ def test_a_non_numeric_thinking_budget_on_adaptive_claude_is_rejected_with_a_400 @pytest.mark.parametrize("drop_globally", [False, True]) @pytest.mark.parametrize("model_info", [None, _OPUS_DECLARING_LOW_AND_HIGH]) @pytest.mark.parametrize("body_sent", [_anthropic_chat_body_sent, _bedrock_converse_request_fields_sent]) -def test_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude( +def test_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude_with_one_warning( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, body_sent: Callable[..., Mapping[str, object]], model_info: Mapping[str, object] | None, drop_globally: bool, @@ -10250,13 +10273,40 @@ def test_drop_params_drops_a_non_numeric_thinking_budget_on_adaptive_claude( monkeypatch.setattr(litellm, "drop_params", drop_globally) request_drop: Final = {} if drop_globally else {"drop_params": True} - dropped: Final = body_sent( - respx_mock, model_info, thinking={"type": "enabled", "budget_tokens": "lots"}, **request_drop - ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + dropped: Final = body_sent( + respx_mock, model_info, thinking={"type": "enabled", "budget_tokens": "lots"}, **request_drop + ) + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] adaptive_alone: Final = body_sent(respx_mock, model_info, thinking={"type": "adaptive"}, **request_drop) assert "output_config" not in adaptive_alone assert dropped == adaptive_alone + assert len(drop_warnings) == 1 + assert all( + fragment in drop_warnings[0] + for fragment in ("thinking.budget_tokens='lots'", "claude-opus-5-5", "(drop_params=True)") + ) + + +@pytest.mark.parametrize("model_info", [None, _KIMI_DEPLOYMENT]) +def test_drop_params_drops_a_non_numeric_thinking_budget_on_hosted_vllm_with_one_warning( + respx_mock: respx.MockRouter, caplog: pytest.LogCaptureFixture, model_info: Mapping[str, object] | None +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + sent: Final = _chat_body_sent( + respx_mock, + "hosted_vllm/my-kimi", + None, + model_info, + thinking={"type": "enabled", "budget_tokens": "lots"}, + drop_params=True, + ) + + assert not {"reasoning_effort", "thinking"} & sent.keys() + drop_warnings: Final = [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] + assert len(drop_warnings) == 1 + assert all(fragment in drop_warnings[0] for fragment in ("thinking.budget_tokens='lots'", "my-kimi")) @pytest.mark.parametrize("model", ["azure_ai/my-kimi", "hosted_vllm/my-kimi"]) diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 8ca678a94ad..e4a2beb3a86 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -12,9 +12,9 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter import litellm from litellm.litellm_core_utils.fallback_generalizations import match_capability_generalizations +from litellm.litellm_core_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts from litellm.llms.openai.chat.gpt_reasoning_transformation import is_gpt_reasoning_series_name from litellm.llms.openai_like.json_loader import JSONProviderRegistry -from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts from litellm.types.utils import SearchFreeRequests REPO_ROOT = Path(__file__).parents[2] diff --git a/tests/unit/test_muse_spark_1_3_model_metadata.py b/tests/unit/test_muse_spark_1_3_model_metadata.py index 049fb988fd5..468294ba91c 100644 --- a/tests/unit/test_muse_spark_1_3_model_metadata.py +++ b/tests/unit/test_muse_spark_1_3_model_metadata.py @@ -6,7 +6,7 @@ import pytest import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import StandardBuiltInToolCostTracking -from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts +from litellm.litellm_core_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts MUSE_SPARK_STANDARD = "meta/muse-spark-1.3" MUSE_SPARK_CONTRIBUTOR = "meta/muse-spark-1.3-contributor" diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index f2195512355..3c0aab8a4b4 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -1966,8 +1966,8 @@ ALIAS_CASES: Final = ( "get_tags_from_request_kwargs", False, ), - ("litellm.router_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner", True), - ("litellm.router_utils.add_retry_fallback_headers", "", "_HiddenParamsHost", "HiddenParamsHost", False), + ("litellm.litellm_core_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner", True), + ("litellm.litellm_core_utils.add_retry_fallback_headers", "", "_HiddenParamsHost", "HiddenParamsHost", False), ( "litellm.litellm_core_utils.batch_utils", "", @@ -1975,7 +1975,7 @@ ALIAS_CASES: Final = ( "get_router_metadata_variable_name", False, ), - ("litellm.router_utils.common_utils", "", "_is_proxy_admin_request", "is_proxy_admin_request", False), + ("litellm.litellm_core_utils.routing_common_utils", "", "_is_proxy_admin_request", "is_proxy_admin_request", False), ( "litellm.router_utils.cooldown_callbacks", "", @@ -2007,7 +2007,7 @@ ALIAS_CASES: Final = ( False, ), ( - "litellm.router_utils.pattern_match_deployments", + "litellm.litellm_core_utils.pattern_match_deployments", "PatternMatchRouter", "_pattern_to_regex", "pattern_to_regex", @@ -2102,7 +2102,7 @@ PROPERTY_CASES: Final = ( ), ("litellm.router", "Router", "_routing_groups", "routing_groups"), ("litellm.router_strategy.adaptive_router.adaptive_router", "AdaptiveRouter", "_state_loaded", "state_loaded"), - ("litellm.router_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner"), + ("litellm.litellm_core_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner"), ) CLASS_PROPERTY_CASES: Final = () diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index e5ff4531f84..819c6b9d4f5 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -4,14 +4,19 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). import asyncio from collections.abc import Awaitable, Callable, Mapping -from datetime import datetime +from datetime import datetime, timezone from typing import Final +from unittest.mock import AsyncMock, MagicMock import pytest +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeRedis as FakeAsyncRedis import litellm from litellm import Router from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -22,6 +27,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( OTPM_RESERVED_KEY, _explicit_max_tokens, _reservation_value, + async_admit_io_token_turn, async_io_token_pre_call_check, async_io_token_reconcile_success, build_io_token_rate_limit_headers, @@ -35,6 +41,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterCacheEnum from litellm.types.utils import ModelResponse, Usage from litellm.utils import get_utc_datetime @@ -115,7 +122,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 10, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -173,7 +180,8 @@ class TestModelRateLimitingCheckIOTokens: ], "max_tokens": 10, "metadata": {}, - } + }, + metadata_channel="metadata", ) with pytest.raises(litellm.RateLimitError) as exc_info: @@ -203,7 +211,8 @@ class TestModelRateLimitingCheckIOTokens: "messages": [{"role": "user", "content": "hi"}], "max_tokens": max_tokens, "metadata": {}, - } + }, + metadata_channel="metadata", ) async def _attempt(): @@ -255,7 +264,8 @@ class TestModelRateLimitingCheckIOTokens: { "max_tokens": 5, "metadata": {}, - } + }, + metadata_channel="metadata", ) async def _attempt(): @@ -320,7 +330,7 @@ class TestModelRateLimitingCheckIOTokens: } request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -363,15 +373,15 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "opus", } request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) # Second zero-estimate call still fits within the itpm=2 limit. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}, metadata_channel="metadata") await check.async_pre_call_check(deployment) # A third exceeds the limit and is rejected. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}, metadata_channel="metadata") with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) @@ -392,7 +402,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 0, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(zero_output_kwargs) + set_io_token_rate_limit_request_kwargs(zero_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] @@ -404,7 +414,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(normal_output_kwargs) + set_io_token_rate_limit_request_kwargs(normal_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] @@ -432,7 +442,7 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "no-max-tokens", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) @@ -462,7 +472,7 @@ class TestModelRateLimitingCheckIOTokens: async def _attempt() -> dict[str, object] | None: request_kwargs: Final = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") try: await check.async_pre_call_check(deployment) except litellm.RateLimitError: @@ -493,7 +503,7 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "roomy", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) @@ -507,7 +517,7 @@ class TestModelRateLimitingCheckIOTokens: async def _reserved(model_info: dict[str, str]) -> int: request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") deployment: Final = { "litellm_params": {"model": "azure/litellm-io-custom-deployment", "otpm": otpm_limit}, "model_info": model_info, @@ -536,7 +546,7 @@ class TestModelRateLimitingCheckIOTokens: } async def _pre_call(request_kwargs: dict[str, object]) -> None: - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) else: @@ -563,12 +573,14 @@ class TestModelRateLimitingCheckIOTokens: "model_name": "opus", } request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == 100 - set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": "hi"}], "metadata": {}}) + set_io_token_rate_limit_request_kwargs( + {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}, metadata_channel="metadata" + ) with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) @@ -585,7 +597,7 @@ class TestModelRateLimitingCheckIOTokens: with_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 300, "metadata": {}} for request_kwargs in (without_cap, with_cap): - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") if use_sync_check: check.pre_call_check(deployment) else: @@ -602,7 +614,7 @@ class TestModelRateLimitingCheckIOTokens: async def reserved(deployment_cap: Mapping[str, int], request_shape: Mapping[str, object]) -> int: request_kwargs: Final = {**request_shape, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await check.async_pre_call_check( { "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, **deployment_cap}, @@ -631,7 +643,7 @@ class TestModelRateLimitingCheckIOTokens: async def reserved(deployment_cap: Mapping[str, int]) -> int: request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") deployment: Final = { "litellm_params": {"model": UNMAPPED_MODEL, "otpm": otpm_limit, **deployment_cap}, "model_info": {"id": f"io-deployment-cap-above-default-{deployment_cap}"}, @@ -667,7 +679,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 10, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") check.pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") @@ -803,7 +815,7 @@ class TestModelRateLimitingCheckIOTokens: # The retry proceeds to stash deployment B's own reservation on the # same dict; it starts clean, unaffected by A's cleared sentinels. - set_io_token_rate_limit_request_kwargs(shared_kwargs) + set_io_token_rate_limit_request_kwargs(shared_kwargs, metadata_channel="metadata") itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b @@ -819,7 +831,7 @@ class TestModelRateLimitingCheckIOTokens: "litellm_metadata": {OTPM_RESERVED_KEY: 7}, "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, } - set_io_token_rate_limit_request_kwargs(forged) + set_io_token_rate_limit_request_kwargs(forged, metadata_channel="metadata") stored = get_io_token_rate_limit_request_kwargs() assert ITPM_RESERVED_KEY not in stored["metadata"] @@ -845,7 +857,7 @@ class TestModelRateLimitingCheckIOTokens: }, } # The router sanitizes the request kwargs before the call runs. - set_io_token_rate_limit_request_kwargs(kwargs) + set_io_token_rate_limit_request_kwargs(kwargs, metadata_channel="metadata") response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), @@ -880,7 +892,8 @@ class TestModelRateLimitingCheckIOTokens: "messages": [{"role": "user", "content": "hello world"}], "max_tokens": 5, "metadata": {}, - } + }, + metadata_channel="metadata", ) with pytest.raises(RuntimeError): @@ -945,7 +958,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") result = await async_io_token_pre_call_check(dual_cache, deployment) @@ -1147,7 +1160,7 @@ class TestModelRateLimitingCheckIOTokens: "max_tokens": 5, "metadata": {}, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): with pytest.raises(litellm.RateLimitError): @@ -1495,15 +1508,15 @@ class TestContextSlotRetention: "messages": [{"role": "user", "content": "x" * 1000}], "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, } - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None assert ITPM_RESERVED_KEY not in kwargs["metadata"] assert ITPM_CACHE_KEY not in kwargs["metadata"] - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is kwargs - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None @pytest.mark.asyncio @@ -1516,7 +1529,7 @@ class TestContextSlotRetention: } ] ) - set_io_token_rate_limit_request_kwargs(None) + set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata") kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} deployment = router.get_deployment_by_model_group_name("plain") assert deployment is not None @@ -1534,7 +1547,7 @@ class TestContextSlotRetention: ], optional_pre_call_checks=["enforce_model_rate_limits"], ) - set_io_token_rate_limit_request_kwargs(None) + set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata") kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} deployment = router.get_deployment_by_model_group_name("limited") assert deployment is not None @@ -1551,7 +1564,7 @@ async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_tok "max_tokens": 5, "metadata": metadata, } - set_io_token_rate_limit_request_kwargs(request_kwargs) + set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") await async_io_token_pre_call_check(DualCache(), deployment) return metadata[ITPM_RESERVED_KEY] @@ -1571,3 +1584,278 @@ async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_tok ) assert named == served + + +def _reply_using(prompt_tokens: int, completion_tokens: int) -> ModelResponse: + return ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + + +def _proxy_success_kwargs(reservation: Mapping[str, object]) -> dict[str, object]: + return { + "litellm_params": {"metadata": {**reservation, "user_api_key_hash": "hashed-key"}}, + "standard_logging_object": { + "hidden_params": {"litellm_model_name": "openai/gpt-5-mini"}, + "metadata": {"user_api_key_hash": "hashed-key"}, + }, + } + + +async def _deliver_success( + check: ModelRateLimitingCheck, delivery: str, kwargs: dict[str, object], reply: object +) -> None: + if delivery == "async": + await check.async_log_success_event(kwargs, reply, None, None) + return + check.log_success_event(kwargs, reply, None, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deliveries", + [("async",), ("sync",), ("async", "async"), ("async", "sync"), ("sync", "async"), ("sync", "sync")], +) +async def test_success_settles_a_reservation_held_only_in_litellm_params_metadata_exactly_once( + deliveries: tuple[str, ...], +) -> None: + dual_cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:itpm:12-34" + otpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:otpm:12-34" + await dual_cache.async_increment_cache(key=itpm_key, value=40, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=4096, ttl=60) + kwargs: Final = _proxy_success_kwargs( + {ITPM_RESERVED_KEY: 40, OTPM_RESERVED_KEY: 4096, ITPM_CACHE_KEY: itpm_key, OTPM_CACHE_KEY: otpm_key} + ) + + for delivery in deliveries: + await _deliver_success(check, delivery, kwargs, _reply_using(prompt_tokens=7, completion_tokens=3)) + + assert await dual_cache.async_get_cache(key=itpm_key) == 7 + assert await dual_cache.async_get_cache(key=otpm_key) == 3 + + +_TURN_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) +_ACTUAL_TOKENS: Final = {"itpm": 5, "otpm": 2} + + +def _turn_deployment(kind: str) -> dict[str, object]: + return {"litellm_params": {"model": "openai/gpt-5-mini", kind: 100}, "model_info": {"id": f"ws-turn-{kind}"}} + + +def _turn_key(kind: str, minute: str = "12-34") -> str: + return f"global_router:ws-turn-{kind}:openai/gpt-5-mini:{kind}:{minute}" + + +@pytest.fixture +def settled_on_the_turn_s_minute(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: _TURN_AT + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used", [100, 101, 250]) +async def test_turn_admission_refuses_once_the_minutes_tokens_reach_the_limit(kind: str, used: int) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage={used}"): + await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used", [0, 1, 99]) +async def test_an_admitted_turn_reserves_nothing_and_counts_its_actual_tokens(kind: str, used: int) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + actual: Final = _ACTUAL_TOKENS + + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + counter_key: Final = reservation[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] + assert isinstance(counter_key, str) + assert await dual_cache.async_get_cache(key=counter_key) == used + await async_io_token_reconcile_success( + dual_cache, _proxy_success_kwargs(reservation), _reply_using(actual["itpm"], actual["otpm"]) + ) + + assert await dual_cache.async_get_cache(key=counter_key) == used + actual[kind] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment", + [ + {"litellm_params": {"model": "openai/gpt-5-mini", "tpm": 1}, "model_info": {"id": "ws-turn-unlimited"}}, + {"litellm_params": {"model": "openai/gpt-5-mini", "otpm": 1}}, + ], +) +async def test_a_turn_without_io_token_counters_is_admitted_with_nothing_to_settle( + deployment: dict[str, object], +) -> None: + assert await async_admit_io_token_turn(DualCache(), deployment, _TURN_AT) == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("used_by_others", [0, 1, 99]) +async def test_a_socket_open_reserves_nothing_and_its_first_turn_counts_its_actual_tokens_on_its_own_minute( + kind: str, used_by_others: int +) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key(kind), value=used_by_others, ttl=60) + socket_open: Final[dict[str, object]] = {"metadata": {}} + set_io_token_rate_limit_request_kwargs(socket_open, opens_responses_socket=True, metadata_channel="metadata") + + assert await async_io_token_pre_call_check(dual_cache, _turn_deployment(kind)) == _turn_deployment(kind) + assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + + opened_with: Final = socket_open["metadata"] + assert isinstance(opened_with, dict) + await async_io_token_reconcile_success( + dual_cache, _proxy_success_kwargs(opened_with), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]) + ) + + assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + _ACTUAL_TOKENS[kind] + + +class _RedisCacheOnFakeServer(RedisCache): + def __init__(self, server: FakeServer) -> None: + super().__init__(host="io-token-turn-redis", connection_pool=FakeRedis(server=server).connection_pool) + self._fake_async_client: Final = FakeAsyncRedis(server=server) + + def init_async_client(self) -> FakeAsyncRedis: + return self._fake_async_client + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_turn_is_refused_on_this_instance_s_count_when_reading_the_shared_counter_fails( + kind: str, +) -> None: + server: Final = FakeServer() + dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnFakeServer(server)) + server.connected = False + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): + await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("delivery", ["async", "sync"]) +async def test_a_turn_finishing_in_the_next_minute_counts_its_actual_tokens_on_that_minute( + kind: str, delivery: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", + lambda: datetime(2026, 1, 1, 12, 35, 10, tzinfo=timezone.utc), + ) + dual_cache: Final = DualCache() + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + await _deliver_success( + ModelRateLimitingCheck(dual_cache=dual_cache), + delivery, + _proxy_success_kwargs(reservation), + _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]), + ) + + assert [await dual_cache.async_get_cache(key=_turn_key(kind, minute)) for minute in ("12-34", "12-35")] == [ + None, + _ACTUAL_TOKENS[kind], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_turn_is_admitted_once_the_shared_counter_drops_below_the_limit_this_instance_last_read( + kind: str, +) -> None: + redis_cache: Final = _RedisCacheOnFakeServer(FakeServer()) + dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) + await redis_cache.async_increment(key=_turn_key(kind), value=99) + + admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "responses_reply", + [ + pytest.param( + ResponsesAPIResponse( + id="resp_cached", + created_at=0, + output=[], + usage=ResponseAPIUsage( + input_tokens=50, + output_tokens=5, + total_tokens=55, + input_tokens_details=InputTokensDetails(cached_tokens=30), + ), + ), + id="usage_object", + ), + pytest.param( + { + "usage": { + "input_tokens": 50, + "output_tokens": 5, + "total_tokens": 55, + "input_tokens_details": {"cached_tokens": 30}, + } + }, + id="usage_dict", + ), + ], +) +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +async def test_a_turn_counts_only_its_uncached_responses_input_toward_itpm(responses_reply: object) -> None: + dual_cache: Final = DualCache() + await dual_cache.async_increment_cache(key=_turn_key("itpm"), value=10, ttl=60) + + reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment("itpm"), _TURN_AT) + await async_io_token_reconcile_success(dual_cache, _proxy_success_kwargs(reservation), responses_reply) + + assert await dual_cache.async_get_cache(key=_turn_key("itpm")) == 10 + 50 - 30 + + +def _shared_counter_answering(answer: int | Exception) -> DualCache: + redis_cache: Final = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(side_effect=[answer]) + return DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_later_turn_is_refused_when_another_instance_filled_the_shared_minute(kind: str) -> None: + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): + await async_admit_io_token_turn(_shared_counter_answering(100), _turn_deployment(kind), _TURN_AT) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_a_later_turn_is_admitted_on_this_instance_s_count_when_the_shared_counter_is_unreachable( + kind: str, +) -> None: + dual_cache: Final = _shared_counter_answering(RedisCircuitBreakerOpenError("circuit open")) + await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=99) + + admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + + assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d4e7cf74650..9039f0fb93b 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -82,6 +82,7 @@ from litellm.router_utils.prompt_caching_cache import PromptCachingCache from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.scheduler import FlowItem from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest +from litellm.router_api import RoutingContext as RouterContext from litellm.types.router import ( CustomRoutingStrategyBase, Deployment, @@ -378,7 +379,7 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): ] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True - assert encrypted_content_callbacks[0].router is router + assert encrypted_content_callbacks[0].ctx is router._core_ctx finally: if router is not None: router.discard() @@ -1506,7 +1507,7 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( assert key == expected_key assert ( - get_deployment_successes_for_current_minute(litellm_router_instance=router, deployment_id="batch-dep") + get_deployment_successes_for_current_minute(ctx=router._core_ctx, deployment_id="batch-dep") == expected_successes ) @@ -2200,7 +2201,11 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): ) def inject_alias_into_kwargs( - deployment: dict, kwargs: dict, function_name: str | None = None, splits_stream_timeout: bool = False + deployment: dict, + kwargs: dict, + function_name: str | None = None, + splits_stream_timeout: bool = False, + opens_responses_socket: bool = False, ) -> None: # Simulate the alias leaking into kwargs (as happens when # _ageneric_api_call_with_fallbacks sets kwargs["model"] = alias before @@ -5136,7 +5141,10 @@ async def test_set_response_headers_is_the_only_complexity_header_source_for_pro @pytest.mark.asyncio -async def test_acompletion_streaming_iterator_adopts_fallback_response_headers(): +@pytest.mark.parametrize("close_fails", [False, True]) +async def test_acompletion_streaming_iterator_adopts_fallback_response_headers( + close_fails: bool, +) -> None: """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must describe the deployment that served the stream, with no value left over from the attempt that failed.""" @@ -5183,6 +5191,7 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() class FallbackStream: def __init__(self): + self.close_attempted: bool = False self._response_headers = {"x-request-id": "req-FALLBACK"} self._hidden_params = { "model_id": "fallback-deployment", @@ -5200,6 +5209,11 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() except StopIteration: raise StopAsyncIteration from None + async def aclose(self) -> None: + self.close_attempted = True + if close_fails: + raise RuntimeError("fallback transport failed during close") + fallback_stream = FallbackStream() with patch.object( router, @@ -5216,12 +5230,14 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() # the very first chunk the fallback produces must already be published under # the fallback's identity: the proxy commits response headers once that chunk # is buffered, so adopting any later is adopting too late - await result.__anext__() + first_chunk: Final = await result.__anext__() + assert first_chunk.choices[0].delta.content == "OK" assert result._response_headers == {"x-request-id": "req-FALLBACK"} assert result._hidden_params["model_id"] == "fallback-deployment" async for _ in result: pass + assert fallback_stream.close_attempted assert result._response_headers == {"x-request-id": "req-FALLBACK"} assert result._hidden_params["model_id"] == "fallback-deployment" assert result._hidden_params["api_base"] == "https://fallback.example" @@ -10088,7 +10104,7 @@ def test_pattern_match_router_remove_deployment(): remove_deployment must drop only the deployment with the given model id and delete patterns whose deployment list becomes empty. """ - from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter pattern_router = PatternMatchRouter() pattern_router.add_pattern( @@ -12913,7 +12929,7 @@ class TestBackgroundResponseCostPollCooldown: is False ) assert self._cooled_down_ids(router) == [] - value = get_deployment_failures_for_current_minute(litellm_router_instance=router, deployment_id="dep-1") + value = get_deployment_failures_for_current_minute(ctx=router._core_ctx, deployment_id="dep-1") assert not value @pytest.mark.asyncio @@ -13003,7 +13019,7 @@ class TestCallerTimeoutCooldown: get_deployment_failures_for_current_minute, ) - return get_deployment_failures_for_current_minute(litellm_router_instance=router, deployment_id="dep-1") + return get_deployment_failures_for_current_minute(ctx=router._core_ctx, deployment_id="dep-1") def _cooled_down_ids(self, router): active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) @@ -17485,7 +17501,7 @@ def test_model_group_info_reasoning_efforts_are_unknown_for_an_operator_declared operator wrote, so a mode they set themselves reads back exactly like one the map supplied. Only a mode the map supplied marks the deployment as known. An off-map deployment carrying an operator mode remains unknown and must keep the whole group's level support unknown.""" - from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + from litellm.litellm_core_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts mapped_model = "openai/gpt-5.6-sol" expected = resolve_supported_reasoning_efforts( @@ -20875,7 +20891,7 @@ class _AllowlistFallbackAccessCheck: self.allowed_models = allowed_models self.checked_models = [] - async def __call__(self, *, model, request_kwargs, llm_router): + async def __call__(self, *, model, request_kwargs, ctx): self.checked_models.append(model) return model in self.allowed_models @@ -20908,7 +20924,7 @@ def _router_with_failing_primary(fallback_access_check) -> Router: class _DenyFallbackBudgetCheck: async def __call__( - self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router + self, *, model: str, request_kwargs: Mapping[str, object], ctx: RouterContext ) -> bool: return False @@ -22005,7 +22021,7 @@ def test_router_retry_skip_stamp_feeds_deployment_ids_to_skip_on_retry( ], ) def test_router_as_retry_skipped_deployment_ids_keeps_only_a_tuple_of_strings(value, expected): - from litellm.router_utils.common_utils import as_retry_skipped_deployment_ids + from litellm.litellm_core_utils.routing_common_utils import as_retry_skipped_deployment_ids assert as_retry_skipped_deployment_ids(value) == expected @@ -22839,7 +22855,7 @@ async def test_router_max_parallel_requests_overflow_is_429_without_cooldown_or_ assert all("capped-deployment" in r.message and "max_parallel_requests=1" in r.message for r in rejected) assert route.call_count == 1 assert sibling_route.call_count == 0 - assert await _async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) == [] + assert await _async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) == [] @pytest.mark.asyncio @@ -22958,7 +22974,7 @@ async def test_router_max_parallel_requests_overflow_takes_the_ordinary_429_fall assert len(results) == 3 assert primary.call_count == 1 assert fallback.call_count == 2 - assert await _async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) == [] + assert await _async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) == [] @pytest.mark.asyncio @@ -28238,11 +28254,11 @@ async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): if sync_mode: cooldown_deployments = get_cooldown_deployments( - litellm_router_instance=router, parent_otel_span=None + ctx=router._core_ctx, parent_otel_span=None ) else: cooldown_deployments = await async_get_cooldown_deployments( - litellm_router_instance=router, parent_otel_span=None + ctx=router._core_ctx, parent_otel_span=None ) assert len(cooldown_deployments) > 0 @@ -29899,7 +29915,7 @@ def test_router_does_not_cool_down_api_connection_errors() -> None: assert ( _is_cooldown_required( - litellm_router_instance=router, + ctx=router._core_ctx, model_id="connection-deployment", exception_status=500, exception_str="litellm.APIConnectionError: connection refused", @@ -31163,7 +31179,7 @@ def _deployment_ids(deployments: list[DeploymentTypedDict]) -> tuple[str, ...]: id="response_polling_provider_lookup", ), pytest.param( - lambda router: _models_in(router, "model-3"), + lambda router: _models_in(router._core_ctx, "model-3"), (Baseline("openai/gpt-3", "deployment-3"),), id="savings_baseline_candidates", ), @@ -31259,7 +31275,7 @@ def test_complexity_router_adaptive_pool_during_a_concurrent_delete_prices_the_l router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) complexity: Final = ComplexityRouter( model_name="hybrid", - litellm_router_instance=router, + ctx=router._core_ctx, complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["model-3"]}}, ) diff --git a/tests/unit/test_router/test_router_endpoints.py b/tests/unit/test_router/test_router_endpoints.py index 815b32b33f7..0c57fb662ac 100644 --- a/tests/unit/test_router/test_router_endpoints.py +++ b/tests/unit/test_router/test_router_endpoints.py @@ -1256,6 +1256,6 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): ] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True - assert encrypted_content_callbacks[0].router is router + assert encrypted_content_callbacks[0].ctx is router._core_ctx finally: router.discard() diff --git a/tests/unit/test_router/test_router_entry_points.py b/tests/unit/test_router/test_router_entry_points.py index 2af7db930c6..366e16a8e91 100644 --- a/tests/unit/test_router/test_router_entry_points.py +++ b/tests/unit/test_router/test_router_entry_points.py @@ -827,3 +827,77 @@ def test_attribute_errors_about_other_objects_pass_through_unchanged(router): other = AttributeError("'dict' object has no attribute 'x'", name="x", obj={}) assert _public_attribute_error(router, "_core_x", "x", other) is other + + +# Router-family helpers whose Router parameter became a RoutingContext keep the old spellings. + + +def test_renamed_router_keyword_still_works_and_warns(router): + from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + get_deployment_failures_for_current_minute, + increment_deployment_failures_for_current_minute, + ) + + deployment_id = router.registry().names(kind="ids")[0] + label = ( + "litellm.router_utils.router_callbacks.track_deployment_metrics." + "increment_deployment_failures_for_current_minute" + ) + before = DEPRECATED_ROUTER_MEMBER_USES[label] + + with pytest.warns(DeprecationWarning, match=r"litellm_router_instance= is deprecated; pass ctx="): + increment_deployment_failures_for_current_minute(litellm_router_instance=router, deployment_id=deployment_id) + + assert DEPRECATED_ROUTER_MEMBER_USES[label] == before + 1 + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert get_deployment_failures_for_current_minute(ctx=router._core_ctx, deployment_id=deployment_id) == 1 + + +def test_router_passed_where_the_context_goes_is_converted_with_a_warning(router): + from litellm.router_utils.cooldown_handlers import get_cooldown_deployments + + with pytest.warns(DeprecationWarning, match="passing a Router is deprecated"): + assert get_cooldown_deployments(router, None) == [] + with pytest.warns(DeprecationWarning, match="a Router as ctx= is deprecated"): + assert get_cooldown_deployments(ctx=router, parent_otel_span=None) == [] + + +def test_renamed_router_keyword_alongside_the_new_one_is_rejected(router): + from litellm.router_utils.cooldown_handlers import get_cooldown_deployments + + with pytest.raises(TypeError, match="got both litellm_router_instance= and ctx="): + get_cooldown_deployments(ctx=router._core_ctx, litellm_router_instance=router, parent_otel_span=None) + + +def test_renamed_router_param_maps_to_a_differently_named_context(router): + from litellm.router_api import RoutingContext + from litellm.router_deprecated import renamed_router_param + + @renamed_router_param("litellm_router", new="litellm_routing_ctx") + def helper(*, litellm_routing_ctx: RoutingContext) -> RoutingContext: + return litellm_routing_ctx + + with pytest.warns(DeprecationWarning, match="litellm_router= is deprecated; pass litellm_routing_ctx="): + assert helper(litellm_router=router) is router._core_ctx + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert helper(litellm_routing_ctx=router._core_ctx) is router._core_ctx + + +def test_renamed_router_attribute_reads_and_writes_the_context(router): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + check = EncryptedContentAffinityCheck(ctx=router._core_ctx) + with pytest.warns(DeprecationWarning, match=r"EncryptedContentAffinityCheck\.router: reading this attribute"): + assert check.router is router._core_ctx + + other = _router() + with pytest.warns(DeprecationWarning, match="setting this attribute is deprecated; set .ctx"): + check.router = other + assert check.ctx is other._core_ctx + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + other.retire() diff --git a/tests/unit/test_router/test_router_fallbacks.py b/tests/unit/test_router/test_router_fallbacks.py index 88d89696297..bccd6290f44 100644 --- a/tests/unit/test_router/test_router_fallbacks.py +++ b/tests/unit/test_router/test_router_fallbacks.py @@ -1006,3 +1006,204 @@ async def test_unknown_default_fallback_raises_after_primary_failure( assert route.call_count == 1 assert "missing" in str(exc_info.value) + + +def _context_error_response_stream( + *, event_type: str = "error", preceding_event: dict[str, object] | None = None, completed: bool = False +) -> httpx.Response: + response: Final = { + "id": "resp_backup" if completed else "resp_primary", + "object": "response", + "created_at": 1700000000, + "model": "backup" if completed else "primary", + "status": "completed" if completed else "in_progress", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + } + error: Final = { + "type": "invalid_request_error", + "code": "context_length_exceeded", + "message": "Your input exceeds the context window of this model. Please adjust your input and try again.", + } + terminal: Final = ( + {"type": "response.completed", "response": response, "sequence_number": 2} + if completed + else {"type": event_type, "error": error, "sequence_number": 2} + if event_type == "error" + else {"type": event_type, "response": {**response, "status": "failed", "error": error}, "sequence_number": 2} + ) + events: Final = ( + {"type": "response.created", "response": response, "sequence_number": 0}, + *((preceding_event,) if preceding_event is not None else ()), + terminal, + ) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_type", ("error", "response.failed")) +@pytest.mark.parametrize("context_fallback", (False, True)) +async def test_responses_context_error_uses_the_configured_fallback_before_output( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, event_type: str, context_fallback: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[{"primary": ["missing"]}] if context_fallback else None) + route: Final = respx_mock.post("https://fallback-migration.local/v1/responses").mock( + side_effect=( + _context_error_response_stream(event_type=event_type), + _context_error_response_stream(completed=True), + ) + ) + stream: Final = await router.aresponses( + model="primary", + input="oversized prompt", + stream=True, + context_window_fallbacks=[{"primary": ["backup"]}] if context_fallback else None, + ) + events: Final = [event async for event in stream] + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "backup") + assert tuple(event.type for event in events) == ("response.created", "response.completed") + assert events[0].response.id == events[-1].response.id + assert events[-1].response.model == "backup" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("barrier", ("text", "tool", "disabled", "unconfigured")) +async def test_responses_context_error_does_not_restart_committed_or_disabled_streams( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, barrier: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[] if barrier == "unconfigured" else None) + preceding: Final = { + "text": { + "type": "response.output_text.delta", + "delta": "already visible", + "item_id": "msg_primary", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + }, + "tool": { + "type": "response.output_item.added", + "item": { + "type": "function_call", "id": "fc_primary", "call_id": "call_primary", + "name": "lookup", "arguments": "", + }, + "output_index": 0, + "sequence_number": 1, + }, + }.get(barrier) + route: Final = respx_mock.post("https://fallback-migration.local/v1/responses").mock( + return_value=_context_error_response_stream(preceding_event=preceding) + ) + stream: Final = await router.aresponses( + model="primary", input="oversized prompt", stream=True, disable_fallbacks=barrier == "disabled" + ) + with pytest.raises(litellm.ContextWindowExceededError): + async for _ in stream: + pass + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", (False, True)) +@pytest.mark.parametrize("emission", ("role", "text", "tool")) +async def test_chat_context_error_falls_back_after_role_only_but_not_after_output( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool, emission: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[{"primary": ["missing"]}]) + delta: Final = ( + {"role": "assistant", "tool_calls": [{ + "index": 0, "id": "call_primary", "type": "function", "function": {"name": "lookup", "arguments": ""} + }]} + if emission == "tool" + else {"role": "assistant", "content": "visible" if emission == "text" else ""} + ) + first_chunk: Final = { + "id": "chatcmpl-primary", "object": "chat.completion.chunk", "created": 1700000000, "model": "primary", + "choices": [{"index": 0, "delta": delta}], + } + error: Final = {"error": { + "type": "invalid_request_error", "code": "context_length_exceeded", + "message": "This model's maximum context length is 128000 tokens. Your messages resulted in 130000 tokens.", + }} + backup: Final = { + "id": "chatcmpl-backup", "object": "chat.completion.chunk", "created": 1700000000, "model": "backup", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "backup answer"}, "finish_reason": "stop"}], + } + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=( + httpx.Response(200, content=f"data: {json.dumps(first_chunk)}\n\ndata: {json.dumps(error)}\n\n".encode()), + httpx.Response(200, content=f"data: {json.dumps(backup)}\n\ndata: [DONE]\n\n".encode()), + ) + ) + request: Final = { + "model": "primary", "messages": [{"role": "user", "content": "oversized prompt"}], "stream": True, + "context_window_fallbacks": [{"primary": ["backup"]}], + } + stream: Final = router.completion(**request) if sync_mode else await router.acompletion(**request) + if emission != "role": + with pytest.raises(litellm.ContextWindowExceededError): + tuple(stream) if sync_mode else tuple([chunk async for chunk in stream]) + assert route.call_count == 1 + return + chunks: Final = tuple(stream) if sync_mode else tuple([chunk async for chunk in stream]) + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "backup") + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "backup answer" + + +@pytest.mark.parametrize( + "wrapped,has_output,generated_content,disqualifying,recoverable", + ( + (False, False, "", False, True), + (False, True, "", False, False), + (True, False, "", False, True), + (True, True, "", False, False), + (True, False, "already visible", False, False), + (True, False, "", True, False), + ), +) +def test_context_stream_fallback_obeys_all_output_receipts( + wrapped: bool, has_output: bool, generated_content: str, disqualifying: bool, recoverable: bool +) -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.router import _context_window_stream_fallback_error + + original: Final = litellm.ContextWindowExceededError(message="too long", model="primary", llm_provider="openai") + error: Final = ( + MidStreamFallbackError( + message="too long", model="primary", llm_provider="openai", original_exception=original, + is_pre_first_chunk=True, generated_content=generated_content, emitted_disqualifying_content=disqualifying, + ) + if wrapped + else original + ) + if not recoverable: + with pytest.raises(litellm.ContextWindowExceededError) as raised: + _context_window_stream_fallback_error(error, "primary", "openai", has_output) + assert raised.value is original + return + recovered: Final = _context_window_stream_fallback_error(error, "primary", "openai", has_output) + assert recovered.original_exception is original + assert recovered.is_pre_first_chunk is True + if wrapped: + assert recovered is error + + +def test_context_stream_fallback_preserves_other_errors_for_existing_continuation_policy() -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.router import _context_window_stream_fallback_error + + error: Final = MidStreamFallbackError( + message="provider unavailable", model="primary", llm_provider="openai", generated_content="already visible" + ) + assert _context_window_stream_fallback_error(error, "primary", "openai", True) is error diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 031f0ac533a..c3814371c26 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -4,13 +4,15 @@ import gc import json import os import weakref -from collections.abc import Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from datetime import datetime, timezone from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm from litellm import Router @@ -538,14 +540,14 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( assert await router.cache.async_get_cache(key=tpm_key) == 17 -async def _kwargs_routed_to(router: Router, model: str, function_name: str) -> dict[str, object]: +async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]: routed: Final[dict[str, object]] = {} async def call(**kwargs: object) -> None: routed.update(kwargs) call.__name__ = function_name - await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call) + await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call, **request) return routed @@ -743,6 +745,363 @@ async def test_responses_websocket_turns_do_not_reserve_deployment_otpm(): await admit_turn() +def _io_limited_router(kind: str, **litellm_params: object) -> Router: + return Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": { + "model": "openai/ws-io-unpriced", + "api_key": "sk-fake", + kind: 10_000, + **litellm_params, + }, + "model_info": {"id": "ws-io-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + +def _minutes_counter_key(kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> str: + return f"global_router:ws-io-dep:{model}:{kind}:{minute.strftime('%H-%M')}" + + +async def _count_minutes_tokens( + router: Router, kind: str, tokens: int, minute: datetime, model: str = "openai/ws-io-unpriced" +) -> None: + await router.cache.async_increment_cache(key=_minutes_counter_key(kind, minute, model), value=tokens, ttl=120) + + +async def _settle_turn( + router: Router, reservation: Mapping[str, object], input_tokens: int, output_tokens: int +) -> None: + from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + rate_limits: Final = next( + check for check in router.optional_callbacks or () if isinstance(check, ModelRateLimitingCheck) + ) + await rate_limits.async_log_success_event( + {"litellm_params": get_litellm_params(litellm_metadata=dict(reservation))}, + ResponsesWebSocketTokenUsageProcessor.create_logging_object( + [ + { + "type": "response.completed", + "response": {"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}}, + } + ] + ), + None, + None, + ) + + +_LONG_OUTPUT_MODEL: Final = "openai/gpt-5.4-mini" + + +def _long_output_router() -> tuple[Router, Mapping[str, int]]: + max_output: Final = litellm.get_model_info(_LONG_OUTPUT_MODEL)["max_output_tokens"] + assert isinstance(max_output, int) + limits: Final = {"itpm": 10_000, "otpm": max_output // 2} + return _io_limited_router("itpm", model=_LONG_OUTPUT_MODEL, otpm=limits["otpm"]), limits + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_opens_without_counting_tokens_until_the_minute_reaches_deployment_io_limit( + kind: str, router_minute_pinned: datetime +): + router, limits = _long_output_router() + + await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) + for counted in ("itpm", "otpm") + ] == [None, None] + await _count_minutes_tokens(router, kind, limits[kind], router_minute_pinned, _LONG_OUTPUT_MODEL) + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit={limits[kind]}, current usage="): + await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + + +@pytest.mark.asyncio +async def test_responses_websocket_first_turn_counts_its_actual_tokens_toward_deployment_itpm_and_otpm( + router_minute_pinned: datetime, +): + router, _ = _long_output_router() + socket_metadata: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()))[ + "litellm_metadata" + ] + assert isinstance(socket_metadata, dict) + + await _settle_turn(router, socket_metadata, input_tokens=20, output_tokens=15) + + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) + for counted in ("itpm", "otpm") + ] == [20, 15] + + +_ADMITTED_LATE_IN_THE_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 50, tzinfo=timezone.utc) +_FINISHED_IN_THE_NEXT_MINUTE: Final = datetime(2026, 1, 1, 12, 1, 15, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +@pytest.mark.parametrize("turn", ["first", "later"]) +async def test_responses_websocket_turn_finishing_in_the_next_minute_counts_its_tokens_toward_that_minute( + kind: str, turn: str, monkeypatch: pytest.MonkeyPatch +): + _pin_router_clock(monkeypatch, _ADMITTED_LATE_IN_THE_MINUTE) + router: Final = _io_limited_router(kind) + socket: Final = await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) + admit_turn: Final = socket["turn_admission"] + assert callable(admit_turn) + reservation: Final = socket["litellm_metadata"] if turn == "first" else await admit_turn() + assert isinstance(reservation, Mapping) + _pin_router_clock(monkeypatch, _FINISHED_IN_THE_NEXT_MINUTE) + + await _settle_turn(router, reservation, input_tokens=10_000, output_tokens=10_000) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage=10000"): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [{}, {"websocket": True}], ids=["plain", "websocket_key"]) +async def test_messages_call_reserves_max_tokens_against_deployment_otpm_whatever_websocket_key_its_body_carries( + body: Mapping[str, object], router_minute_pinned: datetime +): + router: Final = _io_limited_router("otpm") + + for _ in range(3): + await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) + + assert await router.cache.async_get_cache(key=_minutes_counter_key("otpm", router_minute_pinned)) == 9_000 + with pytest.raises(litellm.RateLimitError, match="OTPM limit=10000, current usage=12000"): + await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) + + +async def _realtime_opens_until_refused(monkeypatch: pytest.MonkeyPatch, **socket: object) -> tuple[int, str]: + monkeypatch.setattr(litellm, "callbacks", []) + router: Final = _io_limited_router("otpm", otpm=50) + for opened in range(51): + try: + await _kwargs_routed_to(router, "ws-io", "_arealtime", **socket) + except litellm.RateLimitError as refused: + return opened, str(refused) + pytest.fail("realtime opens never reached the deployment's otpm") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +async def test_realtime_open_reserves_and_is_refused_at_deployment_otpm_like_an_open_without_a_socket( + monkeypatch: pytest.MonkeyPatch, +): + with_socket: Final = await _realtime_opens_until_refused(monkeypatch, websocket=object()) + + assert "OTPM limit=50, current usage=" in with_socket[1] + assert with_socket == await _realtime_opens_until_refused(monkeypatch) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_turn_is_refused_once_the_minutes_tokens_reach_deployment_io_limit( + kind: str, router_minute_pinned: datetime +): + router: Final = _io_limited_router(kind) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) + + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage="): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["itpm", "otpm"]) +async def test_responses_websocket_turn_refused_for_tokens_leaves_the_daily_request_limit_unused( + kind: str, router_minute_pinned: datetime +): + router: Final = _io_limited_router(kind, rpd=2) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) + with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000"): + await admit_turn() + await _count_minutes_tokens(router, kind, -10_000, router_minute_pinned) + + await admit_turn() + + with pytest.raises(litellm.RateLimitError, match="RPD"): + await admit_turn() + + +_PROVIDER_RESPONSE: Final = { + "id": "resp_io", + "object": "response", + "created_at": 1767268830, + "status": "completed", + "model": "ws-io-unpriced", + "output": [ + { + "type": "message", + "id": "msg_io", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_metadata", [None, {"session": "abc", "team": "search"}], ids=["without_metadata", "with_metadata"] +) +async def test_responses_call_sends_the_provider_only_the_client_metadata_and_settles_deployment_io_tokens( + client_metadata: Mapping[str, str] | None, + router_minute_pinned: datetime, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_PROVIDER_RESPONSE) + ) + router: Final = _io_limited_router("itpm", otpm=10_000) + + await router.aresponses( + model="ws-io", + input="hi", + max_output_tokens=50, + **({} if client_metadata is None else {"metadata": dict(client_metadata)}), + ) + await drain_logging_tasks() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata] + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned)) + for kind in ("itpm", "otpm") + ] == [10, 5] + + +@pytest.mark.parametrize( + "client_metadata", [None, {"session": "abc", "team": "search"}], ids=["without_metadata", "with_metadata"] +) +def test_sync_responses_call_sends_the_provider_only_the_client_metadata( + client_metadata: Mapping[str, str] | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_PROVIDER_RESPONSE) + ) + + _io_limited_router("itpm", otpm=10_000).responses( + model="ws-io", + input="hi", + max_output_tokens=50, + **({} if client_metadata is None else {"metadata": dict(client_metadata)}), + ) + + assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata] + + +_ANTHROPIC_PROVIDER_RESPONSE: Final = { + "id": "msg_io", + "type": "message", + "role": "assistant", + "model": "ws-io-unpriced", + "content": [{"type": "text", "text": "Done."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, +} + + +async def _responses_call(router: Router) -> object: + return await router.aresponses(model="ws-io", input="hi", max_output_tokens=50) + + +async def _messages_call(router: Router) -> object: + return await router.aanthropic_messages(model="ws-io", messages=[{"role": "user", "content": "hi"}], max_tokens=50) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "provider_url", "provider_response", "call"), + [ + pytest.param( + "openai/ws-io-unpriced", + "https://api.openai.com/v1/responses", + _PROVIDER_RESPONSE, + _responses_call, + id="responses", + ), + pytest.param( + "anthropic/ws-io-unpriced", + "https://api.anthropic.com/v1/messages", + _ANTHROPIC_PROVIDER_RESPONSE, + _messages_call, + id="messages", + ), + ], +) +async def test_retried_call_settles_deployment_io_tokens_to_the_successful_attempt_s_tokens( + model: str, + provider_url: str, + provider_response: Mapping[str, object], + call: Callable[[Router], Awaitable[object]], + router_minute_pinned: datetime, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + provider: Final = respx_mock.post(provider_url).mock( + side_effect=[ + httpx.Response(500, json={"error": {"message": "upstream failed", "type": "server_error"}}), + httpx.Response(200, json=provider_response), + ] + ) + router: Final = Router( + model_list=[ + { + "model_name": "ws-io", + "litellm_params": {"model": model, "api_key": "sk-fake", "itpm": 10_000, "otpm": 10_000}, + "model_info": {"id": "ws-io-dep"}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + num_retries=1, + retry_after=0, + ) + + await call(router) + await drain_logging_tasks() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert provider.call_count == 2 + assert [ + await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned, model)) + for kind in ("itpm", "otpm") + ] == [10, 5] + + @pytest.mark.asyncio async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment(): """ @@ -1222,10 +1581,17 @@ def _rpm_tpm_router(model_id: str) -> Router: ) +def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None: + monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned) + monkeypatch.setattr( + "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: pinned + ) + + @pytest.fixture def router_minute_pinned(monkeypatch: pytest.MonkeyPatch) -> datetime: pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) - monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned) + _pin_router_clock(monkeypatch, pinned) return pinned @@ -2276,7 +2642,7 @@ def test_deployments_by_pattern(model_list): def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model): from re import Match - from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter pattern_router = PatternMatchRouter() @@ -2461,7 +2827,7 @@ def test_init_auto_router_deployment_success(mock_auto_router, model_list): auto_router_config=None, default_model="gpt-5-mini", embedding_model="text-embedding-3-small", - litellm_router_instance=router, + ctx=router._core_ctx, max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ) diff --git a/tests/unit/test_router_streaming_fallback_metadata.py b/tests/unit/test_router_streaming_fallback_metadata.py index 7e0cb30f1e3..40449d8cc66 100644 --- a/tests/unit/test_router_streaming_fallback_metadata.py +++ b/tests/unit/test_router_streaming_fallback_metadata.py @@ -5,9 +5,9 @@ from unittest.mock import MagicMock import pytest import litellm +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.proxy.proxy_server import _should_include_fallback_errors from litellm.router import Router -from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.types.llms.base import HiddenParams diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0f471b1de71..5cf872f4ea7 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -9269,6 +9269,77 @@ def reset_mock_cache() -> None: _model_cache.flush_cache() +@respx.mock +def test_get_valid_models_does_not_ask_a_timed_out_provider_again_on_the_next_listing(reset_mock_cache: None) -> None: + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock( + side_effect=httpx.ReadTimeout("catalog hung") + ) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-during-an-outage") + + listings: Final = tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + for _ in range(2) + ) + + assert (listings, catalog.call_count) == (([], []), 1) + + +@respx.mock +def test_get_valid_models_lists_a_provider_again_right_after_a_fast_failure(reset_mock_cache: None) -> None: + respx.get("https://openrouter.ai/api/v1/models/user").mock( + side_effect=(httpx.Response(503), httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]})) + ) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-recovers") + + listings: Final = tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + for _ in range(2) + ) + + assert listings == ([], ["openrouter/anthropic/claude-haiku-5-5"]) + + +@respx.mock +def test_get_valid_models_keeps_a_good_listing_when_an_overlapping_listing_fails_after_it( + reset_mock_cache: None, +) -> None: + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-overlapping-listings") + first_request: Final = iter((True, False)) + + def answer_after_an_overlapping_listing(request: httpx.Request) -> httpx.Response: + if next(first_request): + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + raise httpx.ReadTimeout("catalog hung", request=request) + return httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + + respx.get("https://openrouter.ai/api/v1/models/user").mock(side_effect=answer_after_an_overlapping_listing) + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + + assert get_valid_models( + check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment + ) == ["openrouter/anthropic/claude-haiku-5-5"] + + +@respx.mock +def test_get_valid_models_asks_a_provider_once_when_many_listings_start_together(reset_mock_cache: None) -> None: + def slow_catalog(request: httpx.Request) -> httpx.Response: + threading.Event().wait(0.2) + return httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock(side_effect=slow_catalog) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-many-listings") + + def list_models(_: int) -> tuple[str, ...]: + return tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + ) + + with ThreadPoolExecutor(max_workers=8) as pool: + listings: Final = frozenset(pool.map(list_models, range(8))) + + assert (listings, catalog.call_count) == (frozenset({("openrouter/anthropic/claude-haiku-5-5",)}), 1) + + @pytest.fixture def isolated_openai_model_sets(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( diff --git a/tests/unit/types/llms/test_types_llms_openai.py b/tests/unit/types/llms/test_types_llms_openai.py index 1b7a9693f7b..05801c9feec 100644 --- a/tests/unit/types/llms/test_types_llms_openai.py +++ b/tests/unit/types/llms/test_types_llms_openai.py @@ -607,6 +607,23 @@ def test_openai_file_object_accepts_pending_status(): assert file_obj.status == "pending" +@pytest.mark.parametrize("object_type", ({}, {"object": None}), ids=["missing", "null"]) +def test_openai_file_object_without_an_object_type_validates_as_a_file(object_type: dict[str, None]) -> None: + from litellm.types.llms.openai import OpenAIFileObject + + server_file: Final = { + "id": "file-b73f4c", + "bytes": 156, + "created_at": 1793238707, + "filename": "batch_requests.jsonl", + "purpose": "batch", + } + + file_obj: Final = OpenAIFileObject.model_validate({**server_file, **object_type}) + + assert file_obj.model_dump(exclude_none=True) == {**server_file, "object": "file"} + + class TestOpenAIFileObjectBatchGuardrailSerialization: """The proxy-only `litellm_batch_guardrail` key must reach the wire only when something set it.""" diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index e369b7e93ab..22eb9929e7f 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -33,6 +33,7 @@ from litellm.types.litellm_params import ( wire_names, ) from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionUserMessage +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import HeadroomRequestContext from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.router import ( ConfigurableClientsideParamsCustomAuth, @@ -231,6 +232,7 @@ OPTION_NAMES: Final = ( ) AGENTIC_LOOP_STATE_NAMES: Final = ( + "_headroom_request_context", "_agentic_loop_depth", "_agentic_loop_fingerprints", "_agentic_loop_api_surface", @@ -529,6 +531,7 @@ TYPE_HINT_NAMESPACE: Final[Mapping[str, object]] = { "Logging": Logging, "SecretFields": SecretFields, "CompactionState": CompactionState, + "HeadroomRequestContext": HeadroomRequestContext, "RouterWeights": RouterWeights, "AttemptedFallbackTargets": AttemptedFallbackTargets, } diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 9fce93782b2..c14f3e6f036 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -625,23 +625,6 @@ "count": 1 } }, - "src/app/(dashboard)/old-usage/_components/usage.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "max-lines": { - "count": 1 - }, - "prefer-const": { - "count": 6 - }, - "react-hooks/immutability": { - "count": 1 - }, - "react-hooks/purity": { - "count": 1 - } - }, "src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx": { "react-hooks/set-state-in-effect": { "count": 2 @@ -2090,17 +2073,6 @@ "count": 1 } }, - "src/components/view_user_spend.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "prefer-const": { - "count": 3 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, "src/contexts/AuthContext.tsx": { "react-hooks/set-state-in-effect": { "count": 1 diff --git a/ui/litellm-dashboard/public/assets/logos/aiand.svg b/ui/litellm-dashboard/public/assets/logos/aiand.svg new file mode 100644 index 00000000000..486c73b8dbc --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/aiand.svg @@ -0,0 +1,4 @@ + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx index 78091584cde..12ff10bc0c5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx @@ -1,6 +1,7 @@ import { useTerminology } from "@/contexts/TerminologyContext"; import { Play } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { Button } from "@/components/ui/button"; @@ -48,6 +49,7 @@ export function EvaluationSettingsModal({ accessToken, onRunEvaluation, }: EvaluationSettingsModalProps) { + const { t } = useTranslation(); const term = useTerminology(); const [prompt, setPrompt] = useState(DEFAULT_PROMPT); const [schema, setSchema] = useState(DEFAULT_SCHEMA); @@ -94,11 +96,11 @@ export function EvaluationSettingsModal({ !nextOpen && onClose()}> - Evaluation Settings + {t("guardrailsMonitor.evaluationSettings.title")} {guardrailName - ? `${term("Configure AI evaluation for")} ${guardrailName}` - : term("Configure AI evaluation for re-running on logs")} + ? `${term(t("guardrailsMonitor.evaluationSettings.descriptionPrefix"))} ${guardrailName}` + : term(t("guardrailsMonitor.evaluationSettings.descriptionGeneric"))} @@ -106,10 +108,10 @@ export function EvaluationSettingsModal({