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-unit.yml b/.github/workflows/test-unit.yml index 5ec5975d83b..2c05ee827a3 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -577,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/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/_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/constants.py b/litellm/constants.py index c586a5f767a..1a1e66ecf86 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -177,6 +177,8 @@ MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 # Set to 0 to disable truncation. MAX_BASE64_LENGTH_FOR_LOGGING: Final = get_env_int("MAX_BASE64_LENGTH_FOR_LOGGING", 64) BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: Final = 256 * 1024 +BASE64_TRUNCATION_OFFLOAD_THRESHOLD_STRINGS: Final = 256 +BASE64_TRUNCATION_INLINE_COPY_MAX_CHARS: Final = 8 * 1024 * 1024 REDACTED_BY_LITELLM: Final = "redacted-by-litellm" # in-memory stand-in handed to provider converters for redacted arguments; never stored REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}" @@ -634,6 +636,8 @@ STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS: Final = 60 BEDROCK_MIN_THINKING_BUDGET_TOKENS: Final = get_env_int("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024) # Anthropic's Messages API rejects thinking.budget_tokens < 1024. ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: Final = 1024 +ANTHROPIC_MAX_USER_ID_LENGTH: Final = 512 +ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH: Final = 255 ANTHROPIC_BILLING_METADATA_PREFIX: Final = "x-anthropic-billing-header:" REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_float("REPLICATE_POLLING_DELAY_SECONDS", 0.5) DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS: Final = get_env_int("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 128000) @@ -902,6 +906,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", @@ -1162,6 +1167,7 @@ openai_compatible_providers: Final[list[str]] = [ "cortecs", "scx-ai", "prism", + "hostinger_router", "sail", "reka", "streamlake", @@ -2046,6 +2052,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..871fa8227ea 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -3,10 +3,10 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field, replace from functools import partial from types import MappingProxyType -from typing import Final, Literal, TypeAlias +from typing import Annotated, Final, Literal, TypeAlias import httpx -from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError +from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError from pydantic_core import to_jsonable_python from typing_extensions import assert_never @@ -23,6 +23,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( BaseDecisionsConfig, DecisionsResponseValidationError, + catalog_mode, image_media_type, image_param, input_images, @@ -61,7 +62,7 @@ from litellm.types.decisions import ( UnsupportedDecisionsRequest, ) from litellm.types.llms.openai import ChatCompletionUserMessage -from litellm.types.utils import LlmProviders, ProviderSpecificHeader +from litellm.types.utils import LlmProviders, ProviderSpecificHeader, is_decisions_model_mode from litellm.utils import ProviderConfigManager, get_model_cost_key SystemOneQuestions: TypeAlias = Mapping[str, DecisionQuestion | Mapping[str, object]] @@ -218,7 +219,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) @@ -257,19 +258,105 @@ def _request_safety_identifier( return None -def _fields_without_unsupported_media( - fields: RequestFields, +def _bool_or_none(value: object) -> bool | None: + return value if isinstance(value, bool) else None + + +def _str_or_none(value: object) -> str | None: + return value if isinstance(value, str) else None + + +_DeploymentFlag: TypeAlias = Annotated[bool | None, BeforeValidator(_bool_or_none)] + + +class _DeploymentCapabilities(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + mode: Annotated[str | None, BeforeValidator(_str_or_none)] = None + supports_box_questions: _DeploymentFlag = None + supports_audio_input: _DeploymentFlag = None + supports_video_input: _DeploymentFlag = None + + def decisions_media_flag(self, flag: bool | None, model: str, custom_llm_provider: str) -> bool | None: + if flag is None: + return None + mode: Final = self.mode if self.mode is not None else catalog_mode(model, custom_llm_provider) + return flag if mode is None or is_decisions_model_mode(mode) else None + + +@dataclass(frozen=True, slots=True) +class _ResolvedSupport: + box_questions: bool + audio: bool + videos: bool + + +def _takes_audio( + provider_config: BaseDecisionsConfig, model: str, custom_llm_provider: str, deployment: _DeploymentCapabilities +) -> bool: + return provider_config.supports_audio_input( + model, + custom_llm_provider, + deployment.decisions_media_flag(deployment.supports_audio_input, model, custom_llm_provider), + ) + + +def _takes_videos( + provider_config: BaseDecisionsConfig, model: str, custom_llm_provider: str, deployment: _DeploymentCapabilities +) -> bool: + return provider_config.supports_video_input( + model, + custom_llm_provider, + deployment.decisions_media_flag(deployment.supports_video_input, model, custom_llm_provider), + ) + + +def deployment_refuses_media( + model: str, + custom_llm_provider: str, + model_info: Mapping[str, object], + call_kwargs: Mapping[str, object], + *, + audio: bool, + videos: bool, +) -> bool: + if _drops_params(call_kwargs): + return False + provider_config: Final = get_decisions_provider_config(model, custom_llm_provider) + canonical_model: Final = provider_config.canonical_model(model) + deployment: Final = _DeploymentCapabilities.model_validate(model_info) + return (audio and not _takes_audio(provider_config, canonical_model, custom_llm_provider, deployment)) or ( + videos and not _takes_videos(provider_config, canonical_model, custom_llm_provider, deployment) + ) + + +def _resolved_support( provider_config: BaseDecisionsConfig, model: str, custom_llm_provider: str, + headers: Mapping[str, str], kwargs: Mapping[str, object], + fields: RequestFields, +) -> _ResolvedSupport: + deployment: Final = _DeploymentCapabilities.model_validate(kwargs.get("model_info") or {}) + return _ResolvedSupport( + box_questions=provider_config.supports_box_questions( + model, custom_llm_provider, headers, deployment.supports_box_questions + ), + audio=bool(fields.audio) and _takes_audio(provider_config, model, custom_llm_provider, deployment), + videos=bool(fields.videos) and _takes_videos(provider_config, model, custom_llm_provider, deployment), + ) + + +def _fields_without_unsupported_media( + fields: RequestFields, support: _ResolvedSupport, kwargs: Mapping[str, object] ) -> RequestFields: if not _drops_params(kwargs): return fields return replace( fields, - audio=fields.audio if provider_config.supports_audio_input(model, custom_llm_provider) else None, - videos=fields.videos if provider_config.supports_video_input(model, custom_llm_provider) else None, + audio=fields.audio if support.audio else None, + videos=fields.videos if support.videos else None, ) @@ -292,10 +379,8 @@ class _MalformedImages: def _unsupported_params( ir_request: DecisionsIRRequest, provider_config: BaseDecisionsConfig, - model: str, - custom_llm_provider: str, headers: Mapping[str, str], - supports_box_questions: bool, + support: _ResolvedSupport, ) -> tuple[str, ...]: refuses_safety_identifier: Final = ( ir_request.safety_identifier is not None and not provider_config.supports_safety_identifier @@ -304,13 +389,9 @@ def _unsupported_params( None if provider_config.supports_images_in(ir_request.input, headers) else image_param(ir_request.input) ) refuses_assets: Final = bool(ir_request.assets) and not provider_config.supports_assets(headers) - refuses_box_questions: Final = _has_box_question(ir_request.questions) and not supports_box_questions - refuses_audio: Final = bool(ir_request.audio) and not provider_config.supports_audio_input( - model, custom_llm_provider - ) - refuses_videos: Final = bool(ir_request.videos) and not provider_config.supports_video_input( - model, custom_llm_provider - ) + refuses_box_questions: Final = _has_box_question(ir_request.questions) and not support.box_questions + refuses_audio: Final = bool(ir_request.audio) and not support.audio + refuses_videos: Final = bool(ir_request.videos) and not support.videos return ( *(("safety_identifier",) if refuses_safety_identifier else ()), *(() if refused_image_param is None else (refused_image_param,)), @@ -333,7 +414,7 @@ def _without_unsupported_params( ir_request: DecisionsIRRequest, provider_config: BaseDecisionsConfig, headers: Mapping[str, str], - supports_box_questions: bool, + support: _ResolvedSupport, ) -> DecisionsIRRequest: return replace( ir_request, @@ -345,9 +426,7 @@ def _without_unsupported_params( safety_identifier=ir_request.safety_identifier if provider_config.supports_safety_identifier else None, assets=ir_request.assets if provider_config.supports_assets(headers) else MappingProxyType({}), questions=( - ir_request.questions - if supports_box_questions - else tuple(filter(_is_not_box_question, ir_request.questions)) + ir_request.questions if support.box_questions else tuple(filter(_is_not_box_question, ir_request.questions)) ), ) @@ -368,39 +447,20 @@ def _without_malformed_images( return well_formed -class _DeploymentBoxQuestions(BaseModel): - model_config = ConfigDict(frozen=True, from_attributes=True) - - supports_box_questions: StrictBool | None = None - - -def _deployment_supports_box_questions(kwargs: Mapping[str, object]) -> bool | None: - try: - return _DeploymentBoxQuestions.model_validate(kwargs.get("model_info") or {}).supports_box_questions - except ValidationError: - return None - - def _provider_ir_request( ir_request: DecisionsIRRequest, *, provider_config: BaseDecisionsConfig, - model: str, - custom_llm_provider: str, headers: Mapping[str, str], + support: _ResolvedSupport, kwargs: Mapping[str, object], ) -> DecisionsIRRequest | _UnsupportedParams | _NothingLeftAfterDrop: - supports_box_questions: Final = provider_config.supports_box_questions( - model, custom_llm_provider, headers, _deployment_supports_box_questions(kwargs) - ) - unsupported: Final = _unsupported_params( - ir_request, provider_config, model, custom_llm_provider, headers, supports_box_questions - ) + unsupported: Final = _unsupported_params(ir_request, provider_config, headers, support) if not unsupported: return ir_request if not _drops_params(kwargs): return _UnsupportedParams(names=unsupported) - supported: Final = _without_unsupported_params(ir_request, provider_config, headers, supports_box_questions) + supported: Final = _without_unsupported_params(ir_request, provider_config, headers, support) if _leaves_no_input(supported.input): return _NothingLeftAfterDrop(names=unsupported, missing="input to decide on") if not supported.questions: @@ -408,7 +468,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 +526,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,11 +538,12 @@ 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) + support: Final = _resolved_support(provider_config, canonical_model, provider, headers, kwargs, fields) 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 - ).validate(safety_identifier=request_safety_identifier, drops_params=_drops_params(kwargs)) + request: Final = _fields_without_unsupported_media(fields, support, kwargs).validate( + safety_identifier=request_safety_identifier, drops_params=_drops_params(kwargs) + ) except ValidationError as error: raise litellm.BadRequestError( message=f"Invalid Decisions request: {error}", @@ -523,12 +584,7 @@ def prepare_call( llm_provider=provider, ) ir_request: Final = _provider_ir_request( - caller_ir_request, - provider_config=provider_config, - model=canonical_model, - custom_llm_provider=provider, - headers=headers, - kwargs=kwargs, + caller_ir_request, provider_config=provider_config, headers=headers, support=support, kwargs=kwargs ) if isinstance(ir_request, _UnsupportedParams): raise litellm.UnsupportedParamsError( diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index d816bc661f3..39368dee47c 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -14,7 +14,7 @@ import itertools import os import re from collections.abc import Callable, Container, Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar, cast from urllib.parse import urlparse from pydantic import TypeAdapter, ValidationError @@ -102,6 +102,192 @@ def _validated_object_list(value: object) -> list[object] | None: return None +_MessageT: Final = TypeVar("_MessageT") +_CacheControlRewrite: TypeAlias = Callable[[object], ChatCompletionCachedContent | None] +_MARK_CARRYING_MESSAGE_FIELDS: Final = ("content", "tool_calls") + + +def _mark_carrying_blocks(fields: Mapping[str, object]) -> Sequence[object]: + lists: Final = tuple(fields.get(key) for key in _MARK_CARRYING_MESSAGE_FIELDS) + return tuple( + itertools.chain.from_iterable( + _OBJECT_LIST_ADAPTER.validate_python(value) for value in lists if isinstance(value, list) + ) + ) + + +def _block_carries_cache_control(block: object) -> bool: + return isinstance(block, dict) and _CONTENT_PART_ADAPTER.validate_python(block).get("cache_control") is not None + + +_ROLES_SENDING_MESSAGE_MARK_WITH_LIST_CONTENT: Final = frozenset({"tool", "function"}) + + +def _message_mark_reaches_provider(fields: Mapping[str, object]) -> bool: + if fields.get("cache_control") is None: + return False + return ( + not isinstance(fields.get("content"), list) + or fields.get("role") in _ROLES_SENDING_MESSAGE_MARK_WITH_LIST_CONTENT + ) + + +def cache_control_count(message: object) -> int: + if not isinstance(message, dict): + return 0 + fields: Final = _CONTENT_PART_ADAPTER.validate_python(message) + block_marks: Final = sum(1 for block in _mark_carrying_blocks(fields) if _block_carries_cache_control(block)) + return int(_message_mark_reaches_provider(fields)) + block_marks + + +def _fields_with_cache_control(fields: Mapping[str, object], rewrite: _CacheControlRewrite) -> Mapping[str, object]: + others: Final = {key: value for key, value in fields.items() if key != "cache_control"} + if "cache_control" not in fields: + return others + cache_control: Final = rewrite(fields["cache_control"]) + return others if cache_control is None else {**others, "cache_control": cache_control} + + +def block_with_cache_control(block: object, rewrite: _CacheControlRewrite) -> object: + if not isinstance(block, dict): + return block + return _fields_with_cache_control(_CONTENT_PART_ADAPTER.validate_python(block), rewrite) + + +def with_cache_control(message: _MessageT, rewrite: _CacheControlRewrite) -> _MessageT: + if not isinstance(message, dict): + return message + fields: Final = _CONTENT_PART_ADAPTER.validate_python(message) + blocks_rewritten: Final = { + key: ( + [block_with_cache_control(block, rewrite) for block in _OBJECT_LIST_ADAPTER.validate_python(value)] + if key in _MARK_CARRYING_MESSAGE_FIELDS and isinstance(value, list) + else value + ) + for key, value in fields.items() + } + return cast( # cast-ok: the same message shape with each cache marker rewritten or dropped + _MessageT, _fields_with_cache_control(blocks_rewritten, rewrite) + ) + + +def dropped_cache_control(_: object) -> None: + return None + + +def _first_marked_index(value: object) -> int | None: + if not isinstance(value, list): + return None + blocks: Final = _OBJECT_LIST_ADAPTER.validate_python(value) + return next((index for index, block in enumerate(blocks) if _block_carries_cache_control(block)), None) + + +_MarkSlot: TypeAlias = tuple[str, int] +_MESSAGE_MARK_SLOT: Final[_MarkSlot] = ("cache_control", -1) + + +def _block_mark_slots(fields: Mapping[str, object], key: str) -> tuple[_MarkSlot, ...]: + value: Final = fields.get(key) + if not isinstance(value, list): + return () + blocks: Final = _OBJECT_LIST_ADAPTER.validate_python(value) + return tuple((key, index) for index, block in enumerate(blocks) if _block_carries_cache_control(block)) + + +def _mark_slots_in_wire_order(fields: Mapping[str, object]) -> tuple[_MarkSlot, ...]: + content_slots: Final = _block_mark_slots(fields, "content") + tool_call_slots: Final = _block_mark_slots(fields, "tool_calls") + if not _message_mark_reaches_provider(fields): + return content_slots + tool_call_slots + if isinstance(fields.get("content"), list): + return content_slots + (_MESSAGE_MARK_SLOT,) + tool_call_slots + return (_MESSAGE_MARK_SLOT,) + content_slots + tool_call_slots + + +def _blocks_without_marks(key: str, value: object, dropped_slots: frozenset[_MarkSlot]) -> object: + if key not in _MARK_CARRYING_MESSAGE_FIELDS or not isinstance(value, list): + return value + return [ + block_with_cache_control(block, dropped_cache_control) if (key, index) in dropped_slots else block + for index, block in enumerate(_OBJECT_LIST_ADAPTER.validate_python(value)) + ] + + +def _without_oldest_marks(message: _MessageT, dropped: int) -> _MessageT: + if dropped == 0 or not isinstance(message, dict): + return message + fields: Final = _CONTENT_PART_ADAPTER.validate_python(message) + dropped_slots: Final = frozenset(_mark_slots_in_wire_order(fields)[:dropped]) + blocks_unmarked: Final = {key: _blocks_without_marks(key, value, dropped_slots) for key, value in fields.items()} + return cast( # cast-ok: the same message with its oldest wire-order marks dropped + _MessageT, + _fields_with_cache_control(blocks_unmarked, dropped_cache_control) + if _MESSAGE_MARK_SLOT in dropped_slots + else blocks_unmarked, + ) + + +def tool_message_with_single_cache_mark(message: _MessageT) -> _MessageT: + if not isinstance(message, dict): + return message + fields: Final = _CONTENT_PART_ADAPTER.validate_python(message) + content: Final = fields.get("content") + if fields.get("role") != "tool" or not isinstance(content, list): + return message + kept_index: Final = None if fields.get("cache_control") is not None else _first_marked_index(content) + return cast( # cast-ok: the same tool message keeping only the mark its single cache point reads + _MessageT, + { + **fields, + "content": [ + block if index == kept_index else block_with_cache_control(block, dropped_cache_control) + for index, block in enumerate(_OBJECT_LIST_ADAPTER.validate_python(content)) + ], + }, + ) + + +_BudgetSpend: TypeAlias = Callable[[tuple[int, int], int], tuple[int, int]] + + +def _newest_marks_kept(state: tuple[int, int], marks: int) -> tuple[int, int]: + remaining, _ = state + kept: Final = min(marks, remaining) + return remaining - kept, marks - kept + + +def _whole_message_kept(state: tuple[int, int], marks: int) -> tuple[int, int]: + remaining, _ = state + if marks > remaining: + return remaining, marks + return remaining - marks, 0 + + +def _messages_within_budget( + messages: Sequence[_MessageT], reserved_marks: int, spend: _BudgetSpend +) -> Sequence[_MessageT]: + budget: Final = max(MAX_CACHE_CONTROL_BLOCKS - reserved_marks, 0) + newest_first: Final = tuple(reversed(messages)) + budget_after_each: Final = itertools.accumulate( + (cache_control_count(message) for message in newest_first), spend, initial=(budget, 0) + ) + dropped_counts: Final = tuple(dropped for _, dropped in itertools.islice(budget_after_each, 1, None)) + kept_newest_first: Final = tuple( + _without_oldest_marks(message, dropped) for message, dropped in zip(newest_first, dropped_counts) + ) + return tuple(reversed(kept_newest_first)) + + +def newest_marks_within_cache_control_limit(messages: Sequence[_MessageT], reserved_marks: int) -> Sequence[_MessageT]: + return _messages_within_budget(messages, reserved_marks, _newest_marks_kept) + + +def whole_messages_within_cache_control_limit( + messages: Sequence[_MessageT], reserved_marks: int +) -> Sequence[_MessageT]: + return _messages_within_budget(messages, reserved_marks, _whole_message_kept) + + def _part_with_cache_control_renamed(part: object) -> object: if not _accepts_prompt_cache_breakpoint(part): return part @@ -1439,3 +1625,25 @@ class AnthropicCacheControlHook(CustomPromptManagement): llm_router=None, ) return None + + +def messages_within_cache_control_cap( + messages: Sequence[_MessageT], tools: object, cache_control: object = None, request_kwargs: object = None +) -> Sequence[_MessageT]: + """Under drop_params, drop the oldest message hints the provider's cache_control cap has no room for. + + Tool hints and a top-level ``cache_control`` stay on the wire, so they are reserved first. + """ + reserved_marks: Final = AnthropicCacheControlHook.count_external_cache_breakpoints( + _validated_object_list(tools), cache_control, request_kwargs + ) + kept: Final = newest_marks_within_cache_control_limit(messages, reserved_marks) + dropped: Final = sum(map(cache_control_count, messages)) - sum(map(cache_control_count, kept)) + if dropped > 0: + verbose_logger.warning( + "Dropped %d cache_control hint(s) from the oldest messages because drop_params is on and the " + "provider accepts at most %d cache_control blocks per request", + dropped, + MAX_CACHE_CONTROL_BLOCKS, + ) + return kept 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 5c2f0f553a8..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 ( @@ -3022,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: @@ -5209,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/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 33b6d4e1857..85c232ad3f4 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -98,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", 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/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 26949d9bbfb..924603dc7e8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -165,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, @@ -355,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") @@ -6817,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, @@ -6853,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): 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/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 8f58e8b7051..88d84201beb 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -3,13 +3,15 @@ import functools import inspect import re import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast from litellm._logging import format_base64_size, verbose_logger from litellm.constants import ( + BASE64_TRUNCATION_INLINE_COPY_MAX_CHARS, BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS, + BASE64_TRUNCATION_OFFLOAD_THRESHOLD_STRINGS, MAX_BASE64_LENGTH_FOR_LOGGING, ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs @@ -44,6 +46,8 @@ Helper utils used for logging callbacks # Regex matching data-URI base64 content: "data:;base64," # Captures: group(1)=mime_type, group(2)=base64_payload _DATA_URI_RE: Final = re.compile(r"data:([^;,\s]{1,255});base64,([A-Za-z0-9+/=]+)") +_DATA_URI_PREFIX: Final = "data:" +_DATA_URI_BASE64_MARKER: Final = ";base64," # Maximum nesting depth for _truncate_base64_in_value to guard against # pathological payloads. OpenAI message format is typically 3-4 levels deep. @@ -60,15 +64,20 @@ def _base64_data_uri_replacer(match: re.Match[str]) -> str: return f"data:{mime_type};base64,[base64_data truncated: {size_str}]" +def _may_contain_base64_data_uri(value: str) -> bool: + return ";" in value and _DATA_URI_PREFIX in value and _DATA_URI_BASE64_MARKER in value + + def _truncate_base64_in_string(value: str) -> str: """Replace long base64 data-URI payloads in a string with a size placeholder.""" - if MAX_BASE64_LENGTH_FOR_LOGGING <= 0: + if MAX_BASE64_LENGTH_FOR_LOGGING <= 0 or not _may_contain_base64_data_uri(value): return value return _DATA_URI_RE.sub(_base64_data_uri_replacer, value) def _truncate_base64_in_value( value: str | dict[str, object] | list[object] | None, + truncate_string: Callable[[str], str], ) -> str | dict[str, object] | list[object] | None: """Iteratively truncate base64 data URIs in a JSON-like value (str/list/dict). @@ -78,13 +87,13 @@ def _truncate_base64_in_value( # Stack entries: (source_value, depth, parent_container, key_or_index) # We mutate *copies* of dicts/lists in-place via parent references. if isinstance(value, str): - return _truncate_base64_in_string(value) + return truncate_string(value) if not isinstance(value, (dict, list)): return value # Shallow-copy the root so we don't mutate the caller's data. root: Final = {k: v for k, v in value.items()} if isinstance(value, dict) else list(value) - stack: Final[list] = [(root, 0)] + stack: Final[list[tuple[dict[str, object] | list[object], int]]] = [(root, 0)] # mutable-ok: explicit stack while stack: container, depth = stack.pop() @@ -93,19 +102,19 @@ def _truncate_base64_in_value( if isinstance(container, dict): for k, v in container.items(): if isinstance(v, str): - container[k] = _truncate_base64_in_string(v) + container[k] = truncate_string(v) elif isinstance(v, dict): - copy: dict | list = {ck: cv for ck, cv in v.items()} + copy: dict[str, object] | list[object] = {ck: cv for ck, cv in v.items()} container[k] = copy stack.append((copy, depth + 1)) elif isinstance(v, list): copy = list(v) container[k] = copy stack.append((copy, depth + 1)) - elif isinstance(container, list): + else: for i, v in enumerate(container): if isinstance(v, str): - container[i] = _truncate_base64_in_string(v) + container[i] = truncate_string(v) elif isinstance(v, dict): copy = {ck: cv for ck, cv in v.items()} container[i] = copy @@ -127,8 +136,15 @@ def truncate_base64_in_messages( """ if messages is None or MAX_BASE64_LENGTH_FOR_LOGGING <= 0: return messages + return _copy_messages(messages, _truncate_base64_in_string) + + +def _copy_messages( + messages: str | dict[str, object] | list[object], # mutable-ok: same contract as truncate_base64_in_messages + truncate_string: Callable[[str], str], +) -> str | dict[str, object] | list[object] | None: # mutable-ok: same contract as truncate_base64_in_messages try: - return _truncate_base64_in_value(messages) + return _truncate_base64_in_value(messages, truncate_string) except Exception as e: verbose_logger.debug("Failed to truncate base64 in messages: %s", e) return messages @@ -151,19 +167,31 @@ def _iter_string_leaves(value: _StringTree) -> Iterator[str]: pass +def _is_few_plain_strings(leaves: tuple[str, ...]) -> bool: + return len(leaves) < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_STRINGS and not any( + _may_contain_base64_data_uri(leaf) for leaf in leaves + ) + + async def truncate_base64_in_messages_async( messages: str | list | dict | None, # mutable-ok: same contract as truncate_base64_in_messages ) -> str | list | dict | None: # mutable-ok: same contract as truncate_base64_in_messages """ Same result as truncate_base64_in_messages, but payloads whose string content - reaches BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS are scanned in a worker - thread so the regex pass over multi-MB base64 images does not block the event loop. + reaches BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS are copied in a worker thread, + so neither the regex pass over multi-MB base64 images nor the copy of a long + history blocks the event loop. A few strings with no data URI, under + BASE64_TRUNCATION_INLINE_COPY_MAX_CHARS in total, are copied inline, since + checking them for data URIs and copying them costs less than the thread hop. """ if messages is None or MAX_BASE64_LENGTH_FOR_LOGGING <= 0: return messages - total_chars: Final = sum(len(leaf) for leaf in _iter_string_leaves(messages)) + leaves: Final = tuple(_iter_string_leaves(messages)) + total_chars: Final = sum(len(leaf) for leaf in leaves) if total_chars < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: return truncate_base64_in_messages(messages) + if total_chars < BASE64_TRUNCATION_INLINE_COPY_MAX_CHARS and _is_few_plain_strings(leaves): + return _copy_messages(messages, lambda text: text) return await asyncio.to_thread(truncate_base64_in_messages, messages) 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/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 862e61cb631..e733f48ffd7 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -88,17 +88,22 @@ def handle_any_messages_to_chat_completion_str_messages_conversion( return [{"input": str(messages)}] +def _message_with_str_content(message: AllMessageValues) -> AllMessageValues: + texts: Final = convert_content_list_to_str(message=message) + if not texts: + return message + return cast( # cast-ok: the same message with its content flattened to one string + AllMessageValues, {**message, "content": texts} + ) + + def handle_messages_with_content_list_to_str_conversion( messages: list[AllMessageValues], ) -> list[AllMessageValues]: """ - Handles messages with content list conversion + Returns new messages with each content list flattened to a string, leaving the caller's messages untouched """ - for message in messages: - texts = convert_content_list_to_str(message=message) - if texts: - message["content"] = texts - return messages + return [_message_with_str_content(message) for message in messages] def strip_name_from_message(message: AllMessageValues, allowed_name_roles: list[str] = ["user"]) -> AllMessageValues: 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 100% rename from litellm/router_utils/reasoning_effort_capability.py rename to litellm/litellm_core_utils/reasoning_effort_capability.py 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/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0ddfa1210c2..4857973d51d 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -329,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, @@ -506,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 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 9710956c97e..484edd869cb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -277,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: @@ -293,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) @@ -308,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 } @@ -360,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 = [ @@ -1653,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] = {} 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 d773e31c733..25245c347f4 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -167,6 +167,7 @@ _DEFER_LOADING_ADAPTER: Final = TypeAdapter(bool) _TOOL_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _TOOLS_ADAPTER: Final = TypeAdapter(list[object]) _LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_REQUEST_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _NATIVE_OUTPUT_STRING_FORMATS: Final = frozenset( ("date-time", "time", "date", "duration", "email", "hostname", "uri", "ipv4", "ipv6", "uuid") ) @@ -2254,6 +2255,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): headers=headers, ) + @staticmethod + def _messages_within_cache_control_cap( + messages: Sequence[AllMessageValues], optional_params: Mapping[str, object] + ) -> Sequence[AllMessageValues]: + from litellm.integrations.anthropic_cache_control_hook import messages_within_cache_control_cap + + return messages_within_cache_control_cap( + messages, optional_params.get("tools"), optional_params.get("cache_control"), optional_params + ) + def transform_request( self, model: str, @@ -2271,9 +2282,18 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): anthropic_messages_pt, ) - drop_params: Final = litellm.drop_params is True or litellm_params.get("drop_params") is True + drop_params: Final = drops_unsupported_params(litellm_params.get("drop_params")) + capped_messages: Final = ( + list( + self._messages_within_cache_control_cap( + messages, _REQUEST_PARAMS_ADAPTER.validate_python(optional_params) + ) + ) + if drop_params + else messages + ) - if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages): + if "tools" not in optional_params and has_tool_call_blocks(capped_messages): optional_params["tools"], _ = self._map_tools(add_dummy_tool(custom_llm_provider="anthropic")) # Drop thinking param if thinking is enabled but thinking_blocks are missing @@ -2285,9 +2305,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Related issue: https://github.com/BerriAI/litellm/issues/18926 if ( optional_params.get("thinking") is not None - and messages is not None - and last_assistant_with_tool_calls_has_no_thinking_blocks(messages) - and not any_assistant_message_has_thinking_blocks(messages) + and last_assistant_with_tool_calls_has_no_thinking_blocks(capped_messages) + and not any_assistant_message_has_thinking_blocks(capped_messages) ): if litellm.modify_params: optional_params.pop("thinking", None) @@ -2332,15 +2351,18 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _name_forward_map, _name_reverse_map = self._sanitize_tool_names_in_request( optional_params=optional_params, ) - if _name_forward_map: - messages = self._rewrite_tool_names_in_messages(messages, _name_forward_map) + request_messages: Final = ( + self._rewrite_tool_names_in_messages(capped_messages, _name_forward_map) + if _name_forward_map + else capped_messages + ) if _name_reverse_map and isinstance(litellm_params, dict): litellm_params[ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY] = _name_reverse_map # Only the leading system run becomes the top-level system prompt. A later # system message stays in the conversation: hoisting it rewrites the cached # prefix and re-bills the whole history at cache-write pricing (#36559). - leading_system_run, later_messages = split_leading_system_run(messages) + leading_system_run, later_messages = split_leading_system_run(request_messages) anthropic_system_message_list: Final = self.translate_system_message(messages=list(leading_system_run)) # Handling anthropic API Prompt Caching if len(anthropic_system_message_list) > 0: @@ -2362,7 +2384,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): except Exception as e: raise AnthropicError( status_code=400, - message=f"{e}\nReceived Messages={messages}", + message=f"{e}\nReceived Messages={request_messages}", ) # don't use verbose_logger.exception, if exception is raised self.update_headers_with_optional_anthropic_beta( diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index e1f9a7d540f..e0da8943a4a 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -4,6 +4,7 @@ This file contains common utils for anthropic calls. import copy import json +import os import re from collections.abc import Container, Mapping, MutableMapping, Sequence from dataclasses import dataclass @@ -19,6 +20,8 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( + ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH, + ANTHROPIC_MAX_USER_ID_LENGTH, DEFAULT_MODEL_CREATED_AT_TIME, DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, @@ -178,6 +181,40 @@ def tools_without_custom_field(tools: Sequence[object], hoist_defer_loading: boo return [_tool_without_custom_field(tool, hoist_defer_loading, drop_custom) for tool in tools] +def anthropic_pass_through_api_base() -> str: + return os.getenv("ANTHROPIC_API_BASE") or os.getenv("ANTHROPIC_BASE_URL") or "https://api.anthropic.com" + + +def _same_host_and_port(url: httpx.URL, other: str) -> bool: + try: + other_url: Final = httpx.URL(other) + except httpx.InvalidURL: + return False + return (url.host, url.port) == (other_url.host, other_url.port) + + +def anthropic_upstream_user_id_limit(url: str) -> int | None: + upstream: Final = httpx.URL(url) + if not upstream.path.removesuffix("/").endswith("/v1/messages"): + return None + if upstream.host == "api.anthropic.com": + return ANTHROPIC_MAX_USER_ID_LENGTH + return ( + ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH + if _same_host_and_port(upstream, anthropic_pass_through_api_base()) + else None + ) + + +def _has_email_shape(token: str) -> bool: + at_sign: Final = token.find("@", 1) + return at_sign != -1 and token.find(".", at_sign + 2, -1) != -1 + + +def anthropic_accepts_user_id(user_id: str, max_length: int) -> bool: + return len(user_id) <= max_length and user_id.isprintable() and not any(map(_has_email_shape, user_id.split())) + + def is_anthropic_messages_url(url: str) -> bool: """Check whether a URL addresses Anthropic's Messages API.""" parsed_url: Final = urlparse(url) diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 1da7c2efdda..1bf00cbdb2f 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -19,6 +19,13 @@ 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, @@ -31,13 +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, -) 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 ( 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/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/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..c562c8d0f04 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 @@ -12,6 +12,7 @@ from pydantic_core import ErrorDetails from typing_extensions import assert_never import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -297,13 +298,7 @@ def _systemone_question(question: DecisionsIRQuestion) -> Mapping[str, object]: assert_never(question) -def ir_to_systemone_request( - model: str, request: DecisionsIRRequest, *, supports_audio: bool = False, supports_video: bool = False -) -> Mapping[str, object] | UnsupportedDecisionsRequest: - if request.audio and not supports_audio: - return UnsupportedDecisionsRequest(reason="audio input is not supported by this Decisions model") - if request.videos and not supports_video: - return UnsupportedDecisionsRequest(reason="video input is not supported by this Decisions model") +def ir_to_systemone_request(model: str, request: DecisionsIRRequest) -> Mapping[str, object]: state: Final = _systemone_state(request.input) keyed_questions: Final = zip(systemone_keys(request.questions), request.questions, strict=True) images: Final = input_images(request.input) @@ -503,6 +498,16 @@ def ir_to_systemone_response(response: DecisionsIRResponse, request: DecisionsIR ) +def catalog_mode(model: str, custom_llm_provider: str) -> str | None: + try: + key: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)["key"] + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model, which the catalog does not list + return None + entry: Final = GetModelCostMap.loaded_model_cost_map().get(key) + mode: Final = entry.get("mode") if entry is not None else None + return mode if isinstance(mode, str) else None + + def _is_decisions_model(model: str, custom_llm_provider: str) -> bool: try: model_mode: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get("mode") @@ -511,6 +516,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 +544,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}'" @@ -621,12 +645,16 @@ class BaseDecisionsConfig(ABC): model=model, custom_llm_provider=custom_llm_provider ) - def supports_audio_input(self, model: str, custom_llm_provider: str) -> bool: + def supports_audio_input(self, model: str, custom_llm_provider: str, deployment_flag: bool | None) -> bool: + if deployment_flag is not None: + return deployment_flag return _is_decisions_model(model, custom_llm_provider) and litellm.supports_audio_input( model=model, custom_llm_provider=custom_llm_provider ) - def supports_video_input(self, model: str, custom_llm_provider: str) -> bool: + def supports_video_input(self, model: str, custom_llm_provider: str, deployment_flag: bool | None) -> bool: + if deployment_flag is not None: + return deployment_flag return _is_decisions_model(model, custom_llm_provider) and litellm.supports_video_input( model=model, custom_llm_provider=custom_llm_provider ) @@ -637,12 +665,7 @@ class BaseDecisionsConfig(ABC): request: DecisionsIRRequest, custom_llm_provider: str, ) -> Mapping[str, object] | UnsupportedDecisionsRequest: - return ir_to_systemone_request( - self.request_model(model), - request, - supports_audio=self.supports_audio_input(model, custom_llm_provider), - supports_video=self.supports_video_input(model, custom_llm_provider), - ) + return ir_to_systemone_request(self.request_model(model), request) def unwrap_response(self, payload: object) -> object: return payload diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 3800489ec05..21af97e523e 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -25,6 +25,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.core_helpers import ( as_str_mapping, + drops_unsupported_params, filter_exceptions_from_params, filter_internal_params, map_finish_reason, @@ -2129,6 +2130,24 @@ class AmazonConverseConfig(BaseConfig): return data + @staticmethod + def _claude_messages_within_cache_control_cap( + messages: Sequence[AllMessageValues], + capability_model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> Sequence[AllMessageValues]: + if not drops_unsupported_params(litellm_params.get("drop_params")) or "claude" not in capability_model.lower(): + return messages + from litellm.integrations.anthropic_cache_control_hook import ( + messages_within_cache_control_cap, + tool_message_with_single_cache_mark, + ) + + return messages_within_cache_control_cap( + tuple(map(tool_message_with_single_cache_mark, messages)), optional_params.get("tools") + ) + async def _async_transform_request( self, model: str, @@ -2140,7 +2159,13 @@ class AmazonConverseConfig(BaseConfig): base_model: Final = litellm_params.get("base_model") capability_model: Final = bedrock_capability_model(model, base_model) messages, system_content_blocks = self._transform_system_message( - messages, model=capability_model, deployment_model=model + list( + self._claude_messages_within_cache_control_cap( + messages, capability_model, _JSON_OBJECT.validate_python(optional_params), litellm_params + ) + ), + model=capability_model, + deployment_model=model, ) # Convert last user message to guarded_text if guardrailConfig is present @@ -2221,7 +2246,13 @@ class AmazonConverseConfig(BaseConfig): base_model: Final = litellm_params.get("base_model") capability_model: Final = bedrock_capability_model(model, base_model) messages, system_content_blocks = self._transform_system_message( - messages, model=capability_model, deployment_model=model + list( + self._claude_messages_within_cache_control_cap( + messages, capability_model, _JSON_OBJECT.validate_python(optional_params), litellm_params + ) + ), + model=capability_model, + deployment_model=model, ) # Convert last user message to guarded_text if guardrailConfig is present 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/databricks/embed/handler.py b/litellm/llms/databricks/embed/handler.py index 8b2fa768968..05807a118c3 100644 --- a/litellm/llms/databricks/embed/handler.py +++ b/litellm/llms/databricks/embed/handler.py @@ -26,6 +26,7 @@ class DatabricksEmbeddingHandler(OpenAILikeEmbeddingHandler, DatabricksBase): aembedding=None, custom_endpoint: bool | None = None, headers: dict | None = None, + drop_params: bool = False, ) -> EmbeddingResponse: # Check for custom user agent in optional_params or environment # This allows partners building on LiteLLM to set their own telemetry @@ -58,4 +59,5 @@ class DatabricksEmbeddingHandler(OpenAILikeEmbeddingHandler, DatabricksBase): aembedding=aembedding, custom_endpoint=True, headers=headers, + drop_params=drop_params, ) diff --git a/litellm/llms/deepinfra/chat/transformation.py b/litellm/llms/deepinfra/chat/transformation.py index 004128c0184..7caef58dc79 100644 --- a/litellm/llms/deepinfra/chat/transformation.py +++ b/litellm/llms/deepinfra/chat/transformation.py @@ -130,7 +130,8 @@ class DeepInfraConfig(OpenAIGPTConfig): - Input: {"role": "tool", "content": [{"type": "text", "text": "result"}]} - Output: {"role": "tool", "content": "[{\"type\": \"text\", \"text\": \"result\"}]"} """ - for message in messages: + copied_messages: Final = [message.copy() for message in messages] + for message in copied_messages: if message.get("role") == "tool": content = message.get("content") @@ -149,7 +150,7 @@ class DeepInfraConfig(OpenAIGPTConfig): # For complex content, serialize the entire array as JSON string message["content"] = json.dumps(content) - return messages + return copied_messages @overload def _transform_messages( diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 92c6c8439b3..5a9e5745c40 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -23,6 +23,7 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolParam, + ChatCompletionToolParamFunctionChunk, OpenAIChatCompletionToolParam, ) from litellm.types.utils import ( @@ -45,6 +46,7 @@ from ...openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) +from ...openai.common_utils import without_cache_control from ..common_utils import ( FIREROUTER, FireworksAIException, @@ -480,15 +482,23 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return (("response_format", _json_schema_response_format(choice_schema, "choice")),) def _transform_tools(self, tools: list[OpenAIChatCompletionToolParam]) -> list[OpenAIChatCompletionToolParam]: - for tool in tools: - if tool.get("type") != "function": - continue - function = tool["function"] - function.pop("strict", None) - params = function.get("parameters") - if isinstance(params, dict): - unpack_legacy_defs(params) - return tools + return [self._fireworks_tool(tool) for tool in tools] + + @staticmethod + def _fireworks_tool(tool: OpenAIChatCompletionToolParam) -> OpenAIChatCompletionToolParam: + if tool.get("type") != "function": + return tool + function: Final = tool["function"] + params: Final = function.get("parameters") + unstrict_function: Final = cast( # cast-ok: the same function definition minus `strict` + ChatCompletionToolParamFunctionChunk, {key: value for key, value in function.items() if key != "strict"} + ) + new_function: Final[ChatCompletionToolParamFunctionChunk] = ( + {**unstrict_function, "parameters": unpack_legacy_defs(params, copy=True)} + if isinstance(params, dict) + else unstrict_function + ) + return {**tool, "function": new_function} def _transform_messages_helper( self, messages: list[AllMessageValues], model: str, litellm_params: dict @@ -496,10 +506,6 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): """ Strip fields not permitted by FireworksAI from messages. """ - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - filter_value_from_dict, - ) - supports_vision_value: Final = self._get_model_cost_capability_exact(model=model, capability="supports_vision") for message in messages: if message["role"] == "user": @@ -530,9 +536,11 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): model=model, llm_provider="fireworks_ai", ) - filter_value_from_dict(cast(dict, message), "cache_control") - return [strip_litellm_response_only_fields(message) for message in messages] + return [ + strip_litellm_response_only_fields(cast(AllMessageValues, without_cache_control(message))) + for message in messages + ] # Cached index of fireworks_ai/* entries from litellm.model_cost. Building # this index requires a full scan of model_cost (tens of thousands of diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index 23f68335ea9..fd6b81fce8a 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -140,7 +140,8 @@ class GroqChatConfig(OpenAILikeChatConfig): def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: bool = False ) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]: - for idx, message in enumerate(messages): + transformed_messages: Final = list(messages) + for idx, message in enumerate(transformed_messages): """ 1. Don't pass 'null' function_call assistant message to groq - https://github.com/BerriAI/litellm/issues/5839 """ @@ -154,12 +155,12 @@ class GroqChatConfig(OpenAILikeChatConfig): for k, v in _message.items(): if v is not None: new_message[k] = v - messages[idx] = strip_litellm_internal_message_fields(new_message) + transformed_messages[idx] = strip_litellm_internal_message_fields(new_message) if is_async: - return super()._transform_messages(messages=messages, model=model, is_async=True) + return super()._transform_messages(messages=transformed_messages, model=model, is_async=True) else: - return super()._transform_messages(messages=messages, model=model, is_async=False) + return super()._transform_messages(messages=transformed_messages, model=model, is_async=False) def transform_messages( self, diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 9f9503f93b6..3386cd452da 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -2,6 +2,7 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions` """ +import copy import json from collections.abc import Coroutine from typing import Final, Literal, cast, overload @@ -88,7 +89,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): ) -> dict: _tools = non_default_params.pop("tools", None) if _tools is not None: - _tools = remove_additional_properties(_tools) + _tools = remove_additional_properties(copy.deepcopy(_tools)) # pyright: ignore[reportUnknownArgumentType] # tools arrive untyped through the bare-dict params _tools = remove_strict_from_schema(_tools) if isinstance(_tools, list): _tools = self._convert_custom_tools_to_function_tools(_tools) 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 a402e2c4048..1ecef46ffb1 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -25,16 +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, @@ -755,7 +755,8 @@ class MistralConfig(OpenAIGPTConfig): if _name is not None: # Remove name if not a tool message if message["role"] != "tool" or isinstance(_name, str) and len(_name.strip()) == 0: - message.pop("name", None) + message_without_name: Final = {k: v for k, v in message.items() if k != "name"} + return cast(AllMessageValues, message_without_name) # cast-ok: the same message minus name return message diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 9c827822dc1..224674739f5 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -21,6 +21,7 @@ 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 @@ -54,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, diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 755d647da1d..168b562b358 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -235,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( @@ -248,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/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 22af51aaabe..9d551f6f859 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -14,7 +14,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.constants import OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, map_finish_reason from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( extract_reasoning_content, handle_invalid_parallel_tool_calls, @@ -64,7 +64,7 @@ from litellm.types.utils import ( ) from litellm.utils import convert_to_model_response_object, declared_mid_conversation_system_support -from ..common_utils import OpenAIError +from ..common_utils import OpenAIError, without_cache_control from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config if TYPE_CHECKING: @@ -128,6 +128,29 @@ def _litellm_params_str(litellm_params: Mapping[str, object] | None, key: str) - return value if isinstance(value, str) else None +def _drops_params(litellm_params: Mapping[str, object]) -> bool: + return drops_unsupported_params(litellm_params.get("drop_params")) + + +def _names_claude(model: str, litellm_params: Mapping[str, object]) -> bool: + base_model: Final = litellm_params.get("base_model") + return "claude" in model.lower() or (isinstance(base_model, str) and "claude" in base_model.lower()) + + +def _claude_messages_within_cache_control_cap( + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence["ChatCompletionToolParam"] | None, + litellm_params: Mapping[str, object], + request_params: Mapping[str, object], +) -> Sequence[AllMessageValues]: + if not _names_claude(model, litellm_params) or not _drops_params(litellm_params): + return messages + from litellm.integrations.anthropic_cache_control_hook import messages_within_cache_control_cap + + return list(messages_within_cache_control_cap(messages, tools, request_params.get("cache_control"), request_params)) + + def _prompt_cache_options(optional_params: Mapping[str, object]) -> object | None: try: extra_body: Final = _OBJECT_MAPPING.validate_python(optional_params.get("extra_body")) @@ -311,52 +334,58 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): return content_copy async def _async_handle_pdf_url(self, content_item: ChatCompletionFileObjectFile) -> ChatCompletionFileObjectFile: - file_id: Final = content_item.get("file_id") + content_copy: Final = content_item.copy() + file_id: Final = content_copy.get("file_id") if file_id is not None: # check for file id being url done in _handle_pdf_url base64_data: Final = await async_convert_url_to_base64(file_id) - content_item["file_data"] = base64_data - content_item["filename"] = "my_file.pdf" - content_item.pop("file_id") - return content_item + content_copy["file_data"] = base64_data + content_copy["filename"] = "my_file.pdf" + content_copy.pop("file_id") + return content_copy def _common_file_data_check(self, content_item: ChatCompletionFileObjectFile) -> ChatCompletionFileObjectFile: file_data: Final = content_item.get("file_data") filename: Final = content_item.get("filename") if file_data is not None and filename is None: - content_item["filename"] = "my_file.pdf" + return {**content_item, "filename": "my_file.pdf"} return content_item + @staticmethod + def _image_item_for_openai(content_item: ChatCompletionImageObject) -> ChatCompletionImageObject: + image_url: Final = content_item["image_url"] + if isinstance(image_url, str): + return {**content_item, "image_url": {"url": image_url}} + if not isinstance(image_url, dict): + return content_item + new_image_url_obj: Final = cast( + ChatCompletionImageUrlObject, + {k: v for k, v in image_url.items() if k not in {"format", "video_metadata"}}, + ) + return {**content_item, "image_url": new_image_url_obj} + + @staticmethod + def _file_item_for_openai(content_item: ChatCompletionFileObject) -> ChatCompletionFileObject: + litellm_specific_params: Final = {"format", "detail", "video_metadata"} + file_obj: Final = content_item.get("file") + if file_obj is None: + raise litellm.BadRequestError( + message="Content block has type='file' but is missing the required 'file' field", + model=None, + llm_provider="openai", + ) + new_file_obj: Final = ChatCompletionFileObjectFile( + **{k: v for k, v in file_obj.items() if k not in litellm_specific_params} + ) + return {**content_item, "file": new_file_obj} + def _apply_common_transform_content_item( self, content_item: OpenAIMessageContentListBlock, ) -> OpenAIMessageContentListBlock: - litellm_specific_params: Final = {"format", "detail", "video_metadata"} if content_item.get("type") == "image_url": - content_item = cast(ChatCompletionImageObject, content_item) - if isinstance(content_item["image_url"], str): - content_item["image_url"] = { - "url": content_item["image_url"], - } - elif isinstance(content_item["image_url"], dict): - new_image_url_obj: Final = cast( - ChatCompletionImageUrlObject, - {k: v for k, v in content_item["image_url"].items() if k not in {"format", "video_metadata"}}, - ) - content_item["image_url"] = new_image_url_obj - elif content_item.get("type") == "file": - content_item = cast(ChatCompletionFileObject, content_item) - file_obj: Final = content_item.get("file") - if file_obj is None: - raise litellm.BadRequestError( - message="Content block has type='file' but is missing the required 'file' field", - model=None, - llm_provider="openai", - ) - new_file_obj: Final = ChatCompletionFileObjectFile( - **{k: v for k, v in file_obj.items() if k not in litellm_specific_params} - ) - content_item["file"] = new_file_obj - + return self._image_item_for_openai(cast(ChatCompletionImageObject, content_item)) + if content_item.get("type") == "file": + return self._file_item_for_openai(cast(ChatCompletionFileObject, content_item)) return content_item def _transform_content_item( @@ -419,31 +448,34 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): hoisted_messages: Final = hoist_images_from_tool_messages(stripped_messages) async def _async_transform(): - for message in hoisted_messages: - message_content = message.get("content") - message_role = message.get("role") - - if message_role == "user" and message_content and isinstance(message_content, list): - message_content_types = cast(list[OpenAIMessageContentListBlock], message_content) - for i, content_item in enumerate(message_content_types): - message_content_types[i] = await self._async_transform_content_item( - cast(OpenAIMessageContentListBlock, content_item), - ) - return hoisted_messages + return [await self._async_user_message_for_openai(message) for message in hoisted_messages] if is_async: return _async_transform() - else: - for message in hoisted_messages: - message_content = message.get("content") - message_role = message.get("role") - if message_role == "user" and message_content and isinstance(message_content, list): - message_content_types = cast(list[OpenAIMessageContentListBlock], message_content) - for i, content_item in enumerate(message_content): - message_content_types[i] = self._transform_content_item( - cast(OpenAIMessageContentListBlock, content_item) - ) - return hoisted_messages + return [self._user_message_for_openai(message) for message in hoisted_messages] + + @staticmethod + def _user_content_parts(message: AllMessageValues) -> Sequence[OpenAIMessageContentListBlock] | None: + message_content: Final = message.get("content") + if message.get("role") != "user" or not message_content or not isinstance(message_content, list): + return None + return cast(Sequence[OpenAIMessageContentListBlock], message_content) + + def _user_message_for_openai(self, message: AllMessageValues) -> AllMessageValues: + parts: Final = self._user_content_parts(message) + if parts is None: + return message + return cast( # cast-ok: the same user message with each content part rebuilt for OpenAI + AllMessageValues, {**message, "content": [self._transform_content_item(part) for part in parts]} + ) + + async def _async_user_message_for_openai(self, message: AllMessageValues) -> AllMessageValues: + parts: Final = self._user_content_parts(message) + if parts is None: + return message + return cast( # cast-ok: the same user message with each content part rebuilt for OpenAI + AllMessageValues, {**message, "content": [await self._async_transform_content_item(part) for part in parts]} + ) def transform_messages( self, @@ -462,23 +494,15 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): messages: list[AllMessageValues], tools: list["ChatCompletionToolParam"] | None = None, ) -> tuple[list[AllMessageValues], list["ChatCompletionToolParam"] | None]: - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - filter_value_from_dict, - ) from litellm.types.llms.openai import ChatCompletionToolParam - for i, message in enumerate(messages): - messages[i] = cast( - AllMessageValues, - filter_value_from_dict(message, "cache_control"), - ) - if tools is not None: - for i, tool in enumerate(tools): - tools[i] = cast( - ChatCompletionToolParam, - filter_value_from_dict(tool, "cache_control"), - ) - return messages, tools + new_messages: Final = [cast(AllMessageValues, without_cache_control(message)) for message in messages] + new_tools: Final = ( + [cast(ChatCompletionToolParam, without_cache_control(tool)) for tool in tools] + if tools is not None + else None + ) + return new_messages, new_tools def translate_developer_role_to_system_role( self, @@ -527,6 +551,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): tools: Sequence["ChatCompletionToolParam"] | None, prompt_cache_options: object, litellm_params: Mapping[str, object], + request_params: Mapping[str, object], ) -> tuple[Sequence[AllMessageValues], Sequence["ChatCompletionToolParam"] | None]: """ A deployment that reads OpenAI's ``prompt_cache_breakpoint`` (gpt-5.6+ on @@ -543,10 +568,11 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): model, custom_llm_provider, api_base, prompt_cache_options ) if not reads_breakpoints and self._should_preserve_cache_control_for_endpoint(custom_llm_provider, api_base): - return messages, tools - drops_breakpoints: Final = not reads_breakpoints and ( - litellm.drop_params is True or litellm_params.get("drop_params") is True - ) + return ( + _claude_messages_within_cache_control_cap(model, messages, tools, litellm_params, request_params), + tools, + ) + drops_breakpoints: Final = not reads_breakpoints and _drops_params(litellm_params) renamed: Final = ( [AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(message) for message in messages] if reads_breakpoints @@ -561,7 +587,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): if drops_breakpoints else list(tools) ) - return self.remove_cache_control_flag_from_messages_and_tools(model=model, messages=renamed, tools=kept_tools) + sent_messages, sent_tools = self.remove_cache_control_flag_from_messages_and_tools( + model=model, messages=renamed, tools=kept_tools + ) + return ( + _claude_messages_within_cache_control_cap(model, sent_messages, sent_tools, litellm_params, request_params), + sent_tools, + ) def _sanitized_tools_update_for_openai( self, @@ -623,6 +655,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): tools=optional_params.get("tools", []), prompt_cache_options=_prompt_cache_options(optional_params), litellm_params=litellm_params, + request_params=optional_params, ) if tools is not None and len(tools) > 0: optional_params["tools"] = tools @@ -654,6 +687,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): tools=optional_params.get("tools", []), prompt_cache_options=_prompt_cache_options(optional_params), litellm_params=litellm_params, + request_params=optional_params, ) if tools is not None and len(tools) > 0: optional_params["tools"] = tools diff --git a/litellm/llms/openai/chat/o_series_transformation.py b/litellm/llms/openai/chat/o_series_transformation.py index 64285aafdb0..4d56a3b38c5 100644 --- a/litellm/llms/openai/chat/o_series_transformation.py +++ b/litellm/llms/openai/chat/o_series_transformation.py @@ -153,12 +153,13 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): - role: system ==> translate to role 'user' """ _supports_system_messages: Final = supports_system_messages(model, "openai") - for i, message in enumerate(messages): + transformed_messages: Final = list(messages) + for i, message in enumerate(transformed_messages): if message["role"] == "system" and not _supports_system_messages: new_message = ChatCompletionUserMessage(content=message["content"], role="user") - messages[i] = new_message # Replace the old message with the new one + transformed_messages[i] = new_message # Replace the old message with the new one if is_async: - return super()._transform_messages(messages, model, is_async=cast(Literal[True], True)) + return super()._transform_messages(transformed_messages, model, is_async=cast(Literal[True], True)) else: - return super()._transform_messages(messages, model, is_async=cast(Literal[False], False)) + return super()._transform_messages(transformed_messages, model, is_async=cast(Literal[False], False)) diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index c113cb351d8..d476cf62626 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -3,6 +3,7 @@ Common helpers / utils across al OpenAI endpoints """ import asyncio +import copy import hashlib import inspect import json @@ -31,7 +32,7 @@ if TYPE_CHECKING: 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.core_helpers import as_str_mapping, decoded_json_object, iter_items 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 @@ -43,6 +44,47 @@ from litellm.llms.custom_httpx.http_handler import ( ) +def _is_plain_dict(value: object) -> bool: + return isinstance(value, dict) + + +def _is_list(value: object) -> bool: + return isinstance(value, list) + + +def _contains_cache_control(value: object) -> bool: + pending: Final[list[object]] = [value] # mutable-ok: a local worklist, never returned or stored + while pending: + node = pending.pop() + mapping = as_str_mapping(node) + if mapping is not None: + if "cache_control" in mapping: + return True + pending.extend(mapping.values()) + elif _is_list(node): + pending.extend(iter_items(node)) + return False + + +def without_cache_control(item: object) -> object: + """Return `item` with every `cache_control` key removed, without touching the original. + + `filter_value_from_dict` deletes the key in place and recurses into nested dicts and + lists, so the caller's own message, tool or input object has to be copied first: it + may reuse the same list on a provider that does support prompt caching. The copy is + skipped when there is nothing to strip, which keeps the object identity the + pass-through paths rely on and avoids a deep copy on every request. + """ + from litellm.litellm_core_utils.prompt_templates.common_utils import filter_value_from_dict + + if not _is_plain_dict(item) or not _contains_cache_control(item): + return item + copied: Final = copy.deepcopy(item) + if isinstance(copied, dict): + filter_value_from_dict(copied, "cache_control") + return copied + + def _get_client_init_params(cls: type) -> tuple[str, ...]: """Extract __init__ parameter names (excluding 'self') from a class.""" return tuple(p for p in inspect.signature(cls.__init__).parameters if p != "self") diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 4518fe6a77e..369b6959411 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -354,10 +354,10 @@ class OpenAIDecisionsConfig(BaseDecisionsConfig): ) -> str | None: return resolve_openai_bearer_token(api_key=api_key, api_base=api_base, litellm_params=litellm_params) - def supports_audio_input(self, model: str, custom_llm_provider: str) -> bool: + def supports_audio_input(self, model: str, custom_llm_provider: str, deployment_flag: bool | None) -> bool: return False - def supports_video_input(self, model: str, custom_llm_provider: str) -> bool: + def supports_video_input(self, model: str, custom_llm_provider: str, deployment_flag: bool | None) -> bool: return False def transform_decisions_request( diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index bdd9f285823..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, @@ -2182,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 @@ -2372,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, @@ -2384,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( @@ -2409,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/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 9b8515a9ecd..7cca31c5d3f 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -52,7 +52,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS -from ..common_utils import OpenAIError +from ..common_utils import OpenAIError, without_cache_control from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS: Final = 16 @@ -658,13 +658,6 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): input=self._replayable_input(marker_safe_input, reads_breakpoints, maps_anthropic_tool_search), tools=mapped_tools, ) - if marker_safe_input is not validated_input: - self._strip_caller_input_as_if_unmarked( - model=model, - caller_input=validated_input, - reads_breakpoints=reads_breakpoints, - maps_anthropic_tool_search=maps_anthropic_tool_search, - ) object_schema_tools: Final = self._tools_with_object_parameters(model=model, tools=stripped_tools) sanitized_tools: Final = self._sanitized_tool_schemas_for_openai( model=model, tools=object_schema_tools, litellm_params=litellm_params @@ -679,17 +672,6 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return marked_input return input_without_unanswered_anthropic_tool_search_calls(marked_input) - def _strip_caller_input_as_if_unmarked( - self, - model: str, - caller_input: str | ResponseInputParam, - reads_breakpoints: bool, - maps_anthropic_tool_search: bool, - ) -> None: - self.remove_cache_control_flag_from_input_and_tools( - model=model, input=self._replayable_input(caller_input, reads_breakpoints, maps_anthropic_tool_search) - ) - def _input_with_wire_safe_prompt_cache_breakpoints( self, model: str, @@ -800,24 +782,29 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): the chat path. Strips Anthropic-only `cache_control` markers from Responses API input content blocks and tools. - `filter_value_from_dict` mutates each dict in place, so the same - objects are returned. + `filter_value_from_dict` deletes the key in place and recurses, so each + item is copied first: the caller keeps its own input list and may reuse + it on a provider that does support prompt caching. """ - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - filter_value_from_dict, + new_input: Final = ( + cast( # cast-ok: the comprehension rebuilds the input item for item + ResponseInputParam, + [without_cache_control(item) for item in input], + ) + if isinstance(input, list) + else input ) - if isinstance(input, list): - for item in input: - if isinstance(item, dict): - filter_value_from_dict(cast(dict, item), "cache_control") + new_tools: Final = ( + cast( # cast-ok: the comprehension rebuilds the tool list item for item + Sequence[ALL_RESPONSES_API_TOOL_PARAMS], + [without_cache_control(tool) for tool in tools], + ) + if tools is not None + else None + ) - if tools is not None: - for tool in tools: - if isinstance(tool, dict): - filter_value_from_dict(cast(dict, tool), "cache_control") - - return input, tools + return new_input, new_tools def _drop_bridge_minted_reasoning_items(self, input: str | ResponseInputParam) -> str | ResponseInputParam: if self.custom_llm_provider not in _PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING or not isinstance(input, list): diff --git a/litellm/llms/openai_like/README.md b/litellm/llms/openai_like/README.md index c56f58b24c9..fb00b07c986 100644 --- a/litellm/llms/openai_like/README.md +++ b/litellm/llms/openai_like/README.md @@ -125,6 +125,20 @@ The provider inherits all request/response handling from OpenAI's Responses API If `supported_endpoints` is omitted, it defaults to `[]` (only chat completions, which is always enabled for JSON providers). +When only some models serve an endpoint natively, list them under `model_supported_endpoints`. Every other model reaches `/v1/responses` and `/v1/messages` through chat completions: + +```json +{ + "your_provider": { + "supported_endpoints": ["/v1/chat/completions"], + "model_supported_endpoints": { + "model-a": ["/v1/messages"], + "model-b": ["/v1/responses"] + } + } +} +``` + ### How It Works 1. `json_loader.py` checks `supported_endpoints` for `/v1/responses` diff --git a/litellm/llms/openai_like/embedding/handler.py b/litellm/llms/openai_like/embedding/handler.py index 77b6b673707..540feeca6e7 100644 --- a/litellm/llms/openai_like/embedding/handler.py +++ b/litellm/llms/openai_like/embedding/handler.py @@ -3,11 +3,16 @@ ## Allows jina ai embedding calls - which don't allow 'encoding_format' in payload. import json +import re +from collections.abc import Mapping +from types import MappingProxyType from typing import Final import httpx +from pydantic import TypeAdapter import litellm +from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -17,6 +22,57 @@ from litellm.types.utils import EmbeddingResponse from ..common_utils import OpenAILikeBase, OpenAILikeError +_KEY_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[str, str]) +_EXTRA_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_TRANSPORT_HEADER_NAMES: Final = frozenset( + { + "host", + "content-length", + "transfer-encoding", + "connection", + "keep-alive", + "proxy-connection", + "te", + "trailer", + "upgrade", + "expect", + } +) +_HEADER_NAME: Final = re.compile(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]+") +_HEADER_VALUE: Final = re.compile(r"[\t\x20-\x7e]*") + + +def _header_text(name: str, value: object) -> str | None: + if isinstance(value, bool) or not isinstance(value, str | int | float): + return None + if _HEADER_NAME.fullmatch(name) is None: + return None + return str(value).strip(" \t\r\n") + + +def _extra_headers_as_text(extra_headers: Mapping[str, object], model: str, drop_params: bool) -> Mapping[str, str]: + caller_set: Final = { + name: value for name, value in extra_headers.items() if name.lower() not in _TRANSPORT_HEADER_NAMES + } + texts: Final = MappingProxyType({name: _header_text(name, value) for name, value in caller_set.items()}) + refused: Final = sorted(name for name, text in texts.items() if text is None) + if refused and not drop_params: + raise litellm.UnsupportedParamsError( + message=f"extra_headers {', '.join(map(repr, refused))} cannot be sent. A header needs a name made of " + "letters, digits or !#$%&'*+-.^_`|~ and a text or number value. " + "To drop them from the call, set `litellm.drop_params = True`.", + model=model, + ) + named: Final = MappingProxyType({name: text for name, text in texts.items() if text is not None}) + unsendable: Final = sorted(name for name, text in named.items() if _HEADER_VALUE.fullmatch(text) is None) + if unsendable: + verbose_logger.warning( + "extra_headers %s dropped for model %r, since a header value must be printable ASCII on one line", + unsendable, + model, + ) + return MappingProxyType({name: text for name, text in named.items() if name not in unsendable}) + class OpenAILikeEmbeddingHandler(OpenAILikeBase): def __init__(self, **kwargs): @@ -94,6 +150,7 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): aembedding=None, custom_endpoint: bool | None = None, headers: dict | None = None, + drop_params: bool = False, ) -> EmbeddingResponse: api_base, headers = self._validate_environment( api_base=api_base, @@ -103,8 +160,20 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): custom_endpoint=custom_endpoint, ) model = model - filtered_optional_params: Final = {k: v for k, v in optional_params.items() if v not in (None, "")} - data: Final = {"model": model, "input": input, **filtered_optional_params} + data: Final = { + "model": model, + "input": input, + **{k: v for k, v in optional_params.items() if k != "extra_headers" and v not in (None, "")}, + } + extra_headers: Final = _extra_headers_as_text( + _EXTRA_HEADERS_ADAPTER.validate_python(optional_params.get("extra_headers") or {}), model, drop_params + ) + key_headers: Final = _KEY_HEADERS_ADAPTER.validate_python(headers) + extra_header_names: Final = frozenset(name.lower() for name in extra_headers) + request_headers: Final = { + **{name: value for name, value in key_headers.items() if name.lower() not in extra_header_names}, + **extra_headers, + } ## LOGGING logging_obj.pre_call( @@ -123,7 +192,7 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): api_key=api_key, timeout=timeout, client=client, - headers=headers, + headers=request_headers, ) if client is None or isinstance(client, AsyncHTTPHandler): self.client = HTTPHandler(timeout=timeout) @@ -134,7 +203,7 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): try: response: Final = self.client.post( api_base, - headers=headers, + headers=request_headers, data=json.dumps(data), ) diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index 6e99b9e7fec..658e4176f22 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -3,7 +3,7 @@ JSON-based provider configuration loader for OpenAI-compatible providers. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from pathlib import Path from typing import Final @@ -14,6 +14,7 @@ from litellm.types.utils import ProviderSpecificModelInfo _PROVIDER_INFO_ADAPTER: Final = TypeAdapter(ProviderSpecificModelInfo) _SPECIAL_HANDLING_ADAPTER: Final = TypeAdapter(Mapping[str, bool]) +_MODEL_ENDPOINTS_ADAPTER: Final = TypeAdapter(Mapping[str, tuple[str, ...]]) class SimpleProviderConfig: @@ -30,7 +31,10 @@ class SimpleProviderConfig: self.special_handling: Mapping[str, bool] = _SPECIAL_HANDLING_ADAPTER.validate_python( data.get("special_handling", {}) ) - self.supported_endpoints = data.get("supported_endpoints", []) + self.supported_endpoints: Sequence[str] = data.get("supported_endpoints", []) + self.model_supported_endpoints: Mapping[str, tuple[str, ...]] = _MODEL_ENDPOINTS_ADAPTER.validate_python( + data.get("model_supported_endpoints", {}) + ) self.provider_info: ProviderSpecificModelInfo = ( _PROVIDER_INFO_ADAPTER.validate_python(data["provider_info"]) if "provider_info" in data else {} ) @@ -90,6 +94,23 @@ class JSONProviderRegistry: return False return "/v1/responses" in provider.supported_endpoints + @classmethod + def supports_embeddings(cls, slug: str) -> bool: + provider: Final = cls._providers.get(slug) + if provider is None: + return False + return "/v1/embeddings" in provider.supported_endpoints + + @classmethod + def serves_endpoint(cls, slug: str, model: str | None, endpoint: str) -> bool: + provider: Final = cls._providers.get(slug) + if provider is None: + return False + model_endpoints: Final = ( + provider.model_supported_endpoints.get(model.removeprefix(f"{slug}/"), ()) if model else () + ) + return endpoint in provider.supported_endpoints or endpoint in model_endpoints + @classmethod def list_providers(cls) -> list: """List all registered provider slugs""" diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index ad47b334ed4..a71a9c3bb9b 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -211,7 +211,7 @@ "param_mappings": { "max_completion_tokens": "max_tokens" }, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/embeddings"] + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] }, "scx-ai": { "base_url": "https://api.scx.ai/v1", @@ -231,6 +231,16 @@ "api_base_env": "PRISM_API_BASE", "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] }, + "hostinger_router": { + "base_url": "https://router.hostinger.com/v1", + "api_key_env": "HOSTINGER_ROUTER_API_KEY", + "api_base_env": "HOSTINGER_ROUTER_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/embeddings"], + "model_supported_endpoints": { + "claude-sonnet-5": ["/v1/messages"], + "gpt-5.2": ["/v1/responses"] + } + }, "sail": { "base_url": "https://api.sailresearch.com/v1", "api_key_env": "SAIL_API_KEY", diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index cca6815399f..69f3d76c034 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -14,7 +14,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.core_helpers import finite_cost -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_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 @@ -212,19 +212,27 @@ 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"] ) reported_cost: Final = ( finite_cost(float(response_cost)) if isinstance(response_cost, str | int | float) else None @@ -282,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"], @@ -290,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}", @@ -298,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/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index c3d178771a5..c40a86b221b 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -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/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/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index f683ad0e377..a4601af944b 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -110,38 +110,23 @@ def separate_cached_messages( messages: list[AllMessageValues], ) -> tuple[list[AllMessageValues], list[AllMessageValues]]: """ - Returns separated cached and non-cached messages. - - Args: - messages: List of messages to be separated. - - Returns: - Tuple containing: - - cached_messages: List of cached messages. - - non_cached_messages: List of non-cached messages. + Returns the first continuous block of cached messages, plus every system message once + any message is cached, since generateContent rejects a system_instruction beside a + cachedContent, and the remaining messages. """ - cached_messages: list[AllMessageValues] = [] - non_cached_messages: list[AllMessageValues] = [] + filtered_messages: Final = [(idx, message) for idx, message in enumerate(messages) if is_cached_message(message)] + if not filtered_messages: + return [], messages - # Extract cached messages and their indices - filtered_messages: Final[list[tuple[int, AllMessageValues]]] = [] - for idx, message in enumerate(messages): - if is_cached_message(message=message): - filtered_messages.append((idx, message)) - - # Validate only one block of continuous cached messages last_continuous_block_idx: Final = get_first_continuous_block_idx(filtered_messages) - # Separate messages based on the block of cached messages - if filtered_messages and last_continuous_block_idx is not None: - first_cached_idx: Final = filtered_messages[0][0] - last_cached_idx: Final = filtered_messages[last_continuous_block_idx][0] - - cached_messages = messages[first_cached_idx : last_cached_idx + 1] - non_cached_messages = messages[:first_cached_idx] + messages[last_cached_idx + 1 :] - else: - non_cached_messages = messages - - return cached_messages, non_cached_messages + cached_block: Final = range(filtered_messages[0][0], filtered_messages[last_continuous_block_idx][0] + 1) + goes_in_cache: Final = tuple( + idx in cached_block or message.get("role") == "system" for idx, message in enumerate(messages) + ) + return ( + [message for message, cached in zip(messages, goes_in_cache) if cached], + [message for message, cached in zip(messages, goes_in_cache) if not cached], + ) def cached_messages_end_on_supported_turn(cached_messages: Sequence[AllMessageValues]) -> bool: diff --git a/litellm/llms/watsonx/chat/transformation.py b/litellm/llms/watsonx/chat/transformation.py index 071ff7108b9..25981b0f04c 100644 --- a/litellm/llms/watsonx/chat/transformation.py +++ b/litellm/llms/watsonx/chat/transformation.py @@ -4,6 +4,7 @@ Translation from OpenAI's `/chat/completions` endpoint to IBM WatsonX's `/text/c Docs: https://cloud.ibm.com/apidocs/watsonx-ai#text-chat """ +import copy from typing import Final from litellm import verbose_logger @@ -57,7 +58,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): _tools = non_default_params.pop("tools", None) if _tools is not None: # remove 'additionalProperties' from tools - _tools = remove_additional_properties(_tools) + _tools = remove_additional_properties(copy.deepcopy(_tools)) # pyright: ignore[reportUnknownArgumentType] # tools arrive untyped through the bare-dict params # remove 'strict' from tools _tools = remove_strict_from_schema(_tools) if _tools is not None: diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 2b943fa4107..f53b71a5a9a 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,3 +1,4 @@ +import copy from collections.abc import AsyncIterator, Iterator, Mapping from types import MappingProxyType from typing import Any, Final @@ -208,7 +209,8 @@ class XAIChatConfig(OpenAIGPTConfig): tools = [ filtered_tool for filtered_tool in ( - _FILTERED_TOOL.validate_python(filter_value_from_dict(tool, "strict")) for tool in value + _FILTERED_TOOL.validate_python(filter_value_from_dict(copy.deepcopy(tool), "strict")) # pyright: ignore[reportUnknownArgumentType] # tools arrive untyped through the bare-dict params + for tool in value ) if filtered_tool is not None ] diff --git a/litellm/main.py b/litellm/main.py index 5346741e4d1..21ae9657d68 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 ( @@ -6809,6 +6809,23 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: ) +def _embedding_headers(param: str, configured: object, model: str, drop_params: bool) -> Mapping[str, object]: + if not configured: + return {} + if isinstance(configured, Mapping): + return dict(cast(Mapping[str, object], configured)) # cast-ok: header names from untyped request JSON + try: + return dict(cast(Iterable[tuple[str, object]], configured)) # cast-ok: name and value pairs from request JSON + except (TypeError, ValueError): + if drop_params: + return {} + raise litellm.UnsupportedParamsError( + message=f"{param} must be a JSON object of header names to values. " + "To drop it from the call, set `litellm.drop_params = True`.", + model=model, + ) + + # fmt: off # Overload for when aembedding=True (returns coroutine) @@ -6923,11 +6940,10 @@ def embedding( azure_ad_token_provider: Final = kwargs.get("azure_ad_token_provider", None) aembedding: Final[bool | None] = kwargs.get("aembedding", None) extra_headers: Final = kwargs.get("extra_headers", None) - headers = kwargs.get("headers", None) or extra_headers - if headers is None: - headers = {} - if extra_headers is not None: - headers.update(extra_headers) + headers = { + **_embedding_headers("headers", kwargs.get("headers"), model, drop_unsupported_params), + **_embedding_headers("extra_headers", extra_headers, model, drop_unsupported_params), + } # Inject proxy auth headers if configured if litellm.proxy_auth is not None: try: @@ -7091,7 +7107,7 @@ def embedding( # set API KEY api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") - if headers is not None and headers != {}: + if headers: optional_params["extra_headers"] = headers requested_encoding_format: Final = ( @@ -7183,14 +7199,15 @@ def embedding( custom_llm_provider == "openai_like" or custom_llm_provider == "llamafile" or custom_llm_provider == "lm_studio" + or JSONProviderRegistry.supports_embeddings(custom_llm_provider) ): api_base = api_base or litellm.api_base or get_secret_str("OPENAI_LIKE_API_BASE") # set API KEY - if api_key is None: - api_key = api_key or litellm.api_key or litellm.openai_like_key or get_secret_str("OPENAI_LIKE_API_KEY") + if api_key is None and not JSONProviderRegistry.supports_embeddings(custom_llm_provider): + api_key = litellm.api_key or litellm.openai_like_key or get_secret_str("OPENAI_LIKE_API_KEY") - if headers is not None and headers != {}: + if headers: optional_params["extra_headers"] = headers ## EMBEDDING CALL @@ -7205,10 +7222,9 @@ def embedding( optional_params=optional_params, client=client, aembedding=aembedding, + drop_params=drop_unsupported_params, ) elif custom_llm_provider == "oci": - if headers is None: - headers = {} response = base_llm_http_handler.embedding( model=model, input=input, @@ -7233,11 +7249,6 @@ def embedding( or litellm.api_key ) - # Use the merged headers variable (already merged at the top of the function) - # Don't overwrite it with just extra_headers - if headers is None: - headers = {} - response = base_llm_http_handler.embedding( model=model, input=input, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dac7af6ab7a..7eede4a395f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23078,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", @@ -25284,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", @@ -29561,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, @@ -30167,6 +30218,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_pdf_input": true, @@ -30210,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, @@ -30274,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, @@ -30334,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, @@ -30401,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, @@ -31218,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": [ @@ -31398,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, @@ -32405,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, @@ -32479,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, @@ -32543,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, @@ -32603,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, @@ -32670,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, @@ -39622,6 +39687,57 @@ "/v1/images/generations" ] }, + "hostinger_router/gpt-5.2": { + "cache_read_input_token_cost": 0.00000021, + "input_cost_per_token": 0.0000021, + "litellm_provider": "hostinger_router", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0000168, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_vision": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "source": "https://router.hostinger.com/v1/models" + }, + "hostinger_router/claude-sonnet-5": { + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_1hr": 5.28e-06, + "cache_read_input_token_cost": 2.64e-07, + "input_cost_per_token": 2.64e-06, + "litellm_provider": "hostinger_router", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.32e-05, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_response_schema": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/messages"], + "source": "https://router.hostinger.com/v1/models" + }, + "hostinger_router/text-embedding-3-large": { + "input_cost_per_token": 1.716e-07, + "litellm_provider": "hostinger_router", + "max_input_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0, + "source": "https://router.hostinger.com/v1/models" + }, + "hostinger_router/text-embedding-3-small": { + "input_cost_per_token": 3e-08, + "litellm_provider": "hostinger_router", + "max_input_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0, + "source": "https://router.hostinger.com/v1/models" + }, "hyperbolic/NousResearch/Hermes-3-Llama-3.1-70B": { "input_cost_per_token": 1.2e-07, "litellm_provider": "hyperbolic", @@ -47624,6 +47740,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, @@ -51304,6 +51438,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, @@ -51332,6 +51471,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, @@ -52078,6 +52222,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, @@ -52183,6 +52332,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, @@ -52713,6 +52867,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, @@ -59714,12 +59869,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": [ @@ -75398,6 +75554,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, @@ -76498,6 +76655,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, @@ -76543,6 +76701,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, @@ -76654,6 +76813,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": 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/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index d33e61026df..700a35edb06 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1087,6 +1087,32 @@ "generateContent": false } }, + "hostinger_router": { + "display_name": "Hostinger Router (`hostinger_router`)", + "url": "https://docs.litellm.ai/docs/providers/hostinger_router", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": true, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } + }, "crusoe": { "display_name": "Crusoe (`crusoe`)", "url": "https://docs.litellm.ai/docs/providers/crusoe", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index dc8b9415020..41ab7fcaf88 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -58015,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": [ { @@ -58054,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": [ { @@ -58093,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": [ { @@ -58132,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": [ { @@ -58171,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": [ { @@ -58210,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": [ { @@ -58249,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 79a46ef355e..9c0bb93ea16 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -167,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 @@ -525,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). @@ -541,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). @@ -667,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. @@ -937,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, @@ -1485,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) @@ -4837,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: @@ -4851,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, @@ -4902,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, ...]: @@ -4920,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, @@ -5047,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: @@ -5534,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]: """ @@ -5613,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, @@ -5761,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, @@ -5864,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: """ @@ -5900,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]: """ @@ -5921,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 @@ -6084,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, @@ -6122,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, @@ -6388,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, @@ -7180,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 762f552dadc..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, @@ -81,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, @@ -146,12 +148,10 @@ from litellm.responses.streaming_iterator import partial_responses_incomplete_ev 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") @@ -894,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) @@ -2495,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) ) ) @@ -2800,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 @@ -4159,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 @@ -5003,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 a9424018d9a..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 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/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_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 8545e0609db..333687b07cd 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -26,7 +26,12 @@ 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, @@ -94,11 +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.router_utils.common_utils import resolve_model_group_alias 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 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 0d333d4fd08..2951bc02d8c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,6 +12,7 @@ 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 ( @@ -134,6 +135,8 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) ) +FAILURE_SPEND_COUNTER_WAIT_SECONDS: Final = 1.0 + _OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) @@ -235,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] @@ -361,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 @@ -590,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( @@ -623,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, @@ -637,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, @@ -1060,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) @@ -1095,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: ... @@ -1111,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, @@ -1174,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/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/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index dbd77e45d57..2c7453ebb1c 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -1515,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: 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 287475c868b..2d8f9448f3c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -52,7 +52,11 @@ from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.fallback_generalizations import match_routing_generalization -from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers +from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, + anthropic_pass_through_api_base, + merge_anthropic_beta_headers, +) from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, is_azure_body_model_inference_endpoint, @@ -1048,7 +1052,7 @@ async def anthropic_proxy_route( """ [Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion) """ - base_target_url = os.getenv("ANTHROPIC_API_BASE") or os.getenv("ANTHROPIC_BASE_URL") or "https://api.anthropic.com" + base_target_url = anthropic_pass_through_api_base() encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 70d677dfe4c..07353793b86 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -72,7 +72,11 @@ from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.llms.anthropic.common_utils import is_anthropic_messages_url +from litellm.llms.anthropic.common_utils import ( + anthropic_accepts_user_id, + anthropic_upstream_user_id_limit, + is_anthropic_messages_url, +) from litellm.llms.base_llm.managed_resources.utils import ( resolve_passthrough_managed_id_provider, ) @@ -151,6 +155,8 @@ 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, CallTypes, Usage @@ -505,7 +511,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): url: httpx.URL, headers: dict, requested_query_params: dict | None = None, - _parsed_body: dict | None = None, + _parsed_body: Mapping[str, object] | None = None, forward_multipart: bool = False, ) -> httpx.Response: """ @@ -719,8 +725,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): # unenforced. Set after the client merge so a request body cannot supply its own budget. # # Only for the built-in provider routes. `get_model_from_request` returns - # None for a user-defined pass-through, deliberately: its body is forwarded - # verbatim, so `model` there names an UPSTREAM model rather than a + # None for a user-defined pass-through, deliberately: the client's `model` + # goes upstream unchanged, so it names an UPSTREAM model rather than a # LiteLLM-managed one. Enforcement is therefore skipped on those routes, and # charging a counter anyway would track spend that nothing can refuse, and # would attribute it to a budget the operator scoped to a LiteLLM model that @@ -1227,6 +1233,36 @@ def _unfiltered_headers(headers: Mapping[str, str]) -> Mapping[str, str]: return headers +def _json_object(value: object) -> Mapping[str, object]: + return cast("Mapping[str, object]", value) if isinstance(value, Mapping) else MappingProxyType({}) # cast-ok: json + + +def _anthropic_user_id_to_forward( + body: Mapping[str, object], + url: str, + has_pre_call_guardrails: Callable[[Mapping[str, object]], bool], +) -> str | None: + user_id: Final = _json_object(body.get("metadata")).get("user_id") + if not isinstance(user_id, str): + return None + max_length: Final = anthropic_upstream_user_id_limit(url) + if max_length is None: + return None + if not anthropic_accepts_user_id(user_id, max_length): + verbose_proxy_logger.warning("Not forwarding metadata.user_id because the upstream would refuse it") + return None + if has_pre_call_guardrails(body): + verbose_proxy_logger.debug("Not forwarding metadata.user_id because a pre-call guardrail cannot scan it") + return None + return user_id + + +def _with_anthropic_user_id(body: Mapping[str, object] | None, user_id: str | None) -> Mapping[str, object] | None: + if body is None or user_id is None: + return body + return {**body, "metadata": {"user_id": user_id}} + + def _eu_jev_client_body(custom_llm_provider: str | None, body: Mapping[str, object]) -> Mapping[str, object] | None: return MappingProxyType(dict(body)) if custom_llm_provider == "eu_jev" else None @@ -1473,6 +1509,10 @@ async def pass_through_request( request_method=getattr(request, "method", None), cost_per_request=cost_per_request, ) + hooked_body: Final = _json_object(_parsed_body) + client_user_id: Final = _anthropic_user_id_to_forward( + hooked_body, str(url), proxy_logging_obj.request_has_pre_call_guardrails + ) kwargs = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( # rebind-ok: pre-existing rebinding on a rename-only line user_api_key_dict=user_api_key_dict, _parsed_body=_parsed_body, @@ -1640,11 +1680,13 @@ async def pass_through_request( else: logging_url = str(url) + "?" + requested_query_params_str + outgoing_body: Final = cast("Mapping[str, object] | None", _parsed_body) # cast-ok: json + upstream_body: Final = _with_anthropic_user_id(outgoing_body, client_user_id) logging_obj.pre_call( input=[{"role": "user", "content": safe_dumps(_parsed_body)}], api_key="", additional_args={ - "complete_input_dict": _parsed_body, + "complete_input_dict": upstream_body, "api_base": str(logging_url), "headers": redact_credential_headers(upstream_headers), }, @@ -1684,7 +1726,7 @@ async def pass_through_request( url, params=requested_query_params, headers=upstream_headers, - json=_parsed_body, + json=upstream_body, ) ) @@ -1783,7 +1825,7 @@ async def pass_through_request( url=url, headers=upstream_headers, requested_query_params=requested_query_params, - _parsed_body=_parsed_body, + _parsed_body=upstream_body, forward_multipart=is_multipart, ) verbose_proxy_logger.debug("response.headers= %s", response.headers) @@ -2649,6 +2691,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, @@ -2661,6 +2710,7 @@ async def websocket_passthrough_request( 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. @@ -2676,6 +2726,7 @@ async def websocket_passthrough_request( 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 @@ -2686,6 +2737,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) @@ -2804,6 +2856,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() @@ -2873,6 +2926,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: @@ -2999,6 +3054,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( 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..044b8a1de6e 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", @@ -3315,6 +3343,34 @@ ], "default_model_placeholder": "prism/deepseek-v4.1-flash" }, + { + "provider": "HOSTINGER_ROUTER", + "provider_display_name": "Hostinger Router", + "litellm_provider": "hostinger_router", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://router.hostinger.com/v1", + "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": "hostinger_router/claude-sonnet-5" + }, { "provider": "RECRAFT", "provider_display_name": "Recraft", @@ -4257,6 +4313,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/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index f8d5cba4229..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 ), 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 ddaa8d2e69d..c7efbf9dacf 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -147,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 @@ -159,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 ( @@ -277,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 @@ -633,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 @@ -711,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( @@ -2708,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, ) @@ -2733,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( @@ -2766,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 @@ -2782,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, @@ -2890,14 +3031,21 @@ class ProxyLogging: hook to consult, but it judges the payload the same way, so a proxy configured only with one of those still has something to say about every record. """ - if request_metadata.get("_guardrail_pipelines"): + return self.request_has_pre_call_guardrails({"metadata": dict(request_metadata)}, event_type) + + def request_has_pre_call_guardrails( + self, + request_data: Mapping[str, object], + event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call, + ) -> bool: + if _policy_state_metadata(request_data).get("_guardrail_pipelines"): return True caps: Final = ProxyLogging.callback_capabilities() if caps.has_content_enforcer: return True - probe: Final = {"metadata": dict(request_metadata)} return any( - isinstance(callback, CustomGuardrail) and callback.should_run_guardrail(data=probe, event_type=event_type) + isinstance(callback, CustomGuardrail) + and callback.should_run_guardrail(data=request_data, event_type=event_type) for callback in caps.resolved_callbacks ) @@ -3024,7 +3172,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, ...]] = ( @@ -3128,17 +3276,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( @@ -3504,6 +3652,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: @@ -3515,6 +3664,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 ): @@ -4731,7 +4881,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", @@ -4742,6 +4893,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 @@ -5192,6 +5350,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): @@ -5276,6 +5436,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] = [] @@ -5377,6 +5538,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 @@ -5418,6 +5582,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 @@ -6626,6 +6791,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 @@ -6644,6 +6817,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. @@ -6737,7 +6958,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, @@ -6757,6 +6978,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 @@ -6818,7 +7042,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, @@ -6880,7 +7104,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, @@ -6953,7 +7177,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: @@ -6980,7 +7204,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.""" @@ -7002,7 +7226,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, @@ -7209,6 +7433,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 @@ -7399,7 +7625,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: @@ -7433,7 +7659,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()) @@ -9098,8 +9324,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 @@ -9500,7 +9725,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 @@ -9521,17 +9746,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, + ) ) ) @@ -9576,6 +9804,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 @@ -9656,18 +9885,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 8f0a85d67de..2e156d5abba 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -8,7 +8,7 @@ import uuid from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial -from itertools import accumulate, chain, groupby, islice +from itertools import chain, groupby from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -45,7 +45,15 @@ from litellm.constants import ( REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER, ) -from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook +from litellm.integrations.anthropic_cache_control_hook import ( + AnthropicCacheControlHook, + block_with_cache_control, + cache_control_count, + dropped_cache_control, + newest_marks_within_cache_control_limit, + whole_messages_within_cache_control_limit, + with_cache_control, +) from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, @@ -603,10 +611,6 @@ def _validated_cache_control(value: object) -> ChatCompletionCachedContent | Non return None -def _dropped_cache_control(_: object) -> None: - return None - - def _without_cache_control_key(item: object) -> object: if not isinstance(item, dict): return item @@ -694,83 +698,6 @@ def _content_blocks_carrying(content: object, cache_control: ChatCompletionCache return _with_cache_control_on_last_block(_OBJECT_LIST_ADAPTER.validate_python(content), cache_control) -_MARK_CARRYING_MESSAGE_FIELDS: Final = ("content", "tool_calls") - - -def _mark_carrying_blocks(fields: Mapping[str, object]) -> Sequence[object]: - lists: Final = tuple(fields.get(key) for key in _MARK_CARRYING_MESSAGE_FIELDS) - return tuple( - chain.from_iterable(_OBJECT_LIST_ADAPTER.validate_python(value) for value in lists if isinstance(value, list)) - ) - - -def _block_carries_cache_control(block: object) -> bool: - return isinstance(block, dict) and _STR_KEY_DICT_ADAPTER.validate_python(block).get("cache_control") is not None - - -def _cache_control_count(message: object) -> int: - if not isinstance(message, dict): - return 0 - fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) - block_marks: Final = sum(1 for block in _mark_carrying_blocks(fields) if _block_carries_cache_control(block)) - return int(fields.get("cache_control") is not None) + block_marks - - -_CacheControlRewrite: TypeAlias = Callable[[object], ChatCompletionCachedContent | None] - - -def _fields_with_cache_control(fields: Mapping[str, object], rewrite: _CacheControlRewrite) -> Mapping[str, object]: - others: Final = {key: value for key, value in fields.items() if key != "cache_control"} - if "cache_control" not in fields: - return others - cache_control: Final = rewrite(fields["cache_control"]) - return others if cache_control is None else {**others, "cache_control": cache_control} - - -def _block_with_cache_control(block: object, rewrite: _CacheControlRewrite) -> object: - if not isinstance(block, dict): - return block - return _fields_with_cache_control(_STR_KEY_DICT_ADAPTER.validate_python(block), rewrite) - - -def _with_cache_control(message: _MessageT, rewrite: _CacheControlRewrite) -> _MessageT: - if not isinstance(message, dict): - return message - fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) - blocks_rewritten: Final = { - key: ( - [_block_with_cache_control(block, rewrite) for block in _OBJECT_LIST_ADAPTER.validate_python(value)] - if key in _MARK_CARRYING_MESSAGE_FIELDS and isinstance(value, list) - else value - ) - for key, value in fields.items() - } - return cast( # cast-ok: the same message shape with each cache marker rewritten or dropped - _MessageT, _fields_with_cache_control(blocks_rewritten, rewrite) - ) - - -def _kept_within_budget(state: tuple[int, bool], marks: int) -> tuple[int, bool]: - remaining, _ = state - if marks > remaining: - return remaining, False - return remaining - marks, True - - -def _messages_within_cache_control_limit(messages: Sequence[_MessageT], reserved_marks: int) -> Sequence[_MessageT]: - budget: Final = max(MAX_CACHE_CONTROL_BLOCKS - reserved_marks, 0) - newest_first: Final = tuple(reversed(messages)) - budget_after_each: Final = accumulate( - (_cache_control_count(message) for message in newest_first), _kept_within_budget, initial=(budget, True) - ) - kept_flags: Final = tuple(kept for _, kept in islice(budget_after_each, 1, None)) - kept_newest_first: Final = tuple( - message if kept else _with_cache_control(message, _dropped_cache_control) - for message, kept in zip(newest_first, kept_flags) - ) - return tuple(reversed(kept_newest_first)) - - _SYSTEM_MESSAGE_FIELDS: Final = frozenset(("content", "cache_control")) @@ -799,7 +726,7 @@ def _without_tool_block_marks(message: _MessageT) -> _MessageT: if fields is None: return message if fields.get("role") == "tool": - return _with_cache_control(message, _dropped_cache_control) + return with_cache_control(message, dropped_cache_control) tool_calls: Final = fields.get("tool_calls") if not isinstance(tool_calls, list): return message @@ -808,7 +735,7 @@ def _without_tool_block_marks(message: _MessageT) -> _MessageT: { **fields, "tool_calls": [ - _block_with_cache_control(tool_call, _dropped_cache_control) + block_with_cache_control(tool_call, dropped_cache_control) for tool_call in _OBJECT_LIST_ADAPTER.validate_python(tool_calls) ], }, @@ -822,7 +749,7 @@ class _ItemMarkRules: external_marks: int def sanitized(self, message: _MessageT) -> _MessageT: - valid: Final = _with_cache_control(message, _validated_cache_control) + valid: Final = with_cache_control(message, _validated_cache_control) return _without_tool_block_marks(valid) if self.refuses_tool_block_marks else valid def checked(self, message: _MessageT) -> _MessageT: @@ -845,15 +772,15 @@ def _request_messages_for_cache_points(messages: Sequence[_MessageT], rules: _It as_system: Final = tuple(_developer_message_as_system_message(message) for message in messages) if not rules.drop_params: return as_system - return _messages_within_cache_control_limit(tuple(map(rules.checked, as_system)), rules.external_marks) + return newest_marks_within_cache_control_limit(tuple(map(rules.checked, as_system)), rules.external_marks) def _replayed_messages_for_cache_points( replayed_messages: Sequence[_MessageT], current_messages: Sequence[object], rules: _ItemMarkRules ) -> Sequence[_MessageT]: as_system: Final = tuple(_developer_message_as_system_message(message) for message in replayed_messages) - current_marks: Final = sum(_cache_control_count(message) for message in current_messages) - return _messages_within_cache_control_limit( + current_marks: Final = sum(cache_control_count(message) for message in current_messages) + return whole_messages_within_cache_control_limit( tuple(map(rules.sanitized, as_system)), current_marks + rules.external_marks ) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 9259103cb5c..b7b93a20b4e 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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, 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/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 fb91ba6a618..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, ) @@ -2277,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: @@ -2562,7 +2585,8 @@ 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( @@ -2995,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: @@ -3416,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 diff --git a/litellm/router.py b/litellm/router.py index 233e7581a40..6392ea8cd90 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 @@ -233,49 +300,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, @@ -292,7 +317,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, @@ -337,10 +361,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, @@ -392,19 +412,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, @@ -420,11 +427,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 @@ -504,6 +506,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, @@ -519,8 +522,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 @@ -634,6 +635,7 @@ _PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS: Final = frozenset({"anthropic_messages", "aresponses", "responses"}) _video_call_deployment_id: Final[ContextVar[str | None]] = ContextVar("video_call_deployment_id", default=None) _routed_call_type: Final[ContextVar[str | None]] = ContextVar("routed_call_type", default=None) +_DECISIONS_MEDIA_CALL_TYPES: Final = frozenset({"adecisions", "decisions", "asystemone", "systemone"}) _ROUTE_CAPABILITY_RESOLVERS: Final[Mapping[str, Callable[[str, LlmProviders], object | None]]] = MappingProxyType( { "adecisions": ProviderConfigManager.get_provider_decisions_config, @@ -897,6 +899,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 @@ -1592,6 +1611,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 @@ -1867,7 +1887,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, @@ -2241,7 +2261,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 "") @@ -2566,7 +2586,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) @@ -3115,7 +3135,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 ) @@ -3123,7 +3143,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, ) @@ -3672,7 +3692,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, @@ -3845,7 +3865,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() @@ -3917,7 +3949,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, @@ -3985,10 +4019,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()) @@ -4120,13 +4154,16 @@ class Router(DeprecatedRouterMembers): raise e.original_exception from e raise e + def _active_default_litellm_params(self) -> Mapping[str, object]: + default_params: Final = _REQUEST_KWARGS_ADAPTER.validate_python(self._core_default_litellm_params) + return {key: value for key, value in default_params.items() if value is not None} + def _params_sent_to_deployment( self, request_kwargs: Mapping[str, object], deployment: Deployment | None ) -> Mapping[str, object]: """The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies.""" deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} - default_params: Final = _REQUEST_KWARGS_ADAPTER.validate_python(self._core_default_litellm_params) - defaults: Final = {key: value for key, value in default_params.items() if value is not None} + defaults: Final = self._active_default_litellm_params() deployment_tool_choice: Final = ( {"tool_choice": deployment_params["tool_choice"]} if deployment_params.get("tool_choice") is not None @@ -4348,6 +4385,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__() @@ -4375,6 +4416,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): @@ -4401,7 +4443,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( @@ -4572,7 +4625,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, @@ -4658,7 +4715,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)) @@ -4683,7 +4752,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, @@ -4867,7 +4938,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) @@ -5330,7 +5401,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),), @@ -5936,7 +6007,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, @@ -5946,15 +6017,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() @@ -6021,7 +6119,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, @@ -6114,7 +6212,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, @@ -6176,7 +6274,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, @@ -6292,7 +6390,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, @@ -6362,7 +6460,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, @@ -6421,7 +6519,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, ) @@ -6435,7 +6533,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, @@ -6553,7 +6651,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, @@ -6653,7 +6751,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), @@ -7531,7 +7629,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, @@ -7610,7 +7708,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, @@ -7624,7 +7722,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) @@ -7850,7 +7948,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, @@ -8018,7 +8116,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, @@ -8049,7 +8147,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, @@ -8839,7 +8937,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 @@ -8919,7 +9017,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, @@ -9135,7 +9233,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( @@ -9756,7 +9854,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), ) @@ -9936,7 +10034,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 @@ -10047,11 +10145,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, @@ -10357,9 +10455,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) @@ -10387,7 +10483,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) @@ -10445,7 +10541,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"], @@ -10529,7 +10625,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), @@ -10960,7 +11056,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( @@ -11041,7 +11137,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( @@ -11341,7 +11437,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( @@ -13562,7 +13658,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, @@ -14445,7 +14541,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", {}) @@ -14538,7 +14634,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": @@ -14819,6 +14915,62 @@ class Router(DeprecatedRouterMembers): return healthy_deployments return cast(_RouteDeploymentsT, filtered_deployments) # cast-ok: same list shape as the input candidates + @staticmethod + def _decisions_deployment_refuses_media( + deployment: DeploymentTypedDict, group: str, request_kwargs: Mapping[str, object], *, audio: bool, videos: bool + ) -> bool: + from litellm.decisions.call import deployment_refuses_media + + plain_deployment: Final = cast(Mapping[str, object], deployment) # cast-ok: router deployments are plain dicts + model_info: Final = cast(Mapping[str, object], plain_deployment.get("model_info") or {}) # cast-ok: plain dict + deployment_params: Final = cast(Mapping[str, object], deployment["litellm_params"]) # cast-ok: plain dict + try: + model, custom_llm_provider = Router._capability_model_and_provider(deployment["litellm_params"], group) + return deployment_refuses_media( + model, + custom_llm_provider, + model_info, + {**deployment_params, **request_kwargs}, + audio=audio, + videos=videos, + ) + except Exception as e: # noqa: BLE001 # best-effort filter: a deployment it cannot resolve answers for itself + verbose_router_logger.debug( + "litellm.router.py::_decisions_deployment_refuses_media: keeping a deployment of model=%s. Got - %s", + group, + e, + ) + return False + + def _without_decisions_deployments_refusing_media( + self, healthy_deployments: _RouteDeploymentsT, group: str, request_kwargs: Mapping[str, object] | None + ) -> _RouteDeploymentsT: + call_kwargs: Final = request_kwargs or {} + audio: Final = bool(call_kwargs.get("audio")) + videos: Final = bool(call_kwargs.get("videos")) + if ( + _routed_call_type.get() not in _DECISIONS_MEDIA_CALL_TYPES + or not (audio or videos) + or isinstance(healthy_deployments, Mapping) + ): + return healthy_deployments + + sent_kwargs: Final = {**self._active_default_litellm_params(), **call_kwargs} + deployments_taking_media: Final = [ + deployment + for deployment in healthy_deployments + if not Router._decisions_deployment_refuses_media( + cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape + group, + sent_kwargs, + audio=audio, + videos=videos, + ) + ] + if not deployments_taking_media: + return healthy_deployments + return cast(_RouteDeploymentsT, deployments_taking_media) # cast-ok: same list shape as the input candidates + @staticmethod def _without_declared_text_only_deployments( healthy_deployments: _RouteDeploymentsT, @@ -15469,10 +15621,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, ) @@ -15527,7 +15679,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, @@ -15543,10 +15695,14 @@ class Router(DeprecatedRouterMembers): ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + batch_candidates: Final = self._without_decisions_deployments_refusing_media( + self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict + ), model, - messages, cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( @@ -15606,7 +15762,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( @@ -15852,7 +16008,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, ) @@ -16022,7 +16178,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, ) @@ -16390,7 +16546,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, @@ -16753,9 +16909,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, @@ -16788,10 +16942,14 @@ class Router(DeprecatedRouterMembers): ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + batch_candidates: Final = self._without_decisions_deployments_refusing_media( + self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict + ), model, - messages, cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( @@ -16854,7 +17012,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, @@ -16862,7 +17020,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, @@ -16901,7 +17059,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, @@ -17008,9 +17166,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, @@ -17039,7 +17195,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, @@ -17047,7 +17203,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, @@ -17085,7 +17241,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 046fb8c452a..0b328cc3e6b 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -40,11 +40,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, @@ -838,6 +838,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/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/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 a54d66c4beb..ee86ec52327 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -402,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 @@ -440,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 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/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 f8115571dfd..24f4acaf5df 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -106,6 +106,7 @@ from .llms.openai import ( OpenAIRealtimeStreamList, ResponsesAPIResponse, WebSearchOptions, + _JsonValue, ) from .rerank import RerankResponse as RerankResponse @@ -464,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) @@ -3323,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] @@ -3381,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", @@ -3460,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 @@ -4157,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 @@ -4669,6 +4683,7 @@ class LlmProviders(str, Enum): PRISM = "prism" OPENINFER = "openinfer" INFERENCEAPIS = "inferenceapis" + HOSTINGER_ROUTER = "hostinger_router" ZEROGPU = "zerogpu" DARKBLOOM = "darkbloom" META = "meta" @@ -4950,29 +4965,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 2ca8a49f968..f19fa0dbb34 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -363,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, ) @@ -455,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, @@ -5258,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 () @@ -6889,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 @@ -6899,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, ) @@ -6962,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 @@ -7079,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.", @@ -7089,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.", @@ -7371,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), @@ -8724,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, @@ -8738,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, @@ -8791,20 +8832,27 @@ 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, ...]: @@ -10111,7 +10159,7 @@ class ProviderConfigManager: from litellm.llms.openai_like.json_loader import JSONProviderRegistry json_provider: Final = JSONProviderRegistry.get(provider.value) - if json_provider is not None and "/v1/messages" in json_provider.supported_endpoints: + if json_provider is not None and JSONProviderRegistry.serves_endpoint(provider.value, model, "/v1/messages"): from litellm.llms.openai_like.messages.transformation import ( JSONProviderAnthropicMessagesConfig, ) @@ -10250,7 +10298,7 @@ class ProviderConfigManager: return result # Fall back to JSON providers (generic OpenAI-compatible) - if JSONProviderRegistry.exists(provider_str) and JSONProviderRegistry.supports_responses_api(provider_str): + if JSONProviderRegistry.serves_endpoint(provider_str, model, "/v1/responses"): provider_config: Final = JSONProviderRegistry.get(provider_str) if provider_config is not None: return create_responses_config_class(provider_config)() @@ -10470,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 dac7af6ab7a..7eede4a395f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23078,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", @@ -25284,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", @@ -29561,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, @@ -30167,6 +30218,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_pdf_input": true, @@ -30210,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, @@ -30274,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, @@ -30334,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, @@ -30401,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, @@ -31218,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": [ @@ -31398,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, @@ -32405,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, @@ -32479,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, @@ -32543,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, @@ -32603,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, @@ -32670,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, @@ -39622,6 +39687,57 @@ "/v1/images/generations" ] }, + "hostinger_router/gpt-5.2": { + "cache_read_input_token_cost": 0.00000021, + "input_cost_per_token": 0.0000021, + "litellm_provider": "hostinger_router", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0000168, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_vision": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "source": "https://router.hostinger.com/v1/models" + }, + "hostinger_router/claude-sonnet-5": { + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_1hr": 5.28e-06, + "cache_read_input_token_cost": 2.64e-07, + "input_cost_per_token": 2.64e-06, + "litellm_provider": "hostinger_router", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.32e-05, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_response_schema": true, + "supported_endpoints": ["/v1/chat/completions", "/v1/messages"], + "source": "https://router.hostinger.com/v1/models" + }, + "hostinger_router/text-embedding-3-large": { + "input_cost_per_token": 1.716e-07, + "litellm_provider": "hostinger_router", + "max_input_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0, + "source": "https://router.hostinger.com/v1/models" + }, + "hostinger_router/text-embedding-3-small": { + "input_cost_per_token": 3e-08, + "litellm_provider": "hostinger_router", + "max_input_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0, + "source": "https://router.hostinger.com/v1/models" + }, "hyperbolic/NousResearch/Hermes-3-Llama-3.1-70B": { "input_cost_per_token": 1.2e-07, "litellm_provider": "hyperbolic", @@ -47624,6 +47740,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, @@ -51304,6 +51438,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, @@ -51332,6 +51471,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, @@ -52078,6 +52222,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, @@ -52183,6 +52332,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, @@ -52713,6 +52867,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, @@ -59714,12 +59869,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": [ @@ -75398,6 +75554,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, @@ -76498,6 +76655,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, @@ -76543,6 +76701,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, @@ -76654,6 +76813,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index d33e61026df..700a35edb06 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1087,6 +1087,32 @@ "generateContent": false } }, + "hostinger_router": { + "display_name": "Hostinger Router (`hostinger_router`)", + "url": "https://docs.litellm.ai/docs/providers/hostinger_router", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": true, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } + }, "crusoe": { "display_name": "Crusoe (`crusoe`)", "url": "https://docs.litellm.ai/docs/providers/crusoe", diff --git a/pyproject.toml b/pyproject.toml index 6e9e75a5b7f..5c23a64ebb3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -383,6 +383,8 @@ exclude-newer = "3 days" [tool.uv.exclude-newer-package] # TODO: remove after 2026-10-25, once brotli 1.2.1 (CVE-2026-111497) clears the 3-day window brotli = "2026-10-22T14:01:00Z" +# TODO: remove after 2026-10-31, once pypdf 6.21.4 clears the 3-day window +pypdf = "2026-10-28T17:01:00Z" [tool.uv.sources] litellm-proxy-extras = { workspace = 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/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..3a708ce6d0a 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -183,6 +183,9 @@ ignored_function_names = [ "_aresponses_background_stream_can_release", # Tested through the background Responses stream tests in test_router.py "_raise_released_stream_error", # Tested through the background Responses stream content-policy tests in test_router.py "_mid_stream_error_before_output", # Tested through the Responses mid-stream fallback and background stream tests + "_without_decisions_deployments_refusing_media", # Tested through the mixed media-flag group tests in tests/unit/decisions/test_main.py + "_decisions_deployment_refuses_media", # Read by _without_decisions_deployments_refusing_media, tested through the same group tests + "_active_default_litellm_params", # Tested through the router default drop_params test in tests/unit/decisions/test_main.py ] @@ -198,11 +201,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/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_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/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_anthropic_passthrough_wire.py b/tests/integration/providers/test_anthropic_passthrough_wire.py index a791eb9c2a2..0a16520f626 100644 --- a/tests/integration/providers/test_anthropic_passthrough_wire.py +++ b/tests/integration/providers/test_anthropic_passthrough_wire.py @@ -317,7 +317,6 @@ def test_credentialless_proxy_forwards_the_caller_anthropic_credential_and_never def test_anthropic_sdk_metadata_user_id_reaches_upstream(api_key_proxy: Gateway, wire: Wire) -> None: - pytest.skip("BUG: /anthropic/v1/messages drops the native metadata.user_id field before forwarding upstream") with api_key_proxy.scenario() as scenario: key: Final = scenario.key() sent: Final[list[httpx.Request]] = [] 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_decisions_media_deployment_flag_wire.py b/tests/integration/providers/test_decisions_media_deployment_flag_wire.py new file mode 100644 index 00000000000..a2e53ef47ca --- /dev/null +++ b/tests/integration/providers/test_decisions_media_deployment_flag_wire.py @@ -0,0 +1,123 @@ +import json +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_API_KEY: Final = "synthetic-decisions-key" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_INPUT: Final = "Listen to the attached clip." +_AUDIO: Final = ("data:audio/wav;base64,UklGRiQAAABXQVZF",) +_VIDEOS: Final = ("data:video/mp4;base64,AAAAIGZ0eXBpc29t",) +_QUESTION: Final[dict[str, JsonValue]] = { + "type": "predicate", + "name": "banana", + "instructions": "Does the clip name a banana?", +} +_SPEND_QUERY: Final = 'SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id = %s' + +Route = Literal["/v1/decisions", "/v1/systemone"] +Media = Literal["audio", "videos"] + + +@dataclass(frozen=True, slots=True) +class _Upstream: + model: str + + def reply(self, request: Request) -> Reply: + questions: Final = object_value(_JSON_OBJECT.validate_json(request.body)["questions"]) + plain: Final[dict[str, JsonValue]] = { + "model": self.model, + "answers": {name: {"type": "noul", "noul": 0.9} for name in questions}, + "usage": {"input_tokens": 41, "output_tokens": 0}, + } + return Reply(body=json.dumps(plain).encode()) + + +def _unlisted_openrouter_model() -> _Upstream: + return _Upstream(f"openrouter/example/omni-{uuid.uuid4().hex[:12]}") + + +def _decide(gateway: Gateway, model: str, route: Route, media: Mapping[str, JsonValue]) -> httpx.Response: + body: Final[dict[str, JsonValue]] = ( + {"model": model, "input": _INPUT, "questions": [_QUESTION], "num_retries": 0, **media} + if route == "/v1/decisions" + else { + "model": model, + "state": _INPUT, + "questions": {"banana": {"type": "noul", "instructions": _QUESTION["instructions"]}}, + "num_retries": 0, + **media, + } + ) + return gateway.request("POST", route, body) + + +def _sent_media(request: Request) -> dict[str, JsonValue]: + body: Final = _JSON_OBJECT.validate_json(request.body) + return {kind: body[kind] for kind in ("audio", "videos") if kind in body} + + +def _assert_refused(response: httpx.Response, kinds: tuple[Media, ...]) -> None: + assert response.status_code == 400, response.text + assert f"does not support parameters: {list(kinds)}" in response.text, response.text + + +_BOTH: Final[dict[str, JsonValue]] = {"audio": list(_AUDIO), "videos": list(_VIDEOS)} + + +@pytest.mark.parametrize("route", ("/v1/decisions", "/v1/systemone")) +def test_a_deployment_flagged_for_audio_and_video_sends_both_clips_for_a_model_the_cost_map_does_not_list( + gateway: Gateway, route: Route +) -> None: + upstream: Final = _unlisted_openrouter_model() + with wire_server(upstream.reply) as wire, gateway.scenario() as scenario: + flags: Final[dict[str, JsonValue]] = {"supports_audio_input": True, "supports_video_input": True} + flagged: Final = scenario.model(model=upstream.model, api_base=wire.url, api_key=_API_KEY, model_info=flags) + flagged_dropping: Final = scenario.model( + model=upstream.model, api_base=wire.url, api_key=_API_KEY, model_info=flags, drop_params=True + ) + audio_only: Final = scenario.model( + model=upstream.model, api_base=wire.url, api_key=_API_KEY, model_info={"supports_audio_input": True} + ) + unflagged: Final = scenario.model(model=upstream.model, api_base=wire.url, api_key=_API_KEY) + + _assert_refused(_decide(gateway, unflagged, route, _BOTH), ("audio", "videos")) + _assert_refused(_decide(gateway, audio_only, route, _BOTH), ("videos",)) + assert wire.drain() == () + + answered: Final = _decide(gateway, flagged, route, _BOTH) + assert answered.status_code == 200, answered.text + (call,) = wire.drain() + assert _sent_media(call) == _BOTH + rows: Final = eventually( + lambda: read_rows(_SPEND_QUERY, (answered.headers["x-litellm-call-id"],)), + lambda found: len(found) == 1, + seconds=70, + ) + assert (rows[0]["status"], rows[0]["model_group"]) == ("success", flagged) + + kept: Final = _decide(gateway, flagged_dropping, route, _BOTH) + assert kept.status_code == 200, kept.text + (kept_call,) = wire.drain() + assert _sent_media(kept_call) == _BOTH + + +@pytest.mark.parametrize( + "value", + ({"supports_audio_input": True, "supports_video_input": True}, 5, [{"supports_audio_input": True}], "x" * 5000), + ids=("flags", "int", "list", "5kb"), +) +def test_a_caller_sent_model_info_never_turns_media_on(gateway: Gateway, value: JsonValue) -> None: + upstream: Final = _unlisted_openrouter_model() + with wire_server(upstream.reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=upstream.model, api_base=wire.url, api_key=_API_KEY) + _assert_refused(_decide(gateway, model, "/v1/decisions", {**_BOTH, "model_info": value}), ("audio", "videos")) + assert wire.drain() == () 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/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_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/decisions/test_main.py b/tests/unit/decisions/test_main.py index 0c357837d31..789d59bd613 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -6,10 +6,11 @@ import copy import json import re import threading -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from concurrent.futures import wait +from functools import partial from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias import httpx import pytest @@ -23,6 +24,7 @@ import litellm from litellm.constants import AWS_SIGNING_MAX_THREADS from litellm.cost_calculator import get_response_cost_from_hidden_params from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded, refetch_model_cost_map from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR @@ -3702,6 +3704,387 @@ async def test_a_model_info_object_instead_of_a_dict_still_reaches_the_provider( assert response.answers[0].model_dump(mode="json")["name"] == "is_defect" +_AUDIO_CLIPS: Final = ("data:audio/wav;base64,YQ==",) +_VIDEO_CLIPS: Final = ("data:video/mp4;base64,YQ==",) + + +@pytest.mark.asyncio +async def test_a_deployment_that_advertises_media_sends_it_for_a_model_the_cost_map_does_not_list( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + backend: Final = "openrouter/example/omni-preview" + route: Final = respx_mock.post("https://openrouter.ai/api/alpha/decisions").respond(json=_RESPONSE) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "hears", + "litellm_params": {"model": backend, "api_key": "k"}, + "model_info": {"id": "media-hears", "supports_audio_input": True}, + }, + { + "model_name": "hears-and-sees", + "litellm_params": {"model": backend, "api_key": "k"}, + "model_info": {"id": "media-both", "supports_audio_input": True, "supports_video_input": True}, + }, + ] + ) + send: Final = partial( + router.adecisions, input="review", questions=_PREDICATE_QUESTIONS, audio=_AUDIO_CLIPS, videos=_VIDEO_CLIPS + ) + + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("['videos']")): + await send(model="hears") + assert not route.called + + await send(model="hears-and-sees") + + sent: Final = json.loads(route.calls.last.request.content) + assert (sent["audio"], sent["videos"]) == (list(_AUDIO_CLIPS), list(_VIDEO_CLIPS)) + + +@pytest.mark.asyncio +async def test_media_flags_on_a_deployment_whose_own_mode_is_not_decisions_never_send_clips( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://openrouter.ai/api/alpha/decisions").respond(json=_RESPONSE) + advertised: Final = {"supports_audio_input": True, "supports_video_input": True} + router: Final = litellm.Router( + model_list=[ + { + "model_name": mode, + "litellm_params": {"model": f"openrouter/example/omni-{mode}-preview", "api_key": "k"}, + "model_info": {"id": f"media-{mode}", "mode": mode, **advertised}, + } + for mode in ("chat", "decisions") + ] + ) + send: Final = partial( + router.adecisions, input="review", questions=_PREDICATE_QUESTIONS, audio=_AUDIO_CLIPS, videos=_VIDEO_CLIPS + ) + + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("['audio', 'videos']")): + await send(model="chat") + assert not route.called + + await send(model="decisions") + + sent: Final = json.loads(route.calls.last.request.content) + assert (sent["audio"], sent["videos"]) == (list(_AUDIO_CLIPS), list(_VIDEO_CLIPS)) + + +CatalogLoader: TypeAlias = Callable[[Mapping[str, Mapping[str, object]]], Awaitable[None]] + + +@pytest.fixture +def load_catalog_with(monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None) -> Iterator[CatalogLoader]: + async def load(rows: Mapping[str, Mapping[str, object]]) -> None: + catalog: Final = {**litellm.get_model_cost_map(url=""), **rows} + transport: Final = httpx.MockTransport(lambda _: httpx.Response(200, json=catalog)) + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP") + async with httpx.AsyncClient(transport=transport) as client: + loaded: Final = await refetch_model_cost_map(url="https://catalog.example/prices.json", client=client) + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + assert isinstance(loaded, ModelCostMapReloaded) + monkeypatch.setattr(litellm, "model_cost", loaded.model_cost_map) + litellm.get_model_info.cache_clear() + + yield load + litellm.get_model_cost_map(url="") + + +@pytest.mark.asyncio +async def test_media_flags_on_a_deployment_without_its_own_mode_follow_the_mode_the_catalog_lists( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, load_catalog_with: CatalogLoader +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + await load_catalog_with( + { + f"openrouter/example/omni-{mode}": {"litellm_provider": "openrouter", "mode": mode} + for mode in ("chat", "decisions") + } + ) + route: Final = respx_mock.post("https://openrouter.ai/api/alpha/decisions").respond(json=_RESPONSE) + router: Final = litellm.Router( + model_list=[ + { + "model_name": mode, + "litellm_params": {"model": f"openrouter/example/omni-{mode}", "api_key": "k"}, + "model_info": {"id": f"listed-{mode}", "supports_audio_input": True, "supports_video_input": True}, + } + for mode in ("chat", "decisions") + ] + ) + send: Final = partial( + router.adecisions, input="review", questions=_PREDICATE_QUESTIONS, audio=_AUDIO_CLIPS, videos=_VIDEO_CLIPS + ) + + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("['audio', 'videos']")): + await send(model="chat") + assert not route.called + + monkeypatch.setattr(litellm, "drop_params", True) + await send(model="chat") + dropped: Final = json.loads(route.calls.last.request.content) + assert ("audio" in dropped, "videos" in dropped) == (False, False) + + await send(model="decisions") + sent: Final = json.loads(route.calls.last.request.content) + assert (sent["audio"], sent["videos"]) == (list(_AUDIO_CLIPS), list(_VIDEO_CLIPS)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sibling_mode", ("responses", "embedding")) +async def test_another_deployments_mode_never_switches_off_a_deployments_own_media_flags( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None, sibling_mode: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://openrouter.ai/api/alpha/decisions").respond(json=_RESPONSE) + shared_params: Final = {"model": "openrouter/example/omni-shared", "api_key": "k"} + router: Final = litellm.Router( + model_list=[ + {"model_name": "sibling", "litellm_params": shared_params, "model_info": {"id": "s", "mode": sibling_mode}}, + { + "model_name": "flagged", + "litellm_params": shared_params, + "model_info": {"id": "f", "supports_audio_input": True, "supports_video_input": True}, + }, + ] + ) + + await router.adecisions( + model="flagged", input="review", questions=_PREDICATE_QUESTIONS, audio=_AUDIO_CLIPS, videos=_VIDEO_CLIPS + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert (sent["audio"], sent["videos"]) == (list(_AUDIO_CLIPS), list(_VIDEO_CLIPS)) + + +_MEDIA_SURFACES: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType( + { + "adecisions": {"input": "review", "questions": _PREDICATE_QUESTIONS}, + "asystemone": {"state": "review", "questions": _QUESTIONS}, + "systemone": {"state": "review", "questions": _QUESTIONS}, + } +) + + +def _group_member( + group: str, model: str, key: str, weight: int, model_info: Mapping[str, object], **params: object +) -> dict[str, object]: + return { + "model_name": group, + "litellm_params": {"model": model, "api_key": key, "weight": weight, **params}, + "model_info": {"id": f"{group}-{key}", **model_info}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", tuple(_MEDIA_SURFACES)) +async def test_a_group_sends_clips_to_a_deployment_that_takes_them_and_refuses_only_when_none_does( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None, surface: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + monkeypatch.delenv("CLOUDFLARE_API_BASE", raising=False) + monkeypatch.setitem( + litellm.model_cost, + "cloudflare/@cf/cloudflare/clef-omni", + {"litellm_provider": "cloudflare", "mode": "evaluation", "supports_audio_input": True}, + ) + respx_mock.post("https://openrouter.ai/api/alpha/decisions").respond(json=_RESPONSE) + respx_mock.post("https://api.cloudflare.com/client/v4/accounts/acct/ai/run/@cf/cloudflare/clef-omni").respond( + json=_RESPONSE + ) + unlisted: Final = "openrouter/example/omni-preview" + listed: Final = "cloudflare/clef-omni" + router: Final = litellm.Router( + model_list=[ + _group_member("unlisted", unlisted, "refuses", 1, {}), + _group_member("unlisted", unlisted, "takes", 0, {"supports_audio_input": True}), + _group_member("listed", listed, "refuses", 1, {"supports_audio_input": False}), + _group_member("listed", listed, "takes", 0, {}), + _group_member("none", unlisted, "refuses", 1, {}), + _group_member("none", listed, "refuses", 1, {"supports_audio_input": False}), + _group_member("drops", unlisted, "drops", 1, {}, drop_params=True), + _group_member("drops", listed, "takes", 0, {}), + _group_member("unresolvable", "databricks/not a serving endpoint", "unresolvable", 0, {}), + _group_member("unresolvable", listed, "takes", 1, {}), + ] + ) + call: Final = getattr(router, surface) + + async def send(group: str, **media: object) -> object: + response: Final = call(model=group, **_MEDIA_SURFACES[surface], **media) + return await response if asyncio.iscoroutine(response) else response + + for group in ("unlisted", "listed", "unresolvable"): + await send(group, audio=_AUDIO_CLIPS) + assert respx_mock.calls.last.request.headers["authorization"] == "Bearer takes" + assert json.loads(respx_mock.calls.last.request.content)["audio"] == list(_AUDIO_CLIPS) + + await send("drops", audio=_AUDIO_CLIPS) + assert respx_mock.calls.last.request.headers["authorization"] == "Bearer drops" + assert "audio" not in json.loads(respx_mock.calls.last.request.content) + + await send("unlisted") + assert respx_mock.calls.last.request.headers["authorization"] == "Bearer refuses" + + sent_before_refusal: Final = len(respx_mock.calls) + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("['audio']")): + await send("none", audio=_AUDIO_CLIPS) + assert len(respx_mock.calls) == sent_before_refusal + + +def _answer_unless_key_is(rejected_key: str) -> Callable[[httpx.Request], httpx.Response]: + def answer(request: httpx.Request) -> httpx.Response: + if request.headers["authorization"] == f"Bearer {rejected_key}": + return httpx.Response(401, json={"error": {"message": "invalid key"}}) + return httpx.Response(200, json=_RESPONSE) + + return answer + + +@pytest.mark.asyncio +async def test_clips_to_a_group_whose_only_taker_is_cooling_down_are_refused_rather_than_rate_limited( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://openrouter.ai/api/alpha/decisions").mock( + side_effect=_answer_unless_key_is("takes") + ) + unlisted: Final = "openrouter/example/omni-preview" + router: Final = litellm.Router( + model_list=[ + _group_member("g", unlisted, "refuses", 1, {}), + _group_member("g", unlisted, "takes", 1, {"supports_audio_input": True}), + ], + num_retries=0, + cooldown_time=120, + ) + with pytest.raises(litellm.AuthenticationError): + await router.adecisions(model="g-takes", input="review", questions=_PREDICATE_QUESTIONS) + sent_before_clips: Final = len(route.calls) + + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("['audio']")): + await router.adecisions(model="g", input="review", questions=_PREDICATE_QUESTIONS, audio=_AUDIO_CLIPS) + + assert len(route.calls) == sent_before_clips + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", tuple(_MEDIA_SURFACES)) +async def test_a_router_default_drop_params_overrides_a_deployment_when_choosing_who_takes_clips( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None, surface: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + respx_mock.post("https://openrouter.ai/api/alpha/decisions").respond(json=_RESPONSE) + unlisted: Final = "openrouter/example/omni-preview" + router: Final = litellm.Router( + model_list=[ + _group_member("g", unlisted, "drops", 1, {}, drop_params=True), + _group_member("g", unlisted, "takes", 0, {"supports_audio_input": True}), + ], + default_litellm_params={"drop_params": False}, + ) + call: Final = getattr(router, surface) + + async def send(**params: object) -> object: + response: Final = call(model="g", **_MEDIA_SURFACES[surface], audio=_AUDIO_CLIPS, **params) + return await response if asyncio.iscoroutine(response) else response + + await send() + assert respx_mock.calls.last.request.headers["authorization"] == "Bearer takes" + assert json.loads(respx_mock.calls.last.request.content)["audio"] == list(_AUDIO_CLIPS) + + await send(drop_params=True) + assert respx_mock.calls.last.request.headers["authorization"] == "Bearer drops" + assert "audio" not in json.loads(respx_mock.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_a_deployment_that_turns_audio_off_is_refused_or_dropped_even_when_the_cost_map_allows_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setitem( + litellm.model_cost, + "typesafe/jev-omni", + { + "litellm_provider": "typesafe", + "mode": "decisions", + "supports_audio_input": True, + "supports_video_input": True, + }, + ) + route: Final = respx_mock.post("https://example.com/v1/systemone").respond(json=_RESPONSE) + send: Final = partial( + litellm.asystemone, + model="typesafe/jev-omni", + state="test", + questions=_QUESTIONS, + api_base="https://example.com", + api_key="test", + audio=_AUDIO_CLIPS, + videos=_VIDEO_CLIPS, + model_info={"supports_audio_input": False}, + ) + + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("does not support parameters: ['audio']")): + await send() + assert not route.called + + await send(drop_params=True) + + sent: Final = json.loads(route.calls.last.request.content) + assert "audio" not in sent + assert sent["videos"] == list(_VIDEO_CLIPS) + + +@pytest.mark.asyncio +async def test_a_media_flag_that_is_not_a_boolean_leaves_the_cost_map_in_charge_of_that_kind_only( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://example.com/v1/systemone").respond(json=_RESPONSE) + + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("does not support parameters: ['audio']")): + await litellm.asystemone( + model="typesafe/jev-unlisted", + state="test", + questions=_QUESTIONS, + api_base="https://example.com", + api_key="test", + audio=_AUDIO_CLIPS, + videos=_VIDEO_CLIPS, + model_info={"supports_audio_input": "true", "supports_video_input": True}, + ) + + assert not route.called + + +@pytest.mark.asyncio +async def test_an_openai_format_deployment_that_advertises_media_is_still_refused_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("['audio', 'videos']")): + await litellm.adecisions( + model="openai/gpt-6-luna", + input="review", + questions=_PREDICATE_QUESTIONS, + api_key="deployment-key", + audio=_AUDIO_CLIPS, + videos=_VIDEO_CLIPS, + model_info={"supports_audio_input": True, "supports_video_input": True}, + ) + + assert not route.called + + def test_dropping_every_question_answers_400_before_http(respx_mock: respx.MockRouter) -> None: route: Final = respx_mock.post(_PERPLEXITY_DECISIONS_URL).respond(json=_RESPONSE) 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/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/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 72e72714cd4..4ce6aee1619 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_semantic_cache_prompt_from_messages, get_str_from_messages, handle_any_messages_to_chat_completion_str_messages_conversion, + handle_messages_with_content_list_to_str_conversion, hoist_images_from_tool_messages, is_encrypted_reasoning_block, merge_consecutive_system_messages, @@ -2675,3 +2676,16 @@ def test_tpm_projection_preserves_non_media_identity_and_opaque_values() -> None projected, media = split_media_for_tpm(request) assert projected is request assert media == () + + +def test_content_list_flattening_returns_new_messages_and_leaves_the_callers_untouched() -> None: + messages: Final = [ + {"role": "user", "content": [{"type": "text", "text": "Hello"}, {"type": "text", "text": " world"}]}, + {"role": "assistant", "content": "Hi"}, + ] + messages_before: Final = copy.deepcopy(messages) + + flattened: Final = handle_messages_with_content_list_to_str_conversion(messages) + + assert flattened == [{"role": "user", "content": "Hello world"}, {"role": "assistant", "content": "Hi"}] + assert messages == messages_before 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 efb9648cf13..2b1e4a5eed3 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -440,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.""" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index d0e31cdd57a..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. @@ -12574,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 = { @@ -12584,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", @@ -12596,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_logging_utils.py b/tests/unit/litellm_core_utils/test_logging_utils.py index edf0dc7960b..a367c02fd6f 100644 --- a/tests/unit/litellm_core_utils/test_logging_utils.py +++ b/tests/unit/litellm_core_utils/test_logging_utils.py @@ -3,7 +3,10 @@ Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers """ import asyncio, datetime, importlib, litellm, os, pytest_asyncio +import sys import threading +from itertools import chain +from types import FrameType from unittest.mock import MagicMock import pytest @@ -17,7 +20,7 @@ from litellm.litellm_core_utils.logging_utils import ( truncate_base64_in_messages, truncate_base64_in_messages_async, ) -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable, Coroutine, Sequence from datetime import datetime as datetime_assemble_streaming from litellm import( Choices, @@ -28,6 +31,7 @@ from litellm import( ) from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.llms.openai import AllMessageValues from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from typing import Final @@ -99,8 +103,8 @@ class TestTruncateBase64InString: @pytest.mark.parametrize( "text", [ - 'data: {"choices": [{"delta": {"content": "hi"}}]}\n\n' * 50_000, - "data:" * 200_000, + 'data: {"choices": [{"delta": {"content": "hi"}}]}\n\n' * 50_000 + ";base64,", + "data:" * 200_000 + ";base64,", ], ids=["sse_lines", "whitespace_free_prefixes"], ) @@ -241,6 +245,49 @@ def scan_threads(monkeypatch): return threads +def _run_without_yielding(coroutine: Coroutine[object, object, object]) -> object: + try: + coroutine.send(None) + except StopIteration as finished: + return finished.value + coroutine.close() + raise AssertionError("the coroutine yielded to the event loop") + + +def _regex_calls_while(run: Callable[[], object]) -> tuple[object, tuple[str, ...]]: + calls: Final[list[str]] = [] # mutable-ok: sys.setprofile reports each call to a callback that has to record it + previous: Final = sys.getprofile() + + def record(frame: FrameType, event: str, arg: object) -> None: + name: Final = getattr(arg, "__qualname__", "") + if event == "c_call" and name.startswith("Pattern."): + calls.append(name) + + sys.setprofile(record) + try: + result: Final = run() + finally: + sys.setprofile(previous) + return result, tuple(calls) + + +def _tool_result_messages(tool_output: str) -> list[AllMessageValues]: # mutable-ok: the truncator only walks lists + return [ + {"role": "user", "content": "Summarize the rows the tool returned."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "dump_rows", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": tool_output}, + ] + + +_PLAIN_TOOL_OUTPUT_CHARS: Final = 400 * 1024 + + class TestTruncateBase64InMessagesAsync: @pytest.mark.asyncio async def test_large_payload_is_scanned_off_the_event_loop(self, monkeypatch, scan_threads): @@ -268,6 +315,83 @@ class TestTruncateBase64InMessagesAsync: assert scan_threads assert set(scan_threads) == {threading.get_ident()} + @pytest.mark.parametrize( + "row", + [ + '{"id": 7, "name": "widget", "status": "ok"}\n', + "data: {\"choices\": [{\"delta\": {\"content\": \"hi\"}}]}\n\n", + "const total = rows.map((r) => r.price).reduce((a, b) => a + b, 0);\n", + 'data: {"choices": [{"delta": {"content": "a; b"}}]}\n\n', + "Content-Type: text/plain;base64,aGVsbG8=\n", + ], + ids=["json_rows", "sse_lines", "code_with_semicolons", "sse_lines_with_semicolons", "base64_marker_without_data"], + ) + def test_tool_message_without_data_uri_is_copied_without_a_regex_scan(self, row: str): + tool_output: Final = (row * (_PLAIN_TOOL_OUTPUT_CHARS // len(row) + 1))[:_PLAIN_TOOL_OUTPUT_CHARS] + messages: Final = _tool_result_messages(tool_output) + + result, regex_calls = _regex_calls_while(lambda: truncate_base64_in_messages(messages)) + + assert result == messages + assert result is not messages + assert regex_calls == () + + def test_small_data_uri_is_truncated_inline_with_a_regex_scan(self): + payload: Final = "L" * 2_000 + messages: Final = _image_messages(payload) + + result, regex_calls = _regex_calls_while( + lambda: _run_without_yielding(truncate_base64_in_messages_async(messages)) + ) + + assert regex_calls + assert payload not in result[0]["content"][1]["image_url"]["url"] + assert "base64_data truncated" in result[0]["content"][1]["image_url"]["url"] + + def test_large_plain_text_is_copied_inline_without_rescanning_or_a_thread_hop( + self, monkeypatch: pytest.MonkeyPatch, scan_threads: Sequence[int] + ): + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + messages: Final = _tool_result_messages("plain tool output\n" * 100) + + result: Final = _run_without_yielding(truncate_base64_in_messages_async(messages)) + + assert result == messages + assert result is not messages + assert scan_threads == [] + + @pytest.mark.asyncio + async def test_long_plain_text_history_is_copied_off_the_event_loop( + self, monkeypatch: pytest.MonkeyPatch, scan_threads: Sequence[int] + ): + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_STRINGS", 10) + messages: Final = list(chain.from_iterable(_tool_result_messages(f"rows {turn}\n" * 100) for turn in range(5))) + + result: Final = await truncate_base64_in_messages_async(messages) + offload_threads: Final = tuple(scan_threads) + + assert result == messages + assert result is not messages + assert offload_threads + assert threading.get_ident() not in offload_threads + + @pytest.mark.asyncio + async def test_plain_text_over_the_inline_limit_is_copied_off_the_event_loop( + self, monkeypatch: pytest.MonkeyPatch, scan_threads: Sequence[int] + ): + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_INLINE_COPY_MAX_CHARS", 10_000) + messages: Final = _tool_result_messages('data: {"choices": [{"delta": {"content": "a; b"}}]}\n\n' * 400) + + result: Final = await truncate_base64_in_messages_async(messages) + offload_threads: Final = tuple(scan_threads) + + assert result == messages + assert result is not messages + assert offload_threads + assert threading.get_ident() not in offload_threads + @pytest.mark.asyncio async def test_none_and_disabled_truncation_short_circuit(self, monkeypatch, scan_threads): assert await truncate_base64_in_messages_async(None) is None 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_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 9df45b92686..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"]}] diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6df4731e335..4d1699c7fe4 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6879,7 +6879,7 @@ def test_rebuilding_a_priced_stream_bills_the_margin_once(monkeypatch: pytest.Mo 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, ) @@ -6940,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/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index cbb60caa57b..501dd5bd0c6 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1,6 +1,7 @@ import copy import json import logging +from collections.abc import Iterator from itertools import chain from typing import Final from unittest.mock import MagicMock, patch @@ -9849,3 +9850,186 @@ def test_transform_request_tool_config_point_without_tools_drops_the_key(): assert "cache_control_injection_points" not in data assert "tools" not in data assert litellm_params["metadata"] == {} + + +def _five_cache_hint_conversation() -> list[dict[str, object]]: + return [ + {"role": "system", "content": [{"type": "text", "text": "system prompt", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "turn 1", "cache_control": {"type": "ephemeral"}}]}, + {"role": "assistant", "content": [{"type": "text", "text": "answer 1", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "turn 2", "cache_control": {"type": "ephemeral"}}]}, + {"role": "assistant", "content": "answer 2"}, + {"role": "user", "content": [{"type": "text", "text": "turn 3", "cache_control": {"type": "ephemeral"}}]}, + ] + + +def _cached_texts(node: object) -> Iterator[str]: + if isinstance(node, dict): + yield from ((node["text"],) if "cache_control" in node and "text" in node else ()) + yield from chain.from_iterable(map(_cached_texts, node.values())) + elif isinstance(node, list): + yield from chain.from_iterable(map(_cached_texts, node)) + + +@pytest.mark.parametrize( + ("config", "model"), + [ + (AnthropicConfig(), "claude-haiku-5-5"), + (VertexAIAnthropicConfig(), "claude-haiku-5-5"), + (AmazonAnthropicClaudeConfig(), "anthropic.claude-haiku-5-5"), + (AzureAnthropicConfig(), "claude-haiku-5-5"), + ], + ids=["anthropic", "vertex_ai", "bedrock_invoke", "azure_ai"], +) +@pytest.mark.parametrize( + ("drop_params", "expected_cached_texts"), + [ + (True, ["answer 1", "turn 1", "turn 2", "turn 3"]), + ("true", ["answer 1", "turn 1", "turn 2", "turn 3"]), + (False, ["answer 1", "system prompt", "turn 1", "turn 2", "turn 3"]), + ], + ids=[ + "drop_params_keeps_the_newest_four", + "string_drop_params_keeps_the_newest_four", + "without_drop_params_all_five_pass_through", + ], +) +def test_claude_request_over_the_cache_control_cap_drops_the_oldest_hints_only_under_drop_params( + config: AnthropicConfig, + model: str, + drop_params: bool | str, + expected_cached_texts: list[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + messages: Final = _five_cache_hint_conversation() + messages_before: Final = copy.deepcopy(messages) + + body: Final = config.transform_request( + model=model, + messages=messages, + optional_params={}, + litellm_params={"drop_params": drop_params}, + headers={}, + ) + + assert sorted(_cached_texts(body)) == expected_cached_texts + assert json.dumps(body).count('"cache_control"') == len(expected_cached_texts) + assert messages == messages_before + + +def test_claude_cache_control_cap_keeps_a_cached_tool_and_drops_the_oldest_message_hint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + messages: Final = _five_cache_hint_conversation()[1:] + tool: Final = { + "name": "get_weather", + "description": "Get the weather", + "input_schema": {"type": "object", "properties": {}}, + "cache_control": {"type": "ephemeral"}, + } + + body: Final = AnthropicConfig().transform_request( + model="claude-haiku-5-5", + messages=messages, + optional_params={"tools": [tool]}, + litellm_params={}, + headers={}, + ) + + assert body["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert sorted(_cached_texts(body["messages"])) == ["answer 1", "turn 2", "turn 3"] + + +def test_claude_cache_control_cap_drops_a_string_assistant_turns_text_hint_before_its_tool_call_hints( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + hint: Final = {"type": "ephemeral"} + tool_calls: Final = [ + {"id": f"call_{name}", "type": "function", "function": {"name": name, "arguments": "{}"}, "cache_control": hint} + for name in ("lookup", "search") + ] + follow_up: Final = [ + {"type": "text", "text": "follow-up", "cache_control": hint}, + {"type": "text", "text": "and more", "cache_control": hint}, + ] + messages: Final = [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "checking both", "cache_control": hint, "tool_calls": tool_calls}, + {"role": "tool", "tool_call_id": "call_lookup", "content": "found"}, + {"role": "tool", "tool_call_id": "call_search", "content": "none"}, + {"role": "user", "content": follow_up}, + ] + + body: Final = AnthropicConfig().transform_request( + model="claude-haiku-5-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assistant_blocks: Final = body["messages"][1]["content"] + assert [block["type"] for block in assistant_blocks] == ["text", "tool_use", "tool_use"] + assert ["cache_control" in block for block in assistant_blocks] == [False, True, True] + assert json.dumps(body).count('"cache_control"') == 4 + assert json.dumps(body).count('"cache_control"') == 4 + + +def test_claude_cache_control_cap_keeps_the_newest_hints_inside_one_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + messages: Final = [ + { + "role": "user", + "content": [ + *({"type": "text", "text": f"part {n}", "cache_control": {"type": "ephemeral"}} for n in range(1, 6)), + {"type": "text", "text": "Reply with the single word OK."}, + ], + } + ] + messages_before: Final = copy.deepcopy(messages) + + body: Final = AnthropicConfig().transform_request( + model="claude-haiku-5-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert sorted(_cached_texts(body["messages"])) == ["part 2", "part 3", "part 4", "part 5"] + assert json.dumps(body).count('"cache_control"') == 4 + assert messages == messages_before + + +def test_claude_cache_control_cap_does_not_count_a_message_level_hint_that_list_content_never_sends( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + messages: Final = [ + {"role": "system", "content": "system prompt", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "turn 1", "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "content": "answer 1"}, + {"role": "user", "content": "turn 2", "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "content": "answer 2"}, + { + "role": "user", + "content": [{"type": "text", "text": "turn 3", "cache_control": {"type": "ephemeral"}}], + "cache_control": {"type": "ephemeral"}, + }, + ] + + body: Final = AnthropicConfig().transform_request( + model="claude-haiku-5-5", + messages=messages, + optional_params={}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert sorted(_cached_texts(body)) == ["system prompt", "turn 1", "turn 2", "turn 3"] + assert json.dumps(body).count('"cache_control"') == 4 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 7128c04a056..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 diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 0808ce0ec3c..cf0b8a39865 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -56,6 +56,98 @@ def test_is_anthropic_messages_url(url: str, expected: bool) -> None: assert is_anthropic_messages_url(url) is expected +@pytest.mark.parametrize( + ("configured_base", "url", "expected"), + [ + pytest.param(None, "https://api.anthropic.com/v1/messages", 512, id="default-upstream"), + pytest.param(None, "https://api.anthropic.com:443/v1/messages/", 512, id="default-upstream-explicit-port"), + pytest.param("http://relay.test:8080", "https://api.anthropic.com/v1/messages", 512, id="direct-target"), + pytest.param("https://api.anthropic.com", "https://api.anthropic.com/v1/messages", 512, id="configured-default"), + pytest.param("http://relay.test:8080", "http://relay.test:8080/v1/messages", 255, id="configured-upstream"), + pytest.param("http://relay.test:8080/a", "http://relay.test:8080/a/v1/messages", 255, id="configured-prefix"), + pytest.param("https://Relay.Test", "https://relay.test/v1/messages", 255, id="configured-upper-case-host"), + pytest.param("https://relay.test:443", "https://relay.test/v1/messages", 255, id="configured-https-port"), + pytest.param("http://relay.test:80/a", "http://relay.test/a/v1/messages", 255, id="configured-http-port"), + pytest.param("http://relay.test:8080", "http://relay.test:8081/v1/messages", None, id="other-port"), + pytest.param("http://relay.test:notaport", "http://relay.test/v1/messages", None, id="invalid-configured-base"), + pytest.param(None, "https://other-vendor.test/v1/messages", None, id="other-vendor"), + pytest.param(None, "https://api.anthropic.com/v1/messages/count_tokens", None, id="count-tokens"), + ], +) +def test_anthropic_upstream_user_id_limit( + monkeypatch: pytest.MonkeyPatch, configured_base: str | None, url: str, expected: int | None +) -> None: + from litellm.llms.anthropic.common_utils import anthropic_upstream_user_id_limit + + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + if configured_base is not None: + monkeypatch.setenv("ANTHROPIC_API_BASE", configured_base) + + assert anthropic_upstream_user_id_limit(url) == expected + + +def test_anthropic_pass_through_api_base_prefers_the_api_base_variable(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.anthropic.common_utils import anthropic_pass_through_api_base + + monkeypatch.setenv("ANTHROPIC_BASE_URL", "http://fallback.test") + monkeypatch.setenv("ANTHROPIC_API_BASE", "http://preferred.test") + assert anthropic_pass_through_api_base() == "http://preferred.test" + + monkeypatch.delenv("ANTHROPIC_API_BASE") + assert anthropic_pass_through_api_base() == "http://fallback.test" + + +# Each row was sent to https://api.anthropic.com/v1/messages on 2026-10-30, False rows answered 400. +@pytest.mark.parametrize( + ("user_id", "expected"), + [ + pytest.param("user-1", True, id="opaque"), + pytest.param("", True, id="empty"), + pytest.param('{"device_id":"abc","session_id":"def"}', True, id="claude-code-shape"), + pytest.param("1234567890", True, id="digits"), + pytest.param("+14155550100", True, id="phone-shape"), + pytest.param("a@b", True, id="no-dot"), + pytest.param("@b.co", True, id="nothing-before-at-sign"), + pytest.param("a@.co", True, id="nothing-between-at-sign-and-dot"), + pytest.param("a@b.", True, id="nothing-after-dot"), + pytest.param("a @ b.co", True, id="spaces-around-at-sign"), + pytest.param("a\u00a0@b.co", False, id="no-break-space-before-at-sign"), + pytest.param("a@b\n.co", False, id="newline-before-dot"), + pytest.param("line one\nline two", False, id="newline"), + pytest.param("a\tb", False, id="tab"), + pytest.param("a\x1f@b.co", False, id="unit-separator-before-at-sign"), + pytest.param("a@\x1cb.co", False, id="file-separator-after-at-sign"), + pytest.param("a\u200bb", False, id="zero-width-space"), + pytest.param("a@b.co", False, id="email"), + pytest.param("user a@b.co here", False, id="email-inside-text"), + pytest.param('{"device_id":"a@b.co"}', False, id="email-inside-json"), + pytest.param("a@@example.com", False, id="doubled-at-sign"), + pytest.param("@@example.com", False, id="leading-doubled-at-sign"), + pytest.param("a@b@c.co", False, id="two-at-signs"), + pytest.param("a@b..co", False, id="doubled-dot"), + pytest.param("a@.b.co", False, id="dot-after-at-sign"), + pytest.param("a@b.c d", False, id="email-before-space"), + ], +) +def test_anthropic_accepts_user_id(user_id: str, expected: bool) -> None: + from litellm.constants import ANTHROPIC_MAX_USER_ID_LENGTH + from litellm.llms.anthropic.common_utils import anthropic_accepts_user_id + + assert anthropic_accepts_user_id(user_id, ANTHROPIC_MAX_USER_ID_LENGTH) is expected + + +@pytest.mark.parametrize("character", ["a", "\u00e9", "\U0001f600"], ids=["ascii", "two-bytes", "four-bytes"]) +def test_anthropic_accepts_user_id_counts_characters_up_to_the_limit(character: str) -> None: + from litellm.constants import ANTHROPIC_MAX_USER_ID_LENGTH + from litellm.llms.anthropic.common_utils import anthropic_accepts_user_id + + assert anthropic_accepts_user_id(character * ANTHROPIC_MAX_USER_ID_LENGTH, ANTHROPIC_MAX_USER_ID_LENGTH) is True + assert ( + anthropic_accepts_user_id(character * (ANTHROPIC_MAX_USER_ID_LENGTH + 1), ANTHROPIC_MAX_USER_ID_LENGTH) is False + ) + + def test_anthropic_model_not_found_maps_to_not_found_but_authentication_error_does_not() -> None: with respx.mock(assert_all_called=True) as router: router.post("https://api.anthropic.com/v1/messages").mock( @@ -3946,10 +4038,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 +4050,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 +4061,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 +4086,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 +4101,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 +4130,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", 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..041c625eaef 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, @@ -21,7 +26,6 @@ from litellm.types.decisions import ( DecisionsIRScoreAnswer, DecisionsRequestBody, OpenAIDecisionRequestBody, - UnsupportedDecisionsRequest, ) from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -49,6 +53,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)) @@ -219,7 +278,7 @@ def test_a_malformed_systemone_response_names_the_field_at_fault( assert [item["loc"] for item in error.value.errors()] == [field] -def test_each_media_kind_is_gated_by_its_own_support_flag() -> None: +def test_each_media_kind_is_sent_in_its_own_field() -> None: audio: Final = ["data:audio/wav;base64,YQ=="] videos: Final = ["data:video/mp4;base64,YQ=="] audio_only: Final = systemone_request_to_ir(_SYSTEMONE_BODY.validate_python({**_SYSTEMONE_REQUEST, "audio": audio})) @@ -227,22 +286,16 @@ def test_each_media_kind_is_gated_by_its_own_support_flag() -> None: _SYSTEMONE_BODY.validate_python({**_SYSTEMONE_REQUEST, "videos": videos}) ) - assert ir_to_systemone_request("clef-omni", audio_only, supports_audio=True) == { + assert ir_to_systemone_request("clef-omni", audio_only) == { "model": "clef-omni", **_SYSTEMONE_REQUEST, "audio": audio, } - assert ir_to_systemone_request("clef-omni", video_only, supports_video=True) == { + assert ir_to_systemone_request("clef-omni", video_only) == { "model": "clef-omni", **_SYSTEMONE_REQUEST, "videos": videos, } - assert isinstance( - ir_to_systemone_request("clef-omni", audio_only, supports_video=True), UnsupportedDecisionsRequest - ) - assert isinstance( - ir_to_systemone_request("clef-omni", video_only, supports_audio=True), UnsupportedDecisionsRequest - ) _ASSET_REQUEST: Final[Mapping[str, object]] = { diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 524caaba7a1..9ab7060d766 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -11259,3 +11259,86 @@ def test_converse_sends_only_the_header_betas_when_header_and_body_both_carry_be "thinking-binding-controls-2026-08-01", "thinking-display-updates-2026-08-18", ] + + +def _five_cache_hint_conversation() -> list[dict[str, object]]: + return [ + {"role": "system", "content": [{"type": "text", "text": "system prompt", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "turn 1", "cache_control": {"type": "ephemeral"}}]}, + {"role": "assistant", "content": [{"type": "text", "text": "answer 1", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "turn 2", "cache_control": {"type": "ephemeral"}}]}, + {"role": "assistant", "content": "answer 2"}, + {"role": "user", "content": [{"type": "text", "text": "turn 3", "cache_control": {"type": "ephemeral"}}]}, + ] + + +@pytest.mark.parametrize( + ("drop_params", "expected_system_cache_points"), + [(True, 0), ("true", 0), (False, 1)], + ids=[ + "drop_params_drops_the_oldest_hint", + "string_drop_params_drops_the_oldest_hint", + "without_drop_params_all_five_pass_through", + ], +) +def test_converse_claude_request_over_the_cache_control_cap_drops_the_oldest_hints_only_under_drop_params( + drop_params: bool | str, expected_system_cache_points: int, monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + messages: Final = _five_cache_hint_conversation() + messages_before: Final = copy.deepcopy(messages) + + request: Final = AmazonConverseConfig().transform_request( + model="anthropic.claude-haiku-5-5", + messages=messages, + optional_params={}, + litellm_params={"drop_params": drop_params}, + headers={}, + ) + + assert len(list(_cache_points(request["system"]))) == expected_system_cache_points + assert [len(list(_cache_points(message))) for message in request["messages"]] == [1, 1, 1, 0, 1] + assert messages == messages_before + + +def test_converse_claude_tool_message_with_several_marked_blocks_counts_as_one_cache_point_against_the_cap( + monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "policy", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "question", "cache_control": {"type": "ephemeral"}}, + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + {"type": "text", "text": "part 1", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "part 2", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "part 3", "cache_control": {"type": "ephemeral"}}, + ], + }, + ] + messages_before: Final = copy.deepcopy(messages) + + request: Final = AmazonConverseConfig().transform_request( + model="anthropic.claude-haiku-5-5", + messages=messages, + optional_params={ + "tools": [{"type": "function", "function": {"name": "lookup", "parameters": {"type": "object"}}}] + }, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert len(list(_cache_points(request["messages"]))) == 3 + assert messages == messages_before diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 20c0e10cb81..38c69e0ce27 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -1,5 +1,5 @@ -import asyncio, importlib, litellm, litellm.litellm_core_utils.get_model_cost_map as bedrock_govcloud_model_cost_map, pytest +import asyncio, importlib, litellm, pytest from typing import Final @@ -1563,7 +1563,6 @@ def _pr4_bedrock_env(monkeypatch: pytest.MonkeyPatch) -> None: def _use_local_model_cost_map() -> Iterator[None]: with pytest.MonkeyPatch.context() as monkeypatch: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - importlib.reload(bedrock_govcloud_model_cost_map) importlib.reload(litellm) yield 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/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 19234767b19..7a6f9c7dc18 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -2773,3 +2773,66 @@ def test_claude_json_schema_stream_split_across_json_tools_keeps_the_non_stream_ choices: Final = _json_schema_stream_choices(respx_mock, call, split_stream) assert [choice.finish_reason for choice in choices if choice.finish_reason] == [answer.finish_reason] + + +@pytest.mark.parametrize( + ("model", "base_model"), + [("databricks-claude-haiku-5-5", None), ("support-bot-endpoint", "databricks/databricks-claude-haiku-5-5")], + ids=["claude_in_the_endpoint_name", "claude_named_only_by_base_model"], +) +def test_databricks_claude_request_over_the_cache_control_cap_keeps_the_newest_four_hints_under_drop_params( + model: str, base_model: str | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + hint: Final = {"type": "ephemeral"} + messages: Final = [ + {"role": "system", "content": "system prompt", "cache_control": hint}, + {"role": "user", "content": "turn 1", "cache_control": hint}, + {"role": "assistant", "content": "answer 1", "cache_control": hint}, + {"role": "user", "content": "turn 2", "cache_control": hint}, + {"role": "assistant", "content": "answer 2"}, + {"role": "user", "content": "turn 3", "cache_control": hint}, + ] + messages_before: Final = copy.deepcopy(messages) + + body: Final = DatabricksConfig().transform_request( + model=model, + messages=messages, + optional_params={}, + litellm_params={"drop_params": True, "custom_llm_provider": "databricks", "base_model": base_model}, + headers={}, + ) + + assert "cache_control" not in json.dumps(body["messages"][0]) + assert json.dumps(body["messages"]).count('"cache_control"') == 4 + assert messages == messages_before + + +@pytest.mark.parametrize( + "optional_params", + [{"cache_control": {"type": "ephemeral"}}, {"extra_body": {"cache_control": {"type": "ephemeral"}}}], + ids=["top_level_cache_control", "extra_body_cache_control"], +) +def test_databricks_claude_cache_control_cap_keeps_room_for_a_top_level_cache_control_under_drop_params( + optional_params: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + hint: Final = {"type": "ephemeral"} + messages: Final = [ + {"role": "user", "content": "turn 1", "cache_control": hint}, + {"role": "assistant", "content": "answer 1", "cache_control": hint}, + {"role": "user", "content": "turn 2", "cache_control": hint}, + {"role": "assistant", "content": "answer 2", "cache_control": hint}, + {"role": "user", "content": "turn 3"}, + ] + + body: Final = DatabricksConfig().transform_request( + model="databricks-claude-haiku-5-5", + messages=messages, + optional_params=copy.deepcopy(optional_params), + litellm_params={"drop_params": True, "custom_llm_provider": "databricks"}, + headers={}, + ) + + assert "cache_control" not in json.dumps(body["messages"][0]) + assert json.dumps(body["messages"]).count('"cache_control"') == 3 diff --git a/tests/unit/llms/deepinfra/test_deepinfra_chat_transformation.py b/tests/unit/llms/deepinfra/test_deepinfra_chat_transformation.py index 7ec9216df7b..773c1d056f0 100644 --- a/tests/unit/llms/deepinfra/test_deepinfra_chat_transformation.py +++ b/tests/unit/llms/deepinfra/test_deepinfra_chat_transformation.py @@ -1,6 +1,8 @@ import asyncio +import copy import json import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -200,3 +202,38 @@ async def test_deepinfra_tool_message_content_transformation_async(): print(f"✓ Async test passed: {tool_message['content']}") print("\n✅ DeepInfra async tool message transformation test passed!") + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_deepinfra_request_stringifies_tool_content_without_touching_the_callers_message(is_async: bool) -> None: + from litellm.llms.deepinfra.chat.transformation import DeepInfraConfig + + config: Final = DeepInfraConfig() + messages: Final = [ + {"role": "user", "content": "Calculate 10 + 10"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "calculator", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": [{"type": "text", "text": "20"}]}, + ] + messages_before: Final = copy.deepcopy(messages) + request_kwargs: Final = { + "model": "deepseek-ai/DeepSeek-V3.2", + "messages": messages, + "optional_params": {}, + "litellm_params": {}, + "headers": {}, + } + + body: Final = ( + await config.async_transform_request(**request_kwargs) + if is_async + else config.transform_request(**request_kwargs) + ) + + assert body["messages"][2]["content"] == "20" + assert messages == messages_before 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 ed166d206c8..dd352aac66d 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,3 +1,4 @@ +import copy import json import logging from collections.abc import Callable @@ -537,6 +538,26 @@ def test_tool_search_reply_sent_to_fireworks_drops_litellm_response_fields(): } +def test_cache_control_is_stripped_from_the_request_but_kept_on_the_callers_messages(): + hint: Final = {"type": "ephemeral"} + messages: Final = [ + {"role": "system", "content": [{"type": "text", "text": "a long cached prompt", "cache_control": hint}]}, + {"role": "user", "content": "hi", "cache_control": hint}, + ] + caller_messages_before: Final = json.dumps(messages) + + request: Final = FireworksAIConfig().transform_request( + model="accounts/fireworks/models/gpt-oss-120b", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert "cache_control" not in json.dumps(request["messages"]) + assert json.dumps(messages) == caller_messages_before + + # ----------------------------------------------------------------------------- # Regression tests for legacy / OpenAPI $ref defs in tool parameters. # @@ -2477,3 +2498,43 @@ def test_translate_developer_role_follows_the_deployment_model_info_over_the_cos ) assert [message["role"] for message in messages] == expected_roles + + +def test_fireworks_request_rewrites_tools_without_touching_the_callers_tool() -> None: + tools: Final = [ + { + "type": "function", + "function": { + "name": "create_thing", + "strict": True, + "parameters": { + "type": "object", + "properties": {"thing": {"$ref": "#/definitions/Thing"}}, + "definitions": {"Thing": {"type": "object", "properties": {"id": {"type": "string"}}}}, + }, + }, + } + ] + tools_before: Final = copy.deepcopy(tools) + + request: Final = FireworksAIConfig().transform_request( + model="accounts/fireworks/models/gpt-oss-120b", + messages=[{"role": "user", "content": "Make a thing"}], + optional_params={"tools": tools}, + litellm_params={}, + headers={}, + ) + + assert request["tools"] == [ + { + "type": "function", + "function": { + "name": "create_thing", + "parameters": { + "type": "object", + "properties": {"thing": {"type": "object", "properties": {"id": {"type": "string"}}}}, + }, + }, + } + ] + assert tools == tools_before diff --git a/tests/unit/llms/groq/chat/test_groq_chat_transformation.py b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py index 7aa8c65215e..a493a259f39 100644 --- a/tests/unit/llms/groq/chat/test_groq_chat_transformation.py +++ b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py @@ -1,3 +1,4 @@ +import copy import logging from typing import Final from unittest.mock import patch @@ -589,3 +590,19 @@ async def test_streamed_reasoning_reply_can_be_replayed(is_async: bool, with_too } ] assert reply.reasoning_content == "I should check the weather." + + +def test_groq_request_drops_null_assistant_fields_without_touching_the_callers_messages() -> None: + messages: Final = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello", "function_call": None, "tool_calls": None}, + {"role": "user", "content": "Again"}, + ] + messages_before: Final = copy.deepcopy(messages) + + body: Final = GroqChatConfig().transform_request( + model="openai/gpt-oss-120b", messages=messages, optional_params={}, litellm_params={}, headers={} + ) + + assert body["messages"][1] == {"role": "assistant", "content": "Hello"} + assert messages == messages_before diff --git a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index c792d5dffcd..fc5ae3631ff 100644 --- a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -1,4 +1,6 @@ +import copy import json +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -397,3 +399,43 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema(): assert tools[0]["function"]["name"] == "search" assert tools[0]["function"]["description"] == "Search docs" assert tools[0]["function"]["parameters"] == input_schema + + +def _strict_weather_tool() -> dict[str, object]: + return { + "type": "function", + "function": { + "name": "get_weather", + "strict": True, + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + }, + }, + } + + +def test_hosted_vllm_tools_are_rewritten_without_touching_the_callers_tool() -> None: + tools: Final = [_strict_weather_tool()] + tools_before: Final = copy.deepcopy(tools) + + optional_params: Final = HostedVLLMChatConfig().map_openai_params( + non_default_params={"tools": tools}, optional_params={}, model="gpt-oss-120b", drop_params=False + ) + + assert optional_params["tools"] == [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ] + assert tools == tools_before 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 720a1b4f458..af18a9caf3d 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -1,3 +1,4 @@ +import copy import json import logging from collections.abc import Mapping @@ -1212,7 +1213,10 @@ def test_mistral_sends_only_the_text_items_of_a_native_thinking_chunk_beside_an_ "content": [ { "type": "thinking", - "thinking": [{"type": "text", "text": "It has whiskers."}, {"type": "reference", "reference_ids": [1]}], + "thinking": [ + {"type": "text", "text": "It has whiskers."}, + {"type": "reference", "reference_ids": [1]}, + ], }, {"type": "text", "text": "A cat."}, ], @@ -1672,6 +1676,18 @@ 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 test_mistral_request_drops_a_user_name_without_touching_the_callers_message() -> None: + messages: Final = [{"role": "user", "content": "Hi"}, {"role": "user", "content": "", "name": "alice"}] + messages_before: Final = copy.deepcopy(messages) + + body: Final = MistralConfig().transform_request( + model="mistral-medium-latest", messages=messages, optional_params={}, litellm_params={}, headers={} + ) + + assert "name" not in json.dumps(body["messages"]) + assert messages == messages_before + + 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() diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index af1b85cf325..6a80f29ac43 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -856,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. diff --git a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py index 5b85a94b24d..79fe7061297 100644 --- a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py @@ -1824,6 +1824,94 @@ class TestToolSchemaCombinatorFlatteningForOpenAI: assert request["tools"][0] is tool +class TestCacheControlStrippingDoesNotMutateCallerInput: + """ + Stripping cache_control for a provider that cannot use it must not reach back + into the caller's own message and tool objects. + + filter_value_from_dict deletes the key in place and recurses into nested dicts + and lists, and remove_cache_control_flag_from_messages_and_tools assigned the + result back into the caller's list, so one call to any OpenAI-compatible + provider permanently stripped cache_control from a list the caller still holds. + Reusing that list on Anthropic or Bedrock afterwards then silently lost prompt + caching, with no error and full-price billing. + """ + + def setup_method(self): + self.config = OpenAIGPTConfig() + + @pytest.fixture(autouse=True) + def _clean_openai_base_env(self, monkeypatch): + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None, raising=False) + + @staticmethod + def _messages(): + return [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "a long cached system prompt", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": "Hello", + "cache_control": {"type": "ephemeral"}, + }, + ] + + @staticmethod + def _tools(): + return [ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object"}}, + "cache_control": {"type": "ephemeral"}, + } + ] + + def _transform(self, messages, tools): + return self.config.transform_request( + model="gpt-4o", + messages=messages, + optional_params={"tools": tools}, + litellm_params={"custom_llm_provider": "openai", "api_base": None}, + headers={}, + ) + + def test_caller_messages_and_tools_keep_cache_control(self): + messages, tools = self._messages(), self._tools() + messages_before, tools_before = copy.deepcopy(messages), copy.deepcopy(tools) + + request = self._transform(messages, tools) + + assert "cache_control" not in json.dumps(request) + assert messages == messages_before + assert tools == tools_before + + def test_a_later_anthropic_call_still_sees_cache_control(self): + """The user-visible consequence: the same message list reused on a provider + that does support caching must still carry the cache breakpoints.""" + messages = self._messages() + + self._transform(messages, self._tools()) + + anthropic_body = litellm.AnthropicConfig().transform_request( + model="claude-3-5-sonnet-20240620", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "cache_control" in json.dumps(anthropic_body) + + def _developer_message_after_a_user_turn() -> list[dict[str, str]]: return [ {"role": "system", "content": "You are terse."}, @@ -1916,3 +2004,53 @@ def test_completion_leaves_the_callers_closing_developer_message_unchanged_on_a_ ) assert json.loads(chat.calls.last.request.content)["messages"][-1]["role"] == "user" assert messages == sent + + +def _user_turn_with_litellm_only_media_fields() -> list[dict[str, object]]: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe these"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png", "format": "image/png"}}, + {"type": "image_url", "image_url": "https://example.com/dog.png"}, + { + "type": "file", + "file": {"file_data": "data:application/pdf;base64,JVBERi0=", "format": "application/pdf"}, + }, + ], + } + ] + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_openai_request_rewrites_media_parts_without_touching_the_callers_message( + is_async: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None, raising=False) + config: Final = OpenAIGPTConfig() + messages: Final = _user_turn_with_litellm_only_media_fields() + messages_before: Final = copy.deepcopy(messages) + request_kwargs: Final = { + "model": "gpt-5.5", + "messages": messages, + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai", "api_base": None}, + "headers": {}, + } + + body: Final = ( + await config.async_transform_request(**request_kwargs) + if is_async + else config.transform_request(**request_kwargs) + ) + + assert body["messages"][0]["content"] == [ + {"type": "text", "text": "Describe these"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + {"type": "image_url", "image_url": {"url": "https://example.com/dog.png"}}, + {"type": "file", "file": {"file_data": "data:application/pdf;base64,JVBERi0=", "filename": "my_file.pdf"}}, + ] + assert messages == messages_before diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index f5a33f87fd0..41c54bf6cc9 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -1,3 +1,4 @@ +import copy import json import logging from collections.abc import Callable, Mapping @@ -541,6 +542,50 @@ class TestOpenAIResponsesAPIConfig: assert "cache_control" not in result["tools"][0] assert result["tools"][0]["name"] == "get_weather" + def test_transform_does_not_strip_cache_control_from_the_callers_input(self): + """Stripping for OpenAI must not reach back into the caller's own objects. + + `filter_value_from_dict` deletes the key in place and recurses, and + `_validate_input_param` passes plain dict items through by reference, so + the caller's input list used to lose its cache breakpoints. Reusing that + list on Anthropic or Bedrock afterwards then silently lost prompt caching. + Same defect as the Chat Completions path. + """ + input_with_cache_control = [ + { + "role": "system", + "content": [ + { + "type": "input_text", + "text": "a long cached system prompt", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + tools_with_cache_control = [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object"}, + "cache_control": {"type": "ephemeral"}, + } + ] + input_before = copy.deepcopy(input_with_cache_control) + tools_before = copy.deepcopy(tools_with_cache_control) + + result = self.config.transform_responses_api_request( + model=self.model, + input=input_with_cache_control, + response_api_optional_request_params={"tools": tools_with_cache_control}, + litellm_params={}, + headers={}, + ) + + assert "cache_control" not in json.dumps(result) + assert input_with_cache_control == input_before + assert tools_with_cache_control == tools_before + def test_transform_preserves_input_without_cache_control(self): """Inputs without cache_control must pass through unmodified.""" input_clean = [ @@ -4607,10 +4652,10 @@ def test_cache_control_still_marks_a_part_whose_own_marker_is_malformed_under_dr } -def _caller_input_after_transform(model: str, marker: object) -> str: +def _hinted_caller_input(marker: object) -> list[dict[str, object]]: hint: Final = {"type": "ephemeral"} marked: Final = {} if marker is None else {"prompt_cache_breakpoint": marker} - input: Final = [ + return [ { "type": "message", "role": "user", @@ -4626,6 +4671,10 @@ def _caller_input_after_transform(model: str, marker: object) -> str: "content": [{"type": "output_text", "text": "ok", "cache_control": hint, **marked}], }, ] + + +def _caller_input_after_transform(model: str, marker: object) -> str: + input: Final = _hinted_caller_input(marker) _CONFIGS[0].transform_responses_api_request( model=model, input=cast("ResponseInputParam", input), # cast-ok: plain dicts shaped like the SDK's input items @@ -4647,6 +4696,15 @@ def test_a_marker_leaves_the_caller_request_that_fallbacks_and_logs_read_as_an_u assert _caller_input_after_transform(model, marker).replace(marker_json, "") == unmarked +@pytest.mark.parametrize("model", ("gpt-4.1", "gpt-6.1-sol")) +@pytest.mark.parametrize("marker", (None, _VALID_MARKERS[0], "yes"), ids=("unmarked", "explicit", "malformed")) +def test_the_caller_input_keeps_its_cache_control_hints_for_a_fallback_deployment( + monkeypatch: pytest.MonkeyPatch, model: str, marker: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + assert _caller_input_after_transform(model, marker) == json.dumps(_hinted_caller_input(marker)) + + @pytest.mark.asyncio @pytest.mark.usefixtures(httpx_transport.__name__) async def test_aresponses_forwards_previous_response_id_to_openai() -> None: 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_o_series_transformation.py b/tests/unit/llms/openai/test_o_series_transformation.py index e78e1a6f12b..9d39dfda60a 100644 --- a/tests/unit/llms/openai/test_o_series_transformation.py +++ b/tests/unit/llms/openai/test_o_series_transformation.py @@ -1,3 +1,4 @@ +import copy from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -198,3 +199,23 @@ def test_translate_developer_role_keeps_developer_messages_untouched() -> None: ) == messages ) + + +def test_o_series_request_turns_system_into_user_without_touching_the_callers_messages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model: Final = "o-series-test-model-without-system-role" + monkeypatch.setitem( + litellm.model_cost, + model, + {"litellm_provider": "openai", "mode": "chat", "supports_system_messages": False}, + ) + messages: Final = [{"role": "system", "content": "Be terse"}, {"role": "user", "content": "Hi"}] + messages_before: Final = copy.deepcopy(messages) + + body: Final = OpenAIOSeriesConfig().transform_request( + model=model, messages=messages, optional_params={}, litellm_params={}, headers={} + ) + + assert [message["role"] for message in body["messages"]] == ["user", "user"] + assert messages == messages_before 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_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_hostinger_router_provider.py b/tests/unit/llms/openai_like/test_hostinger_router_provider.py new file mode 100644 index 00000000000..cbc16c8567a --- /dev/null +++ b/tests/unit/llms/openai_like/test_hostinger_router_provider.py @@ -0,0 +1,706 @@ +import copy +import json +from typing import Final + +import pytest +import respx + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache + + +def test_hostinger_router_provider_resolution(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("HOSTINGER_ROUTER_API_KEY", "hostinger-test-key") + + model, provider, api_key, api_base = get_llm_provider( + model="hostinger_router/claude-sonnet-5", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "claude-sonnet-5" + assert provider == "hostinger_router" + assert api_key == "hostinger-test-key" + assert api_base == "https://router.hostinger.com/v1" + + +def test_hostinger_router_provider_keeps_explicit_credentials(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("HOSTINGER_ROUTER_API_KEY", "hostinger-env-key") + + _, provider, api_key, api_base = get_llm_provider( + model="hostinger_router/claude-sonnet-5", + custom_llm_provider=None, + api_base="https://router.internal.example/v1", + api_key="hostinger-explicit-key", + ) + + assert provider == "hostinger_router" + assert api_key == "hostinger-explicit-key" + assert api_base == "https://router.internal.example/v1" + + +@pytest.mark.parametrize("model", ["gpt-5-nano", "claude-sonnet-5"]) +def test_hostinger_router_api_base_keeps_unprefixed_model_provider(model: str): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + _, provider_without_api_base, _, _ = get_llm_provider(model=model) + _, provider_with_hostinger_api_base, _, _ = get_llm_provider( + model=model, + api_base="https://router.hostinger.com/v1", + api_key="hostinger-test-key", + ) + + assert provider_with_hostinger_api_base == provider_without_api_base + + +HOSTINGER_ROUTER_MODELS = tuple(sorted(name for name in litellm.model_cost if name.startswith("hostinger_router/"))) +HOSTINGER_ROUTER_CHAT_MODELS = tuple( + name for name in HOSTINGER_ROUTER_MODELS if litellm.get_model_info(name)["mode"] == "chat" +) + + +@pytest.mark.parametrize("model", HOSTINGER_ROUTER_CHAT_MODELS) +def test_hostinger_router_model_cost_and_capabilities(model: str): + from litellm.cost_calculator import cost_per_token + + prompt_cost, completion_cost = cost_per_token( + model=model, + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + custom_llm_provider="hostinger_router", + ) + model_info = litellm.get_model_info(model) + + assert prompt_cost == pytest.approx(model_info["input_cost_per_token"] * 1_000_000) + assert completion_cost == pytest.approx(model_info["output_cost_per_token"] * 1_000_000) + assert model_info["output_cost_per_token"] > 0 + assert model_info["max_tokens"] == model_info["max_output_tokens"] <= model_info["max_input_tokens"] + assert model_info["litellm_provider"] == "hostinger_router" + assert model_info["mode"] == "chat" + if model_info.get("supports_prompt_caching"): + assert 0 < model_info["cache_read_input_token_cost"] < model_info["input_cost_per_token"] + capability_checks: Final = ( + ("supports_function_calling", litellm.supports_function_calling), + ("supports_reasoning", litellm.supports_reasoning), + ("supports_response_schema", litellm.supports_response_schema), + ("supports_vision", litellm.supports_vision), + ) + assert {name: check(model) for name, check in capability_checks} == { + name: model_info.get(name) is True for name, _ in capability_checks + } + + +def test_hostinger_router_chat_completions_request(): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_hostinger", + "object": "chat.completion", + "created": 1_789_550_000, + "model": "claude-sonnet-5", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello from Hostinger Router"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}, + }, + ) + response: Final = litellm.completion( + model="hostinger_router/claude-sonnet-5", + messages=[{"role": "user", "content": "Say hello"}], + api_key="hostinger-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert request.headers["authorization"] == "Bearer hostinger-test-key" + assert body["model"] == "claude-sonnet-5" + assert response.choices[0].message.content == "Hello from Hostinger Router" + + +def test_hostinger_router_chat_completions_streaming_request(): + chunks = [ + 'data: {"id":"chatcmpl_hostinger","object":"chat.completion.chunk","created":1789550000,"model":"claude-sonnet-5","choices":[{"index":0,"delta":{"role":"assistant","content":"Hi"},"finish_reason":null}]}\n\n', + 'data: {"id":"chatcmpl_hostinger","object":"chat.completion.chunk","created":1789550000,"model":"claude-sonnet-5","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":1,"total_tokens":5}}\n\n', + "data: [DONE]\n\n", + ] + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/chat/completions").respond( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(chunks).encode(), + ) + response: Final = litellm.completion( + model="hostinger_router/claude-sonnet-5", + messages=[{"role": "user", "content": "Say hello"}], + stream=True, + api_key="hostinger-test-key", + ) + collected = [chunk for chunk in response] + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert body["stream"] is True + assert any(chunk.choices and chunk.choices[0].delta.content == "Hi" for chunk in collected) + + +def test_hostinger_router_responses_request(): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/responses").respond( + 200, + json={ + "id": "resp_hostinger", + "object": "response", + "created_at": 1_789_550_000, + "model": "gpt-5.2", + "status": "completed", + "output": [ + { + "id": "msg_hostinger", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Hello from Hostinger Router", "annotations": []}], + } + ], + "usage": {"input_tokens": 4, "output_tokens": 3, "total_tokens": 7}, + }, + ) + response: Final = litellm.responses( + model="hostinger_router/gpt-5.2", + input="Say hello", + api_key="hostinger-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert str(request.url) == "https://router.hostinger.com/v1/responses" + assert request.headers["authorization"] == "Bearer hostinger-test-key" + assert body["model"] == "gpt-5.2" + assert body["input"] == "Say hello" + assert response.output[0].content[0].text == "Hello from Hostinger Router" + + +EMBEDDING_RESPONSE: Final = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, +} + + +def test_hostinger_router_embeddings_request(): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + response: Final = litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert request.headers["authorization"] == "Bearer hostinger-test-key" + assert body["model"] == "text-embedding-3-small" + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_hostinger_router_embeddings_send_extra_headers_as_http_headers(): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={"X-Trace-Id": "trace-123"}, + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert request.headers["x-trace-id"] == "trace-123" + assert request.headers["authorization"] == "Bearer hostinger-test-key" + assert request.headers["content-type"] == "application/json" + assert set(body) == {"model", "input"} + + +def test_hostinger_router_embeddings_extra_headers_replace_key_headers_of_any_case(): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={"authorization": "Bearer caller-key", "x-trace-id": "trace-1"}, + ) + + request: Final = route.calls.last.request + assert request.headers.get_list("authorization") == ["Bearer caller-key"] + assert request.headers.get_list("content-type") == ["application/json"] + assert request.headers.get_list("x-trace-id") == ["trace-1"] + + +@pytest.mark.parametrize(("configured", "sent"), [(12345, "12345"), (1.5, "1.5")]) +def test_hostinger_router_embeddings_send_numeric_extra_headers_as_text(configured: int | float, sent: str): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={"X-Org-Id": configured}, + ) + + assert route.calls.last.request.headers.get_list("x-org-id") == [sent] + + +@pytest.mark.parametrize("configured", [True, None, ["a", "b"]]) +def test_hostinger_router_embeddings_refuse_non_text_extra_headers(configured: object): + with respx.mock(assert_all_called=False) as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + with pytest.raises(litellm.UnsupportedParamsError, match="X-Debug") as refused: + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={"X-Debug": configured}, + ) + + assert refused.value.status_code == 400 + assert route.call_count == 0 + + +@pytest.mark.parametrize("configured", [True, None, ["a", "b"]]) +def test_hostinger_router_embeddings_drop_non_text_extra_headers_under_drop_params(configured: object): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={"X-Debug": configured, "X-Trace-Id": "trace-1"}, + drop_params=True, + ) + + request: Final = route.calls.last.request + assert "x-debug" not in request.headers + assert request.headers.get_list("x-trace-id") == ["trace-1"] + assert set(json.loads(request.content)) == {"model", "input"} + + +TRANSPORT_EXTRA_HEADERS: Final = [ + ("Host", "other.example"), + ("content-length", "1"), + ("Transfer-Encoding", "chunked"), + ("Connection", "close"), + ("Keep-Alive", "timeout=5"), + ("Proxy-Connection", "close"), + ("TE", "gzip"), + ("Trailer", "X-Checksum"), + ("Upgrade", "websocket"), + ("Expect", "100-continue"), +] + + +@pytest.mark.parametrize(("name", "configured"), TRANSPORT_EXTRA_HEADERS) +def test_hostinger_router_embeddings_ignore_transport_extra_headers(name: str, configured: str): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={name: configured, "X-Trace-Id": "trace-1"}, + ) + + request: Final = route.calls.last.request + assert request.headers.get_list("host") == ["router.hostinger.com"] + assert request.headers.get_list("content-length") == [str(len(request.content))] + assert "transfer-encoding" not in request.headers + assert configured not in request.headers.get_list(name) + assert request.headers.get_list("x-trace-id") == ["trace-1"] + + +MALFORMED_HEADER_NAMES: Final = ["Host:evil.example", "X Note", "X-Caf\u00e9"] + + +@pytest.mark.parametrize("name", MALFORMED_HEADER_NAMES) +def test_hostinger_router_embeddings_refuse_malformed_extra_header_names(name: str): + with respx.mock(assert_all_called=False) as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + with pytest.raises(litellm.UnsupportedParamsError) as refused: + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={name: "x"}, + ) + + assert refused.value.status_code == 400 + assert repr(name) in refused.value.message + assert route.call_count == 0 + + +@pytest.mark.parametrize("name", MALFORMED_HEADER_NAMES) +def test_hostinger_router_embeddings_drop_malformed_extra_header_names_under_drop_params(name: str): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={name: "x", "X-Trace-Id": "trace-1"}, + drop_params=True, + ) + + request: Final = route.calls.last.request + assert request.headers.get_list("host") == ["router.hostinger.com"] + assert request.headers.get_list("x-trace-id") == ["trace-1"] + assert name.lower() not in request.headers.keys() + + +@pytest.mark.parametrize("configured", ["a\r\nHost: evil.example", "line\nbreak", "caf\u00e9", "a\x7fb"]) +def test_hostinger_router_embeddings_drop_extra_header_values_http_cannot_carry(configured: str): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={"x-litellm-user_api_key_alias": configured, "X-Trace-Id": "trace-1"}, + ) + + request: Final = route.calls.last.request + assert request.headers.get_list("host") == ["router.hostinger.com"] + assert request.headers.get_list("x-trace-id") == ["trace-1"] + assert "x-litellm-user_api_key_alias" not in request.headers + + +@pytest.mark.parametrize( + ("configured", "sent"), + [(" x ", "x"), ("\t", ""), (" a\tb ", "a\tb"), ("env-secret\n", "env-secret"), ("\r\n x \r\n", "x")], +) +def test_hostinger_router_embeddings_trim_whitespace_around_extra_header_values(configured: str, sent: str): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + extra_headers={"X-Rig": configured}, + ) + + assert route.calls.last.request.headers.get_list("x-rig") == [sent] + + +def test_hostinger_router_embeddings_extra_headers_stay_with_their_request(): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "shared", + "litellm_params": { + "model": "hostinger_router/text-embedding-3-small", + "api_key": "hostinger-test-key", + "headers": {"X-Org-Id": "org-1"}, + }, + } + ] + ) + + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + router.embedding( + model="shared", input=["first"], extra_headers={"Authorization": "Bearer caller-a", "X-Leak": "a"} + ) + router.embedding(model="shared", input=["second"]) + + first: Final = route.calls[0].request + second: Final = route.calls[1].request + assert first.headers.get_list("x-leak") == ["a"] + assert second.headers.get_list("authorization") == ["Bearer hostinger-test-key"] + assert "x-leak" not in second.headers + assert second.headers.get_list("x-org-id") == ["org-1"] + + +@pytest.mark.parametrize("param", ["headers", "extra_headers"]) +def test_hostinger_router_embeddings_send_headers_given_as_name_value_pairs(param: str): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + **{param: [["X-Trace-Id", "trace-1"]]}, + ) + + request: Final = route.calls.last.request + assert request.headers.get_list("x-trace-id") == ["trace-1"] + assert request.headers.get_list("authorization") == ["Bearer hostinger-test-key"] + + +@pytest.mark.parametrize("param", ["headers", "extra_headers"]) +def test_hostinger_router_embeddings_refuse_headers_that_are_not_an_object(param: str): + with respx.mock(assert_all_called=False) as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + with pytest.raises(litellm.UnsupportedParamsError, match=rf"\b{param} must be a JSON object") as refused: + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + **{param: "X-Trace-Id: trace-1"}, + ) + + assert refused.value.status_code == 400 + assert route.call_count == 0 + + +@pytest.mark.parametrize("param", ["headers", "extra_headers"]) +def test_hostinger_router_embeddings_drop_headers_that_are_not_an_object_under_drop_params(param: str): + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + litellm.embedding( + model="hostinger_router/text-embedding-3-small", + input=["Say hello"], + api_key="hostinger-test-key", + drop_params=True, + **{param: "X-Trace-Id: trace-1"}, + ) + + request: Final = route.calls.last.request + assert "x-trace-id" not in request.headers + assert request.headers.get_list("authorization") == ["Bearer hostinger-test-key"] + assert set(json.loads(request.content)) == {"model", "input"} + + +def test_hostinger_router_embeddings_never_borrow_another_providers_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("HOSTINGER_ROUTER_API_KEY", raising=False) + monkeypatch.setenv("OPENAI_LIKE_API_KEY", "other-provider-key") + monkeypatch.setenv("OPENAI_API_KEY", "other-provider-key") + + with respx.mock(assert_all_called=False) as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/embeddings").respond(200, json=EMBEDDING_RESPONSE) + with pytest.raises(litellm.BadRequestError, match="Missing API Key"): + litellm.embedding(model="hostinger_router/text-embedding-3-small", input=["Say hello"]) + + assert route.call_count == 0 + + +def test_json_provider_without_embeddings_endpoint_refuses_embeddings(): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + slug: Final = next( + slug + for slug in JSONProviderRegistry.list_providers() + if not JSONProviderRegistry.supports_embeddings(slug) and slug not in litellm.provider_list + ) + + with respx.mock(assert_all_called=False) as upstream: + route: Final = upstream.route().respond(200, json=EMBEDDING_RESPONSE) + with pytest.raises(litellm.BadRequestError): + litellm.embedding(model=f"{slug}/text-embedding-3-small", input=["Say hello"], api_key="test-key") + + assert route.call_count == 0 + + +@pytest.mark.asyncio +async def test_hostinger_router_anthropic_messages_request(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + with respx.mock() as upstream: + route: Final = upstream.post("https://router.hostinger.com/v1/messages").respond( + 200, + json={ + "id": "msg_hostinger", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "Hello from Hostinger Router"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 4, "output_tokens": 3}, + }, + ) + response: Final = await litellm.anthropic.messages.acreate( + model="hostinger_router/claude-sonnet-5", + messages=[{"role": "user", "content": "Say hello"}], + max_tokens=32, + api_key="hostinger-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert str(request.url) == "https://router.hostinger.com/v1/messages" + assert request.headers["authorization"] == "Bearer hostinger-test-key" + assert request.headers["anthropic-version"] == "2023-06-01" + assert body["model"] == "claude-sonnet-5" + assert body["messages"] == [{"role": "user", "content": "Say hello"}] + assert response["content"][0]["text"] == "Hello from Hostinger Router" + + +CHAT_COMPLETION_RESPONSE: Final = { + "id": "chatcmpl_hostinger", + "object": "chat.completion", + "created": 1_789_550_000, + "model": "gpt-5.2", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello from Hostinger Router"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}, +} +UNSUPPORTED_ENDPOINT: Final = {"error": {"message": "Model does not support this endpoint."}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-5.2", "gpt-5-nano"]) +async def test_hostinger_router_messages_bridge_models_without_native_messages( + model: str, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + with respx.mock(assert_all_called=False) as upstream: + native: Final = upstream.post("https://router.hostinger.com/v1/messages").respond( + 403, json=UNSUPPORTED_ENDPOINT + ) + chat: Final = upstream.post("https://router.hostinger.com/v1/chat/completions").respond( + 200, json=CHAT_COMPLETION_RESPONSE + ) + response: Final = await litellm.anthropic.messages.acreate( + model=f"hostinger_router/{model}", + messages=[{"role": "user", "content": "Say hello"}], + max_tokens=32, + api_key="hostinger-test-key", + ) + + assert native.call_count == 0 + assert chat.call_count == 1 + assert json.loads(chat.calls.last.request.content)["model"] == model + assert response["content"][0]["text"] == "Hello from Hostinger Router" + + +def test_hostinger_router_responses_bridge_models_without_native_responses(): + with respx.mock(assert_all_called=False) as upstream: + native: Final = upstream.post("https://router.hostinger.com/v1/responses").respond( + 403, json=UNSUPPORTED_ENDPOINT + ) + chat: Final = upstream.post("https://router.hostinger.com/v1/chat/completions").respond( + 200, json=CHAT_COMPLETION_RESPONSE + ) + response: Final = litellm.responses( + model="hostinger_router/claude-sonnet-5", + input="Say hello", + api_key="hostinger-test-key", + ) + + assert native.call_count == 0 + assert chat.call_count == 1 + assert json.loads(chat.calls.last.request.content)["model"] == "claude-sonnet-5" + assert response.output[0].content[0].text == "Hello from Hostinger Router" + + +ANTHROPIC_MESSAGE_RESPONSE: Final = { + "id": "msg_hostinger", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "Hello from Hostinger Router"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 4, "output_tokens": 3}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sibling_model_info", + [ + {"supported_endpoints": ["/v1/chat/completions", "/v1/messages", "/v1/responses"]}, + {"supported_endpoints": "/v1/messages"}, + ], +) +async def test_hostinger_router_endpoint_opt_in_stays_with_its_own_deployment( + sibling_model_info: dict[str, object], monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + deployment_params: Final = {"model": "hostinger_router/gpt-5-nano", "api_key": "hostinger-test-key"} + router: Final = litellm.Router( + model_list=[ + {"model_name": "opted-in", "litellm_params": deployment_params, "model_info": sibling_model_info}, + {"model_name": "plain", "litellm_params": deployment_params}, + ] + ) + + with respx.mock(assert_all_called=False) as upstream: + native_messages: Final = upstream.post("https://router.hostinger.com/v1/messages").respond( + 403, json=UNSUPPORTED_ENDPOINT + ) + native_responses: Final = upstream.post("https://router.hostinger.com/v1/responses").respond( + 403, json=UNSUPPORTED_ENDPOINT + ) + chat: Final = upstream.post("https://router.hostinger.com/v1/chat/completions").respond( + 200, json=CHAT_COMPLETION_RESPONSE + ) + messages_response: Final = await router.aanthropic_messages( + model="plain", messages=[{"role": "user", "content": "Say hello"}], max_tokens=32 + ) + responses_response: Final = await router.aresponses(model="plain", input="Say hello") + + assert native_messages.call_count == 0 + assert native_responses.call_count == 0 + assert chat.call_count == 2 + assert messages_response["content"][0]["text"] == "Hello from Hostinger Router" + assert responses_response.output[0].content[0].text == "Hello from Hostinger Router" + + +@pytest.mark.asyncio +async def test_hostinger_router_deployment_cache_ttl_opt_in_keeps_one_hour_ttl(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "opted-in", + "litellm_params": { + "model": "hostinger_router/claude-haiku-4-5", + "api_base": "https://router.hostinger.com/v1", + "api_key": "hostinger-test-key", + }, + "model_info": {"supported_endpoints": ["/v1/messages"], "cache_control_ttl": True}, + } + ] + ) + one_hour_cache: Final = {"type": "ephemeral", "ttl": "1h"} + + with respx.mock() as upstream: + native: Final = upstream.post("https://router.hostinger.com/v1/messages").respond( + 200, json=ANTHROPIC_MESSAGE_RESPONSE + ) + await router.aanthropic_messages( + model="opted-in", + messages=[ + {"role": "user", "content": [{"type": "text", "text": "Say hello", "cache_control": one_hour_cache}]} + ], + max_tokens=32, + ) + + body: Final = json.loads(native.calls.last.request.content) + assert body["messages"][0]["content"][0]["cache_control"] == one_hour_cache 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 888c47505a4..6cef65f8a88 100644 --- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -1,5 +1,8 @@ +import copy import json +from collections.abc import Iterator +from itertools import chain from typing import Final from unittest.mock import MagicMock @@ -872,3 +875,172 @@ 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 _five_message_level_cache_hints() -> list[dict[str, object]]: + return [ + {"role": "system", "content": "system prompt", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "turn 1", "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "content": "answer 1", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "turn 2", "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "content": "answer 2"}, + {"role": "user", "content": "turn 3", "cache_control": {"type": "ephemeral"}}, + ] + + +def _cached_texts(node: object) -> Iterator[str]: + if isinstance(node, dict): + yield from ((node["text"],) if "cache_control" in node and "text" in node else ()) + yield from chain.from_iterable(map(_cached_texts, node.values())) + elif isinstance(node, list): + yield from chain.from_iterable(map(_cached_texts, node)) + + +@pytest.mark.parametrize( + ("drop_params", "expected_cached_texts"), + [ + (True, ["answer 1", "turn 1", "turn 2", "turn 3"]), + ("true", ["answer 1", "turn 1", "turn 2", "turn 3"]), + (False, ["answer 1", "system prompt", "turn 1", "turn 2", "turn 3"]), + ], + ids=[ + "drop_params_keeps_the_newest_four", + "string_drop_params_keeps_the_newest_four", + "without_drop_params_all_five_pass_through", + ], +) +def test_openrouter_claude_request_over_the_cache_control_cap_drops_the_oldest_hints_only_under_drop_params( + drop_params: bool | str, expected_cached_texts: list[str], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + messages: Final = _five_message_level_cache_hints() + messages_before: Final = copy.deepcopy(messages) + + body: Final = OpenrouterConfig().transform_request( + model="anthropic/claude-haiku-5.5", + messages=messages, + optional_params={}, + litellm_params={"drop_params": drop_params, "custom_llm_provider": "openrouter"}, + headers={}, + ) + + assert sorted(_cached_texts(body["messages"])) == expected_cached_texts + assert json.dumps(body).count('"cache_control"') == len(expected_cached_texts) + assert messages == messages_before + + +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/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 380f85445e5..5e774b05207 100644 --- a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py @@ -9,6 +9,7 @@ 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 @@ -119,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 @@ -135,12 +137,55 @@ 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", diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 78b9e930001..bc8421c80f8 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1,3 +1,4 @@ +import json from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -2507,3 +2508,46 @@ def test_check_and_create_cache_clamps_negative_creation_tokens() -> None: assert cache_name == "cachedContents/created" assert logging_obj.vertex_ai_cache_creation_input_tokens == 0 + + +def _recording_gemini_client(calls: list[httpx.Request]) -> HTTPHandler: + def route(request: httpx.Request) -> httpx.Response: + calls.append(request) + if request.url.path.endswith("/cachedContents"): + return httpx.Response(200, json={} if request.method == "GET" else _created_cache_body(3000)) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"role": "model", "parts": [{"text": "OK"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 3000, "candidatesTokenCount": 1, "totalTokenCount": 3001}, + }, + ) + + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(route))) + + +def test_a_plain_system_prompt_goes_into_the_gemini_cache_instead_of_beside_it(local_model_cost_map) -> None: + calls: Final[list[httpx.Request]] = [] + litellm.completion( + model="gemini/gemini-flash-latest", + messages=[ + {"role": "system", "content": "You are terse."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "plain system prefix " * 3000, "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "Reply with OK only."}, + ], + }, + ], + api_key="test-key", + client=_recording_gemini_client(calls), + ) + + created: Final = next( + json.loads(call.content) for call in calls if call.method == "POST" and call.url.path.endswith("/cachedContents") + ) + generated: Final = next(json.loads(call.content) for call in calls if call.url.path.endswith(":generateContent")) + assert created["system_instruction"] == {"parts": [{"text": "You are terse."}]} + assert generated["cachedContent"] == "cachedContents/created" + assert "system_instruction" not in generated diff --git a/tests/unit/llms/watsonx/test_watsonx.py b/tests/unit/llms/watsonx/test_watsonx.py index 4afa717debe..31d2875d2b9 100644 --- a/tests/unit/llms/watsonx/test_watsonx.py +++ b/tests/unit/llms/watsonx/test_watsonx.py @@ -1,4 +1,5 @@ import asyncio +import copy import importlib import json from typing import Final, Optional, cast @@ -10,6 +11,7 @@ import pytest import litellm from litellm import completion, embedding from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.watsonx.chat.transformation import IBMWatsonXChatConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -542,3 +544,43 @@ def test_watsonx_deployment_space_id_embedding(monkeypatch, watsonx_embedding_ca # nor space_id or project_id is required by wx.ai API when inferencing deployment assert "project_id" not in json_data and "space_id" not in json_data + + +def _strict_weather_tool() -> dict[str, object]: + return { + "type": "function", + "function": { + "name": "get_weather", + "strict": True, + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + }, + }, + } + + +def test_watsonx_chat_tools_are_rewritten_without_touching_the_callers_tool() -> None: + tools: Final = [_strict_weather_tool()] + tools_before: Final = copy.deepcopy(tools) + + optional_params: Final = IBMWatsonXChatConfig().map_openai_params( + non_default_params={"tools": tools}, optional_params={}, model="ibm/granite-4-h-small", drop_params=False + ) + + assert optional_params["tools"] == [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ] + assert tools == tools_before diff --git a/tests/unit/llms/xai/test_xai_chat_transformation.py b/tests/unit/llms/xai/test_xai_chat_transformation.py index 69e620dbc8a..37086f797fa 100644 --- a/tests/unit/llms/xai/test_xai_chat_transformation.py +++ b/tests/unit/llms/xai/test_xai_chat_transformation.py @@ -1,3 +1,4 @@ +import copy import json import os from typing import Final @@ -581,3 +582,44 @@ def test_xai_grok_4_frequency_penalty_not_supported(model): """ supported_params = XAIChatConfig().get_supported_openai_params(model=model) assert "frequency_penalty" not in supported_params + + +def _strict_weather_tool() -> dict[str, object]: + return { + "type": "function", + "function": { + "name": "get_weather", + "strict": True, + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + }, + }, + } + + +def test_xai_tools_are_rewritten_without_touching_the_callers_tool() -> None: + tools: Final = [_strict_weather_tool()] + tools_before: Final = copy.deepcopy(tools) + + optional_params: Final = XAIChatConfig().map_openai_params( + non_default_params={"tools": tools}, optional_params={}, model="grok-4.3", drop_params=False + ) + + assert optional_params["tools"] == [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + }, + }, + } + ] + assert tools == tools_before 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_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 6b59f6d40dc..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 @@ -5884,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 @@ -5892,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 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/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_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/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/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_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index abaa802de2c..ab6cfcfebfe 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4,8 +4,9 @@ import inspect import json import socket from collections.abc import AsyncIterator, Iterator, Mapping +from datetime import datetime, time, timezone from types import MappingProxyType, SimpleNamespace -from typing import Dict, Final, Optional +from typing import Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -46,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): @@ -4048,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/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 4decee140e4..e5672f52fb6 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 @@ -7020,7 +7020,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 c70e66329ad..3ebd2f32953 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 @@ -31,7 +31,7 @@ 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.constants import ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH, ANTHROPIC_MAX_USER_ID_LENGTH, DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger @@ -77,6 +77,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( 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 @@ -6620,6 +6621,272 @@ async def test_pass_through_body_stays_byte_identical_outside_the_opt_in( assert "x-litellm-call-id" in response.headers +def _echo_upstream_body(upstream_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, content=upstream_request.content, headers={"content-type": "application/json"}, request=upstream_request + ) + + +async def _relayed_body(response: Response) -> bytes: + if isinstance(response, StreamingResponse): + return b"".join([chunk async for chunk in response.body_iterator]) + return bytes(response.body) + + +async def _relay_through_pass_through( + client_body: Mapping[str, object], + target: str, + guardrails_config: dict[str, object] | None, + custom_body: dict[str, object] | None = None, +) -> Mapping[str, object]: + async def receive() -> Mapping[str, object]: + return {"type": "http.request", "body": json.dumps(dict(client_body)).encode()} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/my-claude-route", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + fake_client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(_echo_upstream_body), timeout=None) + try: + with ExitStack() as stack: + _enter_relay_logging_mocks(stack, {}) + stack.enter_context( + patch("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())) + ) + response: Final = await pass_through_request( + request=request, + target=target, + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-relay-test"), + custom_body=custom_body, + guardrails_config=guardrails_config, + ) + relayed: Final = await _relayed_body(response) + finally: + cleanup() + await fake_client.aclose() + + assert response.status_code == 200 + return MappingProxyType(json.loads(relayed)) + + +_MESSAGES_BODY: Final = MappingProxyType( + {"model": "claude-sonnet-4-5", "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]} +) +_ANTHROPIC_MESSAGES: Final = "https://api.anthropic.com/v1/messages" + + +@pytest.fixture +def no_configured_anthropic_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_configured_anthropic_base") +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +@pytest.mark.parametrize( + "client_metadata, target, forwarded_metadata", + [ + ({"user_id": "user-1"}, _ANTHROPIC_MESSAGES, {"user_id": "user-1"}), + ({"user_id": "user-1"}, f"{_ANTHROPIC_MESSAGES}/", {"user_id": "user-1"}), + ({"user_id": "user-1"}, f"{_ANTHROPIC_MESSAGES}?beta=true", {"user_id": "user-1"}), + ( + {"user_id": "user-1", "tags": ["a"], "trace_id": "t", "generation_name": "g", "foo": "bar"}, + _ANTHROPIC_MESSAGES, + {"user_id": "user-1"}, + ), + ({"user_id": "1234567890"}, _ANTHROPIC_MESSAGES, {"user_id": "1234567890"}), + ({"tags": ["a"], "foo": "bar"}, _ANTHROPIC_MESSAGES, None), + ({"user_id": "a@example.com"}, _ANTHROPIC_MESSAGES, None), + ({"user_id": "x" * (ANTHROPIC_MAX_USER_ID_LENGTH + 1)}, _ANTHROPIC_MESSAGES, None), + ({"user_id": 42}, _ANTHROPIC_MESSAGES, None), + ({"user_id": "line one\nline two"}, _ANTHROPIC_MESSAGES, None), + ({"user_id": "a\x1f@example.com"}, _ANTHROPIC_MESSAGES, None), + ({"user_id": "user-1"}, f"{_ANTHROPIC_MESSAGES}/count_tokens", None), + ({"user_id": "user-1"}, f"{_ANTHROPIC_MESSAGES}/batches", None), + ({"user_id": "user-1"}, "https://api.anthropic.com/v1/models", None), + ({"user_id": "user-1"}, "https://other-vendor.test/anthropic/v1/messages", None), + ({"user_id": "user-1"}, "https://other-vendor.test/v1/responses", None), + ({"user_id": "user-1"}, "https://other-vendor.test/v1/chat/completions", None), + ], + ids=[ + "messages", + "messages-trailing-slash", + "messages-with-query", + "user-id-beside-litellm-and-unknown-keys", + "digits-user-id", + "no-user-id", + "email-user-id", + "over-long-user-id", + "integer-user-id", + "newline-user-id", + "control-character-before-at-sign", + "count-tokens", + "batches", + "other-anthropic-path", + "other-vendor-messages", + "other-vendor-responses", + "other-vendor-chat-completions", + ], +) +async def test_pass_through_forwards_only_an_acceptable_user_id_to_anthropic_messages( + stream: bool, client_metadata: object, target: str, forwarded_metadata: Mapping[str, object] | None +) -> None: + client_body: Final = MappingProxyType( + {**_MESSAGES_BODY, "metadata": client_metadata, **({"stream": True} if stream else {})} + ) + body_without_metadata: Final = {k: v for k, v in client_body.items() if k != "metadata"} + + relayed: Final = await _relay_through_pass_through(client_body, target, guardrails_config=None) + + assert relayed == ( + body_without_metadata if forwarded_metadata is None else {**body_without_metadata, "metadata": forwarded_metadata} + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_configured_anthropic_base") +@pytest.mark.parametrize("base_env", ["ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) +@pytest.mark.parametrize( + "target, forwarded", + [ + ("http://gateway.test:8080/anthropic/v1/messages", True), + ("http://gateway.test:9090/v1/messages", False), + ("http://other-gateway.test:8080/v1/messages", False), + ], + ids=["configured-upstream", "other-port", "other-host"], +) +async def test_pass_through_treats_the_configured_anthropic_base_as_anthropic( + monkeypatch: pytest.MonkeyPatch, base_env: str, target: str, forwarded: bool +) -> None: + monkeypatch.setenv(base_env, "http://gateway.test:8080/anthropic") + client_body: Final = MappingProxyType({**_MESSAGES_BODY, "metadata": {"user_id": "user-1"}}) + + relayed: Final = await _relay_through_pass_through(client_body, target, guardrails_config=None) + + assert relayed == (client_body if forwarded else _MESSAGES_BODY) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_configured_anthropic_base") +@pytest.mark.parametrize( + "user_id_length, target, forwarded", + [ + (ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH, "http://gateway.test/v1/messages", True), + (ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH + 1, "http://gateway.test/v1/messages", False), + (ANTHROPIC_MAX_USER_ID_LENGTH, _ANTHROPIC_MESSAGES, True), + ], + ids=["at-the-configured-base-limit", "over-the-configured-base-limit", "direct-anthropic-keeps-its-limit"], +) +async def test_pass_through_caps_the_user_id_for_a_configured_anthropic_compatible_base( + monkeypatch: pytest.MonkeyPatch, user_id_length: int, target: str, forwarded: bool +) -> None: + monkeypatch.setenv("ANTHROPIC_API_BASE", "http://gateway.test") + client_body: Final = MappingProxyType({**_MESSAGES_BODY, "metadata": {"user_id": "u" * user_id_length}}) + + relayed: Final = await _relay_through_pass_through(client_body, target, guardrails_config=None) + + assert relayed == (client_body if forwarded else _MESSAGES_BODY) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_configured_anthropic_base") +@pytest.mark.parametrize( + "default_on, event_hook, guardrails_config, client_fields, stream_scope, forwarded", + [ + (False, GuardrailEventHooks.pre_call, None, {}, None, True), + (False, GuardrailEventHooks.pre_call, {"unregistered-guard": {}}, {}, None, True), + (True, GuardrailEventHooks.post_call, None, {}, None, True), + (True, GuardrailEventHooks.pre_call, None, {}, None, False), + (False, GuardrailEventHooks.pre_call, {"pii-guard": {}}, {}, None, False), + (False, GuardrailEventHooks.pre_call, None, {"litellm_metadata": {"guardrails": ["pii-guard"]}}, None, False), + (False, GuardrailEventHooks.pre_call, {"pii-guard": {}}, {"litellm_metadata": {"guardrails": []}}, None, False), + (True, GuardrailEventHooks.pre_call, None, {"stream": True}, "streaming", False), + (True, GuardrailEventHooks.pre_call, None, {}, "streaming", True), + (True, GuardrailEventHooks.pre_call, None, {"stream": True}, "non_streaming", True), + (True, GuardrailEventHooks.pre_call, None, {}, "non_streaming", False), + ], + ids=[ + "guardrail-not-requested", + "guardrail-not-registered", + "default-on-post-call-guardrail", + "default-on-pre-call-guardrail", + "guardrail-on-the-route", + "guardrail-in-litellm-metadata", + "route-guardrail-beside-an-empty-client-list", + "streaming-guardrail-on-a-stream", + "streaming-guardrail-off-a-stream", + "non-streaming-guardrail-on-a-stream", + "non-streaming-guardrail-off-a-stream", + ], +) +async def test_pass_through_keeps_the_user_id_from_routes_a_guardrail_scans( + default_on: bool, + event_hook: GuardrailEventHooks, + guardrails_config: dict[str, object] | None, + client_fields: Mapping[str, object], + stream_scope: str | None, + forwarded: bool, +) -> None: + guardrail: Final = CustomGuardrail( + guardrail_name="pii-guard", default_on=default_on, event_hook=event_hook, stream_scope=stream_scope + ) + client_body: Final = MappingProxyType({**_MESSAGES_BODY, "metadata": {"user_id": "user-1"}, **client_fields}) + body_without_client_fields: Final = {k: v for k, v in client_body.items() if k not in ("metadata", "litellm_metadata")} + + with patch.object(litellm, "callbacks", [guardrail]): + relayed: Final = await _relay_through_pass_through( + client_body, _ANTHROPIC_MESSAGES, guardrails_config=guardrails_config + ) + + assert relayed == ( + {**body_without_client_fields, "metadata": {"user_id": "user-1"}} if forwarded else body_without_client_fields + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_configured_anthropic_base") +async def test_pass_through_warns_without_the_value_when_it_drops_a_refused_user_id( + caplog: pytest.LogCaptureFixture, +) -> None: + client_body: Final = MappingProxyType({**_MESSAGES_BODY, "metadata": {"user_id": "person@example.com"}}) + + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + await _relay_through_pass_through(client_body, _ANTHROPIC_MESSAGES, guardrails_config=None) + + assert "Not forwarding metadata.user_id" in caplog.text + assert "person@example.com" not in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_configured_anthropic_base") +async def test_pass_through_drops_metadata_that_is_not_an_object() -> None: + client_body: Final = MappingProxyType({**_MESSAGES_BODY, "metadata": "user-1"}) + + relayed: Final = await _relay_through_pass_through(client_body, _ANTHROPIC_MESSAGES, guardrails_config=None) + + assert relayed == _MESSAGES_BODY + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_configured_anthropic_base") +async def test_pass_through_custom_body_replaces_the_client_user_id() -> None: + client_body: Final = MappingProxyType({**_MESSAGES_BODY, "metadata": {"user_id": "user-1"}}) + + relayed: Final = await _relay_through_pass_through( + client_body, _ANTHROPIC_MESSAGES, guardrails_config=None, custom_body=dict(_MESSAGES_BODY) + ) + + assert relayed == _MESSAGES_BODY + + _PARTIAL_RELAY_WARNING_MARKER = "ended before upstream body was fully relayed" @@ -7736,6 +8003,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 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_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 a7596b51e3e..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 @@ -87,7 +88,7 @@ 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(): @@ -10221,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 @@ -10257,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 @@ -10282,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") @@ -10363,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 @@ -10391,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 @@ -10406,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"]) @@ -10472,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 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_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..26ee9f04d01 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 @@ -995,3 +996,21 @@ async def test_guardrails_only_and_skip_guardrails_are_mutually_exclusive( guardrails_only=True, skip_guardrails=True, ) + + +@pytest.mark.parametrize( + "request_data, expected", + [ + ({"metadata": {"_guardrail_pipelines": ["p1"]}}, True), + ({"metadata": {"user_id": "u"}, "litellm_metadata": {"_guardrail_pipelines": ["p1"]}}, True), + ({"metadata": {"_guardrail_pipelines": 5}}, True), + ({"metadata": {"user_id": "u"}}, False), + ], + ids=["pipeline-in-metadata", "pipeline-in-litellm-metadata", "non-list-pipeline-value", "no-pipeline"], +) +def test_request_has_pre_call_guardrails_reads_pipelines_from_both_metadata_buckets( + proxy_logging, monkeypatch, request_data, expected +): + monkeypatch.setattr(litellm, "callbacks", []) + + assert proxy_logging.request_has_pre_call_guardrails(request_data) is expected 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 0341511496c..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 @@ -1595,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_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 4250c78c669..b1047a345d7 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6314,6 +6314,17 @@ def test_marks_past_the_provider_cap_keep_the_newest_under_drop_params(drop_para assert messages[1] == {"role": "user", "content": "turn 0"} +def test_marks_past_the_provider_cap_inside_one_message_keep_the_newest_under_drop_params(): + parts: Final = tuple( + {"type": "input_text", "text": f"part {index}", "cache_control": _EPHEMERAL} for index in range(1, 6) + ) + + messages: Final = _bridged_request({"role": "user", "content": list(parts)}, drop_params=True) + + content: Final = TypeAdapter(list[dict[str, object]]).validate_python(messages[-1]["content"]) + assert ["cache_control" in part for part in content] == [False, True, True, True, True] + + def test_global_drop_params_keeps_the_newest_marks(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(litellm, "drop_params", True) 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/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_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 7235333f59f..6b5c052618b 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.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy from litellm.utils import Rules, function_setup @@ -1567,7 +1568,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 @@ -2060,7 +2061,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_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_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_router.py b/tests/unit/test_router/test_router.py index a6436e47507..7ac27e22a93 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 ) @@ -5140,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.""" @@ -5187,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", @@ -5204,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, @@ -5220,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" @@ -10092,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( @@ -12917,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 @@ -13007,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) @@ -17489,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( @@ -20879,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 @@ -20912,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 @@ -22009,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 @@ -22843,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 @@ -22962,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 @@ -28242,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 @@ -29893,7 +29905,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", @@ -31157,7 +31169,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", ), @@ -31253,7 +31265,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 5f31ded626d..8929c5fd345 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -2614,7 +2614,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() @@ -2799,7 +2799,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/package-lock.json b/ui/litellm-dashboard/package-lock.json index 48def46e3fe..54cd9087c4e 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -2037,6 +2037,13 @@ "url": "https://opencollective.com/libvips" } }, + "node_modules/@isaacs/nonrecursive-glob-match": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/@isaacs/nonrecursive-glob-match/-/nonrecursive-glob-match-1.0.2.tgz", + "integrity": "sha512-GJPHIgPPtFl+nLhFHyF0zZoFcOFLsZWbHFFzo+Zr+PWDXJbl9hHxb2w+9SiF2fHmLgKaXIpUzQo3SbBBZ9WObw==", + "dev": true, + "license": "BlueOak-1.0.0" + }, "node_modules/@jridgewell/gen-mapping": { "version": "0.3.13", "resolved": "https://registry.npmjs.org/@jridgewell/gen-mapping/-/gen-mapping-0.3.13.tgz", @@ -9657,13 +9664,14 @@ } }, "node_modules/minimatch": { - "version": "10.2.4", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.4.tgz", - "integrity": "sha512-oRjTw/97aTBN0RHbYCdtF1MQfvusSIBQM0IZEgzl6426+8jSC0nF1a/GmnVLpfB9yyr6g6FTqWqiZVbxrtaCIg==", + "version": "11.0.2", + "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-11.0.2.tgz", + "integrity": "sha512-pWTgz5CN3tjqyZNSK2tePlSmu9IwM41q8yn8Usdqqwe9TIHh+9dEpLLjF4uApsfPEAZY71qtSiA2B2APEr+1Yw==", "dev": true, "license": "BlueOak-1.0.0", "dependencies": { - "brace-expansion": "^5.0.2" + "@isaacs/nonrecursive-glob-match": "^1.0.2", + "brace-expansion": "^5.0.12" }, "engines": { "node": "18 || 20 || >=22" diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index f65316439bf..0707496a6bc 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -117,7 +117,7 @@ "js-yaml": "4.3.2", "brace-expansion": "5.0.12", "glob": "13.0.0", - "minimatch": "10.2.4", + "minimatch": "11.0.2", "ws": "8.21.0", "braces": "3.0.3", "axios": "1.13.6", diff --git a/ui/litellm-dashboard/public/assets/logos/hostinger.png b/ui/litellm-dashboard/public/assets/logos/hostinger.png new file mode 100644 index 00000000000..9ae3f4bcca1 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/hostinger.png differ 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({