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/.circleci/AGENTS.md b/.circleci/AGENTS.md new file mode 100644 index 00000000000..b3591145a3e --- /dev/null +++ b/.circleci/AGENTS.md @@ -0,0 +1,7 @@ +# .circleci + +The `integration` workflow on `main` runs from a scheduled pipeline in the CircleCI project settings, cron `9 0-8,14-23 * * *` (UTC) on `main` with no pipeline parameters. It has no `triggers:` block in this file because a workflow with one runs only on its schedule, so the API pipelines that set `run_migration_tests` or `routing_parity_base` would stop running it. If main stops getting integration runs, check that the schedule still exists and recreate it with these values + +Jobs clone with the `checkout_over_https` command, never the built-in `checkout` step, because CircleCI falls back to the project's SSH checkout key without its GitHub App access, GitHub rejects that key, and the public repo needs no key. It fetches the job's commit with full history and no file contents (`--filter=blob:none`), like the built-in step did, so `scripts/path_filter.sh` can find a merge base + +A step that checks out files from another commit fetches that commit with `--no-filter`, as `routing_parity` does, because the clone has no file contents and git would otherwise download each missing file in its own request, which took 210 to 896 seconds for the routing parity base diff --git a/.circleci/config.yml b/.circleci/config.yml index 4ce5ea8b6fe..15129f1e5e3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -21,19 +21,30 @@ orbs: win: circleci/windows@5.0 # Add Windows orb commands: - checkout_migration_source: + checkout_over_https: + parameters: + revision: + type: string + default: "" + shell: + type: string + default: /bin/bash -eo pipefail steps: - run: - name: Select the requested migration test revision + name: Check out the public repository over HTTPS + shell: << parameters.shell >> environment: - MIGRATION_SOURCE_SHA: << pipeline.parameters.migration_source_sha >> + REQUESTED_REVISION: << parameters.revision >> command: | - revision="${MIGRATION_SOURCE_SHA:-$CIRCLE_SHA1}" + revision="${REQUESTED_REVISION:-$CIRCLE_SHA1}" [[ "$revision" =~ ^[0-9a-f]{40}$ ]] || exit 1 git init git remote add origin https://github.com/BerriAI/litellm.git - git fetch --depth 1 origin "$revision" - git checkout --detach FETCH_HEAD + for attempt in 1 2 3; do + git fetch --filter=blob:none origin "$revision" && git checkout --detach FETCH_HEAD && break + [[ $attempt -lt 3 ]] || exit 1 + sleep 5 + done skip_if_unrelated_changes: parameters: category: @@ -347,7 +358,7 @@ commands: - when: condition: << parameters.checkout >> steps: - - checkout + - checkout_over_https - setup_google_dns - install_uv - install_rust @@ -377,7 +388,8 @@ jobs: CARGO_HTTP_MULTIPLEXING: "false" CARGO_NET_RETRY: "5" steps: - - checkout + - checkout_over_https: + shell: bash.exe - run: name: Skip job when no windows-release-relevant files changed shell: bash.exe @@ -451,7 +463,7 @@ jobs: password: ${DOCKERHUB_PASSWORD} working_directory: ~/project steps: - - checkout + - checkout_over_https - setup_google_dns - install_uv - install_rust @@ -482,7 +494,7 @@ jobs: working_directory: ~/project steps: - - checkout + - checkout_over_https - skip_if_unrelated_changes - attach_workspace: at: ~/project @@ -519,31 +531,36 @@ jobs: IMAGE_TAG=${CIRCLE_SHA1:-ci} kind load docker-image litellm-ci:${IMAGE_TAG} --name litellm-test - # Run helm lint - run: name: Run helm lint command: | - helm lint ./helm/litellm-helm + helm lint ./helm/litellm -f ./helm/litellm/ci/test-values.yaml + + - run: + name: Start PostgreSQL for the chart + command: | + kubectl apply -f ./helm/litellm/ci/postgres.yaml + kubectl rollout status deployment/litellm-ci-postgres --timeout=3m - # Run helm tests - run: name: Run helm tests command: | IMAGE_TAG=${CIRCLE_SHA1:-ci} - helm install litellm ./helm/litellm-helm -f ./helm/litellm-helm/ci/test-values.yaml \ + helm install litellm ./helm/litellm -f ./helm/litellm/ci/test-values.yaml \ --set image.repository=litellm-ci \ --set image.tag=${IMAGE_TAG} \ - --set image.pullPolicy=Never - if ! kubectl rollout status deployment/litellm --timeout=180s; then - kubectl describe pods -l app.kubernetes.io/name=litellm - kubectl logs -l app.kubernetes.io/name=litellm --tail=200 - exit 1 - fi + --set image.pullPolicy=Never \ + --wait --wait-for-jobs --timeout 10m || { + kubectl get pods -o wide + kubectl describe pods -l app.kubernetes.io/instance=litellm + kubectl logs -l app.kubernetes.io/component=proxy --tail=200 || true + kubectl logs -l app.kubernetes.io/component=migrations --tail=200 || true + exit 1 + } - echo "Printing pod logs..." - kubectl logs $(kubectl get pods -l app.kubernetes.io/name=litellm -o jsonpath="{.items[0].metadata.name}") + echo "Printing proxy pod logs..." + kubectl logs -l app.kubernetes.io/component=proxy --tail=200 - # Run the helm tests helm test litellm --logs # Cleanup @@ -559,7 +576,7 @@ jobs: resource_class: medium working_directory: ~/project steps: - - checkout + - checkout_over_https - run: name: Generate LiteLLM master key command: | @@ -659,7 +676,7 @@ jobs: resource_class: large working_directory: ~/project steps: - - checkout + - checkout_over_https - run: name: Generate LiteLLM master key command: | @@ -740,7 +757,7 @@ jobs: resource_class: large working_directory: ~/project steps: - - checkout + - checkout_over_https - skip_if_unrelated_changes: category: client - setup_google_dns @@ -794,7 +811,7 @@ jobs: # the rest of the suite: nothing else performs a logout. PROXY_LOGOUT_URL: "https://www.example.com" steps: - - checkout + - checkout_over_https - run: name: Generate LiteLLM master key command: | @@ -940,7 +957,7 @@ jobs: # and by migration.serverRootPath.config.ts, which refuses to run without it. SERVER_ROOT_PATH: "/litellm" steps: - - checkout + - checkout_over_https - run: name: Generate LiteLLM master key command: | @@ -1066,11 +1083,12 @@ jobs: - when: condition: << parameters.migration_qualification >> steps: - - checkout_migration_source + - checkout_over_https: + revision: << pipeline.parameters.migration_source_sha >> - unless: condition: << parameters.migration_qualification >> steps: - - checkout + - checkout_over_https - skip_if_unrelated_changes - run: @@ -1086,7 +1104,7 @@ jobs: docker build \ --label org.opencontainers.image.revision="$(git rev-parse HEAD)" \ -t litellm-docker-database:ci \ - -f docker/Dockerfile.database . + -f Dockerfile . fi python3 .circleci/scripts/run_migration_tests.py record-image @@ -1122,7 +1140,8 @@ jobs: MIGRATION_TEST_OUTPUT: /tmp/migration-results PYTHONPATH: tests/e2e steps: - - checkout_migration_source + - checkout_over_https: + revision: << pipeline.parameters.migration_source_sha >> - run: name: Generate LiteLLM master key command: | @@ -1197,7 +1216,7 @@ jobs: resource_class: medium working_directory: ~/project steps: - - checkout + - checkout_over_https - skip_if_unrelated_changes - attach_workspace: at: ~/project @@ -1371,7 +1390,7 @@ jobs: environment: UV_PYTHON_DOWNLOADS: never steps: - - checkout + - checkout_over_https - install_uv - run: name: Install the locked shard verifier and coverage dependencies @@ -1419,7 +1438,7 @@ jobs: ROUTING_PARITY_BASE: << pipeline.parameters.routing_parity_base >> command: | [[ "$ROUTING_PARITY_BASE" =~ ^[0-9a-f]{40}$ ]] || exit 1 - git fetch --depth 1 origin "$ROUTING_PARITY_BASE" + git fetch --depth 1 --no-filter origin "$ROUTING_PARITY_BASE" git rm -r -f --quiet litellm enterprise litellm-proxy-extras git checkout "$ROUTING_PARITY_BASE" -- litellm enterprise litellm-proxy-extras git reset --quiet @@ -1494,7 +1513,7 @@ jobs: environment: DATABASE_URL: postgresql://postgres:postgres@localhost:5432/litellm_test steps: - - checkout + - checkout_over_https - skip_if_unrelated_changes - setup_litellm_test_deps: checkout: false diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 1380b07de53..fb91bcce77f 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -128,6 +128,8 @@ fi sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt" awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt" +export UV_OFFLINE=true + if [ "$mode" = parity ]; then export INTEGRATION_WORKERS=1 else diff --git a/.dockerignore b/.dockerignore index f3a80fee3e4..8459d0cd482 100644 --- a/.dockerignore +++ b/.dockerignore @@ -7,7 +7,6 @@ tests .devcontainer *.tgz log.txt -docker/Dockerfile.* # Claude Flow generated files (must be excluded from Docker build) .claude/ diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index 86a23966671..f815601a29c 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -4,9 +4,14 @@ Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or "## Affected release" with nothing under it must not appear in the final description --> - + No bullets, tables, or code names here. Everything else goes inside the Details block below --> + +[insert handwritten description here]
Details @@ -88,6 +93,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac Otherwise proof must be fully e2e, no mocks, with real provider calls where applicable. Pytest alone is insufficient. The following format applies to live proof Show only the latest proof. Before uses the merge base and After the tested tip. Rerun affected QA after behavior changes. Verified neutral diffs may carry proof with the tested SHA, current tip, equivalence reason and focused checks, following AGENTS.md and /qa Use matching Before/After case headings with numbered commands and observed output. Put shared setup above them. For one case, omit case headings + A reviewer must be able to rerun every case from the body alone, so the shared setup gives the proxy start command, the config.yaml, env var names (never values) and the client code or curl used, all exactly as run. Wrap any block longer than about 20 lines in a collapsed
with a one-line summary ### Before () diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 5dffa78776d..76f8e3ef7e7 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -57,15 +57,6 @@ test_paths: - tests/integration/test_oci_proxy_integration.py dockerfiles: - - reason: >- - The dashboard container is a static Next.js export served by nginx, and the dashboard build - and lint workflows already exercise that output, so building the image adds no signal about it - paths: - - ui/Dockerfile - - reason: >- - An example image under cookbook/ that is documentation rather than a shipped artifact - paths: - - cookbook/litellm-ollama-docker-image/Dockerfile - reason: >- The Rust gateway image compiles the whole workspace in release mode, which is too slow for a per-pull-request job while the gateway binary is still being assembled; the Rust lint, diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 4c4cb399dce..571092263ed 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -214,7 +214,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 48_000_000 + native_size_limit: Final = 55_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), diff --git a/.github/workflows/compat-matrix-image.yml b/.github/workflows/compat-matrix-image.yml index f08792c904c..73a77c1e861 100644 --- a/.github/workflows/compat-matrix-image.yml +++ b/.github/workflows/compat-matrix-image.yml @@ -4,7 +4,7 @@ on: pull_request: paths: - tests/e2e/claude_code/cron_vm/** - - tests/e2e/claude_code/pr_gate_version_resolver.py + - tests/e2e/cli_harness/npm_version_resolver.py - .github/workflows/compat-matrix-image.yml workflow_dispatch: @@ -26,17 +26,24 @@ jobs: with: persist-credentials: false - - name: Build the Render cron image - run: docker build -f tests/e2e/claude_code/cron_vm/Dockerfile -t compat-matrix:${{ github.sha }} tests/e2e - - - name: Resolve and install the Claude Code CLI as the cron user + - name: Install the pinned uv the cron image ships + env: + UV_VERSION: 0.10.9 + UV_SHA256: 20d79708222611fa540b5c9ed84f352bcd3937740e51aacc0f8b15b271c57594 run: | - docker run --rm compat-matrix:${{ github.sha }} bash -c ' - set -euo pipefail - whoami - gh --version - uv --version - version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/claude_code/pr_gate_version_resolver.py)" - /opt/litellm/tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" /tmp/claude-cli - /tmp/claude-cli/claude --version - ' + set -euo pipefail + curl -fsSLo "$RUNNER_TEMP/uv.tar.gz" "https://github.com/astral-sh/uv/releases/download/${UV_VERSION}/uv-x86_64-unknown-linux-gnu.tar.gz" + echo "${UV_SHA256} $RUNNER_TEMP/uv.tar.gz" | sha256sum -c - + mkdir -p "$RUNNER_TEMP/bin" + tar -xzf "$RUNNER_TEMP/uv.tar.gz" -C "$RUNNER_TEMP/bin" --strip-components=1 uv-x86_64-unknown-linux-gnu/uv + echo "$RUNNER_TEMP/bin" >> "$GITHUB_PATH" + + - name: Resolve and install the Claude Code CLI as an unprivileged user + run: | + set -euo pipefail + whoami + gh --version + uv --version + version="$(uv run --no-project --python 3.12 python tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)" + tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" "$RUNNER_TEMP/claude-cli" + "$RUNNER_TEMP/claude-cli/claude" --version diff --git a/.github/workflows/create-rc-branch.yml b/.github/workflows/create-rc-branch.yml index 5269460cc93..08ee8351fda 100644 --- a/.github/workflows/create-rc-branch.yml +++ b/.github/workflows/create-rc-branch.yml @@ -2,7 +2,7 @@ name: Create RC Branch on: schedule: - - cron: "0 3 * * 5" + - cron: "0 3 * * 1" timezone: "America/Los_Angeles" workflow_dispatch: diff --git a/.github/workflows/deploy-dev-on-main.yml b/.github/workflows/deploy-dev-on-main.yml deleted file mode 100644 index 1208a2dbbdb..00000000000 --- a/.github/workflows/deploy-dev-on-main.yml +++ /dev/null @@ -1,30 +0,0 @@ -name: Deploy main to gateway-dev - -on: - push: - branches: [main] - -permissions: {} - -concurrency: - group: deploy-dev-dispatch - cancel-in-progress: true - -jobs: - dispatch: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 5 - steps: - - uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2 - id: token - with: - app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }} - private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }} - owner: BerriAI - repositories: litellm-ops - permission-actions: write - - name: Dispatch the litellm-ops dev deploy - env: - GH_TOKEN: ${{ steps.token.outputs.token }} - run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index 3e2bca8779d..54051104cb9 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -8,31 +8,25 @@ on: - "litellm_**" paths: - Dockerfile - - docker/Dockerfile.non_root - - docker/Dockerfile.database - - migrations/Dockerfile + - .dockerignore + - docker-entrypoint.sh - migrations/run.py - - gateway/Dockerfile - - gateway/launch.py - gateway/main.py + - gateway/launch.py - gateway/routes/allowlist.py - litellm_cli/** - - backend/Dockerfile - backend/main.py + - backend/routes/allowlist.py - litellm-rust/** - litellm/proxy/lens/** - - docker/component_entrypoint.sh - - docker/entrypoint.sh - litellm/proxy/prisma_migration.py - litellm/proxy/admin_mcp.py - litellm/proxy/proxy_server.py - - backend/routes/allowlist.py - pyproject.toml - litellm-proxy-extras/** - tests/integration/database/image/** - uv.lock - ui/litellm-dashboard/package-lock.json - - ui/Dockerfile - ui/nginx.conf - .github/workflows/image-scan.yml - .grype.yaml @@ -46,14 +40,18 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true +env: + IMAGE: litellm-image-scan:${{ github.sha }} + ADMIN_MCP_SCHEMA_IMAGE: litellm-admin-mcp-schema:${{ github.sha }} + jobs: - image-scan: - name: image-scan + build: + name: build runs-on: ubuntu-latest if: >- github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository - timeout-minutes: 30 + timeout-minutes: 60 permissions: contents: read steps: @@ -61,7 +59,129 @@ jobs: with: persist-credentials: false + - name: Build the image + run: docker build --build-arg MATURIN_PEP517_ARGS="--profile ci" --build-arg LITELLM_RELEASE_TAG=v0.0.0-lens-ci -t "$IMAGE" . + + - name: Save the image for the verify matrix + run: docker save "$IMAGE" | zstd -T0 -3 -o "$RUNNER_TEMP/litellm-image.tar.zst" + + - uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: litellm-image + path: ${{ runner.temp }}/litellm-image.tar.zst + retention-days: 1 + compression-level: 0 + + - name: Build the Admin MCP schema image + run: docker build --build-arg MATURIN_PEP517_ARGS="--profile ci" --build-arg LITELLM_RELEASE_TAG=v0.0.0-lens-ci --target builder -t "$ADMIN_MCP_SCHEMA_IMAGE" . + + - name: Save the Admin MCP schema image for the verify matrix + run: docker save "$ADMIN_MCP_SCHEMA_IMAGE" | zstd -T0 -3 -o "$RUNNER_TEMP/litellm-admin-mcp-schema.tar.zst" + + - uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: litellm-admin-mcp-schema + path: ${{ runner.temp }}/litellm-admin-mcp-schema.tar.zst + retention-days: 1 + compression-level: 0 + + verify: + name: verify (${{ matrix.check }}) + needs: build + runs-on: ubuntu-latest + timeout-minutes: 45 + permissions: + contents: read + strategy: + fail-fast: false + matrix: + check: [scan, migrations, gateway, backend, ui, collector] + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + with: + name: litellm-image + path: ${{ runner.temp }} + + - name: Load the image + run: zstd -d --stdout "$RUNNER_TEMP/litellm-image.tar.zst" | docker load + + - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + if: matrix.check == 'migrations' || matrix.check == 'backend' + with: + name: litellm-admin-mcp-schema + path: ${{ runner.temp }} + + - name: Load the Admin MCP schema image + if: matrix.check == 'migrations' || matrix.check == 'backend' + run: zstd -d --stdout "$RUNNER_TEMP/litellm-admin-mcp-schema.tar.zst" | docker load + + - name: Set up Python + if: matrix.check != 'scan' + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Install pytest + if: matrix.check != 'scan' + run: python -m pip install "pytest==9.0.3" + + # The prisma bake must migrate a fresh DB with no egress as an arbitrary + # non-root uid (OpenShift restricted-v2 / air-gapped / readOnlyFilesystem). + # `docker run` as the default uid with network hides a broken bake because + # the migration entrypoint exits 0 even when it applied nothing; asserting + # the schema was created is what catches it. + - name: Verify offline migration as a non-root uid + if: matrix.check == 'migrations' + env: + LITELLM_IMAGE: ${{ env.IMAGE }} + LITELLM_ADMIN_MCP_SCHEMA_IMAGE: ${{ env.ADMIN_MCP_SCHEMA_IMAGE }} + run: | + python -m pytest --noconftest -v \ + tests/integration/database/image/test_offline_image_migration.py \ + tests/integration/database/image/test_image_bedrock_realtime_extra.py \ + tests/integration/database/image/test_image_admin_mcp.py + + - name: Verify the gateway component serves offline as a non-root uid + if: matrix.check == 'gateway' + env: + LITELLM_IMAGE: ${{ env.IMAGE }} + LITELLM_COMPONENT: gateway + LITELLM_IMAGE_COMPONENT: gateway + run: | + python -m pytest --noconftest -v \ + tests/integration/database/image/test_component_image_serves_offline.py \ + tests/integration/database/image/test_image_admin_mcp.py + + - name: Verify the backend component serves offline as a non-root uid + if: matrix.check == 'backend' + env: + LITELLM_IMAGE: ${{ env.IMAGE }} + LITELLM_COMPONENT: backend + LITELLM_IMAGE_COMPONENT: backend + LITELLM_ADMIN_MCP_SCHEMA_IMAGE: ${{ env.ADMIN_MCP_SCHEMA_IMAGE }} + run: | + python -m pytest --noconftest -v \ + tests/integration/database/image/test_component_image_serves_offline.py \ + tests/integration/database/image/test_image_admin_mcp.py + + - name: Verify the ui component serves as an arbitrary uid with a read-only root fs + if: matrix.check == 'ui' + env: + LITELLM_IMAGE: ${{ env.IMAGE }} + run: python -m pytest --noconftest -v tests/integration/database/image/test_ui_image_serves_offline.py + + - name: Verify the collector creates its socket as the image user + if: matrix.check == 'collector' + env: + LITELLM_IMAGE: ${{ env.IMAGE }} + run: python -m pytest --noconftest -v tests/integration/database/image/test_collector_image_serves_socket.py + - name: Download Grype v0.114.0 + if: matrix.check == 'scan' run: | curl -fsSL --retry 3 -o "$RUNNER_TEMP/grype.tar.gz" \ https://github.com/anchore/grype/releases/download/v0.114.0/grype_0.114.0_linux_amd64.tar.gz @@ -69,33 +189,6 @@ jobs: tar xzf "$RUNNER_TEMP/grype.tar.gz" -C "$RUNNER_TEMP" grype chmod +x "$RUNNER_TEMP/grype" - # Dockerfile.non_root is the rootless variant we ship. The other - # Dockerfiles share the same wolfi base and apk set, so OS-layer coverage - # is the same; matrix-scan if those variants ever diverge. - - name: Build runtime image - run: docker build -f docker/Dockerfile.non_root -t litellm-image-scan:${{ github.sha }} . - - - name: Tag the cached builder for Admin MCP schema setup - run: docker build --target builder -f docker/Dockerfile.non_root -t litellm-admin-mcp-schema:${{ github.sha }} . - - # The prisma bake must migrate a fresh DB with no egress as an arbitrary - # non-root uid (OpenShift restricted-v2 / air-gapped / readOnlyRootFilesystem). - # `docker run` as the default uid with network hides a broken bake because - # the migration entrypoint exits 0 even when it applied nothing; asserting - # the schema was created is what catches it. - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Verify offline migration as a non-root uid - env: - LITELLM_IMAGE: litellm-image-scan:${{ github.sha }} - LITELLM_ADMIN_MCP_SCHEMA_IMAGE: litellm-admin-mcp-schema:${{ github.sha }} - run: | - python -m pip install "pytest==9.0.3" - python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v - # Scans the whole shipped artifact: OS/apk plus every language package # baked into the image, including ones no lockfile declares (e.g. prisma's # vendored node engine) that osv-scan cannot see. osv-scan stays the fast @@ -103,175 +196,12 @@ jobs: # free OSS, run as a pinned, checksum-verified binary; no GitHub Action # dependency and no vendor SaaS callout. - name: Scan image for fixable HIGH/CRITICAL CVEs + if: matrix.check == 'scan' env: GRYPE_MATCH_PYTHON_USING_CPES: "true" run: | - "$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \ + "$RUNNER_TEMP/grype" "$IMAGE" \ --config .grype.yaml \ --only-fixed \ --fail-on high \ --output table - - runtime-image: - name: runtime-image (${{ matrix.dockerfile }}) - runs-on: ubuntu-latest - if: >- - github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name == github.repository - timeout-minutes: 45 - permissions: - contents: read - strategy: - fail-fast: false - matrix: - dockerfile: [Dockerfile, docker/Dockerfile.database] - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Build runtime image - run: docker build --build-arg LITELLM_RELEASE_TAG=v0.0.0-lens-ci -f "${{ matrix.dockerfile }}" -t litellm-runtime-scan:${{ github.sha }} . - - - name: Tag the cached builder for Admin MCP schema setup - run: docker build --target builder -f "${{ matrix.dockerfile }}" -t litellm-admin-mcp-schema:${{ github.sha }} . - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Verify offline migration as a non-root uid - env: - LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }} - LITELLM_ADMIN_MCP_SCHEMA_IMAGE: litellm-admin-mcp-schema:${{ github.sha }} - run: | - python -m pip install "pytest==9.0.3" - python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v - - - migrations-image: - name: migrations-image - runs-on: ubuntu-latest - if: >- - github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name == github.repository - timeout-minutes: 30 - permissions: - contents: read - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Build migrations image - run: docker build -f migrations/Dockerfile -t litellm-migrations-scan:${{ github.sha }} . - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Verify offline migration as a non-root uid - env: - LITELLM_IMAGE: litellm-migrations-scan:${{ github.sha }} - LITELLM_MIGRATION_INTERPRETER: python3 - LITELLM_MIGRATION_SCRIPT: /app/run.py - run: | - python -m pip install "pytest==9.0.3" - python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py -v - - gateway-image: - name: gateway-image - runs-on: ubuntu-latest - if: >- - github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name == github.repository - timeout-minutes: 30 - permissions: - contents: read - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Build gateway image - run: docker build -f gateway/Dockerfile -t litellm-gateway-scan:${{ github.sha }} . - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Verify the gateway serves offline as a non-root uid - env: - LITELLM_IMAGE: litellm-gateway-scan:${{ github.sha }} - LITELLM_COMPONENT_PORT: "4000" - LITELLM_IMAGE_COMPONENT: gateway - run: | - python -m pip install "pytest==9.0.3" - python -m pytest --noconftest tests/integration/database/image/test_component_image_serves_offline.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v - - ui-image: - name: ui-image - runs-on: ubuntu-latest - if: >- - github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name == github.repository - timeout-minutes: 30 - permissions: - contents: read - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Build UI image - run: docker build -f ui/Dockerfile -t litellm-ui-scan:${{ github.sha }} . - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Verify the UI serves offline as an arbitrary uid with a read-only root fs - env: - LITELLM_IMAGE: litellm-ui-scan:${{ github.sha }} - run: | - python -m pip install "pytest==9.0.3" - python -m pytest --noconftest tests/integration/database/image/test_ui_image_serves_offline.py -v - - backend-image: - name: backend-image - runs-on: ubuntu-latest - if: >- - github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name == github.repository - timeout-minutes: 30 - permissions: - contents: read - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Build backend image - run: docker build -f backend/Dockerfile -t litellm-backend-scan:${{ github.sha }} . - - - name: Tag the cached builder for Admin MCP schema setup - run: docker build --target builder -f backend/Dockerfile -t litellm-admin-mcp-schema:${{ github.sha }} . - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Verify the backend serves offline as a non-root uid - env: - LITELLM_IMAGE: litellm-backend-scan:${{ github.sha }} - LITELLM_ADMIN_MCP_SCHEMA_IMAGE: litellm-admin-mcp-schema:${{ github.sha }} - LITELLM_COMPONENT_PORT: "4001" - LITELLM_IMAGE_COMPONENT: backend - run: | - python -m pip install "pytest==9.0.3" - python -m pytest --noconftest tests/integration/database/image/test_component_image_serves_offline.py tests/integration/database/image/test_image_admin_mcp.py -v diff --git a/.github/workflows/interactions-spec-canary.yml b/.github/workflows/interactions-spec-canary.yml index a2e4646dce5..89ebdd87a7c 100644 --- a/.github/workflows/interactions-spec-canary.yml +++ b/.github/workflows/interactions-spec-canary.yml @@ -49,9 +49,9 @@ jobs: id: compliance run: | uv run --no-sync pytest tests/unit/interactions/test_openapi_compliance.py \ - --interactions-spec-url=https://ai.google.dev/static/api/interactions.openapi.json \ + --interactions-spec-url=https://aistudio.google.com/api-reference/interactions.openapi.json \ --tb=short -vv -p no:rerunfailures -p no:pytest-retry \ - --allow-hosts=127.0.0.1,::1,localhost,ai.google.dev + --allow-hosts=127.0.0.1,::1,localhost,ai.google.dev,aistudio.google.com tracking-issue: name: tracking-issue diff --git a/.github/workflows/osv-scan.yml b/.github/workflows/osv-scan.yml index 1c7e8b2841a..5016957bd6f 100644 --- a/.github/workflows/osv-scan.yml +++ b/.github/workflows/osv-scan.yml @@ -5,6 +5,9 @@ on: branches: - main - "litellm_**" + - "rc/**" + - "stable/**" + - "lts/**" schedule: - cron: "23 6 * * *" workflow_dispatch: diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index c3388430589..fd1d2973951 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -21,7 +21,7 @@ jobs: pull-requests: read actions: read runs-on: ubuntu-latest - timeout-minutes: 20 + timeout-minutes: 30 services: redis: image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f @@ -208,7 +208,7 @@ jobs: exit 0 fi retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; } - CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)" + CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)" tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli" PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \ diff --git a/.github/workflows/test-merge-smoke.yml b/.github/workflows/test-merge-smoke.yml index 99dbd1c101b..912fc90fbe5 100644 --- a/.github/workflows/test-merge-smoke.yml +++ b/.github/workflows/test-merge-smoke.yml @@ -200,8 +200,36 @@ jobs: - name: Install the test client dependencies run: .github/scripts/uv_sync_with_retries.sh --frozen --no-install-workspace --extra proxy --group dev --python 3.12 - - name: Build the image from the published package - run: docker build -t litellm-build-from-pip:smoke -f docker/build_from_pip/Dockerfile.build_from_pip . + - name: Cache Prisma binaries + timeout-minutes: 3 + uses: ./.github/actions/cache-prisma-binaries + + - name: Install the published proxy into its own venv + run: | + uv --no-config venv --python 3.13 /tmp/litellm-pip + uv --no-config pip install --python /tmp/litellm-pip/bin/python \ + "litellm[proxy,proxy-runtime]==1.83.0" \ + "google-cloud-aiplatform==1.133.0" \ + "google-genai==1.37.0" \ + "anthropic[vertex]==0.84.0" \ + "grpcio==1.78.0" \ + "prometheus-client==0.20.0" \ + "langfuse==2.59.7" \ + "opentelemetry-api==1.28.0" \ + "opentelemetry-sdk==1.28.0" \ + "opentelemetry-exporter-otlp==1.28.0" \ + "ddtrace==4.11.0" \ + "sentry-sdk==2.21.0" \ + "mangum==0.17.0" \ + "azure-ai-contentsafety==1.0.0" \ + "azure-storage-file-datalake==12.20.0" \ + "pypdf==6.7.5" \ + "llm-sandbox==0.3.31" \ + "detect-secrets==1.5.0" \ + "prisma==0.11.0" \ + "openai==2.24.0" + export PATH="/tmp/litellm-pip/bin:$PATH" + /tmp/litellm-pip/bin/prisma generate --schema=schema.prisma - name: Start the fake OpenAI endpoint run: | @@ -219,23 +247,21 @@ jobs: echo "::add-mask::${key}" echo "LITELLM_MASTER_KEY=${key}" >> "${GITHUB_ENV}" - - name: Start the proxy container + - name: Start the published proxy + env: + DATABASE_URL: postgresql://postgres:postgres@localhost:5432/circle_test + LITELLM_LICENSE: bad-license + FAKE_OPENAI_API_BASE: http://localhost:8190 + OTEL_EXPORTER: in_memory + USE_DDTRACE: "True" + DD_TRACE_OPENAI_ENABLED: "False" + GCS_FLUSH_INTERVAL: "1" + LITELLM_LOG: ERROR run: | - docker run -d \ - -p 4000:4000 \ - -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ - -e LITELLM_MASTER_KEY \ - -e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \ - -e OTEL_EXPORTER=in_memory \ - -e USE_DDTRACE=True \ - -e GCS_FLUSH_INTERVAL=1 \ - -e LITELLM_LOG=ERROR \ - --add-host host.docker.internal:host-gateway \ - --name litellm-build-from-pip \ - -v "${GITHUB_WORKSPACE}/docker/build_from_pip/litellm_config.yaml:/app/config.yaml" \ - litellm-build-from-pip:smoke \ - --config /app/config.yaml \ - --port 4000 + cd /tmp/litellm-pip + nohup /tmp/litellm-pip/bin/ddtrace-run /tmp/litellm-pip/bin/litellm \ + --config "${GITHUB_WORKSPACE}/tests/basic_proxy_startup_tests/build_from_pip_config.yaml" \ + --port 4000 > /tmp/litellm-pip.log 2>&1 & for _ in $(seq 1 150); do curl -sf --max-time 5 -o /dev/null http://localhost:4000 && exit 0 sleep 2 @@ -270,7 +296,7 @@ jobs: - name: Show the proxy logs if: always() - run: docker logs litellm-build-from-pip || true + run: cat /tmp/litellm-pip.log || true smoke-passed: name: smoke passed diff --git a/.github/workflows/test-rust-db.yml b/.github/workflows/test-rust-db.yml index bb88b593d5b..80ec256fa6f 100644 --- a/.github/workflows/test-rust-db.yml +++ b/.github/workflows/test-rust-db.yml @@ -42,6 +42,11 @@ concurrency: env: CARGO_TERM_COLOR: always + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" jobs: db-queries: @@ -87,6 +92,11 @@ jobs: workspaces: litellm-rust cache-on-failure: true + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Install Python dependencies working-directory: . run: uv sync --frozen --extra extra-proxy diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 04af770b24f..cf9fed192f6 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -78,6 +78,11 @@ concurrency: env: CARGO_TERM_COLOR: always + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" jobs: rust-lint: @@ -101,6 +106,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - run: cargo clippy --workspace --all-targets --locked -- -D warnings rust-test: @@ -143,6 +153,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Check generated trace contracts working-directory: . run: uv run scripts/generate_trace_types.py --check @@ -189,6 +204,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - run: uv build --wheel --out-dir dist - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 9570b4d8cc8..d0726e0daa0 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -5,6 +5,9 @@ on: branches: - main - "litellm_**" + - "rc/**" + - "stable/**" + - "lts/**" push: branches: - main @@ -42,7 +45,7 @@ jobs: id: rust-bridge-key if: steps.changes.outputs.decision != 'skip' env: - RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml') }} + RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml', 'packaging/litellm-proxy/pyproject.toml') }} run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT" - name: Set up Python @@ -73,9 +76,21 @@ jobs: timeout-minutes: 5 uses: ./.github/actions/cache-cargo-build + - name: Set up sccache + if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Build editable Rust bridge if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' timeout-minutes: 15 + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" run: | mkdir -p rust-bridge-dist uv run --no-project --with maturin==1.15.0 python -c 'import maturin; maturin.build_editable("rust-bridge-dist")' @@ -83,6 +98,18 @@ jobs: mkdir -p rust-bridge-dist/litellm/rust_bridge cp litellm/rust_bridge/_native.abi3.so rust-bridge-dist/litellm/rust_bridge/_native.abi3.so + - name: Build the gateway wheel once for test shards + if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' + timeout-minutes: 15 + working-directory: packaging/litellm-proxy + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" + run: uv run --no-project --with maturin==1.15.0 maturin build --profile dev --strip --out ../../rust-bridge-dist + - name: Upload Rust bridge artifact if: steps.changes.outputs.decision != 'skip' timeout-minutes: 10 @@ -155,6 +182,7 @@ jobs: tests/unit/google_genai tests/unit/router_strategy tests/unit/router_utils + tests/unit/redis tests/unit/proxy/common_utils/test_cache_aware_routing.py tests/unit/enterprise/enterprise_callbacks/send_emails tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py @@ -227,6 +255,7 @@ jobs: tests/unit/llms/test_oss_decision.py tests/unit/llms/test_polling_url_origin_match.py tests/unit/llms/test_predibase_transformation.py + tests/unit/llms/test_subscription_accounts.py tests/unit/llms/a2a tests/unit/llms/aiml tests/unit/llms/aiohttp_openai @@ -267,6 +296,7 @@ jobs: tests/unit/llms/e2b tests/unit/llms/edenai tests/unit/llms/elevenlabs + tests/unit/llms/empiriolabs tests/unit/llms/exa_ai tests/unit/llms/fal_ai tests/unit/llms/fastcrw @@ -367,6 +397,7 @@ jobs: tests/unit/llms/voyage tests/unit/llms/wandb tests/unit/llms/watsonx + tests/unit/llms/webiq tests/unit/llms/xai tests/unit/llms/xinference tests/unit/llms/you_com @@ -402,6 +433,7 @@ jobs: - shard: router test-path: >- + tests/unit/router_rust tests/unit/test_router python-version: "3.12" workers: 4 @@ -428,6 +460,7 @@ jobs: tests/unit/ocr tests/unit/passthrough tests/unit/rag + tests/unit/cookbook tests/unit/realtime_api tests/unit/rerank_api tests/unit/search @@ -558,12 +591,14 @@ jobs: tests/unit/proxy/credential_endpoints tests/unit/proxy/discovery_endpoints tests/unit/proxy/shutdown + tests/unit/proxy/rust_forward tests/unit/proxy/public_endpoints tests/unit/proxy/prompts tests/unit/proxy/rag_endpoints 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" @@ -599,23 +634,32 @@ jobs: job-timeout-minutes: 60 dist: loadscope - - shard: proxy-infra + - shard: proxy-infra-db-spend test-path: |- tests/unit/proxy/db --ignore=tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py --ignore=tests/unit/proxy/db/test_update_daily_tag_spend.py - tests/unit/proxy/middleware - --ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py tests/unit/proxy/spend_tracking --ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py - tests/unit/proxy/pass_through_endpoints - tests/unit/proxy/mcp_server/test_tool_registry.py tests/unit/proxy/common_utils --ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py --ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py --ignore=tests/unit/proxy/common_utils/test_check_responses_cost.py --ignore=tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py --ignore=tests/unit/proxy/common_utils/test_realtime_cache.py + python-version: "3.12" + workers: 4 + reruns: 0 + timeout-minutes: 20 + job-timeout-minutes: 60 + dist: loadscope + + - shard: proxy-infra-routing + test-path: |- + tests/unit/proxy/middleware + --ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py + tests/unit/proxy/pass_through_endpoints + tests/unit/proxy/mcp_server/test_tool_registry.py tests/unit/proxy/enterprise_billing tests/unit/proxy/types_utils tests/unit/proxy/logging_endpoints @@ -958,10 +1002,10 @@ jobs: dist: loadscope steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }} + timeout-minutes: 3 with: persist-credentials: false - fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }} + fetch-depth: 1 - name: Detect relevant changes id: changes @@ -994,27 +1038,37 @@ jobs: with: artifact: ${{ needs.rust-bridge.outputs.artifact }} - - name: Install dependencies and check schema + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' - timeout-minutes: 8 + timeout-minutes: 5 env: RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }} - UNIT_SHARD: ${{ matrix.shard }} - SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }} run: | diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))" fi uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' - if [ "$UNIT_SHARD" = "db-and-spend" ]; then - uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF" - fi + + - name: Check schema.prisma against the merge base + if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'db-and-spend' && github.event_name == 'pull_request' + timeout-minutes: 3 + env: + GH_TOKEN: ${{ github.token }} + PR_NUMBER: ${{ github.event.pull_request.number }} + EVENT_BASE_REF: ${{ github.event.pull_request.base.ref }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/branches/${EVENT_BASE_REF}" --jq .name) || base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" --jq .base.ref) + merge_base=$(gh api "repos/${GITHUB_REPOSITORY}/compare/${base_ref}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha') + test -n "$merge_base" + git fetch --no-tags --depth=1 origin "$merge_base" + uv run --no-sync python scripts/check_schema_compatibility.py --base-commit "$merge_base" - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' @@ -1251,7 +1305,7 @@ jobs: - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' @@ -1300,17 +1354,16 @@ jobs: - name: Run unit tests run: | - for chart in helm/litellm-helm helm/litellm; do - declared="$(grep -h '^suite:' "$chart"/tests/*.yaml | wc -l | tr -d '[:space:]')" - output="$(mktemp)" - helm unittest -f 'tests/*.yaml' "$chart" | tee "$output" - executed="$(sed -n 's/^Test Suites:.*[[:space:]]\([0-9][0-9]*\) total$/\1/p' "$output")" - if [ "$declared" != "$executed" ]; then - echo "::error::$chart declares $declared test suites but helm-unittest ran $executed. Suites are being skipped silently, so their assertions never execute." - exit 1 - fi - echo "$chart: all $declared declared test suites ran" - done + chart=helm/litellm + declared="$(grep -h '^suite:' "$chart"/tests/*.yaml | wc -l | tr -d '[:space:]')" + output="$(mktemp)" + helm unittest -f 'tests/*.yaml' "$chart" | tee "$output" + executed="$(sed -n 's/^Test Suites:.*[[:space:]]\([0-9][0-9]*\) total$/\1/p' "$output")" + if [ "$declared" != "$executed" ]; then + echo "::error::$chart declares $declared test suites but helm-unittest ran $executed. Suites are being skipped silently, so their assertions never execute." + exit 1 + fi + echo "$chart: all $declared declared test suites ran" openai-sdk-compat: name: OpenAI ${{ matrix.openai-version }} compatibility needs: rust-bridge @@ -1454,9 +1507,9 @@ jobs: RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }} run: | if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so fi diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml deleted file mode 100644 index 0ba4df89cc4..00000000000 --- a/.github/workflows/weekly_load_anomaly.yml +++ /dev/null @@ -1,85 +0,0 @@ -name: "Weekly Load Anomaly Check" - -on: - schedule: - - cron: "0 12 * * 6" - workflow_dispatch: - -permissions: - contents: read - -jobs: - weekly-load-anomaly: - if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 45 - services: - postgres: - image: postgres:16.6 - env: - POSTGRES_USER: llmproxy - POSTGRES_PASSWORD: dbpassword9090 - POSTGRES_DB: litellm - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready -U llmproxy" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - env: - DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm - LITELLM_MASTER_KEY: sk-weekly-anomaly-check - ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} - AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }} - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache the Rust build - uses: ./.github/actions/cache-cargo-build - - - name: Install dependencies - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy - - - name: Cache Prisma binaries - uses: ./.github/actions/cache-prisma-binaries - - - name: Generate Prisma client - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Start the proxy - run: | - nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 & - for _ in $(seq 1 90); do - if curl -fs http://localhost:4000/health/liveliness > /dev/null; then - exit 0 - fi - sleep 2 - done - echo "proxy never became live" - tail -n 100 proxy.log - exit 1 - - - name: Run the weekly session anomaly test - env: - E2E_WEEKLY_ANOMALY: "1" - run: | - uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA - - - name: Show proxy log on failure - if: failure() - run: tail -n 300 proxy.log diff --git a/.gitignore b/.gitignore index b08ee92d6af..1b4785d193a 100644 --- a/.gitignore +++ b/.gitignore @@ -61,7 +61,7 @@ ui/litellm-dashboard/next-env.d.ts *.tsbuildinfo ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json -helm/litellm-helm/*.tgz +helm/litellm/charts/*.tgz helm/*.tgz litellm/proxy/vertex_key.json **/.vim/ diff --git a/AGENTS.md b/AGENTS.md index ce4b63173d2..d8de447be90 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -41,16 +41,17 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just drop the section +Run QA proxies with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --num_workers 2 --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`, and swap its `--config` for a copy of `litellm/proxy/dev_config.yaml` that points both `router_settings` (`redis_host`, `redis_port`) and `general_settings.coordination_redis` (`host`, `port`) at one dedicated Redis instance, with credentials when required, so the workers share router state and spend, budget, and rate-limit counters. Check that two worker PIDs are serving, and use the same worker count and Redis setup for the before and after proof. When the PR is about `--run_gunicorn`, `--run_hypercorn`, or `--run_granian`, which `--reload` does not support, drop `--reload`. When it is about behavior without Redis, drop Redis. When the proxy prints `"workers" flag is ignored when reloading is enabled`, that checkout or its uvicorn reloads only one worker, so drop the two-worker check for it. Say in the proof which part you dropped and why + Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at `, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030) -For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce +For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. The body alone must let a reviewer rerun it, with the proxy command, config, env var names and client code exactly as run, and long blocks collapsed in
. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis - don't use "—". Don't use ";" or ":" unless it would be nonsensical not to. Instead, reach for ",", ".", conjunction words, etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them also feels AI-y. A word cap does not penalize you for adding more sentences - don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. - unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose -- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a ".". Exception is that docs (i.e., litellm-docs) should have trailing "." - don't use →. Instead, prefer not to use arrows, and if need be, use -> instead - use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure @@ -90,6 +91,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..426cf986d4c 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 @@ -303,8 +305,8 @@ uv run litellm --config your_config.yaml If you want to build the Docker image yourself: ```bash -# Build using the non-root Dockerfile -docker build -f docker/Dockerfile.non_root -t litellm_dev . +# Build the image (runs as a non-root user by default) +docker build -t litellm_dev . # Generate a master key. Requests send it as the bearer token export LITELLM_MASTER_KEY="sk-$(openssl rand -hex 32)" diff --git a/Dockerfile b/Dockerfile index 4cef63c8d71..0e9337e658b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -76,6 +76,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ COPY pyproject.toml uv.lock ./ COPY enterprise/pyproject.toml enterprise/ COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/ +COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/ # Install third-party dependencies (cached unless pyproject.toml/uv.lock change) RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ @@ -101,7 +102,9 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh # Install project and workspace packages (fast - deps already cached) -RUN uv sync --frozen --no-default-groups --no-editable \ +ARG MATURIN_PEP517_ARGS="" +RUN uv build --package litellm-proxy --wheel --out-dir /wheels +RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ --group admin-mcp \ @@ -111,12 +114,18 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra bedrock-realtime \ --python python3.13 +RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl + RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ prisma generate --schema=./schema.prisma -RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \ - sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh +RUN sed -i 's/\r$//' docker-entrypoint.sh && chmod +x docker-entrypoint.sh + +RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \ + cp -r litellm/proxy/_experimental/out/. /var/lib/litellm/ui/ && \ + cp litellm/proxy/logo.png /var/lib/litellm/assets/logo.png && \ + touch /var/lib/litellm/ui/.litellm_ui_ready FROM $LITELLM_BUILD_IMAGE AS liteadmin-builder COPY --from=uvbin /uv /usr/local/bin/uv @@ -141,24 +150,35 @@ USER root # https://github.com/BerriAI/litellm/issues/33518 RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories -# node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent +RUN for i in 1 2 3; do \ + apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libatomic libevent nginx && break; \ + [ "$i" = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app + ENV PATH="/app/.venv/bin:${PATH}" \ + PYTHONPATH="/app" \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + HOME=/app \ + LITELLM_NON_ROOT=true \ PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \ PRISMA_CLI_QUERY_ENGINE_TYPE=binary \ + PRISMA_SKIP_POSTINSTALL_GENERATE=1 \ + PRISMA_HIDE_UPDATE_MESSAGE=1 \ + PRISMA_ENGINES_CHECKSUM_IGNORE_MISSING=1 \ PRISMA_OFFLINE_MODE=true # Copy only what runtime needs. The application is installed inside the venv; # the rest of the builder's /app is source and build metadata that must not # ship (manifest-scanning tools attribute everything in it to this image). -# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path. COPY --from=builder /app/.venv /app/.venv COPY --from=liteadmin-builder /opt/liteadmin /opt/liteadmin -COPY --from=builder /app/docker /app/docker +COPY --from=builder /app/docker-entrypoint.sh /app/docker-entrypoint.sh COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py # enterprise/ is imported by source path at runtime (proxy_cli puts the @@ -166,21 +186,33 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr # enterprise.enterprise_hooks from it) COPY --from=builder /app/enterprise /app/enterprise COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras -# Prisma CLI + engines are baked under /opt/prisma, a fixed path every -# runtime uid can read and that no cache volume mount shadows. The paths are -# pinned via PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH and recorded into the -# generated client at build time, so `prisma migrate deploy` on a fresh -# database needs no npm and no network access (#33650, #24554). +COPY --from=builder /app/gateway /app/gateway +COPY --from=builder /app/backend /app/backend +COPY --from=builder /app/migrations /app/migrations COPY --from=builder /opt/prisma /opt/prisma +COPY --from=builder /var/lib/litellm /var/lib/litellm +COPY ui/nginx.conf /etc/nginx/nginx.conf +COPY ui/docker-entrypoint.d/40-litellm-runtime-config.sh /app/ui/docker-entrypoint.d/40-litellm-runtime-config.sh -RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ - find /app/.venv -type d -path "*/tornado/test" -delete && \ +RUN find /app/.venv -depth -type d -path "*/tornado/test" -exec rm -rf {} + && \ + mkdir -p /app/.cache /var/run/litellm && \ + chown -R 65532:0 /app /var/lib/litellm /var/run/litellm && \ + chmod -R g=u,g+w /app/.cache /var/lib/litellm /var/run/litellm && \ + PRISMA_PATH="$(python -c 'import os, prisma; print(os.path.dirname(prisma.__file__))')" && \ + PROXY_EXTRAS_PATH="$(python -c 'import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))')" && \ + chmod -R g=u,g+w "$PRISMA_PATH" "$PROXY_EXTRAS_PATH" && \ chmod -R a+rX /opt/prisma && \ test -x /opt/prisma/binaries/node_modules/.bin/prisma && \ test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \ - python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" + ls /opt/prisma/binaries/node_modules/@prisma/engines/query-engine-* >/dev/null && \ + python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \ + python -c "from litellm.rust_bridge import _native" && \ + python -c "import gateway.launch, backend.main" -EXPOSE 4000/tcp +USER 65532:65532 -ENTRYPOINT ["docker/prod_entrypoint.sh"] -CMD ["--port", "4000"] +RUN nginx -t + +EXPOSE 4000/tcp 4001/tcp 3000/tcp + +ENTRYPOINT ["/app/docker-entrypoint.sh"] diff --git a/Makefile b/Makefile index 48df3c8e215..c256d273d8a 100644 --- a/Makefile +++ b/Makefile @@ -341,7 +341,7 @@ test-integration: install-test-deps $(UV_RUN) pytest tests/ test-unit-helm: install-helm-unittest - helm unittest -f 'tests/*.yaml' helm/litellm-helm + helm unittest -f 'tests/*.yaml' helm/litellm # LLM Translation testing targets test-llm-translation: install-test-deps diff --git a/README.md b/README.md index 3f1f21f903e..31a93bf8c67 100644 --- a/README.md +++ b/README.md @@ -398,6 +398,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | | | [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | | +| [EmpirioLabs AI (`empiriolabs`)](https://docs.litellm.ai/docs/providers/empiriolabs) | ✅ | | ✅ | | | | | | | | | [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | | | [eu/jev Decisions API (`eu_jev`)](https://docs.litellm.ai/docs/pass_through/eu_jev) | | | | | | | | | | | | [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | | @@ -444,6 +445,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [OpenAI-like (`openai_like`)](https://docs.litellm.ai/docs/providers/openai_compatible) | | | | ✅ | | | | | | | | [OpenRouter (`openrouter`)](https://docs.litellm.ai/docs/providers/openrouter) | ✅ | ✅ | ✅ | | | | | | | | | [OVHCloud AI Endpoints (`ovhcloud`)](https://docs.litellm.ai/docs/providers/ovhcloud) | ✅ | ✅ | ✅ | | | | | | | | +| [Parallel AI (`parallel_ai`)](https://docs.litellm.ai/docs/search/parallel_ai) | | | | | | | | | | | | [Perplexity AI (`perplexity`)](https://docs.litellm.ai/docs/providers/perplexity) | ✅ | ✅ | ✅ | | | | | | | | | [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | | | [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | | @@ -696,7 +698,7 @@ For companies that need better security, user management and professional suppor [Talk to founders](https://enterprise.litellm.ai/demo) This covers: -- ✅ **Features under the [LiteLLM Commercial License](https://docs.litellm.ai/docs/proxy/enterprise):** +- ✅ **Features under the [LiteLLM Commercial License](https://docs.litellm.ai/docs/enterprise):** - ✅ **Feature Prioritization** - ✅ **Custom Integrations** - ✅ **Professional Support - Dedicated discord + slack** diff --git a/SECURITY.md b/SECURITY.md index d3603448d46..6d0a9a6414d 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -4,23 +4,19 @@ LiteLLM maintains a security announcements mailing list that is open to anyone. Subscribers receive advance notice, typically one to two days, before we release a fix for a particularly severe vulnerability or for any vulnerability exploitable by an unauthenticated attacker. This notice is provided on a best-effort basis -To subscribe, visit [https://berriai.github.io/security-announce-signup/](https://berriai.github.io/security-announce-signup/) +To subscribe, visit [https://security.litellm.ai/mailinglist/](https://security.litellm.ai/mailinglist/) ## Security Vulnerability Reporting Guidelines -> [!WARNING] -> Reports that do not include a video demonstrating the exploit will be closed without review. See [Reproduction Video Requirement](#reproduction-video-requirement) below. -We value the security community's role in protecting our systems and users. To report a security vulnerability: +We value the security community's role in protecting our systems and users. +To report a security vulnerability, please submit a report through our vulnerability disclosure form: -- File a private vulnerability report on GitHub: [Report a vulnerability](https://github.com/BerriAI/litellm/security/advisories/new) -- Include steps to reproduce the issue -- Include a video or screen recording demonstrating the full exploit against a live LiteLLM instance, from initial access through to impact. A terminal recording (for example asciinema) is fine for CLI-only exploits. -- Provide any relevant additional information +[LiteLLM VDP](https://security.litellm.ai/vdp/) -### Reproduction Video Requirement +These reports are submitted to HackerOne and are triaged there, though no account is needed. -A video demonstrating the exploit is required for every report. AI tools have made it easy to produce plausible-sounding vulnerability reports that do not reproduce in practice, and triaging them takes time away from real issues. Reports submitted without a working reproduction video will be closed without review. If you add a video to a closed report, we will reopen and triage it. +**Please do not submit vulnerabilities as draft GHSAs. As of 2026-11-11 new GHSA submissions will be closed.** ### Vulnerability Categories @@ -55,7 +51,9 @@ To qualify for a bounty, reports must include clear reproduction steps, a reprod ### Known Non-Issues -- Attacks that require a misconfiguration on setup (e.g not setting a `master_key` on the proxy configuration), are **explicitly not in scope** and are not considered vulnerable. +- **Misconfigurations:** attacks that require misconfiguration of the application (e.g. not setting a `master_key` on the proxy configuration). +- **Insecure configurations:** attacks that require deliberate configuration changes to disable security functionality. (e.g. IDOR on response IDs requiring configuration of `disable_responses_id_security` to `true`) +- **Social engineering:** attacks that rely on social engineering as their primary mechanism (e.g. convincing a user's agent to install a malicious package we did not publish). ## Security Measures diff --git a/backend/Dockerfile b/backend/Dockerfile deleted file mode 100644 index aad59ba17c6..00000000000 --- a/backend/Dockerfile +++ /dev/null @@ -1,146 +0,0 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 -ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a - -FROM $UV_IMAGE AS uvbin - -# ---------- Rust extension ---------- -# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes -FROM $LITELLM_BUILD_IMAGE AS rust-native - -WORKDIR /app -USER root - -COPY --from=uvbin /uv /uvx /usr/local/bin/ - -RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \ - [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ - sleep 5; \ - done - -ENV UV_PYTHON_DOWNLOADS=0 - -COPY pyproject.toml README.md rust-toolchain.toml ./ -COPY .cargo ./.cargo -COPY litellm-rust ./litellm-rust -COPY model_prices_and_context_window.json ./ - -RUN --mount=type=cache,target=/root/.cache/uv \ - mkdir -p litellm/rust_bridge /wheels && \ - touch litellm/__init__.py litellm/rust_bridge/__init__.py && \ - uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \ - python -c "import maturin; maturin.build_wheel('/wheels')" - -# ---------- Builder ---------- -FROM $LITELLM_BUILD_IMAGE AS builder - -WORKDIR /app -USER root - -COPY --from=uvbin /uv /uvx /usr/local/bin/ - -# nodejs/npm so `prisma generate` uses Wolfi's Node via PRISMA_USE_GLOBAL_NODE -# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi -# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. -RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \ - [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ - sleep 5; \ - done - -# UV_COMPILE_BYTECODE=1 precompiles .pyc at install time → faster cold start. -# UV_LINK_MODE=copy avoids hardlink warnings when uv installs from a -# BuildKit cache mount (different filesystem). -# UV_PYTHON_DOWNLOADS=0 force uv to use the apk-installed CPython instead of -# silently pulling a managed interpreter. -# PRISMA_USE_GLOBAL_NODE explicit (matches default) so an env override can't -# silently re-enable nodeenv's Node download. -ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ - UV_LINK_MODE=copy \ - UV_COMPILE_BYTECODE=1 \ - UV_PYTHON_DOWNLOADS=0 \ - PRISMA_USE_GLOBAL_NODE=true \ - PATH="/app/.venv/bin:${PATH}" - -# Stage 1 — install dependencies only. -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ - --mount=type=bind,source=uv.lock,target=uv.lock \ - --mount=type=bind,source=enterprise/pyproject.toml,target=enterprise/pyproject.toml \ - --mount=type=bind,source=litellm-proxy-extras/pyproject.toml,target=litellm-proxy-extras/pyproject.toml \ - uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ - --extra proxy \ - --extra proxy-runtime \ - --group admin-mcp \ - --extra extra_proxy \ - --extra semantic-router \ - --extra saml \ - --python python3.13 - -# Stage 2 — copy source and install the project + workspace members. -COPY . . - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --frozen --no-install-project --no-default-groups --no-editable \ - --extra proxy \ - --extra proxy-runtime \ - --group admin-mcp \ - --extra extra_proxy \ - --extra semantic-router \ - --extra saml \ - --python python3.13 - -RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \ - uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl - -RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ - -RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ - npm_config_cache=/root/.npm \ - prisma generate --schema=./schema.prisma - -RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh - -# ---------- Runtime ---------- -FROM $LITELLM_RUNTIME_IMAGE AS runtime -ARG LITELLM_RELEASE_TAG="" -ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG} - -USER root - -RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic && break; \ - [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ - sleep 5; \ - done - -# wolfi-base ships an unprivileged `nonroot` account (UID/GID 65532) with -# /home/nonroot. We run the backend as that user -WORKDIR /app -ENV HOME=/home/nonroot \ - PATH="/app/.venv/bin:${PATH}" \ - PYTHONPATH="/app" \ - PYTHONDONTWRITEBYTECODE=1 \ - PYTHONUNBUFFERED=1 \ - PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries - -COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv -COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm -COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise -COPY --from=builder --chown=nonroot:nonroot /app/backend /app/backend -COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh -COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma -COPY --from=builder /opt/prisma /opt/prisma - -RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ - find /app/.venv -type d -path "*/tornado/test" -delete && \ - chmod -R a+rX /opt/prisma && \ - python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" - -USER nonroot - -EXPOSE 4001/tcp - -ENTRYPOINT ["/app/docker/component_entrypoint.sh", "uvicorn", "backend.main:app"] -CMD ["--host", "0.0.0.0", "--port", "4001"] diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 2a7522cf34e..2d3ccf40175 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -55,6 +55,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/cache_settings", "/coordination_redis/", "/cost_tracking", + "/subscription_accounts/", "/cost_optimization/", "/cost/", "/credentials", diff --git a/ci_cd/security_scans_readme.md b/ci_cd/security_scans_readme.md index dd64b01c296..3003ea1d105 100644 --- a/ci_cd/security_scans_readme.md +++ b/ci_cd/security_scans_readme.md @@ -4,6 +4,5 @@ - Trivy scan on `./docs/` (HIGH/CRITICAL/MEDIUM) - Trivy scan on `./ui/` (HIGH/CRITICAL/MEDIUM) -- Grype scan on `Dockerfile.database` (fails on CRITICAL) - Grype scan on main `Dockerfile` (fails on CRITICAL) - Grype CVSS ≥ 4.0 scan on main `Dockerfile` (fails any vulnerabilities with CVSS ≥ 4.0) diff --git a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md index e1f606a2eb7..538a5f8a404 100644 --- a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -46,7 +46,43 @@ For a managed workspace rollout, use OpenAI's [gateway configuration guide](http ## Native compatibility and Pi -For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name. +For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials. A ChatGPT model keeps its upstream name and needs no deployment of its own. + +### Mix ChatGPT and LiteLLM models + +On the native route, Codex's model picker shows OpenAI's catalog first and then the other chat and responses models the virtual key can use. Pick one of those and Codex's requests go through LiteLLM's `/v1/responses` with the key's usual routing, budgets, guardrails, and logging. The proxy removes the ChatGPT bearer and account header from that request, so the other provider and your callbacks never receive them. A name is routed this way when another provider serves it as an exact `model_name`, through a router `model_group_alias`, or as a team's public model name. A ChatGPT model, a name that only a wildcard deployment matches, or a name OpenAI's catalog lists for the caller's ChatGPT account keeps the passthrough, so a model OpenAI adds later still uses the subscription even when a deployment of another provider has the same name. The proxy learns the account's list from the catalog Codex fetches, or asks OpenAI with the caller's login. When it cannot get the list, for example because OpenAI is down, the request times out, or the client is not Codex, every request keeps the passthrough for the next minute, as it did before LiteLLM models were added to the picker. When the key can use many models, the picker lists the ones that fit in Codex's catalog size limit, and the proxy logs a warning with how many it left out. + +Codex's picker sends a `none` reasoning effort for a LiteLLM model whose catalog entry lists no effort levels. With `litellm_settings: drop_params: true`, the proxy drops it for a model that does not support `none`, and the provider applies its own default. A model whose cost map entry or `model_info` sets `supports_none_reasoning_effort: true` keeps it. Without `drop_params`, the effort reaches the provider unchanged. + +Codex also sends its ChatGPT app tools with every turn, which on a LiteLLM model adds a few hundred thousand prompt tokens to each turn. Set `apps = false` under `[features]` in Codex's `config.toml` to turn them off. + +Only a key that may call `/v1/responses` gets these models. A key whose `allowed_routes` covers only the native routes sees OpenAI's catalog unchanged, and every request it sends uses the subscription. + +Codex sends earlier reasoning items back with each turn. The proxy tags the reasoning that a LiteLLM model returns on this route, OpenAI and Azure deployments included. When a session switches to a ChatGPT model, the proxy drops those items before forwarding the request to OpenAI, which would reject them. Messages, tool calls, and ChatGPT's own reasoning are kept, and a LiteLLM model still gets its own reasoning back. + +When the catalog adds LiteLLM models for a key, its `ETag` is a digest of OpenAI's ETag and those models, and that key's responses carry the same value in `X-Models-Etag`. Codex keeps its cached catalog while the two match and fetches it again when OpenAI's catalog or the key's models change. A key whose catalog adds no LiteLLM models keeps OpenAI's ETags. + +### Restrict native models with an access group + +A key with a model allowlist needs a rule for ChatGPT names that OpenAI adds over time. Put the subscription wildcards in an access group under `chatgpt/` names: + +```yaml +model_list: + - model_name: chatgpt/gpt-* + litellm_params: + model: chatgpt/gpt-* + model_info: + access_groups: ["chatgpt-native"] + - model_name: chatgpt/codex-* + litellm_params: + model: chatgpt/codex-* + model_info: + access_groups: ["chatgpt-native"] +``` + +Then add `chatgpt-native` to the key's models, next to any other model names it should reach. On the native route, a key that holds only the `chatgpt/` form of a name is checked as that form, so a request for `gpt-5.6-sol` is allowed when the key may call `chatgpt/gpt-5.6-sol`. The key can call every ChatGPT model that matches the wildcards, including ones the proxy configuration does not name. + +On `/v1` these wildcards match only a request that asks for a `chatgpt/` name, so a request for a plain `gpt-` or `codex-` name still goes where it went before, such as a `*` catch-all deployment. A `/v1` request for a `chatgpt/` name without a ChatGPT login falls back to the proxy's own saved login if it has one, so keep that login unset as described above. ### Configure Pi @@ -71,7 +107,7 @@ Keep Pi's subscription login as the provider credential. The header value names Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person. -The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions. +The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key needs the `chatgpt-native` access group above or that exact name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions. The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate. diff --git a/cookbook/gpt_realtime_translate.py b/cookbook/gpt_realtime_translate.py new file mode 100644 index 00000000000..827c65cdbfe --- /dev/null +++ b/cookbook/gpt_realtime_translate.py @@ -0,0 +1,245 @@ +#!/usr/bin/env python3 + +import argparse +import asyncio +import base64 +import json +import os +import sys +import wave +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from pathlib import Path +from urllib.parse import urlencode, urlsplit, urlunsplit + +import websockets +from websockets.asyncio.client import ClientConnection + +SAMPLE_RATE = 24_000 +CHANNELS = 1 +SAMPLE_WIDTH = 2 +CHUNK_DURATION_SECONDS = 0.1 +CHUNK_BYTES = int(SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH * CHUNK_DURATION_SECONDS) +RESPONSE_TIMEOUT_SECONDS = 30.0 +AUDIO_EVENT_TYPES = frozenset( + { + "session.output_audio.delta", + "response.audio.delta", + "response.output_audio.delta", + } +) +TRANSCRIPT_EVENT_TYPES = frozenset( + { + "session.output_transcript.delta", + "response.text.delta", + "response.output_audio_transcript.delta", + } +) + + +@dataclass(frozen=True, slots=True) +class Settings: + input_wav: Path + output_wav: Path + base_url: str + model: str + target_language: str + trailing_silence_seconds: float + api_key: str + + +def write_stdout(message: str = "", *, end: str = "\n", flush: bool = False) -> None: + sys.stdout.write(f"{message}{end}") + if flush: + sys.stdout.flush() + + +def write_stderr(message: str) -> None: + sys.stderr.write(f"{message}\n") + + +def parse_args(argv: Sequence[str] | None = None) -> Settings | str: + parser = argparse.ArgumentParser( + description="Stream a 24 kHz PCM16 WAV through gpt-realtime-translate and save the translated audio", + ) + parser.add_argument("input_wav", type=Path) + parser.add_argument("--output", type=Path, default=Path("translated.wav")) + parser.add_argument("--base-url", default=os.getenv("LITELLM_BASE_URL", "http://localhost:4000")) + parser.add_argument("--model", default=os.getenv("REALTIME_TRANSLATE_MODEL", "gpt-realtime-translate")) + parser.add_argument("--target-language", default="fr") + parser.add_argument("--trailing-silence", type=float, default=1.5) + parsed = parser.parse_args(argv) + api_key = os.getenv("LITELLM_API_KEY") or os.getenv("OPENAI_API_KEY") + if not api_key: + return "Set LITELLM_API_KEY or OPENAI_API_KEY before running the script" + if parsed.trailing_silence < 0: + return "--trailing-silence must be zero or greater" + return Settings( + input_wav=parsed.input_wav, + output_wav=parsed.output, + base_url=parsed.base_url, + model=parsed.model, + target_language=parsed.target_language, + trailing_silence_seconds=parsed.trailing_silence, + api_key=api_key, + ) + + +def translation_url(base_url: str, model: str) -> str | None: + parsed = urlsplit(base_url.rstrip("/")) + scheme = {"http": "ws", "https": "wss", "ws": "ws", "wss": "wss"}.get(parsed.scheme) + if not scheme or not parsed.netloc: + return None + base_path = parsed.path.rstrip("/") + realtime_path = ( + f"{base_path}/realtime/translations" if base_path.endswith("/v1") else f"{base_path}/v1/realtime/translations" + ) + return urlunsplit((scheme, parsed.netloc, realtime_path, urlencode({"model": model}), "")) + + +def read_pcm16_wav(path: Path) -> bytes | str: + try: + with wave.open(str(path), "rb") as source: + actual_format = ( + source.getnchannels(), + source.getsampwidth(), + source.getframerate(), + source.getcomptype(), + ) + expected_format = (CHANNELS, SAMPLE_WIDTH, SAMPLE_RATE, "NONE") + if actual_format != expected_format: + return ( + f"{path} must be mono, 16-bit PCM, 24 kHz WAV; received " + f"channels={actual_format[0]}, sample_width={actual_format[1]}, " + f"sample_rate={actual_format[2]}, compression={actual_format[3]}" + ) + return source.readframes(source.getnframes()) + except (OSError, EOFError, wave.Error) as exc: + return f"Unable to read {path}: {exc}" + + +def audio_chunks(audio: bytes) -> Iterator[bytes]: + return (audio[offset : offset + CHUNK_BYTES] for offset in range(0, len(audio), CHUNK_BYTES)) + + +def audio_message(audio: bytes) -> str: + return json.dumps( + { + "type": "session.input_audio_buffer.append", + "audio": base64.b64encode(audio).decode("ascii"), + } + ) + + +async def configure_session(connection: ClientConnection, target_language: str) -> str | None: + await connection.send( + json.dumps( + { + "type": "session.update", + "session": {"audio": {"output": {"language": target_language}}}, + } + ) + ) + while True: + raw_event = await asyncio.wait_for(connection.recv(), timeout=20) + event = json.loads(raw_event) + event_type = event.get("type") + if event_type == "session.created": + write_stdout(f"Session: {event.get('session', {}).get('id', 'created')}") + if event_type == "session.updated": + return None + if event_type == "error": + return f"Realtime API error: {json.dumps(event.get('error', event), ensure_ascii=False)}" + + +async def send_audio( + connection: ClientConnection, pcm: bytes, trailing_silence_seconds: float, finished: asyncio.Event +) -> None: + silence = bytes(round(trailing_silence_seconds * SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH)) + try: + for chunk in audio_chunks(pcm + silence): + await connection.send(audio_message(chunk)) + await asyncio.sleep(len(chunk) / (SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH)) + await connection.send(json.dumps({"type": "session.close"})) + finally: + finished.set() + + +async def receive_translation( + connection: ClientConnection, output_path: Path, sender_finished: asyncio.Event +) -> str | None: + audio_received = asyncio.Event() + try: + with wave.open(str(output_path), "wb") as output: + output.setnchannels(CHANNELS) + output.setsampwidth(SAMPLE_WIDTH) + output.setframerate(SAMPLE_RATE) + write_stdout("Translation: ", end="", flush=True) + while True: + try: + raw_event = await asyncio.wait_for(connection.recv(), timeout=RESPONSE_TIMEOUT_SECONDS) + except TimeoutError: + if not sender_finished.is_set(): + continue + return ( + f"The translation stream went quiet for {RESPONSE_TIMEOUT_SECONDS:g} s without closing the session" + ) + event = json.loads(raw_event) + event_type = event.get("type") + if event_type in AUDIO_EVENT_TYPES: + output.writeframes(base64.b64decode(event.get("delta", ""), validate=True)) + audio_received.set() + elif event_type in TRANSCRIPT_EVENT_TYPES: + write_stdout(event.get("delta", event.get("text", "")), end="", flush=True) + elif event_type == "error": + return f"Realtime API error: {json.dumps(event.get('error', event), ensure_ascii=False)}" + elif event_type == "session.closed": + write_stdout() + return None if audio_received.is_set() else "The translation stream ended without translated audio" + except (OSError, wave.Error) as exc: + return f"Unable to write {output_path}: {exc}" + + +async def translate(settings: Settings, pcm: bytes) -> str | None: + url = translation_url(settings.base_url, settings.model) + if not url: + return f"Invalid --base-url: {settings.base_url}" + sender_finished = asyncio.Event() + try: + async with websockets.connect( + url, + additional_headers={"Authorization": f"Bearer {settings.api_key}"}, + proxy=None, + open_timeout=20, + close_timeout=5, + ) as connection: + configuration_error = await configure_session(connection, settings.target_language) + if configuration_error: + return configuration_error + async with asyncio.TaskGroup() as tasks: + receiver = tasks.create_task(receive_translation(connection, settings.output_wav, sender_finished)) + tasks.create_task(send_audio(connection, pcm, settings.trailing_silence_seconds, sender_finished)) + return receiver.result() + except Exception as exc: + return f"Translation failed: {type(exc).__name__}: {exc}" + + +def main(argv: Sequence[str] | None = None) -> int: + settings = parse_args(argv) + if isinstance(settings, str): + write_stderr(settings) + return 2 + pcm = read_pcm16_wav(settings.input_wav) + if isinstance(pcm, str): + write_stderr(pcm) + return 2 + error = asyncio.run(translate(settings, pcm)) + if error: + write_stderr(error) + return 1 + write_stdout(f"Translated audio: {settings.output_wav}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cookbook/litellm-ollama-docker-image/Dockerfile b/cookbook/litellm-ollama-docker-image/Dockerfile deleted file mode 100644 index be237a4df77..00000000000 --- a/cookbook/litellm-ollama-docker-image/Dockerfile +++ /dev/null @@ -1,25 +0,0 @@ -FROM ollama/ollama as ollama - -RUN echo "auto installing llama2" - -# auto install ollama/llama2 -RUN ollama serve & sleep 2 && ollama pull llama2 - -RUN echo "installing litellm" - -RUN apt-get update - -# Install Python -RUN apt-get install -y python3 python3-pip - -# Set the working directory in the container -WORKDIR /app - -# Copy the current directory contents into the container at /app -COPY . /app - -# Install any needed packages specified in requirements.txt - -RUN python3 -m pip install litellm -COPY start.sh /start.sh -ENTRYPOINT [ "/bin/bash", "/start.sh" ] diff --git a/cookbook/litellm-ollama-docker-image/requirements.txt b/cookbook/litellm-ollama-docker-image/requirements.txt deleted file mode 100644 index 9b9181b2360..00000000000 --- a/cookbook/litellm-ollama-docker-image/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -litellm==1.83.14 diff --git a/cookbook/litellm-ollama-docker-image/start.sh b/cookbook/litellm-ollama-docker-image/start.sh deleted file mode 100644 index ecc03ce73c5..00000000000 --- a/cookbook/litellm-ollama-docker-image/start.sh +++ /dev/null @@ -1,2 +0,0 @@ -ollama serve & -litellm \ No newline at end of file diff --git a/cookbook/litellm-ollama-docker-image/test.py b/cookbook/litellm-ollama-docker-image/test.py deleted file mode 100644 index 93b9c6ac4a7..00000000000 --- a/cookbook/litellm-ollama-docker-image/test.py +++ /dev/null @@ -1,35 +0,0 @@ -import openai - -api_base = "http://0.0.0.0:8000" - -openai.api_base = api_base -openai.api_key = "temp-key" -print(openai.api_base) - - -print("LiteLLM: response from proxy with streaming") -response = openai.ChatCompletion.create( - model="ollama/llama2", - messages=[ - { - "role": "user", - "content": "this is a test request, acknowledge that you got it", - } - ], - stream=True, -) - -for chunk in response: - print(f"LiteLLM: streaming response from proxy {chunk}") - -response = openai.ChatCompletion.create( - model="ollama/llama2", - messages=[ - { - "role": "user", - "content": "this is a test request, acknowledge that you got it", - } - ], -) - -print(f"LiteLLM: response from proxy {response}") diff --git a/cookbook/skill_scanner/README.md b/cookbook/skill_scanner/README.md new file mode 100644 index 00000000000..2a68e116002 --- /dev/null +++ b/cookbook/skill_scanner/README.md @@ -0,0 +1,60 @@ +# Manual skill security scans + +This standalone adapter runs NVIDIA SkillSpector on a separate machine. It is not installed in the LiteLLM proxy and does not require a shared Helm or Compose deployment + +Start it with Python 3.12 or newer and uv. The script pins SkillSpector and its HTTP dependencies + +```sh +export SKILL_SCANNER_API_KEY="$(openssl rand -hex 32)" +uv run server.py +``` + +For another host, place the loopback service behind an HTTPS reverse proxy that forwards the Authorization header and enforces a 65 MiB request limit. Run it under an unprivileged account with a writable temporary directory and no production files or credentials. The scanner host needs Git and network access to requested sources and optional model providers. Apply network egress restrictions suitable for those sources + +Configure the same scanner secret in the gateway's environment, then add this to LiteLLM's configuration + +```yaml +general_settings: + skill_scanner: + endpoint: https://scanner.example.com/scan + api_key_env: SKILL_SCANNER_API_KEY + use_llm: false + timeout_seconds: 180 +``` + +The endpoint must be reachable from the gateway. A customer's laptop works while it is reachable through their network, VPN or authenticated HTTPS tunnel. `localhost` refers to the gateway machine, not the admin's browser. HTTP is accepted only for loopback development + +Open Skills, Submissions, and Review as a proxy admin. Click **Scan for threats**, inspect the report, then independently approve or reject the proposal. Scanning is optional and never changes approval permissions or blocks publication. The report shows findings, locations, remediation, risk score, scanner version and whether semantic analysis actually ran. Incomplete reports do not establish safety + +Uploaded ZIPs are sent privately to the scanner with their exact stored bytes. Metadata-only edits reuse the matching stored archive. If that archive is missing or changed, submit a new ZIP before scanning. The displayed SHA-256 identifies that archive, and the scan examines its files rather than proposed catalog metadata. Remote URL reports describe the source at scan time and do not freeze its contents. GitHub branch/tag refs and explicit-ref subdirectories are supported. Other ref/subdirectory combinations require a ZIP so the adapter does not silently scan the wrong revision + +The latest scan result is retained with the submission. Running scans have a deadline. If a gateway restarts or the scanner becomes unavailable, the attempt expires and an admin can retry. Retries replace the latest report. Approving or rejecting during a scan does not cancel it or change the decision + +Set `use_llm: true` to request semantic analysis and configure the provider on the scanner machine using [SkillSpector's provider instructions](https://github.com/NVIDIA/skillspector#llm-analysis). Provider credentials stay on that machine. Static analysis can also call OSV for dependency checks + +## HTTP contract + +The gateway sends an authenticated `POST` to the configured endpoint. The body contains form fields `source` (JSON source descriptor), `use_llm` (boolean), and `expires_at` (timezone-aware ISO timestamp), plus `file` (ZIP) when bytes are available. The service returns HTTP 200 with this JSON shape + +```json +{ + "scanner_version": "2.12.0", + "complete": true, + "scan_mode": "static-only", + "risk_score": 70, + "findings": [{ + "rule_id": "example-rule", + "severity": "HIGH", + "message": "Potential data exfiltration", + "file": "SKILL.md", + "start_line": 8, + "explanation": "Instructions request sending local files to an external host", + "remediation": "Remove the external upload instruction" + }], + "limitations": [] +} +``` + +Reports are limited to 1 MiB and 500 findings. Scanner uploads allow 64 MiB because the gateway can expand compressed skill uploads into stored compatibility ZIPs. The gateway rejects malformed reports, refuses redirects and verifies TLS. Only the configured scanner credential is sent. A timeout, HTTP error or incomplete report remains visible to the reviewer + +The bundled adapter runs at most two child scans concurrently. Each child and its descendants are terminated at the gateway deadline, capped at 570 seconds. Keep the gateway and scanner clocks synchronized. Direct adapter clients may omit `expires_at` to use the cap. Child processes never execute submitted skill scripts. Raise the gateway deadline up to 600 seconds when needed for semantic analysis and configure the reverse proxy timeout to match diff --git a/cookbook/skill_scanner/server.py b/cookbook/skill_scanner/server.py new file mode 100644 index 00000000000..44a1cecb00f --- /dev/null +++ b/cookbook/skill_scanner/server.py @@ -0,0 +1,172 @@ +# /// script +# requires-python = ">=3.12" +# dependencies = [ +# "fastapi==0.162.0", "uvicorn==0.54.0", "python-multipart==0.0.33", +# "skillspector @ git+https://github.com/NVIDIA/skillspector.git@62c2fe8622f580ac292d4078c3b71f790194dbbd", +# ] +# /// +import asyncio +import hmac +import json +import os +import signal +import sys +import tempfile +from contextlib import suppress +from datetime import datetime, timezone +from pathlib import Path +from typing import Annotated, Final +from urllib.parse import quote, urlsplit + +from fastapi import Depends, FastAPI, File, Form, HTTPException, UploadFile +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import AwareDatetime, TypeAdapter, ValidationError + +MAX_UPLOAD: Final = 64 * 1024 * 1024 +MAX_REPORT: Final = 1024 * 1024 +SCAN_TIMEOUT: Final = 570 +SOURCE: Final = TypeAdapter(dict[str, str]) +SLOTS: Final = asyncio.Semaphore(2) +app: Final = FastAPI(title="SkillSpector adapter", docs_url=None, redoc_url=None, openapi_url=None) + + +def authenticate(credentials: Annotated[HTTPAuthorizationCredentials, Depends(HTTPBearer())]) -> None: + token: Final = os.environ.get("SKILL_SCANNER_API_KEY", "") + if not token or not hmac.compare_digest(credentials.credentials, token): + raise HTTPException(401, "Invalid scanner credentials") + + +def source_target(source: dict[str, str]) -> str: + target: Final = ( + "https://github.com/" + source.get("repo", "") if source.get("source") == "github" else source.get("url", "") + ) + parsed: Final = urlsplit(target) + if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment: + raise HTTPException(400, "Remote sources require an HTTPS URL without credentials or a fragment") + ref: Final = source.get("ref") + subdir: Final = source.get("path") + if ref or subdir: + if parsed.hostname != "github.com" or len(parsed.path.strip("/").split("/")) != 2 or not ref: + raise HTTPException(400, "Upload a ZIP to scan this source's exact ref and subdirectory") + suffix: Final = "/" + quote(subdir, safe="/") if subdir else "" + return target.removesuffix(".git").rstrip("/") + "/tree/" + quote(ref, safe="/") + suffix + return target + + +async def worker(target: str, output: Path, use_llm: bool) -> None: + from skillspector.mcp_server import run_scan + + result: Final = await run_scan(target, use_llm=use_llm, allow_local_targets=not target.startswith("https://")) + coverage: Final = result.get("analysis_completeness", {}) + findings: Final = result["findings"] + complete: Final = ( + result.get("execution_successful") is True + and coverage.get("is_complete") is True + and (not use_llm or (result.get("llm_available") is True and result.get("llm_used") is True)) + and len(findings) <= 500 + ) + limitations: Final = tuple(str(value)[:2000] for value in coverage.get("limitations", ())) + report: Final = { + "scanner_version": result["version"], + "complete": complete, + "scan_mode": result["scan_mode"], + "risk_score": result["risk_score"], + "findings": [ + { + "rule_id": finding["id"], + "severity": finding["severity"], + "message": finding.get("pattern") or finding.get("finding") or "Security finding", + "file": finding["location"]["file"], + "start_line": finding["location"].get("start_line"), + "explanation": finding.get("explanation"), + "remediation": finding.get("remediation"), + } + for finding in findings[:500] + ], + "limitations": ( + *limitations, + *(("Only the first 500 findings are displayed.",) if len(findings) > 500 else ()), + ), + } + encoded: Final = json.dumps(report).encode() + if len(encoded) > MAX_REPORT: + raise ValueError("Scanner report exceeded the response limit") + output.write_bytes(encoded) + + +async def run_worker(target: str, output: Path, use_llm: bool, expires_at: datetime | None = None) -> dict[str, object]: + timeout: Final = ( + SCAN_TIMEOUT + if expires_at is None + else min(SCAN_TIMEOUT, (expires_at - datetime.now(timezone.utc)).total_seconds()) + ) + if timeout <= 0: + raise HTTPException(504, "SkillSpector scan deadline expired") + process: Final = await asyncio.create_subprocess_exec( + sys.executable, + str(Path(__file__).resolve()), + "--scan", + target, + str(output), + str(use_llm).lower(), + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + start_new_session=True, + env={ + **os.environ, + "HOME": str(output.parent), + "TMPDIR": str(output.parent), + "GIT_CONFIG_GLOBAL": os.devnull, + "GIT_CONFIG_NOSYSTEM": "1", + "GIT_TERMINAL_PROMPT": "0", + "GIT_ALLOW_PROTOCOL": "https", + }, + ) + try: + code: Final = await asyncio.wait_for(process.wait(), timeout) + if code != 0 or not output.is_file() or output.stat().st_size > MAX_REPORT: + raise HTTPException(502, "SkillSpector could not produce a report for this source") + return TypeAdapter(dict[str, object]).validate_json(output.read_bytes()) + except TimeoutError: + raise HTTPException(504, "SkillSpector scan timed out") from None + finally: + with suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + + +@app.post("/scan", dependencies=[Depends(authenticate)]) +async def scan( + source: Annotated[str, Form(max_length=16384)], + use_llm: Annotated[bool, Form()] = False, + expires_at: Annotated[AwareDatetime | None, Form()] = None, + file: Annotated[UploadFile | None, File()] = None, +) -> dict[str, object]: + if SLOTS.locked(): + raise HTTPException(429, "Scanner is busy; try again shortly") + try: + parsed: Final = SOURCE.validate_json(source) + except ValidationError: + raise HTTPException(400, "Invalid source metadata") from None + async with SLOTS: + with tempfile.TemporaryDirectory(prefix="skill-scan-") as directory: + root: Final = Path(directory) + if file is not None: + content: Final = await file.read(MAX_UPLOAD + 1) + if len(content) > MAX_UPLOAD: + raise HTTPException(413, "Upload exceeds 64 MiB") + archive: Final = root / "skill.zip" + archive.write_bytes(content) + return await run_worker(str(archive), root / "report.json", use_llm, expires_at) + return await run_worker(source_target(parsed), root / "report.json", use_llm, expires_at) + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "--scan": + asyncio.run(worker(sys.argv[2], Path(sys.argv[3]), sys.argv[4] == "true")) + else: + import uvicorn + + if not os.environ.get("SKILL_SCANNER_API_KEY"): + raise SystemExit("Set SKILL_SCANNER_API_KEY before starting the scanner") + uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("PORT", "8099"))) diff --git a/docker-compose.hardened.yml b/docker-compose.hardened.yml deleted file mode 100644 index 84a23faa054..00000000000 --- a/docker-compose.hardened.yml +++ /dev/null @@ -1,44 +0,0 @@ -services: - # Hardened stack: for testing the proxy under non-root, read-only, proxy-enforced constraints. - # Keep this file focused on hardening/QA scenarios; leave the main docker-compose.yml for default dev usage. - litellm: - build: - context: . - dockerfile: docker/Dockerfile.non_root - target: runtime - depends_on: - - squid - user: "101:101" - group_add: - - "2345" - read_only: true - cap_drop: - - ALL - security_opt: - - no-new-privileges:true - tmpfs: - - /app/cache:rw,noexec,nosuid,nodev,size=128m,uid=101,gid=101,mode=1777 - - /app/migrations:rw,noexec,nosuid,nodev,size=64m,uid=101,gid=101,mode=1777 - volumes: - - ./proxy_server_config.yaml:/app/config.yaml:ro - environment: - LITELLM_NON_ROOT: "true" - PRISMA_BINARY_CACHE_DIR: "/app/cache/prisma-python/binaries" - XDG_CACHE_HOME: "/app/cache" - LITELLM_MIGRATION_DIR: "/app/migrations" - HTTP_PROXY: "http://squid:3128" - HTTPS_PROXY: "http://squid:3128" - NO_PROXY: "localhost,127.0.0.1,db" - command: - - "--port" - - "4000" - - "--config" - - "/app/config.yaml" - squid: - image: sameersbn/squid:3.5.27-2 - restart: unless-stopped - ports: - - "3128:3128" - tmpfs: - - /var/spool/squid:rw,noexec,nosuid,nodev,size=64m - - /var/log/squid:rw,noexec,nosuid,nodev,size=16m diff --git a/docker-compose.yml b/docker-compose.yml index 80e1f289aad..b012c3763de 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,52 +1,71 @@ +# LiteLLM plus a Postgres database that stores models, virtual keys and spend +# logs. Used by https://docs.litellm.ai/docs/proxy/docker_quick_start +# +# curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker-compose.yml +# printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env +# docker compose up -d +# +# `docker compose up` runs the published image. From a checkout of the repo, +# `docker compose up --build` builds the Dockerfile next to this file instead. +# `docker compose --profile monitoring up -d` also starts Prometheus on 9090. +# +# Compose reads .env from this directory. Keep it: regenerating LITELLM_SALT_KEY +# makes credentials already stored in the database unreadable. A database +# created without LITELLM_SALT_KEY encrypted them with the master key, so when +# upgrading one, write only LITELLM_MASTER_KEY (the existing one) to .env. For anything +# beyond local evaluation, pin the image to a specific release tag. services: litellm: + image: docker.litellm.ai/berriai/litellm:main-stable build: context: . - args: - target: runtime - image: docker.litellm.ai/berriai/litellm:main-stable - ######################################### - ## Uncomment these lines to start proxy with a config.yaml file ## - # volumes: - # - ./config.yaml:/app/config.yaml - # command: - # - "--config=/app/config.yaml" - ############################################## + target: runtime + pull_policy: missing ports: - - "4000:4000" # Map the container port to the host, change the host port if necessary + - "4000:4000" + # The same image runs every component. The first word of `command` picks + # it: proxy (default), gateway, backend, ui, migrations, metrics, + # collector. See docker/README.md + # command: ["--config", "/app/config.yaml"] + # volumes: + # - ./config.yaml:/app/config.yaml:ro environment: - DATABASE_URL: "postgresql://llmproxy:dbpassword9090@db:5432/litellm" - # Optional: route read-only queries (find_*, count, group_by, query_raw/_first) - # to a separate reader endpoint, e.g. an Aurora reader. Leave unset for - # single-DB deployments. With IAM_TOKEN_DB_AUTH enabled, the reader URL - # is auto-refreshed alongside the writer. - # DATABASE_URL_READ_REPLICA: "postgresql://llmproxy:dbpassword9090@db-reader:5432/litellm" - STORE_MODEL_IN_DB: "True" # allows adding models to proxy via UI + LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY:?set it in .env, see the header of this file} + LITELLM_MIGRATION_DIR: /tmp/litellm-migrations + DATABASE_URL: postgresql://llmproxy:dbpassword9090@db:5432/litellm + STORE_MODEL_IN_DB: "True" env_file: - - .env # Load local .env file + - path: .env + required: false + read_only: true + cap_drop: + - ALL + security_opt: + - no-new-privileges:true + tmpfs: + - /tmp:rw,noexec,nosuid,nodev,size=64m,mode=1777 + - /app/.cache:rw,noexec,nosuid,nodev,size=128m,mode=1777 depends_on: - - db # Indicates that this service depends on the 'db' service, ensuring 'db' starts first - healthcheck: # Defines the health check configuration for the container + db: + condition: service_healthy + healthcheck: test: - CMD-SHELL - - python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:4000/health/liveliness')" # Command to execute for health check - interval: 30s # Perform health check every 30 seconds - timeout: 10s # Health check command times out after 10 seconds - retries: 3 # Retry up to 3 times if health check fails - start_period: 40s # Wait 40 seconds after container start before beginning health checks + - python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:4000/health/liveliness')" + interval: 30s + timeout: 10s + retries: 3 + start_period: 40s db: image: postgres:16 restart: always - container_name: litellm_db environment: POSTGRES_DB: litellm POSTGRES_USER: llmproxy POSTGRES_PASSWORD: dbpassword9090 - ports: - - "5432:5432" volumes: - - postgres_data:/var/lib/postgresql/data # Persists Postgres data across container restarts + - postgres_data:/var/lib/postgresql/data healthcheck: test: ["CMD-SHELL", "pg_isready -d litellm -U llmproxy"] interval: 1s @@ -55,6 +74,8 @@ services: prometheus: image: prom/prometheus + profiles: + - monitoring volumes: - prometheus_data:/prometheus - ./prometheus.yml:/etc/prometheus/prometheus.yml @@ -68,6 +89,5 @@ services: volumes: prometheus_data: - driver: local postgres_data: - name: litellm_postgres_data # Named volume for Postgres data persistence + name: litellm_postgres_data diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh new file mode 100755 index 00000000000..a48aad33324 --- /dev/null +++ b/docker-entrypoint.sh @@ -0,0 +1,140 @@ +#!/bin/sh +# LiteLLM image entrypoint: docker-entrypoint.sh [COMPONENT] [ARGS...] +# +# COMPONENT selects the process this container runs; it can also be given as +# LITELLM_COMPONENT when the command carries no component. Anything after it is +# passed to that process. A first argument starting with "-" (or no argument at +# all) keeps the historical behaviour of running the monolithic proxy. +# +# proxy litellm ARGS everything in one process (default) +# gateway python -m gateway.launch ARGS inference routes, 0.0.0.0:4000 +# backend uvicorn backend.main:app ARGS management routes, 0.0.0.0:4001 +# ui nginx serving the admin UI port 3000 +# migrations python migrations/run.py prisma migrate deploy, then exit +# metrics python -m litellm.proxy.prometheus_metrics_server ARGS +# collector python -m litellm.proxy.collector ARGS +# --admin-agent the bundled LiteAdmin Slack agent (litellm-admin-agent --web ARGS) +# +# gateway and backend get their host and port defaults first, so ARGS such as +# --port 8080 override them. With no DATABASE_URL, proxy assembles DATABASE_URL +# and DATABASE_URL_READ_REPLICA from DATABASE_HOST and the other discrete +# variables the way gateway and backend do. An unknown first word is executed +# as-is (docker run sh), through ddtrace-run like every component when +# USE_DDTRACE=true. PgBouncer is not a component: LITELLM_PGBOUNCER_ENABLED=true +# starts it inside proxy and gateway. Components that write Prometheus samples +# start with an empty PROMETHEUS_MULTIPROC_DIR; metrics, collector and raw +# readers keep the workers' files, other raw commands clear stale samples. +set -eu + +usage() { + awk 'NR == 1 { next } /^#/ { sub(/^# ?/, ""); print; next } { exit }' "$0" >&2 +} + +component="${LITELLM_COMPONENT:-proxy}" +case "${1:-}" in + proxy|gateway|backend|ui|migrations|metrics|collector) + component="$1" + shift + ;; + --admin-agent) + shift + export CONNECTION_AUTH_MODE=native + exec /opt/liteadmin/bin/litellm-admin-agent --web "$@" + ;; + ""|-*) + ;; + help) + usage + exit 0 + ;; + *) + if [ -n "${PROMETHEUS_MULTIPROC_DIR:-}" ]; then + mkdir -p "$PROMETHEUS_MULTIPROC_DIR" + case "$*" in + *litellm.proxy.prometheus_metrics_server*|*litellm.proxy.collector*) ;; + *) rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db "$PROMETHEUS_MULTIPROC_DIR"/litellm_admitted_series_* ;; + esac + fi + component=raw + ;; +esac + +case "$component" in + raw) + ;; + proxy) + if { [ -z "${DATABASE_URL:-}" ] && [ -n "${DATABASE_HOST:-}" ]; } || + { [ -z "${DATABASE_URL_READ_REPLICA:-}" ] && [ -n "${DATABASE_HOST_READ_REPLICA:-}" ]; }; then + set -- python -c 'import os, sys +from litellm.proxy.db.db_url_settings import ( + DatabaseURLSettings, + add_missing_query_params, + connection_params_from_url, + idle_lifetime_params, + translate_libpq_ssl_params, +) +if os.environ.get("USE_AWS_KMS") == "True": + from litellm.secret_managers.aws_secret_manager import decrypt_env_var + os.environ.update(decrypt_env_var()) +settings = DatabaseURLSettings.from_env() +settings.apply_writer_url_to_env() +reader = settings.build_reader_url() +if reader: + os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params( + add_missing_query_params(reader, connection_params_from_url(os.environ.get("DATABASE_URL", ""))) + ) +if settings.max_idle_connection_lifetime is not None: + lifetime = idle_lifetime_params(settings.max_idle_connection_lifetime) + for name in ("DATABASE_URL", "DATABASE_URL_READ_REPLICA"): + if os.environ.get(name): + os.environ[name] = add_missing_query_params(os.environ[name], lifetime) +os.execvp("litellm", ["litellm", *sys.argv[1:]])' "$@" + else + set -- litellm "$@" + fi + ;; + gateway) + set -- python -m gateway.launch --workers "${NUM_WORKERS:-1}" --host 0.0.0.0 --port 4000 "$@" + ;; + backend) + set -- uvicorn backend.main:app --host 0.0.0.0 --port 4001 "$@" + ;; + ui) + sh "$(dirname "$0")/ui/docker-entrypoint.d/40-litellm-runtime-config.sh" + exec nginx -g 'daemon off;' "$@" + ;; + migrations) + set -- python /app/migrations/run.py "$@" + ;; + metrics) + set -- python -m litellm.proxy.prometheus_metrics_server "$@" + ;; + collector) + set -- python -m litellm.proxy.collector "$@" + ;; + *) + echo "docker-entrypoint.sh: unknown LITELLM_COMPONENT '$component'" >&2 + usage + exit 64 + ;; +esac + +if [ -n "${PROMETHEUS_MULTIPROC_DIR:-}" ]; then + case "$component" in + raw) ;; + metrics|collector) mkdir -p "$PROMETHEUS_MULTIPROC_DIR" ;; + *) + mkdir -p "$PROMETHEUS_MULTIPROC_DIR" + rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db "$PROMETHEUS_MULTIPROC_DIR"/litellm_admitted_series_* + ;; + esac +fi + +case "${USE_DDTRACE:-}" in + [Tt][Rr][Uu][Ee]) + export DD_TRACE_OPENAI_ENABLED="False" + exec ddtrace-run "$@" + ;; +esac + +exec "$@" diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database deleted file mode 100644 index 20f810686a3..00000000000 --- a/docker/Dockerfile.database +++ /dev/null @@ -1,168 +0,0 @@ -# syntax=docker/dockerfile:1.7 - -# Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 - -# Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 -ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a -# Pinned by digest like the other base images; bump explicitly on Node upgrades. -ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 -# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) -ARG PGBOUNCER_VERSION=1.26.0 -ARG PGBOUNCER_SHA256=afd25dd61ee6775d37b40629b87ce08736b3e6955f3057bb212e410fbf21c71d - -FROM $UV_IMAGE AS uvbin - -FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder -ARG PGBOUNCER_VERSION -ARG PGBOUNCER_SHA256 -USER root -RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl -WORKDIR /build -RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ - echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ - tar xzf pgbouncer.tar.gz --strip-components=1 && \ - ./configure --prefix=/usr/local --with-openssl=/usr && \ - make -j"$(nproc)" pgbouncer && \ - install -m 0755 pgbouncer /usr/local/bin/pgbouncer - -# Admin UI builder. Pinned to the build platform so the architecture-independent -# Next.js static export compiles once natively even in a multi-arch build, -# instead of once per target arch under QEMU. -FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder - -ENV NEXT_TELEMETRY_DISABLED=1 \ - npm_config_fund=false \ - npm_config_audit=false - -WORKDIR /ui - -COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./ -COPY ui/litellm-dashboard/vendor/ ./vendor/ -RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline - -COPY ui/litellm-dashboard/ ./ -RUN npm run build - -FROM $LITELLM_BUILD_IMAGE AS builder - -WORKDIR /app -USER root - -COPY --from=uvbin /uv /usr/local/bin/uv -COPY --from=uvbin /uvx /usr/local/bin/uvx - -RUN apk add --no-cache \ - bash \ - gcc \ - python-3.13 \ - python-3.13-dev \ - openssl \ - libssl3 \ - openssl-dev \ - nodejs \ - npm \ - libsndfile - -ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ - UV_LINK_MODE=copy \ - UV_PYTHON_DOWNLOADS=0 \ - PATH="/app/.venv/bin:${PATH}" - -# Copy dependency metadata first for layer caching -COPY pyproject.toml uv.lock ./ -COPY enterprise/pyproject.toml enterprise/ -COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/ - -# Install third-party dependencies (cached unless pyproject.toml/uv.lock change) -RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ - --extra proxy \ - --extra proxy-runtime \ - --group admin-mcp \ - --extra extra_proxy \ - --extra semantic-router \ - --extra saml \ - --extra bedrock-realtime \ - --python python3.13 - -# Copy full source tree -COPY . . - -# Replace the committed UI bundle with the one built from this exact source. -# Clearing first drops the committed bundle's content-hashed chunks that COPY -# would otherwise leave behind alongside the fresh ones. -RUN rm -rf litellm/proxy/_experimental/out -COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ - -# Build Admin UI before final sync (applies the enterprise color override when present) -RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh - -# Install project and workspace packages (fast - deps already cached) -RUN uv sync --frozen --no-default-groups --no-editable \ - --extra proxy \ - --extra proxy-runtime \ - --group admin-mcp \ - --extra extra_proxy \ - --extra semantic-router \ - --extra saml \ - --extra bedrock-realtime \ - --python python3.13 - -RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ - npm_config_cache=/root/.npm \ - prisma generate --schema=./schema.prisma - -RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \ - sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh - -FROM $LITELLM_RUNTIME_IMAGE AS runtime -ARG LITELLM_RELEASE_TAG="" -ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG} - -USER root - -# node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent -COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer - -WORKDIR /app -ENV PATH="/app/.venv/bin:${PATH}" \ - PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ - PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \ - PRISMA_CLI_QUERY_ENGINE_TYPE=binary \ - PRISMA_OFFLINE_MODE=true - -# Copy only what runtime needs. The application is installed inside the venv; -# the rest of the builder's /app is source and build metadata that must not -# ship (manifest-scanning tools attribute everything in it to this image). -# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path. -COPY --from=builder /app/.venv /app/.venv -COPY --from=builder /app/docker /app/docker -COPY --from=builder /app/schema.prisma /app/schema.prisma -COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py -# enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; litellm/proxy/hooks resolves -# enterprise.enterprise_hooks from it) -COPY --from=builder /app/enterprise /app/enterprise -COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras -# Prisma CLI + engines are baked under /opt/prisma, a fixed path every -# runtime uid can read and that no cache volume mount shadows (unlike -# /app/.cache or $HOME/.cache in readOnlyRootFilesystem + emptyDir setups). -# The paths are pinned via PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH and -# recorded into the generated client at build time, so `prisma migrate -# deploy` on a fresh database needs no npm and no network access -# (#33650, #24554). -COPY --from=builder /opt/prisma /opt/prisma - -RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ - find /app/.venv -type d -path "*/tornado/test" -delete && \ - chmod -R a+rX /opt/prisma && \ - test -x /opt/prisma/binaries/node_modules/.bin/prisma && \ - test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \ - python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" - -EXPOSE 4000/tcp - -ENTRYPOINT ["docker/prod_entrypoint.sh"] -CMD ["--port", "4000"] diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root deleted file mode 100644 index 90bcc523d9e..00000000000 --- a/docker/Dockerfile.non_root +++ /dev/null @@ -1,208 +0,0 @@ -# syntax=docker/dockerfile:1.7 - -# Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 -ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a -# Pinned by digest like the other base images; bump explicitly on Node upgrades. -ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 -# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) -ARG PGBOUNCER_VERSION=1.26.0 -ARG PGBOUNCER_SHA256=afd25dd61ee6775d37b40629b87ce08736b3e6955f3057bb212e410fbf21c71d - -FROM $UV_IMAGE AS uvbin - -FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder -ARG PGBOUNCER_VERSION -ARG PGBOUNCER_SHA256 -USER root -RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl -WORKDIR /build -RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ - echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ - tar xzf pgbouncer.tar.gz --strip-components=1 && \ - ./configure --prefix=/usr/local --with-openssl=/usr && \ - make -j"$(nproc)" pgbouncer && \ - install -m 0755 pgbouncer /usr/local/bin/pgbouncer - -# Admin UI builder. Pinned to the build platform so the architecture-independent -# Next.js static export compiles once natively even in a multi-arch build, -# instead of once per target arch under QEMU. -FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder - -ENV NEXT_TELEMETRY_DISABLED=1 \ - npm_config_fund=false \ - npm_config_audit=false - -WORKDIR /ui - -COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./ -COPY ui/litellm-dashboard/vendor/ ./vendor/ -RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline - -COPY ui/litellm-dashboard/ ./ -RUN npm run build - -FROM $LITELLM_BUILD_IMAGE AS builder -WORKDIR /app -USER root - -COPY --from=uvbin /uv /usr/local/bin/uv -COPY --from=uvbin /uvx /usr/local/bin/uvx - -RUN for i in 1 2 3; do \ - apk add --no-cache \ - python-3.13 \ - python-3.13-dev \ - gcc \ - rust \ - bash \ - coreutils \ - curl \ - openssl \ - libssl3 \ - libsndfile \ - nodejs \ - npm && break || sleep 5; \ - done - -ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ - UV_LINK_MODE=copy \ - UV_PYTHON_DOWNLOADS=0 \ - PATH="/app/.venv/bin:${PATH}" \ - LITELLM_NON_ROOT=true \ - XDG_CACHE_HOME=/app/.cache - -# Copy dependency metadata first for layer caching -COPY pyproject.toml uv.lock ./ -COPY enterprise/pyproject.toml enterprise/ -COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/ - -# Install third-party dependencies (cached unless pyproject.toml/uv.lock change) -RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ - uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ - --extra proxy \ - --extra proxy-runtime \ - --group admin-mcp \ - --extra extra_proxy \ - --extra semantic-router \ - --extra saml \ - --extra bedrock-realtime \ - --python python3.13 - -# Copy full source tree -COPY . . - -# Replace the committed UI bundle with the one built from this exact source. -# Clearing first drops the committed bundle's content-hashed chunks that COPY -# would otherwise leave behind alongside the fresh ones. -RUN rm -rf litellm/proxy/_experimental/out -COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ - -# Set non-root flag for build time consistency -ENV LITELLM_NON_ROOT=true - -RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \ - cp -r /app/litellm/proxy/_experimental/out/. /var/lib/litellm/ui/ && \ - cp /app/litellm/proxy/logo.png /var/lib/litellm/assets/logo.png && \ - touch /var/lib/litellm/ui/.litellm_ui_ready - -RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ - uv sync --frozen --no-default-groups --no-editable \ - --extra proxy \ - --extra proxy-runtime \ - --group admin-mcp \ - --extra extra_proxy \ - --extra semantic-router \ - --extra saml \ - --extra bedrock-realtime \ - --python python3.13 - -RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ - npm_config_cache=/root/.npm \ - prisma generate --schema=./schema.prisma - -RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \ - sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh - -FROM $LITELLM_RUNTIME_IMAGE AS runtime -ARG LITELLM_RELEASE_TAG="" -ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG} -WORKDIR /app -USER root - -RUN for i in 1 2 3; do \ - apk upgrade --no-cache && break || sleep 5; \ - done && \ - for i in 1 2 3; do \ - apk add --no-cache python-3.13 bash openssl libssl3 tzdata libsndfile nodejs libevent && break || sleep 5; \ - done -COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer - -# Copy only what runtime needs. The application is installed inside the venv; -# the rest of the builder's /app is source and build metadata that must not -# ship (manifest-scanning tools attribute everything in it to this image). -# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path. -COPY --from=builder /app/.venv /app/.venv -COPY --from=builder /app/docker /app/docker -COPY --from=builder /app/schema.prisma /app/schema.prisma -COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py -# enterprise/ is imported by source path at runtime (proxy_cli puts the -# working directory on sys.path; litellm/proxy/hooks resolves -# enterprise.enterprise_hooks from it) -COPY --from=builder /app/enterprise /app/enterprise -COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras -# Prisma CLI + engines are baked under /opt/prisma, a fixed path every runtime -# uid can read and that no cache volume mount shadows (unlike /app/.cache or -# $HOME/.cache under readOnlyRootFilesystem + emptyDir or arbitrary-uid setups). -# PRISMA_CLI_QUERY_ENGINE_TYPE=binary makes the CLI use the baked binary query -# engine directly, so `prisma migrate deploy` on a fresh database needs no npm -# and no network access; without it the CLI looks for the library engine, which -# prisma stopped baking, and falls back to a download that fails offline or as a -# non-writable uid (#33650, #24554). -COPY --from=builder /opt/prisma /opt/prisma -COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui -COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets - -# XDG_CACHE_HOME is intentionally left unset so it falls back to $HOME/.cache -# (/app/.cache, writable by the runtime uid). The prisma bake at the read-only -# /opt/prisma is anchored by PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH, so -# nothing needs XDG to point there; pointing it at the read-only bake would -# deny any XDG-aware library that writes a cache at runtime. -ENV PATH="/app/.venv/bin:${PATH}" \ - PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ - PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \ - PRISMA_CLI_QUERY_ENGINE_TYPE=binary \ - HOME=/app \ - LITELLM_NON_ROOT=true \ - PRISMA_SKIP_POSTINSTALL_GENERATE=1 \ - PRISMA_HIDE_UPDATE_MESSAGE=1 \ - PRISMA_ENGINES_CHECKSUM_IGNORE_MISSING=1 \ - PRISMA_OFFLINE_MODE=true - -RUN mkdir -p /nonexistent /app/.cache /var/lib/litellm/assets /var/lib/litellm/ui && \ - chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent && \ - PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \ - chown -R nobody:nogroup "$PRISMA_PATH" && \ - LITELLM_PKG_MIGRATIONS_PATH="$(python -c 'import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))' 2>/dev/null || echo '')/migrations" && \ - [ -n "$LITELLM_PKG_MIGRATIONS_PATH" ] && chown -R nobody:nogroup "$LITELLM_PKG_MIGRATIONS_PATH" || true && \ - LITELLM_PROXY_EXTRAS_PATH=$(python -c "import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))" 2>/dev/null || echo "") && \ - chgrp -R 0 "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets && \ - [ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chgrp -R 0 "$LITELLM_PROXY_EXTRAS_PATH" || true && \ - chmod -R g=u "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets && \ - [ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g=u "$LITELLM_PROXY_EXTRAS_PATH" || true && \ - chmod -R g+w "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets && \ - [ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w "$LITELLM_PROXY_EXTRAS_PATH" || true && \ - chmod -R g+rX "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets && \ - chmod -R a+rX /opt/prisma && \ - test -x /opt/prisma/binaries/node_modules/.bin/prisma && \ - test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \ - ls /opt/prisma/binaries/node_modules/@prisma/engines/query-engine-* >/dev/null 2>&1 && \ - python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" - -USER 65534 - -EXPOSE 4000/tcp - -ENTRYPOINT ["/app/docker/prod_entrypoint.sh"] -CMD ["--port", "4000"] diff --git a/docker/README.md b/docker/README.md index ac82c6caebf..3ec26219420 100644 --- a/docker/README.md +++ b/docker/README.md @@ -2,15 +2,15 @@ This guide provides instructions for building and running the LiteLLM application using Docker and Docker Compose. -> **Just want to run LiteLLM?** This guide builds from source. To run the published -> image instead, use `docker-compose.quickstart.yml` in this directory — the -> two-service stack (gateway + Postgres) that the -> [Docker quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start) documents: +> **Just want to run LiteLLM?** `docker-compose.yml` in the repository root runs +> the published image with a Postgres database, the stack the +> [Docker quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start) documents, +> and it works on its own outside a checkout: > > ```bash -> curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml +> curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker-compose.yml > printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env -> docker compose -f docker-compose.quickstart.yml up -d +> docker compose up -d > ``` ## Prerequisites @@ -20,33 +20,52 @@ This guide provides instructions for building and running the LiteLLM applicatio ## Building and Running the Application -To build and run the application, you will use the `docker-compose.yml` file located in the root of the project. This file is configured to use the `Dockerfile.non_root` for a secure, non-root container environment. +The same `docker-compose.yml` builds from source when you pass `--build`: it builds the `Dockerfile` in the repository root, the one image LiteLLM ships. `docker-entrypoint.sh` next to it is the only entrypoint script; every container starts through it -### 1. Set the Master Key +## One image, many components -The application requires a `LITELLM_MASTER_KEY` for signing and validating tokens. You must set this key as an environment variable before running the application. +Every LiteLLM container runs the same image. The first word of the container command (or the `LITELLM_COMPONENT` environment variable when the command carries only flags) picks the process the container runs, and everything after it is handed to that process unchanged: -Create a `.env` file in the root of the project and add the following line: +| Component | Runs | Port | +|--------------|-----------------------------------------------------|------| +| `proxy` | `litellm ...` (everything in one process, default) | 4000 | +| `gateway` | `python -m gateway.launch ...` (inference routes) | 4000 | +| `backend` | `uvicorn backend.main:app ...` (management routes) | 4001 | +| `ui` | nginx serving the static admin UI | 3000 | +| `migrations` | `python migrations/run.py`, `prisma migrate deploy` then exit | | +| `metrics` | `python -m litellm.proxy.prometheus_metrics_server ...` | `--port` | +| `collector` | `python -m litellm.proxy.collector ...` | | -``` -LITELLM_MASTER_KEY=your-secret-key +```bash +docker run -p 4000:4000 litellm --config /app/config.yaml # proxy, exactly as before +docker run -p 4000:4000 litellm gateway --port 4000 # componentized data plane +docker run -p 4001:4001 -e LITELLM_COMPONENT=backend litellm # same, chosen through the env +docker run -p 3000:3000 --read-only --tmpfs /tmp litellm ui # admin UI behind nginx +docker run -e DATABASE_URL=... litellm migrations # one-off schema migration job +docker run -it litellm sh # anything else runs verbatim ``` -Replace `your-secret-key` with a strong, randomly generated secret. +PgBouncer is not a separate component: `LITELLM_PGBOUNCER_ENABLED=true` starts an in-container PgBouncer in front of `DATABASE_URL` inside `proxy` and `gateway`. `USE_DDTRACE=true` wraps whichever component runs with `ddtrace-run`, and `PROMETHEUS_MULTIPROC_DIR` is emptied of stale samples before any workers fork (the `metrics` and `collector` sidecars only read it, so their restart keeps the live samples) + +The image runs as uid `65532` (`nonroot` in the Wolfi base) and also works as an arbitrary uid in gid 0, the shape OpenShift `restricted-v2` assigns, because everything it writes at runtime lives under `/app/.cache`, `/var/lib/litellm` and `/tmp`. Mount those (or set `readOnlyRootFilesystem` with emptyDirs there) for a read-only root filesystem. Prisma's CLI and engines are baked under `/opt/prisma`, so migrations need neither network nor a writable home + +### 1. Set the Master and Salt Keys + +The proxy signs virtual keys with `LITELLM_MASTER_KEY` and encrypts stored provider credentials with `LITELLM_SALT_KEY`. Compose reads both from a `.env` file in the directory you run it from, so generate them once: + +```bash +printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env +``` + +Keep the file: regenerating `LITELLM_SALT_KEY` makes credentials already stored in the database unreadable. A database created without `LITELLM_SALT_KEY` encrypted them with the master key, so when upgrading one, write only the existing `LITELLM_MASTER_KEY` to `.env` and leave `LITELLM_SALT_KEY` out. Provider keys such as `OPENAI_API_KEY` go in the same file, and the whole file is passed to the container ### 2. Build and Run the Containers -Once you have set the `LITELLM_MASTER_KEY`, you can build and run the containers using the following command: - ```bash docker compose up -d --build ``` -This command will: - -- Build the Docker image using `Dockerfile.non_root`. -- Start the `litellm`, `litellm_db`, and `prometheus` services in detached mode (`-d`). -- The `--build` flag ensures that the image is rebuilt if there are any changes to the Dockerfile or the application code. +This command builds the image from the root `Dockerfile` and starts the `litellm` and `db` services in detached mode. Without `--build`, `docker compose up` pulls the published `main-stable` image instead. Add `--profile monitoring` to also start Prometheus on port 9090, scraping the proxy with the root `prometheus.yml` ### 3. Verifying the Application is Running @@ -80,13 +99,13 @@ LITELLM_LICENSE="your-enterprise-license" PROXY_BASE_URL=https://gateway.example.com ``` -For the unified source deployment described above, put them in its `.env` file and rebuild: +For the compose deployment described above, put them in its `.env` file and rebuild: ```bash docker compose up -d --build ``` -In componentized deployments, set the flag and license on the backend container and route `/admin/mcp` to the backend service. The gateway component excludes this endpoint. The unified, database, non-root, and backend image builds bundle the connector +In componentized deployments, set the flag and license on the `backend` component and route `/admin/mcp` to it. The `gateway` component excludes this endpoint. The image bundles the connector for every component Hosting is disabled by default. Opting in requires a valid base Enterprise license; an unlicensed opt-in or invalid flag value prevents startup. Enabling it reserves `/admin`, so rename any MCP server alias called `admin` first @@ -96,34 +115,18 @@ Embedded responses default to `full`; selecting `LITELLM_ADMIN_RESPONSE_VIEW=com See the [LiteAdmin MCP guide](https://docs.litellm.ai/docs/proxy/liteadmin_mcp#run-liteadmin-mcp-inside-litellm) for client configuration, tool restrictions, and verification -## Hardened / Offline Testing +## Hardening -To ensure changes are safe for non-root, read-only root filesystems and restricted egress, always validate with the hardened compose file: +The compose file runs the proxy the way a locked-down cluster would: as the image's non-root user with a read-only root filesystem, every capability dropped, `no-new-privileges` set, and tmpfs mounts only at `/tmp` and `/app/.cache`. The image is built for that, so a change that makes the proxy write anywhere else fails here before it fails in Kubernetes. To try an arbitrary uid in gid 0, the shape OpenShift `restricted-v2` assigns, add `user: "101:0"` to the `litellm` service + +Prisma's CLI and engines are baked under `/opt/prisma`, so migrations run without network access. Verify with: ```bash -docker compose -f docker-compose.yml -f docker-compose.hardened.yml build --no-cache -docker compose -f docker-compose.yml -f docker-compose.hardened.yml up -d +docker run --rm --network none --entrypoint prisma docker.litellm.ai/berriai/litellm:main-stable --version ``` -This setup: -- Builds from `docker/Dockerfile.non_root` with Prisma engines and Node toolchain baked into the image. -- Runs the proxy as a non-root user with a read-only rootfs and only writable tmpfs mounts: - - `/app/cache` (Prisma/NPM cache; backing `PRISMA_BINARY_CACHE_DIR`, `NPM_CONFIG_CACHE`, `XDG_CACHE_HOME`) - - `/app/migrations` (Prisma migration workspace; backing `LITELLM_MIGRATION_DIR`) -- Pre-builds and serves the admin UI from read-only paths: - - `/var/lib/litellm/ui` (pre-restructured Next.js UI with `.litellm_ui_ready` marker) - - `/var/lib/litellm/assets` (UI logos and assets) -- Routes all outbound traffic through a local Squid proxy that denies egress, so Prisma migrations must use the cached CLI and engines. - -You should also verify offline Prisma behaviour with: - -```bash -docker run --rm --network none --entrypoint prisma ghcr.io/berriai/litellm:main-stable --version -``` - -This command should succeed (showing engine versions) even with `--network none`, confirming that Prisma binaries are available without network access. - ## Troubleshooting -- **`build_admin_ui.sh: not found`**: This error can occur if the Docker build context is not set correctly. Ensure that you are running the `docker-compose` command from the root of the project. -- **`Master key is not initialized`**: This error means the `LITELLM_MASTER_KEY` environment variable is not set. Make sure you have created a `.env` file in the project root with the `LITELLM_MASTER_KEY` defined. +- **`required variable LITELLM_MASTER_KEY is missing a value`**: Compose did not find a `.env` file with `LITELLM_MASTER_KEY` in the directory you ran it from. Generate one as shown above. +- **`password authentication failed for user "llmproxy"`**: the `litellm_postgres_data` volume was initialised by an older compose file with different credentials. `docker compose down -v` drops it and the next `up` recreates the database. +- **`build_admin_ui.sh: not found`**: the build context is wrong. Run `docker compose` from the root of the repository. diff --git a/docker/build_from_pip/Dockerfile.build_from_pip b/docker/build_from_pip/Dockerfile.build_from_pip deleted file mode 100644 index 25a7ed2fa97..00000000000 --- a/docker/build_from_pip/Dockerfile.build_from_pip +++ /dev/null @@ -1,64 +0,0 @@ -ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.10.9@sha256:10902f58a1606787602f303954cea099626a4adb02acbac4c69920fe9d278f82 -FROM $UV_IMAGE AS uvbin - -FROM python:3.13-slim@sha256:739e7213785e88c0f702dcdc12c0973afcbd606dbf021a589cab77d6b00b579d - -ARG LITELLM_VERSION=1.83.0 - -WORKDIR /app - -COPY --from=uvbin /uv /usr/local/bin/uv -COPY --from=uvbin /uvx /usr/local/bin/uvx - -RUN apt-get update && \ - apt-get install -y --no-install-recommends gcc libffi-dev nodejs npm && \ - rm -rf /var/lib/apt/lists/* - -ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ - UV_LINK_MODE=copy \ - PATH="/app/.venv/bin:${PATH}" - -COPY schema.prisma . - -# This image is specifically for validating/installing the published PyPI -# artifact, not the checked-out source tree. -# Keep the moved proxy-runtime packages explicit until the published PyPI -# artifact includes that extra; newer releases will simply dedupe these. -RUN uv venv --python python && \ - uv pip install --python /app/.venv/bin/python \ - "litellm[proxy,proxy-runtime]==${LITELLM_VERSION}" \ - "google-cloud-aiplatform==1.133.0" \ - "google-genai==1.37.0" \ - "anthropic[vertex]==0.84.0" \ - "grpcio==1.84.1" \ - "prometheus-client==0.20.0" \ - "langfuse==2.59.7" \ - "opentelemetry-api==1.28.0" \ - "opentelemetry-sdk==1.28.0" \ - "opentelemetry-exporter-otlp==1.28.0" \ - "ddtrace==4.11.0" \ - "sentry-sdk==2.21.0" \ - "mangum==0.17.0" \ - "azure-ai-contentsafety==1.0.0" \ - "azure-storage-file-datalake==12.20.0" \ - "pypdf==6.7.5" \ - "llm-sandbox==0.3.31" \ - "detect-secrets==1.5.0" \ - "prisma==0.11.0" \ - "openai==2.24.0" - -RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ - npm_config_cache=/root/.npm \ - prisma generate --schema=./schema.prisma && \ - chmod -R a+rX /opt/prisma && \ - python -c "import sys; from prisma.client import BINARY_PATHS; bad = sorted(p for group in BINARY_PATHS.model_dump().values() for p in group.values() if not p.startswith('/opt/prisma/')); sys.exit('prisma engines baked outside /opt/prisma: %r' % bad) if bad else None" - -ENV PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries - -COPY docker/prod_entrypoint.sh /app/docker/prod_entrypoint.sh -RUN sed -i 's/\r$//' /app/docker/prod_entrypoint.sh && chmod +x /app/docker/prod_entrypoint.sh - -EXPOSE 4000/tcp - -ENTRYPOINT ["/app/docker/prod_entrypoint.sh"] -CMD ["--port", "4000"] diff --git a/docker/build_from_pip/Readme.md b/docker/build_from_pip/Readme.md deleted file mode 100644 index ad043588f33..00000000000 --- a/docker/build_from_pip/Readme.md +++ /dev/null @@ -1,9 +0,0 @@ -# Docker to build LiteLLM Proxy from litellm pip package - -### When to use this ? - -If you need to build LiteLLM Proxy from litellm pip package, you can use this Dockerfile as a reference. - -### Why build from pip package ? - -- If your company has a strict requirement around security / building images you can follow steps outlined here \ No newline at end of file diff --git a/docker/component_entrypoint.sh b/docker/component_entrypoint.sh deleted file mode 100755 index 3cd530e3009..00000000000 --- a/docker/component_entrypoint.sh +++ /dev/null @@ -1,16 +0,0 @@ -#!/bin/sh - -# stale samples from a previous container incarnation would be summed into the aggregate -if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then - mkdir -p "$PROMETHEUS_MULTIPROC_DIR" - rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db "$PROMETHEUS_MULTIPROC_DIR"/litellm_admitted_series_* -fi - -case "$USE_DDTRACE" in - [Tt][Rr][Uu][Ee]) - export DD_TRACE_OPENAI_ENABLED="False" - exec ddtrace-run "$@" - ;; -esac - -exec "$@" diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh deleted file mode 100755 index 003d9b21db8..00000000000 --- a/docker/entrypoint.sh +++ /dev/null @@ -1,16 +0,0 @@ -#!/bin/bash -set -euo pipefail - -REPO_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" -VENV_PYTHON="$REPO_ROOT/.venv/bin/python" -MIGRATION_SCRIPT="$REPO_ROOT/litellm/proxy/prisma_migration.py" - -if [ -x "$VENV_PYTHON" ]; then - "$VENV_PYTHON" "$MIGRATION_SCRIPT" -elif command -v uv >/dev/null 2>&1; then - (cd "$REPO_ROOT" && uv run --no-sync python "$MIGRATION_SCRIPT") -else - python3 "$MIGRATION_SCRIPT" -fi - -echo "Migration script ran successfully!" diff --git a/docker/install_auto_router.sh b/docker/install_auto_router.sh deleted file mode 100755 index 4fedf201b41..00000000000 --- a/docker/install_auto_router.sh +++ /dev/null @@ -1,4 +0,0 @@ -#!/bin/bash -set -euo pipefail - -# semantic-router dependencies are installed via `uv sync`. diff --git a/docker/prod_entrypoint.sh b/docker/prod_entrypoint.sh deleted file mode 100644 index 4386be65a32..00000000000 --- a/docker/prod_entrypoint.sh +++ /dev/null @@ -1,16 +0,0 @@ -#!/bin/sh - -if [ "$1" = "--admin-agent" ]; then - shift - export CONNECTION_AUTH_MODE=native - exec /opt/liteadmin/bin/litellm-admin-agent --web "$@" -fi - -case "$USE_DDTRACE" in - [Tt][Rr][Uu][Ee]) - export DD_TRACE_OPENAI_ENABLED="False" - exec ddtrace-run litellm "$@" - ;; -esac - -exec litellm "$@" diff --git a/docker/tests/nonroot.yaml b/docker/tests/nonroot.yaml deleted file mode 100644 index 36118ca8c59..00000000000 --- a/docker/tests/nonroot.yaml +++ /dev/null @@ -1,18 +0,0 @@ -schemaVersion: 2.0.0 - -metadataTest: - entrypoint: ["docker/prod_entrypoint.sh"] - user: "65534" - workdir: "/app" - -fileExistenceTests: - - name: "Prisma Folder" - path: "/usr/local/lib/python3.13/site-packages/prisma/" - shouldExist: true - uid: 65534 - gid: 65534 - - name: "Prisma Schema" - path: "/usr/local/lib/python3.13/site-packages/prisma/schema.prisma" - shouldExist: true - uid: 65534 - gid: 65534 diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index f920aa7ac13..a0c896b7e05 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -17,6 +17,7 @@ from typing import List, Optional, Union from litellm._logging import verbose_logger from litellm.caching import DualCache +from litellm.exceptions import MissingMessageRoleError from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -86,6 +87,8 @@ class PagerDutyAlerting(SlackAlerting): Record a failure event. Only send an alert to PagerDuty if the configured *failure* threshold is exceeded in the specified window. """ + if isinstance(kwargs.get("exception"), MissingMessageRoleError): + return now = datetime.now(timezone.utc) standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( "standard_logging_object" 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..abc2487a4a5 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -8,8 +8,8 @@ import re from collections.abc import Callable, Mapping, Sequence 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, @@ -20,6 +20,7 @@ from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs from litellm.constants import ( BATCH_LINE_ITEMS_MAX_PART_BYTES, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, @@ -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", @@ -747,8 +766,7 @@ class CheckBatchCost: credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {} _file_content = await afile_content( file_id=cast(str, raw_output_file_id), - _litellm_internal_model_credentials=MappingProxyType(dict(credentials)), - **cast(FileContentCallOptions, credentials), + **cast(FileContentCallOptions, {**credentials, **stored_s3_bucket_kwargs(cast(str, raw_output_file_id))}), ) content_bytes: Final = await _file_content_bytes(_file_content) @@ -848,10 +866,7 @@ class CheckBatchCost: error_file_failed_requests: Final = await count_error_file_failed_requests( response, custom_llm_provider=batch_file_provider, - litellm_params={ - **credentials, - "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), - }, + litellm_params=credentials, ) batch_result: Final = ( output_file_result @@ -1010,11 +1025,7 @@ class CheckBatchCost: custom_llm_provider=deployment_info.litellm_params.custom_llm_provider, ) cost_provider: Final = cast(BatchFileProvider, llm_provider) - credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {} - file_access: Final = { - **credentials, - "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), - } + file_access: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {} output_content: Final = ( await fetch_batch_managed_file_content( response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index a4a40a7a3d7..1c433d04819 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -182,7 +182,7 @@ class ManagedObjectCreatorAttribution: team_id: Final = getattr(job, "team_id", None) rows: Final = _CreatorRows( key=await self._live_key_row(api_key, job.created_by), - team=await self._live_team_row(team_id), + team=await self._live_team_row(team_id, job.created_at), user=await self._user_row(object_id, _creator_user_id(job.created_by)), ) return ResponseCreator( @@ -269,10 +269,15 @@ class ManagedObjectCreatorAttribution: verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}") return None - async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: + async def _live_team_row( + self, team_id: str | None, created_before: datetime | None = None + ) -> prisma_models.LiteLLM_TeamTable | None: if not team_id: return None - return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + team: Final = await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + if team is not None and created_before is not None and team.created_at > created_before: + return None + return team def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5b54e7d72a0..bbb09e75b47 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,9 +3,11 @@ import asyncio import base64 +import copy import json import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta from itertools import chain from types import MappingProxyType @@ -21,6 +23,7 @@ from typing import ( TypedDict, Union, cast, + runtime_checkable, ) from uuid import NAMESPACE_URL, uuid5 @@ -33,7 +36,7 @@ from typing_extensions import ReadOnly, Unpack import litellm from litellm import Router, verbose_logger -from litellm._internal_context import with_service_target +from litellm._internal_context import internal_sub_call, post_response_phase, with_service_target from litellm._uuid import uuid from litellm.caching.caching import DualCache from litellm.constants import ( @@ -46,6 +49,7 @@ from litellm.constants import ( ) from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.internal_call_metadata import requests_background @@ -75,6 +79,7 @@ from litellm.proxy.litellm_pre_call_utils import ( from litellm.proxy.openai_files_endpoints.common_utils import ( BATCH_CREATE_HIDDEN_PARAM, FILE_LIST_CONTINUATION_CHUNK_SIZE, + TeamProviderDeployment, _is_base64_encoded_unified_file_id, apply_unified_file_ids, decode_model_from_file_id, @@ -94,6 +99,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_metadata, ) from litellm.repositories.managed_file_repository import ManagedFileRepository +from litellm.responses.streaming_iterator import spawn_logging_task from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue] AllMessageValues, @@ -123,7 +129,6 @@ class _ManagedFileRetrieve(Protocol): self, *, file_id: str, - _litellm_internal_model_credentials: Mapping[str, object] | None = None, **kwargs: Unpack[FileRetrieveCallOptions], ) -> OpenAIFileObject: ... @@ -374,6 +379,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: ... @@ -386,6 +397,127 @@ class _CursorPageArgs(TypedDict, total=False): class _RouterFileCallKwargs(TypedDict, total=False): client: ReadOnly[AsyncOpenAI | None] custom_llm_provider: ReadOnly[str | None] + stored_s3_bucket_name: ReadOnly[str] + + +_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str]) +_FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"}) +_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id"}) + + +@dataclass(frozen=True, slots=True) +class _ProviderUpload: + model_id: str + provider_file_id: str + upload_call_id: str | None = None + + +def _provider_uploads(file_object: OpenAIFileObject) -> tuple[_ProviderUpload, ...]: + file_hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic file metadata + dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR) + ) + model_file_id_mapping: Final = _MODEL_FILE_ID_MAPPING.validate_python( + file_hidden_params.get("model_file_id_mapping") or {} + ) + upload_call_id: Final = file_hidden_params.get("litellm_call_id") + return tuple( + _ProviderUpload( + model_id=model_id, + provider_file_id=provider_file_id, + upload_call_id=upload_call_id if isinstance(upload_call_id, str) else None, + ) + for model_id, provider_file_id in model_file_id_mapping.items() + ) + + +def _keeps_one_upload_per_target(llm_router: Router, model_id: str) -> bool: + credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id) + return (credentials or {}).get("custom_llm_provider") == LlmProviders.LITELLM_PROXY.value + + +def _is_kept_upload(uploads: Sequence[_ProviderUpload], index: int, per_target_model_ids: frozenset[str]) -> bool: + upload: Final = uploads[index] + other_uploads: Final = uploads[index + 1 :] if upload.model_id in per_target_model_ids else uploads[:index] + return all(other.model_id != upload.model_id for other in other_uploads) + + +def _upload_from_trailing_id_fields(decoded_unified_file_id: str) -> _ProviderUpload: + trailing_fields: Final = decoded_unified_file_id.rpartition(";llm_output_file_id,")[2] + provider_file_id, _, model_id = trailing_fields.partition(";llm_output_file_model_id,") + return _ProviderUpload(model_id=model_id, provider_file_id=provider_file_id) + + +def _carried_upload_missing_from(unified_file_id: str, model_file_ids: Mapping[str, str]) -> _ProviderUpload | None: + decoded: Final = _is_base64_encoded_unified_file_id(unified_file_id) + carried: Final = _upload_from_trailing_id_fields(decoded) if decoded else None + if carried is None or model_file_ids.get(carried.model_id) in (None, carried.provider_file_id): + return None + return carried + + +async def _delete_provider_file( + llm_router: Router, + model_id: str, + provider_file_id: str, + request_data: Mapping[str, object], +) -> None: + delete_data: Final = {k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS} + await llm_router.afile_delete( + model=model_id, + file_id=provider_file_id, + **cast(_RouterFileCallKwargs, {**delete_data, **stored_s3_bucket_kwargs(provider_file_id)}), + ) + + +def _with_own_nested_dicts(request_data: Mapping[str, object]) -> MappingProxyType[str, object]: + return MappingProxyType({k: copy.copy(v) if isinstance(v, dict) else v for k, v in request_data.items()}) + + +async def _delete_redundant_upload( + llm_router: Router, + upload: _ProviderUpload, + delete_request: Mapping[str, object], + parent_call_id: str | None, +) -> None: + try: + with internal_sub_call(parent_call_id): + await _delete_provider_file( + llm_router, upload.model_id, upload.provider_file_id, _with_own_nested_dicts(delete_request) + ) + except Exception as error: + verbose_logger.warning( + f"Could not delete the redundant upload provider_file_id={sanitize_for_log(upload.provider_file_id)} " + f"at deployment model_id={sanitize_for_log(upload.model_id)}: " + f"{type(error).__name__} {sanitize_for_log(error)}" + ) + + +async def _delete_redundant_uploads( + llm_router: Router, + uploads: Sequence[_ProviderUpload], + delete_request: Mapping[str, object], +) -> None: + await asyncio.gather( + *(_delete_redundant_upload(llm_router, upload, delete_request, upload.upload_call_id) for upload in uploads) + ) + + +def _start_deleting_redundant_uploads( + llm_router: Router, + uploads: Sequence[_ProviderUpload], + create_file_request: CreateFileRequest, +) -> None: + if not uploads: + return + delete_request: Final = MappingProxyType( + {k: v for k, v in create_file_request.items() if k not in _FILE_UPLOAD_ONLY_FIELDS} + ) + with post_response_phase(): + spawn_logging_task( + asyncio.get_running_loop(), + _delete_redundant_uploads(llm_router, uploads, delete_request), + task_name="managed file repeat upload cleanup", + ) def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase: @@ -479,6 +611,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 +649,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 +853,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 +884,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 +923,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}, @@ -984,26 +1154,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): *, provider_file_id: str, call_options: Mapping[str, object], - internal_model_credentials: Mapping[str, object] | None = None, ) -> OpenAIFileObject: """Retrieve provider file details with transient retries and SDK retries disabled.""" retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve) retrieve_options: Final = cast( FileRetrieveCallOptions, - {**call_options, "max_retries": 0}, + {**call_options, **stored_s3_bucket_kwargs(provider_file_id), "max_retries": 0}, ) for attempt in range(len(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS) + 1): if attempt > 0: await self._sleep(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS[attempt - 1]) try: - if internal_model_credentials is None: - return await retrieve_file( - file_id=provider_file_id, - **retrieve_options, - ) return await retrieve_file( file_id=provider_file_id, - _litellm_internal_model_credentials=MappingProxyType(dict(internal_model_credentials)), **retrieve_options, ) except Exception as error: @@ -1045,7 +1208,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self._afile_retrieve_with_retries( provider_file_id=provider_file_id, call_options=router_credentials, - internal_model_credentials=router_credentials, ), timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, ) @@ -1267,6 +1429,43 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): user_api_key_dict=owner, ) + async def register_provider_file( + self, + *, + provider_file_id: str, + content_type: str, + deployment: TeamProviderDeployment, + owner: UserAPIKeyAuth, + llm_router: Router | None, + ) -> str | None: + unified_file_id: Final = self.get_unified_output_file_id( + provider_file_id, deployment.deployment_id, deployment.public_model_name, file_type=content_type + ) + provider_file_object, _ = await self._fetch_provider_file_object( + unified_file_id=unified_file_id, + provider_file_id=provider_file_id, + model_id=deployment.deployment_id, + model_name=None, + llm_router=llm_router, + ) + if provider_file_object is None: + return None + try: + await self.store_unified_file_id( + file_id=unified_file_id, + file_object=provider_file_object, + litellm_parent_otel_span=None, + model_mappings={deployment.deployment_id: provider_file_id}, + user_api_key_dict=owner, + ) + except Exception as error: # noqa: BLE001 # the upload succeeded, so the caller keeps the provider id + verbose_logger.warning( + f"Failed to register provider_file_id={sanitize_for_log(provider_file_id)} as a managed file: " + f"{type(error).__name__} {sanitize_for_log(error)}" + ) + return None + return unified_file_id + async def list_user_batches( self, user_api_key_dict: UserAPIKeyAuth, @@ -1598,8 +1797,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if potential_file_id and "llm_output_file_id," in potential_file_id: model_id = self.get_model_id_from_unified_file_id(potential_file_id) if model_id: + provider_file_id: Final = self.get_output_file_id_from_unified_file_id(potential_file_id) data["model"] = model_id - data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id) + data["file_id"] = provider_file_id + data.update(stored_s3_bucket_kwargs(provider_file_id)) elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: @@ -1989,16 +2190,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): target_model_names_list=target_model_names_list, ) - ## STORE MODEL MAPPINGS IN DB - model_mappings: Dict[str, str] = {} - - for file_object in responses: - file_hidden_params = cast( # cast-ok: preserve mapping operations on dynamic file metadata - dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR) - ) - model_file_id_mapping = file_hidden_params.get("model_file_id_mapping") - if model_file_id_mapping and isinstance(model_file_id_mapping, dict): - model_mappings.update(model_file_id_mapping) + uploads: Final = tuple(chain.from_iterable(map(_provider_uploads, responses))) + per_target_model_ids: Final = frozenset( + model_id + for model_id in {upload.model_id for upload in uploads} + if _keeps_one_upload_per_target(llm_router, model_id) + ) + kept_uploads: Final = tuple( + upload for index, upload in enumerate(uploads) if _is_kept_upload(uploads, index, per_target_model_ids) + ) + model_mappings: Final = {upload.model_id: upload.provider_file_id for upload in kept_uploads} await self.store_unified_file_id( file_id=response.id, @@ -2007,6 +2208,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings=model_mappings, user_api_key_dict=user_api_key_dict, ) + _start_deleting_redundant_uploads( + llm_router, + tuple( + upload + for upload in uploads + if upload not in kept_uploads and upload.model_id not in per_target_model_ids + ), + create_file_request, + ) # Emit Prometheus metrics for managed file creation prom_logger = self._get_prometheus_logger() @@ -2085,10 +2295,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") - def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str: + def get_unified_output_file_id( + self, output_file_id: str, model_id: str, model_name: Optional[str], file_type: str = "application/json" + ) -> str: deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id) unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - "application/json", + file_type, str(deterministic_uuid), model_name or "", output_file_id, @@ -2170,6 +2382,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( @@ -2646,23 +2859,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id) if not specific_model_file_id_mapping: return - filtered_data: Final = { - k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials") - } for model_id, model_file_id in specific_model_file_id_mapping.items(): - credentials = llm_router.registry().credentials(model_id, form="with_provider") - delete_data = { - **filtered_data, - **( - {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))} - if credentials is not None - else {} - ), - } - await llm_router.afile_delete( - model=model_id, - file_id=model_file_id, - **cast(_RouterFileCallKwargs, delete_data), + await _delete_provider_file(llm_router, model_id, model_file_id, data) + unrecorded_upload: Final = _carried_upload_missing_from(file_id, specific_model_file_id_mapping) + delete_call_id: Final = data.get("litellm_call_id") + if unrecorded_upload is not None: + await _delete_redundant_upload( + llm_router, unrecorded_upload, data, delete_call_id if isinstance(delete_call_id, str) else None ) async def afile_content( @@ -2690,21 +2893,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): exception_dict = {} for model_id, provider_file_id in specific_model_file_id_mapping.items(): try: - # Cloud-storage providers (e.g. Bedrock S3) validate file ids - # against the deployment's configured bucket, which they only - # trust from this immutable server-side snapshot, never from - # request params. - credentials = llm_router.registry().credentials(model_id, form="with_provider") - if credentials is not None: - data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials))) - else: - data.pop("_litellm_internal_model_credentials", None) return cast( HttpxBinaryResponseContent, await llm_router.afile_content( model=model_id, file_id=provider_file_id, - **cast(_RouterFileCallKwargs, data), + **cast( + _RouterFileCallKwargs, + { + **data, + **( + stored_s3_bucket_kwargs(provider_file_id) + if isinstance(provider_file_id, str) + else {} + ), + }, + ), ), ) except Exception as e: diff --git a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py index a2322d7f650..bb3b712d50b 100644 --- a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py +++ b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py @@ -1,9 +1,9 @@ import enum -from typing import Dict, List, Optional from pydantic import Field from litellm.proxy._types import WebhookEvent +from litellm.types.integrations.model_deprecation_email import MODEL_DEPRECATION_EMAIL_EVENT from litellm.types.llms.base import LiteLLMBaseModel @@ -26,7 +26,7 @@ class SendKeyCreatedEmailEvent(WebhookEvent): class SendKeyRotatedEmailEvent(WebhookEvent): virtual_key: str - key_alias: Optional[str] = None + key_alias: str | None = None """ The virtual key that was rotated this will be sk-123xxx, since we will be emailing this to the user to start using the new key @@ -34,23 +34,33 @@ class SendKeyRotatedEmailEvent(WebhookEvent): class EmailEvent(str, enum.Enum): + model_deprecations = MODEL_DEPRECATION_EMAIL_EVENT virtual_key_created = "Virtual Key Created" new_user_invitation = "New User Invitation" virtual_key_rotated = "Virtual Key Rotated" soft_budget_crossed = "Soft Budget Crossed" max_budget_alert = "Max Budget Alert" + class EmailEventSettings(LiteLLMBaseModel): event: EmailEvent enabled: bool + + class EmailEventSettingsUpdateRequest(LiteLLMBaseModel): - settings: List[EmailEventSettings] + settings: list[EmailEventSettings] + + class EmailEventSettingsResponse(LiteLLMBaseModel): - settings: List[EmailEventSettings] + settings: list[EmailEventSettings] + + class DefaultEmailSettings(LiteLLMBaseModel): """Default settings for email events""" - settings: Dict[EmailEvent, bool] = Field( + + settings: dict[EmailEvent, bool] = Field( default_factory=lambda: { + EmailEvent.model_deprecations: False, # Explicit opt-in for owner retirement reminders. EmailEvent.virtual_key_created: True, # On by default EmailEvent.new_user_invitation: True, # On by default EmailEvent.virtual_key_rotated: True, # On by default @@ -58,10 +68,12 @@ class DefaultEmailSettings(LiteLLMBaseModel): EmailEvent.max_budget_alert: True, # On by default } ) - def to_dict(self) -> Dict[str, bool]: + + def to_dict(self) -> dict[str, bool]: """Convert to dictionary with string keys for storage""" return {event.value: enabled for event, enabled in self.settings.items()} + @classmethod - def get_defaults(cls) -> Dict[str, bool]: + def get_defaults(cls) -> dict[str, bool]: """Get the default settings as a dictionary with string keys""" - return cls().to_dict() \ No newline at end of file + return cls().to_dict() diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 7db2de1f874..64fe03769b8 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.81" +version = "0.1.83" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.81" +version = "0.1.83" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/gateway/AGENTS.md b/gateway/AGENTS.md new file mode 100644 index 00000000000..1bb4013179f --- /dev/null +++ b/gateway/AGENTS.md @@ -0,0 +1,16 @@ +# Gateway image entry point + +`python -m gateway.launch` is how the gateway image starts. `launch.py` hands its arguments to uvicorn's own CLI and serves `main.py`, which trims the proxy app to the LLM data plane + +- Importing `launch.py` stays cheap. Every uvicorn worker imports it again before it can answer its first health ping, so LiteLLM imports wait until `main` or `worker_app` runs +- PgBouncer starts exactly once per pod, before the workers start, and the workers inherit the loopback URL through `DATABASE_URL`. `configure_database` is the only place that does this + +## Rust forwarding + +`launch.py` is an entry point of Rust forwarding, and every forwarding decision lives in `litellm/proxy/rust_forward/`, whose `AGENTS.md` states the contract + +- `main` dispatches on `rust_forward_mode` and does nothing else about forwarding. With forwarding off or refused it serves exactly as before +- `python_server` reads the flags and environment through uvicorn's own CLI parser, so Rust gets the values uvicorn would run with. Never reparse a uvicorn flag by hand +- Under forwarding the supervisor process runs `configure_database`, as `before_start`, once forwarding is known to be possible. The supervised child never runs it and inherits the exported URL +- The supervised child appends `mode.uvicorn_flags` after the operator's flags, so the private listener and `--no-proxy-headers` win +- `worker_app` wraps the app with `peer.wrap`, which is the identity when forwarding is off diff --git a/gateway/Dockerfile b/gateway/Dockerfile deleted file mode 100644 index 658bed831ba..00000000000 --- a/gateway/Dockerfile +++ /dev/null @@ -1,165 +0,0 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7 -ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a -# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) -ARG PGBOUNCER_VERSION=1.26.0 -ARG PGBOUNCER_SHA256=afd25dd61ee6775d37b40629b87ce08736b3e6955f3057bb212e410fbf21c71d - -FROM $UV_IMAGE AS uvbin - -FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder -ARG PGBOUNCER_VERSION -ARG PGBOUNCER_SHA256 -USER root -RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl -WORKDIR /build -RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ - echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ - tar xzf pgbouncer.tar.gz --strip-components=1 && \ - ./configure --prefix=/usr/local --with-openssl=/usr && \ - make -j"$(nproc)" pgbouncer && \ - install -m 0755 pgbouncer /usr/local/bin/pgbouncer - -# ---------- Rust extension ---------- -# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes -FROM $LITELLM_BUILD_IMAGE AS rust-native - -WORKDIR /app -USER root - -COPY --from=uvbin /uv /uvx /usr/local/bin/ - -RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \ - [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ - sleep 5; \ - done - -ENV UV_PYTHON_DOWNLOADS=0 - -COPY pyproject.toml README.md rust-toolchain.toml ./ -COPY .cargo ./.cargo -COPY litellm-rust ./litellm-rust -COPY model_prices_and_context_window.json ./ - -RUN --mount=type=cache,target=/root/.cache/uv \ - mkdir -p litellm/rust_bridge /wheels && \ - touch litellm/__init__.py litellm/rust_bridge/__init__.py && \ - uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \ - python -c "import maturin; maturin.build_wheel('/wheels')" - -# ---------- Builder ---------- -FROM $LITELLM_BUILD_IMAGE AS builder - -WORKDIR /app -USER root - -COPY --from=uvbin /uv /uvx /usr/local/bin/ - -# nodejs/npm so `prisma generate` uses Wolfi's Node via PRISMA_USE_GLOBAL_NODE -# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi -# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. -RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \ - [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ - sleep 5; \ - done - -# UV_COMPILE_BYTECODE=1 precompiles .pyc at install time → faster cold start. -# UV_LINK_MODE=copy avoids hardlink warnings when uv installs from a -# BuildKit cache mount (different filesystem). -# UV_PYTHON_DOWNLOADS=0 force uv to use the apk-installed CPython instead of -# silently pulling a managed interpreter. -# PRISMA_USE_GLOBAL_NODE explicit (matches default) so an env override can't -# silently re-enable nodeenv's Node download. -ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ - UV_LINK_MODE=copy \ - UV_COMPILE_BYTECODE=1 \ - UV_PYTHON_DOWNLOADS=0 \ - PRISMA_USE_GLOBAL_NODE=true \ - PATH="/app/.venv/bin:${PATH}" - -# Stage 1 — install dependencies only. -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ - --mount=type=bind,source=uv.lock,target=uv.lock \ - --mount=type=bind,source=enterprise/pyproject.toml,target=enterprise/pyproject.toml \ - --mount=type=bind,source=litellm-proxy-extras/pyproject.toml,target=litellm-proxy-extras/pyproject.toml \ - uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ - --extra proxy \ - --extra proxy-runtime \ - --group admin-mcp \ - --extra extra_proxy \ - --extra semantic-router \ - --extra bedrock-realtime \ - --python python3.13 - -# Stage 2 — copy source and install the workspace members. -COPY . . - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --frozen --no-install-project --no-default-groups --no-editable \ - --extra proxy \ - --extra proxy-runtime \ - --group admin-mcp \ - --extra extra_proxy \ - --extra semantic-router \ - --extra bedrock-realtime \ - --python python3.13 - -RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \ - uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl - -# PYTHONPATH=/app makes the source tree shadow the installed package, so the -# compiled Rust extension must live next to the source or it is never imported. -RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ - -RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ - npm_config_cache=/root/.npm \ - prisma generate --schema=./schema.prisma - -RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh - -# ---------- Runtime ---------- -FROM $LITELLM_RUNTIME_IMAGE AS runtime - -USER root - -RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic libevent && break; \ - [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ - sleep 5; \ - done - -# wolfi-base ships an unprivileged `nonroot` account (UID/GID 65532) with -# /home/nonroot. We run the proxy as that user. -WORKDIR /app -ENV HOME=/home/nonroot \ - PATH="/app/.venv/bin:${PATH}" \ - PYTHONPATH="/app" \ - PYTHONDONTWRITEBYTECODE=1 \ - PYTHONUNBUFFERED=1 \ - PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries - -COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv -COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm -COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise -COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway -COPY --from=builder --chown=nonroot:nonroot /app/litellm_cli /app/litellm_cli -COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh -COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma -COPY --from=builder /opt/prisma /opt/prisma -COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer - -RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ - find /app/.venv -type d -path "*/tornado/test" -delete && \ - chmod -R a+rX /opt/prisma && \ - python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \ - python -c "import litellm, litellm.rust_bridge._native; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__" - -USER nonroot - -EXPOSE 4000/tcp - -ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh python -m gateway.launch --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"] -CMD ["--host", "0.0.0.0", "--port", "4000"] diff --git a/gateway/launch.py b/gateway/launch.py index 59f7121e513..c8b5b53af92 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -22,15 +22,16 @@ import sys from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final +from typing_extensions import assert_never from uvicorn.main import main as uvicorn_main from litellm_cli.main import without_gc_pauses if TYPE_CHECKING: - from fastapi import FastAPI - from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings + from litellm.proxy.rust_forward.peer import ASGIApp + from litellm.proxy.rust_forward.server import PythonServer GATEWAY_APP: Final = "gateway.launch:worker_app" KEEPALIVE_FLAG: Final = "--timeout-keep-alive" @@ -65,17 +66,40 @@ def pool_database_url( return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth()) -def worker_app() -> "FastAPI": +def worker_app() -> "ASGIApp": with without_gc_pauses(): from gateway.main import app - return app + from litellm.proxy.rust_forward.mode import rust_forward_mode + from litellm.proxy.rust_forward.peer import wrap + return wrap(app, rust_forward_mode(os.environ)) def _serve(argv: Sequence[str]) -> None: uvicorn_main(tuple(argv), prog_name="uvicorn") -def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: +def python_server(argv: Sequence[str]) -> "PythonServer": + """What uvicorn would serve the public port with, read from the same flags and environment it reads.""" + from litellm.proxy.rust_forward.server import PythonServer, Server, UvicornCliParams + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + params: Final = UvicornCliParams.model_validate( + uvicorn_main.make_context("uvicorn", list(uvicorn_argv(argv, os.environ))).params + ) + return PythonServer( + server=Server.UVICORN, + listener=params.listener, + host=params.host, + port=params.port, + keepalive_seconds=params.timeout_keep_alive, + graceful_timeout_seconds=GracefulShutdownManager.get_timeout() + if params.timeout_graceful_shutdown is None + else params.timeout_graceful_shutdown, + trusted_proxies=params.trusted_proxies, + ) + + +def configure_database() -> None: from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, export_pooled_database_url @@ -86,8 +110,59 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) - sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}") if pooled_url is not None: export_pooled_database_url(pooled_url) + + +def _serve_python_alone(argv: Sequence[str], serve: Callable[[Sequence[str]], None]) -> None: + configure_database() serve(uvicorn_argv(argv, os.environ)) +def _python_server_or_exit(argv: Sequence[str]) -> "PythonServer": + """Report ``--help``, ``--version`` and a bad flag the way uvicorn's own CLI does, since it never gets to.""" + import click + + try: + return python_server(argv) + except click.exceptions.Exit as stopped: + sys.exit(stopped.exit_code) + except click.ClickException as error: + error.show() + sys.exit(error.exit_code) + + +def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: + from pydantic import ValidationError + + from litellm.proxy.rust_forward.mode import ( + Mode, + Supervised, + environment_error, + rust_forward_mode, + stop_with_supervisor, + uvicorn_flags, + ) + from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor.preflight import StartupError + + try: + mode: Final = rust_forward_mode(os.environ) + except ValidationError as error: + sys.exit(environment_error(error)) + match mode: + case Mode.DISABLED: + _serve_python_alone(argv, serve) + case Mode.SUPERVISOR: + supervised: Final = supervise(_python_server_or_exit(argv), before_start=configure_database) + if isinstance(supervised, StartupError): + _serve_python_alone(argv, serve) + else: + sys.exit(supervised) + case Supervised(): + stop_with_supervisor(mode) + serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode))) + case _: + assert_never(mode) + + if __name__ == "__main__": main(sys.argv[1:]) diff --git a/helm/litellm-helm/.helmignore b/helm/litellm-helm/.helmignore deleted file mode 100644 index 0e8a0eb36f4..00000000000 --- a/helm/litellm-helm/.helmignore +++ /dev/null @@ -1,23 +0,0 @@ -# Patterns to ignore when building packages. -# This supports shell glob matching, relative path matching, and -# negation (prefixed with !). Only one pattern per line. -.DS_Store -# Common VCS dirs -.git/ -.gitignore -.bzr/ -.bzrignore -.hg/ -.hgignore -.svn/ -# Common backup files -*.swp -*.bak -*.tmp -*.orig -*~ -# Various IDEs -.project -.idea/ -*.tmproj -.vscode/ diff --git a/helm/litellm-helm/Chart.lock b/helm/litellm-helm/Chart.lock deleted file mode 100644 index 1e63f4f9eca..00000000000 --- a/helm/litellm-helm/Chart.lock +++ /dev/null @@ -1,9 +0,0 @@ -dependencies: -- name: postgresql - repository: oci://registry-1.docker.io/bitnamicharts - version: 14.3.1 -- name: redis - repository: oci://registry-1.docker.io/bitnamicharts - version: 18.19.1 -digest: sha256:38962e231f6596b93f82a8412bbe4cf5de696caecf5775dfbbd163383eb1c009 -generated: '2026-10-30T02:13:11.505726000Z' diff --git a/helm/litellm-helm/Chart.yaml b/helm/litellm-helm/Chart.yaml deleted file mode 100644 index a3cb388ffc6..00000000000 --- a/helm/litellm-helm/Chart.yaml +++ /dev/null @@ -1,41 +0,0 @@ -apiVersion: v2 - -# We can't call ourselves just "litellm" because then we couldn't publish to the -# same OCI repository as the "litellm" OCI image -name: litellm-helm -description: Call all LLM APIs using the OpenAI format - -# A chart can be either an 'application' or a 'library' chart. -# -# Application charts are a collection of templates that can be packaged into versioned archives -# to be deployed. -# -# Library charts provide useful utilities or functions for the chart developer. They're included as -# a dependency of application charts to inject those utilities and functions into the rendering -# pipeline. Library charts do not define any templates and therefore cannot be deployed. -type: application - -# This is the chart version. This version number should be incremented each time you make changes -# to the chart and its templates, including the app version. -# Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 1.1.3 - -# This is the version number of the application being deployed. This version number should be -# incremented each time you make changes to the application. Versions are not expected to -# follow Semantic Versioning. They should reflect the version the application is using. -# It is recommended to use it with quotes. -appVersion: v1.85.1 - -annotations: - org.opencontainers.image.source: "https://github.com/BerriAI/litellm" - org.opencontainers.image.url: "https://docs.litellm.ai/" - -dependencies: - - name: "postgresql" - version: "14.3.1" - repository: oci://registry-1.docker.io/bitnamicharts - condition: db.deployStandalone - - name: redis - version: "18.19.1" - repository: oci://registry-1.docker.io/bitnamicharts - condition: redis.enabled diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md deleted file mode 100644 index 694cca05b52..00000000000 --- a/helm/litellm-helm/README.md +++ /dev/null @@ -1,265 +0,0 @@ -# Helm Chart for LiteLLM - -> [!IMPORTANT] -> This is community maintained, Please make an issue if you run into a bug -> We recommend using [Docker or Kubernetes for production deployments](https://docs.litellm.ai/docs/proxy/prod) - -## Prerequisites - -- Kubernetes 1.21+ -- Helm 3.8.0+ - -If `db.deployStandalone` is used: - -- PV provisioner support in the underlying infrastructure - -If `db.useStackgresOperator` is used (not yet implemented): - -- The Stackgres Operator must already be installed in the Kubernetes Cluster. This chart will **not** install the operator if it is missing. - -## Parameters - -### LiteLLM Proxy Deployment Settings - -| Name | Description | Value | -| --------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------- | -| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` | -| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A | -| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A | -| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated on first install and reused on upgrades. | N/A | -| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | -| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | -| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` | -| `image.pullPolicy` | LiteLLM Proxy image pull policy | `IfNotPresent` | -| `image.tag` | Overrides the image tag whose default the latest version of LiteLLM at the time this chart was published. | `""` | -| `imagePullSecrets` | Registry credentials for the LiteLLM and initContainer images. | `[]` | -| `serviceAccount.create` | Whether or not to create a Kubernetes Service Account for this deployment. The default is `false` because LiteLLM has no need to access the Kubernetes API. | `false` | -| `service.type` | Kubernetes Service type (e.g. `LoadBalancer`, `ClusterIP`, etc.) | `ClusterIP` | -| `service.port` | TCP port that the Kubernetes Service will listen on. Also the TCP port within the Pod that the proxy will listen on. | `4000` | -| `livenessProbe.*` | Liveness probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` | -| `readinessProbe.*` | Readiness probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` | -| `startupProbe.*` | Startup probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` | -| `resources.*` | CPU/memory requests and limits for the LiteLLM container. Unset by default; production deployments should set 1 CPU and 4Gi of memory per worker. | `{}` | -| `service.loadBalancerClass` | Optional LoadBalancer implementation class (only used when `service.type` is `LoadBalancer`) | `""` | -| `ingress.labels` | Additional labels for the Ingress resource | `{}` | -| `ingress.*` | See [values.yaml](./values.yaml) for example settings | N/A | -| `proxyConfigMap.create` | When `true`, render a ConfigMap from `.Values.proxy_config` and mount it. | `true` | -| `proxyConfigMap.name` | When `create=false`, name of the existing ConfigMap to mount. | `""` | -| `proxyConfigMap.key` | Key in the ConfigMap that contains the proxy config file. | `"config.yaml"` | -| `proxy_config.*` | See [values.yaml](./values.yaml) for default settings. Rendered into the ConfigMap’s `config.yaml` only when `proxyConfigMap.create=true`. See [example_config_yaml](../../../litellm/proxy/example_config_yaml/) for configuration examples. | `N/A` | -| `extraContainers[]` | An array of additional containers to be deployed as sidecars alongside the LiteLLM Proxy. | -| `pdb.enabled` | Enable a PodDisruptionBudget for the LiteLLM proxy Deployment | `false` | -| `pdb.minAvailable` | Minimum number/percentage of pods that must be available during **voluntary** disruptions (choose **one** of minAvailable/maxUnavailable) | `null` | -| `pdb.maxUnavailable` | Maximum number/percentage of pods that can be unavailable during **voluntary** disruptions (choose **one** of minAvailable/maxUnavailable) | `null` | -| `pdb.annotations` | Extra metadata annotations to add to the PDB | `{}` | -| `pdb.labels` | Extra metadata labels to add to the PDB | `{}` | - -| `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` | -| `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` | -| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` | -| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` | -| `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` | - -#### Example `proxy_config` ConfigMap from values (default): - -``` -proxyConfigMap: - create: true - key: "config.yaml" - -proxy_config: - general_settings: - master_key: os.environ/PROXY_MASTER_KEY - model_list: [] -``` - -The chart ships no models, so add them under `proxy_config.model_list` or from the Admin UI with `store_model_in_db: true` - -``` -proxy_config: - model_list: - - model_name: gpt-5.6 - litellm_params: - model: openai/gpt-5.6 - api_key: os.environ/OPENAI_API_KEY -``` - -#### Example using existing `proxyConfigMap` instead of creating it: - -``` -proxyConfigMap: - create: false - name: my-litellm-config - key: config.yaml - -# proxy_config is ignored in this mode -``` - -#### Example `environmentSecrets` Secret - -``` -apiVersion: v1 -kind: Secret -metadata: - name: litellm-envsecrets -data: - AZURE_OPENAI_API_KEY: TXlTZWN1cmVLM3k= -type: Opaque -``` - -#### Enterprise billable-request metering - -Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on: - -``` -kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key -``` - -``` -billingMetrics: - enabled: true -``` - -Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time. - -### Database Settings - -| Name | Description | Value | -| ------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ | -| `db.useExisting` | Use an existing Postgres database. A Kubernetes Secret object must exist that contains credentials for connecting to the database. An example secret object definition is provided below. | `false` | -| `db.endpoint` | If `db.useExisting` is `true`, this is the IP, Hostname or Service Name of the Postgres server to connect to. | `localhost` | -| `db.database` | If `db.useExisting` is `true`, the name of the existing database to connect to. | `litellm` | -| `db.url` | If `db.useExisting` is `true`, the connection url of the existing database to connect to can be overwritten with this value. | `postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)` | -| `db.secret.name` | If `db.useExisting` is `true`, the name of the Kubernetes Secret that contains credentials. | `postgres` | -| `db.secret.usernameKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the username for authenticating with the Postgres instance. | `username` | -| `db.secret.passwordKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the password associates with the above user. | `password` | -| `db.username` | If `db.useExisting` is `true`, the database user as a plain value. When set, `db.secret.usernameKey` is not read. | `""` | -| `db.port` | If `db.useExisting` is `true` with token auth, the Postgres port used to build the connection URL. Password auth takes the port from `db.url`. | `5432` | -| `db.schema` | With token auth, the Postgres schema holding the LiteLLM tables, rendered as `DATABASE_SCHEMA` when the tables are not in `public`. Ignored with password auth, where the schema stays in `db.url` as `?schema=`. | `""` | -| `db.useIAMAuth` | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with an AWS RDS IAM token instead of a password. See [Token auth](#token-auth-to-an-existing-postgres-rds-iam-or-microsoft-entra-id). | `false` | -| `db.useAzureEntraAuth` | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with a Microsoft Entra ID token instead of a password. Mutually exclusive with `db.useIAMAuth`. | `false` | -| `db.useStackgresOperator` | Not yet implemented. | `false` | -| `db.deployStandalone` | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups. | `true` | -| `postgresql.*` | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml) | -| `postgresql.auth.*` | If `db.deployStandalone` is `true`, care should be taken to ensure the default `password` and `postgres-password` values are **NOT** used. | `NoTaGrEaTpAsSwOrD` | -| `postgresql.image.*` | If `db.deployStandalone` is `true`, the image for the bundled Postgres. Pinned to a `docker.io/bitnamilegacy` build because Bitnami retired the versioned tags under `docker.io/bitnami`. | `bitnamilegacy/postgresql:16.2.0-debian-12-r6` | -| `redis.image.*` | If `redis.enabled` is `true`, the image for the bundled Redis. Pinned to a `docker.io/bitnamilegacy` build for the same reason. | `bitnamilegacy/redis:7.2.4-debian-12-r9` | - -#### Bundled Postgres image - -Bitnami removed the versioned tags from `docker.io/bitnami` and republished the archived builds under `docker.io/bitnamilegacy`, so the image defaults that ship inside the `postgresql` and `redis` subcharts no longer pull. The chart pins both to the `bitnamilegacy` copies of the exact builds those subchart versions were released with, which keeps the on-disk data directory layout unchanged for existing installs. - -Keep `postgresql.image.tag` pinned. `docker.io/bitnami/postgresql` still publishes a floating `latest`, and pointing the bundled Postgres at a different major version starts the server against a data directory it cannot read (`database files are incompatible with server`). There is no in-place way back, so crossing a major version means dumping the database with the old image and restoring it into the new one. The chart refuses to render when the tag is empty or `latest`. - -Those images no longer receive updates. For anything beyond getting started, run Postgres outside the chart and point at it with `db.useExisting`. - -#### Example Postgres `db.useExisting` Secret - -```yaml -apiVersion: v1 -kind: Secret -metadata: - name: postgres -data: - # Password for the "postgres" user - postgres-password: - username: litellm - password: -type: Opaque -``` - -#### Token auth to an existing Postgres (RDS IAM or Microsoft Entra ID) - -When the database accepts short-lived tokens instead of a password (AWS RDS or Aurora with IAM database authentication, or Azure Database for PostgreSQL with Microsoft Entra ID), there is no password to put in a Secret. Turn on `db.useIAMAuth` or `db.useAzureEntraAuth` and the chart stops referencing `db.secret.passwordKey` and `db.url`. It sets `IAM_TOKEN_DB_AUTH` (or `AZURE_POSTGRESQL_AUTH`) and emits `DATABASE_HOST`, `DATABASE_PORT`, `DATABASE_USERNAME` and `DATABASE_NAME` for the proxy, the migrations job and the collector sidecar. Each of them mints its own token at startup and builds `DATABASE_URL` from those parts, and the in-container pgbouncer (`db.connectionPool`) refreshes the token before it expires. - -```yaml -db: - useExisting: true - deployStandalone: false - useIAMAuth: true - endpoint: litellm.cluster-abc123.us-west-1.rds.amazonaws.com - database: litellm - username: litellm_iam # the role granted rds_iam; no db.secret needed -``` - -The pod needs cloud credentials for the token: on EKS use IRSA or Pod Identity on the chart's `serviceAccount` (and `migrationJob.serviceAccountName` when the Helm hook runs before the ServiceAccount exists), on AKS use workload identity. Set `AWS_RDS_REGION` under `envVars` when the RDS hostname is not a standard regional endpoint (the token must be signed for the region the database lives in). TLS options that would otherwise go in `db.url` are set as `DATABASE_SSLMODE` and `DATABASE_SSLROOTCERT` under `envVars`. Leaving `db.username` empty reads the user from `db.secret.usernameKey` as before. - -#### Examples for `environmentSecrets` and `environemntConfigMaps` - -```yaml -# Use config map for not-secret configuration data -apiVersion: v1 -kind: ConfigMap -metadata: - name: litellm-env-configmap -data: - SOME_KEY: someValue - ANOTHER_KEY: anotherValue -``` - -```yaml -# Use secrets for things which are actually secret like API keys, credentials, etc -# Base64 encode the values stored in a Kubernetes Secret: $ pbpaste | base64 | pbcopy -# The --decode flag is convenient: $ pbpaste | base64 --decode - -apiVersion: v1 -kind: Secret -metadata: - name: litellm-env-secret -type: Opaque -data: - SOME_PASSWORD: cDZbUGVXeU5e0ZW # base64 encoded - ANOTHER_PASSWORD: AAZbUGVXeU5e0ZB # base64 encoded -``` - -Source: [GitHub Gist from troyharvey](https://gist.github.com/troyharvey/4506472732157221e04c6b15e3b3f094) - -### Migration Job Settings - -The migration job supports both ArgoCD and Helm hooks to ensure database migrations run at the appropriate time during deployments. - -| Name | Description | Value | -| -------------------------------------- | -------------------------------------------------------------------------------------------------------------------- | ------- | -| `migrationJob.enabled` | Enable or disable the schema migration Job | `true` | -| `migrationJob.backoffLimit` | Backoff limit for Job restarts | `4` | -| `migrationJob.ttlSecondsAfterFinished` | TTL for completed migration jobs | `120` | -| `migrationJob.annotations` | Additional annotations for the migration job pod | `{}` | -| `migrationJob.extraContainers` | Additional containers to run alongside the migration job | `[]` | -| `migrationJob.hooks.argocd.enabled` | Enable ArgoCD hooks for the migration job (uses PreSync hook with BeforeHookCreation delete policy) | `true` | -| `migrationJob.hooks.helm.enabled` | Enable Helm hooks for the migration job (uses pre-install,pre-upgrade hooks with before-hook-creation delete policy) | `false` | -| `migrationJob.hooks.helm.weight` | Helm hook execution order (lower weights executed first). Optional - defaults to "1" if not specified. | N/A | - -## Accessing the Admin UI - -When browsing to the URL published per the settings in `ingress.*`, you will -be prompted for **Admin Configuration**. The **Proxy Endpoint** is the internal -(from the `litellm` pod's perspective) URL published by the `-litellm` -Kubernetes Service. If the deployment uses the default settings for this -service, the **Proxy Endpoint** should be set to `http://-litellm:4000`. - -The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey` -was not provided to the helm command line, the `masterkey` is a randomly -generated string in the `sk-...` format stored in the `-litellm-masterkey` Kubernetes Secret. -The key is generated once on the first install; later `helm upgrade` runs reuse the -value already in that Secret, so upgrading never rotates the master key. - -```bash -kubectl -n litellm get secret -litellm-masterkey -o jsonpath="{.data.masterkey}" -``` - -## Admin UI Limitations - -At the time of writing, the Admin UI is unable to add models. This is because -it would need to update the `config.yaml` file which is a exposed ConfigMap, and -therefore, read-only. This is a limitation of this helm chart, not the Admin UI -itself. - -## Connect Lens - -This chart does not deploy Lens. Deploy Lens from [its own chart](https://github.com/BerriAI/lens/tree/main/helm/lens), then set `lens.enabled: true`, `lens.url`, `lens.publicUrl`, `lens.gatewaySecret.name` and `lens.serviceTokenSecret.name`. The secret keys default to `gateway-secret` and `service-token`. Set `lens.service.name` when that Lens runs in this release's namespace and the chart ingress should route `/lens-ingest` to it. The old `lensWorker` values are no longer read, and enabling them fails the render - -## Shutdown accounting - -`GRACEFUL_SHUTDOWN_TIMEOUT` defaults to 30 seconds. Pending success callbacks share the time remaining since SIGTERM or the first `/health/drain` call. Callback flushing reserves the existing 10-second teardown headroom for cancellation, tracked accounting tasks and persistence. Exhausting this budget can discard callbacks, which is logged - -Set `terminationGracePeriodSeconds` to at least `GRACEFUL_SHUTDOWN_TIMEOUT + 10`, allowing additional time for preStop hooks and database writes. This is not a hard shutdown bound or a guarantee of persistence. Database retries retain their own limits, and cancellation-resistant callbacks can remain alive. The existing preStop hook can also be followed by Uvicorn's separate connection drain, as described in `values.yaml` diff --git a/helm/litellm-helm/charts/postgresql-14.3.1.tgz b/helm/litellm-helm/charts/postgresql-14.3.1.tgz deleted file mode 100644 index e8e2fac0fda..00000000000 Binary files a/helm/litellm-helm/charts/postgresql-14.3.1.tgz and /dev/null differ diff --git a/helm/litellm-helm/charts/redis-18.19.1.tgz b/helm/litellm-helm/charts/redis-18.19.1.tgz deleted file mode 100644 index 4a55a980088..00000000000 Binary files a/helm/litellm-helm/charts/redis-18.19.1.tgz and /dev/null differ diff --git a/helm/litellm-helm/ci/test-values.yaml b/helm/litellm-helm/ci/test-values.yaml deleted file mode 100644 index 33a4df942ee..00000000000 --- a/helm/litellm-helm/ci/test-values.yaml +++ /dev/null @@ -1,15 +0,0 @@ -fullnameOverride: "" -# Disable database deployment and configuration -db: - deployStandalone: false - useExisting: false - -# Test environment variables -envVars: - DD_ENV: "dev_helm" - DD_SERVICE: "litellm" - USE_DDTRACE: "true" - -# Disable migration job since we're not using a database -migrationJob: - enabled: false \ No newline at end of file diff --git a/helm/litellm-helm/templates/NOTES.txt b/helm/litellm-helm/templates/NOTES.txt deleted file mode 100644 index 017bbfa78bd..00000000000 --- a/helm/litellm-helm/templates/NOTES.txt +++ /dev/null @@ -1,23 +0,0 @@ -1. Get the application URL by running these commands: -{{- if .Values.ingress.enabled }} -{{- range $host := .Values.ingress.hosts }} - {{- range .paths }} - http{{ if $.Values.ingress.tls }}s{{ end }}://{{ $host.host }}{{ .path }} - {{- end }} -{{- end }} -{{- else if contains "NodePort" .Values.service.type }} - export NODE_PORT=$(kubectl get --namespace {{ .Release.Namespace }} -o jsonpath="{.spec.ports[0].nodePort}" services {{ include "litellm.fullname" . }}) - export NODE_IP=$(kubectl get nodes --namespace {{ .Release.Namespace }} -o jsonpath="{.items[0].status.addresses[0].address}") - echo http://$NODE_IP:$NODE_PORT -{{- else if contains "LoadBalancer" .Values.service.type }} - NOTE: It may take a few minutes for the LoadBalancer IP to be available. - You can watch the status of by running 'kubectl get --namespace {{ .Release.Namespace }} svc -w {{ include "litellm.fullname" . }}' - export SERVICE_IP=$(kubectl get svc --namespace {{ .Release.Namespace }} {{ include "litellm.fullname" . }} --template "{{"{{ range (index .status.loadBalancer.ingress 0) }}{{.}}{{ end }}"}}") - echo http://$SERVICE_IP:{{ .Values.service.port }} -{{- else if contains "ClusterIP" .Values.service.type }} - export POD_NAME=$(kubectl get pods --namespace {{ .Release.Namespace }} -l "app.kubernetes.io/name={{ include "litellm.name" . }},app.kubernetes.io/instance={{ .Release.Name }}" -o jsonpath="{.items[0].metadata.name}") - export CONTAINER_PORT=$(kubectl get pod --namespace {{ .Release.Namespace }} $POD_NAME -o jsonpath="{.spec.containers[0].ports[0].containerPort}") - echo "Visit http://127.0.0.1:8080 to use your application" - kubectl --namespace {{ .Release.Namespace }} port-forward $POD_NAME 8080:$CONTAINER_PORT -{{- end }} -PDB: {{ if .Values.pdb.enabled }}enabled{{ else }}disabled{{ end }}. Configure via .Values.pdb.* \ No newline at end of file diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl deleted file mode 100644 index 929fc415163..00000000000 --- a/helm/litellm-helm/templates/_helpers.tpl +++ /dev/null @@ -1,425 +0,0 @@ -{{/* -Expand the name of the chart. -*/}} -{{- define "litellm.name" -}} -{{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }} -{{- end }} - -{{/* -Create a default fully qualified app name. -We truncate at 63 chars because some Kubernetes name fields are limited to this (by the DNS naming spec). -If release name contains chart name it will be used as a full name. -*/}} -{{- define "litellm.fullname" -}} -{{- if .Values.fullnameOverride }} -{{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }} -{{- else }} -{{- $name := default .Chart.Name .Values.nameOverride }} -{{- if contains $name .Release.Name }} -{{- .Release.Name | trunc 63 | trimSuffix "-" }} -{{- else }} -{{- printf "%s-%s" .Release.Name $name | trunc 63 | trimSuffix "-" }} -{{- end }} -{{- end }} -{{- end }} - -{{/* -Create chart name and version as used by the chart label. -*/}} -{{- define "litellm.chart" -}} -{{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }} -{{- end }} - -{{/* -Common labels -*/}} -{{- define "litellm.labels" -}} -helm.sh/chart: {{ include "litellm.chart" . }} -{{ include "litellm.selectorLabels" . }} -{{- if .Chart.AppVersion }} -app.kubernetes.io/version: {{ .Chart.AppVersion | quote }} -{{- end }} -app.kubernetes.io/managed-by: {{ .Release.Service }} -{{- end }} - -{{/* -Selector labels -*/}} -{{- define "litellm.selectorLabels" -}} -app.kubernetes.io/name: {{ include "litellm.name" . }} -app.kubernetes.io/instance: {{ .Release.Name }} -{{- end }} - -{{/* -Enterprise billable-request metering. The client certificate identifies the -deployment to LiteLLM's collector, so it is mounted read-only from an existing -Secret rather than passed through the environment. -*/}} -{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}} -{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}} - -{{- define "litellm.billingMetricsEnv" -}} -- name: LITELLM_BILLING_METRICS_ENDPOINT - value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }} -{{- if .Values.billingMetrics.caSecretName }} -- name: LITELLM_BILLING_METRICS_CA_CERT - value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }} -{{- end }} -{{- with .Values.billingMetrics.exportIntervalMs }} -- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS - value: {{ . | quote }} -{{- end }} -{{- end -}} - -{{- define "litellm.billingMetricsVolumes" -}} -- name: billing-metrics-mtls - secret: - secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }} -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - secret: - secretName: {{ .Values.billingMetrics.caSecretName }} -{{- end }} -{{- end -}} - -{{- define "litellm.billingMetricsVolumeMounts" -}} -- name: billing-metrics-mtls - mountPath: {{ include "litellm.billingMetrics.certDir" . }} - readOnly: true -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - mountPath: {{ include "litellm.billingMetrics.caDir" . }} - readOnly: true -{{- end }} -{{- end -}} - -{{/* -Create the name of the service account to use -*/}} -{{- define "litellm.serviceAccountName" -}} -{{- if .Values.serviceAccount.create }} -{{- default (include "litellm.fullname" .) .Values.serviceAccount.name }} -{{- else }} -{{- default "default" .Values.serviceAccount.name }} -{{- end }} -{{- end }} - -{{/* -Create the service account name used by migration jobs. -When Helm hooks are enabled, pre-install/pre-upgrade hooks run before normal resources. -If this chart is creating the ServiceAccount, it is not yet available for the hook job, -so fall back to "default" (or an explicit override) to avoid a cyclic dependency. -*/}} -{{- define "litellm.migrationServiceAccountName" -}} -{{- if and .Values.migrationJob.hooks.helm.enabled .Values.serviceAccount.create }} -{{- default "default" .Values.migrationJob.serviceAccountName }} -{{- else }} -{{- include "litellm.serviceAccountName" . }} -{{- end }} -{{- end }} - -{{/* -Get redis service name. -The bundled Redis subchart only serves sentinel in "replication" architecture -(it rejects standalone + sentinel outright), and in that mode the sentinel -Service is named "-redis", not "-redis-master". -*/}} -{{- define "litellm.redis.serviceName" -}} -{{- if .Values.redis.sentinel.enabled -}} -{{- printf "%s-%s" .Release.Name (default "redis" .Values.redis.nameOverride | trunc 63 | trimSuffix "-") -}} -{{- else -}} -{{- printf "%s-%s-master" .Release.Name (default "redis" .Values.redis.nameOverride | trunc 63 | trimSuffix "-") -}} -{{- end -}} -{{- end -}} - -{{/* -Get redis service port -*/}} -{{- define "litellm.redis.port" -}} -{{- if .Values.redis.sentinel.enabled -}} -{{ .Values.redis.sentinel.service.ports.sentinel }} -{{- else -}} -{{ .Values.redis.master.service.ports.redis }} -{{- end -}} -{{- end -}} - -{{/* -Reject an unpinned image tag for the bundled PostgreSQL. -A floating tag lets a chart upgrade start a newer PostgreSQL major against the -existing PersistentVolumeClaim. The server then refuses to start on a data -directory written by another major version, and the only way back is a dump -taken before the change, which by that point no longer exists. -*/}} -{{- define "litellm.validateBundledPostgresImageTag" -}} -{{- $tag := .Values.postgresql.image.tag | default "" | toString -}} -{{- $digest := .Values.postgresql.image.digest | default "" | toString -}} -{{- if and (eq $digest "") (or (eq $tag "") (eq $tag "latest")) -}} -{{- fail (printf "postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got %q). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore." $tag) -}} -{{- end -}} -{{- end -}} - -{{/* -Whether the existing Postgres is reached with a short-lived token (AWS RDS IAM -or Microsoft Entra ID) instead of a password. -*/}} -{{- define "litellm.db.tokenAuth" -}} -{{- if or .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}true{{- end -}} -{{- end -}} - -{{- define "litellm.validateDbTokenAuth" -}} -{{- if and .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}} -{{- fail "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" -}} -{{- end -}} -{{- if and (include "litellm.db.tokenAuth" .) (or (not .Values.db.useExisting) .Values.db.deployStandalone) -}} -{{- fail "db.useIAMAuth and db.useAzureEntraAuth need db.useExisting: true and db.deployStandalone: false (token auth targets an existing RDS or Azure Database for PostgreSQL; the bundled postgresql subchart only supports password auth)" -}} -{{- end -}} -{{- end -}} - -{{/* -Database env for an existing Postgres, shared by the proxy container, the -collector sidecar and the migrations Job. Under token auth every entrypoint -mints its own token and assembles DATABASE_URL itself, so db.url is left out. -*/}} -{{- define "litellm.existingDbEnv" -}} -{{- $tokenAuth := include "litellm.db.tokenAuth" . -}} -- name: DATABASE_USERNAME - {{- if .Values.db.username }} - value: {{ .Values.db.username | quote }} - {{- else }} - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.usernameKey }} - {{- end }} -{{- if .Values.db.useIAMAuth }} -- name: IAM_TOKEN_DB_AUTH - value: "true" -{{- else if .Values.db.useAzureEntraAuth }} -- name: AZURE_POSTGRESQL_AUTH - value: "true" -{{- else }} -- name: DATABASE_PASSWORD - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.passwordKey }} -{{- end }} -- name: DATABASE_HOST - {{- if .Values.db.secret.endpointKey }} - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.endpointKey }} - {{- else }} - value: {{ .Values.db.endpoint }} - {{- end }} -{{- if $tokenAuth }} -- name: DATABASE_PORT - value: {{ .Values.db.port | quote }} -{{- end }} -- name: DATABASE_NAME - value: {{ .Values.db.database }} -{{- if and $tokenAuth .Values.db.schema }} -- name: DATABASE_SCHEMA - value: {{ .Values.db.schema | quote }} -{{- end }} -{{- if not $tokenAuth }} -- name: DATABASE_URL - value: {{ .Values.db.url | quote }} -{{- end }} -{{- end -}} - -{{/* -Environment shared by the proxy container and the opt-in collector sidecar: -database, pgbouncer, master key, redis, user envVars. Both containers must see -the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer -and the same spend transaction buffer. -*/}} -{{- define "litellm.proxyEnv" -}} -{{- include "litellm.validateDbTokenAuth" . -}} -- name: HOST - value: "{{ .Values.listen | default "0.0.0.0" }}" -- name: PORT - value: {{ .Values.service.port | quote}} -{{- if .Values.db.deployStandalone }} -- name: DATABASE_USERNAME - valueFrom: - secretKeyRef: - name: {{ include "litellm.fullname" . }}-dbcredentials - key: username -- name: DATABASE_PASSWORD - valueFrom: - secretKeyRef: - name: {{ include "litellm.fullname" . }}-dbcredentials - key: password -- name: DATABASE_HOST - value: {{ .Release.Name }}-postgresql -- name: DATABASE_NAME - value: litellm -{{- else if .Values.db.useExisting }} -{{ include "litellm.existingDbEnv" . }} -{{- end }} -{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }} -- name: DATABASE_READER_HOST - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.readReplicaEndpointKey }} -{{- end }} -{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }} -- name: DATABASE_URL_READ_REPLICA - valueFrom: - secretKeyRef: - name: {{ .Values.db.secret.name }} - key: {{ .Values.db.secret.readReplicaUrlKey }} -{{- else if .Values.db.readReplicaUrl }} -- name: DATABASE_URL_READ_REPLICA - value: {{ .Values.db.readReplicaUrl | quote }} -{{- end }} -{{- if .Values.db.connectionPool.enabled }} -- name: LITELLM_PGBOUNCER_ENABLED - value: "true" -- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS - value: {{ .Values.db.connectionPool.maxDbConnections | quote }} -- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN - value: {{ .Values.db.connectionPool.maxClientConn | quote }} -{{- end }} -- name: PROXY_MASTER_KEY - valueFrom: - secretKeyRef: - name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }} - key: {{ .Values.masterkeySecretKey | default "masterkey" }} -{{- if .Values.redis.enabled }} -- name: REDIS_HOST - value: {{ include "litellm.redis.serviceName" . }} -- name: REDIS_PORT - value: {{ include "litellm.redis.port" . | quote }} -- name: REDIS_PASSWORD - valueFrom: - secretKeyRef: - name: {{ include "redis.secretName" .Subcharts.redis }} - key: {{include "redis.secretPasswordKey" .Subcharts.redis }} -{{- end }} -{{- /* - Inject LITELLM_LOG only when envVars does not already define it. -*/}} -{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }} -- name: LITELLM_LOG - value: {{ .Values.logLevel | quote }} -{{- end }} -{{- if .Values.envVars }} -{{- range $key, $val := .Values.envVars }} -- name: {{ $key }} - value: {{ $val | quote }} -{{- end }} -{{- end }} -{{- with .Values.extraEnvVars }} -{{ toYaml . }} -{{- end }} -{{- if .Values.migrationJob.enabled }} -# Schema updates are owned by the dedicated migrations Job; skip -# the proxy's startup `prisma db push` so N replicas don't race -# one DB on every rollout. Placed last (after envVars and -# extraEnvVars) so this override can't be silently shadowed by a -# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env -# semantics — same pattern the migrations Job uses. -- name: DISABLE_SCHEMA_UPDATE - value: "true" -{{- end }} -{{- end -}} - -{{/* -Proxy-only metering and metrics env. The collector sidecar serves no HTTP -traffic, so it gets neither. -*/}} -{{- define "litellm.proxyMetricsEnv" -}} -{{- if .Values.billingMetrics.enabled }} -{{ include "litellm.billingMetricsEnv" . }} -{{- end }} -{{- if .Values.metricsServer.enabled }} -{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }} -{{- fail "metricsServer.port must differ from service.port" }} -{{- end }} -- name: PROMETHEUS_METRICS_PORT - value: {{ .Values.metricsServer.port | quote }} -{{- end }} -{{- end -}} - -{{/* -Directory of the collector's unix socket, shared between the two containers -through an emptyDir. Empty when the sidecar is off or uses 127.0.0.1 TCP. -*/}} -{{- define "litellm.collector.socketDir" -}} -{{- if and .Values.collector.enabled (hasPrefix "unix://" .Values.collector.address) -}} -{{- dir (trimPrefix "unix://" .Values.collector.address) -}} -{{- end -}} -{{- end -}} - -{{- define "litellm.collectorEnv" -}} -- name: LITELLM_COLLECTOR_ENABLED - value: "true" -- name: LITELLM_COLLECTOR_ADDRESS - value: {{ .Values.collector.address | quote }} -- name: LITELLM_COLLECTOR_BUFFER_SIZE - value: {{ .Values.collector.bufferSize | quote }} -- name: LITELLM_COLLECTOR_ON_UNAVAILABLE - value: {{ .Values.collector.onUnavailable | quote }} -- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS - value: {{ .Values.collector.drainTimeoutSeconds | quote }} -{{- end -}} - -{{- define "litellm.gateway.collectorSocketDir" -}} -{{- if and .Values.gateway.collector.enabled (hasPrefix "unix://" .Values.gateway.collector.address) -}} -{{- dir (trimPrefix "unix://" .Values.gateway.collector.address) -}} -{{- end -}} -{{- end -}} - -{{/* -LITELLM_COLLECTOR_* env shared by the producer (gateway container) and the -consumer (collector container), so both agree on the transport and the -shutdown drain window. -*/}} -{{- define "litellm.gateway.collectorEnv" -}} -{{- with .Values.gateway.collector }} -- name: LITELLM_COLLECTOR_ENABLED - value: "true" -- name: LITELLM_COLLECTOR_ADDRESS - value: {{ .address | quote }} -- name: LITELLM_COLLECTOR_BUFFER_SIZE - value: {{ .bufferSize | quote }} -- name: LITELLM_COLLECTOR_ON_UNAVAILABLE - value: {{ .onUnavailable | quote }} -- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS - value: {{ .drainTimeoutSeconds | quote }} -{{- end }} -{{- end -}} - -{{- define "litellm.lensConnectionEnv" -}} -{{- with .Values.lensWorker }} -{{- if or .enabled (and .mode (ne .mode "disabled")) }} -{{- fail "lensWorker was removed; deploy Lens from its own chart and set lens.* instead, see the chart README" }} -{{- end }} -{{- end }} -{{- with .Values.lens }} -{{- if .enabled }} -- name: LITELLM_LENS_URL - value: {{ required "lens.url is required when lens.enabled is true" .url | quote }} -- name: LITELLM_LENS_PUBLIC_URL - value: {{ required "lens.publicUrl is required when lens.enabled is true" .publicUrl | quote }} -- name: LENS_GATEWAY_SECRET - valueFrom: - secretKeyRef: - name: {{ required "lens.gatewaySecret.name is required when lens.enabled is true" .gatewaySecret.name | quote }} - key: {{ .gatewaySecret.key | quote }} -- name: LITELLM_LENS_SERVICE_TOKEN - valueFrom: - secretKeyRef: - name: {{ required "lens.serviceTokenSecret.name is required when lens.enabled is true" .serviceTokenSecret.name | quote }} - key: {{ .serviceTokenSecret.key | quote }} -{{- end }} -{{- end }} -{{- end -}} diff --git a/helm/litellm-helm/templates/configmap-litellm.yaml b/helm/litellm-helm/templates/configmap-litellm.yaml deleted file mode 100644 index 03e4f620206..00000000000 --- a/helm/litellm-helm/templates/configmap-litellm.yaml +++ /dev/null @@ -1,22 +0,0 @@ -{{- if .Values.proxyConfigMap.create }} -{{- $config := deepCopy .Values.proxy_config }} -{{- if and .Values.redis.enabled (dig "coordination" "enabled" true .Values.redis) }} -{{- $generalSettings := (get $config "general_settings") | default dict }} -{{- if not (hasKey $generalSettings "coordination_redis") }} -{{- $coordinationRedis := dict "host" "os.environ/REDIS_HOST" "port" "os.environ/REDIS_PORT" "password" "os.environ/REDIS_PASSWORD" }} -{{- if .Values.redis.sentinel.enabled }} -{{- $sentinelNode := list (include "litellm.redis.serviceName" .) (include "litellm.redis.port" . | int) }} -{{- $coordinationRedis = dict "sentinel_nodes" (list $sentinelNode) "service_name" (default "mymaster" .Values.redis.sentinel.masterSet) "password" "os.environ/REDIS_PASSWORD" }} -{{- end }} -{{- $_ := set $generalSettings "coordination_redis" $coordinationRedis }} -{{- $_ := set $config "general_settings" $generalSettings }} -{{- end }} -{{- end }} -apiVersion: v1 -kind: ConfigMap -metadata: - name: {{ include "litellm.fullname" . }}-config -data: - config.yaml: | -{{ $config | toYaml | indent 6 }} -{{- end }} diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml deleted file mode 100644 index 37e7eaec9e9..00000000000 --- a/helm/litellm-helm/templates/deployment.yaml +++ /dev/null @@ -1,264 +0,0 @@ -apiVersion: apps/v1 -kind: Deployment -metadata: - annotations: - {{- toYaml .Values.deploymentAnnotations | nindent 4 }} - name: {{ include "litellm.fullname" . }} - labels: - {{- include "litellm.labels" . | nindent 4 }} - {{- if .Values.deploymentLabels }} - {{- toYaml .Values.deploymentLabels | nindent 4 }} - {{- end }} -spec: - {{- if and (not .Values.keda.enabled) (not .Values.autoscaling.enabled) }} - replicas: {{ .Values.replicaCount }} - {{- end }} - {{- with .Values.strategy }} - strategy: - {{- toYaml . | nindent 4 }} - {{- end }} - selector: - matchLabels: - {{- include "litellm.selectorLabels" . | nindent 6 }} - {{- if .Values.deploymentMinReadySeconds }} - minReadySeconds: {{ .Values.deploymentMinReadySeconds }} - {{- end }} - template: - metadata: - annotations: - {{- if .Values.proxyConfigMap.create }} - checksum/config: {{ include (print $.Template.BasePath "/configmap-litellm.yaml") . | sha256sum }} - {{- end }} - {{- with .Values.podAnnotations }} - {{- tpl (toYaml .) $ | nindent 8 }} - {{- end }} - labels: - {{- include "litellm.labels" . | nindent 8 }} - {{- with .Values.podLabels }} - {{- toYaml . | nindent 8 }} - {{- end }} - spec: - {{- with .Values.imagePullSecrets }} - imagePullSecrets: - {{- toYaml . | nindent 8 }} - {{- end }} - serviceAccountName: {{ include "litellm.serviceAccountName" . }} - securityContext: - {{- toYaml .Values.podSecurityContext | nindent 8 }} - {{- with .Values.extraInitContainers }} - initContainers: - {{- tpl (toYaml .) $ | nindent 8 }} - {{- end }} - containers: - - name: {{ include "litellm.name" . }} - securityContext: - {{- toYaml .Values.securityContext | nindent 12 }} - image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.image.pullPolicy }} - env: - {{- include "litellm.lensConnectionEnv" . | nindent 12 }} - {{- include "litellm.proxyEnv" . | nindent 12 }} - {{- if .Values.liteadmin.enabled }} - - name: LITELLM_ADMIN_AGENT_URL - value: {{ printf "http://%s-liteadmin:10000" (include "litellm.fullname" . | trunc 53 | trimSuffix "-") | quote }} - - name: ADMIN_AGENT_SERVICE_TOKEN - valueFrom: - secretKeyRef: - name: {{ required "liteadmin.existingSecret is required" .Values.liteadmin.existingSecret }} - key: ADMIN_AGENT_SERVICE_TOKEN - {{- if not (hasKey (default dict .Values.envVars) "PROXY_BASE_URL") }} - - name: PROXY_BASE_URL - value: {{ required "liteadmin.gatewayUrl is required" .Values.liteadmin.gatewayUrl | quote }} - {{- end }} - {{- end }} - {{- include "litellm.proxyMetricsEnv" . | nindent 12 }} - {{- if .Values.collector.enabled }} - {{- include "litellm.collectorEnv" . | nindent 12 }} - {{- end }} - envFrom: - {{- range .Values.environmentSecrets }} - - secretRef: - name: {{ . }} - {{- end }} - {{- range .Values.environmentConfigMaps }} - - configMapRef: - name: {{ . }} - {{- end }} - {{- if .Values.command }} - command: {{ toYaml .Values.command | nindent 12 }} - {{- end }} - {{- if .Values.args }} - args: {{ toYaml .Values.args | nindent 12 }} - {{- else }} - args: - - --config - - /etc/litellm/config.yaml - {{ if .Values.numWorkers }} - - --num_workers - - {{ .Values.numWorkers | quote }} - {{- end }} - {{- end }} - ports: - - name: http - containerPort: {{ .Values.service.port }} - protocol: TCP - {{- if .Values.metricsServer.enabled }} - - name: metrics - containerPort: {{ .Values.metricsServer.port }} - protocol: TCP - {{- end }} - livenessProbe: - httpGet: - path: {{ .Values.livenessProbe.path | quote }} - port: "http" - initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds }} - periodSeconds: {{ .Values.livenessProbe.periodSeconds }} - timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds }} - successThreshold: {{ .Values.livenessProbe.successThreshold }} - failureThreshold: {{ .Values.livenessProbe.failureThreshold }} - readinessProbe: - httpGet: - path: {{ .Values.readinessProbe.path | quote }} - port: "http" - initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds }} - periodSeconds: {{ .Values.readinessProbe.periodSeconds }} - timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds }} - successThreshold: {{ .Values.readinessProbe.successThreshold }} - failureThreshold: {{ .Values.readinessProbe.failureThreshold }} - startupProbe: - httpGet: - path: {{ .Values.startupProbe.path | quote }} - port: "http" - initialDelaySeconds: {{ .Values.startupProbe.initialDelaySeconds }} - periodSeconds: {{ .Values.startupProbe.periodSeconds }} - timeoutSeconds: {{ .Values.startupProbe.timeoutSeconds }} - successThreshold: {{ .Values.startupProbe.successThreshold }} - failureThreshold: {{ .Values.startupProbe.failureThreshold }} - resources: - {{- toYaml .Values.resources | nindent 12 }} - volumeMounts: - - name: litellm-config - mountPath: /etc/litellm/config.yaml - subPath: config.yaml - {{ if .Values.securityContext.readOnlyRootFilesystem }} - - name: tmp - mountPath: /tmp - - name: cache - mountPath: /.cache - - name: npm - mountPath: /.npm - {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} - {{- if include "litellm.collector.socketDir" . }} - - name: collector-socket - mountPath: {{ include "litellm.collector.socketDir" . }} - {{- end }} - {{- with .Values.volumeMounts }} - {{- toYaml . | nindent 12 }} - {{- end }} - {{- with .Values.lifecycle }} - lifecycle: - {{- toYaml . | nindent 12 }} - {{- end }} - {{- if .Values.collector.enabled }} - - name: {{ include "litellm.name" . }}-collector - securityContext: - {{- toYaml .Values.securityContext | nindent 12 }} - image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.image.pullPolicy }} - command: {{ toYaml .Values.collector.command | nindent 12 }} - env: - {{- include "litellm.proxyEnv" . | nindent 12 }} - {{- include "litellm.collectorEnv" . | nindent 12 }} - - name: LITELLM_JOB_ROLE - value: collector - {{- if not (hasKey (default dict .Values.envVars) "CONFIG_FILE_PATH") }} - - name: CONFIG_FILE_PATH - value: /etc/litellm/config.yaml - {{- end }} - envFrom: - {{- range .Values.environmentSecrets }} - - secretRef: - name: {{ . }} - {{- end }} - {{- range .Values.environmentConfigMaps }} - - configMapRef: - name: {{ . }} - {{- end }} - resources: - {{- toYaml .Values.collector.resources | nindent 12 }} - volumeMounts: - - name: litellm-config - mountPath: /etc/litellm/config.yaml - subPath: config.yaml - {{- if include "litellm.collector.socketDir" . }} - - name: collector-socket - mountPath: {{ include "litellm.collector.socketDir" . }} - {{- end }} - {{ if .Values.securityContext.readOnlyRootFilesystem }} - - name: tmp - mountPath: /tmp - - name: cache - mountPath: /.cache - - name: npm - mountPath: /.npm - {{- end }} - {{- with .Values.volumeMounts }} - {{- toYaml . | nindent 12 }} - {{- end }} - {{- end }} - {{- with .Values.extraContainers }} - {{- tpl (toYaml .) $ | nindent 8 }} - {{- end }} - volumes: - {{ if .Values.securityContext.readOnlyRootFilesystem }} - - name: tmp - emptyDir: - sizeLimit: 500Mi - - name: cache - emptyDir: - sizeLimit: 500Mi - - name: npm - emptyDir: - sizeLimit: 500Mi - {{- end }} - - name: litellm-config - configMap: - {{- if .Values.proxyConfigMap.create }} - name: {{ include "litellm.fullname" . }}-config - {{- else }} - name: {{ .Values.proxyConfigMap.name }} - {{- end }} - items: - - key: {{ .Values.proxyConfigMap.key | default "config.yaml" }} - path: "config.yaml" - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} - {{- if include "litellm.collector.socketDir" . }} - - name: collector-socket - emptyDir: - sizeLimit: 1Mi - {{- end }} - {{- with .Values.volumes }} - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.nodeSelector }} - nodeSelector: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.affinity }} - affinity: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.tolerations }} - tolerations: - {{- toYaml . | nindent 8 }} - {{- end }} - terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds | default 90 }} - {{- if .Values.topologySpreadConstraints }} - topologySpreadConstraints: - {{- toYaml .Values.topologySpreadConstraints | nindent 8 }} - {{- end }} diff --git a/helm/litellm-helm/templates/extra-resources.yaml b/helm/litellm-helm/templates/extra-resources.yaml deleted file mode 100644 index 33190d96fc0..00000000000 --- a/helm/litellm-helm/templates/extra-resources.yaml +++ /dev/null @@ -1,6 +0,0 @@ -{{- if .Values.extraResources }} -{{- range .Values.extraResources }} ---- -{{ toYaml . | nindent 0 }} -{{- end }} -{{- end }} \ No newline at end of file diff --git a/helm/litellm-helm/templates/hpa.yaml b/helm/litellm-helm/templates/hpa.yaml deleted file mode 100644 index a651f916d21..00000000000 --- a/helm/litellm-helm/templates/hpa.yaml +++ /dev/null @@ -1,64 +0,0 @@ -{{- if .Values.autoscaling.enabled }} -apiVersion: autoscaling/v2 -kind: HorizontalPodAutoscaler -metadata: - name: {{ include "litellm.fullname" . }} - labels: - {{- include "litellm.labels" . | nindent 4 }} -spec: - scaleTargetRef: - apiVersion: apps/v1 - kind: Deployment - name: {{ include "litellm.fullname" . }} - minReplicas: {{ .Values.autoscaling.minReplicas }} - maxReplicas: {{ .Values.autoscaling.maxReplicas }} - {{- if .Values.autoscaling.behavior }} - behavior: - {{- toYaml .Values.autoscaling.behavior | nindent 4 }} - {{- end }} - metrics: - {{- if .Values.autoscaling.targetCPUUtilizationPercentage }} - {{- if and .Values.collector.enabled .Values.collector.scaleOnProxyContainerCpu }} - - type: ContainerResource - containerResource: - name: cpu - container: {{ include "litellm.name" . }} - target: - type: Utilization - averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }} - {{- else }} - - type: Resource - resource: - name: cpu - target: - type: Utilization - averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }} - {{- end }} - {{- end }} - {{- if .Values.autoscaling.targetMemoryUtilizationPercentage }} - - type: Resource - resource: - name: memory - target: - type: Utilization - averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }} - {{- end }} - {{- with .Values.autoscaling.targetRequestsPerSecond }} - - type: Pods - pods: - metric: - name: litellm_requests_per_second - target: - type: AverageValue - averageValue: {{ toJson . | trimAll "\"" | quote }} - {{- end }} - {{- with .Values.autoscaling.targetTokensPerSecond }} - - type: Pods - pods: - metric: - name: litellm_tokens_per_second - target: - type: AverageValue - averageValue: {{ toJson . | trimAll "\"" | quote }} - {{- end }} -{{- end }} diff --git a/helm/litellm-helm/templates/ingress.yaml b/helm/litellm-helm/templates/ingress.yaml deleted file mode 100644 index 3d6e3f4af7a..00000000000 --- a/helm/litellm-helm/templates/ingress.yaml +++ /dev/null @@ -1,78 +0,0 @@ -{{- if .Values.ingress.enabled -}} -{{- $fullName := include "litellm.fullname" . -}} -{{- $svcPort := .Values.service.port -}} -{{- if and .Values.ingress.className (not (semverCompare ">=1.18-0" .Capabilities.KubeVersion.GitVersion)) }} - {{- if not (hasKey .Values.ingress.annotations "kubernetes.io/ingress.class") }} - {{- $_ := set .Values.ingress.annotations "kubernetes.io/ingress.class" .Values.ingress.className}} - {{- end }} -{{- end }} -{{- if semverCompare ">=1.19-0" .Capabilities.KubeVersion.GitVersion -}} -apiVersion: networking.k8s.io/v1 -{{- else if semverCompare ">=1.14-0" .Capabilities.KubeVersion.GitVersion -}} -apiVersion: networking.k8s.io/v1beta1 -{{- else -}} -apiVersion: extensions/v1beta1 -{{- end }} -kind: Ingress -metadata: - name: {{ $fullName }} - labels: - {{- include "litellm.labels" . | nindent 4 }} - {{- with .Values.ingress.labels }} - {{- toYaml . | nindent 4 }} - {{- end }} - {{- with .Values.ingress.annotations }} - annotations: - {{- toYaml . | nindent 4 }} - {{- end }} -spec: - {{- if and .Values.ingress.className (semverCompare ">=1.18-0" .Capabilities.KubeVersion.GitVersion) }} - ingressClassName: {{ .Values.ingress.className }} - {{- end }} - {{- if .Values.ingress.tls }} - tls: - {{- range .Values.ingress.tls }} - - hosts: - {{- range .hosts }} - - {{ . | quote }} - {{- end }} - secretName: {{ .secretName }} - {{- end }} - {{- end }} - rules: - {{- range .Values.ingress.hosts }} - - host: {{ .host | quote }} - http: - paths: - {{- if and $.Values.lens.enabled $.Values.lens.service.name }} - - path: /lens-ingest - pathType: Prefix - backend: - {{- if semverCompare ">=1.19-0" $.Capabilities.KubeVersion.GitVersion }} - service: - name: {{ $.Values.lens.service.name }} - port: - number: {{ $.Values.lens.service.port }} - {{- else }} - serviceName: {{ $.Values.lens.service.name }} - servicePort: {{ $.Values.lens.service.port }} - {{- end }} - {{- end }} - {{- range .paths }} - - path: {{ .path }} - {{- if and .pathType (semverCompare ">=1.18-0" $.Capabilities.KubeVersion.GitVersion) }} - pathType: {{ .pathType }} - {{- end }} - backend: - {{- if semverCompare ">=1.19-0" $.Capabilities.KubeVersion.GitVersion }} - service: - name: {{ $fullName }} - port: - number: {{ $svcPort }} - {{- else }} - serviceName: {{ $fullName }} - servicePort: {{ $svcPort }} - {{- end }} - {{- end }} - {{- end }} -{{- end }} diff --git a/helm/litellm-helm/templates/keda.yaml b/helm/litellm-helm/templates/keda.yaml deleted file mode 100644 index bf585d0d4be..00000000000 --- a/helm/litellm-helm/templates/keda.yaml +++ /dev/null @@ -1,58 +0,0 @@ -{{- if and .Values.keda.enabled (not .Values.autoscaling.enabled) }} -apiVersion: keda.sh/v1alpha1 -kind: ScaledObject -metadata: - name: {{ include "litellm.fullname" . }} - labels: - {{- include "litellm.labels" . | nindent 4 }} - {{- if .Values.keda.scaledObject.annotations }} - annotations: {{ toYaml .Values.keda.scaledObject.annotations | nindent 4 }} - {{- end }} -spec: - scaleTargetRef: - name: {{ include "litellm.fullname" . }} - pollingInterval: {{ .Values.keda.pollingInterval }} - cooldownPeriod: {{ .Values.keda.cooldownPeriod }} - minReplicaCount: {{ .Values.keda.minReplicas }} - maxReplicaCount: {{ .Values.keda.maxReplicas }} -{{- with .Values.keda.fallback }} - fallback: - failureThreshold: {{ .failureThreshold | default 3 }} - replicas: {{ .replicas | default $.Values.keda.maxReplicas }} -{{- end }} - triggers: -{{- with .Values.keda.triggers }} - {{- toYaml . | nindent 2 }} -{{- end }} -{{- $prom := .Values.keda.prometheus }} -{{- if or $prom.requestsPerSecond $prom.tokensPerSecond }} -{{- if not $prom.serverAddress }} -{{- fail "keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set" }} -{{- end }} -{{- $selector := printf "namespace=%q,job=%q" .Release.Namespace (printf "%s%s" (include "litellm.fullname" .) (ternary "-metrics" "" .Values.metricsServer.enabled)) }} -{{- with $prom.requestsPerSecond }} - - type: prometheus - metadata: - serverAddress: {{ $prom.serverAddress | quote }} - threshold: {{ toJson . | trimAll "\"" | quote }} - query: {{ printf "sum(rate(litellm_proxy_total_requests_metric_total{%s}[1m]))" $selector | quote }} -{{- end }} -{{- with $prom.tokensPerSecond }} - - type: prometheus - metadata: - serverAddress: {{ $prom.serverAddress | quote }} - threshold: {{ toJson . | trimAll "\"" | quote }} - query: {{ printf "sum(rate(litellm_total_tokens_metric_total{%s}[1m]))" $selector | quote }} -{{- end }} -{{- end }} - advanced: - restoreToOriginalReplicaCount: {{ .Values.keda.restoreToOriginalReplicaCount }} -{{- if .Values.keda.behavior }} - horizontalPodAutoscalerConfig: - behavior: -{{- with .Values.keda.behavior }} -{{- toYaml . | nindent 8 }} -{{- end }} - -{{- end }} -{{- end }} diff --git a/helm/litellm-helm/templates/migrations-job.yaml b/helm/litellm-helm/templates/migrations-job.yaml deleted file mode 100644 index 7aa96a59102..00000000000 --- a/helm/litellm-helm/templates/migrations-job.yaml +++ /dev/null @@ -1,122 +0,0 @@ -{{- if .Values.migrationJob.enabled }} -# This job runs the Prisma migrations for the LiteLLM DB. -apiVersion: batch/v1 -kind: Job -metadata: - name: {{ include "litellm.fullname" . }}-migrations - labels: - {{- include "litellm.labels" . | nindent 4 }} - {{- with .Values.migrationJob.jobLabels }} - {{- toYaml . | nindent 4 }} - {{- end }} - annotations: - {{- if .Values.migrationJob.hooks.argocd.enabled }} - argocd.argoproj.io/hook: PreSync - argocd.argoproj.io/hook-delete-policy: BeforeHookCreation - {{- end }} - {{- if .Values.migrationJob.hooks.helm.enabled }} - helm.sh/hook: "pre-install,pre-upgrade" - helm.sh/hook-delete-policy: "before-hook-creation" - helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "1" | quote }} - {{- end }} - checksum/config: {{ toYaml .Values | sha256sum }} - {{- with .Values.migrationJob.jobAnnotations }} - {{- toYaml . | nindent 4 }} - {{- end }} -spec: - template: - metadata: - labels: - {{- include "litellm.labels" . | nindent 8 }} - {{- with .Values.podLabels }} - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.migrationJob.podLabels }} - {{- toYaml . | nindent 8 }} - {{- end }} - annotations: - {{- with .Values.migrationJob.annotations }} - {{- toYaml . | nindent 8 }} - {{- end }} - spec: - {{- with .Values.imagePullSecrets }} - imagePullSecrets: - {{- toYaml . | nindent 8 }} - {{- end }} - serviceAccountName: {{ include "litellm.migrationServiceAccountName" . }} - securityContext: - {{- toYaml .Values.podSecurityContext | nindent 8 }} - {{- with .Values.migrationJob.extraInitContainers }} - initContainers: - {{- tpl (toYaml .) $ | nindent 8 }} - {{- end }} - containers: - - name: prisma-migrations - image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.image.pullPolicy }} - securityContext: - {{- toYaml .Values.securityContext | nindent 12 }} - {{- if .Values.migrationJob.command }} - command: {{ toYaml .Values.migrationJob.command | nindent 12 }} - {{- else }} - command: ["python", "litellm/proxy/prisma_migration.py"] - {{- end }} - {{- if .Values.migrationJob.args }} - args: {{ toYaml .Values.migrationJob.args | nindent 12 }} - {{- else if .Values.migrationJob.command }} - args: [] - {{- end }} - workingDir: "/app" - env: - {{- include "litellm.validateDbTokenAuth" . -}} - {{- if .Values.db.useExisting }} - {{- include "litellm.existingDbEnv" . | nindent 12 }} - {{- else if .Values.db.deployStandalone }} - - name: DATABASE_URL - value: postgresql://{{ .Values.postgresql.auth.username }}:{{ .Values.postgresql.auth.password }}@{{ .Release.Name }}-postgresql/{{ .Values.postgresql.auth.database }} - {{- end }} - {{- if .Values.envVars }} - {{- range $key, $val := .Values.envVars }} - - name: {{ $key }} - value: {{ $val | quote }} - {{- end }} - {{- end }} - {{- with .Values.extraEnvVars }} - {{- toYaml . | nindent 12 }} - {{- end }} - - name: DISABLE_SCHEMA_UPDATE - value: "false" # always run the migration from the Helm PreSync hook, override the value set - {{- with .Values.volumeMounts }} - volumeMounts: - {{- toYaml . | nindent 12 }} - {{- end }} - {{- with .Values.migrationJob.resources }} - resources: - {{- toYaml . | nindent 12 }} - {{- end }} - {{- with .Values.migrationJob.extraContainers }} - {{- tpl (toYaml .) $ | nindent 8 }} - {{- end }} - {{- with .Values.volumes }} - volumes: - {{- toYaml . | nindent 8 }} - {{- end }} - restartPolicy: OnFailure - {{- with .Values.nodeSelector }} - nodeSelector: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.affinity }} - affinity: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.tolerations }} - tolerations: - {{- toYaml . | nindent 8 }} - {{- end }} - ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }} - backoffLimit: {{ .Values.migrationJob.backoffLimit }} - {{- with .Values.migrationJob.activeDeadlineSeconds }} - activeDeadlineSeconds: {{ . }} - {{- end }} -{{- end }} diff --git a/helm/litellm-helm/templates/poddisruptionbudget.yaml b/helm/litellm-helm/templates/poddisruptionbudget.yaml deleted file mode 100644 index 1715b94c1f6..00000000000 --- a/helm/litellm-helm/templates/poddisruptionbudget.yaml +++ /dev/null @@ -1,33 +0,0 @@ -{{- /* -PodDisruptionBudget for LiteLLM proxy -Controlled via .Values.pdb.enabled and .Values.pdb.{minAvailable|maxUnavailable} -Only one of minAvailable / maxUnavailable should be set. If both are set, minAvailable wins. -*/ -}} -{{- if .Values.pdb.enabled }} -apiVersion: policy/v1 -kind: PodDisruptionBudget -metadata: - name: {{ include "litellm.fullname" . }} - labels: - {{- include "litellm.labels" . | nindent 4 }} - {{- with .Values.pdb.labels }} - {{- toYaml . | nindent 4 }} - {{- end }} - {{- with .Values.pdb.annotations }} - annotations: - {{- toYaml . | nindent 4 }} - {{- end }} -spec: - selector: - matchLabels: - {{- /* Match the Deployment selector to target the same pod set */ -}} - {{- include "litellm.selectorLabels" . | nindent 6 }} - {{- if .Values.pdb.minAvailable }} - minAvailable: {{ .Values.pdb.minAvailable }} - {{- else if .Values.pdb.maxUnavailable }} - maxUnavailable: {{ .Values.pdb.maxUnavailable }} - {{- else }} - # Safe default if enabled but not configured - maxUnavailable: 1 - {{- end }} -{{- end }} diff --git a/helm/litellm-helm/templates/secret-dbcredentials.yaml b/helm/litellm-helm/templates/secret-dbcredentials.yaml deleted file mode 100644 index 8ab89a4579e..00000000000 --- a/helm/litellm-helm/templates/secret-dbcredentials.yaml +++ /dev/null @@ -1,13 +0,0 @@ -{{- if .Values.db.deployStandalone -}} -{{- include "litellm.validateBundledPostgresImageTag" . -}} -apiVersion: v1 -kind: Secret -metadata: - name: {{ include "litellm.fullname" . }}-dbcredentials -data: - # Password for the "postgres" user - postgres-password: {{ ( index .Values.postgresql.auth "postgres-password") | default "litellm" | b64enc }} - username: {{ .Values.postgresql.auth.username | default "litellm" | b64enc }} - password: {{ .Values.postgresql.auth.password | default "litellm" | b64enc }} -type: Opaque -{{- end -}} \ No newline at end of file diff --git a/helm/litellm-helm/templates/secret-masterkey.yaml b/helm/litellm-helm/templates/secret-masterkey.yaml deleted file mode 100644 index 60ab4e74c6b..00000000000 --- a/helm/litellm-helm/templates/secret-masterkey.yaml +++ /dev/null @@ -1,12 +0,0 @@ -{{- if not .Values.masterkeySecretName }} -{{- $secretName := printf "%s-masterkey" (include "litellm.fullname" .) }} -{{- $existing := lookup "v1" "Secret" .Release.Namespace $secretName }} -{{- $masterkey := .Values.masterkey | default (dig "data" "masterkey" "" $existing | b64dec) | default (printf "sk-%s" (randAlphaNum 18)) }} -apiVersion: v1 -kind: Secret -metadata: - name: {{ $secretName }} -data: - masterkey: {{ $masterkey | b64enc }} -type: Opaque -{{- end }} diff --git a/helm/litellm-helm/templates/service-metrics.yaml b/helm/litellm-helm/templates/service-metrics.yaml deleted file mode 100644 index 1d23fe39606..00000000000 --- a/helm/litellm-helm/templates/service-metrics.yaml +++ /dev/null @@ -1,17 +0,0 @@ -{{- if .Values.metricsServer.enabled }} -apiVersion: v1 -kind: Service -metadata: - name: {{ include "litellm.fullname" . }}-metrics - labels: - {{- include "litellm.labels" . | nindent 4 }} -spec: - type: ClusterIP - ports: - - port: {{ .Values.metricsServer.port }} - targetPort: metrics - protocol: TCP - name: metrics - selector: - {{- include "litellm.selectorLabels" . | nindent 4 }} -{{- end }} diff --git a/helm/litellm-helm/templates/service.yaml b/helm/litellm-helm/templates/service.yaml deleted file mode 100644 index 11812208929..00000000000 --- a/helm/litellm-helm/templates/service.yaml +++ /dev/null @@ -1,22 +0,0 @@ -apiVersion: v1 -kind: Service -metadata: - name: {{ include "litellm.fullname" . }} - {{- with .Values.service.annotations }} - annotations: - {{- toYaml . | nindent 4 }} - {{- end }} - labels: - {{- include "litellm.labels" . | nindent 4 }} -spec: - type: {{ .Values.service.type }} - {{- if and (eq .Values.service.type "LoadBalancer") .Values.service.loadBalancerClass }} - loadBalancerClass: {{ .Values.service.loadBalancerClass }} - {{- end }} - ports: - - port: {{ .Values.service.port }} - targetPort: http - protocol: TCP - name: http - selector: - {{- include "litellm.selectorLabels" . | nindent 4 }} diff --git a/helm/litellm-helm/templates/serviceaccount.yaml b/helm/litellm-helm/templates/serviceaccount.yaml deleted file mode 100644 index 7655470fa42..00000000000 --- a/helm/litellm-helm/templates/serviceaccount.yaml +++ /dev/null @@ -1,13 +0,0 @@ -{{- if .Values.serviceAccount.create -}} -apiVersion: v1 -kind: ServiceAccount -metadata: - name: {{ include "litellm.serviceAccountName" . }} - labels: - {{- include "litellm.labels" . | nindent 4 }} - {{- with .Values.serviceAccount.annotations }} - annotations: - {{- toYaml . | nindent 4 }} - {{- end }} -automountServiceAccountToken: {{ .Values.serviceAccount.automount }} -{{- end }} diff --git a/helm/litellm-helm/templates/servicemonitor.yaml b/helm/litellm-helm/templates/servicemonitor.yaml deleted file mode 100644 index 68083d0da61..00000000000 --- a/helm/litellm-helm/templates/servicemonitor.yaml +++ /dev/null @@ -1,39 +0,0 @@ -{{- with .Values.serviceMonitor }} -{{- if and (eq .enabled true) }} -apiVersion: monitoring.coreos.com/v1 -kind: ServiceMonitor -metadata: - name: {{ include "litellm.fullname" $ }} - labels: - {{- include "litellm.labels" $ | nindent 4 }} - {{- if .labels }} - {{- toYaml .labels | nindent 4 }} - {{- end }} - {{- if .annotations }} - annotations: - {{- toYaml .annotations | nindent 4 }} - {{- end }} -spec: - selector: - matchLabels: - {{- include "litellm.selectorLabels" $ | nindent 6 }} - namespaceSelector: - matchNames: - # if not set, use the release namespace - {{- if not .namespaceSelector.matchNames }} - - {{ $.Release.Namespace | quote }} - {{- else }} - {{- toYaml .namespaceSelector.matchNames | nindent 4 }} - {{- end }} - endpoints: - - port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }} - path: /metrics/ - interval: {{ .interval }} - scrapeTimeout: {{ .scrapeTimeout }} - scheme: http - {{- if .relabelings }} - relabelings: -{{- toYaml .relabelings | nindent 4 }} - {{- end }} -{{- end }} -{{- end }} diff --git a/helm/litellm-helm/templates/tests/test-connection.yaml b/helm/litellm-helm/templates/tests/test-connection.yaml deleted file mode 100644 index 86a8f66b10b..00000000000 --- a/helm/litellm-helm/templates/tests/test-connection.yaml +++ /dev/null @@ -1,25 +0,0 @@ -apiVersion: v1 -kind: Pod -metadata: - name: "{{ include "litellm.fullname" . }}-test-connection" - labels: - {{- include "litellm.labels" . | nindent 4 }} - annotations: - "helm.sh/hook": test -spec: - containers: - - name: wget - image: busybox - command: ['sh', '-c'] - args: - - | - # Wait for a bit to allow the service to be ready - sleep 10 - # Try multiple times with a delay between attempts - for i in $(seq 1 30); do - wget -T 5 "{{ include "litellm.fullname" . }}:{{ .Values.service.port }}/health/readiness" && exit 0 - echo "Attempt $i failed, waiting..." - sleep 2 - done - exit 1 - restartPolicy: Never \ No newline at end of file diff --git a/helm/litellm-helm/templates/tests/test-env-vars.yaml b/helm/litellm-helm/templates/tests/test-env-vars.yaml deleted file mode 100644 index 9f0277557a4..00000000000 --- a/helm/litellm-helm/templates/tests/test-env-vars.yaml +++ /dev/null @@ -1,43 +0,0 @@ -apiVersion: v1 -kind: Pod -metadata: - name: "{{ include "litellm.fullname" . }}-env-test" - labels: - {{- include "litellm.labels" . | nindent 4 }} - annotations: - "helm.sh/hook": test -spec: - containers: - - name: test - image: busybox - command: ['sh', '-c'] - args: - - | - # Test DD_ENV - if [ "$DD_ENV" != "dev_helm" ]; then - echo "❌ Environment variable DD_ENV mismatch. Expected: dev_helm, Got: $DD_ENV" - exit 1 - fi - echo "✅ Environment variable DD_ENV matches expected value: $DD_ENV" - - # Test DD_SERVICE - if [ "$DD_SERVICE" != "litellm" ]; then - echo "❌ Environment variable DD_SERVICE mismatch. Expected: litellm, Got: $DD_SERVICE" - exit 1 - fi - echo "✅ Environment variable DD_SERVICE matches expected value: $DD_SERVICE" - - # Test USE_DDTRACE - if [ "$USE_DDTRACE" != "true" ]; then - echo "❌ Environment variable USE_DDTRACE mismatch. Expected: true, Got: $USE_DDTRACE" - exit 1 - fi - echo "✅ Environment variable USE_DDTRACE matches expected value: $USE_DDTRACE" - env: - - name: DD_ENV - value: {{ .Values.envVars.DD_ENV | quote }} - - name: DD_SERVICE - value: {{ .Values.envVars.DD_SERVICE | quote }} - - name: USE_DDTRACE - value: {{ .Values.envVars.USE_DDTRACE | quote }} - restartPolicy: Never \ No newline at end of file diff --git a/helm/litellm-helm/templates/tests/test-servicemonitor.yaml b/helm/litellm-helm/templates/tests/test-servicemonitor.yaml deleted file mode 100644 index ef8475339c3..00000000000 --- a/helm/litellm-helm/templates/tests/test-servicemonitor.yaml +++ /dev/null @@ -1,152 +0,0 @@ -{{- if .Values.serviceMonitor.enabled }} -apiVersion: v1 -kind: Pod -metadata: - name: "{{ include "litellm.fullname" . }}-test-servicemonitor" - labels: - {{- include "litellm.labels" . | nindent 4 }} - annotations: - "helm.sh/hook": test -spec: - containers: - - name: test - image: docker.io/bitnamilegacy/kubectl:1.29.2-debian-12-r3 - command: ['sh', '-c'] - args: - - | - set -e - echo "🔍 Testing ServiceMonitor configuration..." - - # Check if ServiceMonitor exists - if ! kubectl get servicemonitor {{ include "litellm.fullname" . }} -n {{ .Release.Namespace }} &>/dev/null; then - echo "❌ ServiceMonitor not found" - exit 1 - fi - echo "✅ ServiceMonitor exists" - - # Get ServiceMonitor YAML - SM=$(kubectl get servicemonitor {{ include "litellm.fullname" . }} -n {{ .Release.Namespace }} -o yaml) - - # Test endpoint configuration - ENDPOINT_PORT=$(echo "$SM" | grep -A 5 "endpoints:" | grep "port:" | awk '{print $2}') - if [ "$ENDPOINT_PORT" != "http" ]; then - echo "❌ Endpoint port mismatch. Expected: http, Got: $ENDPOINT_PORT" - exit 1 - fi - echo "✅ Endpoint port is correctly set to: $ENDPOINT_PORT" - - # Test endpoint path - ENDPOINT_PATH=$(echo "$SM" | grep -A 5 "endpoints:" | grep "path:" | awk '{print $2}') - if [ "$ENDPOINT_PATH" != "/metrics/" ]; then - echo "❌ Endpoint path mismatch. Expected: /metrics/, Got: $ENDPOINT_PATH" - exit 1 - fi - echo "✅ Endpoint path is correctly set to: $ENDPOINT_PATH" - - # Test interval - INTERVAL=$(echo "$SM" | grep "interval:" | awk '{print $2}') - if [ "$INTERVAL" != "{{ .Values.serviceMonitor.interval }}" ]; then - echo "❌ Interval mismatch. Expected: {{ .Values.serviceMonitor.interval }}, Got: $INTERVAL" - exit 1 - fi - echo "✅ Interval is correctly set to: $INTERVAL" - - # Test scrapeTimeout - TIMEOUT=$(echo "$SM" | grep "scrapeTimeout:" | awk '{print $2}') - if [ "$TIMEOUT" != "{{ .Values.serviceMonitor.scrapeTimeout }}" ]; then - echo "❌ ScrapeTimeout mismatch. Expected: {{ .Values.serviceMonitor.scrapeTimeout }}, Got: $TIMEOUT" - exit 1 - fi - echo "✅ ScrapeTimeout is correctly set to: $TIMEOUT" - - # Test scheme - SCHEME=$(echo "$SM" | grep "scheme:" | awk '{print $2}') - if [ "$SCHEME" != "http" ]; then - echo "❌ Scheme mismatch. Expected: http, Got: $SCHEME" - exit 1 - fi - echo "✅ Scheme is correctly set to: $SCHEME" - - {{- if .Values.serviceMonitor.labels }} - # Test custom labels - echo "🔍 Checking custom labels..." - {{- range $key, $value := .Values.serviceMonitor.labels }} - LABEL_VALUE=$(echo "$SM" | grep -A 20 "metadata:" | grep "{{ $key }}:" | awk '{print $2}') - if [ "$LABEL_VALUE" != "{{ $value }}" ]; then - echo "❌ Label {{ $key }} mismatch. Expected: {{ $value }}, Got: $LABEL_VALUE" - exit 1 - fi - echo "✅ Label {{ $key }} is correctly set to: {{ $value }}" - {{- end }} - {{- end }} - - {{- if .Values.serviceMonitor.annotations }} - # Test annotations - echo "🔍 Checking annotations..." - {{- range $key, $value := .Values.serviceMonitor.annotations }} - ANNOTATION_VALUE=$(echo "$SM" | grep -A 10 "annotations:" | grep "{{ $key }}:" | awk '{print $2}') - if [ "$ANNOTATION_VALUE" != "{{ $value }}" ]; then - echo "❌ Annotation {{ $key }} mismatch. Expected: {{ $value }}, Got: $ANNOTATION_VALUE" - exit 1 - fi - echo "✅ Annotation {{ $key }} is correctly set to: {{ $value }}" - {{- end }} - {{- end }} - - {{- if .Values.serviceMonitor.namespaceSelector.matchNames }} - # Test namespace selector - echo "🔍 Checking namespace selector..." - {{- range .Values.serviceMonitor.namespaceSelector.matchNames }} - if ! echo "$SM" | grep -A 5 "namespaceSelector:" | grep -q "{{ . }}"; then - echo "❌ Namespace {{ . }} not found in namespaceSelector" - exit 1 - fi - echo "✅ Namespace {{ . }} found in namespaceSelector" - {{- end }} - {{- else }} - # Test default namespace selector (should be release namespace) - if ! echo "$SM" | grep -A 5 "namespaceSelector:" | grep -q "{{ .Release.Namespace }}"; then - echo "❌ Release namespace {{ .Release.Namespace }} not found in namespaceSelector" - exit 1 - fi - echo "✅ Default namespace selector set to release namespace: {{ .Release.Namespace }}" - {{- end }} - - {{- if .Values.serviceMonitor.relabelings }} - # Test relabelings - echo "🔍 Checking relabelings configuration..." - if ! echo "$SM" | grep -q "relabelings:"; then - echo "❌ Relabelings section not found" - exit 1 - fi - echo "✅ Relabelings section exists" - {{- range .Values.serviceMonitor.relabelings }} - {{- if .targetLabel }} - if ! echo "$SM" | grep -A 50 "relabelings:" | grep -q "targetLabel: {{ .targetLabel }}"; then - echo "❌ Relabeling targetLabel {{ .targetLabel }} not found" - exit 1 - fi - echo "✅ Relabeling targetLabel {{ .targetLabel }} found" - {{- end }} - {{- if .action }} - if ! echo "$SM" | grep -A 50 "relabelings:" | grep -q "action: {{ .action }}"; then - echo "❌ Relabeling action {{ .action }} not found" - exit 1 - fi - echo "✅ Relabeling action {{ .action }} found" - {{- end }} - {{- end }} - {{- end }} - - # Test selector labels match the service - echo "🔍 Checking selector labels match service..." - SVC_LABELS=$(kubectl get svc {{ include "litellm.fullname" . }} -n {{ .Release.Namespace }} -o jsonpath='{.metadata.labels}') - echo "Service labels: $SVC_LABELS" - echo "✅ Selector labels validation passed" - - echo "" - echo "🎉 All ServiceMonitor tests passed successfully!" - serviceAccountName: {{ include "litellm.serviceAccountName" . }} - restartPolicy: Never -{{- end }} - diff --git a/helm/litellm-helm/tests/billing_metrics_tests.yaml b/helm/litellm-helm/tests/billing_metrics_tests.yaml deleted file mode 100644 index 71803df378c..00000000000 --- a/helm/litellm-helm/tests/billing_metrics_tests.yaml +++ /dev/null @@ -1,297 +0,0 @@ -suite: test billingMetrics wiring on the proxy deployment -templates: - - deployment.yaml - - configmap-litellm.yaml - - migrations-job.yaml -tests: - - it: is off by default, adding no env, volume, or mount - template: deployment.yaml - asserts: - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - notContains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - - it: renders the endpoint and the mounted cert paths when enabled - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: /etc/litellm/billing-mtls/tls.crt - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: /etc/litellm/billing-mtls/tls.key - - # The conventional Secret name is the default, so enabling the block is enough. - - it: mounts the default cert secret read-only alongside the config volume - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - - it: honours a secretName override - template: deployment.yaml - set: - billingMetrics: - enabled: true - secretName: my-billing-mtls - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: my-billing-mtls - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - - it: honours an endpoint override - template: deployment.yaml - set: - billingMetrics: - enabled: true - endpoint: https://collector.internal:4318 - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://collector.internal:4318 - - # The production collector presents a public web-PKI certificate, so the CA - # override must stay absent unless a private collector is configured. - - it: omits the CA env, volume, and mount when no caSecretName is set - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca - - notContains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt - - - it: mounts the CA secret when caSecretName is set - template: deployment.yaml - set: - billingMetrics: - enabled: true - caSecretName: billing-ca - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true - - - it: passes the export interval through only when set - template: deployment.yaml - set: - billingMetrics: - enabled: true - exportIntervalMs: 5000 - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS - value: "5000" - - - it: omits the export interval when unset - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS - value: "60000" - - # Kubernetes resolves duplicate env names last-wins, so the chart-owned billing - # entries must render after .Values.envVars or a user could silently redirect - # the metering export. The three billing entries are the last ones emitted here - # (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case). - - it: renders the billing endpoint after envVars so it cannot be shadowed - template: deployment.yaml - set: - migrationJob: - enabled: false - billingMetrics: - enabled: true - envVars: - LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://shadowed.example - - equal: - path: spec.template.spec.containers[0].env[-3] - value: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - equal: - path: spec.template.spec.containers[0].env[-2].name - value: LITELLM_BILLING_METRICS_CLIENT_CERT - - equal: - path: spec.template.spec.containers[0].env[-1].name - value: LITELLM_BILLING_METRICS_CLIENT_KEY - - - it: keeps user-supplied volumes and mounts alongside the billing secret - template: deployment.yaml - set: - billingMetrics: - enabled: true - volumes: - - name: custom-callbacks - configMap: - name: my-callbacks - volumeMounts: - - name: custom-callbacks - mountPath: /app/callbacks - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: custom-callbacks - configMap: - name: my-callbacks - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: custom-callbacks - mountPath: /app/callbacks - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - - it: still mounts the proxy config when enabled - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: litellm-config - mountPath: /etc/litellm/config.yaml - subPath: config.yaml - - # Only the proxy serves billable traffic. The migrations Job must never mount - # the client certificate, and it renders its own env and volumes, so nothing - # stops a future edit from wiring the billing include into it by mistake. - - it: does not touch the migrations job when enabled - template: migrations-job.yaml - set: - billingMetrics: - enabled: true - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - notExists: - path: spec.template.spec.containers[0].volumeMounts - - notExists: - path: spec.template.spec.volumes - - - it: fails loudly when enabled with an emptied secretName - template: deployment.yaml - set: - billingMetrics: - enabled: true - secretName: "" - asserts: - - failedTemplate: - errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key) - - - it: fails loudly when enabled without an endpoint - template: deployment.yaml - set: - billingMetrics: - enabled: true - endpoint: "" - asserts: - - failedTemplate: - errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true diff --git a/helm/litellm-helm/tests/bundled_db_images_tests.yaml b/helm/litellm-helm/tests/bundled_db_images_tests.yaml deleted file mode 100644 index 8f0860c2721..00000000000 --- a/helm/litellm-helm/tests/bundled_db_images_tests.yaml +++ /dev/null @@ -1,94 +0,0 @@ -suite: test bundled database images -templates: - - charts/postgresql/templates/primary/statefulset.yaml - - charts/redis/templates/master/application.yaml - - charts/redis/templates/configmap.yaml - - charts/redis/templates/health-configmap.yaml - - charts/redis/templates/scripts-configmap.yaml - - charts/redis/templates/secret.yaml - - secret-dbcredentials.yaml - - templates/tests/test-servicemonitor.yaml -tests: - - it: should pull the bundled postgres from a repository that still publishes the pinned tag - template: charts/postgresql/templates/primary/statefulset.yaml - set: - db.deployStandalone: true - asserts: - - equal: - path: spec.template.spec.containers[0].image - value: docker.io/bitnamilegacy/postgresql:16.2.0-debian-12-r6 - - - it: should pull the bundled postgres metrics exporter from the same repository - template: charts/postgresql/templates/primary/statefulset.yaml - set: - db.deployStandalone: true - postgresql.metrics.enabled: true - asserts: - - equal: - path: spec.template.spec.containers[1].image - value: docker.io/bitnamilegacy/postgres-exporter:0.15.0-debian-12-r14 - - - it: should run the bundled postgres init container from the same repository - template: charts/postgresql/templates/primary/statefulset.yaml - set: - db.deployStandalone: true - postgresql.volumePermissions.enabled: true - asserts: - - equal: - path: spec.template.spec.initContainers[0].image - value: docker.io/bitnamilegacy/os-shell:12-debian-12-r16 - - - it: should pull the bundled redis from a repository that still publishes the pinned tag - template: charts/redis/templates/master/application.yaml - set: - redis.enabled: true - asserts: - - equal: - path: spec.template.spec.containers[0].image - value: docker.io/bitnamilegacy/redis:7.2.4-debian-12-r9 - - - it: should reject a floating postgres tag that could cross a major version on an existing volume - template: secret-dbcredentials.yaml - set: - db.deployStandalone: true - postgresql.image.tag: latest - asserts: - - failedTemplate: - errorMessage: 'postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got "latest"). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore.' - - - it: should reject an empty postgres tag - template: secret-dbcredentials.yaml - set: - db.deployStandalone: true - postgresql.image.tag: "" - asserts: - - failedTemplate: - errorMessage: 'postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got ""). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore.' - - - it: should accept an empty postgres tag when the image is pinned by digest - template: secret-dbcredentials.yaml - set: - db.deployStandalone: true - postgresql.image.tag: "" - postgresql.image.digest: sha256:0d0e2f1a5b3c4d6e7f8091a2b3c4d5e6f708192a3b4c5d6e7f8091a2b3c4d5e6 - asserts: - - hasDocuments: - count: 1 - - - it: should run the servicemonitor test pod from a pinned image - template: templates/tests/test-servicemonitor.yaml - set: - serviceMonitor.enabled: true - asserts: - - equal: - path: spec.containers[0].image - value: docker.io/bitnamilegacy/kubectl:1.29.2-debian-12-r3 - - - it: should not constrain the postgres tag when the bundled database is not deployed - template: secret-dbcredentials.yaml - set: - db.deployStandalone: false - postgresql.image.tag: latest - asserts: - - hasDocuments: - count: 0 diff --git a/helm/litellm-helm/tests/collector_tests.yaml b/helm/litellm-helm/tests/collector_tests.yaml deleted file mode 100644 index 0340b1161b7..00000000000 --- a/helm/litellm-helm/tests/collector_tests.yaml +++ /dev/null @@ -1,272 +0,0 @@ -suite: test collector sidecar -templates: - - deployment.yaml - - hpa.yaml - - configmap-litellm.yaml -tests: - - it: should run the proxy alone with no collector env by default - template: deployment.yaml - asserts: - - lengthEqual: - path: spec.template.spec.containers - count: 1 - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_COLLECTOR_ENABLED - value: "true" - - notContains: - path: spec.template.spec.volumes - content: - name: collector-socket - any: true - - - it: should add the sidecar on the same image and point both containers at the unix socket - template: deployment.yaml - set: - image.tag: test - db.connectionPool.enabled: true - collector.enabled: true - collector.resources: - requests: - cpu: 500m - memory: 1Gi - limits: - cpu: "1" - memory: 2Gi - asserts: - - lengthEqual: - path: spec.template.spec.containers - count: 2 - - equal: - path: spec.template.spec.containers[1].name - value: litellm-collector - - equal: - path: spec.template.spec.containers[1].image - value: ghcr.io/berriai/litellm:test - - equal: - path: spec.template.spec.containers[1].command - value: [python, -m, litellm.proxy.collector] - - equal: - path: spec.template.spec.containers[1].resources.requests.cpu - value: 500m - - equal: - path: spec.template.spec.containers[1].resources.limits.memory - value: 2Gi - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_COLLECTOR_ENABLED - value: "true" - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_COLLECTOR_ADDRESS - value: unix:///var/run/litellm/collector.sock - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_COLLECTOR_BUFFER_SIZE - value: "1000" - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_COLLECTOR_ON_UNAVAILABLE - value: fallback - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_JOB_ROLE - value: collector - - contains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_JOB_ROLE - value: collector - - contains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_COLLECTOR_ADDRESS - value: unix:///var/run/litellm/collector.sock - - contains: - path: spec.template.spec.containers[1].env - content: - name: CONFIG_FILE_PATH - value: /etc/litellm/config.yaml - - contains: - path: spec.template.spec.containers[1].env - content: - name: DATABASE_HOST - value: RELEASE-NAME-postgresql - - contains: - path: spec.template.spec.containers[1].env - content: - name: DATABASE_PASSWORD - valueFrom: - secretKeyRef: - name: RELEASE-NAME-litellm-dbcredentials - key: password - - contains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_PGBOUNCER_ENABLED - value: "true" - - contains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS - value: "20" - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: collector-socket - mountPath: /var/run/litellm - - contains: - path: spec.template.spec.containers[1].volumeMounts - content: - name: collector-socket - mountPath: /var/run/litellm - - contains: - path: spec.template.spec.containers[1].volumeMounts - content: - name: litellm-config - mountPath: /etc/litellm/config.yaml - subPath: config.yaml - - contains: - path: spec.template.spec.volumes - content: - name: collector-socket - emptyDir: - sizeLimit: 1Mi - - - it: should skip the socket volume and pass the policy through on tcp transport - template: deployment.yaml - set: - collector.enabled: true - collector.address: tcp://127.0.0.1:4100 - collector.onUnavailable: drop - collector.bufferSize: 50 - envVars: - CONFIG_FILE_PATH: /custom/config.yaml - asserts: - - lengthEqual: - path: spec.template.spec.containers - count: 2 - - notContains: - path: spec.template.spec.containers[1].env - content: - name: CONFIG_FILE_PATH - value: /etc/litellm/config.yaml - - contains: - path: spec.template.spec.containers[1].env - content: - name: CONFIG_FILE_PATH - value: /custom/config.yaml - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_COLLECTOR_ADDRESS - value: tcp://127.0.0.1:4100 - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_COLLECTOR_ON_UNAVAILABLE - value: drop - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_COLLECTOR_BUFFER_SIZE - value: "50" - - notContains: - path: spec.template.spec.volumes - content: - name: collector-socket - any: true - - - it: should keep metrics and billing env on the proxy container only - template: deployment.yaml - set: - collector.enabled: true - metricsServer.enabled: true - metricsServer.port: 9090 - billingMetrics.enabled: true - billingMetrics.endpoint: https://metering.example.com - billingMetrics.secretName: billing-mtls - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: PROMETHEUS_METRICS_PORT - value: "9090" - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://metering.example.com - - notContains: - path: spec.template.spec.containers[1].env - content: - name: PROMETHEUS_METRICS_PORT - any: true - - notContains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - any: true - - notContains: - path: spec.template.spec.containers[1].volumeMounts - content: - name: billing-metrics-mtls - any: true - - - it: should give the sidecar the same scratch mounts as the proxy on a read-only root - template: deployment.yaml - set: - collector.enabled: true - securityContext.readOnlyRootFilesystem: true - asserts: - - contains: - path: spec.template.spec.containers[1].volumeMounts - content: - name: npm - mountPath: /.npm - - contains: - path: spec.template.spec.containers[1].volumeMounts - content: - name: cache - mountPath: /.cache - - contains: - path: spec.template.spec.containers[1].volumeMounts - content: - name: tmp - mountPath: /tmp - - - it: should keep the pod-wide cpu metric unless asked to scale on the proxy container - template: hpa.yaml - set: - autoscaling.enabled: true - collector.enabled: true - asserts: - - equal: { path: "spec.metrics[0].type", value: Resource } - - equal: { path: "spec.metrics[0].resource.name", value: cpu } - - - it: should scale on the proxy container's cpu only when opted in - template: hpa.yaml - set: - autoscaling.enabled: true - collector.enabled: true - collector.scaleOnProxyContainerCpu: true - asserts: - - equal: { path: "spec.metrics[0].type", value: ContainerResource } - - equal: { path: "spec.metrics[0].containerResource.name", value: cpu } - - equal: { path: "spec.metrics[0].containerResource.container", value: litellm } - - equal: { path: "spec.metrics[0].containerResource.target.averageUtilization", value: 60 } - - isNull: { path: "spec.metrics[0].resource" } - - - it: should not switch to the container metric while the sidecar is off - template: hpa.yaml - set: - autoscaling.enabled: true - collector.scaleOnProxyContainerCpu: true - asserts: - - equal: { path: "spec.metrics[0].type", value: Resource } diff --git a/helm/litellm-helm/tests/connection_pool_tests.yaml b/helm/litellm-helm/tests/connection_pool_tests.yaml deleted file mode 100644 index 203082f27ba..00000000000 --- a/helm/litellm-helm/tests/connection_pool_tests.yaml +++ /dev/null @@ -1,112 +0,0 @@ -suite: test in-container connection pool -templates: - - deployment.yaml - - configmap-litellm.yaml -tests: - - it: should not emit pgbouncer env vars by default - template: deployment.yaml - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_PGBOUNCER_ENABLED - value: "true" - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS - value: "20" - - - it: should enable the pool with the default sizing when connectionPool.enabled is set - template: deployment.yaml - set: - db.connectionPool.enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_PGBOUNCER_ENABLED - value: "true" - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS - value: "20" - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN - value: "1000" - - - it: should pass custom sizing through as strings next to the worker count - template: deployment.yaml - set: - numWorkers: 4 - db.connectionPool.enabled: true - db.connectionPool.maxDbConnections: 8 - db.connectionPool.maxClientConn: 400 - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS - value: "8" - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN - value: "400" - - contains: - path: spec.template.spec.containers[0].args - content: "4" - - - it: should give the collector sidecar the same pool env as the proxy container - template: deployment.yaml - set: - collector.enabled: true - db.connectionPool.enabled: true - db.connectionPool.maxDbConnections: 8 - db.connectionPool.maxClientConn: 400 - asserts: - - equal: - path: spec.template.spec.containers[1].name - value: litellm-collector - - contains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_PGBOUNCER_ENABLED - value: "true" - - contains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS - value: "8" - - contains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN - value: "400" - - - it: should give the collector sidecar no pool env when the pool is off - template: deployment.yaml - set: - collector.enabled: true - asserts: - - equal: - path: spec.template.spec.containers[1].name - value: litellm-collector - - notContains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_PGBOUNCER_ENABLED - any: true - - notContains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS - any: true - - notContains: - path: spec.template.spec.containers[1].env - content: - name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN - any: true diff --git a/helm/litellm-helm/tests/coordination_redis_tests.yaml b/helm/litellm-helm/tests/coordination_redis_tests.yaml deleted file mode 100644 index 0b58b1e6bc8..00000000000 --- a/helm/litellm-helm/tests/coordination_redis_tests.yaml +++ /dev/null @@ -1,143 +0,0 @@ -suite: test coordination redis -templates: - - configmap-litellm.yaml - - deployment.yaml -tests: - - it: should not render coordination_redis when redis is disabled - template: configmap-litellm.yaml - set: - redis.enabled: false - asserts: - - notMatchRegex: - path: data["config.yaml"] - pattern: coordination_redis - - - it: should not emit redis env vars when redis is disabled - template: deployment.yaml - set: - redis.enabled: false - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: REDIS_HOST - value: RELEASE-NAME-redis-master - any: true - - - it: should render coordination_redis pointing at the bundled redis when enabled - template: configmap-litellm.yaml - set: - redis.enabled: true - asserts: - - matchRegex: - path: data["config.yaml"] - pattern: "coordination_redis:\n host: os.environ/REDIS_HOST\n password: os.environ/REDIS_PASSWORD\n port: os.environ/REDIS_PORT\n" - - matchRegex: - path: data["config.yaml"] - pattern: "master_key: os.environ/PROXY_MASTER_KEY" - - - it: should emit redis env vars backing the coordination_redis os.environ refs - template: deployment.yaml - set: - redis.enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: REDIS_HOST - value: RELEASE-NAME-redis-master - - contains: - path: spec.template.spec.containers[0].env - content: - name: REDIS_PORT - value: "6379" - - contains: - path: spec.template.spec.containers[0].env - content: - name: REDIS_PASSWORD - valueFrom: - secretKeyRef: - name: RELEASE-NAME-redis - key: redis-password - - - it: should not render coordination_redis when coordination is opted out - template: configmap-litellm.yaml - set: - redis.enabled: true - redis.coordination.enabled: false - asserts: - - notMatchRegex: - path: data["config.yaml"] - pattern: coordination_redis - - - it: should keep emitting redis env vars when coordination is opted out - template: deployment.yaml - set: - redis.enabled: true - redis.coordination.enabled: false - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: REDIS_HOST - value: RELEASE-NAME-redis-master - - - it: should not clobber a user supplied coordination_redis block - template: configmap-litellm.yaml - set: - redis.enabled: true - proxy_config.general_settings.coordination_redis: - url: os.environ/COORDINATION_REDIS_URL - asserts: - - matchRegex: - path: data["config.yaml"] - pattern: "coordination_redis:\n url: os.environ/COORDINATION_REDIS_URL\n" - - notMatchRegex: - path: data["config.yaml"] - pattern: "host: os.environ/REDIS_HOST" - - - it: should render sentinel_nodes and service_name in sentinel mode - template: configmap-litellm.yaml - set: - redis.enabled: true - redis.architecture: replication - redis.sentinel.enabled: true - asserts: - # The sentinel Service the redis subchart renders is "-redis", and a - # plain client cannot speak the sentinel protocol, so host/port must not appear - - matchRegex: - path: data["config.yaml"] - pattern: "coordination_redis:\n password: os.environ/REDIS_PASSWORD\n sentinel_nodes:\n - - RELEASE-NAME-redis\n - 26379\n service_name: mymaster\n" - - notMatchRegex: - path: data["config.yaml"] - pattern: "host: os.environ/REDIS_HOST" - - - it: should carry a custom sentinel masterSet into service_name - template: configmap-litellm.yaml - set: - redis.enabled: true - redis.architecture: replication - redis.sentinel.enabled: true - redis.sentinel.masterSet: litellm-master - asserts: - - matchRegex: - path: data["config.yaml"] - pattern: "service_name: litellm-master" - - - it: should point REDIS_HOST at the sentinel service in sentinel mode - template: deployment.yaml - set: - redis.enabled: true - redis.architecture: replication - redis.sentinel.enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: REDIS_HOST - value: RELEASE-NAME-redis - - contains: - path: spec.template.spec.containers[0].env - content: - name: REDIS_PORT - value: "26379" diff --git a/helm/litellm-helm/tests/database_auth_tests.yaml b/helm/litellm-helm/tests/database_auth_tests.yaml deleted file mode 100644 index 2c480196dfb..00000000000 --- a/helm/litellm-helm/tests/database_auth_tests.yaml +++ /dev/null @@ -1,288 +0,0 @@ -suite: test existing database token auth -templates: - - deployment.yaml - - migrations-job.yaml - - configmap-litellm.yaml -tests: - - it: should keep password auth from the secret when no token auth is set - template: deployment.yaml - set: - db.useExisting: true - db.deployStandalone: false - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_USERNAME - valueFrom: - secretKeyRef: - name: postgres - key: username - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_PASSWORD - valueFrom: - secretKeyRef: - name: postgres - key: password - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL - value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME) - - notContains: - path: spec.template.spec.containers[0].env - content: - name: IAM_TOKEN_DB_AUTH - any: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: AZURE_POSTGRESQL_AUTH - any: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_PORT - any: true - - - it: should render RDS IAM auth without a password secret on the proxy and the collector sidecar - template: deployment.yaml - set: - db.useExisting: true - db.deployStandalone: false - db.useIAMAuth: true - db.username: litellm_iam - db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com - db.port: 5433 - collector.enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: IAM_TOKEN_DB_AUTH - value: "true" - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_USERNAME - value: litellm_iam - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_HOST - value: writer.cluster-abc.us-west-1.rds.amazonaws.com - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_PORT - value: "5433" - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_NAME - value: litellm - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_PASSWORD - any: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL - any: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: AZURE_POSTGRESQL_AUTH - any: true - - contains: - path: spec.template.spec.containers[1].env - content: - name: IAM_TOKEN_DB_AUTH - value: "true" - - notContains: - path: spec.template.spec.containers[1].env - content: - name: DATABASE_PASSWORD - any: true - - - it: should give the migrations job the same token auth env as the proxy - template: migrations-job.yaml - set: - migrationJob.enabled: true - db.useExisting: true - db.deployStandalone: false - db.useIAMAuth: true - db.username: litellm_iam - db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: IAM_TOKEN_DB_AUTH - value: "true" - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_USERNAME - value: litellm_iam - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_HOST - value: writer.cluster-abc.us-west-1.rds.amazonaws.com - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_PORT - value: "5432" - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_PASSWORD - any: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL - any: true - - - it: should carry db.schema as DATABASE_SCHEMA under token auth and omit it when empty - template: deployment.yaml - set: - db.useExisting: true - db.deployStandalone: false - db.useIAMAuth: true - db.username: litellm_iam - db.schema: tenant1 - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_SCHEMA - value: tenant1 - - - it: should leave the schema to db.url under password auth even when db.schema is set - template: deployment.yaml - set: - db.useExisting: true - db.deployStandalone: false - db.schema: tenant1 - db.url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1 - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_SCHEMA - any: true - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL - value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1 - - - it: should not emit DATABASE_SCHEMA when db.schema is empty - template: deployment.yaml - set: - db.useExisting: true - db.deployStandalone: false - db.useIAMAuth: true - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_SCHEMA - any: true - - - it: should read the user from the secret when db.username is empty under token auth - template: deployment.yaml - set: - db.useExisting: true - db.deployStandalone: false - db.useIAMAuth: true - db.secret.name: rds-identity - db.secret.usernameKey: user - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_USERNAME - valueFrom: - secretKeyRef: - name: rds-identity - key: user - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_PASSWORD - any: true - - - it: should render Entra auth with AZURE_POSTGRESQL_AUTH and no password - template: deployment.yaml - set: - db.useExisting: true - db.deployStandalone: false - db.useAzureEntraAuth: true - db.username: litellm@example.onmicrosoft.com - db.endpoint: litellm.postgres.database.azure.com - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: AZURE_POSTGRESQL_AUTH - value: "true" - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_USERNAME - value: litellm@example.onmicrosoft.com - - notContains: - path: spec.template.spec.containers[0].env - content: - name: IAM_TOKEN_DB_AUTH - any: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_PASSWORD - any: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL - any: true - - - it: should fail when both token sources are enabled - template: deployment.yaml - set: - migrationJob.enabled: false - db.useExisting: true - db.deployStandalone: false - db.useIAMAuth: true - db.useAzureEntraAuth: true - asserts: - - failedTemplate: - errorMessage: "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" - - - it: should fail when token auth is set while the bundled postgres is still deployed - template: deployment.yaml - set: - migrationJob.enabled: false - db.useExisting: true - db.useIAMAuth: true - asserts: - - failedTemplate: - errorPattern: "need db.useExisting: true and db.deployStandalone: false" - - - it: should fail when token auth is set without useExisting - template: migrations-job.yaml - set: - migrationJob.enabled: true - db.deployStandalone: false - db.useAzureEntraAuth: true - asserts: - - failedTemplate: - errorPattern: "need db.useExisting: true and db.deployStandalone: false" diff --git a/helm/litellm-helm/tests/deployment_command_args_labels_tests.yaml b/helm/litellm-helm/tests/deployment_command_args_labels_tests.yaml deleted file mode 100644 index 6b0d45ebf48..00000000000 --- a/helm/litellm-helm/tests/deployment_command_args_labels_tests.yaml +++ /dev/null @@ -1,68 +0,0 @@ -suite: test deployment command, args, and deploymentLabels -templates: - - deployment.yaml - - configmap-litellm.yaml -tests: - - it: should override args when custom args specified - template: deployment.yaml - set: - args: - - --custom-arg1 - - value1 - - --custom-arg2 - asserts: - - equal: - path: spec.template.spec.containers[0].args - value: - - --custom-arg1 - - value1 - - --custom-arg2 - - it: should set custom command when specified - template: deployment.yaml - set: - command: - - /bin/sh - - -c - asserts: - - equal: - path: spec.template.spec.containers[0].command - value: - - /bin/sh - - -c - - it: should set custom command and args together - template: deployment.yaml - set: - command: - - python - - -u - args: - - my_script.py - - --verbose - asserts: - - equal: - path: spec.template.spec.containers[0].command - value: - - python - - -u - - equal: - path: spec.template.spec.containers[0].args - value: - - my_script.py - - --verbose - - it: should add deploymentLabels to deployment metadata - template: deployment.yaml - set: - deploymentLabels: - environment: production - team: platform - version: v1.2.3 - asserts: - - equal: - path: metadata.labels.environment - value: production - - equal: - path: metadata.labels.team - value: platform - - equal: - path: metadata.labels.version - value: v1.2.3 diff --git a/helm/litellm-helm/tests/deployment_tests.yaml b/helm/litellm-helm/tests/deployment_tests.yaml deleted file mode 100644 index 292eeadfd13..00000000000 --- a/helm/litellm-helm/tests/deployment_tests.yaml +++ /dev/null @@ -1,520 +0,0 @@ -suite: test deployment -templates: - - deployment.yaml - - configmap-litellm.yaml -tests: - - it: should work - template: deployment.yaml - set: - image.tag: test - asserts: - - isKind: - of: Deployment - - matchRegex: - path: metadata.name - pattern: -litellm$ - - equal: - path: spec.template.spec.containers[0].image - value: ghcr.io/berriai/litellm:test - - it: should work with tolerations - template: deployment.yaml - set: - tolerations: - - key: node-role.kubernetes.io/master - operator: Exists - effect: NoSchedule - asserts: - - equal: - path: spec.template.spec.tolerations[0].key - value: node-role.kubernetes.io/master - - equal: - path: spec.template.spec.tolerations[0].operator - value: Exists - - it: should work with affinity - template: deployment.yaml - set: - affinity: - nodeAffinity: - requiredDuringSchedulingIgnoredDuringExecution: - nodeSelectorTerms: - - matchExpressions: - - key: topology.kubernetes.io/zone - operator: In - values: - - antarctica-east1 - asserts: - - equal: - path: spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].key - value: topology.kubernetes.io/zone - - equal: - path: spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].operator - value: In - - equal: - path: spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions[0].values[0] - value: antarctica-east1 - - it: should work without masterkeySecretName or masterkeySecretKey - template: deployment.yaml - set: - masterkeySecretName: "" - masterkeySecretKey: "" - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: PROXY_MASTER_KEY - valueFrom: - secretKeyRef: - name: RELEASE-NAME-litellm-masterkey - key: masterkey - - it: should work with masterkeySecretName and masterkeySecretKey - template: deployment.yaml - set: - masterkeySecretName: my-secret - masterkeySecretKey: my-key - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: PROXY_MASTER_KEY - valueFrom: - secretKeyRef: - name: my-secret - key: my-key - - it: should inject DATABASE_READER_HOST from readReplicaEndpointKey before DATABASE_URL_READ_REPLICA - template: deployment.yaml - set: - db: - deployStandalone: false - useExisting: true - secret: - name: postgres - usernameKey: username - passwordKey: password - readReplicaEndpointKey: reader-host - readReplicaUrl: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)?sslmode=require - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_READER_HOST - valueFrom: - secretKeyRef: - name: postgres - key: reader-host - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL_READ_REPLICA - value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)?sslmode=require - # $(VAR) interpolation only resolves vars defined EARLIER in the env - # array, so the reader host must precede the composed URL - - equal: - path: spec.template.spec.containers[0].env[7].name - value: DATABASE_READER_HOST - - equal: - path: spec.template.spec.containers[0].env[8].name - value: DATABASE_URL_READ_REPLICA - - it: should omit reader host when readReplicaUrl is unset - template: deployment.yaml - set: - db: - deployStandalone: false - useExisting: true - secret: - name: postgres - usernameKey: username - passwordKey: password - readReplicaEndpointKey: reader-host - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_READER_HOST - valueFrom: - secretKeyRef: - name: postgres - key: reader-host - - it: should prefer readReplicaUrlKey over readReplicaEndpointKey composition - template: deployment.yaml - set: - db: - useExisting: true - secret: - name: postgres - usernameKey: username - passwordKey: password - readReplicaUrlKey: reader-url - readReplicaEndpointKey: reader-host - readReplicaUrl: postgresql://ignored - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL_READ_REPLICA - valueFrom: - secretKeyRef: - name: postgres - key: reader-url - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL_READ_REPLICA - value: postgresql://ignored - # the unused reader-host secret ref must be suppressed so a missing - # key can't fail pod creation - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_READER_HOST - valueFrom: - secretKeyRef: - name: postgres - key: reader-host - - it: should work with extraEnvVars - template: deployment.yaml - set: - extraEnvVars: - - name: EXTRA_ENV_VAR - valueFrom: - fieldRef: - fieldPath: metadata.labels['env'] - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: EXTRA_ENV_VAR - valueFrom: - fieldRef: - fieldPath: metadata.labels['env'] - - it: should work with both extraEnvVars and envVars - template: deployment.yaml - set: - envVars: - ENV_VAR: ENV_VAR_VALUE - extraEnvVars: - - name: EXTRA_ENV_VAR - value: EXTRA_ENV_VAR_VALUE - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: ENV_VAR - value: ENV_VAR_VALUE - - contains: - path: spec.template.spec.containers[0].env - content: - name: EXTRA_ENV_VAR - value: EXTRA_ENV_VAR_VALUE - - it: should mount existing configmap when create=false - template: deployment.yaml - set: - proxyConfigMap: - create: false - name: my-litellm-config - key: custom.yaml - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: litellm-config - configMap: - name: my-litellm-config - items: - - key: custom.yaml - path: config.yaml - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: litellm-config - mountPath: /etc/litellm/config.yaml - subPath: config.yaml - - it: should work with lifecycle hooks - template: deployment.yaml - set: - lifecycle: - preStop: - exec: - command: - - /bin/sh - - -c - - echo "Container stopping" - asserts: - - exists: - path: spec.template.spec.containers[0].lifecycle - - equal: - path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[0] - value: /bin/sh - - equal: - path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[1] - value: -c - - equal: - path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[2] - value: echo "Container stopping" - - it: should render background health check settings from proxy_config.general_settings - template: configmap-litellm.yaml - set: - proxy_config.general_settings.background_health_checks: true - proxy_config.general_settings.health_check_interval: 240 - proxy_config.general_settings.health_check_concurrency: 16 - proxy_config.general_settings.health_check_details: false - asserts: - - matchRegex: - path: data["config.yaml"] - pattern: '(?m)^\s*background_health_checks:\s*true$' - - matchRegex: - path: data["config.yaml"] - pattern: '(?m)^\s*health_check_interval:\s*240$' - - matchRegex: - path: data["config.yaml"] - pattern: '(?m)^\s*health_check_concurrency:\s*16$' - - matchRegex: - path: data["config.yaml"] - pattern: '(?m)^\s*health_check_details:\s*false$' - - it: should not render any placeholder models when the user sets proxy_config without model_list - template: configmap-litellm.yaml - set: - proxy_config.general_settings.store_model_in_db: true - asserts: - - notMatchRegex: - path: data["config.yaml"] - pattern: 'model_name:' - - matchRegex: - path: data["config.yaml"] - pattern: 'store_model_in_db: true' - - matchRegex: - path: data["config.yaml"] - pattern: 'master_key: os.environ/PROXY_MASTER_KEY' - - it: should render a user-supplied model_list as the only models - template: configmap-litellm.yaml - set: - proxy_config.model_list: - - model_name: my-model - litellm_params: - model: openai/my-model - api_key: os.environ/OPENAI_API_KEY - asserts: - - matchRegex: - path: data["config.yaml"] - pattern: 'model_list:\n- litellm_params:\n api_key: os.environ/OPENAI_API_KEY\n model: openai/my-model\n model_name: my-model\n$' - - it: should allow overriding liveness, readiness, and startup probes - template: deployment.yaml - set: - livenessProbe: - path: /custom/livez - initialDelaySeconds: 5 - periodSeconds: 15 - timeoutSeconds: 5 - successThreshold: 1 - failureThreshold: 5 - readinessProbe: - path: /custom/readyz - initialDelaySeconds: 10 - periodSeconds: 20 - timeoutSeconds: 6 - successThreshold: 1 - failureThreshold: 6 - startupProbe: - path: /custom/startupz - initialDelaySeconds: 15 - periodSeconds: 25 - timeoutSeconds: 7 - successThreshold: 1 - failureThreshold: 40 - asserts: - - equal: - path: spec.template.spec.containers[0].livenessProbe.httpGet.path - value: /custom/livez - - equal: - path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds - value: 5 - - equal: - path: spec.template.spec.containers[0].readinessProbe.httpGet.path - value: /custom/readyz - - equal: - path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds - value: 6 - - equal: - path: spec.template.spec.containers[0].startupProbe.httpGet.path - value: /custom/startupz - - equal: - path: spec.template.spec.containers[0].startupProbe.failureThreshold - value: 40 - - it: should render container resources from values - template: deployment.yaml - set: - resources: - limits: - cpu: 500m - memory: 2Gi - requests: - cpu: 250m - memory: 1Gi - asserts: - - equal: - path: spec.template.spec.containers[0].resources.limits.cpu - value: 500m - - equal: - path: spec.template.spec.containers[0].resources.limits.memory - value: 2Gi - - equal: - path: spec.template.spec.containers[0].resources.requests.cpu - value: 250m - - equal: - path: spec.template.spec.containers[0].resources.requests.memory - value: 1Gi - - it: should keep default probes and empty resources unchanged - template: deployment.yaml - asserts: - - equal: - path: spec.template.spec.containers[0].livenessProbe.httpGet.path - value: /health/liveliness - - equal: - path: spec.template.spec.containers[0].livenessProbe.initialDelaySeconds - value: 0 - - equal: - path: spec.template.spec.containers[0].livenessProbe.periodSeconds - value: 15 - - equal: - path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds - value: 5 - - equal: - path: spec.template.spec.containers[0].livenessProbe.successThreshold - value: 1 - - equal: - path: spec.template.spec.containers[0].livenessProbe.failureThreshold - value: 5 - - equal: - path: spec.template.spec.containers[0].readinessProbe.httpGet.path - value: /health/readiness - - equal: - path: spec.template.spec.containers[0].readinessProbe.initialDelaySeconds - value: 0 - - equal: - path: spec.template.spec.containers[0].readinessProbe.periodSeconds - value: 10 - - equal: - path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds - value: 5 - - equal: - path: spec.template.spec.containers[0].readinessProbe.successThreshold - value: 1 - - equal: - path: spec.template.spec.containers[0].readinessProbe.failureThreshold - value: 3 - - equal: - path: spec.template.spec.containers[0].startupProbe.httpGet.path - value: /health/readiness - - equal: - path: spec.template.spec.containers[0].startupProbe.initialDelaySeconds - value: 0 - - equal: - path: spec.template.spec.containers[0].startupProbe.periodSeconds - value: 10 - - equal: - path: spec.template.spec.containers[0].startupProbe.timeoutSeconds - value: 5 - - equal: - path: spec.template.spec.containers[0].startupProbe.successThreshold - value: 1 - - equal: - path: spec.template.spec.containers[0].startupProbe.failureThreshold - value: 30 - - equal: - path: spec.template.spec.containers[0].resources - value: {} - - it: should be able to set minReadySeconds - template: deployment.yaml - set: - deploymentMinReadySeconds: 5 - asserts: - - equal: - path: spec.minReadySeconds - value: 5 - - it: should have minReadySeconds absent when deploymentMinReadySeconds is not set - template: deployment.yaml - asserts: - - notExists: - path: spec.minReadySeconds - - it: should work with extraInitContainers - template: deployment.yaml - set: - extraInitContainers: - - name: init-test - image: busybox:latest - command: ["echo", "hello"] - asserts: - - contains: - path: spec.template.spec.initContainers - content: - name: init-test - image: busybox:latest - command: ["echo", "hello"] - - it: should support tpl in extraInitContainers - template: deployment.yaml - set: - image: - repository: ghcr.io/berriai/litellm - tag: test - extraInitContainers: - - name: init-tpl - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" - command: ["echo", "hello"] - asserts: - - contains: - path: spec.template.spec.initContainers - content: - name: init-tpl - image: "ghcr.io/berriai/litellm:test" - command: ["echo", "hello"] - - it: should work with extraContainers - template: deployment.yaml - set: - extraContainers: - - name: sidecar - image: busybox:latest - asserts: - - contains: - path: spec.template.spec.containers - content: - name: sidecar - image: busybox:latest - - it: should support tpl in extraContainers - template: deployment.yaml - set: - image: - repository: ghcr.io/berriai/litellm - tag: test - extraContainers: - - name: sidecar-tpl - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" - asserts: - - contains: - path: spec.template.spec.containers - content: - name: sidecar-tpl - image: "ghcr.io/berriai/litellm:test" - - it: should support tpl in podAnnotations - template: deployment.yaml - set: - image: - repository: ghcr.io/berriai/litellm - tag: test - # Mirrors the real-world scenario this feature unblocks: - # user disables the built-in ConfigMap (and its built-in checksum/config - # annotation) and re-implements checksum/config themselves via tpl. - proxyConfigMap: - create: false - podAnnotations: - checksum/config: "{{ .Values.image.tag }}" - example.com/some-key: "{{ .Values.image.repository }}" - example.com/literal: "plain-string-value" - asserts: - - equal: - path: spec.template.metadata.annotations["checksum/config"] - value: "test" - - equal: - path: spec.template.metadata.annotations["example.com/some-key"] - value: "ghcr.io/berriai/litellm" - - equal: - path: spec.template.metadata.annotations["example.com/literal"] - value: "plain-string-value" diff --git a/helm/litellm-helm/tests/hpa_tests.yaml b/helm/litellm-helm/tests/hpa_tests.yaml deleted file mode 100644 index e446f58c8fe..00000000000 --- a/helm/litellm-helm/tests/hpa_tests.yaml +++ /dev/null @@ -1,144 +0,0 @@ -suite: "hpa" -templates: - - hpa.yaml -tests: - - it: "renders behavior when set" - set: - autoscaling.enabled: true - autoscaling.behavior: - scaleUp: - stabilizationWindowSeconds: 60 - policies: - - type: Pods - value: 2 - periodSeconds: 60 - scaleDown: - stabilizationWindowSeconds: 90 - policies: - - type: Pods - value: 1 - periodSeconds: 60 - asserts: - - isKind: { of: HorizontalPodAutoscaler } - - equal: { path: spec.behavior.scaleUp.stabilizationWindowSeconds, value: 60 } - - equal: { path: spec.behavior.scaleDown.stabilizationWindowSeconds, value: 90 } - - - it: "does not render behavior when not set" - set: - autoscaling.enabled: true - asserts: - - isKind: { of: HorizontalPodAutoscaler } - - isNull: { path: spec.behavior } - - - it: "scales on cpu at the documented 60 percent by default" - set: - autoscaling.enabled: true - asserts: - - isKind: { of: HorizontalPodAutoscaler } - - equal: { path: "spec.metrics[0].resource.name", value: cpu } - - equal: { path: "spec.metrics[0].resource.target.type", value: Utilization } - - equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 60 } - - - it: "does not scale on memory by default" - set: - autoscaling.enabled: true - asserts: - - lengthEqual: { path: spec.metrics, count: 1 } - - - it: "honours an explicit cpu target override" - set: - autoscaling.enabled: true - autoscaling.targetCPUUtilizationPercentage: 75 - asserts: - - equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 75 } - - - it: "renders a memory metric only when a memory target is set" - set: - autoscaling.enabled: true - autoscaling.targetMemoryUtilizationPercentage: 80 - asserts: - - lengthEqual: { path: spec.metrics, count: 2 } - - equal: { path: "spec.metrics[1].resource.name", value: memory } - - equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 } - - - it: "renders no workload metrics by default" - set: - autoscaling.enabled: true - autoscaling.targetMemoryUtilizationPercentage: 80 - asserts: - - lengthEqual: { path: spec.metrics, count: 2 } - - notContains: { path: spec.metrics, content: { type: Pods }, any: true } - - - it: "adds a requests-per-second Pods metric after the cpu metric" - set: - autoscaling.enabled: true - autoscaling.targetRequestsPerSecond: 90 - asserts: - - lengthEqual: { path: spec.metrics, count: 2 } - - equal: { path: "spec.metrics[0].resource.name", value: cpu } - - equal: - path: "spec.metrics[1]" - value: - type: Pods - pods: - metric: { name: litellm_requests_per_second } - target: { type: AverageValue, averageValue: "90" } - - - it: "adds a tokens-per-second Pods metric on its own" - set: - autoscaling.enabled: true - autoscaling.targetTokensPerSecond: 6M - asserts: - - lengthEqual: { path: spec.metrics, count: 2 } - - equal: - path: "spec.metrics[1]" - value: - type: Pods - pods: - metric: { name: litellm_tokens_per_second } - target: { type: AverageValue, averageValue: "6M" } - - notContains: - path: spec.metrics - content: { type: Pods, pods: { metric: { name: litellm_requests_per_second } } } - any: true - - - it: "renders requests, tokens, cpu and memory metrics together" - set: - autoscaling.enabled: true - autoscaling.targetMemoryUtilizationPercentage: 80 - autoscaling.targetRequestsPerSecond: 90 - autoscaling.targetTokensPerSecond: 6000000 - asserts: - - lengthEqual: { path: spec.metrics, count: 4 } - - equal: { path: "spec.metrics[0].resource.name", value: cpu } - - equal: { path: "spec.metrics[1].resource.name", value: memory } - - equal: { path: "spec.metrics[2].pods.metric.name", value: litellm_requests_per_second } - - equal: { path: "spec.metrics[2].pods.target.averageValue", value: "90" } - - equal: { path: "spec.metrics[3].pods.metric.name", value: litellm_tokens_per_second } - - equal: { path: "spec.metrics[3].pods.target.averageValue", value: "6000000" } - - - it: "scales on workload metrics alone when the cpu target is cleared" - set: - autoscaling.enabled: true - autoscaling.targetCPUUtilizationPercentage: null - autoscaling.targetRequestsPerSecond: 90 - autoscaling.targetTokensPerSecond: 6000000 - asserts: - - lengthEqual: { path: spec.metrics, count: 2 } - - notContains: { path: spec.metrics, content: { type: Resource }, any: true } - - equal: { path: "spec.metrics[0].pods.metric.name", value: litellm_requests_per_second } - - equal: { path: "spec.metrics[1].pods.metric.name", value: litellm_tokens_per_second } - - notMatchRegexRaw: { pattern: per_minute } - - - it: "ignores the per-minute keys, which the chart never shipped" - set: - autoscaling.enabled: true - autoscaling.targetRequestsPerMinute: 5400 - autoscaling.targetTokensPerMinute: 360000000 - asserts: - - lengthEqual: { path: spec.metrics, count: 1 } - - notContains: { path: spec.metrics, content: { type: Pods }, any: true } - - - it: "renders no hpa when autoscaling is disabled" - asserts: - - hasDocuments: { count: 0 } diff --git a/helm/litellm-helm/tests/ingress_tests.yaml b/helm/litellm-helm/tests/ingress_tests.yaml deleted file mode 100644 index aad6ecfcee8..00000000000 --- a/helm/litellm-helm/tests/ingress_tests.yaml +++ /dev/null @@ -1,45 +0,0 @@ -suite: Ingress Configuration Tests -templates: - - ingress.yaml -tests: - - it: should not create Ingress by default - asserts: - - hasDocuments: - count: 0 - - - it: should create Ingress when enabled - set: - ingress.enabled: true - asserts: - - hasDocuments: - count: 1 - - isKind: - of: Ingress - - - it: should add custom labels - set: - ingress.enabled: true - ingress.labels: - custom-label: "true" - another-label: "value" - asserts: - - isKind: - of: Ingress - - equal: - path: metadata.labels.custom-label - value: "true" - - equal: - path: metadata.labels.another-label - value: "value" - - - it: should add annotations - set: - ingress.enabled: true - ingress.annotations: - kubernetes.io/ingress.class: "nginx" - asserts: - - isKind: - of: Ingress - - equal: - path: metadata.annotations["kubernetes.io/ingress.class"] - value: "nginx" diff --git a/helm/litellm-helm/tests/keda_tests.yaml b/helm/litellm-helm/tests/keda_tests.yaml deleted file mode 100644 index c9598646223..00000000000 --- a/helm/litellm-helm/tests/keda_tests.yaml +++ /dev/null @@ -1,106 +0,0 @@ -suite: "keda" -templates: - - keda.yaml -release: - name: rel - namespace: llm -tests: - - it: "renders no scaled object by default" - asserts: - - hasDocuments: { count: 0 } - - - it: "passes user triggers through and adds no prometheus triggers by default" - set: - keda.enabled: true - keda.triggers: - - type: cpu - metricType: Utilization - metadata: { value: "60" } - asserts: - - isKind: { of: ScaledObject } - - equal: - path: spec.triggers - value: - - type: cpu - metricType: Utilization - metadata: { value: "60" } - - - it: "scales on release-wide requests per second divided by the per-replica target" - set: - keda.enabled: true - keda.prometheus.serverAddress: http://prometheus-operated.monitoring.svc:9090 - keda.prometheus.requestsPerSecond: 90 - asserts: - - lengthEqual: { path: spec.triggers, count: 1 } - - equal: - path: "spec.triggers[0]" - value: - type: prometheus - metadata: - serverAddress: http://prometheus-operated.monitoring.svc:9090 - threshold: "90" - query: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm"}[1m])) - - - it: "scales on tokens per second on its own" - set: - keda.enabled: true - keda.prometheus.serverAddress: http://prom:9090 - keda.prometheus.tokensPerSecond: 6000000 - asserts: - - lengthEqual: { path: spec.triggers, count: 1 } - - equal: { path: "spec.triggers[0].type", value: prometheus } - - equal: { path: "spec.triggers[0].metadata.threshold", value: "6000000" } - - equal: - path: "spec.triggers[0].metadata.query" - value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm"}[1m])) - - - it: "appends requests and tokens triggers after user triggers and selects the metrics service job" - set: - keda.enabled: true - metricsServer.enabled: true - keda.triggers: - - type: cpu - metricType: Utilization - metadata: { value: "60" } - keda.prometheus.serverAddress: http://prom:9090 - keda.prometheus.requestsPerSecond: 90 - keda.prometheus.tokensPerSecond: 6000000 - asserts: - - lengthEqual: { path: spec.triggers, count: 3 } - - equal: { path: "spec.triggers[0].type", value: cpu } - - equal: { path: "spec.triggers[1].metadata.threshold", value: "90" } - - equal: - path: "spec.triggers[1].metadata.query" - value: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m])) - - equal: { path: "spec.triggers[2].metadata.threshold", value: "6000000" } - - equal: - path: "spec.triggers[2].metadata.query" - value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m])) - - notMatchRegexRaw: { pattern: "\\* *60|per_minute|PerMinute" } - - - it: "ignores the per-minute keys, which the chart never shipped" - set: - keda.enabled: true - keda.prometheus.serverAddress: http://prom:9090 - keda.prometheus.requestsPerMinute: 5400 - keda.prometheus.tokensPerMinute: 360000000 - asserts: - - isKind: { of: ScaledObject } - - isNullOrEmpty: { path: spec.triggers } - - - it: "refuses a workload target without a prometheus server address" - set: - keda.enabled: true - keda.prometheus.requestsPerSecond: 90 - asserts: - - failedTemplate: - errorMessage: keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set - - - it: "yields to the hpa when both autoscalers are enabled" - set: - autoscaling.enabled: true - keda.enabled: true - keda.prometheus.serverAddress: http://prom:9090 - keda.prometheus.requestsPerSecond: 90 - asserts: - - hasDocuments: { count: 0 } diff --git a/helm/litellm-helm/tests/lens_connection_tests.yaml b/helm/litellm-helm/tests/lens_connection_tests.yaml deleted file mode 100644 index e01e49d0ffc..00000000000 --- a/helm/litellm-helm/tests/lens_connection_tests.yaml +++ /dev/null @@ -1,128 +0,0 @@ -suite: lens connection -templates: - - deployment.yaml - - configmap-litellm.yaml - - ingress.yaml -tests: - - it: connects the proxy to an external Lens - template: deployment.yaml - set: - lens.enabled: true - lens.url: http://lens.lens.svc:4318 - lens.publicUrl: https://lens.example.com/lens-ingest - lens.gatewaySecret.name: lens-gateway - lens.serviceTokenSecret.name: lens-service - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_LENS_URL - value: http://lens.lens.svc:4318 - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_LENS_PUBLIC_URL - value: https://lens.example.com/lens-ingest - - contains: - path: spec.template.spec.containers[0].env - content: - name: LENS_GATEWAY_SECRET - valueFrom: - secretKeyRef: - name: lens-gateway - key: gateway-secret - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_LENS_SERVICE_TOKEN - valueFrom: - secretKeyRef: - name: lens-service - key: service-token - - - it: requires lens.url when Lens is enabled - template: deployment.yaml - set: - lens.enabled: true - lens.publicUrl: https://lens.example.com/lens-ingest - lens.gatewaySecret.name: lens-gateway - lens.serviceTokenSecret.name: lens-service - asserts: - - failedTemplate: - errorMessage: lens.url is required when lens.enabled is true - - - it: requires lens.publicUrl when Lens is enabled - template: deployment.yaml - set: - lens.enabled: true - lens.url: http://lens.lens.svc:4318 - lens.gatewaySecret.name: lens-gateway - lens.serviceTokenSecret.name: lens-service - asserts: - - failedTemplate: - errorMessage: lens.publicUrl is required when lens.enabled is true - - - it: requires lens.gatewaySecret.name when Lens is enabled - template: deployment.yaml - set: - lens.enabled: true - lens.url: http://lens.lens.svc:4318 - lens.publicUrl: https://lens.example.com/lens-ingest - lens.serviceTokenSecret.name: lens-service - asserts: - - failedTemplate: - errorMessage: lens.gatewaySecret.name is required when lens.enabled is true - - - it: requires lens.serviceTokenSecret.name when Lens is enabled - template: deployment.yaml - set: - lens.enabled: true - lens.url: http://lens.lens.svc:4318 - lens.publicUrl: https://lens.example.com/lens-ingest - lens.gatewaySecret.name: lens-gateway - asserts: - - failedTemplate: - errorMessage: lens.serviceTokenSecret.name is required when lens.enabled is true - - - it: routes /lens-ingest to the named Lens service ahead of the host paths - template: ingress.yaml - set: - ingress.enabled: true - lens.enabled: true - lens.url: http://lens.lens.svc:4318 - lens.publicUrl: https://lens.example.com/lens-ingest - lens.gatewaySecret.name: lens-gateway - lens.serviceTokenSecret.name: lens-service - lens.service.name: existing-lens - asserts: - - equal: - path: spec.rules[0].http.paths[0].path - value: /lens-ingest - - contains: - path: spec.rules[0].http.paths - content: - path: /lens-ingest - pathType: Prefix - backend: - service: - name: existing-lens - port: - number: 4318 - - - it: leaves /lens-ingest off the ingress without a Lens service name - template: ingress.yaml - set: - ingress.enabled: true - lens.enabled: true - lens.url: http://lens.lens.svc:4318 - lens.publicUrl: https://lens.example.com/lens-ingest - lens.gatewaySecret.name: lens-gateway - lens.serviceTokenSecret.name: lens-service - asserts: - - exists: - path: spec.rules[0].http.paths - - notContains: - path: spec.rules[0].http.paths - content: - path: /lens-ingest - any: true diff --git a/helm/litellm-helm/tests/masterkey-secret_tests.yaml b/helm/litellm-helm/tests/masterkey-secret_tests.yaml deleted file mode 100644 index 296f26755b8..00000000000 --- a/helm/litellm-helm/tests/masterkey-secret_tests.yaml +++ /dev/null @@ -1,71 +0,0 @@ -suite: test masterkey secret -templates: - - secret-masterkey.yaml -tests: - - it: should create a secret if masterkeySecretName is not set. should start with sk-xxxx (base64 encoded as c2st*) - template: secret-masterkey.yaml - set: - masterkeySecretName: "" - asserts: - - isKind: - of: Secret - - matchRegex: - path: data.masterkey - pattern: ^c2st - # Note: The masterkey is generated as "sk-<18-random-chars>" in plain text, - # but stored as base64 encoded in Kubernetes secret (requirement). - # "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern. - - it: should reuse the master key already stored in the cluster instead of generating a new one on upgrade - template: secret-masterkey.yaml - set: - masterkeySecretName: "" - kubernetesProvider: - scheme: - "v1/Secret": - gvr: - version: "v1" - resource: "secrets" - namespaced: true - objects: - - kind: Secret - apiVersion: v1 - metadata: - name: RELEASE-NAME-litellm-masterkey - namespace: NAMESPACE - data: - masterkey: c2stZXhpc3Rpbmcta2V5 - asserts: - - equal: - path: data.masterkey - value: c2stZXhpc3Rpbmcta2V5 - - it: should let an explicit masterkey value override the one already stored in the cluster - template: secret-masterkey.yaml - set: - masterkeySecretName: "" - masterkey: sk-explicit - kubernetesProvider: - scheme: - "v1/Secret": - gvr: - version: "v1" - resource: "secrets" - namespaced: true - objects: - - kind: Secret - apiVersion: v1 - metadata: - name: RELEASE-NAME-litellm-masterkey - namespace: NAMESPACE - data: - masterkey: c2stZXhpc3Rpbmcta2V5 - asserts: - - equal: - path: data.masterkey - value: c2stZXhwbGljaXQ= - - it: should not create a secret if masterkeySecretName is set - template: secret-masterkey.yaml - set: - masterkeySecretName: my-secret - asserts: - - hasDocuments: - count: 0 diff --git a/helm/litellm-helm/tests/metrics_server_tests.yaml b/helm/litellm-helm/tests/metrics_server_tests.yaml deleted file mode 100644 index 085d69ac640..00000000000 --- a/helm/litellm-helm/tests/metrics_server_tests.yaml +++ /dev/null @@ -1,106 +0,0 @@ -suite: separate metrics server -templates: - - configmap-litellm.yaml - - deployment.yaml - - service.yaml - - service-metrics.yaml - - servicemonitor.yaml -tests: - - it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default - asserts: - - notContains: - path: spec.template.spec.containers[0].ports - content: - name: metrics - any: true - template: deployment.yaml - - notContains: - path: spec.template.spec.containers[0].env - content: - name: PROMETHEUS_METRICS_PORT - any: true - template: deployment.yaml - - lengthEqual: - path: spec.ports - count: 1 - template: service.yaml - - hasDocuments: - count: 0 - template: service-metrics.yaml - - - it: should scrape the proxy port when the metrics server is disabled - template: servicemonitor.yaml - set: - serviceMonitor.enabled: true - asserts: - - equal: - path: spec.endpoints[0].port - value: http - - - it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor - set: - metricsServer.enabled: true - metricsServer.port: 4101 - serviceMonitor.enabled: true - service.type: LoadBalancer - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: PROMETHEUS_METRICS_PORT - value: "4101" - template: deployment.yaml - - contains: - path: spec.template.spec.containers[0].ports - content: - name: metrics - containerPort: 4101 - protocol: TCP - template: deployment.yaml - - lengthEqual: - path: spec.ports - count: 1 - template: service.yaml - - equal: - path: spec.type - value: LoadBalancer - template: service.yaml - - equal: - path: metadata.name - value: RELEASE-NAME-litellm-metrics - template: service-metrics.yaml - - equal: - path: spec.type - value: ClusterIP - template: service-metrics.yaml - - equal: - path: spec.ports - value: - - port: 4101 - targetPort: metrics - protocol: TCP - name: metrics - template: service-metrics.yaml - - equal: - path: spec.selector - value: - app.kubernetes.io/name: litellm - app.kubernetes.io/instance: RELEASE-NAME - template: service-metrics.yaml - - equal: - path: spec.endpoints[0].port - value: metrics - template: servicemonitor.yaml - - equal: - path: spec.endpoints[0].path - value: /metrics/ - template: servicemonitor.yaml - - - it: should reject a metrics port equal to the proxy port - template: deployment.yaml - set: - metricsServer.enabled: true - metricsServer.port: 4000 - asserts: - - failedTemplate: - errorMessage: metricsServer.port must differ from service.port diff --git a/helm/litellm-helm/tests/migrations-job_tests.yaml b/helm/litellm-helm/tests/migrations-job_tests.yaml deleted file mode 100644 index bd5dd457f95..00000000000 --- a/helm/litellm-helm/tests/migrations-job_tests.yaml +++ /dev/null @@ -1,432 +0,0 @@ -suite: test migrations job -templates: - - migrations-job.yaml -tests: - - it: should work with envVars - template: migrations-job.yaml - set: - envVars: - TEST_ENV_VAR: "test_value" - ANOTHER_VAR: "another_value" - migrationJob: - enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: TEST_ENV_VAR - value: "test_value" - - contains: - path: spec.template.spec.containers[0].env - content: - name: ANOTHER_VAR - value: "another_value" - - - it: should work with extraEnvVars - template: migrations-job.yaml - set: - extraEnvVars: - - name: EXTRA_ENV_VAR - valueFrom: - fieldRef: - fieldPath: metadata.labels['env'] - - name: SIMPLE_EXTRA_VAR - value: "simple_value" - migrationJob: - enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: EXTRA_ENV_VAR - valueFrom: - fieldRef: - fieldPath: metadata.labels['env'] - - contains: - path: spec.template.spec.containers[0].env - content: - name: SIMPLE_EXTRA_VAR - value: "simple_value" - - - it: should work with both envVars and extraEnvVars - template: migrations-job.yaml - set: - envVars: - ENV_VAR: "env_var_value" - extraEnvVars: - - name: EXTRA_ENV_VAR - value: "extra_env_var_value" - migrationJob: - enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: ENV_VAR - value: "env_var_value" - - contains: - path: spec.template.spec.containers[0].env - content: - name: EXTRA_ENV_VAR - value: "extra_env_var_value" - - - it: should not render when migrations job is disabled - template: migrations-job.yaml - set: - migrationJob: - enabled: false - asserts: - - hasDocuments: - count: 0 - - - it: should still include default env vars - template: migrations-job.yaml - set: - envVars: - CUSTOM_VAR: "custom_value" - migrationJob: - enabled: true - db: - useExisting: true - endpoint: "test-db" - database: "testdb" - url: "postgresql://user:pass@test-db:5432/testdb" - secret: - name: "test-secret" - usernameKey: "username" - passwordKey: "password" - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: DISABLE_SCHEMA_UPDATE - value: "false" - - contains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_HOST - value: "test-db" - - contains: - path: spec.template.spec.containers[0].env - content: - name: CUSTOM_VAR - value: "custom_value" - - - it: should override a user-supplied DISABLE_SCHEMA_UPDATE so the Job always migrates - template: migrations-job.yaml - set: - envVars: - DISABLE_SCHEMA_UPDATE: "true" - migrationJob: - enabled: true - asserts: - # The Job is what owns the schema, so it renders its own - # DISABLE_SCHEMA_UPDATE=false after envVars and extraEnvVars. Kubernetes - # takes the last value for a duplicated name, so the user's "true" cannot - # leave the schema unmigrated. Skipping migrations is migrationJob.enabled. - - equal: - path: spec.template.spec.containers[0].env[-1] - value: - name: DISABLE_SCHEMA_UPDATE - value: "false" - - - it: should not include DATABASE_URL when deployStandalone is false - template: migrations-job.yaml - set: - migrationJob: - enabled: true - db: - deployStandalone: false - useExisting: false - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: DATABASE_URL - - - it: should use default service account for helm hooks when serviceAccount.create is true - template: migrations-job.yaml - set: - migrationJob: - enabled: true - hooks: - helm: - enabled: true - serviceAccount: - create: true - asserts: - - equal: - path: spec.template.spec.serviceAccountName - value: default - - - it: should use migrationJob.serviceAccountName override for helm hooks when serviceAccount.create is true - template: migrations-job.yaml - set: - migrationJob: - enabled: true - serviceAccountName: migration-sa - hooks: - helm: - enabled: true - serviceAccount: - create: true - asserts: - - equal: - path: spec.template.spec.serviceAccountName - value: migration-sa - - - it: should use chart service account when helm hooks are disabled - template: migrations-job.yaml - set: - migrationJob: - enabled: true - hooks: - helm: - enabled: false - serviceAccount: - create: true - name: my-custom-sa - asserts: - - equal: - path: spec.template.spec.serviceAccountName - value: my-custom-sa - - - it: should use pre-existing service account when helm hooks are enabled but serviceAccount.create is false - template: migrations-job.yaml - set: - migrationJob: - enabled: true - hooks: - helm: - enabled: true - serviceAccount: - create: false - name: pre-existing-sa - asserts: - - equal: - path: spec.template.spec.serviceAccountName - value: pre-existing-sa - - it: should work with extraInitContainers - template: migrations-job.yaml - set: - migrationJob: - enabled: true - extraInitContainers: - - name: init-test - image: busybox:latest - command: ["echo", "hello"] - asserts: - - contains: - path: spec.template.spec.initContainers - content: - name: init-test - image: busybox:latest - command: ["echo", "hello"] - - it: should support tpl in extraInitContainers - template: migrations-job.yaml - set: - image: - repository: ghcr.io/berriai/litellm - tag: test - migrationJob: - enabled: true - extraInitContainers: - - name: init-tpl - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" - command: ["echo", "hello"] - asserts: - - contains: - path: spec.template.spec.initContainers - content: - name: init-tpl - image: "ghcr.io/berriai/litellm:test" - command: ["echo", "hello"] - - it: should work with extraContainers - template: migrations-job.yaml - set: - migrationJob: - enabled: true - extraContainers: - - name: sidecar - image: busybox:latest - asserts: - - contains: - path: spec.template.spec.containers - content: - name: sidecar - image: busybox:latest - - it: should support tpl in extraContainers - template: migrations-job.yaml - set: - image: - repository: ghcr.io/berriai/litellm - tag: test - migrationJob: - enabled: true - extraContainers: - - name: sidecar-tpl - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" - asserts: - - contains: - path: spec.template.spec.containers - content: - name: sidecar-tpl - image: "ghcr.io/berriai/litellm:test" - - it: should render the pod-level securityContext from podSecurityContext - template: migrations-job.yaml - set: - migrationJob: - enabled: true - podSecurityContext: - fsGroup: 10000 - runAsUser: 10000 - runAsNonRoot: true - asserts: - - equal: - path: spec.template.spec.securityContext - value: - fsGroup: 10000 - runAsUser: 10000 - runAsNonRoot: true - - it: should keep the pod-level and container-level securityContext separate - template: migrations-job.yaml - set: - migrationJob: - enabled: true - podSecurityContext: - fsGroup: 10000 - securityContext: - allowPrivilegeEscalation: false - readOnlyRootFilesystem: true - asserts: - - equal: - path: spec.template.spec.securityContext - value: - fsGroup: 10000 - - equal: - path: spec.template.spec.containers[0].securityContext - value: - allowPrivilegeEscalation: false - readOnlyRootFilesystem: true - - it: should schedule onto the same nodes as the gateway - template: migrations-job.yaml - set: - migrationJob: - enabled: true - nodeSelector: - karpenter.sh/nodepool: litellm-e2e - tolerations: - - key: workload - operator: Equal - value: litellm-e2e - effect: NoSchedule - asserts: - - equal: - path: spec.template.spec.nodeSelector - value: - karpenter.sh/nodepool: litellm-e2e - - equal: - path: spec.template.spec.tolerations - value: - - key: workload - operator: Equal - value: litellm-e2e - effect: NoSchedule - - - it: bounds the Job with a deadline by default, so a blocked migration cannot stall the release forever - set: - migrationJob: - enabled: true - asserts: - - equal: - path: spec.activeDeadlineSeconds - value: 1800 - - - it: honours an operator-supplied deadline - set: - migrationJob: - enabled: true - activeDeadlineSeconds: 600 - asserts: - - equal: - path: spec.activeDeadlineSeconds - value: 600 - - - it: omits the deadline entirely when it is nulled out, restoring the unbounded behaviour - set: - migrationJob: - enabled: true - activeDeadlineSeconds: null - asserts: - - notExists: - path: spec.activeDeadlineSeconds - - - it: should set custom jobLabels and jobAnnotations on Job metadata - template: migrations-job.yaml - set: - migrationJob: - enabled: true - jobLabels: - environment: production - team: platform - jobAnnotations: - example.com/cost-center: "1234" - asserts: - - equal: - path: metadata.labels.environment - value: production - - equal: - path: metadata.labels.team - value: platform - - equal: - path: metadata.annotations['example.com/cost-center'] - value: "1234" - - - it: should set custom podLabels on Pod template - template: migrations-job.yaml - set: - migrationJob: - enabled: true - podLabels: - custom.io/pod-role: migration - asserts: - - equal: - path: spec.template.metadata.labels['custom.io/pod-role'] - value: migration - - - it: should override container command and args - template: migrations-job.yaml - set: - migrationJob: - enabled: true - command: - - sh - args: - - -c - - echo migrating - asserts: - - equal: - path: spec.template.spec.containers[0].command - value: - - sh - - equal: - path: spec.template.spec.containers[0].args - value: - - -c - - echo migrating - - - it: should clear container args when only command is specified - template: migrations-job.yaml - set: - migrationJob: - enabled: true - command: - - sh - asserts: - - equal: - path: spec.template.spec.containers[0].command - value: - - sh - - equal: - path: spec.template.spec.containers[0].args - value: [] diff --git a/helm/litellm-helm/tests/pdb_tests.yaml b/helm/litellm-helm/tests/pdb_tests.yaml deleted file mode 100644 index 5e042e80bd3..00000000000 --- a/helm/litellm-helm/tests/pdb_tests.yaml +++ /dev/null @@ -1,45 +0,0 @@ -suite: "pdb enabled" -templates: - - poddisruptionbudget.yaml -tests: - - it: "renders a PDB with maxUnavailable=1" - set: - pdb.enabled: true - pdb.maxUnavailable: 1 - asserts: - - hasDocuments: { count: 1 } - - isKind: { of: PodDisruptionBudget } - - equal: { path: apiVersion, value: policy/v1 } - - equal: { path: spec.maxUnavailable, value: 1 } - - equal: - path: spec.selector.matchLabels - value: - app.kubernetes.io/name: litellm - app.kubernetes.io/instance: RELEASE-NAME - ---- -suite: "pdb disabled" -templates: - - poddisruptionbudget.yaml -tests: - - it: "does not render when disabled" - set: - pdb.enabled: false - asserts: - - hasDocuments: { count: 0 } - ---- -suite: "pdb minAvailable precedence" -templates: - - poddisruptionbudget.yaml -tests: - - it: "uses minAvailable when both are set" - set: - pdb.enabled: true - pdb.minAvailable: "50%" - pdb.maxUnavailable: 1 - asserts: - - isKind: { of: PodDisruptionBudget } - - equal: { path: apiVersion, value: policy/v1 } - - equal: { path: spec.minAvailable, value: "50%" } - - isNull: { path: spec.maxUnavailable } diff --git a/helm/litellm-helm/tests/service_tests.yaml b/helm/litellm-helm/tests/service_tests.yaml deleted file mode 100644 index 43ed0180bc8..00000000000 --- a/helm/litellm-helm/tests/service_tests.yaml +++ /dev/null @@ -1,116 +0,0 @@ -suite: Service Configuration Tests -templates: - - service.yaml -tests: - - it: should create a default ClusterIP service - template: service.yaml - asserts: - - isKind: - of: Service - - equal: - path: spec.type - value: ClusterIP - - equal: - path: spec.ports[0].port - value: 4000 - - equal: - path: spec.ports[0].targetPort - value: http - - equal: - path: spec.ports[0].protocol - value: TCP - - equal: - path: spec.ports[0].name - value: http - - isNull: - path: spec.loadBalancerClass - - - it: should create a NodePort service when specified - template: service.yaml - set: - service.type: NodePort - asserts: - - isKind: - of: Service - - equal: - path: spec.type - value: NodePort - - isNull: - path: spec.loadBalancerClass - - - it: should create a LoadBalancer service when specified - template: service.yaml - set: - service.type: LoadBalancer - asserts: - - isKind: - of: Service - - equal: - path: spec.type - value: LoadBalancer - - isNull: - path: spec.loadBalancerClass - - - it: should add loadBalancerClass when specified with LoadBalancer type - template: service.yaml - set: - service.type: LoadBalancer - service.loadBalancerClass: tailscale - asserts: - - isKind: - of: Service - - equal: - path: spec.type - value: LoadBalancer - - equal: - path: spec.loadBalancerClass - value: tailscale - - - it: should not add loadBalancerClass when specified with ClusterIP type - template: service.yaml - set: - service.type: ClusterIP - service.loadBalancerClass: tailscale - asserts: - - isKind: - of: Service - - equal: - path: spec.type - value: ClusterIP - - isNull: - path: spec.loadBalancerClass - - - it: should use custom port when specified - template: service.yaml - set: - service.port: 8080 - asserts: - - equal: - path: spec.ports[0].port - value: 8080 - - - it: should add service annotations when specified - template: service.yaml - set: - service.annotations: - cloud.google.com/load-balancer-type: "Internal" - service.beta.kubernetes.io/aws-load-balancer-internal: "true" - asserts: - - isKind: - of: Service - - equal: - path: metadata.annotations - value: - cloud.google.com/load-balancer-type: "Internal" - service.beta.kubernetes.io/aws-load-balancer-internal: "true" - - - it: should use the correct selector labels - template: service.yaml - asserts: - - isNotNull: - path: spec.selector - - equal: - path: spec.selector - value: - app.kubernetes.io/name: litellm - app.kubernetes.io/instance: RELEASE-NAME diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml deleted file mode 100644 index f51f8bdd173..00000000000 --- a/helm/litellm-helm/values.yaml +++ /dev/null @@ -1,682 +0,0 @@ -# Default values for litellm. -# This is a YAML-formatted file. -# Declare variables to be passed into your templates. - -replicaCount: 1 -liteadmin: - enabled: false - existingSecret: "" - gatewayUrl: "" - model: "" - readOnly: false - storageSize: 1Gi - storageClassName: "" - resources: - requests: - cpu: 100m - memory: 256Mi - limits: - memory: 1Gi -# numWorkers: 2 - -image: - # Bundles the prisma CLI and engines, which is what lets the migrations job - # and the proxy's own schema check run without network access. - repository: ghcr.io/berriai/litellm - pullPolicy: Always - # Overrides the image tag whose default is the chart appVersion. - # tag: "latest" - tag: "" - -imagePullSecrets: [] -nameOverride: "litellm" -fullnameOverride: "" - -serviceAccount: - # Specifies whether a service account should be created - create: false - # Automatically mount a ServiceAccount's API credentials? - automount: true - # Annotations to add to the service account - annotations: {} - # The name of the service account to use. - # If not set and create is true, a name is generated using the fullname template - name: "" - -# annotations for litellm deployment -deploymentAnnotations: {} -deploymentLabels: {} -deploymentMinReadySeconds: 0 - -# annotations for litellm pods -podAnnotations: {} -podLabels: {} - -# -- Deployment strategy configuration -# Example: -# type: RollingUpdate -# rollingUpdate: -# maxUnavailable: 0 -# maxSurge: 1 -strategy: {} - -terminationGracePeriodSeconds: 90 -topologySpreadConstraints: - [] - # - maxSkew: 1 - # topologyKey: kubernetes.io/hostname - # whenUnsatisfiable: DoNotSchedule - # labelSelector: - # matchLabels: - # app: litellm - -# At the time of writing, the litellm docker image requires write access to the -# filesystem on startup so that prisma can install some dependencies. -podSecurityContext: {} -securityContext: - {} - # capabilities: - # drop: - # - ALL - # readOnlyRootFilesystem: false - # runAsNonRoot: true - # runAsUser: 1000 - -# A list of Kubernetes Secret objects that will be exported to the LiteLLM proxy -# pod as environment variables. These secrets can then be referenced in the -# configuration file (or "litellm" ConfigMap) with `os.environ/` -environmentSecrets: - [] - # - litellm-env-secret - -# A list of Kubernetes ConfigMap objects that will be exported to the LiteLLM proxy -# pod as environment variables. The ConfigMap kv-pairs can then be referenced in the -# configuration file (or "litellm" ConfigMap) with `os.environ/` -environmentConfigMaps: - [] - # - litellm-env-configmap - -service: - type: ClusterIP - port: 4000 - # If service type is `LoadBalancer` you can - # optionally specify loadBalancerClass - # loadBalancerClass: tailscale - -# Probes for LiteLLM gateway container -livenessProbe: - path: /health/liveliness - initialDelaySeconds: 0 - periodSeconds: 15 - timeoutSeconds: 5 - successThreshold: 1 - failureThreshold: 5 - -readinessProbe: - path: /health/readiness - initialDelaySeconds: 0 - periodSeconds: 10 - timeoutSeconds: 5 - successThreshold: 1 - failureThreshold: 3 - -startupProbe: - path: /health/readiness - initialDelaySeconds: 0 - periodSeconds: 10 - timeoutSeconds: 5 - successThreshold: 1 - failureThreshold: 30 - -ingress: - enabled: false - className: "nginx" - labels: {} - annotations: - {} - # kubernetes.io/ingress.class: nginx - # kubernetes.io/tls-acme: "true" - hosts: - - host: api.example.local - paths: - - path: / - pathType: ImplementationSpecific - tls: [] - # - secretName: chart-example-tls - # hosts: - # - chart-example.local - -# masterkey: changeit - -# if set, use this secret for the master key; otherwise, autogenerate a new one -masterkeySecretName: "" - -# if set, use this secret key for the master key; otherwise, use the default key -masterkeySecretKey: "" - -# Optional: enterprise billable-request metering. When enabled, the proxy counts -# successful requests to inference, MCP, and A2A endpoints and pushes them to -# LiteLLM's collector over mutual TLS. Requires an enterprise license. -# The client certificate identifies the deployment, so it is mounted read-only -# from an existing Secret and never passed through the environment. -billingMetrics: - enabled: false - endpoint: https://telemetry.litellm.ai # collector to push the counter to - secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key - # Only for private or test collectors whose server certificate is not on the - # public web PKI. The production collector needs no CA override. - caSecretName: "" # existing Secret holding ca.crt - exportIntervalMs: "" # push cadence; the proxy defaults to 60000 - -proxyConfigMap: - # when true, creates a new configmap - create: true - # if create is false and name is set, use existing ConfigMap - # create: false - # name: "" - # key: "config.yaml" - -# The elements within proxy_config are rendered as config.yaml for the proxy -# Examples: https://github.com/BerriAI/litellm/tree/main/litellm/proxy/example_config_yaml -# Reference: https://docs.litellm.ai/docs/proxy/configs -proxy_config: - model_list: [] - general_settings: - master_key: os.environ/PROXY_MASTER_KEY - -# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT) -# so a scrape never runs on an inference worker. Adds a `metrics` port to the -# container and a dedicated ClusterIP `-metrics` Service, and the -# ServiceMonitor scrapes it instead of the proxy port. The separate port has -# no virtual-key auth: keep it off public ingress. Needs the proxy image -# v1.101.0 or newer. -metricsServer: - enabled: false - port: 4001 - -# Opt-in sidecar that runs the post-response spend pipeline (cost calculation, -# spend logs, spend counters, budget reservation reconciliation) so the proxy's -# uvicorn workers only serialise a compact typed event and go back to serving -# inference. Same image and tag as the proxy, second container in the same pod, -# fed over loopback (a unix socket on a shared emptyDir, or 127.0.0.1 TCP). It -# reuses the pod's in-container pgbouncer (db.connectionPool) and the same Redis -# spend transaction buffer, so the per-pod DB connection budget is unchanged. -# Delivery is at-most-once inside the pod: events already handed to the sidecar -# are lost if it crashes before writing them; events the workers could not hand -# over follow onUnavailable. Both containers drain on SIGTERM within -# terminationGracePeriodSeconds -collector: - enabled: false - # unix:////.sock (the becomes a shared emptyDir) or tcp://127.0.0.1: - address: unix:///var/run/litellm/collector.sock - # Events each uvicorn worker holds in memory while the sidecar is slow or restarting - bufferSize: 1000 - # fallback: run the pipeline in the worker when the sidecar is unreachable or the - # buffer is full (spend stays exact, that request costs proxy CPU again) - # drop: count and discard the event instead (spend under-reports) - onUnavailable: fallback - # How long the workers keep pushing buffered events on shutdown, and how long the - # sidecar keeps serving its open connections after SIGTERM - drainTimeoutSeconds: 10 - command: - - python - - -m - - litellm.proxy.collector - # Sized independently of the proxy container; the pipeline is CPU bound - resources: {} - # requests: - # cpu: 500m - # memory: 1Gi - # limits: - # cpu: "1" - # memory: 2Gi - # When autoscaling.enabled, swap the pod-wide cpu Resource metric for an - # autoscaling/v2 ContainerResource metric on the proxy container only, so the - # sidecar's CPU never scales inference replicas. Needs Kubernetes 1.30+ (or the - # HPAContainerMetrics feature gate on 1.27 to 1.29) - scaleOnProxyContainerCpu: false - -resources: - {} - # Unset by default so the chart installs on small clusters such as Minikube, and so an - # upgrade never leaves a running pod Pending. Production deployments should set these. - # A proxy at DB-connected steady state needs about 1 CPU and 4Gi of memory per worker; - # sizing below that gets the pod OOMKilled once traffic and DB connections ramp up. - # Scale both figures with --num_workers, then uncomment the lines below and remove the - # curly braces after 'resources:'. See "Recommended Machine Specifications" in - # https://docs.litellm.ai/docs/proxy/prod. - # requests: - # cpu: "1" - # memory: 4Gi - # limits: - # cpu: "1" - # memory: 4Gi - -autoscaling: - enabled: false - minReplicas: 1 - maxReplicas: 100 - # 60 is the documented recommendation. See "Recommended Machine Specifications" - # in https://docs.litellm.ai/docs/proxy/prod. A new replica clears the startupProbe - # above only after up to failureThreshold x periodSeconds = 300 seconds, so a target - # high enough to trip near saturation adds capacity minutes after it was needed. - targetCPUUtilizationPercentage: 60 - # Deliberately left unset rather than given a value. The prisma query engine's - # resident memory is a high-water mark that ratchets to the pod's worst-ever write - # and is never returned, so a memory target reads the largest write a pod ever did - # rather than what it is doing now, and replicas ratchet up without scaling back in. - # Memory is a floor to provision under 'resources', not a signal to scale on. - # targetMemoryUtilizationPercentage: 80 - # behavior: {} - # Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics - # named `litellm_requests_per_second` and `litellm_tokens_per_second` with an - # AverageValue target, alongside whichever resource targets are set (the HPA - # follows the metric asking for the most replicas). A Prometheus Adapter must - # serve those two names on custom.metrics.k8s.io from the proxy's counters, - # grouped by the scrape target's `pod` label (enable serviceMonitor below so - # every pod is scraped on its own): - # litellm_requests_per_second: - # sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>) - # litellm_tokens_per_second: - # sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>) - # rate() over [1m] is already per second, so no `* 60`. How fast the HPA - # reacts is set by that window, the scrape interval and the HPA sync period - # (15s by default), not by the unit: keep serviceMonitor.interval at 15s or - # faster so a 1m window holds at least 4 samples. averageValue takes SI - # suffixes, so "6M" is six million tokens per second per pod. Tokens are - # counted when a response completes, so TPS trails long streams. - targetRequestsPerSecond: "" - targetTokensPerSecond: "" - -# Autoscaling with keda is mutually exclusive with hpa -keda: - enabled: false - minReplicas: 1 - maxReplicas: 100 - pollingInterval: 30 - cooldownPeriod: 300 - # fallback: - # failureThreshold: 3 - # replicas: 11 - restoreToOriginalReplicaCount: false - scaledObject: - annotations: {} - triggers: [] - # - type: prometheus - # metadata: - # serverAddress: http://:9090 - # metricName: http_requests_total - # threshold: '100' - # query: sum(rate(http_requests_total{deployment="my-deployment"}[2m])) - # First-class Prometheus triggers on the proxy's own request and token - # counters, appended to `triggers`. Each target is the per-second load one - # replica should carry: KEDA divides the release-wide - # `sum(rate([1m]))` by it to pick the replica count. Thresholds - # are plain numbers (KEDA parses them as floats, no SI suffixes). The - # queries select samples by the release namespace and the `job` label the - # chart's ServiceMonitor produces (the metrics Service name), so enable - # serviceMonitor below together with metricsServer: the http port serves - # /metrics/ behind virtual-key auth and answers an unauthenticated scrape - # with 401. Reaction time comes from the [1m] window, the scrape interval - # and pollingInterval above, so keep both at 15s or faster. Tokens are - # counted at completion, so TPS trails long streams. serverAddress is - # required once either target is set. - prometheus: - serverAddress: "" - requestsPerSecond: "" - tokensPerSecond: "" - behavior: {} - # scaleDown: - # stabilizationWindowSeconds: 300 - # policies: - # - type: Pods - # value: 1 - # periodSeconds: 180 - # scaleUp: - # stabilizationWindowSeconds: 300 - # policies: - # - type: Pods - # value: 2 - # periodSeconds: 60 - -# Additional volumes on the output Deployment definition. -volumes: [] -# - name: foo -# secret: -# secretName: mysecret -# optional: false - -# Additional volumeMounts on the output Deployment definition. -volumeMounts: [] -# - name: foo -# mountPath: "/etc/foo" -# readOnly: true - -nodeSelector: {} - -tolerations: [] - -affinity: {} - -db: - # Use an existing postgres server/cluster - useExisting: false - - # How to connect to the existing postgres server/cluster - endpoint: localhost - database: litellm - url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME) - # Authenticate with a short-lived token instead of a password. useIAMAuth is - # AWS RDS / Aurora IAM database authentication (the pod needs AWS credentials, - # e.g. IRSA or EKS Pod Identity, and the database user needs GRANT rds_iam); - # useAzureEntraAuth is Azure Database for PostgreSQL with a Microsoft Entra ID - # identity (workload identity). Mutually exclusive, and both need - # useExisting: true with deployStandalone: false. The chart then sets - # IAM_TOKEN_DB_AUTH or AZURE_POSTGRESQL_AUTH, emits DATABASE_HOST, DATABASE_PORT, - # DATABASE_USERNAME and DATABASE_NAME and skips DATABASE_PASSWORD and db.url: - # the proxy, the migrations job and the collector sidecar each mint their own - # token and build DATABASE_URL from those parts. Set DATABASE_SSLMODE / - # DATABASE_SSLROOTCERT under envVars for TLS options you would otherwise put - # in db.url. - useIAMAuth: false - useAzureEntraAuth: false - # Postgres port, used to build the token-auth URL (password auth takes the - # port from db.url) - port: 5432 - # Postgres schema for the LiteLLM tables, rendered as DATABASE_SCHEMA under - # token auth only (password auth carries it as ?schema= in db.url) - schema: "" - # Database user as a plain value. Empty means read it from - # db.secret.usernameKey. With token auth the user name is the only credential - # the chart needs, so setting it here removes the need for db.secret entirely. - username: "" - secret: - name: postgres - usernameKey: username - passwordKey: password - # Optional: when set, DATABASE_HOST will be sourced from this secret key instead of db.endpoint - endpointKey: "" - # Optional: when set, DATABASE_URL_READ_REPLICA will be sourced from this - # secret key instead of db.readReplicaUrl. Prefer this over the plain - # value: read-replica URLs typically embed credentials, and a value - # written to db.readReplicaUrl ends up visible in the rendered pod spec - # and the Helm release secret. - readReplicaUrlKey: "" - # Optional: when set, a DATABASE_READER_HOST env var is sourced from this - # secret key, so db.readReplicaUrl can compose the reader URL from - # individual secret components, e.g. - # postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME) - # Use this when your secret store holds the bare reader hostname rather - # than a full connection URL. Only takes effect when readReplicaUrl is - # set; ignored when readReplicaUrlKey is set. - readReplicaEndpointKey: "" - - # Optional read-replica routing. When set, the proxy sends read-only - # queries (find_*, count, group_by, query_raw/_first) to this URL while - # writes continue to go to db.url. Useful for Aurora-style clusters with - # separate reader/writer endpoints. Leave empty to keep single-DB behavior. - # When IAM_TOKEN_DB_AUTH is enabled, the reader URL is auto-refreshed - # alongside the writer (host/port/user/db are parsed from this URL once - # at startup; only the IAM token rotates). - # - # If the URL embeds credentials, prefer db.secret.readReplicaUrlKey over - # this field — the plain value is rendered into the pod spec and the - # Helm release secret. This field is intended for credential-less URLs - # only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime). - readReplicaUrl: "" - - # In-container connection pool (PgBouncer, transaction mode) shared by every - # worker in the pod. Without it each --num_workers worker opens its own - # connection_limit connections to Postgres, so a pod's footprint against the - # database's connection ceiling is workers x connection_limit and grows with - # every replica. With it, the pod holds at most maxDbConnections upstream - # connections no matter how many workers run; the workers connect to the pool - # over loopback, with no extra network hop. Migrations still go straight to - # Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so - # a database with a 5000-connection ceiling fits roughly 200 replicas. - connectionPool: - enabled: false - maxDbConnections: 20 - maxClientConn: 1000 - - # Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster. - # The Stackgres Operator must already be installed within the target - # Kubernetes cluster. - # TODO: Stackgres deployment currently unsupported - useStackgresOperator: false - - # Use the Postgres Helm chart to create a single node, stand alone postgres - # instance. See the "postgresql" top level key for additional configuration. - deployStandalone: true - -# Lifecycle hooks for the LiteLLM container -# -# Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks -# the pod NotReady and blocks only until in-flight requests actually finish -# (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always -# waiting the worst-case duration. After SIGTERM the server waits up to -# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels -# them, so set terminationGracePeriodSeconds a few seconds above twice -# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL. -# -# /health/drain is off by default; enable it with -# general_settings.enable_drain_endpoint: true. The kubelet calls preStop -# hooks without proxy credentials, so when the health port is reachable from -# other pods (the common case) also set -# general_settings.drain_endpoint_token (or the DRAIN_ENDPOINT_TOKEN env -# var) and send the same value on the X-Drain-Token header from the hook. -# Calls missing/wrong the token get a 401 and have no side effect. -# Example: -# lifecycle: -# preStop: -# httpGet: -# path: /health/drain -# port: 4000 -# httpHeaders: -# - name: X-Drain-Token -# value: -lifecycle: {} - -# Settings for Bitnami postgresql chart (if db.deployStandalone is true, ignored -# otherwise) -# -# Bitnami retired the versioned tags under docker.io/bitnami and republished the -# archived builds under docker.io/bitnamilegacy, so the subchart's own image -# defaults no longer resolve. The repository below points at the same build the -# subchart was released with, which keeps the on-disk data directory layout -# identical for existing installs. -# -# Keep the tag pinned. docker.io/bitnami still publishes a floating `latest`, -# and starting a newer PostgreSQL major against an existing data directory -# leaves the server refusing to boot ("database files are incompatible with -# server") with no way back other than a dump taken beforehand. Crossing a major -# version is a dump-and-restore, not an image bump. The chart refuses to render -# an unpinned tag for this reason -postgresql: - architecture: standalone - image: - repository: bitnamilegacy/postgresql - tag: 16.2.0-debian-12-r6 - volumePermissions: - image: - repository: bitnamilegacy/os-shell - tag: 12-debian-12-r16 - metrics: - image: - repository: bitnamilegacy/postgres-exporter - tag: 0.15.0-debian-12-r14 - auth: - username: litellm - database: litellm - - # You should override these on the helm command line with - # `--set postgresql.auth.postgres-password=,postgresql.auth.password=` - password: NoTaGrEaTpAsSwOrD - postgres-password: NoTaGrEaTpAsSwOrD - - # A secret is created by this chart (litellm-helm) with the credentials that - # the new Postgres instance should use. - # existingSecret: "" - # secretKeys: - # userPasswordKey: password - -# Redis is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend -# tracking, and the pod lock manager. Enabling this deploys the bundled Redis -# subchart, wires REDIS_HOST / REDIS_PORT / REDIS_PASSWORD into the proxy, and -# renders a `general_settings.coordination_redis` block into the proxy config. -# -# To point at an existing Redis instead, leave `enabled: false` and pass a -# secret for REDIS_HOST, REDIS_PORT, REDIS_PASSWORD or REDIS_URL; the proxy -# falls back to those env vars for coordination. Set `cache: true` in the proxy -# config only if you also want LLM response caching, which is independent of -# coordination -# -# When `redis.sentinel.enabled` is set, the coordination block is rendered with -# `sentinel_nodes` and `service_name` (from `redis.sentinel.masterSet`) instead -# of host/port, because a plain Redis client cannot talk to the sentinel port -# -# The image repositories carry the same bitnamilegacy repoint as postgresql -# above; the versioned tags the subchart ships with are gone from -# docker.io/bitnami -redis: - enabled: false - architecture: standalone - image: - repository: bitnamilegacy/redis - tag: 7.2.4-debian-12-r9 - sentinel: - image: - repository: bitnamilegacy/redis-sentinel - tag: 7.2.4-debian-12-r7 - metrics: - image: - repository: bitnamilegacy/redis-exporter - tag: 1.58.0-debian-12-r4 - volumePermissions: - image: - repository: bitnamilegacy/os-shell - tag: 12-debian-12-r16 - sysctl: - image: - repository: bitnamilegacy/os-shell - tag: 12-debian-12-r16 - kubectl: - image: - repository: bitnamilegacy/kubectl - tag: 1.29.2-debian-12-r3 - coordination: - # Set to false to keep the bundled Redis for response caching only and leave - # `general_settings.coordination_redis` out of the rendered config. A - # `coordination_redis` block you define yourself in `proxy_config` always wins - enabled: true - -# Prisma migration job settings -migrationJob: - enabled: true # Enable or disable the schema migration Job - backoffLimit: 4 # Backoff limit for Job restarts - # Wall-clock budget for the whole Job, shared across every `backoffLimit` - # retry rather than granted per attempt. Without it a migration that blocks - # on the database never fails, and when the Helm hook is enabled the release - # waits on it forever: `helm upgrade` and any GitOps controller driving it - # stop reconciling the whole chart until someone deletes the Job by hand. - # Set to null to opt out and restore the unbounded behaviour. - activeDeadlineSeconds: 1800 - # Optional service account for the migration job. - # Only used when migrationJob.hooks.helm.enabled=true and serviceAccount.create=true. - # In that case, pre-install/pre-upgrade hooks run before normal resources, so this defaults to "default". - serviceAccountName: "" - annotations: {} - jobLabels: {} # Custom labels for the Job metadata - jobAnnotations: {} # Custom annotations for the Job metadata - podLabels: {} # Custom labels for the Job pod template - command: [] # Override container command (defaults to ["python", "litellm/proxy/prisma_migration.py"]) - args: [] # Override container args - ttlSecondsAfterFinished: 120 - resources: {} - # Unset by default. This job runs the database migration and exits, so it does not - # need the steady-state headroom the proxy does; size it from your own migration - # runs rather than from the proxy figures above. - extraContainers: [] - extraInitContainers: [] - - # Hook configuration - hooks: - argocd: - enabled: true - helm: - enabled: false - -# Log level for the litellm proxy (sets LITELLM_LOG in the deployment env). -# Rendered as a direct `env:` entry, which in Kubernetes takes precedence over -# any `envFrom:` source. If you currently source LITELLM_LOG from an -# environmentSecret or environmentConfigMap, set `logLevel: ""` here to -# disable injection — otherwise this value silently overrides your secret / -# configmap entry. -# -# Setting LITELLM_LOG inside `envVars:` below also wins: the template skips -# this injection entirely when envVars already defines LITELLM_LOG. -logLevel: INFO - -# Additional environment variables to be added to the deployment as a map of key-value pairs -envVars: {} - -# USE_DDTRACE: "true" -# Additional environment variables to be added to the deployment as a list of k8s env vars -extraEnvVars: {} - -# if you want to override the container command, you can do so here -command: {} -# if you want to override the container args, you can do so here -args: {} - -# - name: EXTRA_ENV_VAR -# value: EXTRA_ENV_VAR_VALUE -# Additional Kubernetes resources to deploy with litellm -extraResources: [] - -# - apiVersion: v1 -# kind: ConfigMap -# metadata: -# name: my-extra-config -# data: -# foo: bar -# Pod Disruption Budget -pdb: - enabled: false - # Set exactly one of the following. If both are set, minAvailable takes precedence. - minAvailable: null # e.g. "50%" or 1 - maxUnavailable: null # e.g. 1 or "20%" - annotations: {} - labels: {} - -serviceMonitor: - enabled: false - labels: - {} - # test: test - annotations: - {} - # kubernetes.io/test: test - interval: 15s - scrapeTimeout: 10s - relabelings: [] - # - targetLabel: __meta_kubernetes_pod_node_name - # replacement: $1 - # action: replace - namespaceSelector: - matchNames: [] - # - test-namespace - -lens: - enabled: false - url: "" - publicUrl: "" - service: - name: "" - port: 4318 - gatewaySecret: - name: "" - key: gateway-secret - serviceTokenSecret: - name: "" - key: service-token diff --git a/helm/litellm/.helmignore b/helm/litellm/.helmignore new file mode 100644 index 00000000000..6b8d3002acd --- /dev/null +++ b/helm/litellm/.helmignore @@ -0,0 +1,20 @@ +# Patterns to ignore when building packages. +.DS_Store +.git/ +.gitignore +.bzr/ +.bzrignore +.hg/ +.hgignore +.svn/ +*.swp +*.bak +*.tmp +*.orig +*~ +.project +.idea/ +*.tmproj +.vscode/ +/tests/ +/ci/ diff --git a/helm/litellm/Chart.yaml b/helm/litellm/Chart.yaml index e67f5790c7e..7f308a454f5 100644 --- a/helm/litellm/Chart.yaml +++ b/helm/litellm/Chart.yaml @@ -1,8 +1,8 @@ apiVersion: v2 name: litellm -description: LiteLLM componentized — gateway, UI backend, and UI as separate services +description: Deploys LiteLLM from one image, either componentized (gateway, UI backend, and UI as separate Deployments) or as a single monolith Deployment type: application -version: 0.1.0 -appVersion: "0.1.0" +version: 1.0.0 +appVersion: "1.104.0" annotations: org.opencontainers.image.source: "https://github.com/BerriAI/litellm" diff --git a/helm/litellm/README.md b/helm/litellm/README.md index 97b94754402..5ca63d90bcd 100644 --- a/helm/litellm/README.md +++ b/helm/litellm/README.md @@ -1,6 +1,51 @@ -# LiteLLM componentized chart +# litellm Helm chart -This chart deploys separate gateway, backend and UI services +Deploys [LiteLLM](https://github.com/BerriAI/litellm) from one image, `ghcr.io/berriai/litellm` (mirrored at `docker.litellm.ai/berriai/litellm`), in one of two layouts: + +- componentized (default): the `gateway` (LLM data plane, port 4000), the `backend` (management API, port 4001) and the `ui` (static dashboard, port 3000) each get their own Deployment, Service, HPA and PDB so they scale independently +- monolith: `monolith.enabled: true` renders one Deployment and Service running the full proxy (`args: [proxy, ...]`), which serves the gateway routes, the management routes and the Admin UI from a single process + +Both layouts run the same image. The image entrypoint dispatches on the first container argument (`proxy`, `gateway`, `backend`, `ui`, `migrations`, `metrics`, `collector`), so the chart only ever sets `args` and never `command` + +## Requirements + +Kubernetes 1.25+ and Helm 3.8+. Bring your own PostgreSQL (`database.writer.*`) and, optionally, Redis (`redis.*`). The proxy also needs a master key, either an existing Secret named by `masterKey.secretName` or one the chart generates with `masterKey.generate: true` and `masterKey.secretName: ""` + +## Install + +```bash +kubectl create secret generic litellm-master-key-secret --from-literal=master-key=sk-change-me +helm install litellm helm/litellm \ + --set database.writer.host=postgres.example.com \ + --set database.writer.dbname=litellm \ + --set database.writer.passwordSecret.name=litellm-db-secret +``` + +The image tag defaults to the chart's `appVersion`. Override it with `image.tag`, or pin the bytes with `image.digest`, which renders as `repository:tag@digest` + +## Monolith quickstart + +```bash +helm install litellm helm/litellm \ + --set monolith.enabled=true \ + --set masterKey.secretName="" \ + --set masterKey.generate=true \ + --set database.writer.host=postgres.example.com \ + --set database.writer.dbname=litellm \ + --set database.writer.passwordSecret.name=litellm-db-secret +kubectl port-forward svc/litellm-litellm 4000:4000 +``` + +With `monolith.enabled: true`: + +- one Deployment and one Service named `-litellm` render, running `args: [proxy, --port, 4000, --config, /app/config/config.yaml]` plus `monolith.extraArgs` +- the gateway, backend and ui Deployments, Services, HPAs, PDBs and ServiceMonitors are not rendered, whatever `gateway.enabled`, `backend.enabled` and `ui.enabled` say +- every `gateway.*` value configures the monolith pod: `config`, `numWorkers`, `resources`, probes, `securityContext`, `hpa`, `keda`, `pdb`, `metricsServer`, `collector`, `volumes`, `extraEnv`, scheduling. The monolith runs as `serviceAccounts.gateway` +- every `backend.*` and `ui.*` value is ignored +- the Ingress sends every path, built in or from `ingress.extraPaths`, to the monolith Service +- the migrations Job renders exactly as in componentized mode + +## Lens This chart does not deploy Lens. Deploy Lens from [its own chart](https://github.com/BerriAI/lens/tree/main/helm/lens), then set `lens.enabled: true`, `lens.url`, `lens.publicUrl`, `lens.gatewaySecret.name` and `lens.serviceTokenSecret.name`. The secret keys default to `gateway-secret` and `service-token`. Set `lens.service.name` when that Lens runs in this release's namespace and the chart ingress should route `/lens-ingest` to it. The old `lensWorker` values are no longer read, and enabling them fails the render @@ -22,22 +67,114 @@ backend: gatewayUrl: https://gateway.example.com ``` -## Routes to another Service +## Testing -For a shared hostname, `ingress.extraPaths` can route to a Service in the same namespace using `serviceName` and `servicePort`. The port can be a number or a named Service port. Omit `service` when using these fields - -```yaml -ingress: - enabled: true - className: alb - host: gateway.example.com - extraPaths: - - path: /moyai - pathType: Prefix - serviceName: moyai-ingress - servicePort: 8080 +```bash +helm lint helm/litellm +helm unittest -f 'tests/*.yaml' helm/litellm +helm test --logs ``` -The Service and application must already support the path prefix. The chart adds the path to the existing Ingress and leaves its built-in routes in place. It does not create the target Service or rewrite requests +## Migrating from the litellm-helm chart -For AWS ALB, configure the target Service's `alb.ingress.kubernetes.io/healthcheck-path` annotation when the application needs a different health check from LiteLLM. Existing `service` aliases (`gateway`, `backend`, `ui`) remain supported +The `litellm-helm` chart (`oci://ghcr.io/berriai/litellm-helm`) is retired; its published packages stay available for a grace period. Its flat values described one monolith Deployment, so the equivalent install here is `monolith.enabled: true` with the values moved under `gateway.*`. The per-component `gateway.image`, `backend.image`, `ui.image` and `migrations.image` blocks of earlier `helm/litellm` versions are gone too: the chart is a major bump to `1.0.0` and every container uses the top-level `image` + +| litellm-helm value | litellm value | +|---|---| +| (implicit single Deployment) | `monolith.enabled: true` | +| `image.repository` / `image.tag` / `image.pullPolicy` | `image.repository` / `image.tag` / `image.pullPolicy` (`image.digest` is new) | +| `replicaCount` | `gateway.replicaCount` | +| `args` | `monolith.extraArgs` (appended after the chart's proxy arguments). Drop any `--port` argument and its value so the proxy keeps listening on port 4000, matching the container port and health checks. Set `gateway.service.port` to change the external Service port | +| `command` | removed: the image entrypoint dispatcher must stay in place | +| `proxy_config` | `gateway.config.proxy_config` | +| `proxyConfigMap.create: false` + `proxyConfigMap.name` | Set `gateway.config.create: false`, mount your ConfigMap with `gateway.volumes` / `gateway.volumeMounts`, and pass `--config` followed by the mounted file path in `monolith.extraArgs` | +| `masterkeySecretName` / `masterkeySecretKey` | `masterKey.secretName` / `masterKey.secretKey`. The default Secret field changes from `masterkey` to `master-key`. When reusing a Secret with the old default field, explicitly set `masterKey.secretKey: masterkey` | +| `masterkeySecretName: ""` (auto generated Secret) | `masterKey.generate: true` with `masterKey.secretName: ""` (creates a new key; copy the old key first, see below, to keep existing credentials valid) | +| `db.useExisting`, `db.endpoint`, `db.database` | `database.writer.host`, `database.writer.port`, `database.writer.dbname` | +| `db.useIAMAuth` / `db.useAzureEntraAuth` | `database.writer.useIAMAuth` / `database.writer.useAzureEntraAuth` (mutually exclusive). Both token-auth modes still require a username in the Secret referenced by `database.writer.passwordSecret.name` and `usernameKey` | +| `db.schema` | `database.writer.schema` | +| `db.username` | Store the username in the Secret referenced by `database.writer.passwordSecret.name`, under its `usernameKey`. There is no literal username value in the new chart | +| `db.secret.name` / `db.secret.usernameKey` / `db.secret.passwordKey` | `database.writer.passwordSecret.name` / `usernameKey` / `passwordKey` | +| `db.secret.endpointKey` (writer host read from a Secret) | `database.writer.passwordSecret.hostKey` with `database.writer.host` left empty | +| `db.url` (custom URL template) | `database.writer.passwordSecret.urlKey` pointing at a Secret key that holds the complete writer URL, with the other `database.writer.*` connection values left empty | +| `db.readReplicaUrl` (credential-less reader URL) | `database.reader.host`, `database.reader.port`, `database.reader.dbname`, with `database.reader.passwordSecret.*` for its credentials (or `database.reader.useIAMAuth: true`) | +| `db.secret.readReplicaEndpointKey` (reader host read from a Secret) | `database.reader.passwordSecret.hostKey` with `database.reader.host` left empty | +| `db.secret.readReplicaUrlKey` (full reader URL in a Secret) | `database.reader.passwordSecret.urlKey` with the other `database.reader.*` connection values left empty | +| `db.connectionPool.*` | `database.connectionPool.*` | +| `db.deployStandalone: true` / `postgresql.*` | removed: the chart ships no PostgreSQL, point `database.writer.*` at your own | +| `redis.enabled: true` (bundled) | removed: the chart ships no Redis, point `redis.host` at your own | +| external Redis via `envVars` | `redis.host`, `redis.port`, `redis.passwordSecret.*`, `redis.cluster` | +| `envVars` / `extraEnvVars` | `gateway.extraEnv` (list of `name` / `value` or `valueFrom` entries). The retired chart also passed these to the migration Job, so copy any it needs, such as the AWS region and credentials for IAM database auth, into `migrationJob.extraEnv` too | +| `environmentSecrets` | `gateway.envSecrets` | +| `environmentConfigMaps` | `gateway.envConfigMaps` | +| `logLevel` | `gateway.logLevel` | +| `resources` | `gateway.resources` | +| `livenessProbe` / `readinessProbe` / `startupProbe` | `gateway.livenessProbe` / `gateway.readinessProbe` / `gateway.startupProbe`, written as full Kubernetes probes: the old `path: /health/readiness` becomes `httpGet: { path: /health/readiness, port: http }`, the timing fields (`initialDelaySeconds`, `periodSeconds`, `timeoutSeconds`, `successThreshold`, `failureThreshold`) carry over unchanged | +| `securityContext` / `podSecurityContext` | `gateway.securityContext` / `gateway.podSecurityContext`. With `readOnlyRootFilesystem: true` the retired chart added `emptyDir` mounts at `/tmp`, `/.cache` and `/.npm`; this chart does not, so add them through `gateway.volumes` / `gateway.volumeMounts`. The retired chart applied both to the migration Job as well; set `migrationJob.securityContext` / `migrationJob.podSecurityContext` to keep the Job under the same policy | +| `service.*` | `gateway.service.*` | +| `ingress.className`, `ingress.annotations`, `ingress.tls` | same keys (routes to the monolith Service in monolith mode) | +| `ingress.hosts[].host` | `ingress.host`, one host per release; the built-in paths are rendered for that host, so `ingress.hosts[].paths` has no equivalent. For more than one hostname, add a second Ingress through `extraResources` or point a DNS alias at the one host | +| `autoscaling.*` | `gateway.hpa.*`. Defaults change from 100 to 10 maximum replicas and from 60% to 70% CPU utilization, with an 80% memory target now enabled. To keep the old scaling behavior, set `gateway.hpa.maxReplicas: 100`, `gateway.hpa.targetCPUUtilizationPercentage: 60` and `gateway.hpa.targetMemoryUtilizationPercentage: 0` | +| `keda.*` | `gateway.keda.*` with these renames: `minReplicas` / `maxReplicas` are `minReplicaCount` / `maxReplicaCount` (the new default maximum is 10, the old one was 100), `behavior` moves to `advanced.horizontalPodAutoscalerConfig.behavior`, `restoreToOriginalReplicaCount` moves to `advanced.restoreToOriginalReplicaCount`, and `keda.prometheus.requestsPerSecond` / `tokensPerSecond` are `gateway.keda.prometheus.targetRequestsPerSecond` / `targetTokensPerSecond`. `pollingInterval`, `cooldownPeriod`, `fallback`, `triggers` and `scaledObject.annotations` (for example KEDA's pause annotations) keep their names | +| `numWorkers` | `gateway.numWorkers` | +| `pdb.*` | `gateway.pdb.*`, except `labels` and `annotations` are no longer supported. To retain custom PDB metadata, set `gateway.pdb.enabled: false` and define the PDB through `extraResources` | +| `metricsServer.*` | `gateway.metricsServer.*` | +| `serviceMonitor.*` | `gateway.serviceMonitor.*` (`labels`, `annotations`, `interval`, `scrapeTimeout`, `relabelings`, `namespaceSelector.matchNames`); it now requires `gateway.metricsServer.enabled` | +| `collector.*` | `gateway.collector.*`, except `command` is no longer supported and `scaleOnProxyContainerCpu` becomes `scaleOnGatewayContainerCpu`. Its default changes from `false` to `true`. Set `gateway.collector.scaleOnGatewayContainerCpu: false` to keep pod-wide CPU scaling. For a custom collector command, disable the built-in sidecar with `gateway.collector.enabled: false` and define it through `gateway.extraContainers`, with its environment, volumes and volume mounts configured explicitly | +| `billingMetrics.*` | `billingMetrics.*` | +| `migrationJob.*` | `migrationJob.*` | +| `volumes` / `volumeMounts` | `gateway.volumes` / `gateway.volumeMounts`. The retired chart also mounted these in the migration Job; copy any mounts the Job needs, such as a database CA certificate, to `migrationJob.volumes` / `migrationJob.volumeMounts` too | +| `extraContainers` / `extraInitContainers` | `gateway.extraContainers` / `gateway.extraInitContainers` | +| `lifecycle` | `gateway.lifecycle` | +| `strategy` | `gateway.strategy` | +| `deploymentAnnotations` / `deploymentLabels` / `deploymentMinReadySeconds` | `gateway.deploymentAnnotations` / `gateway.deploymentLabels` / `gateway.minReadySeconds` | +| `podAnnotations` / `podLabels` | `gateway.podAnnotations` / `gateway.podLabels`. The retired chart ran `podAnnotations` through `tpl`; this chart renders them verbatim, so a copied template such as `checksum/config: {{ ... }}` stays literal text. Resolve it before install instead, e.g. `--set-string gateway.podAnnotations.checksum/config=$(kubectl get configmap my-config -o yaml \| sha256sum \| cut -d' ' -f1)`. The chart's own proxy ConfigMap already rolls the pods when `gateway.config` changes. The retired chart also put `podLabels` on the migration Job pod; copy any the Job needs, such as a network-policy selector, to `migrationJob.podLabels` | +| `nodeSelector` / `tolerations` / `affinity` / `topologySpreadConstraints` | `gateway.nodeSelector` / `gateway.tolerations` / `gateway.affinity` / `gateway.topologySpreadConstraints`. The retired chart also applied `nodeSelector`, `tolerations` and `affinity` to the migration Job; copy them to `migrationJob.nodeSelector` / `migrationJob.tolerations` / `migrationJob.affinity` too, or a Job that needs tainted or zoned nodes to reach the database stays Pending and blocks the install | +| `terminationGracePeriodSeconds` | `gateway.terminationGracePeriodSeconds` | +| `serviceAccount.*` | `serviceAccounts.gateway.*`. The retired chart also shared an existing ServiceAccount with migrations; set `migrationJob.serviceAccountName` to that existing account too when the Job needs its permissions, such as for IAM database auth | +| `lens.*` | `lens.*`. Deploy Lens from its own chart. The removed `lensWorker.*` values fail the render when enabled | +| `extraResources` | `extraResources` | +| `nameOverride: "litellm"` | `nameOverride: ""` (the chart name is already `litellm`) | +| `fullnameOverride` | `fullnameOverride`. The new default always includes the chart-name suffix, even when the release name already contains it. Set this to the existing Service name to retain its DNS name | + +A minimal migration: + +```yaml +monolith: + enabled: true +masterKey: + secretName: litellm-master-key-secret +database: + writer: + host: postgres.example.com + dbname: litellm + passwordSecret: + name: litellm-db-secret + usernameKey: username + passwordKey: password +gateway: + replicaCount: 2 + config: + proxy_config: + model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + envSecrets: + - litellm-provider-keys +``` + +The monolith Service is named `-litellm`. To retain a different old Service name, set `fullnameOverride` to that existing name. The generated Secret in this chart has `helm.sh/resource-policy: keep` and is reused across upgrades of this chart. The retired chart stored its generated key under a different Secret and data key, and `helm uninstall` deletes that Secret, so copy it into a Secret you own before uninstalling the old release to keep existing credentials valid. Run the following in the old release's namespace, replacing `OLD_MASTER_KEY_SECRET` with the actual existing Secret name. For a release named `litellm` without a `fullnameOverride`, that name is `litellm-masterkey`: + +```bash +set -euo pipefail +OLD_MASTER_KEY_SECRET=litellm-masterkey +old_master_key="$(kubectl get secret "$OLD_MASTER_KEY_SECRET" -o jsonpath='{.data.masterkey}' | base64 -d)" +test -n "$old_master_key" +kubectl create secret generic litellm-master-key-secret \ + --from-literal=master-key="$old_master_key" +unset old_master_key +``` + +Then set `masterKey.secretName: litellm-master-key-secret`, as in the minimal example above. Releases that set `masterkeySecretName` already point `masterKey.secretName` and `masterKey.secretKey` at that Secret diff --git a/helm/litellm/ci/postgres.yaml b/helm/litellm/ci/postgres.yaml new file mode 100644 index 00000000000..aa210e932b0 --- /dev/null +++ b/helm/litellm/ci/postgres.yaml @@ -0,0 +1,58 @@ +# Throwaway PostgreSQL for the CircleCI helm_chart_testing job. Applied with +# kubectl before `helm install`; the chart itself ships no database. +apiVersion: v1 +kind: Secret +metadata: + name: litellm-ci-postgres +type: Opaque +stringData: + username: litellm + password: litellm-ci +--- +apiVersion: v1 +kind: Service +metadata: + name: litellm-ci-postgres +spec: + selector: + app: litellm-ci-postgres + ports: + - port: 5432 + targetPort: 5432 +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: litellm-ci-postgres +spec: + replicas: 1 + selector: + matchLabels: + app: litellm-ci-postgres + template: + metadata: + labels: + app: litellm-ci-postgres + spec: + containers: + - name: postgres + image: postgres:16.10-alpine@sha256:029660641a0cfc575b14f336ba448fb8a75fd595d42e1fa316b9fb4378742297 + ports: + - containerPort: 5432 + env: + - name: POSTGRES_DB + value: litellm + - name: POSTGRES_USER + valueFrom: + secretKeyRef: + name: litellm-ci-postgres + key: username + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: litellm-ci-postgres + key: password + readinessProbe: + exec: + command: [pg_isready, -U, litellm, -d, litellm] + periodSeconds: 2 diff --git a/helm/litellm/ci/test-values.yaml b/helm/litellm/ci/test-values.yaml new file mode 100644 index 00000000000..40c68489ed3 --- /dev/null +++ b/helm/litellm/ci/test-values.yaml @@ -0,0 +1,34 @@ +# Values for the CircleCI helm_chart_testing job: a monolith install on a +# kind cluster from the image built by the pipeline (image.repository / tag / +# pullPolicy are passed with --set), backed by the throwaway PostgreSQL in +# ci/postgres.yaml so the migrations Job and the proxy have a database. +monolith: + enabled: true + +masterKey: + secretName: "" + generate: true + +gateway: + replicaCount: 1 + hpa: + enabled: false + resources: + requests: + cpu: 250m + memory: 512Mi + limits: + memory: 2Gi + config: + create: true + proxy_config: + general_settings: {} + +database: + writer: + host: litellm-ci-postgres + dbname: litellm + passwordSecret: + name: litellm-ci-postgres + usernameKey: username + passwordKey: password diff --git a/helm/litellm/templates/NOTES.txt b/helm/litellm/templates/NOTES.txt index 135dcb876e0..a046e4e5280 100644 --- a/helm/litellm/templates/NOTES.txt +++ b/helm/litellm/templates/NOTES.txt @@ -1,4 +1,16 @@ -LiteLLM componentized — release {{ .Release.Name }} in namespace {{ .Release.Namespace }}. +{{- if .Values.monolith.enabled }} +LiteLLM monolith: release {{ .Release.Name }} in namespace {{ .Release.Namespace }}, image {{ include "litellm.image" . }} + +One Deployment runs the full proxy (gateway routes, management routes and the Admin UI): + - proxy : Service {{ include "litellm.workload.fullname" . }} on port {{ .Values.gateway.service.port }} + +Port-forward: + kubectl -n {{ .Release.Namespace }} port-forward svc/{{ include "litellm.workload.fullname" . }} {{ .Values.gateway.service.port }} + +The monolith takes its configuration from the gateway.* values (config, resources, probes, hpa, pdb, keda, +metricsServer, collector, volumes, env, scheduling). backend.* and ui.* are ignored while monolith.enabled is true +{{- else }} +LiteLLM componentized: release {{ .Release.Name }} in namespace {{ .Release.Namespace }}, image {{ include "litellm.image" . }} Components: {{- if .Values.gateway.enabled }} @@ -16,43 +28,48 @@ Port-forward examples: kubectl -n {{ .Release.Namespace }} port-forward svc/{{ include "litellm.backend.fullname" . }} {{ .Values.backend.service.port }} kubectl -n {{ .Release.Namespace }} port-forward svc/{{ include "litellm.ui.fullname" . }} {{ .Values.ui.service.port }} +Set monolith.enabled=true to run everything as one Deployment instead +{{- end }} + Reminders: - Sensitive values come from Secret references only. Before installing, set: - - masterKey.secretName (Secret with the proxy master key) + - masterKey.secretName (Secret with the proxy master key, or + masterKey.generate: true with secretName "" to let the chart create one) - database.writer.{host,port,dbname} (writer connection pieces) - database.writer.passwordSecret.{name,usernameKey,passwordKey} (Secret holding the writer DB username + password) - database.writer.passwordSecret.{hostKey,urlKey} (optional, read writer host or full URL from that Secret) - - database.writer.useIAMAuth: true (optional — chart sets IAM_TOKEN_DB_AUTH=true and + - database.writer.useIAMAuth: true (optional: chart sets IAM_TOKEN_DB_AUTH=true and omits DATABASE_PASSWORD / DATABASE_URL so the proxy mints the URL from an IAM token at startup) - database.reader.host / database.reader.passwordSecret.{hostKey,urlKey} (optional, enables read-replica routing. The Secret name is required for hostKey, urlKey, and for host with password auth) - - database.reader.useIAMAuth: true (optional, requires database.writer.useIAMAuth: true — + - database.reader.useIAMAuth: true (optional, requires database.writer.useIAMAuth: true; chart emits DATABASE_*_READ_REPLICA env vars and omits DATABASE_PASSWORD_READ_REPLICA / DATABASE_URL_READ_REPLICA so the proxy mints the reader URL from an IAM token at startup) - - redis.passwordSecret.name (optional — set when redis.host is provided and the + - redis.host (optional: the proxy's coordination Redis) + - redis.passwordSecret.name (optional: set when redis.host is provided and the cache requires auth) - redis.passwordSecret.{hostKey,urlKey} (optional, read Redis host or full URL from that Secret) (the Secret name is required when either key is set) - - redis.cluster: true (optional — chart sets REDIS_CLUSTER_NODES from + - redis.cluster: true (optional: chart sets REDIS_CLUSTER_NODES from redis.host / redis.port so the proxy's Cache() constructs a RedisClusterCache; the cluster client discovers remaining nodes from CLUSTER SLOTS) - - Per-component extras (gateway / backend / ui): - - {component}.extraEnv / envConfigMaps / envSecrets (the latter two are lists of resource names → - envFrom configMapRef / secretRef) + - Per-component extras (gateway / backend / ui; gateway.* also drives the monolith): + - {component}.extraEnv / envConfigMaps / envSecrets (the latter two are lists of resource names rendered + as envFrom configMapRef / secretRef) - {component}.logLevel (renders as LITELLM_LOG) - gateway.config.proxy_config (rendered into a ConfigMap and mounted at - /app/config/config.yaml; gateway reads it via - CONFIG_FILE_PATH) + /app/config/config.yaml; read via CONFIG_FILE_PATH) - {component}.pdb.{enabled,minAvailable,maxUnavailable} (per-component PodDisruptionBudget; disabled by - default — with hpa.minReplicas of 1, minAvailable: 1 + default: with hpa.minReplicas of 1, minAvailable: 1 would block node drains) - {component}.topologySpreadConstraints (standard k8s list, e.g. spread replicas across topology.kubernetes.io/zone) - - Enable ingress.enabled=true to dispatch / → ui, gateway data-plane prefixes → gateway, and the catch-all → backend. + - Enable ingress.enabled=true to dispatch / to ui, gateway data-plane prefixes to gateway and the catch-all to + backend; in monolith mode every path goes to the monolith Service diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 18e938564e9..13222d8284e 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -1,11 +1,116 @@ {{/* -Common naming + label helpers shared by gateway, backend, and ui templates. +Common naming + label helpers shared by gateway, backend, ui, and monolith templates. */}} {{- define "litellm.name" -}} {{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" -}} {{- end -}} +{{/* +The one image reference every container in the chart uses. `repository:tag`, +or `repository:tag@digest` when image.digest is set; the tag falls back to +the chart appVersion. +*/}} +{{- define "litellm.image" -}} +{{- $tag := .Values.image.tag | default .Chart.AppVersion -}} +{{- if .Values.image.digest -}} +{{- printf "%s:%s@%s" .Values.image.repository $tag .Values.image.digest -}} +{{- else -}} +{{- printf "%s:%s" .Values.image.repository $tag -}} +{{- end -}} +{{- end -}} + +{{/* +Componentized mode renders the gateway / backend / ui Deployments only when +monolith mode is off. +*/}} +{{- define "litellm.gateway.render" -}} +{{- if and .Values.gateway.enabled (not .Values.monolith.enabled) -}}true{{- end -}} +{{- end -}} + +{{- define "litellm.backend.render" -}} +{{- if and .Values.backend.enabled (not .Values.monolith.enabled) -}}true{{- end -}} +{{- end -}} + +{{- define "litellm.ui.render" -}} +{{- if and .Values.ui.enabled (not .Values.monolith.enabled) -}}true{{- end -}} +{{- end -}} + +{{/* +The proxy Deployment (monolith mode) and the gateway Deployment +(componentized mode) are the same pod spec fed by .Values.gateway, so the +templates that serve both (HPA, KEDA, PDB, metrics Service, ServiceMonitor) +resolve their name, component label and selector through these. +*/}} +{{- define "litellm.proxy.fullname" -}} +{{- include "litellm.fullname" . -}} +{{- end -}} + +{{- define "litellm.proxy.selectorLabels" -}} +app.kubernetes.io/name: {{ include "litellm.name" . }} +app.kubernetes.io/instance: {{ .Release.Name }} +app.kubernetes.io/component: proxy +{{- end -}} + +{{- define "litellm.workload.componentName" -}} +{{- if .Values.monolith.enabled -}}proxy{{- else -}}gateway{{- end -}} +{{- end -}} + +{{- define "litellm.workload.fullname" -}} +{{- if .Values.monolith.enabled -}} +{{- include "litellm.proxy.fullname" . -}} +{{- else -}} +{{- include "litellm.gateway.fullname" . -}} +{{- end -}} +{{- end -}} + +{{- define "litellm.workload.selectorLabels" -}} +{{- if .Values.monolith.enabled -}} +{{- include "litellm.proxy.selectorLabels" . -}} +{{- else -}} +{{- include "litellm.gateway.selectorLabels" . -}} +{{- end -}} +{{- end -}} + +{{- define "litellm.workload.render" -}} +{{- if or .Values.monolith.enabled .Values.gateway.enabled -}}true{{- end -}} +{{- end -}} + +{{/* +Container arguments for the proxy (monolith) container. The image entrypoint +dispatches on the first argument. +*/}} +{{- define "litellm.proxy.args" -}} +- proxy +- --port +- "4000" +{{- if .Values.gateway.config.create }} +- --config +- /app/config/config.yaml +{{- end }} +{{- with .Values.monolith.extraArgs }} +{{ toYaml . }} +{{- end }} +{{- end -}} + +{{/* +Master key Secret reference. `masterKey.secretName` when set, otherwise the +Secret the chart generates when `masterKey.generate` is true. +*/}} +{{- define "litellm.masterKey.generatedSecretName" -}} +{{- printf "%s-masterkey" (include "litellm.fullname" .) -}} +{{- end -}} + +{{- define "litellm.masterKey.secretName" -}} +{{- if .Values.masterKey.secretName -}} +{{- .Values.masterKey.secretName -}} +{{- else if .Values.masterKey.generate -}} +{{- include "litellm.masterKey.generatedSecretName" . -}} +{{- else -}} +{{- fail "masterKey.secretName is required (the chart never accepts an inline master key); set it to an existing Secret or set masterKey.generate: true" -}} +{{- end -}} +{{- end -}} + {{- define "litellm.fullname" -}} {{- if .Values.fullnameOverride -}} {{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" -}} @@ -58,56 +163,27 @@ app.kubernetes.io/managed-by: {{ .Release.Service }} helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }} {{- end -}} -{{/* -Enterprise billable-request metering. Wired into gateway and backend, not the -migrations job. The gateway serves nearly all billable traffic, but the backend -keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a -SpendLogs row, so metering only the gateway would silently drop that traffic. -The client certificate identifies the deployment to LiteLLM's collector, so it is -mounted read-only from an existing Secret rather than passed through the -environment. -*/}} -{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}} -{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}} - {{- define "litellm.billingMetricsEnv" -}} - name: LITELLM_BILLING_METRICS_ENDPOINT value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }} -{{- if .Values.billingMetrics.caSecretName }} -- name: LITELLM_BILLING_METRICS_CA_CERT - value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }} -{{- end }} +- name: LITELLM_BILLING_METRICS_TOKEN_URL + value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-id +- name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-secret {{- with .Values.billingMetrics.exportIntervalMs }} - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: {{ . | quote }} {{- end }} {{- end -}} -{{- define "litellm.billingMetricsVolumes" -}} -- name: billing-metrics-mtls - secret: - secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }} -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - secret: - secretName: {{ .Values.billingMetrics.caSecretName }} -{{- end }} -{{- end -}} - -{{- define "litellm.billingMetricsVolumeMounts" -}} -- name: billing-metrics-mtls - mountPath: {{ include "litellm.billingMetrics.certDir" . }} - readOnly: true -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - mountPath: {{ include "litellm.billingMetrics.caDir" . }} - readOnly: true -{{- end }} -{{- end -}} {{/* Per-component selector labels — used in both Service selectors and Deployment matchLabels. @@ -137,6 +213,7 @@ Each component (gateway, backend, ui) has its own SA config under .Values.serviceAccounts.. When `create` is true and `name` is empty the chart defaults to "-litellm-". When `create` is false the chart uses the provided name, or the namespace `default` SA. +The monolith pod runs as the gateway ServiceAccount. */}} {{- define "litellm.gateway.serviceAccountName" -}} {{- if .Values.serviceAccounts.gateway.create -}} @@ -215,6 +292,22 @@ by the controller rather than declared, so nothing there can collide. {{- toYaml .podLabels }} {{- end -}} +{{/* +LITELLM_MASTER_KEY, and PROXY_MASTER_KEY for configs copied from litellm-helm, +for the app containers only. The migrations Job runs as a +pre-install hook, before the chart's generated master key Secret exists, and +migrations/run.py never reads the key, so the Job must not reference it. +*/}} +{{- define "litellm.masterKeyEnv" -}} +{{- range list "LITELLM_MASTER_KEY" "PROXY_MASTER_KEY" }} +- name: {{ . }} + valueFrom: + secretKeyRef: + name: {{ include "litellm.masterKey.secretName" $ }} + key: {{ $.Values.masterKey.secretKey | default "master-key" }} +{{- end }} +{{- end -}} + {{/* Master-key + database + redis env block — shared by gateway, backend, and the migrations Job. @@ -258,11 +351,6 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets. {{- define "litellm.serverEnv" -}} {{- $root := .root -}} {{- $component := .component -}} -- name: LITELLM_MASTER_KEY - valueFrom: - secretKeyRef: - name: {{ required "masterKey.secretName is required (the chart no longer accepts an inline master key)" $root.Values.masterKey.secretName }} - key: {{ $root.Values.masterKey.secretKey | default "master-key" }} {{- if $component.logLevel }} - name: LITELLM_LOG value: {{ $component.logLevel | quote }} @@ -305,6 +393,11 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets. secretKeyRef: name: {{ required "database.writer.passwordSecret.name is required" .passwordSecret.name }} key: {{ .passwordSecret.usernameKey | default "username" }} +- name: DATABASE_USERNAME + valueFrom: + secretKeyRef: + name: {{ .passwordSecret.name }} + key: {{ .passwordSecret.usernameKey | default "username" }} - name: DATABASE_NAME value: {{ required "database.writer.dbname is required" .dbname | quote }} {{- if .schema }} @@ -502,7 +595,7 @@ harmless no-op for the Job and authoritative for the app pods. {{- end -}} {{/* -In-container PgBouncer env for the gateway container. Under IAM or Entra auth the pooler mints and renews the database token itself. +In-container PgBouncer env for the gateway and proxy containers. Under IAM or Entra auth the pooler mints and renews the database token itself. */}} {{- define "litellm.connectionPoolEnv" -}} {{- with .Values.database.connectionPool -}} @@ -539,7 +632,7 @@ than silently replaced by the fallback. {{- $max := $component.pdb.maxUnavailable -}} {{- $minSet := not (or (kindIs "invalid" $min) (eq (printf "%v" $min) "")) -}} {{- $maxSet := not (or (kindIs "invalid" $max) (eq (printf "%v" $max) "")) -}} -{{- if and $component.enabled $component.pdb $component.pdb.enabled }} +{{- if and .enabled $component.pdb $component.pdb.enabled }} apiVersion: policy/v1 kind: PodDisruptionBudget metadata: @@ -597,6 +690,16 @@ ImplementationSpecific {{- end -}} {{- end -}} +{{/* +Service spec fields shared by the component and monolith Services. +Invoke with the component's `service` dict. +*/}} +{{- define "litellm.service.extras" -}} +{{- if and (eq .type "LoadBalancer") .loadBalancerClass -}} +loadBalancerClass: {{ .loadBalancerClass | quote }} +{{- end }} +{{- end -}} + {{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}} {{/* @@ -655,3 +758,27 @@ shutdown drain window. {{- end }} {{- end }} {{- end -}} +{{- define "litellm.liteadminEnv" -}} +{{- $root := .root -}} +{{- with $root.Values.liteadmin }} +{{- if .enabled }} +- name: LITELLM_ADMIN_AGENT_URL + value: {{ printf "http://%s-liteadmin:10000" (include "litellm.fullname" $root | trunc 53 | trimSuffix "-") | quote }} +- name: ADMIN_AGENT_SERVICE_TOKEN + valueFrom: + secretKeyRef: + name: {{ required "liteadmin.existingSecret is required" .existingSecret }} + key: ADMIN_AGENT_SERVICE_TOKEN +{{- $hasProxyBaseUrl := false }} +{{- range $.component.extraEnv }} +{{- if eq .name "PROXY_BASE_URL" }} +{{- $hasProxyBaseUrl = true }} +{{- end }} +{{- end }} +{{- if not $hasProxyBaseUrl }} +- name: PROXY_BASE_URL + value: {{ required "liteadmin.gatewayUrl is required" .gatewayUrl | quote }} +{{- end }} +{{- end }} +{{- end }} +{{- end -}} diff --git a/helm/litellm/templates/_workload.tpl b/helm/litellm/templates/_workload.tpl new file mode 100644 index 00000000000..8f52cd52f2a --- /dev/null +++ b/helm/litellm/templates/_workload.tpl @@ -0,0 +1,515 @@ +{{/* +The pod that runs the LLM data plane. In componentized mode it is the gateway +Deployment (`args: [gateway, ...]`), in monolith mode the proxy Deployment +(`args: [proxy, ...]`) that also serves the management API and the Admin UI. +Both are configured by .Values.gateway; the name, component label, selector +and args come from the litellm.workload.* helpers. +*/}} +{{- define "litellm.workload.deployment" -}} +{{- $component := include "litellm.workload.componentName" . -}} +{{- $fullname := include "litellm.workload.fullname" . -}} +{{- if and .Values.gateway.hpa.enabled .Values.gateway.keda.enabled }} +{{- fail "gateway.hpa.enabled and gateway.keda.enabled are mutually exclusive: two autoscalers on one Deployment fight over the replica count, so set gateway.hpa.enabled: false when using KEDA" }} +{{- end }} +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ $fullname }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: {{ $component }} + {{- with .Values.gateway.deploymentLabels }} + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with .Values.gateway.deploymentAnnotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + {{- if and (not .Values.gateway.hpa.enabled) (not .Values.gateway.keda.enabled) (not (kindIs "invalid" .Values.gateway.replicaCount)) }} + replicas: {{ .Values.gateway.replicaCount }} + {{- end }} + {{- $minReady := .Values.gateway.minReadySeconds }} + {{- if not (or (kindIs "invalid" $minReady) (eq (printf "%v" $minReady) "")) }} + minReadySeconds: {{ $minReady }} + {{- end }} + {{- with .Values.gateway.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} + selector: + matchLabels: + {{- include "litellm.workload.selectorLabels" . | nindent 6 }} + template: + metadata: + annotations: + {{- if .Values.gateway.config.create }} + checksum/config: {{ include (print $.Template.BasePath "/gateway/configmap.yaml") . | sha256sum }} + {{- end }} + {{- with .Values.gateway.podAnnotations }} + {{- toYaml . | nindent 8 }} + {{- end }} + labels: + {{- include "litellm.workload.selectorLabels" . | nindent 8 }} + {{- with .Values.gateway.podLabels }} + {{- include "litellm.podLabels" (dict "podLabels" . "componentName" "gateway") | nindent 8 }} + {{- end }} + spec: + serviceAccountName: {{ include "litellm.gateway.serviceAccountName" . }} + automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }} + {{- with .Values.gateway.podSecurityContext }} + securityContext: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.gateway.extraInitContainers }} + initContainers: + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} + containers: + - name: {{ $component }} + image: {{ include "litellm.image" . | quote }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + {{- with .Values.gateway.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} + args: + {{- if .Values.monolith.enabled }} + {{- include "litellm.proxy.args" . | nindent 12 }} + {{- else }} + - gateway + - --host + - 0.0.0.0 + - --port + - "4000" + {{- end }} + ports: + - name: http + containerPort: 4000 + protocol: TCP + env: + {{- include "litellm.lensConnectionEnv" . | nindent 12 }} + {{- if and (not .Values.monolith.enabled) .Values.ui.publicUrl }} + - name: LITELLM_CORS_ORIGINS + value: {{ include "litellm.ui.publicOrigin" . | quote }} + {{- end }} + {{- include "litellm.masterKeyEnv" $ | nindent 12 }} + {{- if .Values.monolith.enabled }} + {{- include "litellm.liteadminEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }} + {{- end }} + {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }} + {{- if .Values.gateway.config.create }} + - name: CONFIG_FILE_PATH + value: /app/config/config.yaml + {{- end }} + {{- if .Values.gateway.numWorkers }} + - name: NUM_WORKERS + value: {{ .Values.gateway.numWorkers | quote }} + {{- end }} + {{- if .Values.database.connectionPool.enabled }} + {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} + {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsEnv" . | nindent 12 }} + {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + {{- if eq (int .Values.gateway.metricsServer.port) 4000 }} + {{- fail (printf "gateway.metricsServer.port must differ from the %s port 4000" $component) }} + {{- end }} + - name: PROMETHEUS_MULTIPROC_DIR + value: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + {{- end }} + {{- if .Values.gateway.collector.enabled }} + {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} + {{- end }} + {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} + {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} + volumeMounts: + {{- if .Values.gateway.config.create }} + - name: gateway-config + mountPath: /app/config/config.yaml + subPath: config.yaml + {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: prometheus-multiproc + mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + {{- end }} + {{- if include "litellm.gateway.collectorSocketDir" . }} + - name: collector-socket + mountPath: {{ include "litellm.gateway.collectorSocketDir" . }} + {{- end }} + {{- with .Values.gateway.volumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} + {{- end }} + {{- with .Values.gateway.livenessProbe }} + livenessProbe: + {{- toYaml . | nindent 12 }} + {{- end }} + {{- with .Values.gateway.readinessProbe }} + readinessProbe: + {{- toYaml . | nindent 12 }} + {{- end }} + {{- with .Values.gateway.startupProbe }} + startupProbe: + {{- toYaml . | nindent 12 }} + {{- end }} + {{- with .Values.gateway.lifecycle }} + lifecycle: + {{- toYaml . | nindent 12 }} + {{- end }} + resources: + {{- toYaml .Values.gateway.resources | nindent 12 }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: metrics + image: {{ include "litellm.image" . | quote }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + {{- with .Values.gateway.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} + args: + - metrics + - --port + - {{ .Values.gateway.metricsServer.port | quote }} + env: + - name: PROMETHEUS_MULTIPROC_DIR + value: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + ports: + - name: metrics + containerPort: {{ .Values.gateway.metricsServer.port }} + protocol: TCP + volumeMounts: + - name: prometheus-multiproc + mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + readinessProbe: + tcpSocket: { port: metrics } + periodSeconds: 10 + livenessProbe: + tcpSocket: { port: metrics } + periodSeconds: 15 + failureThreshold: 6 + resources: + {{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }} + {{- end }} + {{- if .Values.gateway.collector.enabled }} + - name: collector + image: {{ include "litellm.image" . | quote }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + {{- with .Values.gateway.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} + args: + - collector + env: + {{- include "litellm.masterKeyEnv" $ | nindent 12 }} + {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }} + {{- if .Values.gateway.config.create }} + - name: CONFIG_FILE_PATH + value: /app/config/config.yaml + {{- end }} + {{- if .Values.database.connectionPool.enabled }} + {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} + {{- end }} + {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} + - name: LITELLM_JOB_ROLE + value: collector + {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} + {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts (include "litellm.gateway.collectorSocketDir" .) }} + volumeMounts: + {{- if .Values.gateway.config.create }} + - name: gateway-config + mountPath: /app/config/config.yaml + subPath: config.yaml + {{- end }} + {{- if include "litellm.gateway.collectorSocketDir" . }} + - name: collector-socket + mountPath: {{ include "litellm.gateway.collectorSocketDir" . }} + {{- end }} + {{- with .Values.gateway.volumeMounts }} + {{- toYaml . | nindent 12 }} + {{- end }} + {{- end }} + resources: + {{- toYaml .Values.gateway.collector.resources | nindent 12 }} + {{- end }} + {{- with .Values.gateway.extraContainers }} + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} + {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} + volumes: + {{- if .Values.gateway.config.create }} + - name: gateway-config + configMap: + name: {{ include "litellm.workload.fullname" . }}-config + {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: prometheus-multiproc + emptyDir: {} + {{- end }} + {{- if include "litellm.gateway.collectorSocketDir" . }} + - name: collector-socket + emptyDir: + sizeLimit: 1Mi + {{- end }} + {{- with .Values.gateway.volumes }} + {{- toYaml . | nindent 8 }} + {{- end }} + {{- end }} + {{- with .Values.gateway.nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.gateway.affinity }} + affinity: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.gateway.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.gateway.topologySpreadConstraints }} + topologySpreadConstraints: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- $gracePeriod := .Values.gateway.terminationGracePeriodSeconds }} + {{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }} + terminationGracePeriodSeconds: {{ $gracePeriod }} + {{- end }} +{{- end -}} + +{{- define "litellm.workload.service" -}} +{{- $component := include "litellm.workload.componentName" . -}} +apiVersion: v1 +kind: Service +metadata: + name: {{ include "litellm.workload.fullname" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: {{ $component }} + {{- with .Values.gateway.service.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + type: {{ .Values.gateway.service.type }} + {{- with include "litellm.service.extras" .Values.gateway.service }} + {{- . | nindent 2 }} + {{- end }} + ports: + - port: {{ .Values.gateway.service.port }} + targetPort: http + protocol: TCP + name: http + selector: + {{- include "litellm.workload.selectorLabels" . | nindent 4 }} +{{- end -}} + +{{- define "litellm.workload.hpa" -}} +{{- $component := include "litellm.workload.componentName" . -}} +apiVersion: autoscaling/v2 +kind: HorizontalPodAutoscaler +metadata: + name: {{ include "litellm.workload.fullname" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: {{ $component }} +spec: + scaleTargetRef: + apiVersion: apps/v1 + kind: Deployment + name: {{ include "litellm.workload.fullname" . }} + minReplicas: {{ .Values.gateway.hpa.minReplicas }} + maxReplicas: {{ .Values.gateway.hpa.maxReplicas }} + metrics: + {{- if .Values.gateway.hpa.targetCPUUtilizationPercentage }} + {{- if and .Values.gateway.collector.enabled .Values.gateway.collector.scaleOnGatewayContainerCpu }} + - type: ContainerResource + containerResource: + name: cpu + container: {{ $component }} + target: + type: Utilization + averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }} + {{- else }} + - type: Resource + resource: + name: cpu + target: + type: Utilization + averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }} + {{- end }} + {{- end }} + {{- if .Values.gateway.hpa.targetMemoryUtilizationPercentage }} + - type: Resource + resource: + name: memory + target: + type: Utilization + averageUtilization: {{ .Values.gateway.hpa.targetMemoryUtilizationPercentage }} + {{- end }} + {{- with .Values.gateway.hpa.targetRequestsPerSecond }} + - type: Pods + pods: + metric: + name: litellm_requests_per_second + target: + type: AverageValue + averageValue: {{ toJson . | trimAll "\"" | quote }} + {{- end }} + {{- with .Values.gateway.hpa.targetTokensPerSecond }} + - type: Pods + pods: + metric: + name: litellm_tokens_per_second + target: + type: AverageValue + averageValue: {{ toJson . | trimAll "\"" | quote }} + {{- end }} + {{- with .Values.gateway.hpa.behavior }} + behavior: + {{- toYaml . | nindent 4 }} + {{- end }} +{{- end -}} + +{{/* +KEDA ScaledObject targeting the workload Deployment. The optional Prometheus +triggers query the same per pod counters the HPA workload metrics use, +scoped to this release's metrics Service through the `job` label the +ServiceMonitor gives every scrape. +*/}} +{{- define "litellm.workload.keda" -}} +{{- $component := include "litellm.workload.componentName" . -}} +{{- $keda := .Values.gateway.keda -}} +{{- $prom := $keda.prometheus -}} +{{- $wantsPrometheus := or $prom.targetRequestsPerSecond $prom.targetTokensPerSecond -}} +{{- if and $wantsPrometheus (not $prom.serverAddress) }} +{{- fail "gateway.keda.prometheus.serverAddress is required when gateway.keda.prometheus.targetRequestsPerSecond or targetTokensPerSecond is set" }} +{{- end }} +{{- if and $wantsPrometheus (not .Values.gateway.metricsServer.enabled) }} +{{- fail "gateway.metricsServer.enabled is required when gateway.keda.prometheus.targetRequestsPerSecond or targetTokensPerSecond is set" }} +{{- end }} +{{- $selector := printf "namespace=%q,job=%q" .Release.Namespace (printf "%s-metrics" (include "litellm.workload.fullname" .)) -}} +apiVersion: keda.sh/v1alpha1 +kind: ScaledObject +metadata: + name: {{ include "litellm.workload.fullname" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: {{ $component }} + {{- with (dig "scaledObject" "annotations" (dict) $keda) }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + scaleTargetRef: + apiVersion: apps/v1 + kind: Deployment + name: {{ include "litellm.workload.fullname" . }} + minReplicaCount: {{ $keda.minReplicaCount }} + maxReplicaCount: {{ $keda.maxReplicaCount }} + pollingInterval: {{ $keda.pollingInterval }} + cooldownPeriod: {{ $keda.cooldownPeriod }} + {{- with $keda.advanced }} + advanced: + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with $keda.fallback }} + {{- $fallback := deepCopy . }} + {{- if not (hasKey $fallback "failureThreshold") }} + {{- $_ := set $fallback "failureThreshold" 3 }} + {{- end }} + {{- if not (hasKey $fallback "replicas") }} + {{- $_ := set $fallback "replicas" $keda.maxReplicaCount }} + {{- end }} + fallback: + {{- toYaml $fallback | nindent 4 }} + {{- end }} + triggers: + {{- with $keda.triggers }} + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with $prom.targetRequestsPerSecond }} + - type: prometheus + metricType: AverageValue + metadata: + serverAddress: {{ $prom.serverAddress | quote }} + query: {{ printf "sum(rate(litellm_proxy_total_requests_metric_total{%s}[1m]))" $selector | quote }} + threshold: {{ toJson . | trimAll "\"" | quote }} + {{- end }} + {{- with $prom.targetTokensPerSecond }} + - type: prometheus + metricType: AverageValue + metadata: + serverAddress: {{ $prom.serverAddress | quote }} + query: {{ printf "sum(rate(litellm_total_tokens_metric_total{%s}[1m]))" $selector | quote }} + threshold: {{ toJson . | trimAll "\"" | quote }} + {{- end }} +{{- end -}} + +{{- define "litellm.workload.serviceMetrics" -}} +{{- $component := include "litellm.workload.componentName" . -}} +apiVersion: v1 +kind: Service +metadata: + name: {{ include "litellm.workload.fullname" . }}-metrics + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: {{ $component }} +spec: + type: ClusterIP + ports: + - port: {{ .Values.gateway.metricsServer.port }} + targetPort: metrics + protocol: TCP + name: metrics + selector: + {{- include "litellm.workload.selectorLabels" . | nindent 4 }} +{{- end -}} + +{{- define "litellm.workload.serviceMonitor" -}} +{{- $component := include "litellm.workload.componentName" . -}} +{{- if not .Values.gateway.metricsServer.enabled }} +{{- fail "gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled: the http port serves /metrics/ behind virtual-key auth, so an unauthenticated scrape gets 401" }} +{{- end }} +apiVersion: monitoring.coreos.com/v1 +kind: ServiceMonitor +metadata: + name: {{ include "litellm.workload.fullname" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: {{ $component }} + {{- with .Values.gateway.serviceMonitor.labels }} + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with .Values.gateway.serviceMonitor.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + selector: + matchLabels: + {{- include "litellm.workload.selectorLabels" . | nindent 6 }} + namespaceSelector: + matchNames: + {{- with .Values.gateway.serviceMonitor.namespaceSelector.matchNames }} + {{- toYaml . | nindent 6 }} + {{- else }} + - {{ .Release.Namespace | quote }} + {{- end }} + endpoints: + - port: metrics + path: /metrics/ + interval: {{ .Values.gateway.serviceMonitor.interval }} + scrapeTimeout: {{ .Values.gateway.serviceMonitor.scrapeTimeout }} + scheme: http + {{- with .Values.gateway.serviceMonitor.relabelings }} + relabelings: + {{- toYaml . | nindent 8 }} + {{- end }} +{{- end -}} diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 7360261dab1..9118c6f3d88 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -1,4 +1,4 @@ -{{- if .Values.backend.enabled }} +{{- if include "litellm.backend.render" . }} {{- $componentHosts := .Values.ingress.componentHosts | default dict -}} {{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}} apiVersion: apps/v1 @@ -8,10 +8,21 @@ metadata: labels: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: backend + {{- with .Values.backend.deploymentLabels }} + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with .Values.backend.deploymentAnnotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} spec: {{- if and (not .Values.backend.hpa.enabled) (not (kindIs "invalid" .Values.backend.replicaCount)) }} replicas: {{ .Values.backend.replicaCount }} {{- end }} + {{- $minReady := .Values.backend.minReadySeconds }} + {{- if not (or (kindIs "invalid" $minReady) (eq (printf "%v" $minReady) "")) }} + minReadySeconds: {{ $minReady }} + {{- end }} {{- with .Values.backend.strategy }} strategy: {{- toYaml . | nindent 4 }} @@ -46,14 +57,24 @@ spec: imagePullSecrets: {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.backend.extraInitContainers }} + initContainers: + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} containers: - name: backend - image: "{{ .Values.backend.image.repository }}:{{ .Values.backend.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.backend.image.pullPolicy }} + image: {{ include "litellm.image" . | quote }} + imagePullPolicy: {{ .Values.image.pullPolicy }} {{- with .Values.backend.securityContext }} securityContext: {{- toYaml . | nindent 12 }} {{- end }} + args: + - backend + - --host + - 0.0.0.0 + - --port + - "4001" ports: - name: http containerPort: 4001 @@ -66,6 +87,8 @@ spec: - name: LITELLM_CORS_ALLOW_CREDENTIALS value: "true" {{- end }} + {{- include "litellm.masterKeyEnv" $ | nindent 12 }} + {{- include "litellm.liteadminEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }} {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }} {{- if .Values.backend.gatewayUrl }} - name: LITELLM_GATEWAY_URL @@ -79,16 +102,13 @@ spec: {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} {{- include "litellm.envFrom" .Values.backend | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.backend.volumeMounts }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config mountPath: /app/config/config.yaml subPath: config.yaml {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} {{- with .Values.backend.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -114,15 +134,12 @@ spec: {{- with .Values.backend.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} - {{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.backend.volumes }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config configMap: - name: {{ include "litellm.gateway.fullname" . }}-config - {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} + name: {{ include "litellm.workload.fullname" . }}-config {{- end }} {{- with .Values.backend.volumes }} {{- toYaml . | nindent 8 }} diff --git a/helm/litellm/templates/backend/hpa.yaml b/helm/litellm/templates/backend/hpa.yaml index a414092fb39..62efaacb830 100644 --- a/helm/litellm/templates/backend/hpa.yaml +++ b/helm/litellm/templates/backend/hpa.yaml @@ -1,4 +1,4 @@ -{{- if and .Values.backend.enabled .Values.backend.hpa.enabled }} +{{- if and (include "litellm.backend.render" .) .Values.backend.hpa.enabled }} apiVersion: autoscaling/v2 kind: HorizontalPodAutoscaler metadata: diff --git a/helm/litellm/templates/backend/poddisruptionbudget.yaml b/helm/litellm/templates/backend/poddisruptionbudget.yaml index 02853ac879c..ed40bda0d7c 100644 --- a/helm/litellm/templates/backend/poddisruptionbudget.yaml +++ b/helm/litellm/templates/backend/poddisruptionbudget.yaml @@ -1,5 +1,6 @@ {{- include "litellm.pdb" (dict "root" $ + "enabled" (include "litellm.backend.render" .) "component" .Values.backend "componentName" "backend" "fullname" (include "litellm.backend.fullname" .) diff --git a/helm/litellm/templates/backend/service.yaml b/helm/litellm/templates/backend/service.yaml index d480c654784..cd11ad95ccf 100644 --- a/helm/litellm/templates/backend/service.yaml +++ b/helm/litellm/templates/backend/service.yaml @@ -1,4 +1,4 @@ -{{- if .Values.backend.enabled }} +{{- if include "litellm.backend.render" . }} apiVersion: v1 kind: Service metadata: @@ -6,8 +6,15 @@ metadata: labels: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: backend + {{- with .Values.backend.service.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} spec: type: {{ .Values.backend.service.type }} + {{- with include "litellm.service.extras" .Values.backend.service }} + {{- . | nindent 2 }} + {{- end }} ports: - port: {{ .Values.backend.service.port }} targetPort: http diff --git a/helm/litellm/templates/extra-resources.yaml b/helm/litellm/templates/extra-resources.yaml new file mode 100644 index 00000000000..bd809ffb696 --- /dev/null +++ b/helm/litellm/templates/extra-resources.yaml @@ -0,0 +1,4 @@ +{{- range .Values.extraResources }} +--- +{{ tpl (toYaml .) $ }} +{{- end }} diff --git a/helm/litellm/templates/gateway/configmap.yaml b/helm/litellm/templates/gateway/configmap.yaml index d262bf25b87..8944aaa0a2f 100644 --- a/helm/litellm/templates/gateway/configmap.yaml +++ b/helm/litellm/templates/gateway/configmap.yaml @@ -2,7 +2,10 @@ apiVersion: v1 kind: ConfigMap metadata: - name: {{ include "litellm.gateway.fullname" . }}-config + name: {{ include "litellm.workload.fullname" . }}-config + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: {{ include "litellm.workload.componentName" . }} data: config.yaml: | {{ .Values.gateway.config.proxy_config | toYaml | indent 6 }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index b0e7f0dcebb..9db08850e26 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -1,252 +1,3 @@ -{{- if .Values.gateway.enabled }} -apiVersion: apps/v1 -kind: Deployment -metadata: - name: {{ include "litellm.gateway.fullname" . }} - labels: - {{- include "litellm.commonLabels" . | nindent 4 }} - app.kubernetes.io/component: gateway -spec: - {{- if and (not .Values.gateway.hpa.enabled) (not (kindIs "invalid" .Values.gateway.replicaCount)) }} - replicas: {{ .Values.gateway.replicaCount }} - {{- end }} - {{- with .Values.gateway.strategy }} - strategy: - {{- toYaml . | nindent 4 }} - {{- end }} - selector: - matchLabels: - {{- include "litellm.gateway.selectorLabels" . | nindent 6 }} - template: - metadata: - annotations: - {{- if .Values.gateway.config.create }} - checksum/config: {{ include (print $.Template.BasePath "/gateway/configmap.yaml") . | sha256sum }} - {{- end }} - {{- with .Values.gateway.podAnnotations }} - {{- toYaml . | nindent 8 }} - {{- end }} - labels: - {{- include "litellm.gateway.selectorLabels" . | nindent 8 }} - {{- with .Values.gateway.podLabels }} - {{- include "litellm.podLabels" (dict "podLabels" . "componentName" "gateway") | nindent 8 }} - {{- end }} - spec: - serviceAccountName: {{ include "litellm.gateway.serviceAccountName" . }} - automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }} - {{- with .Values.gateway.podSecurityContext }} - securityContext: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.imagePullSecrets }} - imagePullSecrets: - {{- toYaml . | nindent 8 }} - {{- end }} - containers: - - name: gateway - image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.gateway.image.pullPolicy }} - {{- with .Values.gateway.securityContext }} - securityContext: - {{- toYaml . | nindent 12 }} - {{- end }} - ports: - - name: http - containerPort: 4000 - protocol: TCP - env: - {{- include "litellm.lensConnectionEnv" . | nindent 12 }} - {{- if .Values.ui.publicUrl }} - - name: LITELLM_CORS_ORIGINS - value: {{ include "litellm.ui.publicOrigin" . | quote }} - {{- end }} - {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }} - {{- if .Values.gateway.config.create }} - - name: CONFIG_FILE_PATH - value: /app/config/config.yaml - {{- end }} - {{- if .Values.gateway.numWorkers }} - - name: NUM_WORKERS - value: {{ .Values.gateway.numWorkers | quote }} - {{- end }} - {{- if .Values.database.connectionPool.enabled }} - {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} - {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsEnv" . | nindent 12 }} - {{- end }} - {{- if .Values.gateway.metricsServer.enabled }} - {{- if eq (int .Values.gateway.metricsServer.port) 4000 }} - {{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }} - {{- end }} - - name: PROMETHEUS_MULTIPROC_DIR - value: {{ include "litellm.gateway.prometheusMultiprocDir" . }} - {{- end }} - {{- if .Values.gateway.collector.enabled }} - {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} - {{- end }} - {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} - volumeMounts: - {{- if .Values.gateway.config.create }} - - name: gateway-config - mountPath: /app/config/config.yaml - subPath: config.yaml - {{- end }} - {{- if .Values.gateway.metricsServer.enabled }} - - name: prometheus-multiproc - mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }} - {{- end }} - {{- if include "litellm.gateway.collectorSocketDir" . }} - - name: collector-socket - mountPath: {{ include "litellm.gateway.collectorSocketDir" . }} - {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} - {{- with .Values.gateway.volumeMounts }} - {{- toYaml . | nindent 12 }} - {{- end }} - {{- end }} - {{- with .Values.gateway.livenessProbe }} - livenessProbe: - {{- toYaml . | nindent 12 }} - {{- end }} - {{- with .Values.gateway.readinessProbe }} - readinessProbe: - {{- toYaml . | nindent 12 }} - {{- end }} - {{- with .Values.gateway.startupProbe }} - startupProbe: - {{- toYaml . | nindent 12 }} - {{- end }} - {{- with .Values.gateway.lifecycle }} - lifecycle: - {{- toYaml . | nindent 12 }} - {{- end }} - resources: - {{- toYaml .Values.gateway.resources | nindent 12 }} - {{- if .Values.gateway.metricsServer.enabled }} - - name: metrics - image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.gateway.image.pullPolicy }} - {{- with .Values.gateway.securityContext }} - securityContext: - {{- toYaml . | nindent 12 }} - {{- end }} - command: - - python - - -m - - litellm.proxy.prometheus_metrics_server - - --port - - {{ .Values.gateway.metricsServer.port | quote }} - env: - - name: PROMETHEUS_MULTIPROC_DIR - value: {{ include "litellm.gateway.prometheusMultiprocDir" . }} - ports: - - name: metrics - containerPort: {{ .Values.gateway.metricsServer.port }} - protocol: TCP - volumeMounts: - - name: prometheus-multiproc - mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }} - readinessProbe: - tcpSocket: { port: metrics } - periodSeconds: 10 - livenessProbe: - tcpSocket: { port: metrics } - periodSeconds: 15 - failureThreshold: 6 - resources: - {{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }} - {{- end }} - {{- if .Values.gateway.collector.enabled }} - - name: collector - image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.gateway.image.pullPolicy }} - {{- with .Values.gateway.securityContext }} - securityContext: - {{- toYaml . | nindent 12 }} - {{- end }} - command: - - python - - -m - - litellm.proxy.collector - env: - {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }} - {{- if .Values.gateway.config.create }} - - name: CONFIG_FILE_PATH - value: /app/config/config.yaml - {{- end }} - {{- if .Values.database.connectionPool.enabled }} - {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} - {{- end }} - {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} - - name: LITELLM_JOB_ROLE - value: collector - {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts (include "litellm.gateway.collectorSocketDir" .) }} - volumeMounts: - {{- if .Values.gateway.config.create }} - - name: gateway-config - mountPath: /app/config/config.yaml - subPath: config.yaml - {{- end }} - {{- if include "litellm.gateway.collectorSocketDir" . }} - - name: collector-socket - mountPath: {{ include "litellm.gateway.collectorSocketDir" . }} - {{- end }} - {{- with .Values.gateway.volumeMounts }} - {{- toYaml . | nindent 12 }} - {{- end }} - {{- end }} - resources: - {{- toYaml .Values.gateway.collector.resources | nindent 12 }} - {{- end }} - {{- with .Values.gateway.extraContainers }} - {{- tpl (toYaml .) $ | nindent 8 }} - {{- end }} - {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} - volumes: - {{- if .Values.gateway.config.create }} - - name: gateway-config - configMap: - name: {{ include "litellm.gateway.fullname" . }}-config - {{- end }} - {{- if .Values.gateway.metricsServer.enabled }} - - name: prometheus-multiproc - emptyDir: {} - {{- end }} - {{- if include "litellm.gateway.collectorSocketDir" . }} - - name: collector-socket - emptyDir: - sizeLimit: 1Mi - {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} - {{- with .Values.gateway.volumes }} - {{- toYaml . | nindent 8 }} - {{- end }} - {{- end }} - {{- with .Values.gateway.nodeSelector }} - nodeSelector: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.gateway.affinity }} - affinity: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.gateway.tolerations }} - tolerations: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- with .Values.gateway.topologySpreadConstraints }} - topologySpreadConstraints: - {{- toYaml . | nindent 8 }} - {{- end }} - {{- $gracePeriod := .Values.gateway.terminationGracePeriodSeconds }} - {{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }} - terminationGracePeriodSeconds: {{ $gracePeriod }} - {{- end }} +{{- if include "litellm.gateway.render" . }} +{{- include "litellm.workload.deployment" . }} {{- end }} diff --git a/helm/litellm/templates/gateway/hpa.yaml b/helm/litellm/templates/gateway/hpa.yaml index e7094e96106..d416a6056bd 100644 --- a/helm/litellm/templates/gateway/hpa.yaml +++ b/helm/litellm/templates/gateway/hpa.yaml @@ -1,65 +1,3 @@ -{{- if and .Values.gateway.enabled .Values.gateway.hpa.enabled }} -apiVersion: autoscaling/v2 -kind: HorizontalPodAutoscaler -metadata: - name: {{ include "litellm.gateway.fullname" . }} - labels: - {{- include "litellm.commonLabels" . | nindent 4 }} - app.kubernetes.io/component: gateway -spec: - scaleTargetRef: - apiVersion: apps/v1 - kind: Deployment - name: {{ include "litellm.gateway.fullname" . }} - minReplicas: {{ .Values.gateway.hpa.minReplicas }} - maxReplicas: {{ .Values.gateway.hpa.maxReplicas }} - metrics: - {{- if .Values.gateway.hpa.targetCPUUtilizationPercentage }} - {{- if and .Values.gateway.collector.enabled .Values.gateway.collector.scaleOnGatewayContainerCpu }} - - type: ContainerResource - containerResource: - name: cpu - container: gateway - target: - type: Utilization - averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }} - {{- else }} - - type: Resource - resource: - name: cpu - target: - type: Utilization - averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }} - {{- end }} - {{- end }} - {{- if .Values.gateway.hpa.targetMemoryUtilizationPercentage }} - - type: Resource - resource: - name: memory - target: - type: Utilization - averageUtilization: {{ .Values.gateway.hpa.targetMemoryUtilizationPercentage }} - {{- end }} - {{- with .Values.gateway.hpa.targetRequestsPerSecond }} - - type: Pods - pods: - metric: - name: litellm_requests_per_second - target: - type: AverageValue - averageValue: {{ toJson . | trimAll "\"" | quote }} - {{- end }} - {{- with .Values.gateway.hpa.targetTokensPerSecond }} - - type: Pods - pods: - metric: - name: litellm_tokens_per_second - target: - type: AverageValue - averageValue: {{ toJson . | trimAll "\"" | quote }} - {{- end }} - {{- with .Values.gateway.hpa.behavior }} - behavior: - {{- toYaml . | nindent 4 }} - {{- end }} +{{- if and (include "litellm.gateway.render" .) .Values.gateway.hpa.enabled }} +{{- include "litellm.workload.hpa" . }} {{- end }} diff --git a/helm/litellm/templates/gateway/keda.yaml b/helm/litellm/templates/gateway/keda.yaml new file mode 100644 index 00000000000..d46085ffef1 --- /dev/null +++ b/helm/litellm/templates/gateway/keda.yaml @@ -0,0 +1,3 @@ +{{- if and (include "litellm.gateway.render" .) .Values.gateway.keda.enabled }} +{{- include "litellm.workload.keda" . }} +{{- end }} diff --git a/helm/litellm/templates/gateway/poddisruptionbudget.yaml b/helm/litellm/templates/gateway/poddisruptionbudget.yaml index 15e89af17d7..313e3cc7535 100644 --- a/helm/litellm/templates/gateway/poddisruptionbudget.yaml +++ b/helm/litellm/templates/gateway/poddisruptionbudget.yaml @@ -1,5 +1,6 @@ {{- include "litellm.pdb" (dict "root" $ + "enabled" (include "litellm.gateway.render" .) "component" .Values.gateway "componentName" "gateway" "fullname" (include "litellm.gateway.fullname" .) diff --git a/helm/litellm/templates/gateway/service-metrics.yaml b/helm/litellm/templates/gateway/service-metrics.yaml index ad9bc05a9fd..e33ae8c63e3 100644 --- a/helm/litellm/templates/gateway/service-metrics.yaml +++ b/helm/litellm/templates/gateway/service-metrics.yaml @@ -1,18 +1,3 @@ -{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }} -apiVersion: v1 -kind: Service -metadata: - name: {{ include "litellm.gateway.fullname" . }}-metrics - labels: - {{- include "litellm.commonLabels" . | nindent 4 }} - app.kubernetes.io/component: gateway -spec: - type: ClusterIP - ports: - - port: {{ .Values.gateway.metricsServer.port }} - targetPort: metrics - protocol: TCP - name: metrics - selector: - {{- include "litellm.gateway.selectorLabels" . | nindent 4 }} +{{- if and (include "litellm.gateway.render" .) .Values.gateway.metricsServer.enabled }} +{{- include "litellm.workload.serviceMetrics" . }} {{- end }} diff --git a/helm/litellm/templates/gateway/service.yaml b/helm/litellm/templates/gateway/service.yaml index 03a4167a0ab..8f3e7ad0673 100644 --- a/helm/litellm/templates/gateway/service.yaml +++ b/helm/litellm/templates/gateway/service.yaml @@ -1,18 +1,3 @@ -{{- if .Values.gateway.enabled }} -apiVersion: v1 -kind: Service -metadata: - name: {{ include "litellm.gateway.fullname" . }} - labels: - {{- include "litellm.commonLabels" . | nindent 4 }} - app.kubernetes.io/component: gateway -spec: - type: {{ .Values.gateway.service.type }} - ports: - - port: {{ .Values.gateway.service.port }} - targetPort: http - protocol: TCP - name: http - selector: - {{- include "litellm.gateway.selectorLabels" . | nindent 4 }} +{{- if include "litellm.gateway.render" . }} +{{- include "litellm.workload.service" . }} {{- end }} diff --git a/helm/litellm/templates/gateway/servicemonitor.yaml b/helm/litellm/templates/gateway/servicemonitor.yaml index e1bafa6e388..3e5d538b4e0 100644 --- a/helm/litellm/templates/gateway/servicemonitor.yaml +++ b/helm/litellm/templates/gateway/servicemonitor.yaml @@ -1,28 +1,3 @@ -{{- if and .Values.gateway.enabled .Values.gateway.serviceMonitor.enabled }} -{{- if not .Values.gateway.metricsServer.enabled }} -{{- fail "gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled: the http port serves /metrics/ behind virtual-key auth, so an unauthenticated scrape gets 401" }} -{{- end }} -apiVersion: monitoring.coreos.com/v1 -kind: ServiceMonitor -metadata: - name: {{ include "litellm.gateway.fullname" . }} - labels: - {{- include "litellm.commonLabels" . | nindent 4 }} - app.kubernetes.io/component: gateway - {{- with .Values.gateway.serviceMonitor.labels }} - {{- toYaml . | nindent 4 }} - {{- end }} -spec: - selector: - matchLabels: - {{- include "litellm.gateway.selectorLabels" . | nindent 6 }} - namespaceSelector: - matchNames: - - {{ .Release.Namespace | quote }} - endpoints: - - port: metrics - path: /metrics/ - interval: {{ .Values.gateway.serviceMonitor.interval }} - scrapeTimeout: {{ .Values.gateway.serviceMonitor.scrapeTimeout }} - scheme: http +{{- if and (include "litellm.gateway.render" .) .Values.gateway.serviceMonitor.enabled }} +{{- include "litellm.workload.serviceMonitor" . }} {{- end }} diff --git a/helm/litellm/templates/ingress.yaml b/helm/litellm/templates/ingress.yaml index 9f778e78087..3fb138aacdb 100644 --- a/helm/litellm/templates/ingress.yaml +++ b/helm/litellm/templates/ingress.yaml @@ -4,6 +4,7 @@ {{- if not (and $componentHosts.ui $componentHosts.backend $componentHosts.gateway) }} {{- fail "ingress.componentHosts.ui, backend, and gateway must all be set together" }} {{- end }} +{{- if not .Values.monolith.enabled }} {{- if not .Values.gateway.enabled }} {{- fail "ingress.componentHosts requires gateway.enabled=true" }} {{- end }} @@ -13,12 +14,14 @@ {{- if not .Values.ui.enabled }} {{- fail "ingress.componentHosts requires ui.enabled=true" }} {{- end }} +{{- end }} {{- if .Values.ingress.host }} {{- fail "ingress.host must be empty when ingress.componentHosts is configured" }} {{- end }} {{- if .Values.ingress.extraPaths }} {{- fail "ingress.extraPaths must be empty when ingress.componentHosts is configured" }} {{- end }} +{{- if not .Values.monolith.enabled }} {{- if not (trim .Values.ui.backendUrl) }} {{- fail "ui.backendUrl is required when ingress.componentHosts is configured" }} {{- end }} @@ -31,24 +34,30 @@ {{- fail "ui.publicUrl is required when ingress.componentHosts is configured" }} {{- end }} {{- end }} +{{- end }} {{- if .Values.ingress.enabled -}} +{{- $monolith := .Values.monolith.enabled -}} {{- $gatewayName := include "litellm.gateway.fullname" . -}} {{- $backendName := include "litellm.backend.fullname" . -}} {{- $uiName := include "litellm.ui.fullname" . -}} {{- $gatewayPort := .Values.gateway.service.port -}} {{- $backendPort := .Values.backend.service.port -}} {{- $uiPort := .Values.ui.service.port -}} +{{- $proxyName := include "litellm.proxy.fullname" . -}} {{- $controller := .Values.ingress.controller | default "alb" -}} {{- if not (has $controller (list "alb" "nginx")) }} {{- fail (printf "ingress.controller: unknown controller %q, expected one of alb, nginx" $controller) }} {{- end }} {{/* Backends addressable from ingress.extraPaths, keyed by the `service` field. + In monolith mode the one proxy Service serves every component, so every + key resolves to it and an entry's `service` only has to be a known name. */}} +{{- $proxyBackend := dict "name" $proxyName "port" $gatewayPort -}} {{- $extraPathBackends := dict - "gateway" (dict "name" $gatewayName "port" $gatewayPort) - "backend" (dict "name" $backendName "port" $backendPort) - "ui" (dict "name" $uiName "port" $uiPort) + "gateway" (ternary $proxyBackend (dict "name" $gatewayName "port" $gatewayPort) $monolith) + "backend" (ternary $proxyBackend (dict "name" $backendName "port" $backendPort) $monolith) + "ui" (ternary $proxyBackend (dict "name" $uiName "port" $uiPort) $monolith) -}} {{/* UI paths (Next.js static export). @@ -122,7 +131,7 @@ at "/" Prefix would swallow the whole backend management API) instead of adding to it. */}} -{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}} +{{- $builtinPathKeys := ternary (list "/|Prefix") (list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix") $monolith -}} {{- if $hasComponentHosts }} apiVersion: networking.k8s.io/v1 kind: Ingress @@ -150,9 +159,9 @@ spec: pathType: Prefix backend: service: - name: {{ $uiName }} + name: {{ $extraPathBackends.ui.name }} port: - number: {{ $uiPort }} + number: {{ $extraPathBackends.ui.port }} - host: {{ $componentHosts.backend | quote }} http: paths: @@ -160,9 +169,9 @@ spec: pathType: Prefix backend: service: - name: {{ $backendName }} + name: {{ $extraPathBackends.backend.name }} port: - number: {{ $backendPort }} + number: {{ $extraPathBackends.backend.port }} - host: {{ $componentHosts.gateway | quote }} http: paths: @@ -179,9 +188,9 @@ spec: pathType: Prefix backend: service: - name: {{ $gatewayName }} + name: {{ $extraPathBackends.gateway.name }} port: - number: {{ $gatewayPort }} + number: {{ $extraPathBackends.gateway.port }} {{- else }} apiVersion: networking.k8s.io/v1 kind: Ingress @@ -207,6 +216,7 @@ spec: {{- end }} http: paths: + {{- if not $monolith }} # --- UI (Next.js static export) --- {{- range $uiPaths }} {{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" .path "pathType" .pathType) }} @@ -249,6 +259,7 @@ spec: port: number: {{ $gatewayPort }} {{- end }} + {{- end }} {{- if and .Values.lens.enabled .Values.lens.service.name }} {{- $builtinPathKeys = append $builtinPathKeys "/lens-ingest|Prefix" }} - path: /lens-ingest @@ -324,6 +335,16 @@ spec: port: {{- toYaml $port | nindent 18 }} {{- end }} + {{- if $monolith }} + # --- Monolith: the proxy serves the UI, the data plane and the management API --- + - path: / + pathType: Prefix + backend: + service: + name: {{ $proxyName }} + port: + number: {{ $gatewayPort }} + {{- else }} # --- Catch-all → backend (management API: /key/*, /user/*, /team/*, ...) --- - path: / pathType: Prefix @@ -332,5 +353,6 @@ spec: name: {{ $backendName }} port: number: {{ $backendPort }} + {{- end }} {{- end }} {{- end }} diff --git a/helm/litellm-helm/templates/liteadmin.yaml b/helm/litellm/templates/liteadmin.yaml similarity index 96% rename from helm/litellm-helm/templates/liteadmin.yaml rename to helm/litellm/templates/liteadmin.yaml index 711edaf6913..bc5f2c9510b 100644 --- a/helm/litellm-helm/templates/liteadmin.yaml +++ b/helm/litellm/templates/liteadmin.yaml @@ -32,7 +32,7 @@ spec: runAsNonRoot: true containers: - name: liteadmin - image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" + image: {{ include "litellm.image" . | quote }} imagePullPolicy: {{ .Values.image.pullPolicy }} args: ["--admin-agent"] securityContext: diff --git a/helm/litellm/templates/migrations-job.yaml b/helm/litellm/templates/migrations-job.yaml index de1cc2b103b..8179c26cbd7 100644 --- a/helm/litellm/templates/migrations-job.yaml +++ b/helm/litellm/templates/migrations-job.yaml @@ -1,7 +1,7 @@ {{- if .Values.migrationJob.enabled -}} # Pre-install / pre-upgrade hook that runs `prisma migrate deploy` against -# the writer database before the gateway and backend Deployments are rolled -# out. Required because the gateway and backend both spin up Prisma at +# the writer database before the application Deployments are rolled out. +# Required because the proxy, gateway and backend all spin up Prisma at # startup and assume the LiteLLM schema (LiteLLM_Config, # LiteLLM_VerificationToken, LiteLLM_SpendLogs, ...) already exists. # @@ -16,9 +16,13 @@ metadata: labels: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: migrations - {{- if or .Values.migrationJob.hooks.helm.enabled .Values.migrationJob.hooks.argocd.enabled }} + {{- with .Values.migrationJob.jobLabels }} + {{- toYaml . | nindent 4 }} + {{- end }} + {{- $helmHook := .Values.migrationJob.hooks.helm.enabled }} + {{- if or $helmHook .Values.migrationJob.hooks.argocd.enabled .Values.migrationJob.jobAnnotations }} annotations: - {{- if .Values.migrationJob.hooks.helm.enabled }} + {{- if $helmHook }} helm.sh/hook: pre-install,pre-upgrade helm.sh/hook-delete-policy: before-hook-creation helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "0" | quote }} @@ -27,6 +31,9 @@ metadata: argocd.argoproj.io/hook: PreSync argocd.argoproj.io/hook-delete-policy: BeforeHookCreation {{- end }} + {{- with .Values.migrationJob.jobAnnotations }} + {{- toYaml . | nindent 4 }} + {{- end }} {{- end }} spec: backoffLimit: {{ .Values.migrationJob.backoffLimit }} @@ -43,6 +50,10 @@ spec: {{- $chartLabels := merge (dict "app.kubernetes.io/component" "migrations") (fromYaml (include "litellm.commonLabels" .)) }} labels: {{- toYaml (merge (deepCopy .Values.migrationJob.podLabels) $chartLabels) | nindent 8 }} + {{- with .Values.migrationJob.annotations }} + annotations: + {{- toYaml . | nindent 8 }} + {{- end }} spec: restartPolicy: Never serviceAccountName: {{ include "litellm.migrations.serviceAccountName" . }} @@ -55,14 +66,29 @@ spec: imagePullSecrets: {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.migrationJob.extraInitContainers }} + initContainers: + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} containers: - name: prisma-migrations - image: "{{ .Values.migrationJob.image.repository }}:{{ .Values.migrationJob.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.migrationJob.image.pullPolicy }} + image: {{ include "litellm.image" . | quote }} + imagePullPolicy: {{ .Values.image.pullPolicy }} {{- with .Values.migrationJob.securityContext }} securityContext: {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.migrationJob.command }} + command: {{ toYaml . | nindent 12 }} + {{- end }} + {{- if .Values.migrationJob.args }} + args: {{ toYaml .Values.migrationJob.args | nindent 12 }} + {{- else if .Values.migrationJob.command }} + args: [] + {{- else }} + args: + - migrations + {{- end }} env: {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.migrationJob) | nindent 12 }} {{- with .Values.migrationJob.volumeMounts }} @@ -73,6 +99,9 @@ spec: resources: {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.migrationJob.extraContainers }} + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} {{- with .Values.migrationJob.volumes }} volumes: {{- toYaml . | nindent 8 }} diff --git a/helm/litellm/templates/monolith/deployment.yaml b/helm/litellm/templates/monolith/deployment.yaml new file mode 100644 index 00000000000..7bad16386e1 --- /dev/null +++ b/helm/litellm/templates/monolith/deployment.yaml @@ -0,0 +1,3 @@ +{{- if .Values.monolith.enabled }} +{{- include "litellm.workload.deployment" . }} +{{- end }} diff --git a/helm/litellm/templates/monolith/hpa.yaml b/helm/litellm/templates/monolith/hpa.yaml new file mode 100644 index 00000000000..ee9180e71bb --- /dev/null +++ b/helm/litellm/templates/monolith/hpa.yaml @@ -0,0 +1,3 @@ +{{- if and .Values.monolith.enabled .Values.gateway.hpa.enabled }} +{{- include "litellm.workload.hpa" . }} +{{- end }} diff --git a/helm/litellm/templates/monolith/keda.yaml b/helm/litellm/templates/monolith/keda.yaml new file mode 100644 index 00000000000..d4054432e38 --- /dev/null +++ b/helm/litellm/templates/monolith/keda.yaml @@ -0,0 +1,3 @@ +{{- if and .Values.monolith.enabled .Values.gateway.keda.enabled }} +{{- include "litellm.workload.keda" . }} +{{- end }} diff --git a/helm/litellm/templates/monolith/poddisruptionbudget.yaml b/helm/litellm/templates/monolith/poddisruptionbudget.yaml new file mode 100644 index 00000000000..2be83e57dfc --- /dev/null +++ b/helm/litellm/templates/monolith/poddisruptionbudget.yaml @@ -0,0 +1,7 @@ +{{- include "litellm.pdb" (dict + "root" $ + "enabled" .Values.monolith.enabled + "component" .Values.gateway + "componentName" "proxy" + "fullname" (include "litellm.proxy.fullname" .) + "selectorLabels" (include "litellm.proxy.selectorLabels" .)) }} diff --git a/helm/litellm/templates/monolith/service-metrics.yaml b/helm/litellm/templates/monolith/service-metrics.yaml new file mode 100644 index 00000000000..07a68c453df --- /dev/null +++ b/helm/litellm/templates/monolith/service-metrics.yaml @@ -0,0 +1,3 @@ +{{- if and .Values.monolith.enabled .Values.gateway.metricsServer.enabled }} +{{- include "litellm.workload.serviceMetrics" . }} +{{- end }} diff --git a/helm/litellm/templates/monolith/service.yaml b/helm/litellm/templates/monolith/service.yaml new file mode 100644 index 00000000000..4f87259983a --- /dev/null +++ b/helm/litellm/templates/monolith/service.yaml @@ -0,0 +1,3 @@ +{{- if .Values.monolith.enabled }} +{{- include "litellm.workload.service" . }} +{{- end }} diff --git a/helm/litellm/templates/monolith/servicemonitor.yaml b/helm/litellm/templates/monolith/servicemonitor.yaml new file mode 100644 index 00000000000..29a51a379f2 --- /dev/null +++ b/helm/litellm/templates/monolith/servicemonitor.yaml @@ -0,0 +1,3 @@ +{{- if and .Values.monolith.enabled .Values.gateway.serviceMonitor.enabled }} +{{- include "litellm.workload.serviceMonitor" . }} +{{- end }} diff --git a/helm/litellm/templates/secret-masterkey.yaml b/helm/litellm/templates/secret-masterkey.yaml new file mode 100644 index 00000000000..3af2ee6f07d --- /dev/null +++ b/helm/litellm/templates/secret-masterkey.yaml @@ -0,0 +1,22 @@ +{{- if and .Values.masterKey.generate (not .Values.masterKey.secretName) }} +{{- $name := include "litellm.masterKey.generatedSecretName" . }} +{{- $key := .Values.masterKey.secretKey | default "master-key" }} +{{- $existing := lookup "v1" "Secret" .Release.Namespace $name }} +{{- $value := "" }} +{{- if and $existing (index $existing.data $key) }} +{{- $value = index $existing.data $key | b64dec }} +{{- else }} +{{- $value = printf "sk-%s" (randAlphaNum 32) }} +{{- end }} +apiVersion: v1 +kind: Secret +metadata: + name: {{ $name }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + annotations: + helm.sh/resource-policy: keep +type: Opaque +data: + {{ $key }}: {{ $value | b64enc }} +{{- end }} diff --git a/helm/litellm/templates/serviceaccount.yaml b/helm/litellm/templates/serviceaccount.yaml index a2fc52f47c0..9b1171a3e90 100644 --- a/helm/litellm/templates/serviceaccount.yaml +++ b/helm/litellm/templates/serviceaccount.yaml @@ -1,5 +1,5 @@ {{- $prev := false -}} -{{- if .Values.serviceAccounts.gateway.create -}} +{{- if and .Values.serviceAccounts.gateway.create (include "litellm.workload.render" .) -}} {{- $prev = true }} apiVersion: v1 kind: ServiceAccount @@ -14,7 +14,7 @@ metadata: {{- end }} automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }} {{- end }} -{{- if .Values.serviceAccounts.backend.create }} +{{- if and .Values.serviceAccounts.backend.create (include "litellm.backend.render" .) }} {{- if $prev }} --- {{- end }} @@ -32,7 +32,7 @@ metadata: {{- end }} automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }} {{- end }} -{{- if .Values.serviceAccounts.ui.create }} +{{- if and .Values.serviceAccounts.ui.create (include "litellm.ui.render" .) }} {{- if $prev }} --- {{- end }} diff --git a/helm/litellm/templates/tests/test-connection.yaml b/helm/litellm/templates/tests/test-connection.yaml new file mode 100644 index 00000000000..bc258c92790 --- /dev/null +++ b/helm/litellm/templates/tests/test-connection.yaml @@ -0,0 +1,53 @@ +{{- if include "litellm.workload.render" . }} +apiVersion: v1 +kind: Pod +metadata: + name: {{ include "litellm.fullname" . }}-test-connection + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: test + annotations: + helm.sh/hook: test + helm.sh/hook-delete-policy: before-hook-creation +spec: + restartPolicy: Never + {{- with .Values.gateway.podSecurityContext }} + securityContext: + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 4 }} + {{- end }} + containers: + - name: health + image: {{ include "litellm.image" . | quote }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + {{- with .Values.gateway.securityContext }} + securityContext: + {{- toYaml . | nindent 8 }} + {{- end }} + env: + - name: LITELLM_READINESS_URL + value: http://{{ include "litellm.workload.fullname" . }}:{{ .Values.gateway.service.port }}/health/readiness + args: + - python + - -c + - | + import json, os, sys, time, urllib.error, urllib.request + url = os.environ["LITELLM_READINESS_URL"] + for attempt in range(1, 31): + try: + with urllib.request.urlopen(url, timeout=5) as response: + body = response.read().decode() + print(url, response.status, body) + db = json.loads(body).get("db") + if db != "connected": + print(f"readiness reports db={db!r}: the proxy is not using the configured database") + sys.exit(1) + sys.exit(0) + except (urllib.error.URLError, OSError) as error: + print(f"attempt {attempt}: {url}: {error}") + time.sleep(2) + sys.exit(1) +{{- end }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index 394eea17451..486e1dae5c8 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -1,6 +1,6 @@ +{{- if include "litellm.ui.render" . }} {{- $componentHosts := .Values.ingress.componentHosts | default dict -}} {{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}} -{{- if .Values.ui.enabled }} apiVersion: apps/v1 kind: Deployment metadata: @@ -8,10 +8,21 @@ metadata: labels: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: ui + {{- with .Values.ui.deploymentLabels }} + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with .Values.ui.deploymentAnnotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} spec: {{- if and (not .Values.ui.hpa.enabled) (not (kindIs "invalid" .Values.ui.replicaCount)) }} replicas: {{ .Values.ui.replicaCount }} {{- end }} + {{- $minReady := .Values.ui.minReadySeconds }} + {{- if not (or (kindIs "invalid" $minReady) (eq (printf "%v" $minReady) "")) }} + minReadySeconds: {{ $minReady }} + {{- end }} {{- with .Values.ui.strategy }} strategy: {{- toYaml . | nindent 4 }} @@ -41,14 +52,20 @@ spec: imagePullSecrets: {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.ui.extraInitContainers }} + initContainers: + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} containers: - name: ui - image: "{{ .Values.ui.image.repository }}:{{ .Values.ui.image.tag | default .Chart.AppVersion }}" - imagePullPolicy: {{ .Values.ui.image.pullPolicy }} + image: {{ include "litellm.image" . | quote }} + imagePullPolicy: {{ .Values.image.pullPolicy }} {{- with .Values.ui.securityContext }} securityContext: {{- toYaml . | nindent 12 }} {{- end }} + args: + - ui ports: - name: http containerPort: 3000 diff --git a/helm/litellm/templates/ui/hpa.yaml b/helm/litellm/templates/ui/hpa.yaml index a9b0b51129e..67e17b269c7 100644 --- a/helm/litellm/templates/ui/hpa.yaml +++ b/helm/litellm/templates/ui/hpa.yaml @@ -1,4 +1,4 @@ -{{- if and .Values.ui.enabled .Values.ui.hpa.enabled }} +{{- if and (include "litellm.ui.render" .) .Values.ui.hpa.enabled }} apiVersion: autoscaling/v2 kind: HorizontalPodAutoscaler metadata: diff --git a/helm/litellm/templates/ui/poddisruptionbudget.yaml b/helm/litellm/templates/ui/poddisruptionbudget.yaml index f7a3a694e9c..fa4bb44612c 100644 --- a/helm/litellm/templates/ui/poddisruptionbudget.yaml +++ b/helm/litellm/templates/ui/poddisruptionbudget.yaml @@ -1,5 +1,6 @@ {{- include "litellm.pdb" (dict "root" $ + "enabled" (include "litellm.ui.render" .) "component" .Values.ui "componentName" "ui" "fullname" (include "litellm.ui.fullname" .) diff --git a/helm/litellm/templates/ui/service.yaml b/helm/litellm/templates/ui/service.yaml index 52b539fa00c..d544d9d6148 100644 --- a/helm/litellm/templates/ui/service.yaml +++ b/helm/litellm/templates/ui/service.yaml @@ -1,4 +1,4 @@ -{{- if .Values.ui.enabled }} +{{- if include "litellm.ui.render" . }} apiVersion: v1 kind: Service metadata: @@ -6,8 +6,15 @@ metadata: labels: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: ui + {{- with .Values.ui.service.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} spec: type: {{ .Values.ui.service.type }} + {{- with include "litellm.service.extras" .Values.ui.service }} + {{- . | nindent 2 }} + {{- end }} ports: - port: {{ .Values.ui.service.port }} targetPort: http diff --git a/helm/litellm/tests/billing_metrics_tests.yaml b/helm/litellm/tests/billing_metrics_tests.yaml index ceba0bd1430..d17b8192fce 100644 --- a/helm/litellm/tests/billing_metrics_tests.yaml +++ b/helm/litellm/tests/billing_metrics_tests.yaml @@ -1,39 +1,30 @@ -suite: test billingMetrics wiring on gateway and backend +suite: test bearer-only billingMetrics wiring templates: - gateway/deployment.yaml - gateway/configmap.yaml - backend/deployment.yaml + - backend/configmap.yaml - migrations-job.yaml values: - ./values/required.yaml tests: - - it: is off by default, adding no env, volume, or mount + - it: is disabled by default template: gateway/deployment.yaml asserts: - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - equal: - path: spec.template.spec.containers[0].volumeMounts - value: - - name: gateway-config - mountPath: /app/config/config.yaml - subPath: config.yaml - notContains: path: spec.template.spec.containers[0].env content: name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics - - it: renders the endpoint and the mounted cert paths when enabled + - it: wires gateway URLs and bearer credentials from Secret keys template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls asserts: - contains: path: spec.template.spec.containers[0].env @@ -43,87 +34,68 @@ tests: - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: /etc/litellm/billing-mtls/tls.crt + name: LITELLM_BILLING_METRICS_TOKEN_URL + value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: /etc/litellm/billing-mtls/tls.key - - - it: mounts the cert secret read-only alongside the config volume - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-id - contains: - path: spec.template.spec.volumes + path: spec.template.spec.containers[0].env content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - # The production collector presents a public web-PKI certificate, so the CA - # override must stay absent unless a private collector is configured. - - it: omits the CA env, volume, and mount when no caSecretName is set - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-secret - notContains: path: spec.template.spec.volumes content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca + name: billing-metrics - notContains: - path: spec.template.spec.containers[0].env + path: spec.template.spec.containers[0].volumeMounts content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt + name: billing-metrics - - it: mounts the CA secret when caSecretName is set - template: gateway/deployment.yaml + - it: wires billing metrics on the backend without volumes + template: backend/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls - caSecretName: billing-ca asserts: - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt - - contains: + name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-secret + - notContains: path: spec.template.spec.volumes content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true + name: billing-metrics - - it: passes the export interval through only when set + - it: does not add billing metrics to migrations + template: migrations-job.yaml + set: + billingMetrics: + enabled: true + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + + - it: renders an optional export interval template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls exportIntervalMs: 5000 asserts: - contains: @@ -132,102 +104,27 @@ tests: name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: "5000" - - it: keeps user-supplied gateway volumes alongside the billing secret + - it: requires the endpoint when enabled template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls - gateway.volumes: - - name: custom-callbacks - configMap: - name: my-callbacks - gateway.volumeMounts: - - name: custom-callbacks - mountPath: /app/callbacks + endpoint: "" asserts: - - contains: - path: spec.template.spec.volumes - content: - name: custom-callbacks - configMap: - name: my-callbacks - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls + - failedTemplate: + errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true - # The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp), - # which writes a SpendLogs row, so it must meter too or that traffic is lost. - - it: meters the backend as well, since it serves the MCP transport - template: backend/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - - it: leaves the backend alone when metering is off - template: backend/deployment.yaml - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - # The migrations job runs prisma and serves no traffic; it must never receive - # the client key. - - it: never mounts the billing cert on the migrations job - template: migrations-job.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - isNull: - path: spec.template.spec.volumes - - # The conventional Secret name is the default, so enabling metering needs no - # secretName at all; the guard below only fires on an explicitly blanked one. - - it: uses the conventional secret name by default + - it: requires the token URL when enabled template: gateway/deployment.yaml set: billingMetrics: enabled: true + tokenUrl: "" asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls + - failedTemplate: + errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true - - it: fails loudly when the secretName is explicitly blanked + - it: requires the secret name when enabled template: gateway/deployment.yaml set: billingMetrics: @@ -235,15 +132,4 @@ tests: secretName: "" asserts: - failedTemplate: - errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key) - - - it: fails loudly when enabled without an endpoint - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - endpoint: "" - secretName: billing-mtls - asserts: - - failedTemplate: - errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true + errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true diff --git a/helm/litellm/tests/collector_tests.yaml b/helm/litellm/tests/collector_tests.yaml index 4ef7e3c8ca4..8e9804dece8 100644 --- a/helm/litellm/tests/collector_tests.yaml +++ b/helm/litellm/tests/collector_tests.yaml @@ -34,7 +34,7 @@ tests: gateway.collector.enabled: true gateway.collector.bufferSize: 250 gateway.collector.onUnavailable: drop - gateway.image.tag: v1.102.0 + image.tag: v1.102.0 gateway.numWorkers: 4 database.connectionPool.enabled: true database.connectionPool.maxDbConnections: 8 @@ -80,14 +80,12 @@ tests: template: gateway/deployment.yaml - equal: path: spec.template.spec.containers[1].image - value: ghcr.io/berriai/litellm-gateway:v1.102.0 + value: ghcr.io/berriai/litellm:v1.102.0 template: gateway/deployment.yaml - equal: - path: spec.template.spec.containers[1].command + path: spec.template.spec.containers[1].args value: - - python - - -m - - litellm.proxy.collector + - collector template: gateway/deployment.yaml - contains: path: spec.template.spec.containers[1].env diff --git a/helm/litellm/tests/component_hosts_tests.yaml b/helm/litellm/tests/component_hosts_tests.yaml index b55b3b694a0..047f1e7006d 100644 --- a/helm/litellm/tests/component_hosts_tests.yaml +++ b/helm/litellm/tests/component_hosts_tests.yaml @@ -316,16 +316,16 @@ tests: path: spec.template.spec.containers[0].env[1].value value: "true" - equal: - path: spec.template.spec.containers[0].env[10].name + path: spec.template.spec.containers[0].env[12].name value: LITELLM_CORS_ORIGINS - equal: - path: spec.template.spec.containers[0].env[10].value + path: spec.template.spec.containers[0].env[12].value value: https://backend-override.example.com - equal: - path: spec.template.spec.containers[0].env[11].name + path: spec.template.spec.containers[0].env[13].name value: LITELLM_CORS_ALLOW_CREDENTIALS - equal: - path: spec.template.spec.containers[0].env[11].value + path: spec.template.spec.containers[0].env[13].value value: "false" - it: exposes the public UI origin to the gateway without enabling credentials @@ -389,10 +389,10 @@ tests: path: spec.template.spec.containers[0].env[0].value value: https://ui.example.com - equal: - path: spec.template.spec.containers[0].env[9].name + path: spec.template.spec.containers[0].env[11].name value: LITELLM_CORS_ORIGINS - equal: - path: spec.template.spec.containers[0].env[9].value + path: spec.template.spec.containers[0].env[11].value value: https://gateway-override.example.com --- diff --git a/helm/litellm/tests/extra_resources_tests.yaml b/helm/litellm/tests/extra_resources_tests.yaml new file mode 100644 index 00000000000..4935b3a5401 --- /dev/null +++ b/helm/litellm/tests/extra_resources_tests.yaml @@ -0,0 +1,45 @@ +suite: test extraResources +templates: + - extra-resources.yaml +values: + - ./values/required.yaml +tests: + - it: renders nothing by default + asserts: + - hasDocuments: + count: 0 + + - it: renders every entry as its own document and templates it against the release + set: + extraResources: + - apiVersion: v1 + kind: ConfigMap + metadata: + name: "{{ .Release.Name }}-extra" + data: + chart: "{{ .Chart.Name }}" + - apiVersion: networking.k8s.io/v1 + kind: NetworkPolicy + metadata: + name: deny-all + spec: + podSelector: {} + policyTypes: + - Ingress + asserts: + - hasDocuments: + count: 2 + - isKind: + of: ConfigMap + documentIndex: 0 + - equal: + path: metadata.name + value: RELEASE-NAME-extra + documentIndex: 0 + - equal: + path: data.chart + value: litellm + documentIndex: 0 + - isKind: + of: NetworkPolicy + documentIndex: 1 diff --git a/helm/litellm/tests/hpa_workload_metrics_tests.yaml b/helm/litellm/tests/hpa_workload_metrics_tests.yaml index a29c53c5ed9..a6c17942a11 100644 --- a/helm/litellm/tests/hpa_workload_metrics_tests.yaml +++ b/helm/litellm/tests/hpa_workload_metrics_tests.yaml @@ -201,6 +201,50 @@ tests: scrapeTimeout: 10s scheme: http + - it: carries litellm-helm annotations, relabelings and namespaceSelector onto the ServiceMonitor + template: gateway/servicemonitor.yaml + release: + namespace: llm + set: + gateway.serviceMonitor.enabled: true + gateway.metricsServer.enabled: true + gateway.serviceMonitor.annotations: + team: platform + gateway.serviceMonitor.relabelings: + - targetLabel: node + sourceLabels: [__meta_kubernetes_pod_node_name] + action: replace + gateway.serviceMonitor.namespaceSelector.matchNames: + - llm + - shared + asserts: + - equal: + path: metadata.annotations + value: + team: platform + - equal: + path: spec.endpoints[0].relabelings + value: + - targetLabel: node + sourceLabels: [__meta_kubernetes_pod_node_name] + action: replace + - equal: + path: spec.namespaceSelector.matchNames + value: + - llm + - shared + + - it: renders no ServiceMonitor annotations or relabelings by default + template: gateway/servicemonitor.yaml + set: + gateway.serviceMonitor.enabled: true + gateway.metricsServer.enabled: true + asserts: + - notExists: + path: metadata.annotations + - notExists: + path: spec.endpoints[0].relabelings + - it: honours a custom scrape interval template: gateway/servicemonitor.yaml set: diff --git a/helm/litellm/tests/ingress_monolith_tests.yaml b/helm/litellm/tests/ingress_monolith_tests.yaml new file mode 100644 index 00000000000..08f7c598260 --- /dev/null +++ b/helm/litellm/tests/ingress_monolith_tests.yaml @@ -0,0 +1,164 @@ +suite: test ingress routing in monolith mode +templates: + - ingress.yaml +values: + - ./values/required.yaml +tests: + - it: monolith mode routes all component hosts to the proxy Service + set: + monolith.enabled: true + ingress.enabled: true + gateway.enabled: false + backend.enabled: false + ui.enabled: false + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + gateway.service.port: 8080 + asserts: + - lengthEqual: + path: spec.rules + count: 3 + - equal: + path: spec.rules[0].http.paths[0].backend.service + value: + name: RELEASE-NAME-litellm + port: + number: 8080 + - equal: + path: spec.rules[1].http.paths[0].backend.service + value: + name: RELEASE-NAME-litellm + port: + number: 8080 + - equal: + path: spec.rules[2].http.paths[0].backend.service + value: + name: RELEASE-NAME-litellm + port: + number: 8080 + + - it: componentized mode still splits traffic across ui, gateway and backend Services + set: + ingress.enabled: true + asserts: + - equal: + path: spec.rules[0].http.paths[0].backend.service.name + value: RELEASE-NAME-litellm-ui + - contains: + path: spec.rules[0].http.paths + content: + path: /v1/chat + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm-gateway + port: + number: 4000 + - equal: + path: spec.rules[0].http.paths[-1].backend.service.name + value: RELEASE-NAME-litellm-backend + + - it: monolith mode routes a single catch-all to the monolith Service and nothing else + set: + monolith.enabled: true + ingress.enabled: true + asserts: + - lengthEqual: + path: spec.rules[0].http.paths + count: 1 + - equal: + path: spec.rules[0].http.paths[0] + value: + path: / + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm + port: + number: 4000 + - notMatchRegexRaw: + pattern: RELEASE-NAME-litellm-(gateway|backend|ui) + + - it: monolith mode uses the gateway service port for the catch-all + set: + monolith.enabled: true + ingress.enabled: true + gateway.service.port: 8080 + asserts: + - equal: + path: spec.rules[0].http.paths[0].backend.service.port.number + value: 8080 + + - it: monolith mode sends every extraPaths entry to the monolith Service regardless of its service field + set: + monolith.enabled: true + ingress.enabled: true + ingress.extraPaths: + - path: /watsonx + - path: /admin + service: backend + pathType: Exact + - path: /dashboard + service: ui + asserts: + - lengthEqual: + path: spec.rules[0].http.paths + count: 4 + - equal: + path: spec.rules[0].http.paths[0] + value: + path: /watsonx + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm + port: + number: 4000 + - equal: + path: spec.rules[0].http.paths[1].backend.service.name + value: RELEASE-NAME-litellm + - equal: + path: spec.rules[0].http.paths[1].pathType + value: Exact + - equal: + path: spec.rules[0].http.paths[2].backend.service.name + value: RELEASE-NAME-litellm + - equal: + path: spec.rules[0].http.paths[3].path + value: / + + - it: monolith mode keeps rejecting an extra path at the root + set: + monolith.enabled: true + ingress.enabled: true + ingress.extraPaths: + - path: / + asserts: + - failedTemplate: + errorPattern: "ingress.extraPaths\\[0\\]: path / is already routed" + + - it: monolith mode keeps the host, TLS and class settings + set: + monolith.enabled: true + ingress.enabled: true + ingress.className: nginx + ingress.controller: nginx + ingress.host: llm.example.com + ingress.tls: + - hosts: + - llm.example.com + secretName: llm-tls + asserts: + - equal: + path: spec.ingressClassName + value: nginx + - equal: + path: spec.rules[0].host + value: llm.example.com + - equal: + path: spec.tls[0].secretName + value: llm-tls + - equal: + path: spec.rules[0].http.paths[0].backend.service.name + value: RELEASE-NAME-litellm diff --git a/helm/litellm/tests/keda_tests.yaml b/helm/litellm/tests/keda_tests.yaml new file mode 100644 index 00000000000..1f48654fa73 --- /dev/null +++ b/helm/litellm/tests/keda_tests.yaml @@ -0,0 +1,266 @@ +suite: test KEDA ScaledObject +templates: + - gateway/keda.yaml + - gateway/hpa.yaml + - gateway/deployment.yaml + - gateway/configmap.yaml +values: + - ./values/required.yaml +release: + name: rel + namespace: llm +tests: + - it: renders no ScaledObject by default + template: gateway/keda.yaml + asserts: + - hasDocuments: + count: 0 + + - it: passes user triggers through and adds no prometheus triggers by default + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.triggers: + - type: cpu + metricType: Utilization + metadata: + value: "60" + templates: + - gateway/keda.yaml + - gateway/hpa.yaml + asserts: + - hasDocuments: + count: 0 + template: gateway/hpa.yaml + - isKind: + of: ScaledObject + template: gateway/keda.yaml + - equal: + path: metadata.name + value: rel-litellm-gateway + template: gateway/keda.yaml + - equal: + path: spec.scaleTargetRef + value: + apiVersion: apps/v1 + kind: Deployment + name: rel-litellm-gateway + template: gateway/keda.yaml + - equal: + path: spec.minReplicaCount + value: 1 + template: gateway/keda.yaml + - equal: + path: spec.maxReplicaCount + value: 10 + template: gateway/keda.yaml + - equal: + path: spec.triggers + value: + - type: cpu + metricType: Utilization + metadata: + value: "60" + template: gateway/keda.yaml + + - it: leaves spec.replicas off the Deployment so KEDA owns the replica count + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.replicaCount: 3 + template: gateway/deployment.yaml + asserts: + - notExists: + path: spec.replicas + + - it: scales on requests per second against the metrics Service job + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.prometheus.serverAddress: http://prometheus-operated.monitoring.svc:9090 + gateway.metricsServer.enabled: true + gateway.keda.prometheus.targetRequestsPerSecond: 90 + template: gateway/keda.yaml + asserts: + - lengthEqual: + path: spec.triggers + count: 1 + - equal: + path: spec.triggers[0] + value: + type: prometheus + metricType: AverageValue + metadata: + serverAddress: http://prometheus-operated.monitoring.svc:9090 + threshold: "90" + query: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm-gateway-metrics"}[1m])) + + - it: scales on tokens per second on its own + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.prometheus.serverAddress: http://prom:9090 + gateway.metricsServer.enabled: true + gateway.keda.prometheus.targetTokensPerSecond: 6000000 + template: gateway/keda.yaml + asserts: + - lengthEqual: + path: spec.triggers + count: 1 + - equal: + path: spec.triggers[0].type + value: prometheus + - equal: + path: spec.triggers[0].metadata.threshold + value: "6000000" + - equal: + path: spec.triggers[0].metadata.query + value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm-gateway-metrics"}[1m])) + + - it: appends the requests and tokens triggers after user triggers + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.metricsServer.enabled: true + gateway.keda.triggers: + - type: cpu + metricType: Utilization + metadata: + value: "60" + gateway.keda.prometheus.serverAddress: http://prom:9090 + gateway.keda.prometheus.targetRequestsPerSecond: 90 + gateway.keda.prometheus.targetTokensPerSecond: 6000000 + template: gateway/keda.yaml + asserts: + - lengthEqual: + path: spec.triggers + count: 3 + - equal: + path: spec.triggers[0].type + value: cpu + - equal: + path: spec.triggers[1].metadata.threshold + value: "90" + - equal: + path: spec.triggers[2].metadata.threshold + value: "6000000" + - notMatchRegexRaw: + pattern: "\\* *60|per_minute|PerMinute" + + - it: fills a partial fallback with the litellm-helm defaults + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.maxReplicaCount: 20 + gateway.keda.fallback: + replicas: 4 + template: gateway/keda.yaml + asserts: + - equal: + path: spec.fallback + value: + failureThreshold: 3 + replicas: 4 + + - it: defaults fallback replicas to maxReplicaCount and keeps an explicit zero threshold + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.maxReplicaCount: 20 + gateway.keda.fallback: + failureThreshold: 0 + template: gateway/keda.yaml + asserts: + - equal: + path: spec.fallback + value: + failureThreshold: 0 + replicas: 20 + + - it: renders advanced, fallback, polling and cooldown settings + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.minReplicaCount: 2 + gateway.keda.maxReplicaCount: 20 + gateway.keda.pollingInterval: 15 + gateway.keda.cooldownPeriod: 120 + gateway.keda.fallback: + failureThreshold: 3 + replicas: 4 + gateway.keda.advanced: + restoreToOriginalReplicaCount: true + template: gateway/keda.yaml + asserts: + - equal: + path: spec.minReplicaCount + value: 2 + - equal: + path: spec.maxReplicaCount + value: 20 + - equal: + path: spec.pollingInterval + value: 15 + - equal: + path: spec.cooldownPeriod + value: 120 + - equal: + path: spec.fallback + value: + failureThreshold: 3 + replicas: 4 + - equal: + path: spec.advanced.restoreToOriginalReplicaCount + value: true + + - it: refuses a workload target without a prometheus server address + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.prometheus.targetRequestsPerSecond: 90 + template: gateway/keda.yaml + asserts: + - failedTemplate: + errorMessage: gateway.keda.prometheus.serverAddress is required when gateway.keda.prometheus.targetRequestsPerSecond or targetTokensPerSecond is set + + - it: refuses a workload target without the metrics sidecar the trigger queries + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.prometheus.serverAddress: http://prom:9090 + gateway.keda.prometheus.targetRequestsPerSecond: 90 + template: gateway/keda.yaml + asserts: + - failedTemplate: + errorMessage: gateway.metricsServer.enabled is required when gateway.keda.prometheus.targetRequestsPerSecond or targetTokensPerSecond is set + + - it: refuses to run the HPA and KEDA against the same Deployment + set: + gateway.hpa.enabled: true + gateway.keda.enabled: true + template: gateway/deployment.yaml + asserts: + - failedTemplate: + errorPattern: gateway.hpa.enabled and gateway.keda.enabled are mutually exclusive + + - it: renders scaledObject annotations, such as the KEDA pause annotation, on the ScaledObject + template: gateway/keda.yaml + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.scaledObject.annotations: + autoscaling.keda.sh/paused-replicas: "2" + asserts: + - equal: + path: metadata.annotations + value: + autoscaling.keda.sh/paused-replicas: "2" + + - it: renders no ScaledObject annotations by default + template: gateway/keda.yaml + set: + gateway.hpa.enabled: false + gateway.keda.enabled: true + asserts: + - notExists: + path: metadata.annotations diff --git a/helm/litellm/tests/liteadmin_tests.yaml b/helm/litellm/tests/liteadmin_tests.yaml new file mode 100644 index 00000000000..429896ba9b4 --- /dev/null +++ b/helm/litellm/tests/liteadmin_tests.yaml @@ -0,0 +1,149 @@ +suite: test the LiteAdmin worker and its proxy wiring +templates: + - liteadmin.yaml + - gateway/configmap.yaml + - gateway/deployment.yaml + - backend/deployment.yaml + - monolith/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: renders nothing when liteadmin is disabled + template: liteadmin.yaml + asserts: + - hasDocuments: + count: 0 + + - it: enabled renders a Deployment, Service and PVC on the shared image + set: + liteadmin.enabled: true + liteadmin.existingSecret: liteadmin-secret + liteadmin.gatewayUrl: http://gateway.example.com + liteadmin.model: gpt-4o + template: liteadmin.yaml + asserts: + - hasDocuments: + count: 3 + - isKind: + of: Deployment + documentIndex: 0 + - isKind: + of: Service + documentIndex: 1 + - isKind: + of: PersistentVolumeClaim + documentIndex: 2 + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:1.104.0 + documentIndex: 0 + - equal: + path: spec.template.spec.containers[0].args + value: ["--admin-agent"] + documentIndex: 0 + + - it: refuses to render without liteadmin.existingSecret + set: + liteadmin.enabled: true + liteadmin.gatewayUrl: http://gateway.example.com + liteadmin.model: gpt-4o + templates: + - backend/deployment.yaml + template: backend/deployment.yaml + asserts: + - failedTemplate: + errorPattern: liteadmin.existingSecret is required + + - it: wires the liteadmin env onto the backend pod in componentized mode + set: + liteadmin.enabled: true + liteadmin.existingSecret: liteadmin-secret + liteadmin.gatewayUrl: http://gateway.example.com + liteadmin.model: gpt-4o + template: backend/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_ADMIN_AGENT_URL + value: http://RELEASE-NAME-litellm-liteadmin:10000 + - contains: + path: spec.template.spec.containers[0].env + content: + name: ADMIN_AGENT_SERVICE_TOKEN + valueFrom: + secretKeyRef: + name: liteadmin-secret + key: ADMIN_AGENT_SERVICE_TOKEN + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROXY_BASE_URL + value: http://gateway.example.com + + - it: leaves the liteadmin env off the gateway pod in componentized mode + set: + liteadmin.enabled: true + liteadmin.existingSecret: liteadmin-secret + liteadmin.gatewayUrl: http://gateway.example.com + liteadmin.model: gpt-4o + template: gateway/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_ADMIN_AGENT_URL + - notContains: + path: spec.template.spec.containers[0].env + content: + name: ADMIN_AGENT_SERVICE_TOKEN + + - it: wires the liteadmin env onto the monolith pod in monolith mode + set: + monolith.enabled: true + liteadmin.enabled: true + liteadmin.existingSecret: liteadmin-secret + liteadmin.gatewayUrl: http://gateway.example.com + liteadmin.model: gpt-4o + template: monolith/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_ADMIN_AGENT_URL + value: http://RELEASE-NAME-litellm-liteadmin:10000 + - contains: + path: spec.template.spec.containers[0].env + content: + name: ADMIN_AGENT_SERVICE_TOKEN + valueFrom: + secretKeyRef: + name: liteadmin-secret + key: ADMIN_AGENT_SERVICE_TOKEN + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROXY_BASE_URL + value: http://gateway.example.com + + - it: a user-supplied PROXY_BASE_URL in extraEnv wins over liteadmin.gatewayUrl + set: + liteadmin.enabled: true + liteadmin.existingSecret: liteadmin-secret + liteadmin.gatewayUrl: http://gateway.example.com + liteadmin.model: gpt-4o + backend.extraEnv: + - name: PROXY_BASE_URL + value: https://custom.example.com + template: backend/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: PROXY_BASE_URL + value: http://gateway.example.com + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROXY_BASE_URL + value: https://custom.example.com diff --git a/helm/litellm/tests/masterkey_secret_tests.yaml b/helm/litellm/tests/masterkey_secret_tests.yaml new file mode 100644 index 00000000000..8dd261053e4 --- /dev/null +++ b/helm/litellm/tests/masterkey_secret_tests.yaml @@ -0,0 +1,156 @@ +suite: test the generated master key Secret +templates: + - secret-masterkey.yaml + - gateway/deployment.yaml + - backend/deployment.yaml + - gateway/configmap.yaml + - monolith/deployment.yaml + - migrations-job.yaml +values: + - ./values/required.yaml +tests: + - it: renders no Secret when the operator supplies masterKey.secretName + template: secret-masterkey.yaml + asserts: + - hasDocuments: + count: 0 + + - it: refuses to render the workload without a Secret name when masterKey.generate is off + set: + masterKey.secretName: "" + template: gateway/deployment.yaml + asserts: + - failedTemplate: + errorPattern: masterKey.secretName is required + + - it: generates a Secret whose key starts with sk- (base64 c2st) and keeps it across uninstall + set: + masterKey.secretName: "" + masterKey.generate: true + template: secret-masterkey.yaml + asserts: + - isKind: + of: Secret + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-masterkey + - equal: + path: metadata.annotations["helm.sh/resource-policy"] + value: keep + - matchRegex: + path: data["master-key"] + pattern: ^c2st + + - it: stores the value under masterKey.secretKey + set: + masterKey.secretName: "" + masterKey.generate: true + masterKey.secretKey: LITELLM_MASTER_KEY + template: secret-masterkey.yaml + asserts: + - matchRegex: + path: data.LITELLM_MASTER_KEY + pattern: ^c2st + + - it: reuses the master key already stored in the cluster instead of generating a new one on upgrade + set: + masterKey.secretName: "" + masterKey.generate: true + template: secret-masterkey.yaml + kubernetesProvider: + scheme: + "v1/Secret": + gvr: + version: "v1" + resource: "secrets" + namespaced: true + objects: + - kind: Secret + apiVersion: v1 + metadata: + name: RELEASE-NAME-litellm-masterkey + namespace: NAMESPACE + data: + master-key: c2stZXhpc3Rpbmcta2V5 + asserts: + - equal: + path: data["master-key"] + value: c2stZXhpc3Rpbmcta2V5 + + - it: points the app containers at the generated Secret and keeps it out of the pre-install migrations Job + set: + masterKey.secretName: "" + masterKey.generate: true + monolith.enabled: true + templates: + - monolith/deployment.yaml + - migrations-job.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_MASTER_KEY + valueFrom: + secretKeyRef: + name: RELEASE-NAME-litellm-masterkey + key: master-key + template: monolith/deployment.yaml + - equal: + path: metadata.annotations["helm.sh/hook"] + value: pre-install,pre-upgrade + template: migrations-job.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_MASTER_KEY + any: true + template: migrations-job.yaml + + - it: an explicit masterKey.secretName wins over masterKey.generate + set: + masterKey.secretName: operator-secret + masterKey.generate: true + templates: + - secret-masterkey.yaml + - gateway/deployment.yaml + asserts: + - hasDocuments: + count: 0 + template: secret-masterkey.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_MASTER_KEY + valueFrom: + secretKeyRef: + name: operator-secret + key: master-key + template: gateway/deployment.yaml + + - it: exposes the key as PROXY_MASTER_KEY too, so a proxy_config copied from litellm-helm keeps resolving + set: + monolith.enabled: true + templates: + - monolith/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROXY_MASTER_KEY + valueFrom: + secretKeyRef: + name: litellm-master-key-secret + key: master-key + - it: exposes PROXY_MASTER_KEY on the componentized gateway and backend + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROXY_MASTER_KEY + valueFrom: + secretKeyRef: + name: litellm-master-key-secret + key: master-key diff --git a/helm/litellm/tests/metrics_server_tests.yaml b/helm/litellm/tests/metrics_server_tests.yaml index 0e7d9d9e9ee..1a48b7be671 100644 --- a/helm/litellm/tests/metrics_server_tests.yaml +++ b/helm/litellm/tests/metrics_server_tests.yaml @@ -38,7 +38,7 @@ tests: gateway.metricsServer.enabled: true gateway.metricsServer.port: 4101 gateway.service.type: LoadBalancer - gateway.image.tag: v1.101.0 + image.tag: v1.101.0 asserts: - contains: path: spec.template.spec.containers[0].env @@ -58,14 +58,12 @@ tests: template: gateway/deployment.yaml - equal: path: spec.template.spec.containers[1].image - value: ghcr.io/berriai/litellm-gateway:v1.101.0 + value: ghcr.io/berriai/litellm:v1.101.0 template: gateway/deployment.yaml - equal: - path: spec.template.spec.containers[1].command + path: spec.template.spec.containers[1].args value: - - python - - -m - - litellm.proxy.prometheus_metrics_server + - metrics - --port - "4101" template: gateway/deployment.yaml diff --git a/helm/litellm/tests/migration_job_tests.yaml b/helm/litellm/tests/migration_job_tests.yaml index 2ebb1b44926..86251e8a1d2 100644 --- a/helm/litellm/tests/migration_job_tests.yaml +++ b/helm/litellm/tests/migration_job_tests.yaml @@ -91,7 +91,7 @@ tests: app.kubernetes.io/name: litellm app.kubernetes.io/instance: RELEASE-NAME app.kubernetes.io/managed-by: Helm - helm.sh/chart: litellm-0.1.0 + helm.sh/chart: litellm-1.0.0 app.kubernetes.io/component: migrations - it: renders pod-level and container-level securityContext in their own scopes @@ -254,3 +254,76 @@ tests: path: spec.template.spec.nodeSelector - isNull: path: spec.template.spec.tolerations + + - it: sets jobLabels and jobAnnotations on the Job metadata + set: + migrationJob.jobLabels: + environment: production + migrationJob.jobAnnotations: + example.com/cost-center: "1234" + asserts: + - equal: + path: metadata.labels.environment + value: production + - equal: + path: metadata.annotations['example.com/cost-center'] + value: "1234" + - equal: + path: metadata.annotations['helm.sh/hook'] + value: pre-install,pre-upgrade + + - it: renders jobAnnotations when both hooks are off + set: + migrationJob.hooks.helm.enabled: false + migrationJob.jobAnnotations: + example.com/cost-center: "1234" + asserts: + - equal: + path: metadata.annotations + value: + example.com/cost-center: "1234" + + - it: runs the image entrypoint with the migrations component by default + asserts: + - notExists: + path: spec.template.spec.containers[0].command + - equal: + path: spec.template.spec.containers[0].args + value: [migrations] + + - it: overrides the container command and args + set: + migrationJob.command: [sh] + migrationJob.args: [-c, echo migrating] + asserts: + - equal: + path: spec.template.spec.containers[0].command + value: [sh] + - equal: + path: spec.template.spec.containers[0].args + value: [-c, echo migrating] + + - it: clears the container args when only command is set + set: + migrationJob.command: [sh] + asserts: + - equal: + path: spec.template.spec.containers[0].command + value: [sh] + - equal: + path: spec.template.spec.containers[0].args + value: [] + + - it: puts migrationJob.annotations on the Job pod, as litellm-helm did + set: + migrationJob.annotations: + sidecar.istio.io/inject: "false" + asserts: + - equal: + path: spec.template.metadata.annotations + value: + sidecar.istio.io/inject: "false" + - it: renders no pod annotations by default + asserts: + - notExists: + path: spec.template.metadata.annotations diff --git a/helm/litellm/tests/monolith_tests.yaml b/helm/litellm/tests/monolith_tests.yaml new file mode 100644 index 00000000000..cb8aca608ec --- /dev/null +++ b/helm/litellm/tests/monolith_tests.yaml @@ -0,0 +1,539 @@ +suite: test monolith mode +templates: + - monolith/deployment.yaml + - monolith/service.yaml + - monolith/hpa.yaml + - monolith/keda.yaml + - monolith/poddisruptionbudget.yaml + - monolith/service-metrics.yaml + - monolith/servicemonitor.yaml + - gateway/deployment.yaml + - gateway/service.yaml + - gateway/hpa.yaml + - gateway/keda.yaml + - gateway/poddisruptionbudget.yaml + - gateway/service-metrics.yaml + - gateway/servicemonitor.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - backend/service.yaml + - backend/hpa.yaml + - backend/poddisruptionbudget.yaml + - ui/deployment.yaml + - ui/service.yaml + - ui/hpa.yaml + - ui/poddisruptionbudget.yaml + - serviceaccount.yaml + - migrations-job.yaml +values: + - ./values/required.yaml +tests: + - it: monolith ignores a leftover invalid UI public URL + set: + monolith.enabled: true + ui.publicUrl: not-an-origin + template: monolith/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_CORS_ORIGINS + any: true + + - it: renders no monolith resources by default + templates: + - monolith/deployment.yaml + - monolith/service.yaml + - monolith/hpa.yaml + - monolith/poddisruptionbudget.yaml + - monolith/service-metrics.yaml + - monolith/servicemonitor.yaml + asserts: + - hasDocuments: + count: 0 + + - it: componentized mode renders one Deployment per component + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + - ui/deployment.yaml + asserts: + - hasDocuments: + count: 1 + - isKind: + of: Deployment + + - it: monolith mode renders exactly one Deployment, running the proxy dispatcher + set: + monolith.enabled: true + templates: + - monolith/deployment.yaml + - gateway/deployment.yaml + - backend/deployment.yaml + - ui/deployment.yaml + asserts: + - hasDocuments: + count: 1 + template: monolith/deployment.yaml + - hasDocuments: + count: 0 + template: gateway/deployment.yaml + - hasDocuments: + count: 0 + template: backend/deployment.yaml + - hasDocuments: + count: 0 + template: ui/deployment.yaml + - isKind: + of: Deployment + template: monolith/deployment.yaml + - equal: + path: metadata.name + value: RELEASE-NAME-litellm + template: monolith/deployment.yaml + - equal: + path: metadata.labels["app.kubernetes.io/component"] + value: proxy + template: monolith/deployment.yaml + - equal: + path: spec.template.spec.containers[0].name + value: proxy + template: monolith/deployment.yaml + - equal: + path: spec.template.spec.containers[0].args[0] + value: proxy + template: monolith/deployment.yaml + - equal: + path: spec.template.spec.containers[0].args + value: + - proxy + - --port + - "4000" + - --config + - /app/config/config.yaml + template: monolith/deployment.yaml + - notExists: + path: spec.template.spec.containers[0].command + template: monolith/deployment.yaml + + - it: monolith appends monolith.extraArgs after the chart's own proxy arguments + set: + monolith.enabled: true + monolith.extraArgs: + - --detailed_debug + - --run_granian + template: monolith/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].args + value: + - proxy + - --port + - "4000" + - --config + - /app/config/config.yaml + - --detailed_debug + - --run_granian + + - it: monolith drops the --config pair when the chart does not create the ConfigMap + set: + monolith.enabled: true + gateway.config.create: false + template: monolith/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].args + value: + - proxy + - --port + - "4000" + - notExists: + path: spec.template.spec.volumes + + - it: monolith mounts the proxy config ConfigMap the chart renders under the monolith name + set: + monolith.enabled: true + gateway.config.proxy_config: + model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + templates: + - monolith/deployment.yaml + - gateway/configmap.yaml + asserts: + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-config + template: gateway/configmap.yaml + - contains: + path: spec.template.spec.volumes + content: + name: gateway-config + configMap: + name: RELEASE-NAME-litellm-config + template: monolith/deployment.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: CONFIG_FILE_PATH + value: /app/config/config.yaml + template: monolith/deployment.yaml + + - it: monolith mode renders exactly one Service and no component Services + set: + monolith.enabled: true + templates: + - monolith/service.yaml + - gateway/service.yaml + - backend/service.yaml + - ui/service.yaml + asserts: + - hasDocuments: + count: 1 + template: monolith/service.yaml + - hasDocuments: + count: 0 + template: gateway/service.yaml + - hasDocuments: + count: 0 + template: backend/service.yaml + - hasDocuments: + count: 0 + template: ui/service.yaml + - equal: + path: metadata.name + value: RELEASE-NAME-litellm + template: monolith/service.yaml + - equal: + path: spec.selector + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: proxy + template: monolith/service.yaml + - equal: + path: spec.ports[0].port + value: 4000 + template: monolith/service.yaml + + - it: monolith Service selector matches the monolith pod labels + set: + monolith.enabled: true + templates: + - monolith/deployment.yaml + asserts: + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: proxy + - equal: + path: spec.template.metadata.labels["app.kubernetes.io/component"] + value: proxy + + - it: monolith reuses gateway resources, replicas, numWorkers and probes + set: + monolith.enabled: true + gateway.replicaCount: 3 + gateway.hpa.enabled: false + gateway.numWorkers: 6 + gateway.resources: + requests: + cpu: "2" + memory: 4Gi + template: monolith/deployment.yaml + asserts: + - equal: + path: spec.replicas + value: 3 + - equal: + path: spec.template.spec.containers[0].resources.requests.cpu + value: "2" + - contains: + path: spec.template.spec.containers[0].env + content: + name: NUM_WORKERS + value: "6" + - equal: + path: spec.template.spec.containers[0].readinessProbe.httpGet.path + value: /health/readiness + - equal: + path: spec.template.spec.containers[0].livenessProbe.httpGet.path + value: /health/liveliness + + - it: monolith takes the rollout settings a litellm-helm install carried from gateway.minReadySeconds and gateway.strategy + set: + monolith.enabled: true + gateway.minReadySeconds: 60 + gateway.strategy: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 0 + maxSurge: 1 + backend.minReadySeconds: 5 + template: monolith/deployment.yaml + asserts: + - equal: + path: spec.minReadySeconds + value: 60 + - equal: + path: spec.strategy + value: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 0 + maxSurge: 1 + + - it: monolith ignores backend and ui replica overrides + set: + monolith.enabled: true + gateway.hpa.enabled: false + gateway.replicaCount: 2 + backend.replicaCount: 9 + ui.replicaCount: 9 + template: monolith/deployment.yaml + asserts: + - equal: + path: spec.replicas + value: 2 + + - it: monolith HPA targets the monolith Deployment and the gateway HPA is not rendered + set: + monolith.enabled: true + gateway.hpa.enabled: true + templates: + - monolith/hpa.yaml + - gateway/hpa.yaml + - backend/hpa.yaml + - ui/hpa.yaml + asserts: + - hasDocuments: + count: 1 + template: monolith/hpa.yaml + - hasDocuments: + count: 0 + template: gateway/hpa.yaml + - hasDocuments: + count: 0 + template: backend/hpa.yaml + - hasDocuments: + count: 0 + template: ui/hpa.yaml + - equal: + path: spec.scaleTargetRef + value: + apiVersion: apps/v1 + kind: Deployment + name: RELEASE-NAME-litellm + template: monolith/hpa.yaml + + - it: monolith KEDA ScaledObject targets the monolith Deployment + set: + monolith.enabled: true + gateway.hpa.enabled: false + gateway.keda.enabled: true + gateway.keda.prometheus.serverAddress: http://prometheus:9090 + gateway.metricsServer.enabled: true + gateway.keda.prometheus.targetRequestsPerSecond: 50 + templates: + - monolith/keda.yaml + - gateway/keda.yaml + asserts: + - hasDocuments: + count: 1 + template: monolith/keda.yaml + - hasDocuments: + count: 0 + template: gateway/keda.yaml + - equal: + path: spec.scaleTargetRef.name + value: RELEASE-NAME-litellm + template: monolith/keda.yaml + - equal: + path: spec.triggers[0].metadata.query + value: sum(rate(litellm_proxy_total_requests_metric_total{namespace="NAMESPACE",job="RELEASE-NAME-litellm-metrics"}[1m])) + template: monolith/keda.yaml + + - it: monolith PDB selects the monolith pods + set: + monolith.enabled: true + gateway.pdb.enabled: true + gateway.pdb.maxUnavailable: 1 + templates: + - monolith/poddisruptionbudget.yaml + - gateway/poddisruptionbudget.yaml + - backend/poddisruptionbudget.yaml + - ui/poddisruptionbudget.yaml + asserts: + - hasDocuments: + count: 1 + template: monolith/poddisruptionbudget.yaml + - hasDocuments: + count: 0 + template: gateway/poddisruptionbudget.yaml + - hasDocuments: + count: 0 + template: backend/poddisruptionbudget.yaml + - hasDocuments: + count: 0 + template: ui/poddisruptionbudget.yaml + - equal: + path: spec.maxUnavailable + value: 1 + template: monolith/poddisruptionbudget.yaml + - equal: + path: spec.selector.matchLabels["app.kubernetes.io/component"] + value: proxy + template: monolith/poddisruptionbudget.yaml + + - it: monolith metrics sidecar, metrics Service and ServiceMonitor follow the monolith name + set: + monolith.enabled: true + gateway.metricsServer.enabled: true + gateway.serviceMonitor.enabled: true + templates: + - monolith/deployment.yaml + - monolith/service-metrics.yaml + - monolith/servicemonitor.yaml + - gateway/service-metrics.yaml + - gateway/servicemonitor.yaml + asserts: + - hasDocuments: + count: 0 + template: gateway/service-metrics.yaml + - hasDocuments: + count: 0 + template: gateway/servicemonitor.yaml + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-metrics + template: monolith/service-metrics.yaml + - equal: + path: spec.selector["app.kubernetes.io/component"] + value: proxy + template: monolith/service-metrics.yaml + - equal: + path: spec.selector.matchLabels["app.kubernetes.io/component"] + value: proxy + template: monolith/servicemonitor.yaml + - equal: + path: spec.template.spec.containers[1].name + value: metrics + template: monolith/deployment.yaml + - equal: + path: spec.template.spec.containers[1].args[0] + value: metrics + template: monolith/deployment.yaml + + - it: monolith collector sidecar runs the collector dispatcher from the same image + set: + monolith.enabled: true + gateway.collector.enabled: true + template: monolith/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: collector + - equal: + path: spec.template.spec.containers[1].args + value: + - collector + - equal: + path: spec.template.spec.containers[1].image + value: ghcr.io/berriai/litellm:1.104.0 + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_COLLECTOR_ENABLED + value: "true" + + - it: monolith renders only the gateway ServiceAccount and runs the pod with it + set: + monolith.enabled: true + serviceAccounts.gateway.create: true + serviceAccounts.backend.create: true + serviceAccounts.ui.create: true + templates: + - serviceaccount.yaml + - monolith/deployment.yaml + asserts: + - hasDocuments: + count: 1 + template: serviceaccount.yaml + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-gateway + template: serviceaccount.yaml + - equal: + path: spec.template.spec.serviceAccountName + value: RELEASE-NAME-litellm-gateway + template: monolith/deployment.yaml + + - it: the migrations Job renders the same way in monolith mode + set: + monolith.enabled: true + template: migrations-job.yaml + asserts: + - isKind: + of: Job + - equal: + path: spec.template.spec.containers[0].args + value: + - migrations + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:1.104.0 + + - it: monolith gets the writer user under both names the proxy CLI and the componentized entrypoints read + set: + monolith.enabled: true + template: monolith/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_USER + valueFrom: + secretKeyRef: + name: litellm-writer-secret + key: username + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_USERNAME + valueFrom: + secretKeyRef: + name: litellm-writer-secret + key: username + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_NAME + value: litellm + + - it: disabling every component without monolith renders no application workload + set: + gateway.enabled: false + backend.enabled: false + ui.enabled: false + templates: + - monolith/deployment.yaml + - gateway/deployment.yaml + - backend/deployment.yaml + - ui/deployment.yaml + asserts: + - hasDocuments: + count: 0 + + - it: monolith renders even when gateway.enabled is false + set: + monolith.enabled: true + gateway.enabled: false + templates: + - monolith/deployment.yaml + - monolith/service.yaml + asserts: + - hasDocuments: + count: 1 diff --git a/helm/litellm/tests/no_bundled_datastores_tests.yaml b/helm/litellm/tests/no_bundled_datastores_tests.yaml new file mode 100644 index 00000000000..dbafa5a9a11 --- /dev/null +++ b/helm/litellm/tests/no_bundled_datastores_tests.yaml @@ -0,0 +1,54 @@ +suite: the chart ships no PostgreSQL or Redis of its own +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - migrations-job.yaml +tests: + - it: rejects the retired postgresql block instead of silently ignoring it + set: + postgresql.enabled: true + postgresql.auth.password: super-secret + asserts: + - failedTemplate: {} + + - it: rejects the retired redis.enabled flag instead of silently ignoring it + set: + redis.enabled: true + asserts: + - failedTemplate: {} + + - it: wires the writer straight from database.writer with no subchart in between + template: gateway/deployment.yaml + set: + database.writer.host: postgres.example.com + database.writer.dbname: proxydb + database.writer.passwordSecret.name: pg-secret + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_HOST + value: postgres.example.com + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_NAME + value: proxydb + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_PASSWORD + valueFrom: + secretKeyRef: + name: pg-secret + key: password + + - it: keeps the migrations Job a pre-install hook whenever the Helm hook is enabled + template: migrations-job.yaml + set: + database.writer.host: postgres.example.com + database.writer.dbname: litellm + asserts: + - equal: + path: metadata.annotations["helm.sh/hook"] + value: pre-install,pre-upgrade diff --git a/helm/litellm/tests/redis_env_tests.yaml b/helm/litellm/tests/redis_env_tests.yaml index 79a14dfc6e5..849f13e5438 100644 --- a/helm/litellm/tests/redis_env_tests.yaml +++ b/helm/litellm/tests/redis_env_tests.yaml @@ -225,10 +225,10 @@ tests: name: REDIS_CLUSTER_NODES value: '[{"host":"$(REDIS_HOST)","port":6379}]' - equal: - path: spec.template.spec.containers[0].env[8].name + path: spec.template.spec.containers[0].env[10].name value: REDIS_HOST - equal: - path: spec.template.spec.containers[0].env[11].name + path: spec.template.spec.containers[0].env[13].name value: REDIS_CLUSTER_NODES - it: gateway gets REDIS_URL from urlKey without discrete Redis env diff --git a/helm/litellm/tests/rollout_strategy_tests.yaml b/helm/litellm/tests/rollout_strategy_tests.yaml index b12e2073c7c..af6e41d038e 100644 --- a/helm/litellm/tests/rollout_strategy_tests.yaml +++ b/helm/litellm/tests/rollout_strategy_tests.yaml @@ -64,3 +64,39 @@ tests: - notExists: path: spec.strategy template: ui/deployment.yaml + + - it: leaves minReadySeconds to the Kubernetes default when unset + asserts: + - notExists: + path: spec.minReadySeconds + + - it: renders the configured minReadySeconds on each deployment + set: + gateway.minReadySeconds: 60 + backend.minReadySeconds: 30 + ui.minReadySeconds: 10 + asserts: + - equal: + path: spec.minReadySeconds + value: 60 + template: gateway/deployment.yaml + - equal: + path: spec.minReadySeconds + value: 30 + template: backend/deployment.yaml + - equal: + path: spec.minReadySeconds + value: 10 + template: ui/deployment.yaml + + - it: renders an explicit minReadySeconds of 0 instead of dropping it + set: + gateway.minReadySeconds: 0 + asserts: + - equal: + path: spec.minReadySeconds + value: 0 + template: gateway/deployment.yaml + - notExists: + path: spec.minReadySeconds + template: backend/deployment.yaml diff --git a/helm/litellm/tests/shared_image_tests.yaml b/helm/litellm/tests/shared_image_tests.yaml new file mode 100644 index 00000000000..1e7837b1973 --- /dev/null +++ b/helm/litellm/tests/shared_image_tests.yaml @@ -0,0 +1,264 @@ +suite: test the shared image and the entrypoint dispatcher args +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - ui/deployment.yaml + - monolith/deployment.yaml + - migrations-job.yaml + - templates/tests/test-connection.yaml +values: + - ./values/required.yaml +tests: + - it: every componentized container runs the shared image at the chart appVersion + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + - ui/deployment.yaml + - migrations-job.yaml + - templates/tests/test-connection.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:1.104.0 + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:1.104.0 + template: backend/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:1.104.0 + template: ui/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:1.104.0 + template: migrations-job.yaml + - equal: + path: spec.containers[0].image + value: ghcr.io/berriai/litellm:1.104.0 + template: templates/tests/test-connection.yaml + + - it: image.repository and image.tag are applied to every container, sidecars included + set: + image.repository: registry.example.com/litellm + image.tag: v9.9.9 + image.pullPolicy: Always + gateway.metricsServer.enabled: true + gateway.collector.enabled: true + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + - ui/deployment.yaml + - migrations-job.yaml + - templates/tests/test-connection.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].image + value: registry.example.com/litellm:v9.9.9 + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].image + value: registry.example.com/litellm:v9.9.9 + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[2].image + value: registry.example.com/litellm:v9.9.9 + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[0].imagePullPolicy + value: Always + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: registry.example.com/litellm:v9.9.9 + template: backend/deployment.yaml + - equal: + path: spec.template.spec.containers[0].imagePullPolicy + value: Always + template: backend/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: registry.example.com/litellm:v9.9.9 + template: ui/deployment.yaml + - equal: + path: spec.template.spec.containers[0].imagePullPolicy + value: Always + template: ui/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: registry.example.com/litellm:v9.9.9 + template: migrations-job.yaml + - equal: + path: spec.template.spec.containers[0].imagePullPolicy + value: Always + template: migrations-job.yaml + - equal: + path: spec.containers[0].image + value: registry.example.com/litellm:v9.9.9 + template: templates/tests/test-connection.yaml + + - it: the monolith and its sidecars run the shared image + set: + monolith.enabled: true + image.tag: v9.9.9 + gateway.metricsServer.enabled: true + gateway.collector.enabled: true + template: monolith/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:v9.9.9 + - equal: + path: spec.template.spec.containers[1].image + value: ghcr.io/berriai/litellm:v9.9.9 + - equal: + path: spec.template.spec.containers[2].image + value: ghcr.io/berriai/litellm:v9.9.9 + + - it: image.digest renders as repository:tag@digest everywhere + set: + image.tag: v9.9.9 + image.digest: sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + - ui/deployment.yaml + - migrations-job.yaml + - templates/tests/test-connection.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:v9.9.9@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:v9.9.9@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + template: backend/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:v9.9.9@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + template: ui/deployment.yaml + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:v9.9.9@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + template: migrations-job.yaml + - equal: + path: spec.containers[0].image + value: ghcr.io/berriai/litellm:v9.9.9@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + template: templates/tests/test-connection.yaml + + - it: image.digest with an empty tag falls back to the appVersion tag + set: + image.digest: sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + monolith.enabled: true + template: monolith/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].image + value: ghcr.io/berriai/litellm:1.104.0@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef + + - it: componentized containers select their component through args and never set command + set: + gateway.metricsServer.enabled: true + gateway.collector.enabled: true + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + - ui/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].args + value: + - gateway + - --host + - 0.0.0.0 + - --port + - "4000" + template: gateway/deployment.yaml + - notExists: + path: spec.template.spec.containers[0].command + template: gateway/deployment.yaml + - notExists: + path: spec.template.spec.containers[1].command + template: gateway/deployment.yaml + - notExists: + path: spec.template.spec.containers[2].command + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[0].args + value: + - backend + - --host + - 0.0.0.0 + - --port + - "4001" + template: backend/deployment.yaml + - notExists: + path: spec.template.spec.containers[0].command + template: backend/deployment.yaml + - equal: + path: spec.template.spec.containers[0].args + value: + - ui + template: ui/deployment.yaml + - notExists: + path: spec.template.spec.containers[0].command + template: ui/deployment.yaml + + - it: the migrations Job runs the migrations dispatcher and nothing else + template: migrations-job.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].args + value: + - migrations + - notExists: + path: spec.template.spec.containers[0].command + + - it: the Helm test pod polls the workload Service readiness endpoint through the image entrypoint + template: templates/tests/test-connection.yaml + asserts: + - isKind: + of: Pod + - equal: + path: metadata.annotations["helm.sh/hook"] + value: test + - equal: + path: spec.containers[0].env[0] + value: + name: LITELLM_READINESS_URL + value: http://RELEASE-NAME-litellm-gateway:4000/health/readiness + - equal: + path: spec.containers[0].args[0] + value: python + - equal: + path: spec.containers[0].args[1] + value: -c + - matchRegex: + path: spec.containers[0].args[2] + pattern: LITELLM_READINESS_URL + - matchRegex: + path: spec.containers[0].args[2] + pattern: 'db != "connected"' + - matchRegex: + path: spec.containers[0].args[2] + pattern: sys\.exit\(1\) + - notExists: + path: spec.containers[0].command + + - it: the Helm test pod targets the monolith Service in monolith mode + set: + monolith.enabled: true + template: templates/tests/test-connection.yaml + asserts: + - equal: + path: spec.containers[0].env[0].value + value: http://RELEASE-NAME-litellm:4000/health/readiness + + - it: the removed per-component image blocks are rejected by the values schema + set: + gateway.image.tag: v1 + template: gateway/deployment.yaml + asserts: + - failedTemplate: {} diff --git a/helm/litellm/values.schema.json b/helm/litellm/values.schema.json new file mode 100644 index 00000000000..3aafbe68fe6 --- /dev/null +++ b/helm/litellm/values.schema.json @@ -0,0 +1,183 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "LiteLLM chart values", + "type": "object", + "additionalProperties": true, + "properties": { + "image": { + "type": "object", + "additionalProperties": false, + "required": [ + "repository" + ], + "properties": { + "repository": { + "type": "string", + "minLength": 1 + }, + "tag": { + "type": "string" + }, + "pullPolicy": { + "type": "string", + "enum": [ + "Always", + "IfNotPresent", + "Never" + ] + }, + "digest": { + "type": "string", + "pattern": "^(sha256:[a-f0-9]{64})?$" + } + } + }, + "imagePullSecrets": { + "type": "array", + "items": { + "type": "object", + "required": [ + "name" + ], + "properties": { + "name": { + "type": "string" + } + } + } + }, + "monolith": { + "type": "object", + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean" + }, + "extraArgs": { + "type": "array", + "items": { + "type": "string" + } + } + } + }, + "gateway": { + "type": "object", + "additionalProperties": true, + "properties": { + "enabled": { + "type": "boolean" + }, + "replicaCount": { + "type": [ + "integer", + "null" + ], + "minimum": 0 + }, + "numWorkers": { + "type": [ + "integer", + "null" + ], + "minimum": 1 + }, + "image": false + } + }, + "backend": { + "type": "object", + "additionalProperties": true, + "properties": { + "enabled": { + "type": "boolean" + }, + "replicaCount": { + "type": [ + "integer", + "null" + ], + "minimum": 0 + }, + "image": false + } + }, + "ui": { + "type": "object", + "additionalProperties": true, + "properties": { + "enabled": { + "type": "boolean" + }, + "replicaCount": { + "type": [ + "integer", + "null" + ], + "minimum": 0 + }, + "image": false + } + }, + "migrationJob": { + "type": "object", + "additionalProperties": true, + "properties": { + "enabled": { + "type": "boolean" + }, + "image": false + } + }, + "masterKey": { + "type": "object", + "additionalProperties": true, + "properties": { + "secretName": { + "type": "string" + }, + "secretKey": { + "type": "string" + }, + "generate": { + "type": "boolean" + } + } + }, + "postgresql": false, + "redis": { + "type": "object", + "additionalProperties": true, + "properties": { + "enabled": false, + "cluster": { + "type": "boolean" + }, + "host": { + "type": "string" + }, + "port": { + "type": [ + "integer", + "string" + ] + } + } + }, + "ingress": { + "type": "object", + "additionalProperties": true, + "properties": { + "enabled": { + "type": "boolean" + } + } + }, + "extraResources": { + "type": "array", + "items": { + "type": "object" + } + } + } +} diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 92a48153c5a..9bfc3ca32e9 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -1,10 +1,64 @@ # LiteLLM helm chart values +# +# Two deployment modes, both from the single image below: +# +# componentized (default): `gateway.enabled`, `backend.enabled` and +# `ui.enabled` each render their own Deployment, Service, HPA and PDB, so +# the LLM data plane, the management API and the static dashboard scale +# independently. Their containers run `args: [gateway|backend|ui, ...]`. +# +# monolith: `monolith.enabled: true` renders ONE Deployment and Service +# running the full proxy (`args: [proxy, ...]`): gateway routes, +# management routes and the Admin UI served by the proxy itself. The +# gateway, backend and ui Deployments and Services are not rendered and +# the Ingress sends every path to the monolith Service. The monolith is +# configured through the `gateway.*` values (config, resources, probes, +# securityContext, hpa, pdb, keda, collector, metricsServer, volumes, +# env), so a monolith install only flips this flag. See `monolith` below +# for the precedence rules. nameOverride: "" fullnameOverride: "" +# The one image every container in this chart runs: gateway, backend, ui, +# monolith, the migrations Job, the metrics and collector sidecars and the +# Helm test pod. Its entrypoint dispatches on the first container argument +# (proxy, gateway, backend, ui, migrations, metrics, collector), which is why +# the templates only ever set `args` and never `command`: overriding the +# entrypoint would drop the wrappers it applies (ddtrace, Prometheus +# multiprocess dir cleanup). `tag` defaults to the chart's appVersion. When +# `digest` is set the image renders as `repository:tag@digest`, so the tag is +# informational and the digest pins the bytes. Also mirrored at +# docker.litellm.ai/berriai/litellm. +image: + repository: ghcr.io/berriai/litellm + tag: "" + pullPolicy: IfNotPresent + digest: "" + imagePullSecrets: [] +# Monolith mode. When enabled the chart renders a single Deployment named +# `-litellm` running `args: [proxy, --port, 4000, --config, +# /app/config/config.yaml]` (the --config pair only when gateway.config.create +# is true) plus `extraArgs`, and a Service of the same name. Precedence when +# this is on: +# - gateway.enabled / backend.enabled / ui.enabled are ignored: no +# component Deployment, Service, HPA, PDB or ServiceMonitor is rendered +# - every `backend.*` and `ui.*` value is ignored, including their +# ServiceAccounts; the monolith runs as serviceAccounts.gateway +# - every `gateway.*` value applies to the monolith pod instead of a gateway +# pod: config, numWorkers, resources, probes, securityContext, +# hpa / keda / pdb, metricsServer, collector, volumes, env, scheduling +# - ingress.extraPaths entries keep their `service` field but every path, +# built in or extra, targets the monolith Service +# - the migrations Job renders exactly as in componentized mode +monolith: + enabled: false + # Extra CLI arguments appended after the chart's own proxy arguments, e.g. + # ["--detailed_debug"] or ["--run_granian"]. + extraArgs: [] + # Optional Ingress wiring the three component Services behind a single L7 # entrypoint. Required when serving the static UI bundle over the network. ingress: @@ -86,10 +140,8 @@ serviceAccounts: # LiteLLM_VerificationToken, LiteLLM_SpendLogs, ...). Disable if your # pipeline runs migrations out-of-band. # -# Uses a dedicated `litellm-migrations` image (prisma CLI + the migration -# files from `litellm-proxy-extras`) instead of the backend image, so the -# Job doesn't drag in the rest of the proxy and doesn't run `prisma -# generate` — the migration engine doesn't need the generated client. +# Runs the shared image with `args: [migrations]`, which the entrypoint +# dispatches to the prisma migration runner without importing the proxy. migrationJob: enabled: true # Which controller is responsible for running the Job. @@ -150,8 +202,18 @@ migrationJob: # pod. Same shape as gateway.podSecurityContext / gateway.securityContext. podSecurityContext: {} securityContext: {} + # Extra labels and annotations on the Job object itself. + jobLabels: {} + jobAnnotations: {} # Extra pod labels on the Job pod, merged into the chart's common labels. podLabels: {} + # Annotations on the Job pod, e.g. sidecar.istio.io/inject: "false" so an + # injected sidecar cannot keep the pod running after the migration exits. + annotations: {} + # Override the container command and args. The default runs the image + # entrypoint with `migrations`; a command with no args clears the args. + command: [] + args: [] # Additional volumes on the Job pod and volumeMounts on its container, e.g. # the writable scratch space a read-only root filesystem needs. volumes: [] @@ -163,42 +225,39 @@ migrationJob: nodeSelector: {} tolerations: [] affinity: {} - image: - repository: ghcr.io/berriai/litellm-migrations - tag: "" # defaults to .Chart.AppVersion - pullPolicy: IfNotPresent + # Extra init containers and sidecars on the Job pod, rendered through `tpl`. + extraInitContainers: [] + extraContainers: [] # Extra env appended to the migration container. The migration entrypoint # uses the v2 resolver by default (no diff-and-force recovery — avoids the # schema thrashing seen during rolling deploys). To opt back into the v1 # resolver, append `- name: USE_V2_MIGRATION_RESOLVER` / `value: "false"`. extraEnv: [] -# Required: a master key used by gateway + backend to mint/verify proxy tokens. -# Must reference an existing Secret. +# Master key used by the proxy to mint and verify virtual keys. Reference an +# existing Secret through `secretName` / `secretKey`, or set `generate: true` +# and leave `secretName` empty to have the chart create +# `-litellm-masterkey` with a random `sk-...` key under `secretKey`. +# The generated Secret is looked up on upgrade so the key never rotates by +# accident (`helm template` cannot look anything up and renders a fresh key +# each time; that is expected). The chart never accepts an inline key. masterKey: secretName: litellm-master-key-secret # name of a Secret containing the master key secretKey: master-key + generate: false # Optional: enterprise billable-request metering. When enabled, the gateway and -# backend count successful requests to inference, MCP, and A2A endpoints and push -# them to LiteLLM's collector over mutual TLS. Both components serve billable -# routes: the backend keeps the named-server MCP transport. Requires an -# enterprise license. The client certificate identifies the deployment, so it is -# mounted read-only from an existing Secret and never passed through the env. +# backend count successful requests to inference, MCP, and A2A endpoints. Requires +# an enterprise license. billingMetrics: enabled: false - endpoint: https://telemetry.litellm.ai # collector to push the counter to - # An existing Secret holding the client certificate under tls.crt and its key - # under tls.key, usually created from the onboarding artifact. The default is - # the conventional name, so the common path is to create that Secret and set - # enabled: true. Override only if yours is named differently. - secretName: litellm-billing-metrics-mtls - # Only for private or test collectors whose server certificate is not on the - # public web PKI. The production collector needs no CA override. - caSecretName: "" # existing Secret holding ca.crt - exportIntervalMs: "" # push cadence; the proxy defaults to 60000 + endpoint: https://telemetry.litellm.ai + tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token + secretName: litellm-billing-metrics + exportIntervalMs: "" -# External Postgres connection. +# Postgres connection. The chart ships no database; point `writer` at your +# own PostgreSQL (managed or self hosted) and its credentials Secret. database: writer: host: "" @@ -258,24 +317,23 @@ database: maxDbConnections: 20 maxClientConn: 1000 -# Optional Redis. Leave host empty to disable. +# Redis is the proxy's coordination store: cross-pod tpm/rpm rate limits, +# spend tracking, and the pod lock manager. The chart emits REDIS_HOST / +# REDIS_PORT / REDIS_PASSWORD, which the proxy picks up through its +# coordination Redis env fallback. Response caching is separate and off unless +# you enable it in `proxy_config.litellm_settings.cache`. # -# This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend -# tracking, and the pod lock manager. The chart emits REDIS_HOST / REDIS_PORT / -# REDIS_PASSWORD, which the proxy picks up through its coordination Redis env -# fallback. Response caching is separate and off unless you enable it in -# `proxy_config.litellm_settings.cache`. +# Set `host` (and `passwordSecret.name` when auth is required). Set +# `cluster: true` for Redis Cluster mode (AWS ElastiCache Cluster, self-hosted +# Redis Cluster): the chart emits REDIS_CLUSTER_NODES from `host` / `port` as +# the single seed and the client discovers the rest from CLUSTER SLOTS. The +# chart ships no Redis of its own. # # For full control, define `general_settings.coordination_redis` in # `proxy_config` (host/port/password/username/url/ssl/startup_nodes/ # sentinel_nodes/sentinel_password/service_name, each accepting os.environ/VAR # refs). An explicit block overrides these env vars. # -# Set `cluster: true` for Redis Cluster mode (e.g. AWS ElastiCache Cluster, -# self-hosted Redis Cluster). The chart emits REDIS_CLUSTER_NODES from -# `host` / `port` as the single seed; the cluster client discovers the -# remaining nodes from CLUSTER SLOTS at startup. -# # For a TLS Redis endpoint set `ssl: true`, which the chart emits as REDIS_SSL. # `sslCertReqs` (none / optional / required) becomes REDIS_SSL_CERT_REQS and # `sslCaCerts` (path to a CA bundle mounted in the pod, e.g. via @@ -296,12 +354,17 @@ redis: hostKey: "" # Secret key containing the Redis host urlKey: "" # Secret key containing the complete Redis URL +# Arbitrary extra Kubernetes manifests rendered verbatim with the release, +# e.g. a NetworkPolicy or an ExternalSecret the proxy needs. +extraResources: [] + # ---------- gateway (LLM data plane) ---------- +# In monolith mode this whole block configures the monolith pod instead. gateway: enabled: true logLevel: INFO - # Number of uvicorn worker processes per gateway pod. Sets NUM_WORKERS, - # consumed by the gateway image entrypoint. Default is 1. + # Number of worker processes per pod. Sets NUM_WORKERS, read by the gateway + # launcher and by the proxy's --num_workers default. Default is 1. numWorkers: 1 extraEnv: [] # Add extra environment variables to the gateway envConfigMaps: [] # Add extra environment variables to the gateway from config maps @@ -320,7 +383,6 @@ gateway: # scrape never runs on an inference worker. Adds a `metrics` port to the pod # and a dedicated ClusterIP `-metrics` Service; point your scrape # config at it. The port has no virtual-key auth: keep it off public ingress. - # Needs the gateway image v1.101.0 or newer. metricsServer: enabled: false port: 4001 @@ -338,8 +400,13 @@ gateway: serviceMonitor: enabled: false labels: {} + annotations: {} interval: 15s scrapeTimeout: 10s + relabelings: [] + # Empty scrapes the release namespace. + namespaceSelector: + matchNames: [] # Opt-in `collector` sidecar (same image, `python -m litellm.proxy.collector`) # that runs the post-response spend pipeline (cost calculation, spend logs, # spend counters, budget reservation reconciliation) so the uvicorn workers @@ -376,13 +443,13 @@ gateway: # ContainerResource metric of the `gateway` container only, so the # sidecar's CPU never drives inference replicas. Needs Kubernetes 1.30+. scaleOnGatewayContainerCpu: true - image: - repository: ghcr.io/berriai/litellm-gateway - tag: "" # defaults to .Chart.AppVersion - pullPolicy: IfNotPresent service: type: ClusterIP port: 4000 + annotations: {} + # For LoadBalancer Services on clusters with several load balancer + # implementations. + loadBalancerClass: "" resources: requests: cpu: "1" @@ -456,6 +523,41 @@ gateway: # counted when a response completes, so TPS trails long streams. targetRequestsPerSecond: "" targetTokensPerSecond: "" + # KEDA ScaledObject as an alternative to the HPA above. Mutually exclusive + # with hpa.enabled (two autoscalers on one Deployment fight each other), so + # set `hpa.enabled: false` when turning this on. Needs KEDA installed in the + # cluster. `triggers` are rendered verbatim; the two optional Prometheus + # workload targets below add prometheus triggers on the same counters the + # HPA workload metrics use, scraped through the metrics Service (enable + # metricsServer and serviceMonitor). + keda: + enabled: false + minReplicaCount: 1 + maxReplicaCount: 10 + pollingInterval: 30 + cooldownPeriod: 300 + # Rendered on the ScaledObject, e.g. autoscaling.keda.sh/paused-replicas: "2" + # to pause scaling at a fixed replica count + scaledObject: + annotations: {} + # Optional; rendered verbatim under spec.advanced + advanced: {} + # Optional; rendered under spec.fallback. A missing `failureThreshold` + # defaults to 3 and a missing `replicas` to maxReplicaCount, as in litellm-helm. + fallback: {} + # Custom triggers, rendered verbatim, e.g. + # - type: cpu + # metricType: Utilization + # metadata: + # value: "70" + triggers: [] + prometheus: + # Prometheus server URL, required when either workload target is set. + serverAddress: "" + # Per pod averages, same units as hpa.targetRequestsPerSecond / + # hpa.targetTokensPerSecond. Empty leaves the trigger out. + targetRequestsPerSecond: "" + targetTokensPerSecond: "" # PodDisruptionBudget for the gateway pods. Set exactly one of # `minAvailable` / `maxUnavailable` (minAvailable wins if both are set; # enabling without either falls back to `maxUnavailable: 1`). Disabled by @@ -465,6 +567,13 @@ gateway: enabled: false minAvailable: "" maxUnavailable: "" + # Extra annotations and labels on the Deployment object itself (not the + # pods), e.g. for a GitOps controller or a reloader. + deploymentAnnotations: {} + deploymentLabels: {} + # Seconds a new pod must be ready before the rollout treats it as + # available. Empty inherits the Kubernetes default of 0. + minReadySeconds: "" podAnnotations: {} # Extra pod labels, merged into the chart's selector labels. Do not # re-declare `app.kubernetes.io/name` / `instance` / `component` here: they @@ -491,7 +600,16 @@ gateway: # egress proxy. Rendered through `tpl`, so entries may reference chart # values and release metadata. extraContainers: [] + # Init containers, rendered through `tpl` the same way. + extraInitContainers: [] # Container lifecycle hooks (postStart / preStop) for the gateway container. + # + # Prefer the proxy's /health/drain preStop hook over a fixed `sleep`: it + # marks the pod NotReady and blocks only until in-flight requests finish + # (bounded by GRACEFUL_SHUTDOWN_TIMEOUT). It is off by default; enable it + # with general_settings.enable_drain_endpoint: true and, when the health + # port is reachable from other pods, set general_settings.drain_endpoint_token + # and send the same value on the X-Drain-Token header from the hook. lifecycle: {} # Grace period the kubelet allows between SIGTERM and SIGKILL. Leave empty # to inherit the Kubernetes default of 30s. Set it a few seconds above the @@ -522,13 +640,11 @@ backend: volumes: [] # Additional volumeMounts on the backend container. volumeMounts: [] - image: - repository: ghcr.io/berriai/litellm-backend - tag: "" - pullPolicy: IfNotPresent service: type: ClusterIP port: 4001 + annotations: {} + loadBalancerClass: "" resources: requests: cpu: "1" @@ -565,12 +681,16 @@ backend: enabled: false minAvailable: "" maxUnavailable: "" + deploymentAnnotations: {} + deploymentLabels: {} + minReadySeconds: "" podAnnotations: {} # Same shape as the gateway blocks of the same name. podLabels: {} podSecurityContext: {} securityContext: {} extraContainers: [] + extraInitContainers: [] lifecycle: {} terminationGracePeriodSeconds: "" nodeSelector: {} @@ -590,13 +710,13 @@ ui: volumes: [] # Additional volumeMounts on the ui container. volumeMounts: [] - image: - repository: ghcr.io/berriai/litellm-ui - tag: "" - pullPolicy: IfNotPresent service: type: ClusterIP port: 3000 + annotations: {} + loadBalancerClass: "" + # The dashboard expects to know where to reach the backend API. Set this to + # the externally-routable URL (typically the ingress host + /api or similar). # The backend's public origin, passed to the UI only when ingress.componentHosts is configured. backendUrl: "" # Public origin of the UI, used for backend and gateway CORS. @@ -634,17 +754,19 @@ ui: enabled: false minAvailable: "" maxUnavailable: "" + deploymentAnnotations: {} + deploymentLabels: {} + minReadySeconds: "" podAnnotations: {} # Same shape as the gateway blocks of the same name. The nginx runtime # writes its pid, cache, and proxy temp files under /tmp, so it boots as # any (arbitrary, non-root) uid; `securityContext.readOnlyRootFilesystem: - # true` here needs an emptyDir volume mounted over /tmp. Images before - # the /tmp move instead need emptyDirs over /var/cache/nginx and /run to - # run as a non-root uid at all. + # true` here needs an emptyDir volume mounted over /tmp. podLabels: {} podSecurityContext: {} securityContext: {} extraContainers: [] + extraInitContainers: [] lifecycle: {} terminationGracePeriodSeconds: "" nodeSelector: {} @@ -653,6 +775,24 @@ ui: # Same shape as gateway.topologySpreadConstraints. topologySpreadConstraints: [] +# Optional: bundled LiteAdmin worker. Renders one Deployment, Service and PVC +# running the shared image with `args: ["--admin-agent"]`, and points the pod +# that serves the /liteadmin routes (the backend in componentized mode, the +# monolith in monolith mode) at it via LITELLM_ADMIN_AGENT_URL. +liteadmin: + enabled: false + existingSecret: "" + gatewayUrl: "" + model: "" + readOnly: false + storageSize: 1Gi + storageClassName: "" + resources: + requests: + cpu: 100m + memory: 256Mi + limits: + memory: 1Gi lens: enabled: false url: "" diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261028000000_subscription_accounts/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261028000000_subscription_accounts/migration.sql new file mode 100644 index 00000000000..e2788c85be9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261028000000_subscription_accounts/migration.sql @@ -0,0 +1,21 @@ +SET LOCAL lock_timeout = '10s'; + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_SubscriptionAccountTable" ( + "subscription_account_id" TEXT NOT NULL, + "custom_llm_provider" TEXT NOT NULL, + "account_id" TEXT NOT NULL, + "label" TEXT, + "monthly_fee" DOUBLE PRECISION NOT NULL, + "currency" TEXT NOT NULL, + "billing_period_start" TEXT NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_SubscriptionAccountTable_pkey" PRIMARY KEY ("subscription_account_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_SubscriptionAccountTable_provider_account_key" ON "LiteLLM_SubscriptionAccountTable"("custom_llm_provider", "account_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql new file mode 100644 index 00000000000..4965ed3cfc4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql @@ -0,0 +1,9 @@ +SET LOCAL lock_timeout = '10s'; + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BackgroundResponseSettlement" ( + "response_id" TEXT NOT NULL, + "settled_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_BackgroundResponseSettlement_pkey" PRIMARY KEY ("response_id") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql new file mode 100644 index 00000000000..b00cf892b68 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql @@ -0,0 +1,19 @@ +SET LOCAL lock_timeout = '10s'; +CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogDelivery" ( + "delivery_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text, + "request_id" TEXT NOT NULL, + "state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending', + "attempts" INTEGER NOT NULL DEFAULT 0, + "next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "applied_at" TIMESTAMP(3), + "last_error" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "accounting_payload" JSONB NOT NULL, + "is_batch" BOOLEAN NOT NULL +); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_next_attempt_at_idx" + ON "LiteLLM_SpendLogDelivery" ("state", "next_attempt_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_applied_at_idx" + ON "LiteLLM_SpendLogDelivery" ("state", "applied_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_request_id_idx" + ON "LiteLLM_SpendLogDelivery" ("request_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104000000_skill_zip_upload/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104000000_skill_zip_upload/migration.sql new file mode 100644 index 00000000000..acbe87515d9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104000000_skill_zip_upload/migration.sql @@ -0,0 +1,9 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginUpload" ( + "plugin_name" TEXT NOT NULL, + "content" BYTEA NOT NULL, + CONSTRAINT "LiteLLM_ClaudeCodePluginUpload_pkey" PRIMARY KEY ("plugin_name"), + CONSTRAINT "LiteLLM_ClaudeCodePluginUpload_plugin_name_fkey" FOREIGN KEY ("plugin_name") + REFERENCES "LiteLLM_ClaudeCodePluginTable"("name") ON DELETE CASCADE ON UPDATE CASCADE +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104010000_drop_accounting_event_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104010000_drop_accounting_event_tables/migration.sql new file mode 100644 index 00000000000..cba6a18368f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104010000_drop_accounting_event_tables/migration.sql @@ -0,0 +1,5 @@ +SET LOCAL lock_timeout = '10s'; +DROP TABLE IF EXISTS "LiteLLM_AccountingDelivery"; +DROP TABLE IF EXISTS "LiteLLM_AccountingEvent"; +DROP TYPE IF EXISTS "AccountingDeliveryTarget"; +DROP TYPE IF EXISTS "AccountingEventKind"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql new file mode 100644 index 00000000000..ecfa7321d87 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql @@ -0,0 +1,16 @@ +SET LOCAL lock_timeout = '10s'; +CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion" ( + "id" TEXT NOT NULL, + "plugin_id" TEXT NOT NULL, + "version" TEXT, + "description" TEXT, + "manifest_json" TEXT, + "published_at" TIMESTAMP(3), + "published_by" TEXT, + "archive_available" BOOLEAN NOT NULL DEFAULT false, + "content" BYTEA, + CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_pkey" PRIMARY KEY ("id"), + CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_plugin_id_fkey" FOREIGN KEY ("plugin_id") REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE +); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_version_key" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "version"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_published_at_idx" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "published_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql new file mode 100644 index 00000000000..e7975258ca9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql @@ -0,0 +1,27 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE TABLE IF NOT EXISTS "LiteLLM_SkillSubmission" ( + "id" TEXT NOT NULL, + "plugin_name" TEXT NOT NULL, + "operation" TEXT NOT NULL, + "status" TEXT NOT NULL DEFAULT 'pending_review', + "proposal_json" TEXT NOT NULL, + "archive_content" BYTEA, + "submitted_by" TEXT NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "base_plugin_id" TEXT, + "base_updated_at" TIMESTAMP(3), + "reviewed_by" TEXT, + "reviewed_at" TIMESTAMP(3), + "rejection_reason" TEXT, + CONSTRAINT "LiteLLM_SkillSubmission_pkey" PRIMARY KEY ("id") +); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_created_at_id_idx" + ON "LiteLLM_SkillSubmission"("submitted_by", "created_at" DESC, "id" DESC); +CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_status_created_at_id_idx" + ON "LiteLLM_SkillSubmission"("submitted_by", "status", "created_at" DESC, "id" DESC); +CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_status_created_at_id_idx" + ON "LiteLLM_SkillSubmission"("status", "created_at" DESC, "id" DESC); +CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_created_at_id_idx" + ON "LiteLLM_SkillSubmission"("created_at" DESC, "id" DESC); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql new file mode 100644 index 00000000000..6b70cca4a9a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql @@ -0,0 +1,3 @@ +SET LOCAL lock_timeout = '10s'; + +ALTER TABLE "LiteLLM_SkillSubmission" ADD COLUMN IF NOT EXISTS "scan_json" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql new file mode 100644 index 00000000000..43e26b324a0 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql @@ -0,0 +1,37 @@ +SET LOCAL lock_timeout = '10s'; + +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_attribute + WHERE attrelid = to_regclass('"LiteLLM_ClaudeCodePluginTable"') + AND attname = 'anonymous_download_count' AND NOT attisdropped + ) THEN + ALTER TABLE "LiteLLM_ClaudeCodePluginTable" + ADD COLUMN IF NOT EXISTS "anonymous_download_count" BIGINT NOT NULL DEFAULT 0; + END IF; + + IF to_regclass('"LiteLLM_ClaudeCodePluginDownload"') IS NULL THEN + CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload" ( + "id" BIGSERIAL NOT NULL, + "plugin_id" TEXT NOT NULL, + "downloaded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "version" TEXT, + "digest" TEXT NOT NULL, + "authenticated" BOOLEAN NOT NULL, + "user_id" TEXT, + "user_email" TEXT, + "key_alias" TEXT, + "team_id" TEXT, + "team_alias" TEXT, + CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_pkey" PRIMARY KEY ("id"), + CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_plugin_id_fkey" FOREIGN KEY ("plugin_id") + REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE + ); + END IF; + + IF to_regclass('"LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx"') IS NULL THEN + CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx" + ON "LiteLLM_ClaudeCodePluginDownload"("plugin_id", "id"); + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql new file mode 100644 index 00000000000..6b77194e3b4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql @@ -0,0 +1,7 @@ +SET LOCAL lock_timeout = '10s'; + +-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_tpm_limit" JSONB; + +-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_rpm_limit" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110120000_add_daily_spend_guardrail_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110120000_add_daily_spend_guardrail_spend/migration.sql new file mode 100644 index 00000000000..c2c264325a6 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110120000_add_daily_spend_guardrail_spend/migration.sql @@ -0,0 +1,22 @@ +SET LOCAL lock_timeout = '10s'; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyGlobalSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2965174114b..5a5d56cc6e2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -19,6 +19,8 @@ model LiteLLM_BudgetTable { rpm_limit BigInt? tpd_limit BigInt? model_max_budget Json? + model_tpm_limit Json? + model_rpm_limit Json? budget_duration String? budget_reset_at DateTime? allowed_models String[] @default([]) // per-member model scope; empty = inherit team models @@ -881,6 +883,7 @@ model LiteLLM_DailyUserSpend { prompt_caching_savings_spend Float @default(0.0) gateway_injected_caching_savings_spend Float @default(0.0) autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) spend Float @default(0.0) api_requests BigInt @default(0) successful_requests BigInt @default(0) @@ -917,6 +920,7 @@ model LiteLLM_DailyGlobalSpend { prompt_caching_savings_spend Float @default(0.0) gateway_injected_caching_savings_spend Float @default(0.0) autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) spend Float @default(0.0) api_requests BigInt @default(0) successful_requests BigInt @default(0) @@ -950,6 +954,7 @@ model LiteLLM_DailyOrganizationSpend { prompt_caching_savings_spend Float @default(0.0) gateway_injected_caching_savings_spend Float @default(0.0) autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) spend Float @default(0.0) api_requests BigInt @default(0) successful_requests BigInt @default(0) @@ -988,6 +993,7 @@ model LiteLLM_DailyEndUserSpend { prompt_caching_savings_spend Float @default(0.0) gateway_injected_caching_savings_spend Float @default(0.0) autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) spend Float @default(0.0) api_requests BigInt @default(0) successful_requests BigInt @default(0) @@ -1025,6 +1031,7 @@ model LiteLLM_DailyAgentSpend { prompt_caching_savings_spend Float @default(0.0) gateway_injected_caching_savings_spend Float @default(0.0) autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) spend Float @default(0.0) api_requests BigInt @default(0) successful_requests BigInt @default(0) @@ -1062,6 +1069,7 @@ model LiteLLM_DailyTeamSpend { prompt_caching_savings_spend Float @default(0.0) gateway_injected_caching_savings_spend Float @default(0.0) autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) spend Float @default(0.0) api_requests BigInt @default(0) successful_requests BigInt @default(0) @@ -1102,6 +1110,7 @@ model LiteLLM_DailyTagSpend { prompt_caching_savings_spend Float @default(0.0) gateway_injected_caching_savings_spend Float @default(0.0) autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) spend Float @default(0.0) api_requests BigInt @default(0) successful_requests BigInt @default(0) @@ -1615,9 +1624,77 @@ model LiteLLM_ClaudeCodePluginTable { updated_at DateTime? @default(now()) @updatedAt created_by String? + uploaded_archive LiteLLM_ClaudeCodePluginUpload? + versions LiteLLM_ClaudeCodePluginVersion[] + downloads LiteLLM_ClaudeCodePluginDownload[] + anonymous_download_count BigInt @default(0) + @@map("LiteLLM_ClaudeCodePluginTable") } +model LiteLLM_ClaudeCodePluginVersion { + id String @id @default(uuid()) + plugin_id String + version String? + description String? + manifest_json String? + published_at DateTime? + published_by String? + archive_available Boolean @default(false) + content Bytes? + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade) + + @@unique([plugin_id, version]) + @@index([plugin_id, published_at]) +} + + +model LiteLLM_SkillSubmission { + id String @id @default(uuid()) + plugin_name String + operation String + status String @default("pending_review") + proposal_json String + archive_content Bytes? + submitted_by String + created_at DateTime @default(now()) + base_plugin_id String? + base_updated_at DateTime? + reviewed_by String? + reviewed_at DateTime? + rejection_reason String? + scan_json String? + + @@index([submitted_by, created_at(sort: Desc), id(sort: Desc)]) + @@index([submitted_by, status, created_at(sort: Desc), id(sort: Desc)]) + @@index([status, created_at(sort: Desc), id(sort: Desc)]) + @@index([created_at(sort: Desc), id(sort: Desc)]) +} + +model LiteLLM_ClaudeCodePluginDownload { + id BigInt @id @default(autoincrement()) + plugin_id String + downloaded_at DateTime @default(now()) + version String? + digest String + authenticated Boolean + user_id String? + user_email String? + key_alias String? + team_id String? + team_alias String? + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade) + + @@index([plugin_id, id]) +} + +model LiteLLM_ClaudeCodePluginUpload { + plugin_name String @id + content Bytes + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_name], references: [name], onDelete: Cascade, onUpdate: Cascade) +} + + // User/team-scoped memory store with a GLOBAL unique key. // `value` is a string (typically markdown/text meant for LLM context); // `metadata` is an optional JSON envelope for structured tags without schema changes. @@ -1985,6 +2062,11 @@ model LiteLLM_BackgroundInteractionSettlement { @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } +model LiteLLM_BackgroundResponseSettlement { + response_id String @id + settled_at DateTime @default(now()) +} + model LiteLLM_Lens { reviews LiteLLM_LensReview[] id String @id @@ -2062,72 +2144,41 @@ model LiteLLM_TelemetryReport { @@index([created_at]) } -enum AccountingEventKind { - inbound_usage - outbound_usage - cost -} - enum AccountingDeliveryState { pending applied parked } -enum AccountingDeliveryTarget { - LiteLLM_VerificationToken - LiteLLM_UserTable - LiteLLM_TeamTable - LiteLLM_TeamMembership - LiteLLM_OrganizationTable - LiteLLM_OrganizationMembership - LiteLLM_EndUserTable - LiteLLM_AgentsTable - LiteLLM_TagTable - LiteLLM_ProjectTable - LiteLLM_ModelAccessGroupBudgetTable - LiteLLM_BudgetTable - LiteLLM_BudgetWindowSpend - LiteLLM_DailyUserSpend - LiteLLM_DailyTeamSpend - LiteLLM_DailyOrganizationSpend - LiteLLM_DailyEndUserSpend - LiteLLM_DailyAgentSpend - LiteLLM_DailyTagSpend - LiteLLM_DailyModelUsage - LiteLLM_DailyToolSpend - LiteLLM_AutoRouterDailySpend +model LiteLLM_SpendLogDelivery { + delivery_id String @id @default(dbgenerated("(gen_random_uuid())::text")) + request_id String + state AccountingDeliveryState @default(pending) + accounting_payload Json + is_batch Boolean + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + created_at DateTime @default(now()) + + @@index([state, next_attempt_at]) + @@index([state, applied_at]) + @@index([request_id]) } -// Durable accounting ledger. One row per usage or cost event; payload holds the -// private event envelope versioned by schema_version. Usage rows retain request-scoped -// pricing context; cost rows include resolved rates and the finalized pricing snapshot. -model LiteLLM_AccountingEvent { - event_id String @id - kind AccountingEventKind - source_event_id String? - schema_version Int - payload_hash String - occurred_at DateTime - recorded_at DateTime @default(now()) - payload Json +model LiteLLM_SubscriptionAccountTable { + subscription_account_id String @id @default(uuid()) + custom_llm_provider String + account_id String + label String? + monthly_fee Float + currency String + billing_period_start String + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String - @@index([recorded_at]) - @@index([kind, source_event_id]) -} - -// One row per (event, target table, version): pending work for the projector and, -// once applied in the same transaction as the target write, the receipt that stops a retry. -model LiteLLM_AccountingDelivery { - event_id String - target AccountingDeliveryTarget - version Int - state AccountingDeliveryState @default(pending) - attempts Int @default(0) - next_attempt_at DateTime @default(now()) - applied_at DateTime? - last_error String? - - @@id([event_id, target, version]) - @@index([target, state, next_attempt_at]) + @@unique([custom_llm_provider, account_id], map: "LiteLLM_SubscriptionAccountTable_provider_account_key") } diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index a12e1e27cfc..1c0c708ae7a 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.112" +version = "0.4.114" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -30,7 +30,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.112" +version = "0.4.114" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/.agents/skills/rust-gateway/SKILL.md b/litellm-rust/.agents/skills/rust-gateway/SKILL.md index 4fee0b54a68..923c9bf5ab6 100644 --- a/litellm-rust/.agents/skills/rust-gateway/SKILL.md +++ b/litellm-rust/.agents/skills/rust-gateway/SKILL.md @@ -7,7 +7,7 @@ metadata: # Rust gateway -`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching +`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. `gateway-forward` is the reverse proxy that relays every request to the Python proxy and has no routes of its own. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching - [surfaces](references/surfaces.md) - Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose diff --git a/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md b/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md new file mode 100644 index 00000000000..5ead8f287ab --- /dev/null +++ b/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md @@ -0,0 +1,53 @@ +--- +name: rust-python-boundary +description: Decide where Python-shaped input becomes typed Rust in litellm-rust. Use when reading config.yaml values, os.environ/ references, REDIS_* or other environment layering, get_secret fallbacks, or kwargs and objects handed over from Python, when coercing text to a bool, int or float the way Python does, or when adding a dependency on litellm-python-compat +metadata: + internal: true +--- + +# Python boundary + +Python-shaped input becomes typed Rust in exactly two crates. Every crate past them handles typed values, so no crate re-guesses how Python would have read a value + +## The two boundaries + +- `crates/config` reads config.yaml, which is Python-shaped: `os.environ/` references, scalars that Helm or the environment turned into text, raw kwargs maps. It also owns the environment layering and `get_secret` fallbacks Python applies while it loads that config. It has no Python runtime, so it reproduces Python value semantics with `litellm-python-compat` +- `crates/python-bridge` receives live Python objects, such as `RedisCache.redis_kwargs`. It reads them through PyO3 with the real protocol (`is_truthy()`, `str()`, `extract`), so it never imitates Python +- Both produce the same typed input, defined in the data crate the consumer already depends on (`redis-types` for Redis), so a downstream crate cannot tell which host it runs under + +## Python value semantics + +Reading a value the way Python itself reads it. This lives only at the two boundaries + +- Truthiness and `x or fallback`, `str(x)`, `int(text)` and `float(text)` +- Each surface's true and false text: `_str_to_bool`, pydantic's lax bool, `get_secret`'s `true` and `false` +- `os.environ/` references and `REDIS_` style environment layering + +A boundary knows which surface a value came from. `general_settings.coordination_redis` is a pydantic model and `cache_params` goes through `_str_to_bool`, so the boundary applies that surface's rule exactly instead of refusing the text two surfaces disagree on + +Each typed field keeps where its value came from (the key, `REDIS_` or the fallback name), so an error or warning downstream still names it + +## What stays downstream + +A crate that stands in for a Python library keeps that library's semantics, because they are its compatibility contract and not value coercion. `redis-resolve` stands in for redis-py, so it keeps + +- Which target wins, and the defaults redis-py applies +- How `parse_url` reads a url, including `int()` on `?db=` and redis-py's `to_bool` on the TLS query parameters +- Which kwargs redis-py accepts, and whether the Rust client reads, warns about or refuses each one + +Its input is typed. A rule it applies reads typed fields, never a `serde_json::Value` it has to coerce first. A value it only carries back to Python untouched, like `passthrough`, may stay untyped + +## litellm-python-compat + +The rule is about where config and kwargs get coerced, not about who may depend on the crate + +- Its data formats (`pickle`, `json.dumps`, `repr`) are fine anywhere Rust reads or writes state Python also reads or writes, such as cache values in a shared Redis +- Its coercions (`truthy`, `number`) belong at the two boundaries, and in a crate that stands in for a Python library when that library calls them itself + +## Existing code + +Some code still coerces Python-shaped values past the boundary, for example the `REDIS_*` merge in `gateway/src/redis/layers.rs` and the flag and number readers in `redis-resolve/src/options.rs`. Do not add more. When you change one, move the coercion to its boundary instead of extending it in place + +## Tests + +Pin a boundary's coercion with named cases next to the Python function it mirrors, in the Rust test and in its Python twin, like the Redis contract cases in `redis-resolve/tests/interpret.rs` and `tests/unit/test_redis.py`. A downstream crate's tests build its typed input directly and never feed it raw text to coerce diff --git a/litellm-rust/.agents/skills/rust-redaction/SKILL.md b/litellm-rust/.agents/skills/rust-redaction/SKILL.md new file mode 100644 index 00000000000..2e906e10607 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-redaction/SKILL.md @@ -0,0 +1,43 @@ +--- +name: rust-redaction +description: Keep secrets out of Debug output in litellm-rust with SecretValue and veil. Use when a struct or enum holds a password, token, key, credential URL or credential provider, when writing or reviewing a Debug impl, or when a type with secrets gets logged or traced +metadata: + internal: true +--- + +# Rust redaction + +Every secret in a Debug output is hidden the same way, so a reviewer can tell at a glance that a field is covered and a new field cannot slip through a hand-written impl + +## Where SecretValue lives + +`SecretValue` lives in `crates/secret-value` and is imported as `litellm_secret_value::SecretValue` everywhere. No crate re-exports it, so there is one import path. Depend on `litellm-secret-value` directly. Do not pull in `litellm-auth-types` or `litellm-auth` just to get it + +`crates/secret-value` holds only the type. Credential acquisition (`TokenProvider`, `CredentialPlan`, provider param tables) stays in `litellm-auth-types`, and generic clients such as `litellm-redis-client` must not depend on it + +## Rules + +- Never hand-write `impl Debug` to hide a value, and never invent markers like `""`, `""` or `.map(|_| "...")`. A hand-written impl silently prints any field added later +- A secret that crosses a crate boundary or lives in config is typed as `SecretValue`. It redacts itself, compares in constant time and only gives up the plaintext through `expose()`. A type whose only secrets are `SecretValue` fields derives plain `Debug` + ```rust + #[derive(Clone, Debug, PartialEq)] + pub struct HostEndpoint { + pub host: String, + pub port: u16, + pub username: Option, + pub password: Option, + pub tls: bool, + pub db: i64, + } + ``` +- Keep a secret as `SecretValue` until the last conversion into the third-party type (redis-rs, reqwest, an SDK) and call `expose()` there. Calling `expose()` early spreads plaintext `String`s through our own types +- A secret that has to stay a raw type, because a serde wire shape or a third-party struct owns the field, is covered by deriving `veil::Redact` on the holder and marking the field `#[redact(with = "[REDACTED]")]`. `Redact` replaces `Debug`, so do not derive both +- A URL that can carry userinfo is a secret as a whole. Do not parse it to show the host. Keep it as a `SecretValue` +- A bag of unparsed settings that may hold credentials (`additional_fields`, raw kwargs) is redacted as a whole with `#[redact(with = "[REDACTED]")]` +- An opaque handle over `Arc` whose trait has no `Debug` gets a newtype whose `Debug` prints only the type name with `finish_non_exhaustive()`. That impl shows nothing, so it is not a redaction impl and the first rule does not apply. A type holding such a handle can derive `Debug` +- Never put a secret in a tracing field or an error message. Log the type with `?value`, or log only non-secret fields. See the rust-tracing and rust-error skills + +## Tests + +- One test per type that holds a secret, built the way production builds it, asserting the plaintext is absent from `format!("{value:?}")` and `[REDACTED]` is present +- Use a distinctive plaintext (`"credential-value"`) so the absence check cannot pass by accident diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 9e7a55eb232..5fcf05f61f7 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,11 +1,17 @@ # Rust workspace rules +For keeping secrets out of Debug output, follow [.agents/skills/rust-redaction/SKILL.md](.agents/skills/rust-redaction/SKILL.md) + For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md) For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md) +For where config.yaml values, environment layering and kwargs from Python become typed Rust, follow [.agents/skills/rust-python-boundary/SKILL.md](.agents/skills/rust-python-boundary/SKILL.md) + +The process-wide rustls `CryptoProvider` (ring) is installed once by the host. rust-proxy does it at the top of `main` through `litellm_gateway::install_crypto_provider`, and python-bridge must do the same in its module init before it builds any rustls config. Library crates never install it. The AWS SDK links aws-lc-rs next to ring, so without an installed default the first rustls config a library builds, such as redis-rs TLS, panics + For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e959f77917b..ed658e074ee 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -8,6 +8,17 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures 0.2.17", +] + [[package]] name = "ahash" version = "0.8.12" @@ -547,7 +558,7 @@ dependencies = [ "http 1.4.2", "http-body 0.4.6", "hyper 0.14.32", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.24.2", "hyper-rustls 0.27.9", "hyper-util", @@ -736,7 +747,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "itoa", "matchit", @@ -883,6 +894,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" dependencies = [ "fastrand", + "tokio", ] [[package]] @@ -982,6 +994,15 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "block-padding" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93" +dependencies = [ + "generic-array", +] + [[package]] name = "bollard" version = "0.20.2" @@ -1000,7 +1021,7 @@ dependencies = [ "home", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-named-pipe", "hyper-rustls 0.27.9", "hyper-util", @@ -1134,6 +1155,15 @@ dependencies = [ "rustversion", ] +[[package]] +name = "cbc" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6" +dependencies = [ + "cipher", +] + [[package]] name = "cc" version = "1.3.0" @@ -1210,6 +1240,16 @@ dependencies = [ "half", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common 0.1.7", + "inout", +] + [[package]] name = "clap" version = "4.6.7" @@ -1774,6 +1814,15 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "des" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e" +dependencies = [ + "cipher", +] + [[package]] name = "digest" version = "0.10.7" @@ -2255,7 +2304,7 @@ dependencies = [ "chrono", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.27.9", "hyper-util", "ring", @@ -2391,7 +2440,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "lazy_static", "opentelemetry", "opentelemetry-semantic-conventions", @@ -2833,9 +2882,9 @@ dependencies = [ [[package]] name = "hyper" -version = "1.10.1" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +checksum = "2c3e324da4c95177d6291d4c8730197c0d1822f8a9766814a4a44fa5ab797c9c" dependencies = [ "atomic-waker", "bytes", @@ -2860,7 +2909,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fab3637d6b04a8037af8a266fdf6cf92ea957e8c53981a2bf6136572531025bf" dependencies = [ "hex", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -2889,7 +2938,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ "http 1.4.2", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "rustls 0.23.42", "rustls-native-certs", @@ -2904,7 +2953,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" dependencies = [ - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -2923,7 +2972,7 @@ dependencies = [ "futures-util", "http 1.4.2", "http-body 1.1.0", - "hyper 1.10.1", + "hyper 1.12.0", "ipnet", "libc", "percent-encoding", @@ -2942,7 +2991,7 @@ checksum = "986c5ce3b994526b3cd75578e62554abd09f0899d6206de48b3e96ab34ccc8c7" dependencies = [ "hex", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -3105,6 +3154,16 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "block-padding", + "generic-array", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -3380,6 +3439,7 @@ dependencies = [ "litellm-auth-azure", "litellm-auth-gcp", "litellm-auth-types", + "litellm-secret-value", ] [[package]] @@ -3388,15 +3448,19 @@ version = "0.1.0" dependencies = [ "aws-config", "aws-credential-types", + "aws-runtime", "aws-sdk-sts", "aws-sigv4", "aws-smithy-runtime-api", + "aws-smithy-types", "aws-types", "bytes", "futures-util", "litellm-auth-types", "litellm-error", "litellm-http", + "litellm-secret-value", + "litellm-secrets", "moka", "reqwest", "rstest", @@ -3405,6 +3469,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "url", "veil", ] @@ -3416,6 +3481,8 @@ dependencies = [ "azure_identity", "litellm-auth-types", "litellm-error", + "litellm-secret-value", + "litellm-secrets", "moka", "rstest", "serde_json", @@ -3429,18 +3496,22 @@ dependencies = [ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ + "axum", "base64 0.22.1", "gcp_auth", "google-cloud-auth", "http 1.4.2", "litellm-auth-types", "litellm-error", + "litellm-secret-value", "moka", "rstest", + "rustls 0.23.42", "serde", "serde_json", "sha2 0.10.9", "tempfile", + "time", "tokio", "veil", ] @@ -3450,10 +3521,10 @@ name = "litellm-auth-types" version = "0.1.0" dependencies = [ "litellm-error", + "litellm-secret-value", "rstest", "serde", "serde_json", - "subtle", "thiserror 2.0.19", "tokio", "veil", @@ -3523,6 +3594,7 @@ dependencies = [ "litellm-cache-testing", "litellm-error", "litellm-http", + "litellm-secret-value", "percent-encoding", "reqwest", "rstest", @@ -3565,37 +3637,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "litellm-cache-redis" -version = "0.1.0" -dependencies = [ - "litellm-cache", - "litellm-cache-testing", - "litellm-error", - "r2d2", - "redis", - "redis-test", - "rstest", - "serde_json", - "thiserror 2.0.19", - "tokio", -] - -[[package]] -name = "litellm-cache-redis-semantic" -version = "0.1.0" -dependencies = [ - "litellm-cache", - "litellm-cache-redis", - "litellm-cache-testing", - "redis", - "redis-test", - "rstest", - "serde_json", - "sha2 0.10.9", - "tokio", -] - [[package]] name = "litellm-cache-response" version = "0.1.0" @@ -3603,11 +3644,8 @@ dependencies = [ "litellm-cache", "litellm-cache-gcs", "litellm-cache-memory", - "litellm-cache-redis", "litellm-http", "py_literal", - "redis", - "redis-test", "rstest", "serde", "serde_json", @@ -3620,19 +3658,15 @@ dependencies = [ name = "litellm-cache-s3" version = "0.1.0" dependencies = [ - "aws-credential-types", "aws-sdk-s3", - "aws-smithy-runtime-api", "aws-smithy-types", - "aws-types", "futures-util", - "http 1.4.2", "litellm-auth-aws", "litellm-cache", "litellm-cache-testing", "litellm-error", "litellm-http", - "reqwest", + "litellm-s3", "rstest", "serde_json", "tokio", @@ -3646,32 +3680,14 @@ dependencies = [ "litellm-cache", ] -[[package]] -name = "litellm-cache-valkey-semantic" -version = "0.1.0" -dependencies = [ - "litellm-cache", - "litellm-cache-redis", - "litellm-cache-testing", - "litellm-error", - "redis", - "redis-test", - "rstest", - "serde_json", - "sha2 0.10.9", - "thiserror 2.0.19", - "tokio", - "uuid", -] - [[package]] name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ - "litellm-auth", "litellm-error", "litellm-host", "litellm-host-python", + "litellm-secret-value", "proptest", "pyo3", "rstest", @@ -3692,9 +3708,9 @@ dependencies = [ name = "litellm-config" version = "0.1.0" dependencies = [ - "litellm-auth-types", "litellm-error", "litellm-router-types", + "litellm-secret-value", "rstest", "serde", "serde_json", @@ -3702,8 +3718,7 @@ dependencies = [ "strum", "tempfile", "thiserror 2.0.19", - "tracing", - "tracing-subscriber", + "veil", ] [[package]] @@ -3804,11 +3819,12 @@ dependencies = [ "envy", "futures-util", "http-body-util", - "litellm-auth-types", "litellm-circuit-breaker", "litellm-config", + "litellm-core-utils", "litellm-error", "litellm-gateway-auth", + "litellm-gateway-forward", "litellm-gateway-inference", "litellm-gateway-mcp", "litellm-gateway-ui", @@ -3816,10 +3832,15 @@ dependencies = [ "litellm-inference", "litellm-llms", "litellm-redis-client", + "litellm-redis-connect", + "litellm-redis-resolve", "litellm-redis-testing", + "litellm-redis-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", + "rustls 0.23.42", "serde", "serde_json", "tempfile", @@ -3830,6 +3851,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "veil", ] [[package]] @@ -3839,9 +3861,9 @@ dependencies = [ "axum", "axum-login", "futures-util", - "litellm-auth-types", "litellm-config", "litellm-error", + "litellm-secret-value", "litellm-secrets", "rstest", "serde", @@ -3854,6 +3876,29 @@ dependencies = [ "veil", ] +[[package]] +name = "litellm-gateway-forward" +version = "0.1.0" +dependencies = [ + "axum", + "bytes", + "envy", + "futures-util", + "http-body-util", + "httparse", + "hyper 1.12.0", + "hyper-util", + "proptest", + "rstest", + "serde", + "thiserror 2.0.19", + "tokio", + "tokio-tungstenite", + "tokio-util", + "tower-service", + "tracing", +] + [[package]] name = "litellm-gateway-inference" version = "0.1.0" @@ -3879,6 +3924,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-router", + "litellm-secret-value", "litellm-secrets", "rstest", "serde", @@ -3895,9 +3941,9 @@ name = "litellm-gateway-management" version = "0.1.0" dependencies = [ "base64 0.22.1", - "litellm-auth-types", "litellm-error", "litellm-gateway-auth", + "litellm-secret-value", "rand 0.8.7", "rstest", "thiserror 2.0.19", @@ -3912,9 +3958,9 @@ dependencies = [ "base64 0.22.1", "futures-util", "http 1.4.2", - "litellm-auth-types", "litellm-config", "litellm-error", + "litellm-secret-value", "litellm-secrets", "moka", "rmcp", @@ -3942,9 +3988,9 @@ dependencies = [ "futures-util", "governor", "jsonwebtoken", - "litellm-auth-types", "litellm-error", "litellm-gateway-auth", + "litellm-secret-value", "rand 0.8.7", "rstest", "serde", @@ -3964,8 +4010,8 @@ name = "litellm-host" version = "0.1.0" dependencies = [ "futures-util", - "litellm-auth", "litellm-coroutine", + "litellm-secret-value", "rstest", "serde_json", "thiserror 2.0.19", @@ -4094,6 +4140,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", @@ -4124,6 +4171,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "reqwest", @@ -4154,6 +4202,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "mime_guess", @@ -4188,6 +4237,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "reqwest", @@ -4207,6 +4257,7 @@ dependencies = [ "litellm-auth", "litellm-http", "litellm-inference", + "litellm-secret-value", "litellm-secrets", "tokio", ] @@ -4254,6 +4305,7 @@ dependencies = [ "litellm-llms-types", "litellm-python-compat", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "reqwest", "rstest", @@ -4291,12 +4343,14 @@ dependencies = [ "jsonschema", "litellm-cost", "litellm-error", + "litellm-tracing", "macro_rules_attribute", "rstest", "schemars 1.2.2", "serde", "serde_json", "serde_with", + "sha1 0.10.7", "strum", "thiserror 2.0.19", "time", @@ -4306,9 +4360,13 @@ dependencies = [ name = "litellm-model-catalog-refresh" version = "0.1.0" dependencies = [ + "aws-credential-types", + "backon", + "bytes", "litellm-error", "litellm-http", - "litellm-model-catalog", + "litellm-s3", + "litellm-tracing", "reqwest", "rstest", "thiserror 2.0.19", @@ -4320,18 +4378,19 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "async-trait", "aws-sdk-secretsmanager", "bytes", "criterion", "fancy-regex 0.19.2", "futures-util", + "indexmap 2.14.0", "litellm-auth", "litellm-auth-aws", "litellm-cache", - "litellm-cache-memory", - "litellm-cache-redis", "litellm-cache-response", "litellm-callbacks-legacy-python", + "litellm-circuit-breaker", "litellm-core-utils", "litellm-error", "litellm-host", @@ -4346,7 +4405,16 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-model-catalog-refresh", + "litellm-python-compat", + "litellm-redis-client", + "litellm-redis-connect", + "litellm-redis-resolve", + "litellm-redis-store", + "litellm-redis-types", + "litellm-router-core", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "litellm-secrets-aws", "litellm-secrets-types", @@ -4357,9 +4425,9 @@ dependencies = [ "prost", "pyo3", "pyo3-async-runtimes", - "redis", "reqwest", "rstest", + "rustls 0.23.42", "serde", "serde_json", "serde_with", @@ -4395,22 +4463,110 @@ version = "0.1.0" dependencies = [ "futures-util", "litellm-circuit-breaker", + "litellm-python-compat", "litellm-redis-testing", + "litellm-redis-types", + "litellm-secret-value", + "pem", + "pkcs8", "redis", "rstest", + "rustls 0.23.42", + "rustls-native-certs", + "socket2 0.6.5", "thiserror 2.0.19", "tokio", + "tracing", +] + +[[package]] +name = "litellm-redis-connect" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-auth-aws", + "litellm-auth-azure", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-circuit-breaker", + "litellm-core-utils", + "litellm-error", + "litellm-http", + "litellm-redis-client", + "litellm-redis-resolve", + "litellm-redis-testing", + "litellm-redis-types", + "litellm-secret-value", + "rstest", + "rustls 0.23.42", + "serde", + "serde_json", + "strum", + "thiserror 2.0.19", + "time", + "tokio", + "tracing", + "url", + "veil", +] + +[[package]] +name = "litellm-redis-resolve" +version = "0.1.0" +dependencies = [ + "litellm-python-compat", + "litellm-redis-types", + "litellm-secret-value", + "rstest", + "serde", + "serde_json", + "strum", + "thiserror 2.0.19", + "tracing", + "url", +] + +[[package]] +name = "litellm-redis-store" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-cache", + "litellm-cache-memory", + "litellm-circuit-breaker", + "litellm-python-compat", + "litellm-redis-client", + "litellm-redis-testing", + "litellm-secret-value", + "litellm-tracing", + "rstest", + "strum", + "thiserror 2.0.19", + "tokio", + "tracing", ] [[package]] name = "litellm-redis-testing" version = "0.1.0" dependencies = [ + "pkcs8", + "rcgen", + "tempfile", "testcontainers-modules", "thiserror 2.0.19", "tokio", ] +[[package]] +name = "litellm-redis-types" +version = "0.1.0" +dependencies = [ + "litellm-secret-value", + "rstest", + "url", +] + [[package]] name = "litellm-router" version = "0.1.0" @@ -4426,18 +4582,64 @@ dependencies = [ [[package]] name = "litellm-router-core" version = "0.1.0" +dependencies = [ + "arc-swap", + "async-trait", + "jiff", + "litellm-cache", + "litellm-router-types", + "parking_lot", + "rand 0.8.7", + "rstest", + "serde", + "serde_json", + "serde_with", + "sha2 0.10.9", + "strum", + "thiserror 2.0.19", + "tokio", + "tracing", +] [[package]] name = "litellm-router-types" version = "0.1.0" dependencies = [ "litellm-auth-types", + "litellm-secret-value", "rstest", "serde", "serde_json", "serde_with", ] +[[package]] +name = "litellm-s3" +version = "0.1.0" +dependencies = [ + "aws-config", + "aws-credential-types", + "aws-sdk-s3", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "http 1.4.2", + "litellm-http", + "reqwest", + "rstest", + "thiserror 2.0.19", +] + +[[package]] +name = "litellm-secret-value" +version = "0.1.0" +dependencies = [ + "rstest", + "serde", + "subtle", + "veil", +] + [[package]] name = "litellm-secrets" version = "0.1.0" @@ -4451,6 +4653,7 @@ dependencies = [ "litellm-error", "litellm-http", "litellm-python-compat", + "litellm-secret-value", "litellm-secrets-aws", "litellm-secrets-azure", "litellm-secrets-cyberark", @@ -4479,6 +4682,7 @@ dependencies = [ "base64 0.22.1", "litellm-auth-aws", "litellm-core-utils", + "litellm-secret-value", "litellm-secrets-types", "litellm-tracing", "rstest", @@ -4498,6 +4702,7 @@ dependencies = [ "litellm-auth-types", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "percent-encoding", "reqwest", @@ -4519,6 +4724,7 @@ dependencies = [ "base64 0.22.1", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "litellm-tracing", "moka", @@ -4549,6 +4755,7 @@ dependencies = [ "litellm-auth-types", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "moka", "percent-encoding", @@ -4570,6 +4777,7 @@ dependencies = [ "litellm-core-utils", "litellm-error", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "reqwest", "rstest", @@ -4589,7 +4797,7 @@ dependencies = [ name = "litellm-secrets-types" version = "0.1.0" dependencies = [ - "litellm-auth-types", + "litellm-secret-value", "moka", "rstest", "serde", @@ -5258,6 +5466,16 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" +[[package]] +name = "pbkdf2" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" +dependencies = [ + "digest 0.10.7", + "hmac 0.12.1", +] + [[package]] name = "pem" version = "4.0.0" @@ -5368,6 +5586,22 @@ dependencies = [ "spki", ] +[[package]] +name = "pkcs5" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" +dependencies = [ + "aes", + "cbc", + "der", + "des", + "pbkdf2", + "scrypt", + "sha2 0.10.9", + "spki", +] + [[package]] name = "pkcs8" version = "0.10.2" @@ -5375,6 +5609,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" dependencies = [ "der", + "pkcs5", + "rand_core 0.6.4", "spki", ] @@ -5562,6 +5798,7 @@ version = "0.29.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b" dependencies = [ + "indexmap 2.14.0", "libc", "once_cell", "portable-atomic", @@ -5752,17 +5989,6 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" -[[package]] -name = "r2d2" -version = "0.8.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" -dependencies = [ - "log", - "parking_lot", - "scheduled-thread-pool", -] - [[package]] name = "rand" version = "0.8.7" @@ -5936,18 +6162,6 @@ dependencies = [ "xxhash-rust", ] -[[package]] -name = "redis-test" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca" -dependencies = [ - "rand 0.9.5", - "redis", - "socket2 0.6.5", - "tempfile", -] - [[package]] name = "redox_syscall" version = "0.5.18" @@ -6050,7 +6264,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.27.9", "hyper-util", "js-sys", @@ -6406,6 +6620,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "salsa20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" +dependencies = [ + "cipher", +] + [[package]] name = "same-file" version = "1.0.6" @@ -6424,15 +6647,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "scheduled-thread-pool" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" -dependencies = [ - "parking_lot", -] - [[package]] name = "schemars" version = "0.9.0" @@ -6478,6 +6692,17 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scrypt" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0516a385866c09368f0b5bcd1caff3366aace790fcd46e2bb032697bb172fd1f" +dependencies = [ + "pbkdf2", + "salsa20", + "sha2 0.10.9", +] + [[package]] name = "sct" version = "0.7.1" @@ -7540,6 +7765,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -7619,7 +7845,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-timeout", "hyper-util", "percent-encoding", @@ -8195,9 +8421,9 @@ dependencies = [ [[package]] name = "want" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +checksum = "ec4cdd0dd910afe868b7ef477227d8d538b46b3075031afee8a9f2acb0a2ed0b" dependencies = [ "try-lock", ] @@ -8579,7 +8805,7 @@ dependencies = [ "futures", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "log", "once_cell", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 49a33e08456..a00a8ec2446 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" } litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } litellm-router-types = { path = "crates/router-types" } +litellm-router-core = { path = "crates/router-core" } litellm-tracing = { path = "crates/tracing" } litellm-spend-clickhouse = { path = "crates/spend-clickhouse" } litellm-storage-clickhouse = { path = "crates/storage-clickhouse" } @@ -25,6 +26,7 @@ litellm-inference-ocr = { path = "crates/inference-ocr" } litellm-inference-testing = { path = "crates/inference-testing" } litellm-gateway-mcp = { path = "crates/gateway-mcp" } litellm-gateway = { path = "crates/gateway" } +litellm-gateway-forward = { path = "crates/gateway-forward" } litellm-gateway-inference = { path = "crates/gateway-inference" } litellm-gateway-auth = { path = "crates/gateway-auth" } litellm-gateway-management = { path = "crates/gateway-management" } @@ -37,6 +39,7 @@ litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framer = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } litellm-auth-types = { path = "crates/auth-types" } +litellm-secret-value = { path = "crates/secret-value" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } @@ -56,11 +59,10 @@ litellm-db-testing = { path = "crates/db-testing" } litellm-cache = { path = "crates/cache" } litellm-cache-azure-blob = { path = "crates/cache-azure-blob" } litellm-cache-memory = { path = "crates/cache-memory" } -litellm-cache-redis = { path = "crates/cache-redis" } litellm-cache-s3 = { path = "crates/cache-s3" } +litellm-s3 = { path = "crates/s3" } litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-disk = { path = "crates/cache-disk" } -litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" } litellm-cache-response = { path = "crates/cache-response" } litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" } litellm-cache-testing = { path = "crates/cache-testing" } @@ -72,17 +74,25 @@ litellm-host-python = { path = "crates/host-python" } litellm-python-compat = { path = "crates/python-compat" } litellm-circuit-breaker = { path = "crates/circuit-breaker" } litellm-redis-client = { path = "crates/redis-client" } +litellm-redis-connect = { path = "crates/redis-connect" } +litellm-redis-store = { path = "crates/redis-store" } litellm-redis-testing = { path = "crates/redis-testing" } +litellm-redis-resolve = { path = "crates/redis-resolve" } +litellm-redis-types = { path = "crates/redis-types" } litellm-cost = { path = "crates/cost" } litellm-model-catalog = { path = "crates/model-catalog" } litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" } +arc-swap = "1.9.2" +async-trait = "0.1.91" askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] } axum-login = "0.18.0" tower-sessions = { version = "0.14.0", features = ["memory-store"] } bytes = "1" +backon = { version = "1.6.0", default-features = false, features = ["std", "tokio-sleep"] } +aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rt-tokio"] } http = "1" google-cloud-auth = { version = "1.16.0", default-features = false } jsonwebtoken = { version = "11.1.0", default-features = false } @@ -96,7 +106,6 @@ macro_rules_attribute = "0.2.3" schemars = "1" reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] } redis = "1.8.0" -redis-test = "1.0.4" testcontainers-modules = "0.15.0" qdrant-client = { version = "1.19.0", default-features = false } uuid = { version = "1", features = ["v4"] } @@ -152,3 +161,9 @@ panic = "unwind" debug = false incremental = false strip = "symbols" + +[profile.ci] +inherits = "release" +opt-level = 0 +lto = false +codegen-units = 256 diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index e4578777a08..c8ed7bb0f14 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -11,14 +11,16 @@ litellm-error.workspace = true litellm-http.workspace = true reqwest.workspace = true -moka = { workspace = true, features = ["sync"] } +moka.workspace = true +litellm-secret-value.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } sha2.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true -aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] } +aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] } aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] } aws-sigv4 = "1.5.1" @@ -26,6 +28,9 @@ aws-types = "1.4.0" aws-smithy-runtime-api = "1.13.0" [dev-dependencies] +litellm-secrets.workspace = true +aws-smithy-types = "1.6.4" +aws-runtime = { version = "1.9.4", default-features = false } rstest.workspace = true litellm-http = { workspace = true, features = ["test-support"] } bytes.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 81aa4b068bc..e6e32754b4f 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1,13 +1,16 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; -use moka::sync::Cache; -use sha2::{Digest, Sha256}; +use aws_credential_types::provider::{SharedCredentialsProvider, error::CredentialsError}; +use litellm_secret_value::SecretValue; +use moka::{Expiry, future::Cache}; +use std::{future::Future, pin::Pin, sync::Arc}; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ - SignableBody, SignableRequest, SigningParams, SigningSettings, sign, + SignableBody, SignableRequest, SignatureLocation, SigningInstructions, SigningParams, + SigningSettings, sign, }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; @@ -23,72 +26,106 @@ use super::constants::{ }; use crate::error::{CredentialFlow, Error, SigningStage}; -const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); -const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); +const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(600); +const EXPIRY_BUFFER: Duration = Duration::from_secs(60); #[derive(Clone)] pub struct AwsAuthService { - static_credentials: Cache, - ambient_credentials: Cache, + providers: Cache, + credentials: Cache, + loader: Arc, } impl Default for AwsAuthService { fn default() -> Self { + Self::new(Arc::new(NativeProviderLoader)) + } +} + +impl AwsAuthService { + fn new(loader: Arc) -> Self { Self { - static_credentials: Cache::builder() + providers: Cache::builder().max_capacity(200).build(), + credentials: Cache::builder() .max_capacity(200) - .time_to_live(STATIC_CREDENTIALS_TTL) - .build(), - ambient_credentials: Cache::builder() - .max_capacity(200) - .time_to_live(AMBIENT_CREDENTIALS_TTL) + .expire_after(CredentialExpiry) .build(), + loader, } } } -fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { - match flow { - AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL), - AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL), - AwsAuthFlow::WebIdentity { .. } - | AwsAuthFlow::AssumeRole { .. } - | AwsAuthFlow::Profile { .. } - | AwsAuthFlow::SessionToken { .. } => None, +struct CredentialExpiry; + +impl Expiry for CredentialExpiry { + fn expire_after_create( + &self, + _: &AwsAuthConfig, + credentials: &Credentials, + _: std::time::Instant, + ) -> Option { + Some(credentials.expiry().map_or(CREDENTIAL_CACHE_TTL, |expiry| { + expiry + .duration_since(SystemTime::now()) + .unwrap_or_default() + .saturating_sub(EXPIRY_BUFFER) + .min(CREDENTIAL_CACHE_TTL) + })) } } -#[derive(Clone, Debug, Default, PartialEq, Eq)] +type ProviderFuture<'a> = + Pin> + Send + 'a>>; + +trait ProviderLoader: Send + Sync { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_>; +} + +struct NativeProviderLoader; + +impl ProviderLoader for NativeProviderLoader { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_> { + Box::pin(build_provider(config, flow)) + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] pub struct AwsAuthConfig { - pub access_key_id: Option, - pub secret_access_key: Option, - pub session_token: Option, + pub access_key_id: Option, + pub secret_access_key: Option, + pub session_token: Option, pub region_name: Option, pub session_name: Option, pub profile_name: Option, pub role_name: Option, - pub web_identity_token: Option, + pub web_identity_token: Option, pub sts_endpoint: Option, - pub external_id: Option, + pub external_id: Option, } impl AwsAuthConfig { fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option + Sync)) -> Self { Self { - access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)), + access_key_id: self + .access_key_id + .or_else(|| env_lookup(AWS_ACCESS_KEY_ID).map(SecretValue::new)), secret_access_key: self .secret_access_key - .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)), - session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)), + .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY).map(SecretValue::new)), + session_token: self + .session_token + .or_else(|| env_lookup(AWS_SESSION_TOKEN).map(SecretValue::new)), region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)), session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)), profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)), role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)), web_identity_token: self .web_identity_token - .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)), + .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN).map(SecretValue::new)), sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)), - external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)), + external_id: self + .external_id + .or_else(|| env_lookup(AWS_EXTERNAL_ID).map(SecretValue::new)), } } } @@ -96,52 +133,31 @@ impl AwsAuthConfig { #[derive(Clone, Debug, PartialEq, Eq)] pub enum AwsAuthFlow { WebIdentity { - token: String, + token: SecretValue, role: String, session_name: String, }, AssumeRole { role: String, session_name: Option, + web_identity_role: Option, }, Profile { name: String, }, SessionToken { - access_key_id: String, - secret_access_key: String, - session_token: String, + access_key_id: SecretValue, + secret_access_key: SecretValue, + session_token: SecretValue, }, StaticKeys { - access_key_id: String, - secret_access_key: String, + access_key_id: SecretValue, + secret_access_key: SecretValue, region_name: String, }, DefaultChain, } -fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String { - let mut hasher = Sha256::new(); - hasher.update(format!("{config:?}:{flow:?}")); - format!("{:x}", hasher.finalize()) -} - -impl AwsAuthService { - fn get_cached_credentials(&self, key: &str) -> Option { - self.static_credentials - .get(key) - .or_else(|| self.ambient_credentials.get(key)) - } - - fn set_cached_credentials(&self, key: String, credentials: Credentials, ttl: Duration) { - if ttl == STATIC_CREDENTIALS_TTL { - self.static_credentials.insert(key, credentials); - } else { - self.ambient_credentials.insert(key, credentials); - } - } -} - fn role_identity(arn: &str) -> Option<(&str, &str, &str)> { let mut parts = arn.splitn(6, ':'); let ("arn", partition, _, _, account, resource) = ( @@ -170,6 +186,9 @@ pub fn classify_auth( config: AwsAuthConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> AwsAuthFlow { + let web_identity_role = env_lookup(AWS_WEB_IDENTITY_TOKEN_FILE) + .filter(|token_file| !token_file.is_empty()) + .and_then(|_| env_lookup(AWS_ROLE_ARN)); let config = config.with_environment(env_lookup); if let (Some(token), Some(role), Some(session_name)) = ( config.web_identity_token.clone(), @@ -186,6 +205,7 @@ pub fn classify_auth( return AwsAuthFlow::AssumeRole { role, session_name: config.session_name.clone(), + web_identity_role, }; } if let Some(name) = config.profile_name { @@ -216,6 +236,30 @@ pub fn classify_auth( AwsAuthFlow::DefaultChain } +async fn profile_provider( + name: &str, + region: Option, + config: aws_config::provider_config::ProviderConfig, + builder: aws_config::profile::credentials::Builder, +) -> Result { + let region = match region { + Some(region) => Some(aws_types::region::Region::new(region)), + None => { + aws_config::default_provider::region::DefaultRegionChain::builder() + .profile_name(name) + .build() + .region() + .await + } + }; + Ok(SharedCredentialsProvider::new( + builder + .configure(&config.with_region(region)) + .profile_name(name) + .build(), + )) +} + impl AwsAuthService { pub async fn resolve_credentials( &self, @@ -224,183 +268,240 @@ impl AwsAuthService { ) -> Result { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); - match flow { - AwsAuthFlow::SessionToken { - access_key_id, - secret_access_key, - session_token, - } => Ok(Credentials::new( - access_key_id, - secret_access_key, - Some(session_token), - None, - "litellm-static-session", - )), - AwsAuthFlow::StaticKeys { - access_key_id, - secret_access_key, - region_name, - } => { - let flow = AwsAuthFlow::StaticKeys { - access_key_id: access_key_id.clone(), - secret_access_key: secret_access_key.clone(), - region_name, - }; - let key = cache_key(&resolved, &flow); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let credentials = Credentials::new( - access_key_id, - secret_access_key, - None, - None, - "litellm-static", - ); - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL), - ); - Ok(credentials) - } - AwsAuthFlow::Profile { name } => { - let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() - .profile_name(name) - .build(); + self.fresh_credentials(&resolved, &flow).await + } + + /// The credentials botocore's default chain finds in `region`, without LiteLLM's Bedrock settings + /// (`AWS_ROLE_NAME`, `AWS_PROFILE_NAME`, `AWS_SESSION_NAME`) that `resolve_credentials` also reads. + pub async fn default_chain_credentials(&self, region: &str) -> Result { + let config = AwsAuthConfig { + region_name: Some(region.to_owned()), + ..Default::default() + }; + self.fresh_credentials(&config, &AwsAuthFlow::DefaultChain) + .await + } + + async fn fresh_credentials( + &self, + resolved: &AwsAuthConfig, + flow: &AwsAuthFlow, + ) -> Result { + let cached = self.cached_credentials(resolved, flow).await?; + if fresh(&cached) { + return Ok(cached); + } + self.credentials.invalidate(resolved).await; + let refetched = self.cached_credentials(resolved, flow).await?; + if fresh(&refetched) { + return Ok(refetched); + } + self.credentials.invalidate(resolved).await; + Err(Error::ExpiredCredentials) + } + + async fn cached_credentials( + &self, + resolved: &AwsAuthConfig, + flow: &AwsAuthFlow, + ) -> Result { + self.credentials + .try_get_with(resolved.clone(), async { + let provider = self + .providers + .try_get_with( + resolved.clone(), + self.loader.load(resolved.clone(), flow.clone()), + ) + .await + .map_err(|error| (*error).clone())?; provider .provide_credentials() .await - .map_err(Error::credentials(CredentialFlow::Profile)) - } - AwsAuthFlow::AssumeRole { role, session_name } => { - if is_already_running_as_role(&role, &resolved).await? { - let ambient_flow = AwsAuthFlow::DefaultChain; - let key = cache_key(&resolved, &ambient_flow); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::DefaultChain))?; - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - return Ok(credentials); - } - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name.clone() { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint.clone() { - loader = loader.endpoint_url(endpoint); - } - if let (Some(access_key_id), Some(secret_access_key)) = - (resolved.access_key_id, resolved.secret_access_key) - { - loader = loader.credentials_provider(Credentials::new( - access_key_id, - secret_access_key, - resolved.session_token, - None, - "litellm-role-source", - )); - } - let sdk_config = loader.load().await; - let builder = aws_config::sts::AssumeRoleProvider::builder(role); - let builder = match session_name { - Some(name) => builder.session_name(name), - None => builder.session_name(default_session_name()), - }; - let builder = match resolved.external_id { - Some(id) => builder.external_id(id), - None => builder, - }; - let provider = builder.configure(&sdk_config).build().await; - provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::AssumeRole)) - } - AwsAuthFlow::WebIdentity { - token, - role, - session_name, - } => { - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint { - loader = loader.endpoint_url(endpoint); - } - let sdk_config = loader.load().await; - let client = aws_sdk_sts::Client::new(&sdk_config); - let response = client - .assume_role_with_web_identity() - .role_arn(role) - .role_session_name(session_name) - .web_identity_token(token) - .send() - .await - .map_err(Error::credentials(CredentialFlow::WebIdentity))?; - let credentials = response - .credentials() - .ok_or(Error::MissingWebIdentityCredentials)?; - let expiration = - SystemTime::try_from(*credentials.expiration()).map_err(|source| { - Error::WebIdentityExpiration(litellm_error::Source::new(source)) - })?; - Ok(Credentials::new( - credentials.access_key_id(), - credentials.secret_access_key(), - Some(credentials.session_token().to_string()), - Some(expiration), - "litellm-web-identity", - )) - } - AwsAuthFlow::DefaultChain => { - let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::DefaultChain))?; - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&AwsAuthFlow::DefaultChain) - .unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - Ok(credentials) + .map_err(Error::credentials(flow.credential_flow())) + }) + .await + .map_err(|error| (*error).clone()) + } +} + +fn fresh(credentials: &Credentials) -> bool { + credentials + .expiry() + .is_none_or(|expiry| expiry > SystemTime::now() + EXPIRY_BUFFER) +} + +impl AwsAuthFlow { + fn credential_flow(&self) -> CredentialFlow { + match self { + Self::Profile { .. } => CredentialFlow::Profile, + Self::AssumeRole { .. } => CredentialFlow::AssumeRole, + Self::WebIdentity { .. } => CredentialFlow::WebIdentity, + Self::DefaultChain | Self::StaticKeys { .. } | Self::SessionToken { .. } => { + CredentialFlow::DefaultChain } } } } -async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result { +async fn build_provider( + resolved: AwsAuthConfig, + flow: AwsAuthFlow, +) -> Result { + match flow { + AwsAuthFlow::SessionToken { + access_key_id, + secret_access_key, + session_token, + } => Ok(SharedCredentialsProvider::new(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + Some(session_token.expose().to_owned()), + None, + "litellm-static-session", + ))), + AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + .. + } => Ok(SharedCredentialsProvider::new(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + None, + None, + "litellm-static", + ))), + AwsAuthFlow::Profile { name } => { + profile_provider( + &name, + resolved.region_name, + aws_config::provider_config::ProviderConfig::without_region(), + aws_config::profile::ProfileFileCredentialsProvider::builder(), + ) + .await + } + AwsAuthFlow::DefaultChain => { + let builder = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder(); + let builder = match resolved.region_name { + Some(region) => builder.region(aws_types::region::Region::new(region)), + None => builder, + }; + Ok(SharedCredentialsProvider::new(builder.build().await)) + } + AwsAuthFlow::AssumeRole { + role, + session_name, + web_identity_role, + } => { + if is_already_running_as_role(&role, web_identity_role.as_deref(), &resolved).await? { + return Ok(SharedCredentialsProvider::new( + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await, + )); + } + let sdk_config = sdk_config(resolved.clone()).await; + let builder = aws_config::sts::AssumeRoleProvider::builder(role) + .session_name(session_name.unwrap_or_else(default_session_name)); + let builder = match resolved.external_id { + Some(id) => builder.external_id(id.expose()), + None => builder, + }; + Ok(SharedCredentialsProvider::new( + builder.configure(&sdk_config).build().await, + )) + } + AwsAuthFlow::WebIdentity { + token, + role, + session_name, + } => Ok(SharedCredentialsProvider::new(WebIdentityProvider { + client: aws_sdk_sts::Client::new(&sdk_config(resolved).await), + token, + role, + session_name, + })), + } +} + +async fn sdk_config(config: AwsAuthConfig) -> aws_types::SdkConfig { + let loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + let loader = match config.region_name { + Some(region) => loader.region(aws_types::region::Region::new(region)), + None => loader, + }; + let loader = match config.sts_endpoint { + Some(endpoint) => loader.endpoint_url(endpoint), + None => loader, + }; + let loader = match (config.access_key_id, config.secret_access_key) { + (Some(key), Some(secret)) => loader.credentials_provider(Credentials::new( + key.expose(), + secret.expose(), + config.session_token.map(|token| token.expose().to_owned()), + None, + "litellm-role-source", + )), + _ => loader, + }; + loader.load().await +} + +#[derive(Debug)] +struct WebIdentityProvider { + client: aws_sdk_sts::Client, + token: SecretValue, + role: String, + session_name: String, +} + +impl ProvideCredentials for WebIdentityProvider { + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async { + let response = self + .client + .assume_role_with_web_identity() + .role_arn(&self.role) + .role_session_name(&self.session_name) + .web_identity_token(self.token.expose()) + .send() + .await + .map_err(CredentialsError::provider_error)?; + let credentials = response.credentials().ok_or_else(|| { + CredentialsError::provider_error(Error::MissingWebIdentityCredentials) + })?; + let expiration = SystemTime::try_from(*credentials.expiration()) + .map_err(CredentialsError::provider_error)?; + Ok(Credentials::new( + credentials.access_key_id(), + credentials.secret_access_key(), + Some(credentials.session_token().to_owned()), + Some(expiration), + "litellm-web-identity", + )) + }) + } +} + +async fn is_already_running_as_role( + role: &str, + web_identity_role: Option<&str>, + config: &AwsAuthConfig, +) -> Result { + if config.access_key_id.is_some() || config.secret_access_key.is_some() { + return Ok(false); + } if role_identity(role).is_none() { return Ok(false); } - if let (Ok(current_role), Ok(token_file)) = ( - std::env::var(AWS_ROLE_ARN), - std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), - ) && !token_file.is_empty() - { - return Ok(same_role_arns(role, ¤t_role)); + if let Some(current_role) = web_identity_role { + return Ok(same_role_arns(role, current_role)); } let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); @@ -473,16 +574,6 @@ pub fn sign_headers( headers: &HeaderMap, scope: SigningScope<'_>, ) -> Result { - let identity: Identity = scope.credentials.clone().into(); - let params = v4::SigningParams::builder() - .identity(&identity) - .region(scope.region) - .name(scope.service) - .time(scope.time) - .settings(SigningSettings::default()) - .build() - .map(SigningParams::from) - .map_err(Error::signing(SigningStage::Parameters))?; let header_refs = headers .iter() .map(|(name, value)| { @@ -494,16 +585,14 @@ pub fn sign_headers( }) }) .collect::, _>>()?; - let request = SignableRequest::new( + let instructions = sign_request( method, url, header_refs.into_iter(), - SignableBody::Bytes(body), - ) - .map_err(Error::signing(SigningStage::Request))?; - let (instructions, _) = sign(request, ¶ms) - .map_err(Error::signing(SigningStage::Signature))? - .into_parts(); + body, + SigningSettings::default(), + scope, + )?; instructions .headers() .map(|(name, value)| { @@ -516,6 +605,60 @@ pub fn sign_headers( .collect() } +pub fn presign_url( + url: &str, + expires_in: Duration, + scope: SigningScope<'_>, +) -> Result { + let mut signed = url::Url::parse(url).map_err(Error::signing(SigningStage::Request))?; + let mut settings = SigningSettings::default(); + settings.signature_location = SignatureLocation::QueryParams; + settings.expires_in = Some(expires_in); + let instructions = sign_request( + "GET", + signed.as_str(), + std::iter::empty(), + &[], + settings, + scope, + )?; + let (signature, auth): (Vec<_>, Vec<_>) = instructions + .params() + .iter() + .map(|(name, value)| (*name, value.as_ref())) + .partition(|(name, _)| *name == "X-Amz-Signature"); + signed + .query_pairs_mut() + .extend_pairs(auth.into_iter().chain(signature)); + Ok(signed.into()) +} + +fn sign_request<'a>( + method: &'a str, + url: &'a str, + headers: impl Iterator, + body: &'a [u8], + settings: SigningSettings, + scope: SigningScope<'_>, +) -> Result { + let identity: Identity = scope.credentials.clone().into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(scope.region) + .name(scope.service) + .time(scope.time) + .settings(settings) + .build() + .map(SigningParams::from) + .map_err(Error::signing(SigningStage::Parameters))?; + let request = SignableRequest::new(method, url, headers, SignableBody::Bytes(body)) + .map_err(Error::signing(SigningStage::Request))?; + let (instructions, _) = sign(request, ¶ms) + .map_err(Error::signing(SigningStage::Signature))? + .into_parts(); + Ok(instructions) +} + /// Model-id and region parsing shared by every Bedrock route. pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { let mut stripped = model; @@ -586,16 +729,26 @@ pub fn aws_auth_config( env_lookup: &dyn Fn(&str) -> Option, ) -> AwsAuthConfig { AwsAuthConfig { - access_key_id: params.resolve(&AwsParams::ACCESS_KEY_ID, env_lookup), - secret_access_key: params.resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup), - session_token: params.resolve(&AwsParams::SESSION_TOKEN, env_lookup), + access_key_id: params + .resolve(&AwsParams::ACCESS_KEY_ID, env_lookup) + .map(SecretValue::new), + secret_access_key: params + .resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup) + .map(SecretValue::new), + session_token: params + .resolve(&AwsParams::SESSION_TOKEN, env_lookup) + .map(SecretValue::new), region_name: params.resolve(&AwsParams::REGION, env_lookup), session_name: params.resolve(&AwsParams::SESSION_NAME, env_lookup), profile_name: params.resolve(&AwsParams::PROFILE_NAME, env_lookup), role_name: params.resolve(&AwsParams::ROLE_NAME, env_lookup), - web_identity_token: params.resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup), + web_identity_token: params + .resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup) + .map(SecretValue::new), sts_endpoint: params.resolve(&AwsParams::STS_ENDPOINT, env_lookup), - external_id: params.resolve(&AwsParams::EXTERNAL_ID, env_lookup), + external_id: params + .resolve(&AwsParams::EXTERNAL_ID, env_lookup) + .map(SecretValue::new), } } @@ -660,10 +813,133 @@ mod tests { use crate::constants::{AWS_REGION, BEDROCK_SERVICE}; use reqwest::header::CONTENT_TYPE; + #[cfg(unix)] + #[rstest::fixture] + fn process_profile() -> aws_runtime::env_config::file::EnvConfigFiles { + let credentials = serde_json::json!({ + "Version": 1, + "AccessKeyId": "process-access", + "SecretAccessKey": "process-secret", + "SessionToken": "process-session", + }); + aws_runtime::env_config::file::EnvConfigFiles::builder() + .include_default_config_file(false) + .include_default_credentials_file(false) + .with_contents( + aws_runtime::env_config::file::EnvConfigFileKind::Config, + format!( + "[profile smoke]\ncredential_process = /usr/bin/printf '%s' '{credentials}'\n" + ), + ) + .build() + } + + #[cfg(unix)] + #[rstest::rstest] + #[tokio::test] + async fn profile_process_credentials_can_be_loaded( + process_profile: aws_runtime::env_config::file::EnvConfigFiles, + ) { + let provider = profile_provider( + "smoke", + Some("us-west-2".into()), + aws_config::provider_config::ProviderConfig::without_region(), + aws_config::profile::ProfileFileCredentialsProvider::builder() + .profile_files(process_profile), + ) + .await + .unwrap(); + let credentials = provider.provide_credentials().await.unwrap(); + + assert_eq!(credentials.access_key_id(), "process-access"); + assert_eq!(credentials.secret_access_key(), "process-secret"); + assert_eq!(credentials.session_token(), Some("process-session")); + } + + #[derive(Debug)] + struct StsConnector; + + impl aws_smithy_runtime_api::client::http::HttpConnector for StsConnector { + fn call( + &self, + request: aws_smithy_runtime_api::client::orchestrator::HttpRequest, + ) -> aws_smithy_runtime_api::client::http::HttpConnectorFuture { + assert!(request.uri().contains("sts.us-west-2.")); + aws_smithy_runtime_api::client::http::HttpConnectorFuture::new(async { + Ok( + aws_smithy_runtime_api::client::orchestrator::HttpResponse::new( + 200_u16.try_into().unwrap(), + aws_smithy_types::body::SdkBody::from( + r#" + + role-accessrole-secret + role-session2100-01-01T00:00:00Z + arn:aws:sts::123456789012:assumed-role/smoke/session + smoke:session + smoke"#, + ), + ), + ) + }) + } + } + + #[rstest::rstest] + #[tokio::test] + async fn role_profile_uses_region_when_acquiring_sts_credentials() { + use aws_smithy_runtime_api::client::http::{SharedHttpConnector, http_client_fn}; + let files = aws_runtime::env_config::file::EnvConfigFiles::builder() + .include_default_config_file(false) + .include_default_credentials_file(false) + .with_contents( + aws_runtime::env_config::file::EnvConfigFileKind::Config, + "[profile smoke]\nrole_arn = arn:aws:iam::123456789012:role/smoke\nsource_profile = source\n\n[profile source]\naws_access_key_id = source-access\naws_secret_access_key = source-secret\n", + ).build(); + let config = aws_config::provider_config::ProviderConfig::without_region() + .with_http_client(http_client_fn(|_, _| { + SharedHttpConnector::new(StsConnector) + })); + let credentials = profile_provider( + "smoke", + Some("us-west-2".into()), + config, + aws_config::profile::ProfileFileCredentialsProvider::builder().profile_files(files), + ) + .await + .unwrap(); + let credentials = credentials.provide_credentials().await.unwrap(); + assert_eq!(credentials.access_key_id(), "role-access"); + assert_eq!(credentials.secret_access_key(), "role-secret"); + assert_eq!(credentials.session_token(), Some("role-session")); + assert!(credentials.expiry().is_some()); + } + fn no_env(_: &str) -> Option { None } + #[rstest::rstest] + #[case::access_key(Some("source-access"), None)] + #[case::secret_key(None, Some("source-secret"))] + #[case::source_pair(Some("source-access"), Some("source-secret"))] + #[tokio::test] + async fn explicit_source_keys_cannot_reuse_an_ambient_role( + #[case] access_key: Option<&str>, + #[case] secret_key: Option<&str>, + ) { + let role = "arn:aws:iam::123456789012:role/target"; + let config = AwsAuthConfig { + access_key_id: access_key.map(SecretValue::new), + secret_access_key: secret_key.map(SecretValue::new), + ..Default::default() + }; + assert!( + !is_already_running_as_role(role, Some(role), &config) + .await + .unwrap() + ); + } + #[test] fn secret_names_cover_environment_reads() { let seen = std::sync::Arc::new(std::sync::Mutex::new( @@ -706,7 +982,10 @@ mod tests { }; assert_eq!( - aws_auth_config(¶ms, &env).secret_access_key.as_deref(), + aws_auth_config(¶ms, &env) + .secret_access_key + .as_ref() + .map(SecretValue::expose), expected ); assert_eq!( @@ -831,6 +1110,53 @@ mod tests { ); } + #[rstest::rstest] + #[case::running_with_web_identity("/var/run/token", Some("arn:aws:iam::1:role/current"))] + #[case::empty_token_file("", None)] + #[tokio::test] + async fn assume_role_reads_the_ambient_role_from_the_secret_snapshot( + #[case] token_file: &str, + #[case] expected: Option<&str>, + ) { + use litellm_secrets::source::SecretSource; + + struct WebIdentitySecrets<'a>(&'a str); + + impl SecretSource for WebIdentitySecrets<'_> { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> futures_util::future::BoxFuture< + 'a, + Result, litellm_secrets::Error>, + > { + Box::pin(async move { + Ok(match name { + AWS_WEB_IDENTITY_TOKEN_FILE => Some(SecretValue::new(self.0)), + AWS_ROLE_ARN => Some(SecretValue::new("arn:aws:iam::1:role/current")), + _ => None, + }) + }) + } + } + + let snapshot = WebIdentitySecrets(token_file) + .resolve(AwsParams::secret_names()) + .await + .unwrap(); + let config = AwsAuthConfig { + role_name: Some("arn:aws:iam::1:role/target".into()), + ..Default::default() + }; + let AwsAuthFlow::AssumeRole { + web_identity_role, .. + } = classify_auth(config, &|name| snapshot.get(name)) + else { + panic!("a role name selects AssumeRole"); + }; + assert_eq!(web_identity_role.as_deref(), expected); + } + #[tokio::test] async fn static_credentials_do_not_use_network() { let credentials = AwsAuthService::default() @@ -849,110 +1175,231 @@ mod tests { assert_eq!(credentials.session_token(), None); } - #[test] - fn cache_policy_matches_python_flows() { - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::StaticKeys { - access_key_id: "ak".into(), - secret_access_key: "sk".into(), - region_name: "us-east-1".into(), - }), - Some(STATIC_CREDENTIALS_TTL) - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::DefaultChain), - Some(AMBIENT_CREDENTIALS_TTL) - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::SessionToken { - access_key_id: "ak".into(), - secret_access_key: "sk".into(), - session_token: "token".into(), - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::Profile { - name: "profile".into() - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::AssumeRole { - role: "arn:aws:iam::123456789012:role/demo".into(), - session_name: None, - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::WebIdentity { - token: "token".into(), - role: "arn:aws:iam::123456789012:role/demo".into(), - session_name: "session".into(), - }), - None - ); + struct CountingLoader { + loads: Arc, + calls: Arc, + short_lifetime: Duration, + short_calls: usize, + } + + impl ProviderLoader for CountingLoader { + fn load(&self, config: AwsAuthConfig, _: AwsAuthFlow) -> ProviderFuture<'_> { + let loads = self.loads.clone(); + let provider = CountingProvider { + calls: self.calls.clone(), + short_lifetime: self.short_lifetime, + short_calls: self.short_calls, + key: config.access_key_id.unwrap(), + }; + Box::pin(async move { + loads.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(SharedCredentialsProvider::new(provider)) + }) + } + } + + #[derive(Debug)] + struct CountingProvider { + calls: Arc, + short_lifetime: Duration, + short_calls: usize, + key: SecretValue, + } + + impl ProvideCredentials for CountingProvider { + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async move { + tokio::task::yield_now().await; + let call = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let lifetime = if call < self.short_calls { + self.short_lifetime + } else { + Duration::from_secs(3600) + }; + Ok(Credentials::new( + self.key.expose(), + "credential-value", + Some(format!("session-{call}")), + Some(SystemTime::now() + lifetime), + "counting", + )) + }) + } + } + + #[derive(Default)] + struct RecordingLoader { + loads: std::sync::Mutex>, + } + + impl ProviderLoader for RecordingLoader { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_> { + self.loads.lock().unwrap().push((config, flow)); + Box::pin(async { + Ok(SharedCredentialsProvider::new(Credentials::new( + "ak", + "sk", + None, + None, + "recording", + ))) + }) + } } #[rstest::rstest] - fn cache_round_trip_preserves_credentials() { - let auth = AwsAuthService::default(); - let key = format!("cache-test-{}", std::process::id()); - let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test"); - auth.set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL); + #[tokio::test] + async fn default_chain_credentials_do_not_share_a_configured_role_in_the_same_region() { + let loader = Arc::new(RecordingLoader::default()); + let auth = AwsAuthService::new(loader.clone()); + let bedrock_env = |name: &str| match name { + AWS_ROLE_NAME => Some("arn:aws:iam::123456789012:role/bedrock".to_string()), + AWS_REGION_NAME => Some("us-east-1".to_string()), + _ => None, + }; + auth.resolve_credentials(AwsAuthConfig::default(), &bedrock_env) + .await + .unwrap(); + + auth.default_chain_credentials("us-east-1").await.unwrap(); + auth.default_chain_credentials("us-east-1").await.unwrap(); + + let loads = loader.loads.lock().unwrap(); + assert_eq!(loads.len(), 2); assert_eq!( - auth.get_cached_credentials(&key) - .map(|value| value.access_key_id().to_string()), - Some("cache-ak".to_string()) + loads[1], + ( + AwsAuthConfig { + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + AwsAuthFlow::DefaultChain + ) ); } #[rstest::rstest] #[tokio::test] - async fn cloned_services_reuse_credentials_but_independent_services_do_not() { - let auth = AwsAuthService::default(); + async fn near_expiry_credentials_are_refreshed_without_rebuilding_the_provider() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let loads = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: loads.clone(), + calls: calls.clone(), + short_lifetime: Duration::from_secs(30), + short_calls: 1, + })); let config = AwsAuthConfig { - access_key_id: Some("configured-key".into()), - secret_access_key: Some("configured-secret".into()), - region_name: Some("us-east-1".into()), - ..AwsAuthConfig::default() + access_key_id: Some("identity".into()), + ..Default::default() }; - let flow = classify_auth(config.clone(), &no_env); - let cached = Credentials::new("cached-key", "cached-secret", None, None, "test"); - auth.set_cached_credentials( - cache_key(&config, &flow), - cached.clone(), - STATIC_CREDENTIALS_TTL, - ); - - let reused = auth + let first = auth + .resolve_credentials(config.clone(), &no_env) + .await + .unwrap(); + let second = auth .clone() .resolve_credentials(config.clone(), &no_env) .await .unwrap(); - let independent = AwsAuthService::default() - .resolve_credentials(config.clone(), &no_env) - .await - .unwrap(); - let different = AwsAuthConfig { - access_key_id: Some("different-key".into()), - ..config.clone() - }; - let other_identity = auth - .resolve_credentials(different.clone(), &no_env) - .await - .unwrap(); + assert_eq!(first.session_token(), Some("session-1")); + assert_eq!(first.session_token(), second.session_token()); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } - assert_eq!(reused.access_key_id(), cached.access_key_id()); - assert_eq!(reused.secret_access_key(), cached.secret_access_key()); + #[rstest::rstest] + #[case::recovered_by_one_refetch(1, Some("session-1"))] + #[case::still_expired_after_refetch(2, None)] + #[tokio::test] + async fn expired_credentials_are_refetched_once( + #[case] short_calls: usize, + #[case] session: Option<&str>, + ) { + use std::sync::atomic::{AtomicUsize, Ordering}; + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: Arc::new(AtomicUsize::new(0)), + calls: calls.clone(), + short_lifetime: Duration::ZERO, + short_calls, + })); + let config = AwsAuthConfig { + access_key_id: Some("identity".into()), + ..Default::default() + }; + let result = auth.resolve_credentials(config, &no_env).await; assert_eq!( - Some(independent.access_key_id()), - config.access_key_id.as_deref() + result.as_ref().ok().and_then(Credentials::session_token), + session ); - assert_eq!( - Some(other_identity.access_key_id()), - different.access_key_id.as_deref() + assert!(session.is_some() || matches!(result, Err(Error::ExpiredCredentials))); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[tokio::test] + async fn concurrent_calls_share_acquisition_and_distinct_identities_stay_isolated() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let loads = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: loads.clone(), + calls: calls.clone(), + short_lifetime: Duration::from_secs(3600), + short_calls: 0, + })); + let first = AwsAuthConfig { + access_key_id: Some("first".into()), + ..Default::default() + }; + let second = AwsAuthConfig { + access_key_id: Some("second".into()), + ..Default::default() + }; + let (a, b, c) = tokio::join!( + auth.resolve_credentials(first.clone(), &no_env), + auth.resolve_credentials(first, &no_env), + auth.resolve_credentials(second, &no_env) ); + let a = a.unwrap(); + let b = b.unwrap(); + let c = c.unwrap(); + assert_eq!(a.session_token(), b.session_token()); + assert_ne!(a.access_key_id(), c.access_key_id()); + assert_eq!(loads.load(Ordering::SeqCst), 2); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest::rstest] + fn auth_config_debug_never_discloses_credentials() { + let config = AwsAuthConfig { + access_key_id: Some("credential-value".into()), + secret_access_key: Some("credential-value".into()), + session_token: Some("credential-value".into()), + web_identity_token: Some("credential-value".into()), + external_id: Some("credential-value".into()), + ..Default::default() + }; + let output = format!("{config:?}"); + assert!(!output.contains("credential-value")); + assert!(output.contains("[REDACTED]")); + } + + #[rstest::rstest] + #[case::session(AwsAuthFlow::SessionToken { access_key_id: "credential-value".into(), secret_access_key: "credential-value".into(), session_token: "credential-value".into() })] + #[case::static_keys(AwsAuthFlow::StaticKeys { access_key_id: "credential-value".into(), secret_access_key: "credential-value".into(), region_name: "region".into() })] + #[case::web_identity(AwsAuthFlow::WebIdentity { token: "credential-value".into(), role: "role".into(), session_name: "session".into() })] + fn auth_flow_debug_never_discloses_credentials(#[case] flow: AwsAuthFlow) { + let output = format!("{flow:?}"); + assert!(!output.contains("credential-value")); + assert!(output.contains("[REDACTED]")); } #[test] @@ -1142,8 +1589,8 @@ mod tests { let credentials = AwsAuthService::default() .resolve_credentials( AwsAuthConfig { - access_key_id: Some(access_key_id), - secret_access_key: Some(secret_access_key), + access_key_id: Some(access_key_id.into()), + secret_access_key: Some(secret_access_key.into()), region_name: Some("us-west-2".to_string()), ..Default::default() }, @@ -1191,4 +1638,44 @@ mod tests { failures.join("; ") ); } + + #[rstest::rstest] + #[case::provisioned( + None, + "https://cache-1/?Action=connect&User=iam-user", + "cache-1/?Action=connect&User=iam-user&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIDEXAMPLE%2F20261029%2Fus-east-1%2Felasticache%2Faws4_request&X-Amz-Date=20261029T120000Z&X-Amz-Expires=900&X-Amz-SignedHeaders=host&X-Amz-Signature=dfa58492a6ae9ba065088724170a52437329d03f3e9bc9a4c90a80cc8a18e2f6" + )] + #[case::serverless_with_a_session_token( + Some("session-token"), + "https://cache-1/?Action=connect&User=iam-user&ResourceType=ServerlessCache", + "cache-1/?Action=connect&User=iam-user&ResourceType=ServerlessCache&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIDEXAMPLE%2F20261029%2Fus-east-1%2Felasticache%2Faws4_request&X-Amz-Date=20261029T120000Z&X-Amz-Expires=900&X-Amz-SignedHeaders=host&X-Amz-Security-Token=session-token&X-Amz-Signature=e48bdbb3bf9d5e8f1d325bf33d5260e333df1f6b8868e8edd1e4c3cc01b4c3f9" + )] + fn presigned_urls_match_botocore_sigv4_query_auth( + #[case] session_token: Option<&str>, + #[case] url: &str, + #[case] expected: &str, + ) { + // Expected values from botocore 1.43.117 SigV4QueryAuth with get_current_datetime patched to + // 2026-10-29T12:00:00Z (generated 2026-10-29), the signer Python's ElastiCache IAM provider uses + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + session_token.map(str::to_owned), + None, + "test", + ); + let signed = presign_url( + url, + Duration::from_secs(900), + SigningScope { + region: "us-east-1", + service: "elasticache", + credentials: &credentials, + time: UNIX_EPOCH + Duration::from_secs(1_793_275_200), + }, + ) + .unwrap(); + + assert_eq!(signed.strip_prefix("https://"), Some(expected)); + } } diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index 789fde156bd..0e4662eda9e 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -11,8 +11,8 @@ pub const AWS_SESSION_NAME: &str = AwsParams::SESSION_NAME.env[0]; pub const AWS_PROFILE_NAME: &str = AwsParams::PROFILE_NAME.env[0]; pub const AWS_ROLE_NAME: &str = AwsParams::ROLE_NAME.env[0]; pub const AWS_WEB_IDENTITY_TOKEN: &str = AwsParams::WEB_IDENTITY_TOKEN.env[0]; -pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN"; -pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +pub const AWS_ROLE_ARN: &str = AwsParams::ROLE_ARN_ENV; +pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV; pub const AWS_STS_ENDPOINT: &str = AwsParams::STS_ENDPOINT.env[0]; pub const AWS_EXTERNAL_ID: &str = AwsParams::EXTERNAL_ID.env[0]; pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK"; diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index 5f88fc8a171..5445de1e6c3 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -2,6 +2,8 @@ use litellm_error::Source; #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { + #[error("AWS credential provider returned expired credentials")] + ExpiredCredentials, #[error("could not load AWS {flow} credentials")] Credentials { flow: CredentialFlow, diff --git a/litellm-rust/crates/auth-aws/src/lib.rs b/litellm-rust/crates/auth-aws/src/lib.rs index bd0f866371f..c8c5b5e8653 100644 --- a/litellm-rust/crates/auth-aws/src/lib.rs +++ b/litellm-rust/crates/auth-aws/src/lib.rs @@ -1,9 +1,11 @@ mod aws; pub mod constants; mod error; +mod provider; mod signer; pub use aws::*; pub use aws_credential_types::Credentials; pub use error::{CredentialFlow, Error, SigningStage}; +pub use provider::AwsCredentialsProvider; pub use signer::SigV4Signer; diff --git a/litellm-rust/crates/auth-aws/src/provider.rs b/litellm-rust/crates/auth-aws/src/provider.rs new file mode 100644 index 00000000000..6cb2325fc11 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/provider.rs @@ -0,0 +1,135 @@ +use std::sync::Arc; + +use aws_credential_types::Credentials; +use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; + +use crate::{AwsAuthConfig, AwsAuthService}; + +type Environment = Arc Option + Send + Sync>; + +#[derive(Clone)] +pub struct AwsCredentialsProvider { + auth: AwsAuthService, + config: Box, + environment: Environment, +} + +impl AwsCredentialsProvider { + /// Explicit keys in `config` as given, otherwise LiteLLM's `AWS_*` settings read through `environment` + pub fn configured( + auth: AwsAuthService, + config: AwsAuthConfig, + environment: impl Fn(&str) -> Option + Send + Sync + 'static, + ) -> Self { + Self { + auth, + config: Box::new(config), + environment: Arc::new(environment), + } + } + + async fn credentials(&self) -> Result { + match (&self.config.access_key_id, &self.config.secret_access_key) { + (Some(access_key_id), Some(secret_access_key)) => Ok(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + self.config + .session_token + .as_ref() + .map(|token| token.expose().to_owned()), + None, + "litellm-configured", + )), + _ => { + self.auth + .resolve_credentials( + AwsAuthConfig::clone(&self.config), + self.environment.as_ref(), + ) + .await + } + } + } +} + +impl ProvideCredentials for AwsCredentialsProvider { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::new(async { + self.credentials() + .await + .map_err(CredentialsError::provider_error) + }) + } +} + +impl std::fmt::Debug for AwsCredentialsProvider { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("AwsCredentialsProvider") + .finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn explicit(session_token: Option<&str>) -> AwsAuthConfig { + AwsAuthConfig { + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), + session_token: session_token.map(Into::into), + region_name: Some("us-east-1".to_string()), + ..Default::default() + } + } + + fn ambient_token(name: &str) -> Option { + (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()) + } + + fn environment_keys(name: &str) -> Option { + match name { + "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()), + "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()), + "AWS_SESSION_TOKEN" => Some("env-token".to_string()), + _ => None, + } + } + + #[rstest] + #[case::explicit_keys_ignore_an_ambient_session_token( + explicit(None), ambient_token, ("key", "secret", None) + )] + #[case::explicit_keys_keep_their_session_token( + explicit(Some("t")), ambient_token, ("key", "secret", Some("t")) + )] + #[case::environment_keys_resolve_with_their_session_token( + AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token")) + )] + #[tokio::test] + async fn configured_credentials_read_the_environment_they_are_given( + #[case] config: AwsAuthConfig, + #[case] environment: fn(&str) -> Option, + #[case] expected: (&str, &str, Option<&str>), + ) { + let credentials = + AwsCredentialsProvider::configured(AwsAuthService::default(), config, environment) + .provide_credentials() + .await + .unwrap(); + assert_eq!( + ( + credentials.access_key_id(), + credentials.secret_access_key(), + credentials.session_token(), + ), + expected + ); + } +} diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index dff917d7c13..2d9ad9ba123 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -75,7 +75,8 @@ impl RequestSigner for SigV4Signer { ) .map_err(|error| match error { Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error), - Error::Credentials { .. } + Error::ExpiredCredentials + | Error::Credentials { .. } | Error::MissingWebIdentityCredentials | Error::WebIdentityExpiration(_) | Error::Signing { .. } => litellm_http::Error::signature(error), diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index f99972f813b..36c10a74c80 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -6,18 +6,21 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true +litellm-error.workspace = true moka.workspace = true serde_json.workspace = true sha2.workspace = true strum.workspace = true +tokio.workspace = true url.workspace = true azure_core = "1.0.0" azure_identity = { version = "1.0.0", features = ["tokio"] } [dev-dependencies] -litellm-error.workspace = true +litellm-secrets.workspace = true rstest.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs index 9ed505e4160..6319406ce9c 100644 --- a/litellm-rust/crates/auth-azure/src/lib.rs +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -3,5 +3,5 @@ mod native; mod resolve; mod types; -pub use resolve::{AzureAuthService, SECRET_NAMES}; +pub use resolve::{AzureAuthService, AzureIdentity, SECRET_NAMES}; pub use types::{AzureAuthInputs, ConfigValue}; diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 9fda67746d1..03021428ea7 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -13,7 +13,8 @@ use azure_identity::{ use sha2::{Digest, Sha256}; use litellm_auth_types::{Error, ErrorDetail}; -use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; +use litellm_auth_types::{InputSource, ResolvedCredential, Sourced}; +use litellm_secret_value::SecretValue; use super::credential_provider_cache::{ AzureCredentialProviderCache, AzureCredentialProviderCacheKey, @@ -58,6 +59,7 @@ pub(crate) enum NativeAzureRequest { pub(crate) struct ValidatedAzureRequest { request: NativeAzureRequest, credential_source: InputSource, + probe_imds: bool, } impl ValidatedAzureRequest { @@ -67,9 +69,18 @@ impl ValidatedAzureRequest { Ok(Self { request, credential_source, + probe_imds: false, }) } + pub(crate) fn with_imds_probe(self, enabled: bool) -> Self { + Self { + probe_imds: enabled + && matches!(self.request, NativeAzureRequest::ManagedIdentity { .. }), + ..self + } + } + pub(crate) fn credential_source(&self) -> InputSource { self.credential_source } @@ -86,9 +97,34 @@ impl ValidatedAzureRequest { } } +pub(crate) enum CredentialAttempt { + Resolved(ResolvedCredential), + Unavailable(Error), + Failed(Error), +} + +impl CredentialAttempt { + pub(crate) fn into_result(self) -> Result { + match self { + Self::Resolved(credential) => Ok(credential), + Self::Unavailable(error) | Self::Failed(error) => Err(error), + } + } +} + pub(crate) struct NativeAzureTokenAcquirer { cache: AzureCredentialProviderCache, transport: Option, + imds_available: tokio::sync::Mutex, +} + +enum ImdsAvailability { + Unknown, + Available, + Unavailable { + retry_at: tokio::time::Instant, + error: Error, + }, } impl Default for NativeAzureTokenAcquirer { @@ -102,6 +138,7 @@ impl NativeAzureTokenAcquirer { Self { cache: AzureCredentialProviderCache::new(cache_capacity), transport: None, + imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown), } } @@ -113,10 +150,51 @@ impl NativeAzureTokenAcquirer { Self { cache: AzureCredentialProviderCache::new(cache_capacity), transport: Some(transport), + imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown), } } - pub(crate) async fn acquire( + pub(crate) async fn acquire(&self, request: ValidatedAzureRequest) -> CredentialAttempt { + let discover_imds = request.probe_imds; + if discover_imds && let Err(error) = self.discover_imds().await { + return CredentialAttempt::Unavailable(error); + } + match self.acquire_token(request).await { + Ok(credential) => CredentialAttempt::Resolved(credential), + Err(error) => { + let missing_identity = discover_imds && unassigned_identity(&error); + if missing_identity { + CredentialAttempt::Unavailable(error) + } else { + CredentialAttempt::Failed(error) + } + } + } + } + + async fn discover_imds(&self) -> Result<(), Error> { + let mut availability = self.imds_available.lock().await; + match &*availability { + ImdsAvailability::Available => return Ok(()), + ImdsAvailability::Unavailable { retry_at, error } + if tokio::time::Instant::now() < *retry_at => + { + return Err(error.clone()); + } + ImdsAvailability::Unknown | ImdsAvailability::Unavailable { .. } => {} + } + let result = probe_imds(self.transport.clone().unwrap_or_default()).await; + *availability = match &result { + Ok(()) => ImdsAvailability::Available, + Err(error) => ImdsAvailability::Unavailable { + retry_at: tokio::time::Instant::now() + Duration::from_secs(30), + error: error.clone(), + }, + }; + result + } + + async fn acquire_token( &self, request: ValidatedAzureRequest, ) -> Result { @@ -147,6 +225,54 @@ impl NativeAzureTokenAcquirer { } } +fn unassigned_identity(error: &(dyn std::error::Error + 'static)) -> bool { + std::iter::successors(Some(error), |cause| cause.source()).any(|cause| { + cause + .downcast_ref::() + .or_else(|| { + cause + .downcast_ref::()? + .downcast_ref::() + }) + .is_some_and(|error| { + error.http_status() == Some(azure_core::http::StatusCode::BadRequest) + }) + }) +} + +async fn probe_imds(transport: azure_core::http::Transport) -> Result<(), Error> { + use azure_core::http::{Context, Method, Request, Url}; + + let endpoint = + Url::parse("http://169.254.169.254/metadata/identity/oauth2/token?api-version=2019-08-01") + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed("Azure IMDS probe endpoint", error)) + })?; + let mut request = Request::new(endpoint, Method::Get); + let response = tokio::time::timeout( + Duration::from_secs(1), + transport.send(&Context::new(), &mut request), + ) + .await + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error)) + })? + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error)) + })?; + // Headerless IMDS discovery expects 400, checked 2026-11-01 + // https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service#errors-and-debugging + if response.status() != azure_core::http::StatusCode::BadRequest { + return Err(Error::CredentialAcquisition( + ErrorDetail::UnexpectedStatus { + operation: "Azure IMDS availability probe", + status: u16::from(response.status()), + }, + )); + } + Ok(()) +} + impl NativeAzureRequest { fn scope(&self) -> &str { match self { @@ -492,13 +618,15 @@ impl ClientAssertion for StaticAssertion { #[cfg(test)] mod tests { use std::sync::{Arc, Mutex}; + use std::time::Duration; use azure_core::http::headers::Headers; use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport}; use azure_core::{Bytes, Result}; use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; - use litellm_auth_types::{InputSource, SecretValue, Sourced}; + use litellm_auth_types::{InputSource, Sourced}; + use litellm_secret_value::SecretValue; fn deployment(value: T) -> Sourced { Sourced::new(value, InputSource::Deployment) @@ -567,6 +695,296 @@ mod tests { } } + #[derive(Debug)] + struct AvailableImdsClient { + probe_status: StatusCode, + probe_delay: Duration, + recover_after: Option, + response_status: Option, + probes: std::sync::atomic::AtomicUsize, + } + + impl Default for AvailableImdsClient { + fn default() -> Self { + Self { + probe_status: StatusCode::BadRequest, + probe_delay: Duration::ZERO, + recover_after: None, + response_status: None, + probes: std::sync::atomic::AtomicUsize::new(0), + } + } + } + + impl HttpClient for AvailableImdsClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + request: &'life1 Request, + ) -> std::pin::Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if request + .headers() + .get_optional_str(&azure_core::http::headers::HeaderName::from_static( + "metadata", + )) + .is_none() + { + let probe = self + .probes + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if self.recover_after.is_some_and(|failures| probe >= failures) { + return Ok(AsyncRawResponse::from_bytes( + StatusCode::BadRequest, + Headers::new(), + "metadata header required", + )); + } + tokio::time::sleep(self.probe_delay).await; + return Ok(AsyncRawResponse::from_bytes( + self.probe_status, + Headers::new(), + "metadata header required", + )); + } + if let Some(status) = self.response_status { + return Ok(AsyncRawResponse::from_bytes( + status, + Headers::new(), + "identity endpoint failure", + )); + } + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + let expires = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 3600; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + Headers::new(), + serde_json::json!({ + "access_token": "managed-token", "expires_on": expires.to_string(), + "token_type": "Bearer", "resource": "https://service.test", + }) + .to_string(), + )) + }) + } + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn available_imds_retains_full_sdk_token_acquisition() { + let acquirer = NativeAzureTokenAcquirer::with_transport( + 4, + Transport::new(Arc::new(AvailableImdsClient { + probe_status: StatusCode::BadRequest, + response_status: None, + ..Default::default() + })), + ); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + let credential = acquirer.acquire(request).await.into_result().unwrap(); + assert_eq!(credential.secret().expose(), "managed-token"); + } + + #[rstest::rstest] + #[case::another_metadata_service(StatusCode::NotFound)] + #[case::token_required(StatusCode::Unauthorized)] + #[case::proxy_rejection(StatusCode::Forbidden)] + #[case::unavailable(StatusCode::ServiceUnavailable)] + #[tokio::test(start_paused = true)] + async fn unrelated_metadata_responses_do_not_start_managed_identity_acquisition( + #[case] status: StatusCode, + ) { + let client = Arc::new(AvailableImdsClient { + probe_status: status, + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + assert!(matches!( + acquirer.acquire(request.clone()).await, + super::CredentialAttempt::Unavailable(_) + )); + assert!(matches!( + acquirer.acquire(request).await, + super::CredentialAttempt::Unavailable(_) + )); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn an_imds_timeout_is_not_repeated_on_the_next_request() { + let client = Arc::new(AvailableImdsClient { + probe_delay: Duration::from_secs(10), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + assert!(matches!( + acquirer.acquire(request.clone()).await, + super::CredentialAttempt::Unavailable(_) + )); + let second = tokio::time::timeout(Duration::from_millis(10), acquirer.acquire(request)) + .await + .unwrap(); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[rstest::fixture] + fn imds_request() -> ValidatedAzureRequest { + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true) + } + + #[rstest::rstest] + #[case::timeout(StatusCode::BadRequest, Duration::from_secs(10))] + #[case::service_unavailable(StatusCode::ServiceUnavailable, Duration::ZERO)] + #[tokio::test(start_paused = true)] + async fn failed_imds_discovery_recovers_without_restarting_the_acquirer( + imds_request: ValidatedAzureRequest, + #[case] status: StatusCode, + #[case] delay: Duration, + ) { + let client = Arc::new(AvailableImdsClient { + probe_status: status, + probe_delay: delay, + recover_after: Some(1), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request.clone()), + ); + assert!(matches!(first, super::CredentialAttempt::Unavailable(_))); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + + tokio::time::advance(Duration::from_secs(30)).await; + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request.clone()), + ); + assert_eq!( + first.into_result().unwrap().secret().expose(), + "managed-token" + ); + assert_eq!( + second.into_result().unwrap().secret().expose(), + "managed-token" + ); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + + tokio::time::advance(Duration::from_secs(60)).await; + assert_eq!( + acquirer + .acquire(imds_request) + .await + .into_result() + .unwrap() + .secret() + .expose(), + "managed-token", + ); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn a_persistent_imds_outage_reuses_each_failed_probe( + imds_request: ValidatedAzureRequest, + ) { + let client = Arc::new(AvailableImdsClient { + probe_delay: Duration::from_secs(10), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + assert!(matches!( + acquirer.acquire(imds_request.clone()).await, + super::CredentialAttempt::Unavailable(_), + )); + tokio::time::advance(Duration::from_secs(30)).await; + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request), + ); + assert!(matches!(first, super::CredentialAttempt::Unavailable(_))); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[case::default_identity_not_assigned(true, StatusCode::BadRequest, true)] + #[case::explicit_identity_not_assigned(false, StatusCode::BadRequest, false)] + #[case::default_identity_forbidden(true, StatusCode::Forbidden, false)] + #[case::default_identity_service_failure(true, StatusCode::InternalServerError, false)] + #[tokio::test(start_paused = true)] + async fn only_unavailable_default_identities_allow_fallback( + #[case] discovery: bool, + #[case] status: StatusCode, + #[case] unavailable: bool, + ) { + let acquirer = NativeAzureTokenAcquirer::with_transport( + 4, + Transport::new(Arc::new(AvailableImdsClient { + probe_status: StatusCode::BadRequest, + response_status: Some(status), + ..Default::default() + })), + ); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(discovery); + let outcome = acquirer.acquire(request).await; + let is_unavailable = matches!(outcome, super::CredentialAttempt::Unavailable(_)); + let error = outcome.into_result().unwrap_err(); + assert_eq!( + is_unavailable, + unavailable, + "{}", + litellm_error::Report::new(&error) + ); + } + + #[rstest::rstest] #[tokio::test] async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() { let transport = Arc::new(RecordingTokenClient::default()); @@ -580,8 +998,12 @@ mod tests { "https://login.test", ); - let first = acquirer.acquire(request.clone()).await.unwrap(); - let second = acquirer.acquire(request).await.unwrap(); + let first = acquirer + .acquire(request.clone()) + .await + .into_result() + .unwrap(); + let second = acquirer.acquire(request).await.into_result().unwrap(); assert_eq!(first.secret().expose(), "native-token"); assert_eq!(second.secret().expose(), "native-token"); @@ -597,6 +1019,7 @@ mod tests { ); } + #[rstest::rstest] #[tokio::test] async fn credential_provider_cache_isolates_every_client_secret_identity_field() { let transport = Arc::new(RecordingTokenClient::default()); @@ -643,10 +1066,10 @@ mod tests { ), ]; - acquirer.acquire(base.clone()).await.unwrap(); - acquirer.acquire(base).await.unwrap(); + acquirer.acquire(base.clone()).await.into_result().unwrap(); + acquirer.acquire(base).await.into_result().unwrap(); for request in variants.into_iter().skip(1) { - acquirer.acquire(request).await.unwrap(); + acquirer.acquire(request).await.into_result().unwrap(); } assert_eq!(transport.requests.lock().unwrap().len(), 6); diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 25b71741342..9ff4eda540a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,7 +1,10 @@ +use crate::native::CredentialAttempt; use litellm_auth_types::{ - CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, - SecretValue, Sourced, TokenProviderHandle, + CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced, + TokenProviderHandle, }; +use litellm_secret_value::SecretValue; + use litellm_auth_types::{Error, ErrorDetail}; use std::future::Future; use std::pin::Pin; @@ -10,25 +13,31 @@ use std::sync::Arc; use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE}; -const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; -const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID"; -const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID"; -const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET"; -const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE"; -const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST"; -const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL"; -const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; +#[derive(strum::IntoStaticStr, strum::VariantArray, strum::VariantNames)] +enum AzureSecretName { + #[strum(serialize = "AZURE_AD_TOKEN")] + AdToken, + #[strum(serialize = "AZURE_TENANT_ID")] + TenantId, + #[strum(serialize = "AZURE_CLIENT_ID")] + ClientId, + #[strum(serialize = "AZURE_CLIENT_SECRET")] + ClientSecret, + #[strum(serialize = "AZURE_SCOPE")] + Scope, + #[strum(serialize = "AZURE_AUTHORITY_HOST")] + AuthorityHost, + #[strum(serialize = "AZURE_CREDENTIAL")] + Credential, + #[strum(serialize = "AZURE_FEDERATED_TOKEN_FILE")] + FederatedTokenFile, + #[strum(serialize = "IDENTITY_ENDPOINT")] + IdentityEndpoint, + #[strum(serialize = "MSI_ENDPOINT")] + MsiEndpoint, +} -pub const SECRET_NAMES: &[&str] = &[ - AZURE_AD_TOKEN_ENV, - AZURE_TENANT_ID_ENV, - AZURE_CLIENT_ID_ENV, - AZURE_CLIENT_SECRET_ENV, - AZURE_SCOPE_ENV, - AZURE_AUTHORITY_HOST_ENV, - AZURE_CREDENTIAL_ENV, - AZURE_FEDERATED_TOKEN_FILE_ENV, -]; +pub const SECRET_NAMES: &[&str] = ::VARIANTS; #[derive(Clone, Debug)] pub(crate) enum AzureCredentialPlan { @@ -55,14 +64,14 @@ trait AzureTokenAcquirer: Send + Sync { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>>; + ) -> Pin + Send + '_>>; } impl AzureTokenAcquirer for NativeAzureTokenAcquirer { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>> { + ) -> Pin + Send + '_>> { Box::pin(NativeAzureTokenAcquirer::acquire(self, request)) } } @@ -81,12 +90,36 @@ impl AzureAuthService { Self { native } } + pub async fn identity_token( + &self, + identity: &AzureIdentity, + scope: &str, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, Error> { + let plan = identity_plan(identity, scope, env_lookup)?; + let inputs = AzureAuthInputs::default_credential_for_scope(scope); + Ok(self + .run_plan(plan, &inputs, env_lookup) + .await? + .map(Sourced::into_value)) + } + pub async fn get_azure_ad_token( &self, inputs: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result>, Error> { - match select_auth_plan(inputs, env_lookup)? { + self.run_plan(select_auth_plan(inputs, env_lookup)?, inputs, env_lookup) + .await + } + + async fn run_plan( + &self, + plan: AzureCredentialPlan, + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result>, Error> { + match plan { AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)), AzureCredentialPlan::Caller(caller) => { let credential = caller.acquire().await?; @@ -121,6 +154,7 @@ impl AzureAuthService { self.native .acquire(request) .await + .into_result() .map(|credential| Sourced::new(credential, source)) .map(Some) } @@ -129,6 +163,7 @@ impl AzureAuthService { self.native .acquire(request) .await + .into_result() .map(|credential| Some(Sourced::new(credential, source))) } AzureCredentialPlan::Chain(requests) => { @@ -136,8 +171,16 @@ impl AzureAuthService { for request in requests { let source = request.credential_source(); match self.native.acquire(request).await { - Ok(credential) => return Ok(Some(Sourced::new(credential, source))), - Err(error) => failures.push(error), + CredentialAttempt::Resolved(credential) => { + return Ok(Some(Sourced::new(credential, source))); + } + CredentialAttempt::Unavailable(error) => failures.push(error), + CredentialAttempt::Failed(error) => { + return Err(Error::CredentialChain { + last: Box::new(error), + earlier: failures, + }); + } } } match failures.pop() { @@ -153,39 +196,161 @@ impl AzureAuthService { } } +/// The identity azure-identity builds for one explicit credential class, without LiteLLM's provider settings +/// (`AZURE_AD_TOKEN`, `AZURE_CREDENTIAL`, `AZURE_SCOPE`) that `get_azure_ad_token` also reads. +#[derive(Clone, Debug)] +pub enum AzureIdentity { + ClientSecret { + tenant_id: String, + client_id: String, + client_secret: SecretValue, + }, + ManagedIdentity { + client_id: String, + }, + Default, +} + +fn identity_plan( + identity: &AzureIdentity, + scope: &str, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let deployment = |value: &str| Sourced::new(value.to_owned(), InputSource::Deployment); + let environment = |name: &str| { + env_lookup(name) + .filter(|value| !value.trim().is_empty()) + .map(|value| Sourced::new(value, InputSource::Environment)) + }; + let scope = deployment(scope); + let authority = environment(AzureSecretName::AuthorityHost.into()); + match identity { + AzureIdentity::ClientSecret { + tenant_id, + client_id, + client_secret, + } => Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( + NativeAzureRequest::ClientSecret { + tenant_id: deployment(tenant_id), + client_id: deployment(client_id), + client_secret: Sourced::new(client_secret.clone(), InputSource::Deployment), + scope, + authority, + }, + )?)), + AzureIdentity::ManagedIdentity { client_id } => Ok(AzureCredentialPlan::Native( + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: Some(deployment(client_id)), + scope, + selection_source: InputSource::Deployment, + })?, + )), + AzureIdentity::Default => { + let tenant_id = environment(AzureSecretName::TenantId.into()); + let client_id = environment(AzureSecretName::ClientId.into()); + let from_environment = match ( + tenant_id.clone(), + client_id.clone(), + environment(AzureSecretName::ClientSecret.into()), + ) { + (Some(tenant_id), Some(client_id), Some(client_secret)) => { + Some(NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret: client_secret.map(SecretValue::new), + scope: scope.clone(), + authority: authority.clone(), + }) + } + _ => None, + }; + let probe_imds = env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none() + && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none(); + let AzureCredentialPlan::Chain(chain) = select_native_plan( + Some(Sourced::new( + AzureCredentialType::DefaultAzureCredential, + InputSource::Deployment, + )), + tenant_id, + client_id, + environment(AzureSecretName::FederatedTokenFile.into()), + scope, + authority, + InputSource::Environment, + )? + else { + unreachable!("DefaultAzureCredential always plans a chain") + }; + let environment_request = from_environment + .map(ValidatedAzureRequest::new) + .transpose()?; + Ok(AzureCredentialPlan::Chain( + environment_request + .into_iter() + .chain(chain) + .map(|request| request.with_imds_probe(probe_imds)) + .collect(), + )) + } + } +} + pub(crate) fn select_auth_plan( inputs: &AzureAuthInputs, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup); - let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup); - let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup); - let client_secret = - configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup); - let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup) - .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment)); - let authority = configured_string( - &inputs.azure_authority_host, - AZURE_AUTHORITY_HOST_ENV, + let token = configured_secret( + &inputs.azure_ad_token, + AzureSecretName::AdToken.into(), env_lookup, ); - let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup) - .map(|value| { - value - .value() - .parse::() - .map(|selector| Sourced::new(selector, value.source())) - }) - .transpose() - .map_err(|error| { - Error::InvalidConfiguration(ErrorDetail::failed( - "Azure credential selector parsing", - error, - )) - })?; + let tenant_id = configured_string( + &inputs.tenant_id, + AzureSecretName::TenantId.into(), + env_lookup, + ); + let client_id = configured_string( + &inputs.client_id, + AzureSecretName::ClientId.into(), + env_lookup, + ); + let client_secret = configured_secret( + &inputs.client_secret, + AzureSecretName::ClientSecret.into(), + env_lookup, + ); + let scope = configured_string( + &inputs.azure_scope, + AzureSecretName::Scope.into(), + env_lookup, + ) + .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment)); + let authority = configured_string( + &inputs.azure_authority_host, + AzureSecretName::AuthorityHost.into(), + env_lookup, + ); + let selector = configured_string( + &inputs.azure_credential, + AzureSecretName::Credential.into(), + env_lookup, + ) + .map(|value| { + value + .value() + .parse::() + .map(|selector| Sourced::new(selector, value.source())) + }) + .transpose() + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed( + "Azure credential selector parsing", + error, + )) + })?; let federated_token_file = configured_string( &inputs.federated_token_file, - AZURE_FEDERATED_TOKEN_FILE_ENV, + AzureSecretName::FederatedTokenFile.into(), env_lookup, ); @@ -235,7 +400,12 @@ pub(crate) fn select_auth_plan( return Ok(AzureCredentialPlan::Missing); } - select_native_plan( + let probe_imds = selector + .as_ref() + .is_none_or(|selector| *selector.value() == AzureCredentialType::DefaultAzureCredential) + && env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none() + && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none(); + match select_native_plan( selector, tenant_id, client_id, @@ -243,7 +413,15 @@ pub(crate) fn select_auth_plan( scope, authority, inputs.enable_azure_ad_token_refresh.source(), - ) + )? { + AzureCredentialPlan::Chain(requests) => Ok(AzureCredentialPlan::Chain( + requests + .into_iter() + .map(|request| request.with_imds_probe(probe_imds)) + .collect(), + )), + plan => Ok(plan), + } } fn select_native_plan( @@ -500,13 +678,17 @@ mod tests { use litellm_auth_types::ResolvedCredential; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, - CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, + CredentialResolver, CredentialResolverHandle, InputSource, Sourced, }; + use litellm_secret_value::SecretValue; #[derive(Debug)] struct FileResolver; + use crate::native::CredentialAttempt; + struct ChainAcquirer { + terminal_on: Option<&'static str>, requests: Mutex>, succeed_on: Option<&'static str>, } @@ -515,20 +697,25 @@ mod tests { fn acquire( &self, request: ValidatedAzureRequest, - ) -> std::pin::Pin> + Send + '_>> - { + ) -> std::pin::Pin + Send + '_>> { let kind = request.kind(); self.requests.lock().unwrap().push(kind); Box::pin(async move { if self.succeed_on == Some(kind) { - Ok(ResolvedCredential::AccessToken { + CredentialAttempt::Resolved(ResolvedCredential::AccessToken { token: SecretValue::new("chain-token"), expires_on: None, }) } else { - Err(Error::CredentialAcquisition( - litellm_auth_types::ErrorDetail::Unresolved { subject: kind }, - )) + let error = + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::Unresolved { + subject: kind, + }); + if self.terminal_on == Some(kind) { + CredentialAttempt::Failed(error) + } else { + CredentialAttempt::Unavailable(error) + } } }) } @@ -687,9 +874,210 @@ mod tests { assert_eq!(resolved, Some(SecretValue::new("rotated-assertion"))); } + #[derive(Debug)] + struct UnavailableImdsClient; + + impl azure_core::http::HttpClient for UnavailableImdsClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + _request: &'life1 azure_core::http::Request, + ) -> std::pin::Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(std::future::pending()) + } + } + + struct OfflineChainAcquirer(super::NativeAzureTokenAcquirer); + + impl AzureTokenAcquirer for OfflineChainAcquirer { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> std::pin::Pin + Send + '_>> { + Box::pin(async move { + if request.kind() == "developer-tools" { + return CredentialAttempt::Resolved(ResolvedCredential::AccessToken { + token: SecretValue::new("cli-token"), + expires_on: None, + }); + } + self.0.acquire(request).await + }) + } + } + + #[rstest::rstest] + #[case::default_off_azure("DefaultAzureCredential", None, true)] + #[case::explicit_managed_identity("ManagedIdentityCredential", None, false)] + #[case::app_service("DefaultAzureCredential", Some("http://app-service.test"), false)] + #[tokio::test(start_paused = true)] + async fn only_default_imds_discovery_is_bounded( + #[case] selector: &str, + #[case] identity_endpoint: Option<&str>, + #[case] expect_cli: bool, + ) { + let native = super::NativeAzureTokenAcquirer::with_transport( + 4, + azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)), + ); + let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native))); + let inputs = AzureAuthInputs { + azure_credential: crate::types::ConfigValue::Value(Sourced::new( + selector.into(), + InputSource::Deployment, + )), + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + let environment = |name: &str| { + (name == "IDENTITY_ENDPOINT") + .then_some(identity_endpoint) + .flatten() + .map(str::to_string) + }; + let outcome = tokio::time::timeout( + std::time::Duration::from_secs(2), + service.get_azure_ad_token(&inputs, &environment), + ) + .await; + + assert_eq!(outcome.is_ok(), expect_cli); + if expect_cli { + assert_eq!( + outcome.unwrap().unwrap().unwrap().value().secret().expose(), + "cli-token" + ); + } + } + + #[rstest::rstest] + #[case::app_service(Some("IDENTITY_ENDPOINT"), false)] + #[case::legacy_app_service(Some("MSI_ENDPOINT"), false)] + #[case::off_azure(None, true)] + #[tokio::test(start_paused = true)] + async fn default_chain_uses_endpoints_from_the_secret_snapshot( + #[case] endpoint: Option<&str>, + #[case] expect_cli: bool, + #[values(false, true)] consumer_identity: bool, + ) { + use litellm_secrets::source::SecretSource; + + struct EndpointSecrets<'a>(Option<&'a str>); + + impl SecretSource for EndpointSecrets<'_> { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> std::pin::Pin< + Box< + dyn Future, litellm_secrets::Error>> + + Send + + 'a, + >, + > { + Box::pin(async move { + Ok((self.0 == Some(name)).then(|| SecretValue::new("http://app-service.test"))) + }) + } + } + + let snapshot = EndpointSecrets(endpoint) + .resolve(SECRET_NAMES) + .await + .unwrap(); + let native = super::NativeAzureTokenAcquirer::with_transport( + 4, + azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)), + ); + let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native))); + let lookup = |name: &str| snapshot.get(name); + let scope = "https://service.test/.default"; + let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), async { + if consumer_identity { + service + .identity_token(&super::AzureIdentity::Default, scope, &lookup) + .await + } else { + service + .get_azure_ad_token( + &AzureAuthInputs::default_credential_for_scope(scope), + &lookup, + ) + .await + .map(|token| token.map(Sourced::into_value)) + } + }) + .await; + + assert_eq!(outcome.is_ok(), expect_cli); + if expect_cli { + assert_eq!( + outcome.unwrap().unwrap().unwrap().secret().expose(), + "cli-token" + ); + } + } + + #[rstest::rstest] + #[case::litellm_provider_settings_are_ignored( + &[("AZURE_AD_TOKEN", "provider-token"), ("AZURE_CREDENTIAL", "ClientSecretCredential")], + &["managed-identity", "developer-tools"] + )] + #[case::environment_client_secret_comes_first( + &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_CLIENT_SECRET", "secret")], + &["client-secret", "managed-identity", "developer-tools"] + )] + #[case::aks_workload_identity( + &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_FEDERATED_TOKEN_FILE", "/var/run/token")], + &["workload-identity", "managed-identity", "developer-tools"] + )] + #[tokio::test] + async fn the_default_identity_follows_azure_identity_defaults( + #[case] environment: &[(&str, &str)], + #[case] expected: &[&str], + ) { + let acquirer = Arc::new(ChainAcquirer { + terminal_on: None, + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + + let credential = service + .identity_token( + &super::AzureIdentity::Default, + "https://redis.azure.com/.default", + &lookup, + ) + .await + .unwrap() + .unwrap(); + + assert_eq!(credential.secret().expose(), "chain-token"); + assert_eq!(*acquirer.requests.lock().unwrap(), expected); + } + + #[rstest::rstest] #[tokio::test] async fn default_chain_uses_declared_order_and_stops_after_success() { let acquirer = Arc::new(ChainAcquirer { + terminal_on: None, requests: Mutex::new(Vec::new()), succeed_on: Some("developer-tools"), }); @@ -712,9 +1100,13 @@ mod tests { ); } + #[rstest::rstest] + #[case::all_unavailable(false)] + #[case::terminal_failure_after_unavailable(true)] #[tokio::test] - async fn chain_reports_each_acquisition_failure() { + async fn chain_reports_each_acquisition_failure(#[case] terminal: bool) { let acquirer = Arc::new(ChainAcquirer { + terminal_on: terminal.then_some("developer-tools"), requests: Mutex::new(Vec::new()), succeed_on: None, }); @@ -741,6 +1133,27 @@ mod tests { ); } + #[rstest::rstest] + #[tokio::test] + async fn a_configured_identity_failure_does_not_fall_back_to_a_developer_identity() { + let acquirer = Arc::new(ChainAcquirer { + terminal_on: Some("managed-identity"), + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let inputs = AzureAuthInputs { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + let error = service + .get_azure_ad_token(&inputs, &|_| None) + .await + .unwrap_err(); + assert_eq!(*acquirer.requests.lock().unwrap(), ["managed-identity"]); + assert!(matches!(error, Error::CredentialChain { earlier, .. } if earlier.is_empty())); + } + #[derive(Debug)] struct CallerToken(&'static str); diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index a042937a047..6f8290645cf 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -1,8 +1,10 @@ use std::collections::BTreeMap; use litellm_auth_types::{ - CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle, + CredentialResolverHandle, Error, InputSource, Sourced, TokenProviderHandle, }; +use litellm_secret_value::SecretValue; + use serde_json::{Map, Value}; use strum::EnumString; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index 7844b35ecc2..90aab8d281c 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -6,9 +6,10 @@ license.workspace = true repository.workspace = true [features] -google-sdk = ["dep:google-cloud-auth", "dep:http"] +google-sdk = ["dep:http"] [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true base64.workspace = true @@ -20,10 +21,14 @@ tokio.workspace = true veil.workspace = true gcp_auth = "0.12.7" -google-cloud-auth = { workspace = true, optional = true } +google-cloud-auth.workspace = true http = { workspace = true, optional = true } [dev-dependencies] +time = { workspace = true, features = ["formatting"] } +axum = { workspace = true, features = ["json"] } litellm-error.workspace = true rstest.workspace = true +rustls.workspace = true tempfile.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index 8ac1e0f54ba..5976ef82344 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -1,20 +1,11 @@ -use std::{future::Future, pin::Pin, sync::Arc}; +use std::sync::Arc; -use gcp_auth::{CustomServiceAccount, TokenProvider}; -use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential}; -use moka::future::Cache; -use serde_json::Value; -use sha2::{Digest, Sha256}; +use litellm_auth_types::{CredentialPlacement, Error, http::apply_credential}; use crate::{ - config::{ - CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, - is_base64_encoded_json, names_a_credentials_file, non_empty_env, - }, - constants::{ - CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV, - VERTEXAI_API_KEY_ENV, - }, + config::{VertexConfig, credential_source, get_vertex_ai_project, non_empty_env}, + constants::{VERTEX_AI_API_KEY_ENV, VERTEXAI_API_KEY_ENV}, + identity::{GcpAuth, GcpProviderLoader, GcpTokenSource}, }; pub struct VertexEnvironment { @@ -22,28 +13,15 @@ pub struct VertexEnvironment { pub project_id: String, } -struct VertexAccessToken { - token: String, - project_id: String, -} - -#[derive(Clone)] +#[derive(Clone, Default)] pub struct VertexAuth { - providers: Cache>, - loader: Arc, -} - -impl Default for VertexAuth { - fn default() -> Self { - Self::new(Arc::new(GcpProviderLoader)) - } + identity: GcpAuth, } impl VertexAuth { - pub fn new(loader: Arc) -> Self { + pub fn new(loader: Arc) -> Self { Self { - providers: Cache::builder().max_capacity(64).build(), - loader, + identity: GcpAuth::new(loader), } } @@ -52,7 +30,10 @@ impl VertexAuth { config: &VertexConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { - self.load_provider(config, env_lookup).await?.token().await + self.identity + .access_token(credential_source(config, env_lookup)) + .await + .map(|token| token.expose().to_owned()) } pub async fn validate_environment( @@ -77,10 +58,18 @@ impl VertexAuth { (true, _) => None, (false, Some(token)) => Some(token), (false, None) => { - let access = self.get_access_token(config, env_lookup).await?; + let provider = self.load_provider(config, env_lookup).await?; + let (token, project_id) = match project_id { + Some(project_id) => (provider.token().await?, project_id), + None => tokio::try_join!(provider.token(), provider.project_id())?, + }; return Ok(VertexEnvironment { - headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, - project_id: project_id.unwrap_or(access.project_id), + headers: apply_credential( + headers, + token.expose(), + CredentialPlacement::Bearer, + )?, + project_id, }); } }; @@ -104,248 +93,171 @@ impl VertexAuth { }) } - async fn get_access_token( - &self, - config: &VertexConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result { - let provider = self.load_provider(config, env_lookup).await?; - let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?; - Ok(VertexAccessToken { token, project_id }) - } - async fn load_provider( &self, config: &VertexConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result, Error> { - let source = credential_source(config, env_lookup); - let key = source.cache_key(); - self.providers - .try_get_with(key, self.loader.load(source)) + ) -> Result, Error> { + self.identity + .load_provider(credential_source(config, env_lookup)) .await - .map_err(|error| (*error).clone()) } } -pub trait VertexTokenSource: Send + Sync { - fn project_id(&self) -> VertexAuthFuture<'_, String>; - fn token(&self) -> VertexAuthFuture<'_, String>; -} - -pub trait VertexProviderLoader: Send + Sync { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; -} - -pub type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; - -struct GcpTokenSource(Arc); - -impl VertexTokenSource for GcpTokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .project_id() - .await - .map(|project| project.to_string()) - .map_err(auth_acquisition_error) - }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .token(&[CLOUD_PLATFORM_SCOPE]) - .await - .map(|token| token.as_str().to_string()) - .map_err(auth_acquisition_error) - }) - } -} - -struct GcpProviderLoader; - -impl VertexProviderLoader for GcpProviderLoader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - Box::pin(async move { - let provider: Arc = match source { - CredentialSource::Inline(configured) => Arc::new( - CustomServiceAccount::from_json(validate_request_credentials( - configured.expose(), - )?) - .map_err(auth_acquisition_error)?, - ), - CredentialSource::Trusted(configured) => { - Arc::new(load_trusted_service_account(configured.expose())?) - } - CredentialSource::ApplicationCredentials(path) => { - Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?) - } - CredentialSource::Adc => { - gcp_auth::provider().await.map_err(auth_acquisition_error)? - } - }; - Ok(Arc::new(GcpTokenSource(provider)) as Arc) - }) - } -} - -fn validate_request_credentials(configured: &str) -> Result<&str, Error> { - let token_uri = serde_json::from_str::(configured) - .ok() - .and_then(|credentials| { - credentials - .get("token_uri") - .and_then(Value::as_str) - .map(str::to_string) - }); - if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { - return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { - field: "request-controlled vertex_credentials token_uri".into(), - expected: "the canonical Google OAuth token endpoint", - })); - } - Ok(configured) -} - -impl CredentialSource { - fn cache_key(&self) -> CredentialCacheKey { - match self { - Self::Inline(configured) => { - CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) - } - Self::Trusted(configured) => { - CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) - } - Self::ApplicationCredentials(path) => { - CredentialCacheKey::ApplicationCredentials(path.clone()) - } - Self::Adc => CredentialCacheKey::Adc, - } - } -} - -#[derive(Clone, Debug, Hash, PartialEq, Eq)] -enum CredentialCacheKey { - Inline([u8; 32]), - Trusted([u8; 32]), - ApplicationCredentials(String), - Adc, -} - -fn load_trusted_service_account(configured: &str) -> Result { - if is_base64_encoded_json(configured) { - return Err(Error::InvalidConfiguration( - ErrorDetail::Base64EncodedJson { - field: "vertex_credentials", - }, - )); - } - if names_a_credentials_file(configured) { - CustomServiceAccount::from_file(configured) - } else { - CustomServiceAccount::from_json(configured) - } - .map_err(auth_acquisition_error) -} - -fn auth_acquisition_error(error: gcp_auth::Error) -> Error { - Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( - "Vertex AI credentials", - error, - )) -} - #[cfg(test)] mod tests { - use litellm_auth_types::SecretValue; + use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + }; - use super::*; + use crate::{ + CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth, + VertexConfig, get_vertex_ai_location, secret_names, + }; + use litellm_secret_value::SecretValue; + use rstest::rstest; + use serde_json::{Value, json}; - #[rstest::rstest] - #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] - #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] - #[case::missing_endpoint("{}", false)] - fn request_credentials_require_canonical_token_endpoint( - #[case] credentials: &str, - #[case] accepted: bool, - ) { - assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + struct FakeProvider { + project: Option<&'static str>, + calls: Arc, } - #[rstest::rstest] - #[case::missing_file( - "/nonexistent/vertex-key.json", - "failed to open application credentials file" - )] - #[case::malformed_json( - " {\"type\": \"service_account\",", - "failed to deserialize ApplicationCredentials" - )] - #[case::quoted_json( - "'{\"type\": \"service_account\"}'", - "failed to deserialize ApplicationCredentials" - )] - #[tokio::test] - async fn trusted_credentials_are_read_as_a_file_unless_they_are_json( - #[case] configured: &str, - #[case] expected: &str, - ) { - let Err(error) = GcpProviderLoader - .load(CredentialSource::Trusted(SecretValue::new(configured))) - .await - else { - panic!("{configured:?} should not load"); - }; - let report = litellm_error::Report::new(&error).to_string(); - assert!(report.contains(expected), "{report}"); - } + impl GcpTokenSource for FakeProvider { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { + self.project.map(str::to_owned).ok_or( + litellm_auth_types::Error::CredentialAcquisition( + litellm_auth_types::ErrorDetail::Unresolved { + subject: "credential project", + }, + ), + ) + }) + } - #[rstest::rstest] - #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")] - #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")] - #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")] - #[tokio::test] - async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) { - assert!( - trusted_load_error(configured) - .await - .contains("is base64-encoded JSON, not the JSON itself") - ); - } - - async fn trusted_load_error(configured: &str) -> String { - match GcpProviderLoader - .load(CredentialSource::Trusted(SecretValue::new(configured))) - .await - { - Ok(_) => panic!("{configured:?} should not load"), - Err(error) => litellm_error::Report::new(&error).to_string(), + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok(SecretValue::new("adc-token")) }) } } - #[rstest::rstest] - #[tokio::test] - async fn a_relative_key_file_named_like_json_is_read_as_a_file() { - let key = tempfile::Builder::new() - .prefix("{vertex") - .suffix(".json") - .tempfile_in(".") - .unwrap(); - std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); - let relative_name = key.path().file_name().unwrap().to_str().unwrap(); - let absolute_path = std::fs::canonicalize(key.path()).unwrap(); - assert_eq!( - trusted_load_error(relative_name).await, - trusted_load_error(absolute_path.to_str().unwrap()).await - ); + struct FakeLoader { + loads: Arc, + provider: Arc, } - #[rstest::rstest] - fn inline_and_trusted_credentials_never_share_a_cache_entry() { - assert_ne!( - CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), - CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() - ); + impl GcpProviderLoader for FakeLoader { + fn load(&self, _source: CredentialSource) -> GcpAuthFuture<'_, Arc> { + let loads = self.loads.clone(); + let provider = self.provider.clone(); + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok(provider) + }) + } + } + + fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() + } + + fn auth(calls: Arc, loads: Arc) -> VertexAuth { + let provider: Arc = Arc::new(FakeProvider { + calls, + project: Some("adc-project"), + }); + VertexAuth::new(Arc::new(FakeLoader { loads, provider })) + } + + #[rstest] + #[tokio::test] + async fn secret_names_cover_environment_reads() { + let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) + .await + .unwrap(); + get_vertex_ai_location(&VertexConfig::default(), &env); + let allowed = secret_names() + .iter() + .map(|name| (*name).to_owned()) + .collect::>(); + assert!(seen.lock().unwrap().is_subset(&allowed)); + } + + #[rstest] + #[tokio::test] + async fn explicit_token_and_header_do_not_acquire_adc() { + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); + let configured = config(json!({"vertex_project":"project-1"})); + let explicit = auth + .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) + .await + .unwrap(); + assert_eq!(explicit.headers[0].1, "Bearer access-token"); + let existing = auth + .validate_environment( + vec![("authorization".into(), "Bearer existing".into())], + None, + &configured, + &|_| None, + ) + .await + .unwrap(); + assert_eq!(existing.headers[0].1, "Bearer existing"); + assert_eq!(loads.load(Ordering::SeqCst), 0); + } + + #[rstest] + #[tokio::test] + async fn provider_is_reused_across_authentication_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(calls.clone(), loads.clone()); + let config = VertexConfig::default(); + let environment = |_: &str| None; + let (first, second) = tokio::try_join!( + auth.validate_environment(Vec::new(), None, &config, &environment), + auth.validate_environment(Vec::new(), None, &config, &environment), + ) + .unwrap(); + assert_eq!(first.project_id, "adc-project"); + assert_eq!(first.headers[0].1, "Bearer adc-token"); + assert_eq!(first.project_id, second.project_id); + assert_eq!(first.headers, second.headers); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 4); + } + + #[rstest] + #[tokio::test] + async fn an_explicit_project_does_not_require_project_metadata_in_the_credentials() { + let auth = VertexAuth::new(Arc::new(FakeLoader { + loads: Arc::new(AtomicUsize::new(0)), + provider: Arc::new(FakeProvider { + calls: Arc::new(AtomicUsize::new(0)), + project: None, + }), + })); + let configured = VertexConfig::new(None, Some("configured-project".into()), None); + let environment = auth + .validate_environment(Vec::new(), None, &configured, &|_| None) + .await + .unwrap(); + assert_eq!(environment.project_id, "configured-project"); + assert_eq!(environment.headers[0].1, "Bearer adc-token"); } } diff --git a/litellm-rust/crates/auth-gcp/src/config.rs b/litellm-rust/crates/auth-gcp/src/config.rs index 5444897aaaf..048e40b0a2f 100644 --- a/litellm-rust/crates/auth-gcp/src/config.rs +++ b/litellm-rust/crates/auth-gcp/src/config.rs @@ -4,7 +4,9 @@ use base64::{ Engine, alphabet, engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig}, }; -use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced}; +use litellm_auth_types::{Error, ErrorDetail, InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use serde_json::{Map, Value}; use litellm_auth_types::VertexParams; @@ -144,6 +146,20 @@ pub enum CredentialSource { Adc, } +impl CredentialSource { + pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option) -> Self { + VertexParams::CREDENTIALS + .resolve(&|_| None, env_lookup) + .map(|configured| Self::Trusted(SecretValue::new(configured))) + .unwrap_or_else(|| Self::application_default(env_lookup)) + } + + pub fn application_default(env_lookup: &dyn Fn(&str) -> Option) -> Self { + non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) + .map_or(Self::Adc, Self::ApplicationCredentials) + } +} + pub(crate) fn credential_source( config: &VertexConfig, env_lookup: &dyn Fn(&str) -> Option, @@ -156,12 +172,7 @@ pub(crate) fn credential_source( } }; } - if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) { - return CredentialSource::Trusted(SecretValue::new(configured)); - } - non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) - .map(CredentialSource::ApplicationCredentials) - .unwrap_or(CredentialSource::Adc) + CredentialSource::from_vertex_environment(env_lookup) } fn optional_credentials( @@ -317,6 +328,24 @@ mod tests { )); } + #[rstest::rstest] + fn application_default_ignores_vertex_credentials() { + let env = |name: &str| match name { + VERTEXAI_CREDENTIALS_ENV => Some("environment-json".to_string()), + GOOGLE_APPLICATION_CREDENTIALS_ENV => Some("adc.json".to_string()), + _ => None, + }; + assert!( + matches!(CredentialSource::application_default(&env), CredentialSource::ApplicationCredentials(path) if path == "adc.json") + ); + assert!(matches!( + CredentialSource::application_default( + &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into()) + ), + CredentialSource::Adc + )); + } + #[rstest::rstest] fn params_become_a_deployment_sourced_config() { let params = VertexParams { diff --git a/litellm-rust/crates/auth-gcp/src/identity.rs b/litellm-rust/crates/auth-gcp/src/identity.rs new file mode 100644 index 00000000000..5144221b218 --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/identity.rs @@ -0,0 +1,445 @@ +use crate::{ + config::{CredentialSource, is_base64_encoded_json, names_a_credentials_file}, + constants::{CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT}, +}; +use gcp_auth::{CustomServiceAccount, TokenProvider}; +use litellm_auth_types::{Error, ErrorDetail}; +use litellm_secret_value::SecretValue; +use moka::future::Cache; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use std::{future::Future, pin::Pin, sync::Arc}; + +#[derive(Clone)] +pub struct GcpAuth { + providers: Cache>, + loader: Arc, +} + +impl Default for GcpAuth { + fn default() -> Self { + Self::new(Arc::new(NativeProviderLoader)) + } +} + +impl GcpAuth { + pub fn new(loader: Arc) -> Self { + Self { + providers: Cache::builder().max_capacity(64).build(), + loader, + } + } + + pub async fn access_token(&self, source: CredentialSource) -> Result { + self.load_provider(source).await?.token().await + } + + pub(crate) async fn load_provider( + &self, + source: CredentialSource, + ) -> Result, Error> { + self.providers + .try_get_with(source.cache_key(), self.loader.load(source)) + .await + .map_err(|error| (*error).clone()) + } +} + +pub trait GcpTokenSource: Send + Sync { + fn project_id(&self) -> GcpAuthFuture<'_, String>; + fn token(&self) -> GcpAuthFuture<'_, SecretValue>; +} + +pub trait GcpProviderLoader: Send + Sync { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc>; +} + +pub type GcpAuthFuture<'a, T> = Pin> + Send + 'a>>; + +struct ServiceAccountTokenSource(Arc); + +impl GcpTokenSource for ServiceAccountTokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .project_id() + .await + .map(|project| project.to_string()) + .map_err(auth_acquisition_error) + }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.0 + .token(&[CLOUD_PLATFORM_SCOPE]) + .await + .map(|token| SecretValue::new(token.as_str())) + .map_err(auth_acquisition_error) + }) + } +} + +struct SdkTokenSource { + credentials: google_cloud_auth::credentials::AccessTokenCredentials, + project_id: Option, +} + +impl GcpTokenSource for SdkTokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async move { + self.project_id + .clone() + .ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "GCP resource project", + })) + }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.credentials + .access_token() + .await + .map(|token| SecretValue::new(token.token)) + .map_err(auth_acquisition_error) + }) + } +} + +async fn load_application_credentials(path: &str) -> Result, Error> { + let contents = tokio::fs::read_to_string(path) + .await + .map_err(auth_acquisition_error)?; + load_credentials_json(&contents) +} + +fn load_credentials_json(contents: &str) -> Result, Error> { + let configured: Value = serde_json::from_str(contents).map_err(auth_acquisition_error)?; + let kind = configured.get("type").and_then(Value::as_str); + let project_id = configured + .get("project_id") + .and_then(Value::as_str) + .map(str::to_string); + let credentials = match kind { + Some("authorized_user") => { + google_cloud_auth::credentials::user_account::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("impersonated_service_account") => { + google_cloud_auth::credentials::impersonated::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("external_account") => { + google_cloud_auth::credentials::external_account::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("service_account") | None => { + let provider = + CustomServiceAccount::from_json(contents).map_err(auth_acquisition_error)?; + return Ok(Arc::new(ServiceAccountTokenSource(Arc::new(provider)))); + } + Some(_) => { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "Google application credentials type".into(), + expected: "a supported Google credential type", + })); + } + } + .map_err(auth_acquisition_error)?; + Ok(Arc::new(SdkTokenSource { + credentials, + project_id, + })) +} + +struct NativeProviderLoader; + +impl GcpProviderLoader for NativeProviderLoader { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc> { + Box::pin(async move { + let provider: Arc = match source { + CredentialSource::Inline(configured) => Arc::new( + CustomServiceAccount::from_json(validate_request_credentials( + configured.expose(), + )?) + .map_err(auth_acquisition_error)?, + ), + CredentialSource::Trusted(configured) => { + return load_trusted_credentials(configured.expose()).await; + } + CredentialSource::ApplicationCredentials(path) => { + return load_application_credentials(&path).await; + } + CredentialSource::Adc => { + gcp_auth::provider().await.map_err(auth_acquisition_error)? + } + }; + Ok(Arc::new(ServiceAccountTokenSource(provider)) as Arc) + }) + } +} + +fn validate_request_credentials(configured: &str) -> Result<&str, Error> { + let token_uri = serde_json::from_str::(configured) + .ok() + .and_then(|credentials| { + credentials + .get("token_uri") + .and_then(Value::as_str) + .map(str::to_string) + }); + if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "request-controlled vertex_credentials token_uri".into(), + expected: "the canonical Google OAuth token endpoint", + })); + } + Ok(configured) +} + +impl CredentialSource { + fn cache_key(&self) -> CredentialCacheKey { + match self { + Self::Inline(configured) => { + CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) + } + Self::Trusted(configured) => { + CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) + } + Self::ApplicationCredentials(path) => { + CredentialCacheKey::ApplicationCredentials(path.clone()) + } + Self::Adc => CredentialCacheKey::Adc, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +enum CredentialCacheKey { + Inline([u8; 32]), + Trusted([u8; 32]), + ApplicationCredentials(String), + Adc, +} + +async fn load_trusted_credentials(configured: &str) -> Result, Error> { + if is_base64_encoded_json(configured) { + return Err(Error::InvalidConfiguration( + ErrorDetail::Base64EncodedJson { + field: "vertex_credentials", + }, + )); + } + if names_a_credentials_file(configured) { + load_application_credentials(configured).await + } else { + load_credentials_json(configured) + } +} + +fn auth_acquisition_error(error: impl std::error::Error + Send + Sync + 'static) -> Error { + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( + "Google credentials", + error, + )) +} + +#[cfg(test)] +mod tests { + use litellm_secret_value::SecretValue; + + use super::*; + + #[rstest::rstest] + #[case::authorized_user("authorized_user", 2)] + #[case::external_account("external_account", 2)] + #[case::impersonated_service_account("impersonated_service_account", 4)] + #[tokio::test] + async fn explicit_application_credentials_refresh_expired_tokens( + #[case] kind: &str, + #[case] expected_requests: usize, + #[values(false, true)] inline: bool, + ) { + use axum::{Json, Router, routing::post}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let _ = rustls::crypto::ring::default_provider().install_default(); + let requests = Arc::new(AtomicUsize::new(0)); + let recorded = requests.clone(); + let impersonations = requests.clone(); + let app = Router::new().route("/token", post(move || { + let recorded = recorded.clone(); + async move { + Json(serde_json::json!({ + "access_token": format!("token-{}", recorded.fetch_add(1, Ordering::SeqCst)), + "expires_in": 3600, + "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + })) + } + })); + let app = app.route("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken", post(move || { + let recorded = impersonations.clone(); + async move { + let expiration = (time::OffsetDateTime::now_utc() + time::Duration::hours(1)).format(&time::format_description::well_known::Rfc3339).unwrap(); + Json(serde_json::json!({ + "accessToken": format!("impersonated-{}", recorded.fetch_add(1, Ordering::SeqCst)), + "expireTime": expiration, + })) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let file = tempfile::NamedTempFile::new().unwrap(); + let subject_token = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(subject_token.path(), "subject-token").unwrap(); + let configured = if kind == "external_account" { + serde_json::json!({ + "type": "external_account", "audience": "//iam.googleapis.com/projects/test/locations/global/workloadIdentityPools/test/providers/test", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": format!("http://{address}/token"), + "credential_source": { "file": subject_token.path().to_str().unwrap() }, "quota_project_id": "quota-project" + }) + } else { + let source = serde_json::json!({ + "type": "authorized_user", "client_id": "user-client", "client_secret": "user-secret", + "refresh_token": "user-refresh", "quota_project_id": "quota-project", "token_uri": format!("http://{address}/token") + }); + if kind == "impersonated_service_account" { + serde_json::json!({ + "type": kind, "source_credentials": source, "quota_project_id": "quota-project", + "service_account_impersonation_url": format!("http://{address}/v1/projects/-/serviceAccounts/test-principal:generateAccessToken") + }) + } else { + source + } + }; + std::fs::write(file.path(), configured.to_string()).unwrap(); + let source = if inline { + CredentialSource::Trusted(SecretValue::new(configured.to_string())) + } else { + CredentialSource::ApplicationCredentials(file.path().to_str().unwrap().into()) + }; + let provider = NativeProviderLoader.load(source).await.unwrap(); + + assert!(matches!( + provider.project_id().await, + Err(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "GCP resource project" + })) + )); + let original = provider.token().await.unwrap(); + tokio::time::pause(); + tokio::time::advance(std::time::Duration::from_secs(3601)).await; + tokio::time::resume(); + let refreshed = provider.token().await.unwrap(); + assert_ne!(original, refreshed); + assert_eq!(requests.load(Ordering::SeqCst), expected_requests); + server.abort(); + } + + #[rstest::rstest] + #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] + #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] + #[case::missing_endpoint("{}", false)] + fn request_credentials_require_canonical_token_endpoint( + #[case] credentials: &str, + #[case] accepted: bool, + ) { + assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + } + + #[rstest::rstest] + #[case::missing_file("/nonexistent/vertex-key.json", "Google credentials")] + #[case::malformed_json(" {\"type\": \"service_account\",", "Google credentials")] + #[case::quoted_json("'{\"type\": \"service_account\"}'", "Google credentials")] + #[tokio::test] + async fn trusted_credentials_are_read_as_a_file_unless_they_are_json( + #[case] configured: &str, + #[case] expected: &str, + ) { + let Err(error) = NativeProviderLoader + .load(CredentialSource::Trusted(SecretValue::new(configured))) + .await + else { + panic!("{configured:?} should not load"); + }; + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(expected), "{report}"); + } + + #[rstest::rstest] + #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")] + #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")] + #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")] + #[tokio::test] + async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) { + assert!( + trusted_load_error(configured) + .await + .contains("is base64-encoded JSON, not the JSON itself") + ); + } + + async fn trusted_load_error(configured: &str) -> String { + match NativeProviderLoader + .load(CredentialSource::Trusted(SecretValue::new(configured))) + .await + { + Ok(_) => panic!("{configured:?} should not load"), + Err(error) => litellm_error::Report::new(&error).to_string(), + } + } + + #[rstest::rstest] + #[tokio::test] + async fn a_relative_key_file_named_like_json_is_read_as_a_file() { + let key = tempfile::Builder::new() + .prefix("{vertex") + .suffix(".json") + .tempfile_in(".") + .unwrap(); + std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); + let relative_name = key.path().file_name().unwrap().to_str().unwrap(); + let absolute_path = std::fs::canonicalize(key.path()).unwrap(); + assert_eq!( + trusted_load_error(relative_name).await, + trusted_load_error(absolute_path.to_str().unwrap()).await + ); + } + + #[rstest::rstest] + #[tokio::test] + async fn unknown_application_credential_types_are_rejected_without_service_account_fallback() { + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + file.path(), + r#"{"type":"unsupported","private_key":"credential-value"}"#, + ) + .unwrap(); + let result = GcpAuth::default() + .access_token(CredentialSource::ApplicationCredentials( + file.path().to_str().unwrap().into(), + )) + .await; + let error = result.unwrap_err(); + assert!(matches!(error, Error::InvalidConfiguration(_))); + assert!( + !litellm_error::Report::new(&error) + .to_string() + .contains("credential-value") + ); + } + + #[rstest::rstest] + fn inline_and_trusted_credentials_never_share_a_cache_entry() { + assert_ne!( + CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), + CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() + ); + } +} diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 2b389d45926..03c660e6eb8 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,16 +1,16 @@ mod auth; mod config; pub mod constants; +mod identity; #[cfg(feature = "google-sdk")] mod sdk; -pub use auth::{ - VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource, -}; +pub use auth::{VertexAuth, VertexEnvironment}; pub use config::{ CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project, get_vertex_ai_project_from_credentials, }; pub use constants::secret_names; +pub use identity::{GcpAuth, GcpAuthFuture, GcpProviderLoader, GcpTokenSource}; #[cfg(feature = "google-sdk")] pub use sdk::GoogleCredentials; diff --git a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs b/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs deleted file mode 100644 index 5991749ab8a..00000000000 --- a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs +++ /dev/null @@ -1,120 +0,0 @@ -use std::{ - collections::{BTreeMap, BTreeSet}, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, -}; - -use litellm_auth_gcp::{ - CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader, - VertexTokenSource, get_vertex_ai_location, secret_names, -}; -use rstest::rstest; -use serde_json::{Value, json}; - -struct FakeProvider { - calls: Arc, -} - -impl VertexTokenSource for FakeProvider { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-project".into()) }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-token".into()) }) - } -} - -struct FakeLoader { - loads: Arc, - provider: Arc, -} - -impl VertexProviderLoader for FakeLoader { - fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - let loads = self.loads.clone(); - let provider = self.provider.clone(); - Box::pin(async move { - loads.fetch_add(1, Ordering::SeqCst); - Ok(provider) - }) - } -} - -fn config(value: Value) -> VertexConfig { - VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) - .unwrap() -} - -fn auth(calls: Arc, loads: Arc) -> VertexAuth { - let provider: Arc = Arc::new(FakeProvider { calls }); - VertexAuth::new(Arc::new(FakeLoader { loads, provider })) -} - -#[rstest] -#[tokio::test] -async fn secret_names_cover_environment_reads() { - let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); - let recorded = seen.clone(); - let env = |name: &str| { - recorded.lock().unwrap().insert(name.to_string()); - None - }; - let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); - auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) - .await - .unwrap(); - get_vertex_ai_location(&VertexConfig::default(), &env); - assert!( - seen.lock() - .unwrap() - .iter() - .all(|name| secret_names().contains(&name.as_str())) - ); -} - -#[rstest] -#[tokio::test] -async fn explicit_token_and_header_do_not_acquire_adc() { - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); - let configured = config(json!({"vertex_project":"project-1"})); - let explicit = auth - .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) - .await - .unwrap(); - assert_eq!(explicit.headers[0].1, "Bearer access-token"); - let existing = auth - .validate_environment( - vec![("authorization".into(), "Bearer existing".into())], - None, - &configured, - &|_| None, - ) - .await - .unwrap(); - assert_eq!(existing.headers[0].1, "Bearer existing"); - assert_eq!(loads.load(Ordering::SeqCst), 0); -} - -#[rstest] -#[tokio::test] -async fn provider_is_reused_across_authentication_calls() { - let calls = Arc::new(AtomicUsize::new(0)); - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(calls.clone(), loads.clone()); - for _ in 0..2 { - let environment = auth - .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None) - .await - .unwrap(); - assert_eq!(environment.project_id, "adc-project"); - assert_eq!(environment.headers[0].1, "Bearer adc-token"); - } - assert_eq!(loads.load(Ordering::SeqCst), 1); - assert_eq!(calls.load(Ordering::SeqCst), 4); -} diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml index 400f2640221..f607f0d1da0 100644 --- a/litellm-rust/crates/auth-types/Cargo.toml +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -7,9 +7,9 @@ repository.workspace = true [dependencies] litellm-error.workspace = true +litellm-secret-value.workspace = true serde.workspace = true serde_json.workspace = true -subtle.workspace = true thiserror.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/auth-types/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs index 4f69c6fffc8..36e21f7794a 100644 --- a/litellm-rust/crates/auth-types/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -5,7 +5,8 @@ use std::sync::Arc; use veil::Redact; -use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle}; +use crate::{Error, ResolvedCredential, TokenProviderHandle}; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { @@ -104,8 +105,8 @@ mod tests { CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; - use crate::SecretValue; use crate::{Error, ErrorDetail}; + use litellm_secret_value::SecretValue; #[derive(Debug)] struct HostResolver; diff --git a/litellm-rust/crates/auth-types/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs index a8ba7b64e03..a2e5629bf3f 100644 --- a/litellm-rust/crates/auth-types/src/error.rs +++ b/litellm-rust/crates/auth-types/src/error.rs @@ -92,6 +92,11 @@ impl Error { #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ErrorDetail { + #[error("{operation} returned HTTP {status}")] + UnexpectedStatus { + operation: &'static str, + status: u16, + }, #[error("{field} must be {expected}")] InvalidType { field: String, diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index 99651977867..ab3bd3febbf 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -5,7 +5,6 @@ mod error; pub mod http; mod params; mod policy; -mod secret; mod token; use serde::{Deserialize, Serialize}; @@ -55,5 +54,4 @@ pub use error::{Error, ErrorDetail}; pub use http::CredentialPlacement; pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; -pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth-types/src/params/aws.rs b/litellm-rust/crates/auth-types/src/params/aws.rs index 8075fc15e9f..de21fbf94c9 100644 --- a/litellm-rust/crates/auth-types/src/params/aws.rs +++ b/litellm-rust/crates/auth-types/src/params/aws.rs @@ -53,6 +53,9 @@ pub struct AwsParams { } impl AwsParams { + pub const ROLE_ARN_ENV: &str = "AWS_ROLE_ARN"; + pub const WEB_IDENTITY_TOKEN_FILE_ENV: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; + pub const ACCESS_KEY_ID: ParamSpec = spec( "access key id", ParamKind::Credential, @@ -149,6 +152,10 @@ impl AwsParams { AwsParams::SPECS .iter() .flat_map(|spec| spec.env.iter().copied()) + .chain([ + AwsParams::ROLE_ARN_ENV, + AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV, + ]) .fold(Vec::new(), |mut names, name| { if !names.contains(&name) { names.push(name); diff --git a/litellm-rust/crates/auth-types/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs index 8fff59e284a..2232350965a 100644 --- a/litellm-rust/crates/auth-types/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -71,7 +71,8 @@ impl ProviderAuthPolicy { #[cfg(test)] mod tests { use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; - use crate::{CredentialPlacement, ResolvedCredential, SecretValue}; + use crate::{CredentialPlacement, ResolvedCredential}; + use litellm_secret_value::SecretValue; const RULES: &[CredentialRule] = &[CredentialRule { kind: CredentialPlanKind::Static, diff --git a/litellm-rust/crates/auth-types/src/token.rs b/litellm-rust/crates/auth-types/src/token.rs index a3126052065..b948dda44fc 100644 --- a/litellm-rust/crates/auth-types/src/token.rs +++ b/litellm-rust/crates/auth-types/src/token.rs @@ -5,7 +5,8 @@ use std::time::SystemTime; use veil::Redact; -use crate::{Error, SecretValue}; +use crate::Error; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, PartialEq, Eq)] pub enum ResolvedCredential { diff --git a/litellm-rust/crates/auth-types/tests/aws_params.rs b/litellm-rust/crates/auth-types/tests/aws_params.rs index ed9a4f78ec2..b438984c199 100644 --- a/litellm-rust/crates/auth-types/tests/aws_params.rs +++ b/litellm-rust/crates/auth-types/tests/aws_params.rs @@ -3,21 +3,13 @@ use rstest::rstest; use serde_json::{Map, Value}; #[rstest] -fn fields_name_every_param_once_and_in_declaration_order() { +fn every_named_param_roundtrips_without_loss() { let filled: Map = AwsParams::fields() .map(|name| (name.to_string(), Value::from(format!("value-of-{name}")))) .collect(); let typed = AwsParams::from_optional_params(&filled); let serialized = serde_json::to_value(&typed).unwrap(); - assert_eq!( - serialized - .as_object() - .unwrap() - .keys() - .map(String::as_str) - .collect::>(), - AwsParams::fields().collect::>() - ); + assert_eq!(filled.len(), AwsParams::fields().count()); assert_eq!(serialized, Value::Object(filled.clone())); assert_eq!( serde_json::from_value::(Value::Object(filled)).unwrap(), diff --git a/litellm-rust/crates/auth-types/tests/token.rs b/litellm-rust/crates/auth-types/tests/token.rs index 7f7027942b0..39908ae7b8a 100644 --- a/litellm-rust/crates/auth-types/tests/token.rs +++ b/litellm-rust/crates/auth-types/tests/token.rs @@ -8,9 +8,9 @@ use std::{ time::{Duration, SystemTime}, }; -use litellm_auth_types::{ - Error, ErrorDetail, ResolvedCredential, SecretValue, TokenProviderHandle, -}; +use litellm_auth_types::{Error, ErrorDetail, ResolvedCredential, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use rstest::rstest; fn credential(index: usize, access_token: bool) -> ResolvedCredential { diff --git a/litellm-rust/crates/auth/Cargo.toml b/litellm-rust/crates/auth/Cargo.toml index ee4900ebcc2..674ae65a59f 100644 --- a/litellm-rust/crates/auth/Cargo.toml +++ b/litellm-rust/crates/auth/Cargo.toml @@ -16,3 +16,6 @@ litellm-auth-types.workspace = true litellm-auth-aws = { workspace = true, optional = true } litellm-auth-azure = { workspace = true, optional = true } litellm-auth-gcp = { workspace = true, optional = true } + +[dev-dependencies] +litellm-secret-value.workspace = true diff --git a/litellm-rust/crates/auth/tests/facade.rs b/litellm-rust/crates/auth/tests/facade.rs index f1092b15def..5d29bac1072 100644 --- a/litellm-rust/crates/auth/tests/facade.rs +++ b/litellm-rust/crates/auth/tests/facade.rs @@ -1,7 +1,8 @@ use litellm_auth::{ CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, - ProviderAuthPolicy, ResolvedCredential, SecretValue, + ProviderAuthPolicy, ResolvedCredential, }; +use litellm_secret_value::SecretValue; const RULES: &[CredentialRule] = &[CredentialRule { kind: CredentialPlanKind::Static, diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index dbd8c9ac908..878db4b8010 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true futures-util.workspace = true litellm-auth-gcp.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index a6b08445bc8..b8c6720c234 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -8,7 +8,7 @@ use litellm_cache::{ use litellm_http::Client; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; -use crate::{GcpTokenSource, TokenSource}; +use crate::{GcsTokenSource, TokenSource}; pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; @@ -62,7 +62,7 @@ pub struct GcsCache { impl GcsCache { pub fn new(config: GcsConfig, client: Client, codec: S) -> Self { - let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone())); + let token = Arc::new(GcsTokenSource::new(config.path_service_account.clone())); Self::with_token_source(config, client, codec, token) } diff --git a/litellm-rust/crates/cache-gcs/src/lib.rs b/litellm-rust/crates/cache-gcs/src/lib.rs index cbb61cf0685..d40bca081a4 100644 --- a/litellm-rust/crates/cache-gcs/src/lib.rs +++ b/litellm-rust/crates/cache-gcs/src/lib.rs @@ -2,4 +2,4 @@ mod cache; mod token; pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix}; -pub use token::{GcpTokenSource, StaticTokenSource, TokenSource}; +pub use token::{GcsTokenSource, StaticTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index 079ec44dbce..4f639e7b865 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -1,35 +1,38 @@ use std::{future::Future, pin::Pin}; -use litellm_auth_gcp::{VertexAuth, VertexConfig}; -use litellm_auth_types::{InputSource, SecretValue, Sourced}; +use litellm_auth_gcp::{CredentialSource, GcpAuth}; +use litellm_secret_value::SecretValue; + use litellm_cache::Error; pub trait TokenSource: Send + Sync + 'static { fn bearer_token(&self) -> Pin> + Send + '_>>; } -pub struct GcpTokenSource { - auth: VertexAuth, - config: VertexConfig, +pub struct GcsTokenSource { + auth: GcpAuth, + credentials: Option, } -impl GcpTokenSource { +impl GcsTokenSource { pub fn new(path_service_account: Option) -> Self { - let credentials = path_service_account - .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment)); Self { - auth: VertexAuth::default(), - config: VertexConfig::new(credentials, None, None), + auth: GcpAuth::default(), + credentials: path_service_account.filter(|path| !path.trim().is_empty()), } } } -impl TokenSource for GcpTokenSource { +impl TokenSource for GcsTokenSource { fn bearer_token(&self) -> Pin> + Send + '_>> { Box::pin(async move { self.auth - .access_token(&self.config, &|name| std::env::var(name).ok()) + .access_token(match &self.credentials { + Some(path) => CredentialSource::Trusted(SecretValue::new(path.clone())), + None => CredentialSource::application_default(&|name| std::env::var(name).ok()), + }) .await + .map(|token| token.expose().to_owned()) .map_err(Error::unavailable) }) } diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 73fbf14b07a..809ffe3044e 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -2,7 +2,10 @@ use std::{ cmp::Reverse, collections::{BinaryHeap, HashMap, HashSet}, hash::Hash, - sync::{Arc, Mutex, MutexGuard}, + sync::{ + Arc, Mutex, MutexGuard, + atomic::{AtomicUsize, Ordering}, + }, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -31,7 +34,7 @@ struct CacheState { pub struct InMemoryCache { state: Mutex>, - max_size_in_memory: usize, + max_size_in_memory: AtomicUsize, default_ttl: Duration, max_entry_bytes: Option, measure_value: Option>, @@ -74,7 +77,9 @@ impl InMemoryCache { expirations: HashMap::new(), expiration_heap: BinaryHeap::new(), }), - max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + max_size_in_memory: AtomicUsize::new( + max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + ), default_ttl: default_ttl .filter(|ttl| !ttl.is_zero()) .unwrap_or(DEFAULT_TTL), @@ -94,7 +99,7 @@ impl InMemoryCache { value: V, ttl: Option, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(CacheWrite::Disabled); } let now = (self.now)(); @@ -129,7 +134,7 @@ impl InMemoryCache { pub fn evict_cache(&self) -> Result<(), Error> { let now = (self.now)(); let mut state = self.lock()?; - Self::evict(&mut state, self.max_size_in_memory, now, None); + Self::evict(&mut state, self.max_size_in_memory(), now, None); Ok(()) } @@ -170,7 +175,16 @@ impl InMemoryCache { } pub fn max_size_in_memory(&self) -> usize { - self.max_size_in_memory + self.max_size_in_memory.load(Ordering::Relaxed) + } + + /// `DualCache.update_in_memory_max_size`: `None` restores the default. Entries already held stay + /// until the next write evicts down to the new size. + pub fn update_max_size_in_memory(&self, max_size_in_memory: Option) { + self.max_size_in_memory.store( + max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + Ordering::Relaxed, + ); } pub fn max_entry_bytes(&self) -> Option { @@ -200,6 +214,23 @@ impl InMemoryCache { .collect()) } + /// Every live entry, for callers that inspect the whole tier. + pub fn entries(&self) -> Result, Error> { + let now = (self.now)(); + let state = self.lock()?; + Ok(state + .values + .iter() + .filter(|(key, _)| { + state + .expirations + .get(*key) + .is_none_or(|expiration| *expiration >= now) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect()) + } + pub fn delete_cache(&self, key: &str) -> Result<(), Error> { let mut state = self.lock()?; Self::remove(&mut state, key); @@ -276,7 +307,7 @@ impl InMemoryCache { ttl: Option, now: Duration, ) -> Result { - Self::evict(state, self.max_size_in_memory, now, None); + Self::evict(state, self.max_size_in_memory(), now, None); if !self.check_value_size(&value)? { return Ok(CacheWrite::TooLarge); } @@ -300,12 +331,12 @@ where eligible: &[V], context: ExactCacheContext, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(candidate); } let now = (self.now)(); let mut state = self.lock()?; - Self::evict(&mut state, self.max_size_in_memory, now, Some(key)); + Self::evict(&mut state, self.max_size_in_memory(), now, Some(key)); let existing = state .values .get(key) @@ -335,7 +366,7 @@ impl CounterCache for InMemoryCache { amount: f64, context: ExactCacheContext, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(amount); } let now = (self.now)(); @@ -408,7 +439,7 @@ where values: Vec, ttl: Option, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(values); } let now = (self.now)(); diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml deleted file mode 100644 index 6a5c794bde7..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/Cargo.toml +++ /dev/null @@ -1,19 +0,0 @@ -[package] -name = "litellm-cache-redis-semantic" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true - -[dependencies] -litellm-cache.workspace = true -litellm-cache-redis.workspace = true -redis = { workspace = true, features = ["tls-rustls"] } -sha2.workspace = true - -[dev-dependencies] -litellm-cache-testing.workspace = true -redis-test.workspace = true -rstest.workspace = true -serde_json.workspace = true -tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs deleted file mode 100644 index 42fc6ae585f..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ /dev/null @@ -1,404 +0,0 @@ -use std::{ - sync::Arc, - time::{Duration, SystemTime, UNIX_EPOCH}, -}; - -use litellm_cache::{ - BaseCache, CacheCodec, Error, SemanticCacheContext, - semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context}, -}; -use litellm_cache_redis::{ - RedisTopology, - connection::{ConnectionRef, Connections}, -}; -use sha2::{Digest, Sha256}; - -use crate::{ - RedisSemanticConfig, - index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD}, - reply::{bytes_field, first_document, number_field, string_field}, -}; - -struct Inner { - index: Index, - distance_threshold: f64, - clock: fn() -> f64, -} - -impl Inner { - fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self { - Self { - index: Index::new(config.index_name), - distance_threshold: 1.0 - f64::from(config.similarity_threshold), - clock, - } - } - - fn store( - &self, - connection: &mut ConnectionRef<'_>, - tag: &str, - response: Vec, - prompt: &str, - vector: &[f32], - ttl: Option, - ) -> Result<(), Error> { - let index = self.index.ensure(connection, vector.len())?; - let entry_id = entry_id(prompt, tag); - let hash_key = format!("{index}:{entry_id}"); - redis::cmd("HSET") - .arg(&hash_key) - .arg("entry_id") - .arg(&entry_id) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(response) - .arg(VECTOR_FIELD) - .arg(vector_buffer(vector)) - .arg("inserted_at") - .arg(format!("{}", (self.clock)())) - .arg("updated_at") - .arg(format!("{}", (self.clock)())) - .arg(CACHE_KEY_FIELD) - .arg(tag) - .query::<()>(connection) - .map_err(Error::unavailable)?; - if let Some(ttl) = ttl { - redis::cmd("EXPIRE") - .arg(&hash_key) - .arg(ttl_seconds(ttl)) - .query::<()>(connection) - .map_err(Error::unavailable)?; - } - Ok(()) - } - - fn lookup( - &self, - connection: &mut ConnectionRef<'_>, - tag: &str, - vector: &[f32], - ) -> Result>, Error> { - let index = self.index.ensure(connection, vector.len())?; - let query = format!( - "(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]", - escape_tag(tag) - ); - let result = redis::cmd("FT.SEARCH") - .arg(&index) - .arg(query) - .arg("RETURN") - .arg(8) - .arg("entry_id") - .arg("prompt") - .arg("response") - .arg("inserted_at") - .arg("updated_at") - .arg("metadata") - .arg(CACHE_KEY_FIELD) - .arg("vector_distance") - .arg("SORTBY") - .arg("vector_distance") - .arg("ASC") - .arg("DIALECT") - .arg(2) - .arg("LIMIT") - .arg(0) - .arg(1) - .arg("PARAMS") - .arg(2) - .arg("vector") - .arg(vector_buffer(vector)) - .query::(connection) - .map_err(Error::unavailable)?; - let Some(fields) = first_document(&result) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) { - return Ok(SemanticLookup::miss(Some(0.0))); - } - // redisvl's range query only returns entries within the distance threshold, so a - // farther hit reads as no result. - let Some(distance) = number_field(fields, "vector_distance") - .filter(|distance| *distance <= self.distance_threshold) - else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let Some(response) = bytes_field(fields, "response") else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - Ok(SemanticLookup { - value: Some(response), - similarity: Some(1.0 - distance), - }) - } -} - -/// `RedisSemanticCache`: a redisvl-compatible semantic index on Redis Stack. Values go through -/// the injected codec, so the response layer decides what a cached entry is. -pub struct RedisSemanticCache { - connections: Arc>, - embedder: E, - codec: S, - inner: Arc, -} - -impl RedisSemanticCache { - pub fn new( - url: &str, - embedder: E, - codec: S, - config: RedisSemanticConfig, - ) -> Result { - Ok(Self { - connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), - embedder, - codec, - inner: Arc::new(Inner::new(config, timestamp)), - }) - } -} - -impl RedisSemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub fn with_connection( - connection: C, - embedder: E, - codec: S, - config: RedisSemanticConfig, - ) -> Self { - Self { - connections: Arc::new(Connections::fixed(connection)), - embedder, - codec, - inner: Arc::new(Inner::new(config, timestamp)), - } - } - - pub fn with_clock(self, clock: fn() -> f64) -> Self { - let config = RedisSemanticConfig { - index_name: self.index_name().to_owned(), - similarity_threshold: self.similarity_threshold(), - }; - Self { - inner: Arc::new(Inner::new(config, clock)), - ..self - } - } - - pub fn embedder(&self) -> &E { - &self.embedder - } - - pub fn index_name(&self) -> &str { - self.inner.index.name() - } - - pub fn similarity_threshold(&self) -> f32 { - (1.0 - self.inner.distance_threshold) as f32 - } - - fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { - context.scope.as_deref().unwrap_or(key) - } - - fn decode(&self, lookup: SemanticLookup>) -> Result, Error> { - Ok(SemanticLookup { - value: lookup - .value - .map(|bytes| self.codec.decode(&bytes)) - .transpose()?, - similarity: lookup.similarity, - }) - } -} - -impl BaseCache for RedisSemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Value = S::Value; - type Context = SemanticCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl - } - - fn set_cache( - &self, - key: &str, - value: Self::Value, - context: &Self::Context, - ) -> Result<(), Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(()); - }; - let response = self.codec.encode(&value)?; - let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let tag = Self::tag(key, context); - self.connections.execute(|connection| { - self.inner - .store(connection, tag, response, &prompt, &vector, context.ttl) - }) - } - - fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { - self.get_cache_with_similarity(key, context) - .map(|lookup| lookup.value) - } - - async fn async_set_cache( - &self, - key: &str, - value: Self::Value, - context: Self::Context, - ) -> Result<(), Error> { - let Some(prompt) = prompt_from_context(&context) else { - return Ok(()); - }; - let response = self.codec.encode(&value)?; - let vector = self - .embedder - .async_embed(&prompt, context.metadata.as_ref()) - .await?; - let tag = Self::tag(key, &context).to_owned(); - let inner = Arc::clone(&self.inner); - Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - inner.store(connection, &tag, response, &prompt, &vector, context.ttl) - }) - .await - } - - async fn async_get_cache( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - self.async_get_cache_with_similarity(key, context) - .await - .map(|lookup| lookup.value) - } -} - -/// Python stamps a similarity of `0.0` when there is no prompt or no hit in the key's scope. -impl SemanticCache for RedisSemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn get_cache_with_similarity( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let tag = Self::tag(key, context); - let lookup = self - .connections - .execute(|connection| self.inner.lookup(connection, tag, &vector))?; - self.decode(lookup) - } - - async fn async_get_cache_with_similarity( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let vector = self - .embedder - .async_embed(&prompt, context.metadata.as_ref()) - .await?; - let tag = Self::tag(key, context).to_owned(); - let inner = Arc::clone(&self.inner); - let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - inner.lookup(connection, &tag, &vector) - }) - .await?; - self.decode(lookup) - } -} - -fn timestamp() -> f64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_secs_f64()) - .unwrap_or_default() -} - -fn entry_id(prompt: &str, tag: &str) -> String { - let mut digest = Sha256::new(); - digest.update(prompt.as_bytes()); - digest.update(CACHE_KEY_FIELD.as_bytes()); - digest.update(tag.as_bytes()); - format!("{:x}", digest.finalize()) -} - -fn vector_buffer(vector: &[f32]) -> Vec { - vector - .iter() - .flat_map(|component| component.to_le_bytes()) - .collect() -} - -fn escape_tag(value: &str) -> String { - let mut escaped = String::with_capacity(value.len()); - for ch in value.chars() { - if matches!( - ch, - ',' | '.' - | '<' - | '>' - | '{' - | '}' - | '[' - | ']' - | '\\' - | '"' - | '\'' - | ':' - | ';' - | '!' - | '@' - | '#' - | '$' - | '%' - | '^' - | '&' - | '*' - | '(' - | ')' - | '-' - | '+' - | '=' - | '~' - | '|' - | '/' - | ' ' - | '?' - ) { - escaped.push('\\'); - } - escaped.push(ch); - } - escaped -} - -fn ttl_seconds(ttl: Duration) -> u64 { - ttl.as_secs() - .saturating_add(u64::from(ttl.subsec_nanos() > 0)) - .max(1) -} diff --git a/litellm-rust/crates/cache-redis-semantic/src/config.rs b/litellm-rust/crates/cache-redis-semantic/src/config.rs deleted file mode 100644 index 6b810628be7..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/config.rs +++ /dev/null @@ -1,8 +0,0 @@ -/// `RedisSemanticCache.DEFAULT_REDIS_INDEX_NAME`. -pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; - -#[derive(Clone, Debug)] -pub struct RedisSemanticConfig { - pub index_name: String, - pub similarity_threshold: f32, -} diff --git a/litellm-rust/crates/cache-redis-semantic/src/index.rs b/litellm-rust/crates/cache-redis-semantic/src/index.rs deleted file mode 100644 index 23401938a0e..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/index.rs +++ /dev/null @@ -1,209 +0,0 @@ -use std::sync::OnceLock; - -use litellm_cache::Error; -use litellm_cache_redis::connection::ConnectionRef; - -use crate::reply::{number_value, string_value}; - -pub(crate) const CACHE_KEY_FIELD: &str = "litellm_cache_key"; -pub(crate) const VECTOR_FIELD: &str = "prompt_vector"; - -/// The redisvl `SemanticCache` index, resolved once per cache: the configured name when its -/// schema fits, else `_isolated`, recreated when that one is stale too. -pub(crate) struct Index { - name: String, - resolved: OnceLock, -} - -impl Index { - pub(crate) fn new(name: String) -> Self { - Self { - name, - resolved: OnceLock::new(), - } - } - - pub(crate) fn name(&self) -> &str { - &self.name - } - - pub(crate) fn ensure( - &self, - connection: &mut ConnectionRef<'_>, - dims: usize, - ) -> Result { - if let Some(name) = self.resolved.get() { - return Ok(name.clone()); - } - let name = match index_compatible(connection, &self.name, dims)? { - Some(true) => self.name.clone(), - Some(false) => self.isolated(connection, dims)?, - None => match create_index(connection, &self.name, dims) { - Ok(()) => self.name.clone(), - Err(error) => match index_compatible(connection, &self.name, dims)? { - Some(true) => self.name.clone(), - Some(false) => self.isolated(connection, dims)?, - None => return Err(error), - }, - }, - }; - let _ = self.resolved.set(name.clone()); - Ok(name) - } - - fn isolated(&self, connection: &mut ConnectionRef<'_>, dims: usize) -> Result { - let name = format!("{}_isolated", self.name); - match index_compatible(connection, &name, dims)? { - Some(true) => Ok(name), - Some(false) => { - redis::cmd("FT.DROPINDEX") - .arg(&name) - .query::<()>(connection) - .map_err(Error::unavailable)?; - create_index(connection, &name, dims)?; - Ok(name) - } - None => { - create_index(connection, &name, dims)?; - Ok(name) - } - } - } -} - -fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> { - redis::cmd("FT.CREATE") - .arg(name) - .arg("ON") - .arg("HASH") - .arg("PREFIX") - .arg(1) - .arg(name) - .arg("SCORE") - .arg(1.0) - .arg("SCHEMA") - .arg("prompt") - .arg("TEXT") - .arg("WEIGHT") - .arg(1) - .arg("response") - .arg("TEXT") - .arg("WEIGHT") - .arg(1) - .arg("inserted_at") - .arg("NUMERIC") - .arg("updated_at") - .arg("NUMERIC") - .arg(VECTOR_FIELD) - .arg("VECTOR") - .arg("FLAT") - .arg(6) - .arg("TYPE") - .arg("FLOAT32") - .arg("DIM") - .arg(dims) - .arg("DISTANCE_METRIC") - .arg("COSINE") - .arg(CACHE_KEY_FIELD) - .arg("TAG") - .arg("SEPARATOR") - .arg(",") - .query::<()>(connection) - .map_err(Error::unavailable) -} - -fn index_compatible( - connection: &mut ConnectionRef<'_>, - name: &str, - dims: usize, -) -> Result, Error> { - let info = match redis::cmd("FT.INFO") - .arg(name) - .query::(connection) - { - Ok(info) => info, - Err(error) if unknown_index(&error) => return Ok(None), - Err(error) => return Err(Error::unavailable(error)), - }; - Ok(Some(schema_compatible(&info, dims))) -} - -fn unknown_index(error: &redis::RedisError) -> bool { - error - .code() - .is_some_and(|code| code.eq_ignore_ascii_case("unknown")) - && error - .detail() - .is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index")) -} - -struct Attribute { - name: Option, - field_type: Option, - dim: Option, - data_type: Option, - distance_metric: Option, -} - -fn attribute(value: &redis::Value) -> Option { - let redis::Value::Array(pairs) = value else { - return None; - }; - let mut attribute = Attribute { - name: None, - field_type: None, - dim: None, - data_type: None, - distance_metric: None, - }; - for pair in pairs.as_chunks::<2>().0 { - match string_value(&pair[0]).as_deref() { - Some("identifier") => attribute.name = string_value(&pair[1]), - Some("type") => attribute.field_type = string_value(&pair[1]), - Some("dim") => attribute.dim = number_value(&pair[1]), - Some("data_type") => attribute.data_type = string_value(&pair[1]), - Some("distance_metric") => attribute.distance_metric = string_value(&pair[1]), - _ => {} - } - } - Some(attribute) -} - -fn schema_compatible(info: &redis::Value, dims: usize) -> bool { - let redis::Value::Array(entries) = info else { - return false; - }; - let attributes = entries - .as_chunks::<2>() - .0 - .iter() - .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes")) - .map(|pair| &pair[1]); - let Some(redis::Value::Array(attributes)) = attributes else { - return false; - }; - let fields = attributes.iter().filter_map(attribute).collect::>(); - let has_field = |name: &str, field_type: &str| { - fields.iter().any(|field| { - field.name.as_deref() == Some(name) && field.field_type.as_deref() == Some(field_type) - }) - }; - has_field("prompt", "TEXT") - && has_field("response", "TEXT") - && has_field("inserted_at", "NUMERIC") - && has_field("updated_at", "NUMERIC") - && has_field(CACHE_KEY_FIELD, "TAG") - && fields.iter().any(|field| { - field.name.as_deref() == Some(VECTOR_FIELD) - && field.field_type.as_deref() == Some("VECTOR") - && field.dim == Some(dims as f64) - && field - .data_type - .as_deref() - .is_some_and(|data| data.eq_ignore_ascii_case("float32")) - && field - .distance_metric - .as_deref() - .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine")) - }) -} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs deleted file mode 100644 index 251323df6cf..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/lib.rs +++ /dev/null @@ -1,7 +0,0 @@ -mod cache; -mod config; -mod index; -mod reply; - -pub use cache::RedisSemanticCache; -pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig}; diff --git a/litellm-rust/crates/cache-redis-semantic/src/reply.rs b/litellm-rust/crates/cache-redis-semantic/src/reply.rs deleted file mode 100644 index 24cbd573cde..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/reply.rs +++ /dev/null @@ -1,57 +0,0 @@ -pub(crate) fn string_value(value: &redis::Value) -> Option { - match value { - redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), - redis::Value::SimpleString(text) => Some(text.clone()), - redis::Value::VerbatimString { text, .. } => Some(text.clone()), - _ => None, - } -} - -pub(crate) fn number_value(value: &redis::Value) -> Option { - match value { - redis::Value::Int(number) => Some(*number as f64), - redis::Value::Double(number) => Some(*number), - _ => string_value(value).and_then(|text| text.parse().ok()), - } -} - -pub(crate) fn first_document(result: &redis::Value) -> Option<&[redis::Value]> { - let redis::Value::Array(items) = result else { - return None; - }; - let [count, _document_id, fields, ..] = items.as_slice() else { - return None; - }; - if !matches!(count, redis::Value::Int(count) if *count > 0) { - return None; - } - match fields { - redis::Value::Array(fields) => Some(fields.as_slice()), - _ => None, - } -} - -fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> { - fields - .as_chunks::<2>() - .0 - .iter() - .find(|pair| string_value(&pair[0]).as_deref() == Some(name)) - .map(|pair| &pair[1]) -} - -pub(crate) fn string_field(fields: &[redis::Value], name: &str) -> Option { - field_value(fields, name).and_then(string_value) -} - -pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option { - field_value(fields, name).and_then(number_value) -} - -pub(crate) fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { - match field_value(fields, name)? { - redis::Value::BulkString(bytes) => Some(bytes.clone()), - redis::Value::SimpleString(text) => Some(text.clone().into_bytes()), - _ => None, - } -} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs deleted file mode 100644 index cc42f4c3924..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ /dev/null @@ -1,778 +0,0 @@ -mod support; - -use std::time::Duration; - -use litellm_cache::{ - BaseCache, Error, JsonCodec, SemanticCacheContext, - semantic::{SemanticCache, SemanticLookup}, -}; -use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig}; -use redis_test::{MockCmd, MockRedisConnection}; -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; -use sha2::{Digest, Sha256}; -use support::FakeEmbedder; - -const INDEX: &str = DEFAULT_INDEX_NAME; -const PROMPT: &str = "hello prompt"; -const CLOCK: fn() -> f64 = || 1700000000.5; -const VECTOR: [f32; 3] = [0.1, 0.2, 0.3]; - -type MockCache = RedisSemanticCache, MockRedisConnection>; - -#[fixture] -fn config() -> RedisSemanticConfig { - RedisSemanticConfig { - index_name: INDEX.into(), - similarity_threshold: 0.9, - } -} - -#[fixture] -fn entry() -> Value { - json!({"timestamp": 1.0, "response": {"answer": "yes"}}) -} - -#[fixture] -fn context() -> SemanticCacheContext { - messages_context(vec![json!({"role": "user", "content": PROMPT})]) -} - -fn messages_context(messages: Vec) -> SemanticCacheContext { - SemanticCacheContext { - messages: Some(Value::Array(messages)), - ..Default::default() - } -} - -fn cache(commands: Vec, embedder: FakeEmbedder) -> MockCache { - RedisSemanticCache::with_connection( - MockRedisConnection::new(commands).assert_all_commands_consumed(), - embedder, - JsonCodec::new(), - config(), - ) - .with_clock(CLOCK) -} - -fn encoded(value: &Value) -> Vec { - serde_json::to_vec(value).unwrap() -} - -fn vector_bytes(vector: &[f32]) -> Vec { - vector - .iter() - .flat_map(|component| component.to_le_bytes()) - .collect() -} - -fn entry_id(prompt: &str, tag: &str) -> String { - let mut digest = Sha256::new(); - digest.update(prompt.as_bytes()); - digest.update(b"litellm_cache_key"); - digest.update(tag.as_bytes()); - format!("{:x}", digest.finalize()) -} - -fn s(value: &str) -> redis::Value { - redis::Value::BulkString(value.as_bytes().to_vec()) -} - -fn unknown_index_error() -> redis::RedisError { - support::server_error("Unknown index name") -} - -fn info_missing(index: &str) -> MockCmd { - MockCmd::new( - redis::cmd("FT.INFO").arg(index), - Err::(unknown_index_error()), - ) -} - -fn info(index: &str, value: redis::Value) -> MockCmd { - MockCmd::new(redis::cmd("FT.INFO").arg(index), Ok(value)) -} - -fn attribute(name: &str, field_type: &str, extra: Vec) -> redis::Value { - let mut parts = vec![ - s("identifier"), - s(name), - s("attribute"), - s(name), - s("type"), - s(field_type), - ]; - parts.extend(extra); - redis::Value::Array(parts) -} - -fn index_info(attributes: Vec) -> redis::Value { - redis::Value::Array(vec![ - s("index_name"), - s(INDEX), - s("attributes"), - redis::Value::Array(attributes), - ]) -} - -fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value { - attribute( - "prompt_vector", - "VECTOR", - vec![ - s("algorithm"), - s("FLAT"), - s("data_type"), - s(data_type), - s("dim"), - redis::Value::Int(dims), - s("distance_metric"), - s(distance_metric), - ], - ) -} - -fn info_with_vector(vector: redis::Value) -> redis::Value { - index_info(vec![ - attribute("prompt", "TEXT", vec![]), - attribute("response", "TEXT", vec![]), - attribute("inserted_at", "NUMERIC", vec![]), - attribute("updated_at", "NUMERIC", vec![]), - vector, - attribute("litellm_cache_key", "TAG", vec![]), - ]) -} - -fn compatible_info(dims: i64) -> redis::Value { - info_with_vector(vector_attribute_with(dims, "FLOAT32", "COSINE")) -} - -fn unscoped_info(dims: i64) -> redis::Value { - index_info(vec![ - attribute("prompt", "TEXT", vec![]), - attribute("response", "TEXT", vec![]), - attribute("inserted_at", "NUMERIC", vec![]), - attribute("updated_at", "NUMERIC", vec![]), - vector_attribute_with(dims, "FLOAT32", "COSINE"), - ]) -} - -fn create_index(name: &str, dims: usize) -> MockCmd { - MockCmd::new(create_index_command(name, dims), Ok("OK")) -} - -fn create_index_command(name: &str, dims: usize) -> redis::Cmd { - let mut command = redis::cmd("FT.CREATE"); - command - .arg(name) - .arg("ON") - .arg("HASH") - .arg("PREFIX") - .arg(1) - .arg(name) - .arg("SCORE") - .arg(1.0) - .arg("SCHEMA") - .arg("prompt") - .arg("TEXT") - .arg("WEIGHT") - .arg(1) - .arg("response") - .arg("TEXT") - .arg("WEIGHT") - .arg(1) - .arg("inserted_at") - .arg("NUMERIC") - .arg("updated_at") - .arg("NUMERIC") - .arg("prompt_vector") - .arg("VECTOR") - .arg("FLAT") - .arg(6) - .arg("TYPE") - .arg("FLOAT32") - .arg("DIM") - .arg(dims) - .arg("DISTANCE_METRIC") - .arg("COSINE") - .arg("litellm_cache_key") - .arg("TAG") - .arg("SEPARATOR") - .arg(","); - command -} - -fn hset(index: &str, prompt: &str, tag: &str, vector: &[f32], value: &Value) -> MockCmd { - MockCmd::new( - redis::cmd("HSET") - .arg(format!("{index}:{}", entry_id(prompt, tag))) - .arg("entry_id") - .arg(entry_id(prompt, tag)) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(value)) - .arg("prompt_vector") - .arg(vector_bytes(vector)) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg(tag), - Ok(7), - ) -} - -fn search( - index: &str, - tag: &str, - vector: &[f32], - reply: redis::RedisResult, -) -> MockCmd { - let mut command = redis::cmd("FT.SEARCH"); - command - .arg(index) - .arg(format!( - "(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]" - )) - .arg("RETURN") - .arg(8) - .arg("entry_id") - .arg("prompt") - .arg("response") - .arg("inserted_at") - .arg("updated_at") - .arg("metadata") - .arg("litellm_cache_key") - .arg("vector_distance") - .arg("SORTBY") - .arg("vector_distance") - .arg("ASC") - .arg("DIALECT") - .arg(2) - .arg("LIMIT") - .arg(0) - .arg(1) - .arg("PARAMS") - .arg(2) - .arg("vector") - .arg(vector_bytes(vector)); - MockCmd::new(command, reply) -} - -fn hit(tag: &str, distance: &str, response: Vec) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::Int(1), - s("litellm_semantic_cache_index:stored-id"), - redis::Value::Array(vec![ - s("entry_id"), - s("stored-id"), - s("prompt"), - s(PROMPT), - s("response"), - redis::Value::BulkString(response), - s("inserted_at"), - s("1700000000.5"), - s("updated_at"), - s("1700000000.5"), - s("litellm_cache_key"), - s(tag), - s("vector_distance"), - s(distance), - ]), - ]) -} - -fn empty_result() -> redis::Value { - redis::Value::Array(vec![redis::Value::Int(0)]) -} - -#[rstest] -#[case::creates_index_and_expires(false, Some(Duration::from_secs(5)), Some(5))] -#[case::existing_index_without_ttl(true, None, None)] -#[case::fractional_ttl_rounds_up(true, Some(Duration::from_millis(1500)), Some(2))] -fn store_writes_the_redisvl_hash( - #[case] index_exists: bool, - #[case] ttl: Option, - #[case] expire: Option, - entry: Value, - context: SemanticCacheContext, -) { - let hash_key = format!("{INDEX}:{}", entry_id(PROMPT, "key1")); - let mut commands = if index_exists { - vec![info(INDEX, compatible_info(3))] - } else { - vec![info_missing(INDEX), create_index(INDEX, 3)] - }; - commands.push(hset(INDEX, PROMPT, "key1", &VECTOR, &entry)); - commands.extend( - expire.map(|seconds| MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(seconds), Ok(1))), - ); - let cache = cache(commands, FakeEmbedder::new(&[])); - - cache - .set_cache("key1", entry, &SemanticCacheContext { ttl, ..context }) - .unwrap(); -} - -#[rstest] -#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))] -#[case::above_distance_threshold("key1", "0.5", None, Ok(None))] -#[case::other_cache_key("other", "0.05", None, Ok(None))] -#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))] -fn lookup_applies_threshold_scope_and_codec( - #[case] stored_tag: &str, - #[case] distance: &str, - #[case] response: Option<&[u8]>, - #[case] expected: Result, Error>, - entry: Value, - context: SemanticCacheContext, -) { - let response = response.map_or_else(|| encoded(&entry), <[u8]>::to_vec); - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - search( - INDEX, - "key1", - &VECTOR, - Ok(hit(stored_tag, distance, response)), - ), - ], - FakeEmbedder::new(&[]), - ); - - let kind = |result: Result, Error>| { - result.map_err(|error| std::mem::discriminant(&error)) - }; - assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected)); -} - -#[rstest] -#[case::hit(context(), Some(hit("key1", "0.05", encoded(&entry()))), Some(entry()), Some(1.0 - 0.05))] -#[case::beyond_distance_threshold(context(), Some(hit("key1", "0.5", encoded(&entry()))), None, Some(0.0))] -#[case::no_results(context(), Some(empty_result()), None, Some(0.0))] -#[case::other_cache_key(context(), Some(hit("other", "0.05", encoded(&entry()))), None, Some(0.0))] -#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))] -#[tokio::test] -async fn lookup_reports_python_semantic_similarity( - #[case] context: SemanticCacheContext, - #[case] reply: Option, - #[case] value: Option, - #[case] similarity: Option, - #[values(false, true)] use_async: bool, -) { - let commands = reply.map_or_else(Vec::new, |reply| { - vec![ - info(INDEX, compatible_info(3)), - search(INDEX, "key1", &VECTOR, Ok(reply)), - ] - }); - let cache = cache(commands, FakeEmbedder::new(&[])); - - let lookup = if use_async { - cache - .async_get_cache_with_similarity("key1", &context) - .await - } else { - cache.get_cache_with_similarity("key1", &context) - }; - - assert_eq!(lookup, Ok(SemanticLookup { value, similarity })); -} - -#[rstest] -#[tokio::test] -async fn missing_prompt_is_a_noop_that_never_embeds(entry: Value) { - let embedder = FakeEmbedder::new(&[]); - let calls = embedder.calls.clone(); - let cache = cache(Vec::new(), embedder); - let context = SemanticCacheContext::default(); - - cache.set_cache("key1", entry.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key1", &context).unwrap(), None); - cache - .async_set_cache("key1", entry, context.clone()) - .await - .unwrap(); - assert_eq!(cache.async_get_cache("key1", &context).await.unwrap(), None); - assert!(calls.lock().unwrap().is_empty()); -} - -#[rstest] -fn scope_overrides_key_as_filter_tag(entry: Value, context: SemanticCacheContext) { - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - hset(INDEX, PROMPT, "scope-a", &VECTOR, &entry), - search( - INDEX, - "scope\\-a", - &VECTOR, - Ok(hit("scope-a", "0.05", encoded(&entry))), - ), - ], - FakeEmbedder::new(&[]), - ); - let context = SemanticCacheContext { - scope: Some("scope-a".into()), - ..context - }; - - cache.set_cache("key1", entry.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(entry)); -} - -#[rstest] -#[case::unscoped_schema(unscoped_info(3))] -#[case::wrong_distance_metric(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2")))] -#[case::wrong_data_type(info_with_vector(vector_attribute_with(3, "FLOAT64", "COSINE")))] -fn incompatible_schema_falls_back_to_isolated_index( - #[case] base_info: redis::Value, - entry: Value, - context: SemanticCacheContext, -) { - let isolated = format!("{INDEX}_isolated"); - let cache = cache( - vec![ - info(INDEX, base_info), - info_missing(&isolated), - create_index(&isolated, 3), - hset(&isolated, PROMPT, "key1", &VECTOR, &entry), - ], - FakeEmbedder::new(&[]), - ); - - cache.set_cache("key1", entry, &context).unwrap(); -} - -#[rstest] -fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticCacheContext) { - let cache = cache( - vec![ - info_missing(INDEX), - MockCmd::new( - create_index_command(INDEX, 3), - Err::<&str, _>(support::server_error("Index already exists")), - ), - info(INDEX, compatible_info(3)), - hset(INDEX, PROMPT, "key1", &VECTOR, &entry), - ], - FakeEmbedder::new(&[]), - ); - - cache.set_cache("key1", entry, &context).unwrap(); -} - -#[rstest] -#[case::punctuation_and_spaces("a:b, c|d", "a\\:b\\,\\ c\\|d")] -#[case::braces_and_dots("{x}.y", "\\{x\\}\\.y")] -#[case::plain("key1", "key1")] -fn tag_special_characters_are_escaped_in_search_filter( - #[case] tag: &str, - #[case] escaped: &str, - context: SemanticCacheContext, -) { - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - search(INDEX, escaped, &VECTOR, Ok(empty_result())), - ], - FakeEmbedder::new(&[]), - ); - - assert_eq!(cache.get_cache(tag, &context).unwrap(), None); -} - -#[rstest] -#[case::content_parts( - SemanticCacheContext { - messages: Some(json!([ - {"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}, - {"role": "assistant", "content": "reply"}, - ])), - ..Default::default() - }, - "firstsecondreply" -)] -#[case::responses_string_input( - SemanticCacheContext { input: Some(json!(" plain input ")), ..Default::default() }, - "plain input" -)] -#[case::responses_nested_input( - SemanticCacheContext { - input: Some(json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"])), - ..Default::default() - }, - "nested\ntail" -)] -#[case::responses_output_text( - SemanticCacheContext { input: Some(json!({"output_text": " result text "})), ..Default::default() }, - "result text" -)] -#[case::search_results( - messages_context(vec![json!({ - "role": "user", - "content": "question", - "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], - })]), - "questionsrctfound{\"a\":1}" -)] -#[case::empty_messages_fall_back_to_input( - SemanticCacheContext { messages: Some(json!([])), input: Some(json!("fallback")), ..Default::default() }, - "fallback" -)] -fn prompt_extraction_matches_python_message_and_input_shapes( - #[case] context: SemanticCacheContext, - #[case] prompt: &str, -) { - let embedder = FakeEmbedder::new(&[]); - let calls = embedder.calls.clone(); - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - search(INDEX, "key1", &VECTOR, Ok(empty_result())), - ], - embedder, - ); - - cache.get_cache("key1", &context).unwrap(); - - assert_eq!(*calls.lock().unwrap(), vec![(prompt.to_owned(), None)]); -} - -#[rstest] -#[case(None)] -#[case(Some(Duration::from_secs(9)))] -fn ttl_passes_through_context_only(#[case] ttl: Option) { - let cache = cache(Vec::new(), FakeEmbedder::new(&[])); - - assert_eq!( - cache.get_ttl(&SemanticCacheContext { - ttl, - ..Default::default() - }), - ttl - ); -} - -#[rstest] -#[tokio::test] -async fn async_paths_embed_with_metadata_then_run_blocking_redis_work( - entry: Value, - context: SemanticCacheContext, -) { - let embedder = FakeEmbedder::new(&[]); - let calls = embedder.calls.clone(); - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - hset(INDEX, PROMPT, "key1", &VECTOR, &entry), - search( - INDEX, - "key1", - &VECTOR, - Ok(hit("key1", "0.05", encoded(&entry))), - ), - ], - embedder, - ); - let context = SemanticCacheContext { - metadata: Some(json!({"tenant": "team"})), - ..context - }; - - cache - .async_set_cache("key1", entry.clone(), context.clone()) - .await - .unwrap(); - assert_eq!( - cache.async_get_cache("key1", &context).await.unwrap(), - Some(entry) - ); - assert_eq!( - *calls.lock().unwrap(), - vec![(PROMPT.to_owned(), context.metadata.clone()); 2] - ); -} - -#[rstest] -fn accessors_report_the_config(config: RedisSemanticConfig) { - let cache = cache(Vec::new(), FakeEmbedder::new(&[])); - - assert_eq!(cache.index_name(), config.index_name); - assert!((cache.similarity_threshold() - config.similarity_threshold).abs() < 1e-6); -} - -#[rstest] -fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) { - // Pins parity with Python's `_isolated` + overwrite=True flow. - let prompt = "shared prompt"; - let isolated = format!("{INDEX}_isolated"); - let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); - - let vector_a = vec![0.1f32; 8]; - let worker_a = cache( - vec![ - info_missing(INDEX), - create_index(INDEX, 8), - hset(INDEX, prompt, "key1", &vector_a, &entry), - ], - FakeEmbedder::new(&[(prompt, &vector_a)]), - ); - worker_a - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - - let vector_b = vec![0.2f32; 4]; - let worker_b = cache( - vec![ - info(INDEX, compatible_info(8)), - info_missing(&isolated), - create_index(&isolated, 4), - hset(&isolated, prompt, "key1", &vector_b, &entry), - search( - &isolated, - "key1", - &vector_b, - Ok(hit("key1", "0.0", encoded(&entry))), - ), - search( - &isolated, - "key1", - &vector_b, - Err(redis::RedisError::from(( - redis::ErrorKind::Extension, - "Vector dimension mismatch", - ))), - ), - ], - FakeEmbedder::new(&[(prompt, &vector_b)]), - ); - worker_b - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - assert_eq!( - worker_b.get_cache("key1", &context()).unwrap(), - Some(entry.clone()) - ); - - let vector_c = vec![0.3f32; 16]; - let worker_c = cache( - vec![ - info(INDEX, compatible_info(8)), - info(&isolated, compatible_info(4)), - MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")), - create_index(&isolated, 16), - hset(&isolated, prompt, "key1", &vector_c, &entry), - ], - FakeEmbedder::new(&[(prompt, &vector_c)]), - ); - worker_c - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - - assert!(matches!( - worker_b.get_cache("key1", &context()), - Err(Error::Unavailable(Some(_))) - )); -} - -#[fixture] -fn redis_stack_url() -> Option { - std::env::var("LITELLM_REDIS_STACK_URL").ok() -} - -fn live_cache( - url: &str, - index_name: &str, - prompt: &str, - vector: Vec, -) -> RedisSemanticCache> { - RedisSemanticCache::new( - url, - FakeEmbedder::new(&[(prompt, vector.as_slice())]), - JsonCodec::::new(), - RedisSemanticConfig { - index_name: index_name.to_owned(), - similarity_threshold: 0.9, - }, - ) - .unwrap() -} - -#[rstest] -fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option, entry: Value) { - let Some(url) = redis_stack_url else { - return; - }; - // Pins parity with Python's `_isolated` + overwrite=True flow. - let base = format!("rust_semantic_shared_{}", std::process::id()); - let isolated = format!("{base}_isolated"); - let prompt = "shared live prompt"; - let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); - - let worker_a = live_cache(&url, &base, prompt, vec![0.1f32; 8]); - worker_a - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - - let worker_b = live_cache(&url, &base, prompt, vec![0.2f32; 4]); - worker_b - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - assert_eq!( - worker_b.get_cache("key1", &context()).unwrap(), - Some(entry.clone()) - ); - - let worker_c = live_cache(&url, &base, prompt, vec![0.3f32; 16]); - worker_c - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - - assert!(matches!( - worker_b.get_cache("key1", &context()), - Err(Error::Unavailable(Some(_))) - )); - - let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); - for index in [&base, &isolated] { - let _: Result<(), _> = redis::cmd("FT.DROPINDEX") - .arg(index) - .arg("DD") - .query(&mut connection); - } -} - -#[rstest] -fn live_store_lookup_and_ttl_against_redis_stack(redis_stack_url: Option, entry: Value) { - let Some(url) = redis_stack_url else { - return; - }; - let prompt = "rust semantic cache live prompt"; - let index_name = format!("rust_semantic_test_{}", std::process::id()); - let cache = live_cache(&url, &index_name, prompt, vec![0.1, 0.2, 0.3, 0.4]); - let context = SemanticCacheContext { - ttl: Some(Duration::from_secs(120)), - ..messages_context(vec![json!({"role": "user", "content": prompt})]) - }; - - cache - .set_cache("live-key", entry.clone(), &context) - .unwrap(); - assert_eq!(cache.get_cache("live-key", &context).unwrap(), Some(entry)); - assert_eq!(cache.get_cache("other-key", &context).unwrap(), None); - - let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); - let ttl: i64 = redis::Commands::ttl( - &mut connection, - format!("{index_name}:{}", entry_id(prompt, "live-key")), - ) - .unwrap(); - assert!( - ttl > 0, - "expected stored hash to carry an expiry, got {ttl}" - ); -} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/contract.rs b/litellm-rust/crates/cache-redis-semantic/tests/contract.rs deleted file mode 100644 index fb3cda9616e..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/tests/contract.rs +++ /dev/null @@ -1,63 +0,0 @@ -mod support; - -use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding}; -use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig}; -use litellm_cache_testing as contract; -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; -use support::FakeSearch; - -type Cache = RedisSemanticCache, FakeSearch>; - -const PREFIX: &str = "contract:"; - -#[fixture] -fn cache() -> Cache { - RedisSemanticCache::with_connection( - FakeSearch::default(), - PreparedEmbedding(vec![0.6, 0.8]), - JsonCodec::new(), - RedisSemanticConfig { - index_name: DEFAULT_INDEX_NAME.into(), - similarity_threshold: 0.9, - }, - ) -} - -#[fixture] -fn context() -> SemanticCacheContext { - SemanticCacheContext { - messages: Some(json!([{"role": "user", "content": "contract prompt"}])), - ..Default::default() - } -} - -#[rstest] -#[tokio::test] -async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) { - contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await; -} - -#[rstest] -#[tokio::test] -async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) { - contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await; -} - -#[rstest] -#[tokio::test] -async fn overwrite_replaces(cache: Cache, context: SemanticCacheContext) { - contract::overwrite_replaces(&cache, context, PREFIX, json!(1), json!({"b": 2})).await; -} - -#[rstest] -#[tokio::test] -async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) { - contract::pipeline_writes_every_entry( - &cache, - context, - PREFIX, - vec![json!("a"), json!(2), json!({"c": true})], - ) - .await; -} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs deleted file mode 100644 index f112dbe96a7..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs +++ /dev/null @@ -1,302 +0,0 @@ -#![allow(dead_code)] - -use std::{ - collections::{BTreeMap, HashMap}, - sync::{Arc, Mutex}, -}; - -use litellm_cache::{Error, semantic::Embedder}; -use serde_json::Value; - -pub type EmbedCalls = Arc)>>>; - -/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every -/// prompt with its metadata. -pub struct FakeEmbedder { - vectors: HashMap>, - pub calls: EmbedCalls, -} - -impl FakeEmbedder { - pub fn new(vectors: &[(&str, &[f32])]) -> Self { - Self { - vectors: vectors - .iter() - .map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec())) - .collect(), - calls: EmbedCalls::default(), - } - } -} - -impl Embedder for FakeEmbedder { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - self.calls - .lock() - .unwrap() - .push((prompt.to_owned(), metadata.cloned())); - Ok(self - .vectors - .get(prompt) - .cloned() - .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) - } - - async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - self.embed(prompt, metadata) - } -} - -struct FakeIndex { - prefix: Vec, - dims: usize, - vector_field: String, -} - -#[derive(Default)] -struct SearchState { - indexes: HashMap, - hashes: BTreeMap, BTreeMap>>, -} - -/// An in-memory Redis Stack speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache -/// sends, with exact cosine KNN over the hashes under an index prefix. -#[derive(Clone, Default)] -pub struct FakeSearch { - state: Arc>, -} - -impl FakeSearch { - fn run(&self, args: Vec>) -> redis::RedisResult { - let mut state = self.state.lock().unwrap(); - let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned(); - match text(0).to_uppercase().as_str() { - "FT.CREATE" => { - let name = text(1); - if state.indexes.contains_key(&name) { - return Err(server_error("Index already exists")); - } - let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes()); - let prefix = args[position("PREFIX").unwrap() + 2].clone(); - let dims = text(position("DIM").unwrap() + 1).parse().unwrap(); - let vector_field = text(position("VECTOR").unwrap() - 1); - state.indexes.insert( - name, - FakeIndex { - prefix, - dims, - vector_field, - }, - ); - Ok(redis::Value::Okay) - } - "FT.INFO" => { - let index = state - .indexes - .get(&text(1)) - .ok_or_else(|| server_error("Unknown index name"))?; - Ok(index_info(index)) - } - "FT.DROPINDEX" => { - state.indexes.remove(&text(1)); - Ok(redis::Value::Okay) - } - "HSET" => { - let hash = state.hashes.entry(args[1].clone()).or_default(); - for pair in args[2..].chunks(2) { - hash.insert( - String::from_utf8_lossy(&pair[0]).into_owned(), - pair[1].clone(), - ); - } - Ok(redis::Value::Int(((args.len() - 2) / 2) as i64)) - } - "EXPIRE" => Ok(redis::Value::Int(i64::from( - state.hashes.contains_key(&args[1]), - ))), - "FT.SEARCH" => { - let index = state - .indexes - .get(&text(1)) - .ok_or_else(|| server_error("no such index"))?; - let query = text(2); - let tag = query_tag(&query); - let params = args.iter().position(|arg| arg == b"PARAMS").unwrap(); - let vector = floats(&args[params + 3]); - let best = state - .hashes - .iter() - .filter(|(key, _)| key.starts_with(&index.prefix)) - .filter(|(_, fields)| { - fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes()) - }) - .filter_map(|(key, fields)| { - let stored = floats(fields.get(&index.vector_field)?); - (stored.len() == index.dims) - .then(|| (key, fields, 1.0 - cosine(&vector, &stored))) - }) - .min_by(|left, right| left.2.total_cmp(&right.2)); - let Some((key, fields, distance)) = best else { - return Ok(redis::Value::Array(vec![redis::Value::Int(0)])); - }; - let mut reply = fields - .iter() - .filter(|(name, _)| **name != index.vector_field) - .flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)]) - .collect::>(); - reply.extend([ - bulk(b"vector_distance"), - bulk(distance.to_string().as_bytes()), - ]); - Ok(redis::Value::Array(vec![ - redis::Value::Int(1), - bulk(key), - redis::Value::Array(reply), - ])) - } - "PING" => Ok(redis::Value::SimpleString("PONG".into())), - _ => Err(server_error("unsupported command")), - } - } -} - -impl redis::ConnectionLike for FakeSearch { - fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { - let mut commands = parse_commands(command); - self.run(commands.remove(0)) - } - - fn req_packed_commands( - &mut self, - commands: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - let replies = parse_commands(commands) - .into_iter() - .map(|args| self.run(args)) - .collect::>>()?; - Ok(replies.into_iter().skip(offset).take(count).collect()) - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} - -pub fn server_error(reply: &str) -> redis::RedisError { - redis::parse_redis_value(format!("-{reply}\r\n").as_bytes()) - .expect("error reply") - .extract_error() - .expect_err("server error") -} - -fn bulk(bytes: &[u8]) -> redis::Value { - redis::Value::BulkString(bytes.to_vec()) -} - -fn index_info(index: &FakeIndex) -> redis::Value { - let attribute = |name: &str, field_type: &str| { - redis::Value::Array(vec![ - bulk(b"identifier"), - bulk(name.as_bytes()), - bulk(b"type"), - bulk(field_type.as_bytes()), - ]) - }; - redis::Value::Array(vec![ - bulk(b"attributes"), - redis::Value::Array(vec![ - attribute("prompt", "TEXT"), - attribute("response", "TEXT"), - attribute("inserted_at", "NUMERIC"), - attribute("updated_at", "NUMERIC"), - attribute("litellm_cache_key", "TAG"), - redis::Value::Array(vec![ - bulk(b"identifier"), - bulk(index.vector_field.as_bytes()), - bulk(b"type"), - bulk(b"VECTOR"), - bulk(b"dim"), - redis::Value::Int(index.dims as i64), - bulk(b"data_type"), - bulk(b"FLOAT32"), - bulk(b"distance_metric"), - bulk(b"COSINE"), - ]), - ]), - ]) -} - -/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed. -fn query_tag(query: &str) -> String { - let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len(); - let mut tag = String::new(); - let mut characters = query[start..].chars(); - while let Some(character) = characters.next() { - match character { - '\\' => tag.extend(characters.next()), - '}' => break, - character => tag.push(character), - } - } - tag -} - -fn floats(bytes: &[u8]) -> Vec { - bytes - .as_chunks::<4>() - .0 - .iter() - .map(|chunk| f32::from_le_bytes(*chunk)) - .collect() -} - -fn cosine(left: &[f32], right: &[f32]) -> f64 { - let dot = left - .iter() - .zip(right) - .map(|(left, right)| f64::from(*left) * f64::from(*right)) - .sum::(); - let norm = |vector: &[f32]| { - vector - .iter() - .map(|value| f64::from(*value).powi(2)) - .sum::() - .sqrt() - }; - dot / (norm(left) * norm(right)) -} - -/// Splits a packed RESP request into each command's arguments. -fn parse_commands(mut bytes: &[u8]) -> Vec>> { - let line = |bytes: &mut &[u8]| { - let end = bytes - .windows(2) - .position(|window| window == b"\r\n") - .unwrap(); - let text = String::from_utf8(bytes[1..end].to_vec()).unwrap(); - *bytes = &bytes[end + 2..]; - text.parse::().unwrap() - }; - let mut commands = Vec::new(); - while !bytes.is_empty() { - let count = line(&mut bytes); - let mut args = Vec::with_capacity(count); - for _ in 0..count { - let length = line(&mut bytes); - args.push(bytes[..length].to_vec()); - bytes = &bytes[length + 2..]; - } - commands.push(args); - } - commands -} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml deleted file mode 100644 index f7621299506..00000000000 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ /dev/null @@ -1,20 +0,0 @@ -[package] -name = "litellm-cache-redis" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true - -[dependencies] -litellm-cache.workspace = true -litellm-error.workspace = true -redis = { workspace = true, features = ["cluster", "tls-rustls"] } -r2d2 = "0.8.10" -thiserror.workspace = true -tokio.workspace = true - -[dev-dependencies] -litellm-cache-testing.workspace = true -redis-test.workspace = true -rstest.workspace = true -serde_json.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs deleted file mode 100644 index b873a2c0282..00000000000 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ /dev/null @@ -1,192 +0,0 @@ -use std::{ - sync::{Arc, OnceLock}, - time::Duration, -}; - -use litellm_cache::{BatchEntry, CacheCodec, Error}; - -use crate::{ - connection::{ConnectionRef, Connections}, - topology::RedisTopology, -}; - -const DEFAULT_TTL: Duration = Duration::from_secs(600); - -pub struct RedisCache { - pub(crate) connections: Arc>, - pub(crate) default_ttl: Duration, - pub(crate) codec: S, - pub(crate) namespace: Option, - pub(crate) topology: RedisTopology, - /// The server's major version, read from `INFO` once, like Python's `redis_version`. - pub(crate) major_version: Arc>, -} - -impl RedisCache { - pub fn new(url: &str, default_ttl: Option, codec: S) -> Result { - Self::connect(url, &RedisTopology::Standalone, default_ttl, codec) - } - - pub fn connect( - url: &str, - topology: &RedisTopology, - default_ttl: Option, - codec: S, - ) -> Result { - let connections = Connections::open(url, topology)?; - Ok(Self { - connections: Arc::new(connections), - default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), - codec, - namespace: None, - topology: topology.clone(), - major_version: Arc::default(), - }) - } -} - -impl RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { - Self { - connections: Arc::new(Connections::fixed(connection)), - default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), - codec, - namespace: None, - topology: RedisTopology::Standalone, - major_version: Arc::default(), - } - } - - pub fn with_namespace(self, namespace: Option) -> Self { - Self { - namespace: namespace.filter(|value| !value.is_empty()), - ..self - } - } - - pub fn namespace(&self) -> Option<&str> { - self.namespace.as_deref() - } - - pub fn topology(&self) -> &RedisTopology { - &self.topology - } - - pub(crate) fn namespaced_key(&self, key: &str) -> String { - namespaced_key(self.namespace.as_deref(), key) - } - - pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec { - keys.iter().map(|key| self.namespaced_key(key)).collect() - } - - /// Whole seconds for `ttl`, falling back to the default TTL like Python's `get_ttl`. - pub(crate) fn ttl_or_default(&self, ttl: Option) -> u64 { - ttl_seconds(ttl.unwrap_or(self.default_ttl)) - } - - pub(crate) fn execute( - &self, - operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, - ) -> Result { - self.connections.execute(operation) - } - - /// `_parse_redis_major_version`: the major version from `INFO`, or - /// `DEFAULT_REDIS_MAJOR_VERSION` when `INFO` fails or its version does not parse. The first - /// answer is kept, as Python reads `redis_version` once at construction. - pub(crate) async fn major_version(&self) -> u32 { - if let Some(version) = self.major_version.get() { - return *version; - } - let info = self - .run(|connection| connection.node_text(&redis::cmd("INFO"))) - .await; - let version = info - .ok() - .and_then(|info| parse_major_version(&info)) - .unwrap_or_else(default_major_version); - *self.major_version.get_or_init(|| version) - } - - pub(crate) async fn run(&self, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, - { - Connections::run_blocking(Arc::clone(&self.connections), operation).await - } - - pub(crate) fn namespaced_pattern(&self) -> Result { - let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?; - let escaped: String = namespace - .chars() - .flat_map(|ch| { - if matches!(ch, '*' | '?' | '[' | ']' | '\\') { - vec!['\\', ch] - } else { - vec![ch] - } - }) - .collect(); - Ok(format!("{escaped}:*")) - } - - pub(crate) fn decode_response(&self, value: redis::Value) -> Result, Error> { - match value { - redis::Value::Nil => Ok(None), - redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some), - redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some), - _ => Err(Error::InvalidEntry(None)), - } - } - - pub(crate) fn decode_batch_response( - &self, - value: redis::Value, - ) -> Result, Error> { - match self.decode_response(value) { - Ok(Some(value)) => Ok(BatchEntry::Hit(value)), - Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), - Err(error) => Err(error), - } - } -} - -pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String { - match namespace { - Some(namespace) if !key.starts_with(&format!("{namespace}:")) => { - format!("{namespace}:{key}") - } - _ => key.into(), - } -} - -pub(crate) fn ttl_seconds(ttl: Duration) -> u64 { - ttl.as_secs() - .saturating_add(u64::from(ttl.subsec_nanos() > 0)) - .max(1) -} - -fn parse_major_version(info: &str) -> Option { - let version = info - .lines() - .find_map(|line| line.trim().strip_prefix("redis_version:"))? - .trim(); - match version.split_once('.') { - Some((major, _)) => major.parse().ok(), - None => version.parse::().ok().map(|major| major as u32), - } -} - -fn default_major_version() -> u32 { - std::env::var("DEFAULT_REDIS_MAJOR_VERSION") - .ok() - .and_then(|value| value.parse().ok()) - .unwrap_or(7) -} diff --git a/litellm-rust/crates/cache-redis/src/claim.rs b/litellm-rust/crates/cache-redis/src/claim.rs deleted file mode 100644 index c57a1b03395..00000000000 --- a/litellm-rust/crates/cache-redis/src/claim.rs +++ /dev/null @@ -1,109 +0,0 @@ -use litellm_cache::{CacheCodec, ClaimCache, Error, ExactCacheContext}; -use redis::Commands; - -use crate::{cache::RedisCache, connection::ConnectionRef}; - -const CLAIM_SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", - "elseif current ~= ARGV[1] then return 0; end; ", - "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ", - "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" -); -const CLAIM_ATTEMPTS: usize = 8; - -fn stored_bytes(value: redis::Value) -> Result>, Error> { - match value { - redis::Value::Nil => Ok(None), - redis::Value::BulkString(bytes) => Ok(Some(bytes)), - redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())), - _ => Err(Error::InvalidEntry(None)), - } -} - -/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's -/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the -/// bytes that decision was made on, retried when another claimant wins the race. -fn claim( - connection: &mut ConnectionRef<'_>, - codec: &S, - key: &str, - candidate: S::Value, - eligible: &[S::Value], - ttl: u64, -) -> Result -where - S::Value: PartialEq, -{ - let payload = codec.encode(&candidate)?; - if payload.is_empty() { - return Err(Error::InvalidEntry(None)); - } - for _ in 0..CLAIM_ATTEMPTS { - let current = stored_bytes( - connection - .get::<_, redis::Value>(key) - .map_err(Error::unavailable)?, - )? - .filter(|bytes| !bytes.is_empty()); - let existing = current - .as_deref() - .and_then(|bytes| codec.decode(bytes).ok()) - .filter(|existing| eligible.is_empty() || eligible.contains(existing)); - let refresh = existing - .as_ref() - .is_some_and(|existing| !eligible.is_empty() || *existing == candidate); - let write: &[u8] = if existing.is_some() { b"" } else { &payload }; - let applied = redis::cmd("EVAL") - .arg(CLAIM_SCRIPT) - .arg(1) - .arg(key) - .arg(current.as_deref().unwrap_or_default()) - .arg(ttl) - .arg(write) - .arg(u8::from(refresh)) - .query::(connection) - .map_err(Error::unavailable)?; - if applied { - return Ok(existing.unwrap_or(candidate)); - } - } - Err(Error::unavailable(ClaimContended)) -} - -#[derive(Debug, thiserror::Error)] -#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")] -struct ClaimContended; - -impl ClaimCache for RedisCache -where - S: CacheCodec + Clone + 'static, - S::Value: PartialEq, - C: redis::ConnectionLike + Send + 'static, -{ - fn claim_cache( - &self, - key: &str, - candidate: S::Value, - eligible: &[S::Value], - context: ExactCacheContext, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - self.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl)) - } - - async fn async_claim_cache( - &self, - key: &str, - candidate: S::Value, - eligible: Vec, - context: ExactCacheContext, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - let codec = self.codec.clone(); - self.run(move |connection| claim(connection, &codec, &key, candidate, &eligible, ttl)) - .await - } -} diff --git a/litellm-rust/crates/cache-redis/src/connection.rs b/litellm-rust/crates/cache-redis/src/connection.rs deleted file mode 100644 index 99befa16ffa..00000000000 --- a/litellm-rust/crates/cache-redis/src/connection.rs +++ /dev/null @@ -1,513 +0,0 @@ -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; - -use litellm_cache::Error; -use redis::{ - ConnectionAddr, ConnectionInfo, IntoConnectionInfo, - cluster::{ - ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress, - }, - cluster_routing::{ - MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo, - }, -}; - -use crate::topology::{RedisNode, RedisTopology}; - -pub(crate) const REDIS_TIMEOUT: Duration = Duration::from_secs(5); -const REDIS_POOL_SIZE: u32 = 16; - -#[allow(private_interfaces)] -pub enum Connections { - Pool(r2d2::Pool), - Cluster(r2d2::Pool), - Fixed(Mutex), -} - -impl Connections -where - C: redis::ConnectionLike + Send + 'static, -{ - pub fn execute( - &self, - operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, - ) -> Result { - match self { - Self::Pool(pool) => { - let mut pooled = pool.get().map_err(Error::unavailable)?; - let result = operation(&mut ConnectionRef::Node(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable(_))); - result - } - Self::Cluster(pool) => { - let mut pooled = pool.get().map_err(Error::unavailable)?; - let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable(_))); - result - } - Self::Fixed(connection) => { - let mut connection = connection.lock()?; - operation(&mut ConnectionRef::Node(&mut *connection)) - } - } - } - - pub async fn run_blocking(connections: Arc, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, - { - tokio::task::spawn_blocking(move || connections.execute(operation)) - .await - .map_err(Error::unavailable)? - } - - pub fn fixed(connection: C) -> Self { - Self::Fixed(Mutex::new(connection)) - } - - pub fn open(url: &str, topology: &RedisTopology) -> Result { - match topology { - RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)), - RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool( - ClusterConnectionManager::open(url, startup_nodes)?, - )?)), - } - } - - /// Closes every idle pooled connection; the next operation opens a fresh one. Connections - /// checked out right now return to the pool, and a caller-owned connection stays open. - pub fn disconnect(&self) { - match self { - Self::Pool(pool) => close_idle(pool), - Self::Cluster(pool) => close_idle(pool), - Self::Fixed(_) => {} - } - } -} - -fn pool(manager: M) -> Result, Error> { - r2d2::Pool::builder() - .max_size(REDIS_POOL_SIZE) - .min_idle(Some(0)) - .connection_timeout(REDIS_TIMEOUT) - .test_on_check_out(false) - .build(manager) - .map_err(Error::unavailable) -} - -fn close_idle(pool: &r2d2::Pool) -where - M: r2d2::ManageConnection>, -{ - let mut idle = Vec::new(); - while let Some(mut connection) = pool.try_get() { - connection.failed = true; - idle.push(connection); - } -} - -pub(crate) struct PooledConnection { - connection: C, - failed: bool, -} - -/// Pools connections without a checkout PING, which would double every operation's round trips. -/// A timed-out command leaves its reply on the socket while redis still reports the connection -/// open, so any connection whose operation failed is discarded instead of being reused. -pub(crate) struct ConnectionManager(redis::Client); - -impl ConnectionManager { - fn open(url: &str) -> Result { - redis::Client::open(url) - .map(Self) - .map_err(Error::unavailable) - } -} - -impl r2d2::ManageConnection for ConnectionManager { - type Connection = PooledConnection; - type Error = redis::RedisError; - - fn connect(&self) -> Result { - let connection = self.0.get_connection()?; - connection.set_read_timeout(Some(REDIS_TIMEOUT))?; - connection.set_write_timeout(Some(REDIS_TIMEOUT))?; - Ok(PooledConnection { - connection, - failed: false, - }) - } - - fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> { - redis::cmd("PING").query::(&mut connection.connection)?; - Ok(()) - } - - fn has_broken(&self, connection: &mut Self::Connection) -> bool { - connection.failed || !redis::ConnectionLike::is_open(&connection.connection) - } -} - -#[derive(Debug, thiserror::Error)] -enum TopologyFault { - #[error("Redis Cluster needs at least one startup node")] - NoStartupNodes, - #[error("Redis Cluster nodes need a TCP address")] - NonTcpAddress, - #[error("Redis Cluster answered a fan-out command without a per-node map")] - FanOutReply, - #[error("scan target does not match the connection topology")] - ScanTarget, -} - -pub(crate) struct ClusterConnectionManager(ClusterClient); - -impl ClusterConnectionManager { - fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { - if startup_nodes.is_empty() { - return Err(Error::unavailable(TopologyFault::NoStartupNodes)); - } - let info = url.into_connection_info().map_err(Error::unavailable)?; - let nodes = startup_nodes - .iter() - .map(|node| node_info(&info, node)) - .collect::, _>>()?; - ClusterClientBuilder::new(nodes) - .connection_timeout(REDIS_TIMEOUT) - .response_timeout(REDIS_TIMEOUT) - .build() - .map(Self) - .map_err(Error::unavailable) - } -} - -fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result { - let addr = match info.addr() { - ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port), - ConnectionAddr::TcpTls { - insecure, - tls_params, - .. - } => ConnectionAddr::TcpTls { - host: node.host.clone(), - port: node.port, - insecure: *insecure, - tls_params: tls_params.clone(), - }, - _ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)), - }; - Ok(info.clone().set_addr(addr)) -} - -impl r2d2::ManageConnection for ClusterConnectionManager { - type Connection = PooledConnection; - type Error = redis::RedisError; - - fn connect(&self) -> Result { - let connection = self.0.get_connection()?; - connection.set_read_timeout(Some(REDIS_TIMEOUT))?; - connection.set_write_timeout(Some(REDIS_TIMEOUT))?; - Ok(PooledConnection { - connection, - failed: false, - }) - } - - fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> { - redis::cmd("PING").query::(&mut connection.connection)?; - Ok(()) - } - - fn has_broken(&self, connection: &mut Self::Connection) -> bool { - connection.failed || !redis::ConnectionLike::is_open(&connection.connection) - } -} - -pub enum ConnectionRef<'a> { - Node(&'a mut dyn redis::ConnectionLike), - Cluster(&'a mut ClusterConnection), -} - -impl redis::ConnectionLike for ConnectionRef<'_> { - fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { - match self { - Self::Node(connection) => connection.req_packed_command(cmd), - Self::Cluster(connection) => connection.req_packed_command(cmd), - } - } - - fn req_packed_commands( - &mut self, - cmd: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - match self { - Self::Node(connection) => connection.req_packed_commands(cmd, offset, count), - Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count), - } - } - - fn get_db(&self) -> i64 { - match self { - Self::Node(connection) => connection.get_db(), - Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection), - } - } - - fn supports_pipelining(&self) -> bool { - match self { - Self::Node(connection) => connection.supports_pipelining(), - Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection), - } - } - - fn check_connection(&mut self) -> bool { - match self { - Self::Node(connection) => connection.check_connection(), - Self::Cluster(connection) => connection.check_connection(), - } - } - - fn is_open(&self) -> bool { - match self { - Self::Node(connection) => connection.is_open(), - Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection), - } - } -} - -impl ConnectionRef<'_> { - /// Runs `pipeline` and decodes its non-ignored replies as `T`. A cluster connection refuses - /// `Pipeline::query`, so there a transaction goes to its keys' slot as one MULTI/EXEC and - /// anything else is split per node by `ClusterPipeline`; either way the raw replies are - /// handed back to `pipeline` to decode. - pub(crate) fn query_pipeline( - &mut self, - pipeline: &redis::Pipeline, - ) -> Result { - match self { - Self::Node(connection) => pipeline.query(*connection), - Self::Cluster(connection) if pipeline.is_transaction() => { - redis::ConnectionLike::req_packed_commands( - *connection, - &pipeline.get_packed_pipeline(), - pipeline.len() + 1, - 1, - ) - .and_then(|replies| pipeline.query(&mut Replies(Some(replies)))) - } - Self::Cluster(connection) => { - let mut cluster = ClusterPipeline::with_capacity(pipeline.len()); - for command in pipeline.cmd_iter() { - cluster.add_command(command.clone()); - } - cluster - .query(connection) - .and_then(|replies| pipeline.query(&mut Replies(Some(replies)))) - } - } - .map_err(Error::unavailable) - } - - pub(crate) fn scan( - &mut self, - pattern: &str, - count: usize, - mut visit: impl FnMut(&mut Self, Vec) -> Result, - ) -> Result<(), Error> { - let pages = match self { - Self::Node(connection) => { - let page = scan_command(0, pattern, count) - .query::(*connection) - .map_err(Error::unavailable)?; - vec![(None, page)] - } - Self::Cluster(connection) => connection - .route_command( - &scan_command(0, pattern, count), - RoutingInfo::MultiNode(( - MultipleNodeRoutingInfo::AllMasters, - Some(ResponsePolicy::Special), - )), - ) - .map_err(Error::unavailable) - .and_then(primary_pages)? - .into_iter() - .map(|(node, page)| (Some(node), page)) - .collect(), - }; - for (node, (mut cursor, mut keys)) in pages { - loop { - if !visit(self, keys)? { - return Ok(()); - } - if cursor == 0 { - break; - } - (cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?; - } - } - Ok(()) - } - - pub(crate) fn ping(&mut self) -> Result { - let command = redis::cmd("PING"); - match self { - Self::Node(connection) => command - .query::(*connection) - .map(|response| response == "PONG"), - Self::Cluster(connection) => connection - .route_command( - &command, - RoutingInfo::MultiNode(( - MultipleNodeRoutingInfo::AllNodes, - Some(ResponsePolicy::AllSucceeded), - )), - ) - .map(|_| true), - } - } - - pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result { - match self { - Self::Node(connection) => command.query(*connection).map_err(Error::unavailable), - Self::Cluster(connection) => { - let value = connection - .route_command( - command, - RoutingInfo::MultiNode(( - MultipleNodeRoutingInfo::AllNodes, - Some(ResponsePolicy::Special), - )), - ) - .map_err(Error::unavailable)?; - let redis::Value::Map(entries) = value else { - return Err(Error::unavailable(TopologyFault::FanOutReply)); - }; - let mut replies = entries - .into_iter() - .map(|(node, reply)| { - Ok(( - redis::from_redis_value::(node).map_err(Error::unavailable)?, - redis::from_redis_value::(reply).map_err(Error::unavailable)?, - )) - }) - .collect::, Error>>()?; - replies.sort(); - Ok(replies - .into_iter() - .map(|(_, reply)| reply) - .collect::>() - .join("\n")) - } - } - } - - pub(crate) fn flushall(&mut self) -> Result<(), Error> { - let command = redis::cmd("FLUSHALL"); - match self { - Self::Node(connection) => command.query(*connection).map_err(Error::unavailable), - Self::Cluster(connection) => connection - .route_command( - &command, - RoutingInfo::MultiNode(( - MultipleNodeRoutingInfo::AllMasters, - Some(ResponsePolicy::AllSucceeded), - )), - ) - .map(|_| ()) - .map_err(Error::unavailable), - } - } - - fn scan_page( - &mut self, - node: Option<&NodeAddress>, - cursor: u64, - pattern: &str, - count: usize, - ) -> Result { - let command = scan_command(cursor, pattern, count); - match (self, node) { - (Self::Node(connection), None) => { - command.query(*connection).map_err(Error::unavailable) - } - (Self::Cluster(connection), Some(node)) => connection - .route_command( - &command, - RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress { - host: node.host().to_string(), - port: node.port(), - }), - ) - .map_err(Error::unavailable) - .and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)), - _ => Err(Error::unavailable(TopologyFault::ScanTarget)), - } - } -} - -type ScanPage = (u64, Vec); - -fn primary_pages(value: redis::Value) -> Result, Error> { - let redis::Value::Map(entries) = value else { - return Err(Error::unavailable(TopologyFault::FanOutReply)); - }; - entries - .into_iter() - .map(|(node, page)| { - let node = redis::from_redis_value::(node).map_err(Error::unavailable)?; - let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?; - let page = redis::from_redis_value::(page).map_err(Error::unavailable)?; - Ok((node, page)) - }) - .collect() -} - -fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd { - let mut command = redis::cmd("SCAN"); - command - .cursor_arg(cursor) - .arg("MATCH") - .arg(pattern) - .arg("COUNT") - .arg(count); - command -} - -/// Hands already received pipeline replies to `Pipeline::query`, so it applies its own -/// ignore and error handling to replies a cluster pipeline gathered from several nodes. -struct Replies(Option>); - -impl redis::ConnectionLike for Replies { - fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult { - Err((redis::ErrorKind::Client, "replies hold a pipeline only").into()) - } - - fn req_packed_commands( - &mut self, - _: &[u8], - _: usize, - _: usize, - ) -> redis::RedisResult> { - self.0 - .take() - .ok_or_else(|| (redis::ErrorKind::Client, "replies were already read").into()) - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} diff --git a/litellm-rust/crates/cache-redis/src/counter.rs b/litellm-rust/crates/cache-redis/src/counter.rs deleted file mode 100644 index a4beb9f5712..00000000000 --- a/litellm-rust/crates/cache-redis/src/counter.rs +++ /dev/null @@ -1,205 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{ - BoundedCounterCache, CacheCodec, CountReadCache, CounterCache, Error, ExactCacheContext, - IncrementOperation, -}; - -use crate::{ - cache::{RedisCache, ttl_seconds}, - connection::ConnectionRef, - store::mget, -}; - -const INCREMENT_SCRIPT: &str = concat!( - "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", - "if redis.call('TTL', KEYS[1]) == -1 then ", - "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" -); -const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( - "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", - "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ", - "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", - "return count" -); -const SET_MAX_SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if current == false or tonumber(current) < tonumber(ARGV[1]) then ", - "redis.call('SET', KEYS[1], ARGV[1]); ", - "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", - "return ARGV[1]; end; return current" -); - -impl CounterCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn increment_cache( - &self, - key: &str, - amount: f64, - context: ExactCacheContext, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - self.execute(|connection| increment(connection, key, amount, ttl, false)) - } - - /// Python `_incrbyfloat_with_ttl`: without `refresh_ttl` the TTL is set only on a key that - /// has none, in one atomic script; with it, every increment re-arms the TTL. - async fn async_increment( - &self, - key: &str, - amount: f64, - context: ExactCacheContext, - refresh_ttl: bool, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - self.run(move |connection| increment(connection, key, amount, ttl, refresh_ttl)) - .await - } - - async fn async_increment_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - if operations.is_empty() { - return Ok(Vec::new()); - } - let mut pipeline = redis::pipe(); - for operation in operations { - let key = self.namespaced_key(&operation.key); - pipeline.cmd("INCRBYFLOAT").arg(&key).arg(operation.amount); - if let Some(ttl) = operation.ttl { - pipeline - .cmd("EXPIRE") - .arg(key) - .arg(ttl_seconds(ttl)) - .ignore(); - } - } - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - } -} - -fn increment( - connection: &mut ConnectionRef<'_>, - key: String, - amount: f64, - ttl: u64, - refresh_ttl: bool, -) -> Result { - if !refresh_ttl { - return redis::cmd("EVAL") - .arg(INCREMENT_SCRIPT) - .arg(1) - .arg(key) - .arg(amount) - .arg(ttl) - .query(connection) - .map_err(Error::unavailable); - } - connection - .query_pipeline( - redis::pipe() - .cmd("INCRBYFLOAT") - .arg(&key) - .arg(amount) - .cmd("EXPIRE") - .arg(&key) - .arg(ttl) - .ignore(), - ) - .map(|(value,)| value) -} - -impl CountReadCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn batch_get_counts(&self, keys: &[String]) -> Result>, Error> { - let keys = self.namespaced_keys(keys); - self.execute(|connection| mget(connection, keys))? - .into_iter() - .map(count) - .collect() - } - - async fn async_batch_get_counts(&self, keys: Vec) -> Result>, Error> { - let keys = self.namespaced_keys(&keys); - self.run(move |connection| mget(connection, keys)) - .await? - .into_iter() - .map(count) - .collect() - } -} - -fn count(value: redis::Value) -> Result, Error> { - redis::from_redis_value(value).map_err(Error::invalid_entry) -} - -impl BoundedCounterCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result { - let key = self.namespaced_key(key); - let ttl = ttl_seconds(ttl); - self.execute(|connection| increment_with_floor(connection, key, amount, ttl)) - } - - async fn async_increment_with_floor( - &self, - key: &str, - amount: i64, - ttl: Duration, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = ttl_seconds(ttl); - self.run(move |connection| increment_with_floor(connection, key, amount, ttl)) - .await - } - - async fn async_set_max( - &self, - key: &str, - value: f64, - ttl: Option, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(ttl); - self.run(move |connection| { - redis::cmd("EVAL") - .arg(SET_MAX_SCRIPT) - .arg(1) - .arg(key) - .arg(value) - .arg(ttl) - .query(connection) - .map_err(Error::unavailable) - }) - .await - } -} - -fn increment_with_floor( - connection: &mut ConnectionRef<'_>, - key: String, - amount: i64, - ttl: u64, -) -> Result { - redis::cmd("EVAL") - .arg(INCREMENT_WITH_FLOOR_SCRIPT) - .arg(1) - .arg(key) - .arg(amount) - .arg(ttl) - .query(connection) - .map_err(Error::unavailable) -} diff --git a/litellm-rust/crates/cache-redis/src/keys.rs b/litellm-rust/crates/cache-redis/src/keys.rs deleted file mode 100644 index 2e8a9e60e65..00000000000 --- a/litellm-rust/crates/cache-redis/src/keys.rs +++ /dev/null @@ -1,63 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache}; - -use crate::cache::RedisCache; - -impl TtlCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_get_ttl(&self, key: &str) -> Result, Error> { - let key = self.namespaced_key(key); - let ttl = self - .run(move |connection| { - redis::cmd("TTL") - .arg(key) - .query::(connection) - .map_err(Error::unavailable) - }) - .await?; - Ok(u64::try_from(ttl).ok().map(Duration::from_secs)) - } -} - -impl RefreshTtlCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_refresh_ttl(&self, key: &str, ttl: Option) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(ttl); - self.run(move |connection| { - redis::cmd("EXPIRE") - .arg(key) - .arg(ttl) - .query(connection) - .map_err(Error::unavailable) - }) - .await - } -} - -impl ScanCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { - let pattern = format!("{}*", self.namespaced_key(pattern)); - self.run(move |connection| { - let mut matches = Vec::new(); - connection.scan(&pattern, count, |_, keys| { - matches.extend(keys); - Ok(matches.len() < count) - })?; - matches.truncate(count); - Ok(matches) - }) - .await - } -} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs deleted file mode 100644 index 037e39e5d40..00000000000 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ /dev/null @@ -1,15 +0,0 @@ -mod cache; -mod claim; -pub mod connection; -mod counter; -mod keys; -mod lifecycle; -mod queue; -mod script; -mod store; -mod topology; - -pub use cache::RedisCache; -pub use queue::{RedisLpopOperation, RedisLpopResult, RedisRpushOperation}; -pub use script::{RedisArg, RedisScript}; -pub use topology::{RedisNode, RedisTopology}; diff --git a/litellm-rust/crates/cache-redis/src/lifecycle.rs b/litellm-rust/crates/cache-redis/src/lifecycle.rs deleted file mode 100644 index 54fa8ecc3bf..00000000000 --- a/litellm-rust/crates/cache-redis/src/lifecycle.rs +++ /dev/null @@ -1,86 +0,0 @@ -use litellm_cache::{ - CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache, - DisconnectCache, Error, PingCache, -}; - -use crate::{cache::RedisCache, topology::RedisTopology}; - -impl PingCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn sync_ping(&self) -> Result { - self.execute(|connection| connection.ping().map_err(Error::unavailable)) - } - - async fn ping(&self) -> Result { - self.run(|connection| connection.ping().map_err(Error::unavailable)) - .await - } -} - -impl ConnectionCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - /// Python `RedisCache.test_connection`, or `RedisClusterCache.test_connection` for a - /// cluster topology, which differs only in its messages. - async fn test_connection(&self) -> Result { - let label = match self.topology { - RedisTopology::Standalone => "Redis", - RedisTopology::Cluster { .. } => "Redis Cluster", - }; - let ping = self - .run(|connection| Ok(connection.ping().map_err(Error::unavailable))) - .await - .and_then(|ping| ping) - .map_err(|error| litellm_error::Report::new(&error).to_string()); - Ok(match ping { - Ok(true) => CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: format!("{label} connection test successful"), - error: None, - }, - Ok(false) => CacheConnectionResult { - status: CacheConnectionStatus::Failed, - message: format!("{label} ping returned False"), - error: None, - }, - Err(error) => CacheConnectionResult { - status: CacheConnectionStatus::Failed, - message: format!("{label} connection failed: {error}"), - error: Some(error), - }, - }) - } -} - -impl DisconnectCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn disconnect(&self) -> Result<(), Error> { - self.connections.disconnect(); - Ok(()) - } -} - -impl ClientInfoCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type ClientList = String; - type Info = String; - - fn client_list(&self) -> Result { - self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST"))) - } - - fn info(&self) -> Result { - self.execute(|connection| connection.node_text(&redis::cmd("INFO"))) - } -} diff --git a/litellm-rust/crates/cache-redis/src/queue.rs b/litellm-rust/crates/cache-redis/src/queue.rs deleted file mode 100644 index d55e4f49187..00000000000 --- a/litellm-rust/crates/cache-redis/src/queue.rs +++ /dev/null @@ -1,226 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{CacheCodec, Error, PopOperation, PushOperation, QueueCache, SetCache}; - -use crate::{cache::RedisCache, script::RedisArg}; - -pub type RedisRpushOperation = PushOperation; -pub type RedisLpopOperation = PopOperation; - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum RedisLpopResult { - Missing, - Value(Vec), - Values(Vec>), -} - -impl SetCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type SetValue = RedisArg; - type SetResult = usize; - - async fn async_set_cache_sadd( - &self, - key: &str, - values: Vec, - ttl: Option, - ) -> Result { - if values.is_empty() { - return Err(Error::InvalidEntry(None)); - } - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(ttl); - let mut pipeline = redis::pipe(); - pipeline - .cmd("SADD") - .arg(&key) - .arg(values) - .cmd("EXPIRE") - .arg(&key) - .arg(ttl) - .ignore(); - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - .map(|(added,)| added) - } -} - -impl QueueCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type QueueValue = RedisArg; - type PopResult = RedisLpopResult; - - async fn async_rpush(&self, key: &str, values: Vec) -> Result { - if values.is_empty() { - return Err(Error::InvalidEntry(None)); - } - let key = self.namespaced_key(key); - self.run(move |connection| { - redis::cmd("RPUSH") - .arg(key) - .arg(values) - .query(connection) - .map_err(Error::unavailable) - }) - .await - } - - async fn async_rpush_and_trim( - &self, - key: &str, - values: Vec, - max_len: usize, - ) -> Result { - if values.is_empty() { - return Err(Error::InvalidEntry(None)); - } - let key = self.namespaced_key(key); - let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len); - let mut pipeline = redis::pipe(); - pipeline - .atomic() - .cmd("RPUSH") - .arg(&key) - .arg(values) - .cmd("LTRIM") - .arg(&key) - .arg(start) - .arg(-1) - .ignore(); - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - .map(|(length,)| length) - } - - async fn async_rpush_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - if operations.is_empty() { - return Ok(Vec::new()); - } - let mut pipeline = redis::pipe(); - for operation in operations { - if operation.values.is_empty() { - return Err(Error::InvalidEntry(None)); - } - pipeline - .cmd("RPUSH") - .arg(self.namespaced_key(&operation.key)) - .arg(operation.values); - } - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - } - - async fn async_lpop(&self, key: &str, count: Option) -> Result { - if let Some(count) = count - && self.major_version().await < 7 - { - return self.lpop_one_at_a_time(key, count).await; - } - let command = lpop(self.namespaced_key(key), count); - let value = self - .run(move |connection| { - command - .query::(connection) - .map_err(Error::unavailable) - }) - .await?; - lpop_result(value, count.is_some()) - } - - async fn async_lpop_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - if operations.is_empty() { - return Ok(Vec::new()); - } - if operations.iter().any(|operation| operation.count.is_some()) - && self.major_version().await < 7 - { - let mut results = Vec::with_capacity(operations.len()); - for operation in &operations { - results.push(self.async_lpop(&operation.key, operation.count).await?); - } - return Ok(results); - } - let multiple = operations - .iter() - .map(|operation| operation.count.is_some()) - .collect::>(); - let mut pipeline = redis::pipe(); - for operation in operations { - pipeline.add_command(lpop(self.namespaced_key(&operation.key), operation.count)); - } - self.run(move |connection| connection.query_pipeline::>(&pipeline)) - .await? - .into_iter() - .zip(multiple) - .map(|(value, multiple)| lpop_result(value, multiple)) - .collect() - } -} - -fn lpop(key: String, count: Option) -> redis::Cmd { - let mut command = redis::cmd("LPOP"); - command.arg(key); - if let Some(count) = count { - command.arg(count); - } - command -} - -fn redis_bytes(value: redis::Value) -> Result, Error> { - match value { - redis::Value::BulkString(bytes) => Ok(bytes), - redis::Value::SimpleString(text) => Ok(text.into_bytes()), - _ => Err(Error::InvalidEntry(None)), - } -} - -impl RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - /// `handle_lpop_count_for_older_redis_versions`: `count` single-`LPOP` pipelines, keeping - /// only the values actually popped. - async fn lpop_one_at_a_time(&self, key: &str, count: usize) -> Result { - let key = self.namespaced_key(key); - let mut values = Vec::new(); - for _ in 0..count { - let mut pipeline = redis::pipe(); - pipeline.add_command(lpop(key.clone(), None)); - let replies = self - .run(move |connection| connection.query_pipeline::>(&pipeline)) - .await?; - for reply in replies { - if reply != redis::Value::Nil { - values.push(redis_bytes(reply)?); - } - } - } - Ok(RedisLpopResult::Values(values)) - } -} - -fn lpop_result(value: redis::Value, multiple: bool) -> Result { - match value { - redis::Value::Nil => Ok(RedisLpopResult::Missing), - redis::Value::Array(values) if multiple => values - .into_iter() - .map(redis_bytes) - .collect::, _>>() - .map(RedisLpopResult::Values), - value if !multiple => redis_bytes(value).map(RedisLpopResult::Value), - _ => Err(Error::InvalidEntry(None)), - } -} diff --git a/litellm-rust/crates/cache-redis/src/script.rs b/litellm-rust/crates/cache-redis/src/script.rs deleted file mode 100644 index 09987437ed4..00000000000 --- a/litellm-rust/crates/cache-redis/src/script.rs +++ /dev/null @@ -1,136 +0,0 @@ -use std::sync::Arc; - -use litellm_cache::{CacheCodec, CacheScript, Error, ScriptCache}; - -use crate::{ - cache::{RedisCache, namespaced_key}, - connection::Connections, -}; - -#[derive(Clone, Debug, PartialEq)] -pub enum RedisArg { - Bytes(Vec), - Integer(i64), - Float(f64), -} - -impl From<&str> for RedisArg { - fn from(value: &str) -> Self { - Self::Bytes(value.as_bytes().to_vec()) - } -} - -impl From for RedisArg { - fn from(value: String) -> Self { - Self::Bytes(value.into_bytes()) - } -} - -impl From> for RedisArg { - fn from(value: Vec) -> Self { - Self::Bytes(value) - } -} - -impl From for RedisArg { - fn from(value: i64) -> Self { - Self::Integer(value) - } -} - -impl From for RedisArg { - fn from(value: f64) -> Self { - Self::Float(value) - } -} - -impl redis::ToRedisArgs for RedisArg { - fn write_redis_args(&self, out: &mut W) - where - W: ?Sized + redis::RedisWrite, - { - match self { - Self::Bytes(value) => value.write_redis_args(out), - Self::Integer(value) => value.write_redis_args(out), - Self::Float(value) => value.write_redis_args(out), - } - } -} - -pub struct RedisScript { - connections: Arc>, - namespace: Option, - source: String, -} - -impl CacheScript for RedisScript -where - C: redis::ConnectionLike + Send + 'static, -{ - type Argument = RedisArg; - type Output = redis::Value; - - async fn invoke( - &self, - keys: Vec, - arguments: Vec, - ) -> Result { - let keys = keys - .into_iter() - .map(|key| namespaced_key(self.namespace.as_deref(), &key)) - .collect::>(); - let source = self.source.clone(); - Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - eval(connection, &source, keys, arguments) - }) - .await - } -} - -impl RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub async fn async_eval( - &self, - script: String, - keys: Vec, - arguments: Vec, - ) -> Result { - let keys = self.namespaced_keys(&keys); - self.run(move |connection| eval(connection, &script, keys, arguments)) - .await - } -} - -fn eval( - connection: &mut impl redis::ConnectionLike, - script: &str, - keys: Vec, - arguments: Vec, -) -> Result { - redis::cmd("EVAL") - .arg(script) - .arg(keys.len()) - .arg(keys) - .arg(arguments) - .query(connection) - .map_err(Error::unavailable) -} - -impl ScriptCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Script = RedisScript; - - fn async_register_script(&self, source: String) -> Self::Script { - RedisScript { - connections: Arc::clone(&self.connections), - namespace: self.namespace.clone(), - source, - } - } -} diff --git a/litellm-rust/crates/cache-redis/src/store.rs b/litellm-rust/crates/cache-redis/src/store.rs deleted file mode 100644 index daadc940ca9..00000000000 --- a/litellm-rust/crates/cache-redis/src/store.rs +++ /dev/null @@ -1,232 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheCodec, DeleteCache, Error, - ExactCacheContext, FlushAllCache, FlushCache, TtlPipelineCache, -}; -use redis::Commands; - -use crate::{cache::RedisCache, connection::ConnectionRef}; - -impl BaseCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Value = S::Value; - type Context = ExactCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl.or(Some(self.default_ttl)) - } - - fn set_cache( - &self, - key: &str, - value: Self::Value, - context: &ExactCacheContext, - ) -> Result<(), Error> { - let payload = self.codec.encode(&value)?; - let ttl = self.ttl_or_default(context.ttl); - let key = self.namespaced_key(key); - self.execute(|connection| { - connection - .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(Error::unavailable) - }) - } - - fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { - let key = self.namespaced_key(key); - let value = self.execute(|connection| { - connection - .get::<_, redis::Value>(key) - .map_err(Error::unavailable) - })?; - self.decode_response(value) - } - - async fn async_set_cache( - &self, - key: &str, - value: Self::Value, - context: ExactCacheContext, - ) -> Result<(), Error> { - let payload = self.codec.encode(&value)?; - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - self.run(move |connection| { - connection - .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(Error::unavailable) - }) - .await - } - - async fn async_get_cache( - &self, - key: &str, - _: &ExactCacheContext, - ) -> Result, Error> { - let key = self.namespaced_key(key); - let value = self - .run(move |connection| { - connection - .get::<_, redis::Value>(key) - .map_err(Error::unavailable) - }) - .await?; - self.decode_response(value) - } - - async fn async_set_cache_pipeline( - &self, - cache_list: Vec<(String, Self::Value)>, - context: ExactCacheContext, - ) -> Result<(), Error> { - self.async_set_cache_pipeline_with_ttls( - cache_list - .into_iter() - .map(|(key, value)| (key, value, context.ttl)) - .collect(), - ) - .await - } -} - -impl TtlPipelineCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_set_cache_pipeline_with_ttls( - &self, - entries: Vec<(String, Self::Value, Option)>, - ) -> Result<(), Error> { - if entries.is_empty() { - return Ok(()); - } - let mut pipeline = redis::pipe(); - for (key, value, ttl) in entries { - pipeline - .cmd("SETEX") - .arg(self.namespaced_key(&key)) - .arg(self.ttl_or_default(ttl)) - .arg(self.codec.encode(&value)?) - .ignore(); - } - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - } -} - -impl BatchCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn batch_get_cache( - &self, - keys: &[String], - _: &ExactCacheContext, - ) -> Result>, Error> { - let keys = self.namespaced_keys(keys); - self.execute(|connection| mget(connection, keys))? - .into_iter() - .map(|value| self.decode_batch_response(value)) - .collect() - } - - async fn async_batch_get_cache( - &self, - keys: Vec, - _: ExactCacheContext, - ) -> Result>, Error> { - let keys = self.namespaced_keys(&keys); - self.run(move |connection| mget(connection, keys)) - .await? - .into_iter() - .map(|value| self.decode_batch_response(value)) - .collect() - } -} - -pub(crate) fn mget( - connection: &mut ConnectionRef<'_>, - keys: Vec, -) -> Result, Error> { - redis::cmd("MGET") - .arg(keys) - .query(connection) - .map_err(Error::unavailable) -} - -impl DeleteCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn delete_cache(&self, key: &str) -> Result<(), Error> { - let key = self.namespaced_key(key); - self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable)) - } - - async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - let key = self.namespaced_key(key); - self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable)) - .await - } -} - -impl BulkDeleteCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn delete_cache_keys(&self, keys: Vec) -> Result { - if keys.is_empty() { - return Ok(0); - } - let keys = self.namespaced_keys(&keys); - self.run(move |connection| connection.del(keys).map_err(Error::unavailable)) - .await - } -} - -impl FlushCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn flush_cache(&self) -> Result<(), Error> { - let pattern = self.namespaced_pattern()?; - self.execute(|connection| flush_matching(connection, &pattern)) - } - - async fn async_flush_cache(&self) -> Result<(), Error> { - let pattern = self.namespaced_pattern()?; - self.run(move |connection| flush_matching(connection, &pattern)) - .await - } -} - -impl FlushAllCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn flushall(&self) -> Result<(), Error> { - self.execute(|connection| connection.flushall()) - } -} - -fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> { - connection.scan(pattern, 1000, |connection, keys| { - if !keys.is_empty() { - connection - .del::<_, usize>(keys) - .map_err(Error::unavailable)?; - } - Ok(true) - }) -} diff --git a/litellm-rust/crates/cache-redis/src/topology.rs b/litellm-rust/crates/cache-redis/src/topology.rs deleted file mode 100644 index 7f4ee48b222..00000000000 --- a/litellm-rust/crates/cache-redis/src/topology.rs +++ /dev/null @@ -1,14 +0,0 @@ -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RedisNode { - pub host: String, - pub port: u16, -} - -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub enum RedisTopology { - #[default] - Standalone, - Cluster { - startup_nodes: Vec, - }, -} diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs deleted file mode 100644 index 9ec85251733..00000000000 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ /dev/null @@ -1,1164 +0,0 @@ -mod support; - -use std::time::Duration; - -use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheCodec, - CacheConnectionStatus, CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, - CountReadCache, CounterCache, DeleteCache, DisconnectCache, Error, ExactCacheContext, - FlushAllCache, FlushCache, IncrementOperation, JsonCodec, PingCache, QueueCache, - RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache, TtlPipelineCache, get_cache, - set_cache, -}; -use litellm_cache_redis::{ - RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, -}; -use redis_test::{MockCmd, MockRedisConnection}; -use rstest::{fixture, rstest}; -use serde_json::json; -use support::TaggedByteCodec; - -type Mocked = RedisCache; - -fn mock(commands: Vec) -> MockRedisConnection { - MockRedisConnection::new(commands).assert_all_commands_consumed() -} - -fn tagged(commands: Vec) -> Mocked { - RedisCache::with_connection(mock(commands), None, TaggedByteCodec(42)) -} - -fn json_cache(commands: Vec) -> Mocked> { - RedisCache::with_connection(mock(commands), None, JsonCodec::new()) -} - -fn team(commands: Vec) -> Mocked> { - json_cache(commands).with_namespace(Some("team".into())) -} - -#[fixture] -fn context() -> ExactCacheContext { - ExactCacheContext::default() -} - -fn scan(pattern: &str, cursor: u64, count: usize, reply: redis::Value) -> MockCmd { - MockCmd::new( - redis::cmd("SCAN") - .cursor_arg(cursor) - .arg("MATCH") - .arg(pattern) - .arg("COUNT") - .arg(count), - Ok(reply), - ) -} - -#[rstest] -fn constructor_rejects_invalid_urls() { - assert!(RedisCache::new("not a redis url", None, JsonCodec::::new()).is_err()); -} - -#[rstest] -#[case::zero_rounds_up_to_one(Some(Duration::ZERO), 1)] -#[case::fractions_round_up(Some(Duration::from_millis(1500)), 2)] -#[case::whole_seconds_are_kept(Some(Duration::from_secs(15)), 15)] -#[case::missing_ttl_uses_default(None, 600)] -fn writes_round_ttls_up_to_positive_seconds(#[case] ttl: Option, #[case] seconds: u64) { - let cache = tagged(vec![MockCmd::new( - redis::cmd("SETEX") - .arg("key") - .arg(seconds) - .arg([42u8, 7].as_slice()), - Ok("OK"), - )]); - cache - .set_cache("key", 7, &ExactCacheContext { ttl }) - .unwrap(); -} - -#[rstest] -fn generic_helpers_use_the_injected_codec_and_ttl() { - let cache = tagged(vec![ - MockCmd::new( - redis::cmd("SETEX") - .arg("counter") - .arg(2) - .arg([42u8, 7].as_slice()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), - ]); - let context = ExactCacheContext { - ttl: Some(Duration::from_millis(1500)), - }; - set_cache(&cache, "counter", 7, &context).unwrap(); - assert_eq!(get_cache(&cache, "counter", &context).unwrap(), Some(7)); -} - -#[rstest] -fn commands_round_trip_entries_and_delete_only_namespaced_keys(context: ExactCacheContext) { - let value = json!({"deployment": "model-a", "cooldown_seconds": 30}); - let payload = JsonCodec::::new() - .encode(&value) - .unwrap(); - let cache = json_cache(vec![ - MockCmd::new( - redis::cmd("SETEX") - .arg("litellm-cache:key") - .arg(600) - .arg(payload.clone()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)), - MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), - ]) - .with_namespace(Some("litellm-cache".into())); - - cache.set_cache("key", value.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key", &context).unwrap(), Some(value)); - cache.delete_cache("key").unwrap(); -} - -#[rstest] -#[tokio::test] -async fn async_operations_preserve_codec_ttl_and_missing_values(context: ExactCacheContext) { - let cache = RedisCache::with_connection( - mock(vec![ - MockCmd::new( - redis::cmd("SETEX") - .arg("counter") - .arg(9) - .arg([42u8, 7].as_slice()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), - MockCmd::new( - redis::cmd("SETEX") - .arg("batch") - .arg(2) - .arg([42u8, 8].as_slice()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)), - MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)), - ]), - Some(Duration::from_secs(9)), - TaggedByteCodec(42), - ); - cache - .batch_cache_write("counter", 7, context.clone()) - .await - .unwrap(); - assert_eq!( - cache.async_get_cache("counter", &context).await.unwrap(), - Some(7) - ); - cache - .async_set_cache_pipeline( - vec![("batch".into(), 8)], - ExactCacheContext { - ttl: Some(Duration::from_millis(1500)), - }, - ) - .await - .unwrap(); - cache.async_delete_cache("counter").await.unwrap(); - assert_eq!( - cache.async_get_cache("counter", &context).await.unwrap(), - None - ); -} - -#[rstest] -#[tokio::test] -async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCacheContext) { - let cache = tagged(vec![ - MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), - MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), - ]); - assert_eq!( - cache.set_cache("invalid", 255, &context), - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache.async_set_cache("invalid", 255, context.clone()).await, - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache - .async_set_cache_pipeline( - vec![("valid".into(), 7), ("invalid".into(), 255)], - context.clone(), - ) - .await, - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache - .async_set_cache_pipeline_with_ttls(vec![ - ("valid".into(), 7, None), - ("invalid".into(), 255, None), - ]) - .await, - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache.get_cache("invalid", &context), - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache.async_get_cache("invalid", &context).await, - Err(Error::InvalidEntry(None)) - ); -} - -#[rstest] -#[case::bare_key("key")] -#[case::already_prefixed("team:key")] -fn namespaces_are_added_once(#[case] key: &str, context: ExactCacheContext) { - let cache = team(vec![MockCmd::new( - redis::cmd("GET").arg("team:key"), - Ok(redis::Value::Nil), - )]); - assert_eq!(cache.get_cache(key, &context).unwrap(), None); -} - -#[rstest] -#[case::empty(Some(String::new()))] -#[case::missing(None)] -fn empty_namespaces_leave_keys_unprefixed( - #[case] namespace: Option, - context: ExactCacheContext, -) { - let cache = json_cache(vec![MockCmd::new( - redis::cmd("GET").arg("key"), - Ok(redis::Value::Nil), - )]) - .with_namespace(namespace); - assert_eq!(cache.namespace(), None); - assert_eq!(cache.get_cache("key", &context).unwrap(), None); -} - -#[rstest] -#[tokio::test] -async fn flush_requires_a_namespace() { - let unscoped = json_cache(Vec::new()); - assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush)); - assert_eq!( - unscoped.async_flush_cache().await, - Err(Error::UnscopedFlush) - ); -} - -#[rstest] -#[case::plain_namespace("litellm-cache", "litellm-cache:*", "litellm-cache:key")] -#[case::glob_metacharacters_are_escaped("team*", "team\\*:*", "team*:key")] -fn flush_scans_and_deletes_only_namespaced_keys( - #[case] namespace: &str, - #[case] pattern: &str, - #[case] key: &str, -) { - let cache = json_cache(vec![ - scan(pattern, 0, 1000, redis_test::redis_value!(["0", [key]])), - MockCmd::new(redis::cmd("DEL").arg(key), Ok(1u32)), - ]) - .with_namespace(Some(namespace.into())); - cache.flush_cache().unwrap(); -} - -#[rstest] -#[tokio::test] -async fn async_flush_deletes_each_scan_page_separately() { - let cache = team(vec![ - scan( - "team:*", - 0, - 1000, - redis_test::redis_value!(["7", ["team:a", "team:b"]]), - ), - MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)), - scan( - "team:*", - 7, - 1000, - redis_test::redis_value!(["0", ["team:c"]]), - ), - MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)), - ]); - cache.async_flush_cache().await.unwrap(); -} - -#[rstest] -fn flushall_ignores_the_namespace() { - team(vec![MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK"))]) - .flushall() - .unwrap(); -} - -#[rstest] -#[tokio::test] -async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries( - context: ExactCacheContext, -) { - let cache = tagged(vec![MockCmd::new( - redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"), - Ok(vec![ - redis::Value::BulkString(vec![42, 7]), - redis::Value::Nil, - redis::Value::BulkString(vec![99, 7]), - ]), - )]); - - assert_eq!( - cache - .async_batch_get_cache(vec!["hit".into(), "miss".into(), "invalid".into()], context) - .await - .unwrap(), - vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid] - ); -} - -#[rstest] -#[tokio::test] -async fn ttl_pipeline_keeps_each_entry_ttl_and_defaults_missing_ones() { - let mut pipeline = redis::pipe(); - pipeline - .cmd("SETEX") - .arg("ns:team_id:t1") - .arg(60u64) - .arg(r#"{"team_id":"t1"}"#) - .cmd("SETEX") - .arg("ns:u1") - .arg(7u64) - .arg(r#"{"user_id":"u1"}"#) - .cmd("SETEX") - .arg("ns:org_id:o1") - .arg(300u64) - .arg(r#"{"a":1}"#); - let cache = RedisCache::with_connection( - mock(vec![MockCmd::with_values( - pipeline, - Ok(vec!["OK", "OK", "OK"]), - )]), - Some(Duration::from_secs(300)), - JsonCodec::::new(), - ) - .with_namespace(Some("ns".into())); - - cache - .async_set_cache_pipeline_with_ttls(vec![ - ( - "team_id:t1".into(), - json!({"team_id": "t1"}), - Some(Duration::from_secs(60)), - ), - ( - "u1".into(), - json!({"user_id": "u1"}), - Some(Duration::from_secs(7)), - ), - ("org_id:o1".into(), json!({"a": 1}), None), - ]) - .await - .unwrap(); -} - -#[rstest] -#[tokio::test] -async fn empty_pipelines_skip_the_round_trip(context: ExactCacheContext) { - let cache = json_cache(Vec::new()); - cache - .async_set_cache_pipeline(Vec::new(), context) - .await - .unwrap(); - cache - .async_set_cache_pipeline_with_ttls(Vec::new()) - .await - .unwrap(); - assert_eq!(cache.delete_cache_keys(Vec::new()).await.unwrap(), 0); - assert_eq!( - cache.async_rpush_pipeline(Vec::new()).await.unwrap(), - Vec::::new() - ); - assert_eq!( - cache.async_lpop_pipeline(Vec::new()).await.unwrap(), - Vec::::new() - ); - assert_eq!( - cache.async_increment_pipeline(Vec::new()).await.unwrap(), - Vec::::new() - ); -} - -#[rstest] -#[tokio::test] -async fn count_reads_parse_integers_and_keep_missing_counters() { - let mget = || { - MockCmd::new( - redis::cmd("MGET").arg("team:count").arg("team:missing"), - Ok(redis_test::redis_value!(["7", nil])), - ) - }; - let cache = team(vec![mget(), mget()]); - let keys = vec!["count".to_string(), "missing".to_string()]; - - assert_eq!(cache.batch_get_counts(&keys).unwrap(), [Some(7), None]); - assert_eq!( - cache.async_batch_get_counts(keys).await.unwrap(), - [Some(7), None] - ); -} - -#[rstest] -#[tokio::test] -async fn pings_run_on_both_paths() { - let cache = team(vec![ - MockCmd::new(redis::cmd("PING"), Ok("PONG")), - MockCmd::new(redis::cmd("PING"), Ok("PONG")), - ]); - assert!(cache.sync_ping().unwrap()); - assert!(cache.ping().await.unwrap()); -} - -#[rstest] -#[case::remaining(12, Some(Duration::from_secs(12)))] -#[case::no_expiry(-1, None)] -#[case::missing(-2, None)] -#[tokio::test] -async fn ttl_reads_hide_negative_replies(#[case] reply: i64, #[case] ttl: Option) { - let cache = team(vec![MockCmd::new( - redis::cmd("TTL").arg("team:key"), - Ok(reply), - )]); - assert_eq!(cache.async_get_ttl("key").await.unwrap(), ttl); -} - -#[rstest] -#[case::explicit_ttl(Some(Duration::from_secs(30)), 30, 1, true)] -#[case::default_ttl(None, 600, 1, true)] -#[case::missing_key(Some(Duration::from_secs(30)), 30, 0, false)] -#[tokio::test] -async fn refresh_ttl_expires_existing_keys_only( - #[case] ttl: Option, - #[case] seconds: u64, - #[case] reply: i64, - #[case] refreshed: bool, -) { - let cache = team(vec![MockCmd::new( - redis::cmd("EXPIRE").arg("team:key").arg(seconds), - Ok(reply), - )]); - assert_eq!( - cache.async_refresh_ttl("key", ttl).await.unwrap(), - refreshed - ); -} - -#[rstest] -#[tokio::test] -async fn scan_stops_at_count_and_bulk_delete_reports_existing_keys() { - let cache = team(vec![ - scan( - "team:job-*", - 0, - 25, - redis_test::redis_value!(["4", ["team:job-a"]]), - ), - scan( - "team:job-*", - 4, - 25, - redis_test::redis_value!(["0", ["team:job-b"]]), - ), - MockCmd::new( - redis::cmd("DEL").arg("team:job-a").arg("team:job-b"), - Ok(2u32), - ), - ]); - assert_eq!( - cache.async_scan_iter("job-", 25).await.unwrap(), - ["team:job-a", "team:job-b"] - ); - assert_eq!( - cache - .delete_cache_keys(vec!["job-a".into(), "job-b".into()]) - .await - .unwrap(), - 2 - ); -} - -#[rstest] -#[tokio::test] -async fn sets_add_members_and_arm_the_default_ttl() { - let mut pipeline = redis::pipe(); - pipeline - .cmd("SADD") - .arg("team:members") - .arg("a") - .arg("b") - .cmd("EXPIRE") - .arg("team:members") - .arg(600u64) - .ignore(); - let cache = team(vec![MockCmd::with_values( - pipeline, - Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]), - )]); - assert_eq!( - cache - .async_set_cache_sadd("members", vec!["a".into(), "b".into()], None) - .await - .unwrap(), - 2 - ); - assert_eq!( - cache - .async_set_cache_sadd("members", Vec::new(), None) - .await, - Err(Error::InvalidEntry(None)) - ); -} - -#[rstest] -#[tokio::test] -async fn queues_push_and_pop_namespaced_lists() { - let cache = team(vec![ - MockCmd::new( - redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"), - Ok(2u32), - ), - MockCmd::new( - redis::cmd("INFO"), - Ok("# Server\r\nredis_version:7.2.4\r\n"), - ), - MockCmd::new( - redis::cmd("LPOP").arg("team:queue").arg(2usize), - Ok(redis_test::redis_value!(["a", "b"])), - ), - MockCmd::new(redis::cmd("LPOP").arg("team:queue"), Ok("c")), - ]); - assert_eq!( - cache - .async_rpush("queue", vec!["a".into(), "b".into()]) - .await - .unwrap(), - 2 - ); - assert_eq!( - cache.async_rpush("queue", Vec::new()).await, - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache.async_lpop("queue", Some(2)).await.unwrap(), - RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) - ); - assert_eq!( - cache.async_lpop("queue", None).await.unwrap(), - RedisLpopResult::Value(b"c".to_vec()) - ); -} - -/// Python `RedisCache.async_lpop` checks `redis_version` from `INFO` and, below major version 7, -/// pops a counted batch as `count` single-command `LPOP` pipelines, dropping `None` replies. -#[rstest] -#[tokio::test] -async fn counted_lpop_falls_back_to_single_pops_below_redis_7() { - let single_pop = || { - let mut pipeline = redis::pipe(); - pipeline.cmd("LPOP").arg("team:queue"); - pipeline - }; - let cache = team(vec![ - MockCmd::new( - redis::cmd("INFO"), - Ok("# Server\r\nredis_version:6.2.14\r\n"), - ), - MockCmd::with_values(single_pop(), Ok(vec![redis_test::redis_value!("a")])), - MockCmd::with_values(single_pop(), Ok(vec![redis_test::redis_value!("b")])), - MockCmd::with_values(single_pop(), Ok(vec![redis::Value::Nil])), - ]); - - assert_eq!( - cache.async_lpop("queue", Some(3)).await.unwrap(), - RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) - ); -} - -/// Python keeps `redis_version = "Unknown"` when `INFO` fails and then assumes -/// `DEFAULT_REDIS_MAJOR_VERSION` (7), so a counted pop is one `LPOP key count`. The version is read -/// once: the second pop sends no second `INFO`. -#[rstest] -#[tokio::test] -async fn counted_lpop_assumes_redis_7_when_info_fails() { - let cache = team(vec![ - MockCmd::new( - redis::cmd("INFO"), - Err::(redis::RedisError::from((redis::ErrorKind::Io, "down"))), - ), - MockCmd::new( - redis::cmd("LPOP").arg("team:queue").arg(2usize), - Ok(redis_test::redis_value!(["a", "b"])), - ), - MockCmd::new( - redis::cmd("LPOP").arg("team:queue").arg(1usize), - Ok(redis_test::redis_value!(["c"])), - ), - ]); - - assert_eq!( - cache.async_lpop("queue", Some(2)).await.unwrap(), - RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) - ); - assert_eq!( - cache.async_lpop("queue", Some(1)).await.unwrap(), - RedisLpopResult::Values(vec![b"c".to_vec()]) - ); -} - -fn push_and_trim(start: i64) -> redis::Pipeline { - let mut pipeline = redis::pipe(); - pipeline - .atomic() - .cmd("RPUSH") - .arg("ns:buf") - .arg("c") - .arg("d") - .cmd("LTRIM") - .arg("ns:buf") - .arg(start) - .arg(-1); - pipeline -} - -#[rstest] -#[case::keeps_newest_entries(3, -3)] -#[case::zero_keeps_everything(0, 0)] -#[tokio::test] -async fn rpush_and_trim_runs_push_and_trim_in_one_transaction( - #[case] max_len: usize, - #[case] start: i64, -) { - let cache = json_cache(vec![MockCmd::with_values( - push_and_trim(start), - Ok(vec![redis::Value::Array(vec![ - redis::Value::Int(4), - redis::Value::Okay, - ])]), - )]) - .with_namespace(Some("ns".into())); - - assert_eq!( - cache - .async_rpush_and_trim("buf", vec!["c".into(), "d".into()], max_len) - .await - .unwrap(), - 4 - ); -} - -#[rstest] -#[tokio::test] -async fn rpush_and_trim_raises_when_a_queued_command_fails() { - let wrong_type = redis::parse_redis_value( - b"-WRONGTYPE Operation against a key holding the wrong kind of value\r\n", - ) - .unwrap(); - let cache = json_cache(vec![MockCmd::with_values( - push_and_trim(-3), - Ok(vec![redis::Value::Array(vec![ - wrong_type, - redis::Value::Okay, - ])]), - )]) - .with_namespace(Some("ns".into())); - - assert!(matches!( - cache - .async_rpush_and_trim("buf", vec!["c".into(), "d".into()], 3) - .await, - Err(Error::Unavailable(Some(_))) - )); - assert_eq!( - cache.async_rpush_and_trim("buf", Vec::new(), 3).await, - Err(Error::InvalidEntry(None)) - ); -} - -#[rstest] -#[tokio::test] -async fn scripts_and_eval_namespace_their_keys() { - let eval = || { - MockCmd::new( - redis::cmd("EVAL") - .arg("return KEYS[1]") - .arg(1usize) - .arg("team:key"), - Ok("team:key"), - ) - }; - let cache = team(vec![eval(), eval()]); - assert_eq!( - cache - .async_eval("return KEYS[1]".into(), vec!["key".into()], Vec::new()) - .await - .unwrap(), - redis::Value::BulkString(b"team:key".to_vec()) - ); - assert_eq!( - cache - .async_register_script("return KEYS[1]".into()) - .invoke(vec!["key".into()], Vec::new()) - .await - .unwrap(), - redis::Value::BulkString(b"team:key".to_vec()) - ); -} - -#[rstest] -fn client_list_and_info_return_server_text() { - let cache = team(vec![ - MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")), - MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")), - ]); - assert_eq!(cache.client_list().unwrap(), "id=1"); - assert_eq!(cache.info().unwrap(), "redis_version:7"); -} - -#[rstest] -#[tokio::test] -async fn pipelines_preserve_operation_order() { - let mut rpush_pipeline = redis::pipe(); - rpush_pipeline - .cmd("RPUSH") - .arg("team:a") - .arg("one") - .cmd("RPUSH") - .arg("team:b") - .arg("two"); - let mut lpop_pipeline = redis::pipe(); - lpop_pipeline - .cmd("LPOP") - .arg("team:a") - .arg(2usize) - .cmd("LPOP") - .arg("team:b"); - let queue = team(vec![ - MockCmd::with_values( - rpush_pipeline, - Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]), - ), - MockCmd::new( - redis::cmd("INFO"), - Ok("# Server\r\nredis_version:7.2.4\r\n"), - ), - MockCmd::with_values( - lpop_pipeline, - Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]), - ), - ]); - - assert_eq!( - queue - .async_rpush_pipeline(vec![ - RedisRpushOperation { - key: "a".into(), - values: vec![RedisArg::from("one")], - }, - RedisRpushOperation { - key: "b".into(), - values: vec![RedisArg::from("two")], - }, - ]) - .await - .unwrap(), - [1, 2] - ); - assert_eq!( - queue - .async_lpop_pipeline(vec![ - RedisLpopOperation { - key: "a".into(), - count: Some(2), - }, - RedisLpopOperation { - key: "b".into(), - count: None, - }, - ]) - .await - .unwrap(), - [ - RedisLpopResult::Values(vec![b"one".to_vec()]), - RedisLpopResult::Missing, - ] - ); -} - -/// Below major version 7 a counted `LPOP` is unsupported, so a pipeline that mixes counted and -/// plain pops runs each operation through `async_lpop`: `count` single-`LPOP` pipelines for the -/// counted ones, a bare `LPOP` for the rest. No `LPOP key count` reaches the connection. -#[rstest] -#[tokio::test] -async fn lpop_pipeline_pops_one_at_a_time_below_redis_7() { - let single_pop = |key: &str| { - let mut pipeline = redis::pipe(); - pipeline.cmd("LPOP").arg(key); - pipeline - }; - let cache = team(vec![ - MockCmd::new( - redis::cmd("INFO"), - Ok("# Server\r\nredis_version:6.2.14\r\n"), - ), - MockCmd::with_values( - single_pop("team:a"), - Ok(vec![redis_test::redis_value!("one")]), - ), - MockCmd::with_values( - single_pop("team:a"), - Ok(vec![redis_test::redis_value!("two")]), - ), - MockCmd::new(redis::cmd("LPOP").arg("team:b"), Ok(redis::Value::Nil)), - ]); - - assert_eq!( - cache - .async_lpop_pipeline(vec![ - RedisLpopOperation { - key: "a".into(), - count: Some(2), - }, - RedisLpopOperation { - key: "b".into(), - count: None, - }, - ]) - .await - .unwrap(), - [ - RedisLpopResult::Values(vec![b"one".to_vec(), b"two".to_vec()]), - RedisLpopResult::Missing, - ] - ); -} - -#[rstest] -#[tokio::test] -async fn increment_pipeline_expires_only_operations_with_a_ttl() { - let mut increment_pipeline = redis::pipe(); - increment_pipeline - .cmd("INCRBYFLOAT") - .arg("team:counter") - .arg(1.5f64) - .cmd("EXPIRE") - .arg("team:counter") - .arg(10u64) - .ignore() - .cmd("INCRBYFLOAT") - .arg("team:counter") - .arg(2.0f64); - let counters = team(vec![MockCmd::with_values( - increment_pipeline, - Ok(vec![ - redis::Value::BulkString(b"1.5".to_vec()), - redis::Value::Int(1), - redis::Value::BulkString(b"3.5".to_vec()), - ]), - )]); - assert_eq!( - counters - .async_increment_pipeline(vec![ - IncrementOperation { - key: "counter".into(), - amount: 1.5, - ttl: Some(Duration::from_secs(10)), - }, - IncrementOperation { - key: "counter".into(), - amount: 2.0, - ttl: None, - }, - ]) - .await - .unwrap(), - [1.5, 3.5] - ); -} - -const INCREMENT_SCRIPT: &str = concat!( - "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", - "if redis.call('TTL', KEYS[1]) == -1 then ", - "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" -); -const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( - "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", - "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ", - "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", - "return count" -); -const SET_MAX_SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if current == false or tonumber(current) < tonumber(ARGV[1]) then ", - "redis.call('SET', KEYS[1], ARGV[1]); ", - "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", - "return ARGV[1]; end; return current" -); - -fn increment_script(amount: f64) -> MockCmd { - MockCmd::new( - redis::cmd("EVAL") - .arg(INCREMENT_SCRIPT) - .arg(1) - .arg("counter") - .arg(amount) - .arg(600), - Ok("4.5"), - ) -} - -#[rstest] -#[tokio::test] -async fn increments_keep_an_existing_ttl_in_one_atomic_script(context: ExactCacheContext) { - let cache = RedisCache::with_connection( - mock(vec![increment_script(2.5), increment_script(2.5)]), - None, - JsonCodec::::new(), - ); - - assert_eq!( - cache - .increment_cache("counter", 2.5, context.clone()) - .unwrap(), - 4.5 - ); - assert_eq!( - cache - .async_increment("counter", 2.5, context, false) - .await - .unwrap(), - 4.5 - ); -} - -#[rstest] -#[case::explicit_ttl(Some(Duration::from_secs(60)), 60u64)] -#[case::default_ttl(None, 600u64)] -#[tokio::test] -async fn refresh_ttl_increments_rearm_the_ttl_in_the_same_round_trip( - #[case] ttl: Option, - #[case] seconds: u64, -) { - let mut pipeline = redis::pipe(); - pipeline - .cmd("INCRBYFLOAT") - .arg("ns:spend:key:k") - .arg(1.5f64) - .cmd("EXPIRE") - .arg("ns:spend:key:k") - .arg(seconds); - let cache = json_cache(vec![MockCmd::with_values( - pipeline, - Ok(vec![ - redis::Value::BulkString(b"1.5".to_vec()), - redis::Value::Int(1), - ]), - )]) - .with_namespace(Some("ns".into())); - - assert_eq!( - cache - .async_increment("spend:key:k", 1.5, ExactCacheContext { ttl }, true) - .await - .unwrap(), - 1.5 - ); -} - -#[rstest] -#[tokio::test] -async fn counter_repairs_are_atomic_and_use_default_ttl() { - let floor = || { - MockCmd::new( - redis::cmd("EVAL") - .arg(INCREMENT_WITH_FLOOR_SCRIPT) - .arg(1) - .arg("team:counter") - .arg(-2i64) - .arg(30u64), - Ok(0i64), - ) - }; - let cache = team(vec![ - floor(), - floor(), - MockCmd::new( - redis::cmd("EVAL") - .arg(SET_MAX_SCRIPT) - .arg(1) - .arg("team:counter") - .arg(4.5f64) - .arg(600u64), - Ok("4.5"), - ), - ]); - - assert_eq!( - cache - .increment_with_floor("counter", -2, Duration::from_secs(30)) - .unwrap(), - 0 - ); - assert_eq!( - cache - .async_increment_with_floor("counter", -2, Duration::from_secs(30)) - .await - .unwrap(), - 0 - ); - assert_eq!( - cache.async_set_max("counter", 4.5, None).await.unwrap(), - 4.5 - ); -} - -const CLAIM_SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", - "elseif current ~= ARGV[1] then return 0; end; ", - "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ", - "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" -); - -fn claim_eval(expected: &str, write: &str, refresh: bool, applied: i64) -> MockCmd { - MockCmd::new( - redis::cmd("EVAL") - .arg(CLAIM_SCRIPT) - .arg(1) - .arg("pin") - .arg(expected) - .arg(600) - .arg(write) - .arg(u8::from(refresh)), - Ok(applied), - ) -} - -#[rstest] -#[tokio::test] -async fn claims_match_eligible_values_written_by_another_encoder(context: ExactCacheContext) { - let python_payload = r#"{"model_id": "a", "deployment": "east"}"#; - let stored = json!({"deployment": "east", "model_id": "a"}); - let cache = json_cache(vec![ - MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)), - claim_eval(python_payload, "", true, 1), - ]); - - assert_eq!( - cache - .async_claim_cache( - "pin", - json!({"model_id": "b"}), - vec![stored.clone()], - context - ) - .await - .unwrap(), - stored - ); -} - -#[rstest] -fn claims_retry_when_the_key_changes_and_replace_ineligible_winners(context: ExactCacheContext) { - let candidate = json!({"model_id": "b"}); - let payload = r#"{"model_id":"b"}"#; - let cache = json_cache(vec![ - MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)), - claim_eval("", payload, false, 0), - MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)), - claim_eval(r#"{"model_id":"gone"}"#, payload, false, 1), - ]); - - assert_eq!( - cache - .claim_cache( - "pin", - candidate.clone(), - &[json!({"model_id": "a"})], - context - ) - .unwrap(), - candidate - ); -} - -#[rstest] -fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl( - context: ExactCacheContext, -) { - let stored = r#"{"model_id": "a"}"#; - let cache = json_cache(vec![ - MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)), - claim_eval(stored, "", false, 1), - ]); - - assert_eq!( - cache - .claim_cache("pin", json!({"model_id": "b"}), &[], context) - .unwrap(), - json!({"model_id": "a"}) - ); -} - -#[rstest] -#[tokio::test] -async fn test_connection_reports_success_with_the_python_message() { - let cache = team(vec![MockCmd::new(redis::cmd("PING"), Ok("PONG"))]); - - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Success); - assert_eq!(result.message, "Redis connection test successful"); - assert_eq!(result.error, None); -} - -#[rstest] -#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", None)] -#[case::connection_refused( - Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))), - "Redis connection failed:", - Some("connection refused") -)] -#[tokio::test] -async fn test_connection_failures_use_the_python_result_contract( - #[case] reply: redis::RedisResult<&'static str>, - #[case] message: &str, - #[case] cause: Option<&str>, -) { - let cache = json_cache(vec![MockCmd::new(redis::cmd("PING"), reply)]); - - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Failed); - assert!(result.message.starts_with(message), "{}", result.message); - assert_eq!(result.error.is_some(), cause.is_some()); - assert!(cause.is_none_or(|cause| { - result - .error - .as_deref() - .is_some_and(|error| error.contains(cause)) - })); -} - -#[rstest] -#[tokio::test] -async fn disconnect_keeps_a_caller_owned_connection_usable() { - let cache = team(vec![MockCmd::new(redis::cmd("PING"), Ok("PONG"))]); - cache.disconnect().await.unwrap(); - assert!(cache.ping().await.unwrap()); -} - -#[rstest] -#[tokio::test] -async fn disconnect_drains_an_idle_pool_without_connecting() { - let cache = RedisCache::new("redis://127.0.0.1:1", None, JsonCodec::::new()).unwrap(); - cache.disconnect().await.unwrap(); -} diff --git a/litellm-rust/crates/cache-redis/tests/cluster.rs b/litellm-rust/crates/cache-redis/tests/cluster.rs deleted file mode 100644 index e1ba01bcb04..00000000000 --- a/litellm-rust/crates/cache-redis/tests/cluster.rs +++ /dev/null @@ -1,563 +0,0 @@ -//! Tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a comma -//! separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them. - -mod support; - -use std::{collections::HashSet, time::Duration}; - -use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheConnectionStatus, - CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, CounterCache, DeleteCache, - DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec, - PingCache, QueueCache, RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache, - TtlPipelineCache, -}; -use litellm_cache_redis::{ - RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation, - RedisTopology, -}; -use redis::cluster_routing::Slot; -use rstest::{fixture, rstest}; -use serde_json::json; -use support::{JsonCache, cluster_cache, cluster_url}; - -type Counter = RedisCache>; - -#[fixture] -fn cache(#[default("cache")] label: &str) -> Option { - cluster_cache(label, Duration::from_secs(120), JsonCodec::new()) -} - -#[fixture] -fn counter(#[default("counter")] label: &str) -> Option { - cluster_cache(label, Duration::from_secs(60), JsonCodec::new()) -} - -#[fixture] -fn context() -> ExactCacheContext { - ExactCacheContext::default() -} - -fn multi_slot_keys(count: usize) -> Vec { - let keys: Vec = (0..count).map(|index| format!("key-{index}")).collect(); - let slots: HashSet = keys.iter().map(Slot::for_key).collect(); - assert!(slots.len() > 1, "keys must span multiple slots"); - keys -} - -fn seconds(seconds: u64) -> Option { - Some(Duration::from_secs(seconds)) -} - -#[rstest] -#[case::no_startup_nodes("redis://127.0.0.1:7000", Vec::new())] -#[case::unix_socket_url( - "redis+unix:///tmp/redis.sock", - vec![RedisNode { host: "127.0.0.1".into(), port: 7000 }] -)] -fn constructor_rejects_unusable_cluster_configs( - #[case] url: &str, - #[case] startup_nodes: Vec, -) { - let error = JsonCache::connect( - url, - &RedisTopology::Cluster { startup_nodes }, - None, - JsonCodec::new(), - ) - .err(); - assert!(matches!(error, Some(Error::Unavailable(_)))); -} - -#[rstest] -#[tokio::test] -async fn single_key_operations_round_trip_with_ttl_rounding( - #[with("single")] cache: Option, -) { - let Some(cache) = cache else { return }; - let context = ExactCacheContext { - ttl: Some(Duration::from_millis(1500)), - }; - let keys = multi_slot_keys(12); - for (index, key) in keys.iter().enumerate() { - cache - .set_cache(key, json!({ "index": index }), &context) - .unwrap(); - } - for (index, key) in keys.iter().enumerate() { - assert_eq!( - cache.get_cache(key, &context).unwrap(), - Some(json!({ "index": index })) - ); - } - assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(2)); - assert!( - cache - .async_refresh_ttl(&keys[0], seconds(40)) - .await - .unwrap() - ); - assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(40)); - cache.delete_cache(&keys[0]).unwrap(); - assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None); - assert!(!cache.async_refresh_ttl(&keys[0], None).await.unwrap()); - assert!(cache.sync_ping().unwrap()); - cache.async_flush_cache().await.unwrap(); -} - -#[rstest] -#[tokio::test] -async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries( - #[with("batch")] cache: Option, - context: ExactCacheContext, -) { - let Some(cache) = cache else { return }; - let keys = multi_slot_keys(40); - for (index, key) in keys.iter().enumerate() { - if index % 5 == 0 { - continue; - } - cache - .async_set_cache(key, json!(index), context.clone()) - .await - .unwrap(); - } - let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()]) - .unwrap() - .get_connection() - .unwrap(); - let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]); - redis::cmd("SET") - .arg(&malformed) - .arg("not json") - .exec(&mut raw) - .unwrap(); - - let entries = cache - .async_batch_get_cache(keys.clone(), context.clone()) - .await - .unwrap(); - assert_eq!(entries.len(), keys.len()); - for (index, entry) in entries.iter().enumerate() { - let expected = if index == 1 { - BatchEntry::Invalid - } else if index % 5 == 0 { - BatchEntry::Miss - } else { - BatchEntry::Hit(json!(index)) - }; - assert_eq!(*entry, expected, "entry {index}"); - } - assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), entries); - - cache.delete_cache_keys(keys.clone()).await.unwrap(); - let entries = cache.async_batch_get_cache(keys, context).await.unwrap(); - assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss)); -} - -#[rstest] -#[tokio::test] -async fn pipelines_group_by_slot_and_return_results_in_submission_order( - #[with("pipeline")] cache: Option, - counter: Option, - context: ExactCacheContext, -) { - let (Some(cache), Some(counter)) = (cache, counter) else { - return; - }; - let keys = multi_slot_keys(30); - cache - .async_set_cache_pipeline( - keys.iter() - .enumerate() - .map(|(index, key)| (key.clone(), json!(index))) - .collect(), - context.clone(), - ) - .await - .unwrap(); - let hits = cache - .async_batch_get_cache(keys.clone(), context.clone()) - .await - .unwrap(); - assert!( - hits.iter() - .enumerate() - .all(|(index, entry)| *entry == BatchEntry::Hit(json!(index))) - ); - - cache - .async_set_cache_pipeline_with_ttls( - keys.iter() - .enumerate() - .map(|(index, key)| (key.clone(), json!(index), seconds(index as u64 + 10))) - .collect(), - ) - .await - .unwrap(); - for (index, key) in keys.iter().enumerate() { - assert_eq!( - cache.async_get_ttl(key).await.unwrap(), - seconds(index as u64 + 10), - "{key}" - ); - } - - let queues: Vec = keys.iter().map(|key| format!("queue:{key}")).collect(); - let pushed = cache - .async_rpush_pipeline( - queues - .iter() - .enumerate() - .map(|(index, key)| RedisRpushOperation { - key: key.clone(), - values: (0..=index) - .map(|value| RedisArg::Integer(value as i64)) - .collect(), - }) - .collect(), - ) - .await - .unwrap(); - assert_eq!(pushed, (1..=keys.len()).collect::>()); - let popped = cache - .async_lpop_pipeline( - queues - .iter() - .enumerate() - .map(|(index, key)| RedisLpopOperation { - key: key.clone(), - count: (index % 2 == 0).then_some(2), - }) - .collect(), - ) - .await - .unwrap(); - for (index, result) in popped.into_iter().enumerate() { - match result { - RedisLpopResult::Value(value) => { - assert_eq!(index % 2, 1, "queue {index}"); - assert_eq!(value, b"0"); - } - RedisLpopResult::Values(values) => { - assert_eq!(index % 2, 0, "queue {index}"); - let expected: Vec> = (0..=index) - .take(2) - .map(|value| value.to_string().into_bytes()) - .collect(); - assert_eq!(values, expected); - } - other => panic!("queue {index}: {other:?}"), - } - } - - let counters: Vec = keys.iter().map(|key| format!("counter:{key}")).collect(); - let totals = counter - .async_increment_pipeline( - counters - .iter() - .enumerate() - .map(|(index, key)| IncrementOperation { - key: key.clone(), - amount: index as f64 + 0.5, - ttl: (index % 3 == 0).then_some(Duration::from_secs(30)), - }) - .collect(), - ) - .await - .unwrap(); - let expected: Vec = (0..keys.len()).map(|index| index as f64 + 0.5).collect(); - assert_eq!(totals, expected); - assert_eq!( - counter.async_get_ttl(&counters[0]).await.unwrap(), - seconds(30) - ); - assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None); - counter.async_flush_cache().await.unwrap(); - cache.async_flush_cache().await.unwrap(); -} - -#[rstest] -#[tokio::test] -async fn rpush_and_trim_is_one_transaction_on_the_key_slot( - #[with("trim")] cache: Option, -) { - let Some(cache) = cache else { return }; - let values = |values: &[&str]| values.iter().map(|value| RedisArg::from(*value)).collect(); - assert_eq!( - cache - .async_rpush_and_trim("buf", values(&["a", "b"]), 3) - .await - .unwrap(), - 2 - ); - assert_eq!( - cache - .async_rpush_and_trim("buf", values(&["c", "d"]), 3) - .await - .unwrap(), - 4 - ); - assert_eq!( - cache.async_lpop("buf", Some(10)).await.unwrap(), - RedisLpopResult::Values(vec![b"b".to_vec(), b"c".to_vec(), b"d".to_vec()]) - ); - cache.async_flush_cache().await.unwrap(); -} - -#[rstest] -#[tokio::test] -async fn scan_and_scoped_flush_cover_every_primary( - #[with("flush")] cache: Option, - #[from(cache)] - #[with("other")] - other: Option, - context: ExactCacheContext, -) { - let (Some(cache), Some(other)) = (cache, other) else { - return; - }; - let keys = multi_slot_keys(60); - for key in &keys { - cache - .async_set_cache(key, json!(true), context.clone()) - .await - .unwrap(); - other - .async_set_cache(key, json!(true), context.clone()) - .await - .unwrap(); - } - let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap(); - scanned.sort(); - let mut expected: Vec = keys - .iter() - .map(|key| format!("{}:{key}", cache.namespace().unwrap())) - .collect(); - expected.sort(); - assert_eq!(scanned, expected); - assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7); - - cache.flush_cache().unwrap(); - let flushed = cache - .async_batch_get_cache(keys.clone(), context.clone()) - .await - .unwrap(); - assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss)); - let kept = other.async_batch_get_cache(keys, context).await.unwrap(); - assert!( - kept.iter() - .all(|entry| *entry == BatchEntry::Hit(json!(true))) - ); - other.async_flush_cache().await.unwrap(); -} - -fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> { - let mut connection = startup.get_connection().unwrap(); - let nodes: String = redis::cmd("CLUSTER") - .arg("NODES") - .query(&mut connection) - .unwrap(); - let mut counts: Vec<(String, u64)> = nodes - .lines() - .map(|line| { - let address = line.split_whitespace().nth(1).unwrap(); - let address = address.split('@').next().unwrap(); - let mut node = redis::Client::open(format!("redis://{address}")) - .unwrap() - .get_connection() - .unwrap(); - let stats: String = redis::cmd("INFO") - .arg("commandstats") - .query(&mut node) - .unwrap(); - let calls = stats - .lines() - .find_map(|stat| stat.strip_prefix("cmdstat_ping:calls=")) - .and_then(|rest| rest.split(',').next()) - .map_or(0, |calls| calls.parse().unwrap()); - (address.to_string(), calls) - }) - .collect(); - counts.sort(); - counts -} - -#[rstest] -#[tokio::test] -async fn ping_reaches_every_node(#[with("ping")] cache: Option) { - let Some(cache) = cache else { return }; - let startup = redis::Client::open(cluster_url()).unwrap(); - let before = ping_calls_per_node(&startup); - assert!(before.len() >= 2, "{before:?}"); - assert!(cache.ping().await.unwrap()); - let after = ping_calls_per_node(&startup); - for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) { - assert!(calls_after > calls_before, "{node} was not pinged"); - } - assert!(cache.sync_ping().unwrap()); - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Success); - assert_eq!(result.message, "Redis Cluster connection test successful"); -} - -#[rstest] -#[tokio::test] -async fn disconnect_closes_idle_connections_and_reconnects_on_demand( - #[with("disconnect")] cache: Option, -) { - let Some(cache) = cache else { return }; - assert!(cache.ping().await.unwrap()); - cache.disconnect().await.unwrap(); - assert!(cache.ping().await.unwrap()); -} - -#[rstest] -#[case::keep_existing_ttl(false)] -#[case::refresh_ttl(true)] -#[tokio::test] -async fn increments_refresh_the_ttl_only_when_asked( - counter: Option, - #[case] refresh_ttl: bool, -) { - let Some(counter) = counter else { return }; - let context = ExactCacheContext { ttl: seconds(60) }; - counter - .async_set_cache("spend", 0.0, ExactCacheContext { ttl: seconds(600) }) - .await - .unwrap(); - assert_eq!( - counter - .async_increment("spend", 1.5, context.clone(), refresh_ttl) - .await - .unwrap(), - 1.5 - ); - assert_eq!( - counter - .async_increment("spend", 2.0, context, refresh_ttl) - .await - .unwrap(), - 3.5 - ); - let ttl = counter.async_get_ttl("spend").await.unwrap().unwrap(); - assert_eq!(ttl <= Duration::from_secs(60), refresh_ttl, "{ttl:?}"); - counter.async_flush_cache().await.unwrap(); -} - -#[rstest] -#[tokio::test] -async fn counters_claims_scripts_and_sets_work_on_the_cluster( - counter: Option, - #[with("claim")] cache: Option, - context: ExactCacheContext, -) { - let (Some(counter), Some(cache)) = (counter, cache) else { - return; - }; - assert_eq!( - counter - .increment_cache("spend", 1.5, context.clone()) - .unwrap(), - 1.5 - ); - assert_eq!( - counter - .async_increment("spend", 2.0, context.clone(), false) - .await - .unwrap(), - 3.5 - ); - assert_eq!( - counter - .increment_with_floor("budget", -3, Duration::from_secs(30)) - .unwrap(), - 0 - ); - assert_eq!( - counter - .async_increment_with_floor("budget", 7, Duration::from_secs(30)) - .await - .unwrap(), - 7 - ); - assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0); - assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0); - counter.flush_cache().unwrap(); - - let owner = json!("owner-a"); - let rival = json!("owner-b"); - assert_eq!( - cache - .claim_cache("lock", owner.clone(), &[], context.clone()) - .unwrap(), - owner - ); - assert_eq!( - cache - .async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone()) - .await - .unwrap(), - owner - ); - assert_eq!( - cache - .claim_cache("lock", rival.clone(), &[], context.clone()) - .unwrap(), - owner - ); - assert_eq!( - cache - .async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone()) - .await - .unwrap(), - rival - ); - - let script = cache - .async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into()); - let reply = script - .invoke( - vec!["scripted".into()], - vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)], - ) - .await - .unwrap(); - assert_eq!(reply, redis::Value::Okay); - assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), seconds(5)); - let evaluated = cache - .async_eval( - "return redis.call('GET', KEYS[1])".into(), - vec!["scripted".into()], - Vec::new(), - ) - .await - .unwrap(); - assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec())); - - assert_eq!( - cache - .async_set_cache_sadd( - "members", - vec![ - RedisArg::Bytes(b"a".to_vec()), - RedisArg::Bytes(b"b".to_vec()) - ], - seconds(9), - ) - .await - .unwrap(), - 2 - ); - assert_eq!(cache.async_get_ttl("members").await.unwrap(), seconds(9)); - - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Success); - assert!(cache.ping().await.unwrap()); - let info = cache.info().unwrap(); - assert!(info.matches("redis_version").count() > 1, "{info}"); - assert!(cache.client_list().unwrap().contains("id=")); - cache.async_flush_cache().await.unwrap(); - assert_eq!(cache.async_get_ttl("members").await.unwrap(), None); - assert_eq!(cache.get_cache("lock", &context).unwrap(), None); -} diff --git a/litellm-rust/crates/cache-redis/tests/contract.rs b/litellm-rust/crates/cache-redis/tests/contract.rs deleted file mode 100644 index 85d043b90b2..00000000000 --- a/litellm-rust/crates/cache-redis/tests/contract.rs +++ /dev/null @@ -1,95 +0,0 @@ -//! The shared cache contracts, run against the in-process fake connection and, when -//! `LITELLM_TEST_REDIS_CLUSTER_NODES` is set, against a live Redis Cluster. - -mod support; - -use std::time::Duration; - -use litellm_cache::{ExactCacheContext, JsonCodec}; -use litellm_cache_testing as contract; -use rstest::rstest; -use serde_json::json; -use support::{JsonCache, cluster_cache, fake_cache}; - -const PREFIX: &str = "contract:"; - -#[derive(Clone, Copy, Debug)] -enum Contract { - HitAndMiss, - SyncAsyncEquivalence, - OverwriteReplaces, - PipelineWritesEveryEntry, - BatchPreservesOrder, - DeleteRemovesKey, - FlushClears, - CounterAccumulates, -} - -#[derive(Clone, Copy, Debug)] -enum Server { - Fake, - Cluster, -} - -async fn check(contract: Contract, cache: &JsonCache) -where - C: redis::ConnectionLike + Send + 'static, -{ - let context = ExactCacheContext::default(); - match contract { - Contract::HitAndMiss => { - contract::hit_and_miss(cache, context, PREFIX, json!({"answer": 42})).await - } - Contract::SyncAsyncEquivalence => { - contract::sync_async_equivalence(cache, context, PREFIX, json!("first"), json!([2])) - .await - } - Contract::OverwriteReplaces => { - contract::overwrite_replaces(cache, context, PREFIX, json!(1), json!({"b": 2})).await - } - Contract::PipelineWritesEveryEntry => { - contract::pipeline_writes_every_entry( - cache, - context, - PREFIX, - vec![json!("a"), json!(2), json!({"c": true})], - ) - .await - } - Contract::BatchPreservesOrder => { - contract::batch_preserves_order(cache, context, PREFIX, json!("first"), json!(2)).await - } - Contract::DeleteRemovesKey => { - contract::delete_removes_key(cache, context, PREFIX, json!("value")).await - } - Contract::FlushClears => { - contract::flush_clears(cache, context, PREFIX, json!("value")).await - } - Contract::CounterAccumulates => contract::counter_accumulates(cache, context, PREFIX).await, - } -} - -#[rstest] -#[case::hit_and_miss(Contract::HitAndMiss)] -#[case::sync_async_equivalence(Contract::SyncAsyncEquivalence)] -#[case::overwrite_replaces(Contract::OverwriteReplaces)] -#[case::pipeline_writes_every_entry(Contract::PipelineWritesEveryEntry)] -#[case::batch_preserves_order(Contract::BatchPreservesOrder)] -#[case::delete_removes_key(Contract::DeleteRemovesKey)] -#[case::flush_clears(Contract::FlushClears)] -#[case::counter_accumulates(Contract::CounterAccumulates)] -#[tokio::test] -async fn redis_satisfies_the_cache_contract( - #[case] contract: Contract, - #[values(Server::Fake, Server::Cluster)] server: Server, -) { - match server { - Server::Fake => check(contract, &fake_cache("contract")).await, - Server::Cluster => { - let label = format!("{contract:?}"); - if let Some(cache) = cluster_cache(&label, Duration::from_secs(120), JsonCodec::new()) { - check(contract, &cache).await; - } - } - } -} diff --git a/litellm-rust/crates/cache-redis/tests/support/mod.rs b/litellm-rust/crates/cache-redis/tests/support/mod.rs deleted file mode 100644 index b04e1c29f3b..00000000000 --- a/litellm-rust/crates/cache-redis/tests/support/mod.rs +++ /dev/null @@ -1,231 +0,0 @@ -#![allow(dead_code)] - -use std::{ - collections::BTreeMap, - time::{Duration, SystemTime, UNIX_EPOCH}, -}; - -use litellm_cache::{CacheCodec, Error, JsonCodec}; -use litellm_cache_redis::{RedisCache, RedisNode, RedisTopology}; - -/// Encodes a byte behind a tag, so a value written with another tag decodes as invalid. -pub struct TaggedByteCodec(pub u8); - -impl CacheCodec for TaggedByteCodec { - type Value = u8; - - fn encode(&self, value: &u8) -> Result, Error> { - if *value > 127 { - return Err(Error::InvalidEntry(None)); - } - Ok(vec![self.0, *value]) - } - - fn decode(&self, bytes: &[u8]) -> Result { - match bytes { - [tag, value] if *tag == self.0 => Ok(*value), - _ => Err(Error::InvalidEntry(None)), - } - } -} - -/// A stateful in-process stand-in for a Redis server that understands the string commands the -/// shared contracts exercise, so they run without a live server. TTLs are accepted and ignored. -#[derive(Default)] -pub struct FakeRedis { - strings: BTreeMap, Vec>, -} - -impl FakeRedis { - fn run(&mut self, command: Vec>) -> redis::RedisResult { - let name = String::from_utf8_lossy(&command[0]).to_ascii_uppercase(); - let args = &command[1..]; - Ok(match name.as_str() { - "PING" => redis::Value::SimpleString("PONG".into()), - "SET" | "SETEX" => { - let value = if name == "SET" { &args[1] } else { &args[2] }; - self.strings.insert(args[0].clone(), value.clone()); - redis::Value::Okay - } - "GET" => self.get(&args[0]), - "MGET" => redis::Value::Array(args.iter().map(|key| self.get(key)).collect()), - "DEL" => { - let removed = args - .iter() - .filter(|key| self.strings.remove(*key).is_some()) - .count(); - redis::Value::Int(removed as i64) - } - "SCAN" => { - let pattern = &args[2]; - let keys = self - .strings - .keys() - .filter(|key| glob(pattern, key)) - .map(|key| redis::Value::BulkString(key.clone())) - .collect(); - redis::Value::Array(vec![ - redis::Value::BulkString(b"0".to_vec()), - redis::Value::Array(keys), - ]) - } - "EVAL" if args[0].windows(11).any(|window| window == b"INCRBYFLOAT") => { - self.increment_by_float(&args[2], &args[3]) - } - "INCRBYFLOAT" => self.increment_by_float(&args[0], &args[1]), - _ => { - return Err(redis::RedisError::from(( - redis::ErrorKind::Client, - "unsupported command", - name, - ))); - } - }) - } - - fn get(&self, key: &[u8]) -> redis::Value { - self.strings.get(key).map_or(redis::Value::Nil, |value| { - redis::Value::BulkString(value.clone()) - }) - } - - fn increment_by_float(&mut self, key: &[u8], amount: &[u8]) -> redis::Value { - let current = self - .strings - .get(key) - .map_or(0.0, |value| parse_float(value)); - let total = format!("{}", current + parse_float(amount)); - self.strings - .insert(key.to_vec(), total.clone().into_bytes()); - redis::Value::BulkString(total.into_bytes()) - } -} - -fn parse_float(bytes: &[u8]) -> f64 { - std::str::from_utf8(bytes).unwrap().parse().unwrap() -} - -/// Redis `MATCH` globbing for `*`, `?` and backslash escapes. -fn glob(pattern: &[u8], key: &[u8]) -> bool { - match pattern.split_first() { - None => key.is_empty(), - Some((b'*', rest)) => (0..=key.len()).any(|skip| glob(rest, &key[skip..])), - Some((b'?', rest)) => !key.is_empty() && glob(rest, &key[1..]), - Some((b'\\', [escaped, rest @ ..])) => { - key.first() == Some(escaped) && glob(rest, &key[1..]) - } - Some((literal, rest)) => key.first() == Some(literal) && glob(rest, &key[1..]), - } -} - -/// Splits RESP request bytes into the commands they carry. -fn commands(mut bytes: &[u8]) -> Vec>> { - fn line<'a>(bytes: &mut &'a [u8]) -> &'a [u8] { - let end = bytes - .windows(2) - .position(|window| window == b"\r\n") - .unwrap(); - let (line, rest) = bytes.split_at(end); - *bytes = &rest[2..]; - line - } - fn length(line: &[u8]) -> usize { - std::str::from_utf8(&line[1..]).unwrap().parse().unwrap() - } - let mut commands = Vec::new(); - while !bytes.is_empty() { - let count = length(line(&mut bytes)); - let command = (0..count) - .map(|_| { - let size = length(line(&mut bytes)); - let (argument, rest) = bytes.split_at(size); - bytes = &rest[2..]; - argument.to_vec() - }) - .collect(); - commands.push(command); - } - commands -} - -impl redis::ConnectionLike for FakeRedis { - fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { - let command = commands(cmd).into_iter().next().unwrap(); - self.run(command) - } - - fn req_packed_commands( - &mut self, - cmd: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - let replies = commands(cmd) - .into_iter() - .map(|command| self.run(command)) - .collect::>>()?; - Ok(replies.into_iter().skip(offset).take(count).collect()) - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} - -pub type JsonCache = RedisCache, C>; - -pub fn fake_cache(namespace: &str) -> JsonCache { - RedisCache::with_connection(FakeRedis::default(), None, JsonCodec::new()) - .with_namespace(Some(namespace.into())) -} - -/// Startup nodes from `LITELLM_TEST_REDIS_CLUSTER_NODES` (`host:port,host:port`); tests that -/// need a live cluster skip when it is unset. -pub fn cluster_topology() -> Option { - let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?; - let startup_nodes = nodes - .split(',') - .map(|node| { - let (host, port) = node.trim().rsplit_once(':').expect("host:port"); - RedisNode { - host: host.to_string(), - port: port.parse().expect("port"), - } - }) - .collect(); - Some(RedisTopology::Cluster { startup_nodes }) -} - -pub fn cluster_url() -> String { - std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL") - .unwrap_or_else(|_| "redis://127.0.0.1:7000".into()) -} - -pub fn unique_namespace(label: &str) -> String { - let nanos = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_nanos(); - format!("cluster-test:{label}:{nanos}") -} - -pub fn cluster_cache( - label: &str, - default_ttl: Duration, - codec: S, -) -> Option> { - let topology = cluster_topology()?; - Some( - RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec) - .expect("cluster connection") - .with_namespace(Some(unique_namespace(label))), - ) -} diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml index 8f4570edc6f..46ef6ab92a8 100644 --- a/litellm-rust/crates/cache-response/Cargo.toml +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -16,9 +16,6 @@ sha2.workspace = true litellm-cache-gcs.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-memory.workspace = true -litellm-cache-redis.workspace = true -redis.workspace = true -redis-test.workspace = true rstest.workspace = true tokio.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/cache-response/tests/connection.rs b/litellm-rust/crates/cache-response/tests/connection.rs index bc24cf56846..f64057f6f1c 100644 --- a/litellm-rust/crates/cache-response/tests/connection.rs +++ b/litellm-rust/crates/cache-response/tests/connection.rs @@ -2,79 +2,11 @@ mod support; use std::{sync::Arc, time::Duration}; -use litellm_cache::CacheConnectionStatus; use litellm_cache_memory::InMemoryCache; -use litellm_cache_response::{ - CacheEntry, ConnectionProbe, ExactResponseCache, ResponseCache, ResponseCacheRequest, -}; -use redis_test::MockCmd; +use litellm_cache_response::{CacheEntry, ExactResponseCache, ResponseCache, ResponseCacheRequest}; use rstest::rstest; use serde_json::json; -use support::{keyed, memory, redis, request}; - -#[rstest] -#[case::reachable( - Ok("PONG"), - CacheConnectionStatus::Success, - "Redis connection test successful", - false -)] -#[case::unexpected_reply( - Ok("NOPE"), - CacheConnectionStatus::Failed, - "Redis ping returned False", - false -)] -#[case::connection_refused( - Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))), - CacheConnectionStatus::Failed, - "Redis connection failed:", - true -)] -#[tokio::test] -async fn connection_backends_are_reachable_as_a_probe( - #[case] reply: redis::RedisResult<&'static str>, - #[case] status: CacheConnectionStatus, - #[case] message: &str, - #[case] has_error: bool, -) { - let probe: Arc = - Arc::new(redis(vec![MockCmd::new(redis::cmd("PING"), reply)], None)); - - let result = probe.test_connection().await.unwrap(); - assert_eq!(result.status, status); - assert!(result.message.starts_with(message), "{}", result.message); - assert_eq!(result.error.is_some(), has_error); -} - -#[rstest] -#[tokio::test] -async fn one_service_serves_both_the_exact_cache_and_its_probe(request: ResponseCacheRequest) { - let service = Arc::new(redis( - vec![ - MockCmd::new(redis::cmd("PING"), Ok("PONG")), - MockCmd::new( - redis::cmd("GET").arg("tenant:key"), - Ok(br#"{"timestamp":100.0,"response":{"ok":true}}"#.to_vec()), - ), - ], - Some("tenant"), - )); - let probe: Arc = service.clone(); - let exact: Arc = service; - - assert_eq!( - probe.test_connection().await.unwrap().status, - CacheConnectionStatus::Success - ); - assert_eq!( - exact - .async_lookup(&request, Duration::from_secs(100)) - .await - .unwrap(), - Some(json!({"ok": true})) - ); -} +use support::{keyed, memory, request}; /// The in-memory backend has no `test_connection`, as in Python, and still serves every response /// operation. diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 70ae259d33a..659ba065cf2 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -17,10 +17,9 @@ use litellm_cache_response::{ CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest, WriteBuffer, cache_key, }; -use redis_test::MockCmd; use rstest::{fixture, rstest}; use serde_json::{Value, json}; -use support::{keyed, memory, redis, request}; +use support::{keyed, memory, request}; type Memory = Arc>>; @@ -331,79 +330,6 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent( ); } -#[rstest] -#[case::python_sync_literal( - br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice() -)] -#[case::python_async_json(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice())] -#[tokio::test] -async fn redis_consumer_reads_python_sync_and_async_envelopes( - request: ResponseCacheRequest, - #[case] stored: &[u8], - #[values(false, true)] asynchronous: bool, -) { - let cache = redis( - vec![MockCmd::new( - redis::cmd("GET").arg("tenant:key"), - Ok(stored.to_vec()), - )], - Some("tenant"), - ); - let now = Duration::from_secs(101); - let found = if asynchronous { - cache.async_lookup(&request, now).await.unwrap() - } else { - cache.lookup(&request, now).unwrap() - }; - assert_eq!(found, Some(json!({"ok": true, "text": "cached"}))); -} - -#[rstest] -#[case::object( - json!({"ok": true, "text": "cached"}), - br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice() -)] -#[case::array(json!([1, 2]), br#"{"timestamp":100.0,"response":"[1,2]"}"#.as_slice())] -#[tokio::test] -async fn redis_consumer_writes_python_compatible_json( - request: ResponseCacheRequest, - #[case] response: Value, - #[case] wire: &[u8], -) { - let cache = redis( - vec![MockCmd::new( - redis::cmd("SETEX").arg("tenant:key").arg(600).arg(wire), - Ok("OK"), - )], - Some("tenant"), - ); - cache - .async_store(&request, response, Duration::from_secs(100)) - .await - .unwrap(); -} - -#[rstest] -#[tokio::test] -async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis( - mut request: ResponseCacheRequest, -) { - let cache = redis( - vec![MockCmd::new( - redis::cmd("GET").arg("tenant:key"), - Ok(b"invalid".to_vec()), - )], - None, - ); - request.controls.no_cache = true; - assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None); - request.controls.no_cache = false; - assert_eq!( - cache.async_lookup(&request, Duration::ZERO).await.unwrap(), - None - ); -} - #[rstest] #[tokio::test] async fn captured_service_keeps_the_selected_backend_for_background_writes( @@ -541,27 +467,6 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses( assert_eq!(memory.lookup(&requests[2], now).unwrap(), None); } -#[rstest] -#[tokio::test] -async fn batch_lookup_with_no_readable_request_skips_the_backend( - #[values(false, true)] asynchronous: bool, -) { - let cache = redis(Vec::new(), None); - let mut request = keyed("key"); - request.controls.no_cache = true; - let requests = [request.clone(), request]; - let partial = if asynchronous { - cache - .async_lookup_batch(&requests, Duration::ZERO) - .await - .unwrap() - } else { - cache.lookup_batch(&requests, Duration::ZERO).unwrap() - }; - assert_eq!(partial.values, vec![None, None]); - assert_eq!(partial.missing_indices, vec![0, 1]); -} - #[rstest] #[tokio::test] async fn deferred_entries_keep_the_time_they_were_produced( diff --git a/litellm-rust/crates/cache-response/tests/support/mod.rs b/litellm-rust/crates/cache-response/tests/support/mod.rs index b992a937b70..f84da303659 100644 --- a/litellm-rust/crates/cache-response/tests/support/mod.rs +++ b/litellm-rust/crates/cache-response/tests/support/mod.rs @@ -1,15 +1,9 @@ use std::{sync::Arc, time::Duration}; use litellm_cache_memory::InMemoryCache; -use litellm_cache_redis::RedisCache; -use litellm_cache_response::{ - CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, -}; -use redis_test::{MockCmd, MockRedisConnection}; +use litellm_cache_response::{CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheRequest}; use rstest::fixture; -pub type MockedRedis = RedisCache; - #[fixture] pub fn memory() -> Arc>> { Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new( @@ -29,12 +23,3 @@ pub fn keyed(key: &str) -> ResponseCacheRequest { ..Default::default() }) } - -/// A Redis response cache that must receive exactly `commands`, in order. -pub fn redis(commands: Vec, namespace: Option<&str>) -> ResponseCache { - let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); - ResponseCache::new(Arc::new( - RedisCache::with_connection(connection, None, ResponseCacheCodec) - .with_namespace(namespace.map(str::to_owned)), - )) -} diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml index d51c282d4d7..b0382d62b44 100644 --- a/litellm-rust/crates/cache-s3/Cargo.toml +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -9,14 +9,10 @@ repository.workspace = true litellm-http.workspace = true litellm-cache.workspace = true litellm-auth-aws.workspace = true -aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] } -aws-credential-types = "1.3.0" -aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] } -aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] } -aws-types = "1.6.0" +litellm-s3.workspace = true +aws-sdk-s3.workspace = true +aws-smithy-types = "1.6.0" futures-util.workspace = true -http.workspace = true -reqwest.workspace = true tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs deleted file mode 100644 index 659390a5659..00000000000 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ /dev/null @@ -1,58 +0,0 @@ -use aws_credential_types::{ - Credentials as AwsCredentials, - provider::{ProvideCredentials, error::CredentialsError, future}, -}; -use litellm_auth_aws::{AwsAuthConfig, AwsAuthService}; - -#[derive(Clone)] -pub struct S3Credentials { - auth: AwsAuthService, - config: AwsAuthConfig, - env: fn(&str) -> Option, -} - -impl S3Credentials { - pub fn new(config: AwsAuthConfig) -> Self { - Self::with_env(config, |name| std::env::var(name).ok()) - } - - pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option) -> Self { - Self { - auth: AwsAuthService::default(), - config, - env, - } - } -} - -impl ProvideCredentials for S3Credentials { - fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> - where - Self: 'a, - { - future::ProvideCredentials::new(async { - if let (Some(access_key_id), Some(secret_access_key)) = ( - self.config.access_key_id.clone(), - self.config.secret_access_key.clone(), - ) { - return Ok(AwsCredentials::new( - access_key_id, - secret_access_key, - self.config.session_token.clone(), - None, - "litellm-s3-cache", - )); - } - self.auth - .resolve_credentials(self.config.clone(), &self.env) - .await - .map_err(CredentialsError::provider_error) - }) - } -} - -impl std::fmt::Debug for S3Credentials { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("S3Credentials").finish_non_exhaustive() - } -} diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index 5369bca0f18..7d5f3f50418 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -5,20 +5,18 @@ use std::{ }; use aws_sdk_s3::{ - config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation}, + config::{RequestChecksumCalculation, ResponseChecksumValidation}, error::SdkError, primitives::ByteStream, }; use aws_smithy_types::{DateTime, date_time::Format}; use futures_util::future::try_join_all; -use litellm_auth_aws::AwsAuthConfig; +use litellm_auth_aws::AwsCredentialsProvider; use litellm_cache::{ BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache, }; use tokio::runtime::Handle; -use crate::{auth::S3Credentials, transport::ReqwestHttpClient}; - pub struct S3Endpoint { pub url: String, } @@ -28,7 +26,7 @@ pub struct S3CacheConfig { pub key_prefix: String, pub region: String, pub endpoint: Option, - pub auth: AwsAuthConfig, + pub credentials: AwsCredentialsProvider, } pub struct S3Cache { @@ -49,17 +47,14 @@ impl S3Cache { runtime: Handle, ) -> Self { let endpoint_url: Option = config.endpoint.map(|endpoint| endpoint.url); - let base = aws_sdk_s3::Config::builder() - .behavior_version(BehaviorVersion::latest()) - .region(Region::new(config.region.clone())) - .http_client(ReqwestHttpClient(http)) - .credentials_provider(S3Credentials::new(config.auth)) - .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) - .response_checksum_validation(ResponseChecksumValidation::WhenRequired); - let builder = match &endpoint_url { - Some(url) => base.endpoint_url(url).force_path_style(true), - None => base, - }; + let builder = litellm_s3::config( + config.region.clone(), + endpoint_url.as_deref(), + http, + config.credentials, + ) + .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) + .response_checksum_validation(ResponseChecksumValidation::WhenRequired); Self { client: aws_sdk_s3::Client::from_conf(builder.build()), codec, diff --git a/litellm-rust/crates/cache-s3/src/lib.rs b/litellm-rust/crates/cache-s3/src/lib.rs index 9be0210aad3..43006532441 100644 --- a/litellm-rust/crates/cache-s3/src/lib.rs +++ b/litellm-rust/crates/cache-s3/src/lib.rs @@ -1,6 +1,3 @@ -mod auth; mod cache; -mod transport; -pub use auth::S3Credentials; pub use cache::{S3Cache, S3CacheConfig, S3Endpoint}; diff --git a/litellm-rust/crates/cache-s3/tests/auth.rs b/litellm-rust/crates/cache-s3/tests/auth.rs deleted file mode 100644 index 59c5aa2e091..00000000000 --- a/litellm-rust/crates/cache-s3/tests/auth.rs +++ /dev/null @@ -1,57 +0,0 @@ -use aws_credential_types::provider::ProvideCredentials; -use litellm_auth_aws::AwsAuthConfig; -use litellm_cache_s3::S3Credentials; -use rstest::rstest; - -fn explicit(session_token: Option<&str>) -> AwsAuthConfig { - AwsAuthConfig { - access_key_id: Some("key".to_string()), - secret_access_key: Some("secret".to_string()), - session_token: session_token.map(str::to_string), - region_name: Some("us-east-1".to_string()), - ..Default::default() - } -} - -fn ambient_token(name: &str) -> Option { - (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()) -} - -fn environment_keys(name: &str) -> Option { - match name { - "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()), - "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()), - "AWS_SESSION_TOKEN" => Some("env-token".to_string()), - _ => None, - } -} - -#[rstest] -#[case::explicit_keys_ignore_an_ambient_session_token( - explicit(None), ambient_token, ("key", "secret", None) -)] -#[case::explicit_keys_keep_their_session_token( - explicit(Some("t")), ambient_token, ("key", "secret", Some("t")) -)] -#[case::environment_keys_resolve_with_their_session_token( - AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token")) -)] -#[tokio::test] -async fn credentials_resolve( - #[case] config: AwsAuthConfig, - #[case] env: fn(&str) -> Option, - #[case] expected: (&str, &str, Option<&str>), -) { - let credentials = S3Credentials::with_env(config, env) - .provide_credentials() - .await - .unwrap(); - assert_eq!( - ( - credentials.access_key_id(), - credentials.secret_access_key(), - credentials.session_token(), - ), - expected - ); -} diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs index 7378dde21bf..7b62d0f7208 100644 --- a/litellm-rust/crates/cache-s3/tests/support/mod.rs +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -5,7 +5,7 @@ use std::{ sync::{Arc, Mutex}, }; -use litellm_auth_aws::AwsAuthConfig; +use litellm_auth_aws::{AwsAuthConfig, AwsAuthService, AwsCredentialsProvider}; use litellm_cache::JsonCodec; use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint}; use serde_json::Value; @@ -22,12 +22,16 @@ pub fn config(endpoint: &str) -> S3CacheConfig { endpoint: Some(S3Endpoint { url: endpoint.to_string(), }), - auth: AwsAuthConfig { - access_key_id: Some("key".to_string()), - secret_access_key: Some("secret".to_string()), - region_name: Some("us-east-1".to_string()), - ..Default::default() - }, + credentials: AwsCredentialsProvider::configured( + AwsAuthService::default(), + AwsAuthConfig { + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + |_| None, + ), } } diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml deleted file mode 100644 index 5b4e99de382..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml +++ /dev/null @@ -1,22 +0,0 @@ -[package] -name = "litellm-cache-valkey-semantic" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true - -[dependencies] -litellm-cache.workspace = true -litellm-cache-redis.workspace = true -redis = { workspace = true, features = ["tls-rustls"] } -sha2.workspace = true -thiserror.workspace = true -uuid = { version = "1", features = ["v4"] } - -[dev-dependencies] -litellm-error.workspace = true -litellm-cache-testing.workspace = true -redis-test.workspace = true -rstest.workspace = true -serde_json.workspace = true -tokio.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/src/cache.rs b/litellm-rust/crates/cache-valkey-semantic/src/cache.rs deleted file mode 100644 index 947c0edc479..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/cache.rs +++ /dev/null @@ -1,245 +0,0 @@ -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; - -use litellm_cache::{ - BaseCache, CacheCodec, Error, SemanticCacheContext, - semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context}, -}; -use litellm_cache_redis::{RedisTopology, connection::Connections}; - -use crate::{ - ValkeySemanticConfig, - index::IndexState, - search::{embedding_bytes, scope_tag, search_document, write_document}, -}; - -/// `ValkeySemanticCache`: a semantic cache on valkey-search's TAG + VECTOR index. Values go -/// through the injected codec, so the response layer decides what a cached entry is. -pub struct ValkeySemanticCache { - connections: Arc>, - embedder: E, - codec: S, - config: ValkeySemanticConfig, - index_dimension: Arc>>, -} - -impl ValkeySemanticCache { - pub fn new( - url: &str, - embedder: E, - codec: S, - config: ValkeySemanticConfig, - ) -> Result { - Ok(Self { - connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), - embedder, - codec, - config, - index_dimension: Arc::new(Mutex::new(None)), - }) - } -} - -impl ValkeySemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub fn with_connection( - connection: C, - embedder: E, - codec: S, - config: ValkeySemanticConfig, - ) -> Self { - Self { - connections: Arc::new(Connections::fixed(connection)), - embedder, - codec, - config, - index_dimension: Arc::new(Mutex::new(None)), - } - } - - pub fn similarity_threshold(&self) -> f64 { - self.config.similarity_threshold - } - - pub fn index_name(&self) -> &str { - &self.config.index_name - } - - fn index_state(&self) -> IndexState { - IndexState { - name: self.config.index_name.clone(), - prefix: format!("{}:", self.config.index_name), - dimension: Arc::clone(&self.index_dimension), - similarity_threshold: self.config.similarity_threshold, - } - } - - fn decode(&self, lookup: SemanticLookup>) -> Result, Error> { - Ok(SemanticLookup { - value: lookup - .value - .map(|bytes| self.codec.decode(&bytes)) - .transpose()?, - similarity: lookup.similarity, - }) - } -} - -impl ValkeySemanticCache -where - E: Embedder, - S: CacheCodec + Clone, - C: redis::ConnectionLike + Send + 'static, -{ - /// The same index and connections behind a different embedder. - pub fn with_embedder(&self, embedder: E2) -> ValkeySemanticCache { - ValkeySemanticCache { - connections: Arc::clone(&self.connections), - embedder, - codec: self.codec.clone(), - config: self.config.clone(), - index_dimension: Arc::clone(&self.index_dimension), - } - } -} - -impl BaseCache for ValkeySemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Value = S::Value; - type Context = SemanticCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl - } - - fn set_cache( - &self, - key: &str, - value: Self::Value, - context: &Self::Context, - ) -> Result<(), Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(()); - }; - let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let response = self.codec.encode(&value)?; - let index = self.index_state(); - self.connections.execute(|connection| { - write_document( - connection, - &index, - &scope_tag(key), - &prompt, - response, - embedding_bytes(&embedding), - self.get_ttl(context), - ) - }) - } - - fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { - self.get_cache_with_similarity(key, context) - .map(|lookup| lookup.value) - } - - async fn async_set_cache( - &self, - key: &str, - value: Self::Value, - context: Self::Context, - ) -> Result<(), Error> { - let Some(prompt) = prompt_from_context(&context) else { - return Ok(()); - }; - let embedding = self - .embedder - .async_embed(&prompt, context.metadata.as_ref()) - .await?; - let response = self.codec.encode(&value)?; - let index = self.index_state(); - let scope = scope_tag(key); - Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - write_document( - connection, - &index, - &scope, - &prompt, - response, - embedding_bytes(&embedding), - context.ttl, - ) - }) - .await - } - - async fn async_get_cache( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - self.async_get_cache_with_similarity(key, context) - .await - .map(|lookup| lookup.value) - } -} - -/// Python stamps a similarity of `0.0` when there is no prompt or no document in the key's -/// scope, and the closest document's similarity even when it misses the threshold. -impl SemanticCache for ValkeySemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn get_cache_with_similarity( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let index = self.index_state(); - let lookup = self.connections.execute(|connection| { - search_document( - connection, - &index, - &scope_tag(key), - embedding_bytes(&embedding), - ) - })?; - self.decode(lookup) - } - - async fn async_get_cache_with_similarity( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let embedding = self - .embedder - .async_embed(&prompt, context.metadata.as_ref()) - .await?; - let index = self.index_state(); - let scope = scope_tag(key); - let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - search_document(connection, &index, &scope, embedding_bytes(&embedding)) - }) - .await?; - self.decode(lookup) - } -} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/config.rs b/litellm-rust/crates/cache-valkey-semantic/src/config.rs deleted file mode 100644 index c417652c27b..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/config.rs +++ /dev/null @@ -1,8 +0,0 @@ -/// `ValkeySemanticCache.DEFAULT_VALKEY_INDEX_NAME`. -pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; - -#[derive(Clone, Debug, PartialEq)] -pub struct ValkeySemanticConfig { - pub similarity_threshold: f64, - pub index_name: String, -} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/index.rs b/litellm-rust/crates/cache-valkey-semantic/src/index.rs deleted file mode 100644 index 3dfff32ff7c..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/index.rs +++ /dev/null @@ -1,117 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use litellm_cache::Error; -use litellm_cache_redis::connection::ConnectionRef; - -use crate::search::value_text; - -/// The valkey-search index one cache writes to, with the dimension it was last ensured for. -#[derive(Clone)] -pub(crate) struct IndexState { - pub(crate) name: String, - pub(crate) prefix: String, - pub(crate) dimension: Arc>>, - pub(crate) similarity_threshold: f64, -} - -/// `_ensure_index_sync` / `_ensure_index_async`: create the TAG + HNSW index once per dimension, -/// and accept an existing index unless it reports a different dimension. -pub(crate) fn ensure_index( - connection: &mut ConnectionRef<'_>, - index: &IndexState, - dimension: usize, -) -> Result<(), Error> { - if index - .dimension - .lock()? - .is_some_and(|existing| existing == dimension) - { - return Ok(()); - } - let create = redis::cmd("FT.CREATE") - .arg(&index.name) - .arg("ON") - .arg("HASH") - .arg("PREFIX") - .arg(1) - .arg(&index.prefix) - .arg("SCHEMA") - .arg("litellm_cache_key") - .arg("TAG") - .arg("embedding") - .arg("VECTOR") - .arg("HNSW") - .arg(6) - .arg("TYPE") - .arg("FLOAT32") - .arg("DIM") - .arg(dimension) - .arg("DISTANCE_METRIC") - .arg("COSINE") - .query::(connection) - .map(|_| ()); - if let Err(error) = create { - if !index_exists(&error) { - return Err(Error::unavailable(error)); - } - let info = redis::cmd("FT.INFO") - .arg(&index.name) - .query::(connection) - .map_err(Error::unavailable)?; - if let Some(existing) = - index_dimension_from_info(&info).filter(|existing| *existing != dimension) - { - return Err(Error::unavailable(DimensionMismatch { - existing, - dimension, - })); - } - } - *index.dimension.lock()? = Some(dimension); - Ok(()) -} - -fn index_exists(error: &redis::RedisError) -> bool { - error.code().is_some() - && error - .detail() - .is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists")) -} - -#[derive(Debug, thiserror::Error)] -#[error("index exists with dimension {existing}, not {dimension}")] -struct DimensionMismatch { - existing: usize, - dimension: usize, -} - -/// `_extract_index_dim`: flatten each attribute one level and read the value after -/// `dimensions`. -fn index_dimension_from_info(value: &redis::Value) -> Option { - let redis::Value::Array(values) = value else { - return None; - }; - let attributes = values.windows(2).find_map(|pair| { - (value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1]) - })?; - let redis::Value::Array(fields) = attributes else { - return None; - }; - fields.iter().find_map(|field| { - let redis::Value::Array(values) = field else { - return None; - }; - let values = values - .iter() - .flat_map(|value| match value { - redis::Value::Array(values) => values.as_slice(), - _ => std::slice::from_ref(value), - }) - .collect::>(); - values.windows(2).find_map(|pair| { - (value_text(pair[0]).as_deref() == Some("dimensions")) - .then(|| value_text(pair[1]).and_then(|value| value.parse().ok())) - .flatten() - }) - }) -} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs deleted file mode 100644 index f2f62bc95bd..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ /dev/null @@ -1,7 +0,0 @@ -mod cache; -mod config; -mod index; -mod search; - -pub use cache::ValkeySemanticCache; -pub use config::{DEFAULT_INDEX_NAME, ValkeySemanticConfig}; diff --git a/litellm-rust/crates/cache-valkey-semantic/src/search.rs b/litellm-rust/crates/cache-valkey-semantic/src/search.rs deleted file mode 100644 index c5e4e748b58..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/search.rs +++ /dev/null @@ -1,161 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{Error, semantic::SemanticLookup}; -use litellm_cache_redis::connection::ConnectionRef; -use sha2::{Digest, Sha256}; -use uuid::Uuid; - -use crate::index::{IndexState, ensure_index}; - -/// `_scope_tag`: valkey-search TAG fields cannot match arbitrary keys verbatim, so scopes are -/// the key's lowercase SHA-256. -pub(crate) fn scope_tag(key: &str) -> String { - let digest = Sha256::digest(key.as_bytes()); - digest.iter().map(|byte| format!("{byte:02x}")).collect() -} - -pub(crate) fn embedding_bytes(embedding: &[f32]) -> Vec { - embedding - .iter() - .flat_map(|value| value.to_le_bytes()) - .collect() -} - -/// `HSET` a fresh `:` document, then `EXPIRE` it when a TTL is set. -pub(crate) fn write_document( - connection: &mut ConnectionRef<'_>, - index: &IndexState, - scope: &str, - prompt: &str, - response: Vec, - vector: Vec, - ttl: Option, -) -> Result<(), Error> { - ensure_index(connection, index, vector.len() / size_of::())?; - let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4()); - let mut pipeline = redis::pipe(); - pipeline - .cmd("HSET") - .arg(&document) - .arg("litellm_cache_key") - .arg(scope) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(response) - .arg("embedding") - .arg(vector) - .ignore(); - if let Some(ttl) = ttl { - pipeline - .cmd("EXPIRE") - .arg(&document) - .arg(ttl.as_secs()) - .ignore(); - } - pipeline.query::<()>(connection).map_err(Error::unavailable) -} - -/// The KNN-1 search within `scope`: the closest document's similarity, and its stored response -/// when that similarity reaches the threshold. No document reads as a similarity of `0.0`. -pub(crate) fn search_document( - connection: &mut ConnectionRef<'_>, - index: &IndexState, - scope: &str, - vector: Vec, -) -> Result>, Error> { - ensure_index(connection, index, vector.len() / size_of::())?; - let query = - format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); - let response = redis::cmd("FT.SEARCH") - .arg(&index.name) - .arg(query) - .arg("PARAMS") - .arg(2) - .arg("vec") - .arg(vector) - .arg("RETURN") - .arg(2) - .arg("response") - .arg("vector_distance") - .arg("DIALECT") - .arg(2) - .query::(connection) - .map_err(Error::unavailable)?; - let Some(fields) = search_fields(response)? else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let field = |name: &str| { - fields - .iter() - .find_map(|(field, value)| (field == name).then(|| value.clone())) - .ok_or(Error::InvalidEntry(None)) - }; - let response = field("response")?; - let similarity = 1.0 - parse_f64(&field("vector_distance")?)?; - Ok(SemanticLookup { - value: (similarity >= index.similarity_threshold).then_some(response), - similarity: Some(similarity), - }) -} - -type SearchFields = Vec<(String, Vec)>; - -fn search_fields(value: redis::Value) -> Result, Error> { - let redis::Value::Array(values) = value else { - return Err(Error::InvalidEntry(None)); - }; - let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?; - if total <= 0 || values.len() < 3 { - return Ok(None); - } - let redis::Value::Array(fields) = &values[2] else { - return Err(Error::InvalidEntry(None)); - }; - let (pairs, remainder) = fields.as_chunks::<2>(); - if !remainder.is_empty() { - return Err(Error::InvalidEntry(None)); - } - pairs - .iter() - .map(|pair| { - Ok(( - value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?, - value_bytes(&pair[1])?, - )) - }) - .collect::, Error>>() - .map(Some) -} - -fn parse_i64(value: &redis::Value) -> Result { - value_text(value) - .ok_or(Error::InvalidEntry(None))? - .parse() - .map_err(Error::invalid_entry) -} - -fn parse_f64(value: &[u8]) -> Result { - std::str::from_utf8(value) - .map_err(Error::invalid_entry)? - .parse() - .map_err(Error::invalid_entry) -} - -pub(crate) fn value_text(value: &redis::Value) -> Option { - match value { - redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), - redis::Value::SimpleString(value) => Some(value.clone()), - redis::Value::Int(value) => Some(value.to_string()), - _ => None, - } -} - -fn value_bytes(value: &redis::Value) -> Result, Error> { - match value { - redis::Value::BulkString(bytes) => Ok(bytes.clone()), - redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()), - redis::Value::Int(value) => Ok(value.to_string().into_bytes()), - _ => Err(Error::InvalidEntry(None)), - } -} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs b/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs deleted file mode 100644 index 10e32c0183b..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs +++ /dev/null @@ -1,430 +0,0 @@ -mod support; - -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; - -use litellm_cache::{ - BaseCache, Error, JsonCodec, SemanticCacheContext, - semantic::{PreparedEmbedding, SemanticCache, SemanticLookup}, -}; -use litellm_cache_valkey_semantic::{ - DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig, -}; -use redis_test::MockRedisConnection; -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; -use support::{EmbedCalls, FakeEmbedder, RecordingConnection}; - -type Requests = Arc>>>; -type RecordingCache = ValkeySemanticCache, RecordingConnection>; - -const KEY_SCOPE: &str = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"; - -struct Recording { - cache: RecordingCache, - requests: Requests, - calls: EmbedCalls, -} - -impl Recording { - fn text(&self) -> String { - self.requests - .lock() - .unwrap() - .iter() - .map(|request| String::from_utf8_lossy(request).into_owned()) - .collect::>() - .join("\n") - } -} - -fn config() -> ValkeySemanticConfig { - ValkeySemanticConfig { - similarity_threshold: 0.8, - index_name: "test".into(), - } -} - -fn recording(replies: impl IntoIterator>) -> Recording { - let connection = RecordingConnection::new(replies); - let requests = connection.requests(); - let embedder = FakeEmbedder::new(&[]); - let calls = Arc::clone(&embedder.calls); - Recording { - cache: ValkeySemanticCache::with_connection( - connection, - embedder, - JsonCodec::new(), - config(), - ), - requests, - calls, - } -} - -#[fixture] -fn entry() -> Value { - json!({"timestamp": 1.0, "response": {"answer": "ok"}}) -} - -#[fixture] -fn context() -> SemanticCacheContext { - SemanticCacheContext { - messages: Some(json!([{"role": "user", "content": "hello"}])), - metadata: Some(json!({"source": "test"})), - ..Default::default() - } -} - -fn ok() -> redis::RedisResult { - Ok(redis::Value::SimpleString("OK".into())) -} - -fn already_exists() -> redis::RedisResult { - Err(support::server_error("Index already exists")) -} - -fn bulk(value: &[u8]) -> redis::Value { - redis::Value::BulkString(value.to_vec()) -} - -fn search_hit(response: &[u8], distance: &str) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::Int(1), - bulk(b"test:document"), - redis::Value::Array(vec![ - bulk(b"response"), - bulk(response), - bulk(b"vector_distance"), - bulk(distance.as_bytes()), - ]), - ]) -} - -fn encoded(value: &Value) -> Vec { - serde_json::to_vec(value).unwrap() -} - -/// `FT.INFO` with the vector field's dimension nested one level down, as valkey-search reports. -fn nested_dimension_info(dimension: i64) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::SimpleString("attributes".into()), - redis::Value::Array(vec![redis::Value::Array(vec![ - redis::Value::SimpleString("embedding".into()), - redis::Value::Array(vec![ - redis::Value::SimpleString("dimensions".into()), - redis::Value::Int(dimension), - ]), - ])]), - ]) -} - -/// `FT.INFO` with `dimensions` as a sibling string of the identifier. -fn flat_dimension_info(dimension: &str) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::SimpleString("attributes".into()), - redis::Value::Array(vec![redis::Value::Array(vec![ - redis::Value::SimpleString("identifier".into()), - redis::Value::SimpleString("embedding".into()), - redis::Value::Array(vec![ - redis::Value::SimpleString("dimensions".into()), - redis::Value::SimpleString(dimension.into()), - ]), - ])]), - ]) -} - -#[rstest] -#[case::string_content(json!([{"content": "hello"}]), None, Some("hello"))] -#[case::text_parts(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))] -#[case::non_object_parts_are_skipped(json!([{"content": ["raw", {"text": "hello"}]}]), None, Some("hello"))] -#[case::search_results(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))] -#[case::responses_string_input(json!([]), Some(json!(" hello ")), Some("hello"))] -#[case::responses_item_input(json!([]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))] -#[case::blank_input(json!([]), Some(json!(" ")), None)] -fn prompt_shapes_follow_redis_semantic_extraction( - #[case] messages: Value, - #[case] input: Option, - #[case] expected: Option<&str>, -) { - let recording = recording([ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))]); - let context = SemanticCacheContext { - messages: Some(messages), - input, - ..Default::default() - }; - - assert_eq!(recording.cache.get_cache("key", &context).unwrap(), None); - - let prompts = recording - .calls - .lock() - .unwrap() - .iter() - .map(|(prompt, _)| prompt.clone()) - .collect::>(); - assert_eq!( - prompts, - expected.into_iter().map(str::to_owned).collect::>() - ); - assert_eq!( - recording.requests.lock().unwrap().is_empty(), - expected.is_none() - ); -} - -#[rstest] -#[case::key("key", KEY_SCOPE)] -#[case::empty_key("", "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")] -fn documents_are_scoped_by_the_keys_sha256( - #[case] key: &str, - #[case] scope: &str, - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([ok()]); - - recording.cache.set_cache(key, entry, &context).unwrap(); - - let text = recording.text(); - assert!(text.contains(&format!("test:{scope}:"))); - assert!(text.contains(&format!("litellm_cache_key\r\n$64\r\n{scope}\r\n"))); -} - -#[rstest] -#[case::no_ttl(None, None)] -#[case::whole_seconds(Some(Duration::from_secs(5)), Some("5"))] -#[case::fractional_seconds_truncate(Some(Duration::from_millis(1900)), Some("1"))] -fn set_writes_hset_and_expires_only_with_a_ttl( - #[case] ttl: Option, - #[case] expire: Option<&str>, - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([ok()]); - - recording - .cache - .set_cache( - "key", - entry, - &SemanticCacheContext { - ttl, - ..context.clone() - }, - ) - .unwrap(); - - let text = recording.text(); - assert!(text.contains("FT.CREATE")); - assert!(text.contains("HSET")); - let expire_seconds = text - .split_once("EXPIRE\r\n") - .and_then(|(_, rest)| rest.split("\r\n").nth(3)); - assert_eq!(expire_seconds, expire); - assert_eq!( - *recording.calls.lock().unwrap(), - vec![("hello".to_owned(), context.metadata)] - ); -} - -#[rstest] -fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: SemanticCacheContext) { - let recording = recording([ok()]); - - recording - .cache - .set_cache("key", entry.clone(), &context) - .unwrap(); - recording.cache.set_cache("key", entry, &context).unwrap(); - - let text = recording.text(); - assert_eq!(text.matches("FT.CREATE").count(), 1); - assert_eq!(text.matches("HSET").count(), 2); -} - -#[rstest] -#[case::nested_matching(nested_dimension_info(3), Ok(()))] -#[case::nested_mismatch( - nested_dimension_info(2), - Err("cache is unavailable: index exists with dimension 2, not 3") -)] -#[case::flat_matching(flat_dimension_info("3"), Ok(()))] -#[case::flat_mismatch( - flat_dimension_info("2"), - Err("cache is unavailable: index exists with dimension 2, not 3") -)] -#[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))] -fn existing_index_dimension_must_match_embedding( - #[case] info: redis::Value, - #[case] expected: Result<(), &str>, - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([already_exists(), Ok(info)]); - - let result = recording.cache.set_cache("key", entry, &context); - assert_eq!( - result.map_err(|error| litellm_error::Report::new(&error).to_string()), - expected.map_err(str::to_owned) - ); -} - -#[rstest] -#[case::create_failure(Err(redis::RedisError::from((redis::ErrorKind::Io, "boom"))))] -fn index_creation_failures_are_unavailable( - #[case] reply: redis::RedisResult, - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([reply]); - - assert!(matches!( - recording.cache.set_cache("key", entry, &context), - Err(Error::Unavailable(Some(_))) - )); -} - -#[rstest] -#[case::within_threshold(search_hit(&encoded(&entry()), "0.1"), Ok(Some(entry())))] -#[case::at_threshold(search_hit(&encoded(&entry()), "0.2"), Ok(Some(entry())))] -#[case::beyond_threshold(search_hit(&encoded(&entry()), "0.5"), Ok(None))] -#[case::zero_documents(redis::Value::Array(vec![redis::Value::Int(0)]), Ok(None))] -#[case::missing_response( - redis::Value::Array(vec![ - redis::Value::Int(1), - bulk(b"document"), - redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]), - ]), - Err(Error::InvalidEntry(None)) -)] -#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))] -#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))] -fn get_applies_threshold_and_decodes_entry( - #[case] reply: redis::Value, - #[case] expected: Result, Error>, - context: SemanticCacheContext, -) { - let recording = recording([ok(), Ok(reply)]); - - let kind = |result: Result, Error>| { - result.map_err(|error| std::mem::discriminant(&error)) - }; - assert_eq!( - kind(recording.cache.get_cache("key", &context)), - kind(expected) - ); -} - -#[rstest] -#[case::hit(context(), Some(search_hit(&encoded(&entry()), "0.1")), Some(entry()), Some(1.0 - 0.1))] -#[case::below_threshold(context(), Some(search_hit(&encoded(&entry()), "0.5")), None, Some(1.0 - 0.5))] -#[case::no_results(context(), Some(redis::Value::Array(vec![redis::Value::Int(0)])), None, Some(0.0))] -#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))] -#[tokio::test] -async fn lookup_reports_python_semantic_similarity( - #[case] context: SemanticCacheContext, - #[case] reply: Option, - #[case] value: Option, - #[case] similarity: Option, - #[values(false, true)] use_async: bool, -) { - let searched = reply.is_some(); - let recording = recording(reply.map_or_else(Vec::new, |reply| vec![ok(), Ok(reply)])); - - let lookup = if use_async { - recording - .cache - .async_get_cache_with_similarity("key", &context) - .await - } else { - recording.cache.get_cache_with_similarity("key", &context) - }; - - assert_eq!(lookup, Ok(SemanticLookup { value, similarity })); - assert_eq!(recording.text().contains("FT.SEARCH"), searched); -} - -#[rstest] -#[tokio::test] -async fn missing_prompt_does_not_touch_valkey(entry: Value) { - let cache = ValkeySemanticCache::with_connection( - MockRedisConnection::new([]).assert_all_commands_consumed(), - FakeEmbedder::new(&[]), - JsonCodec::new(), - config(), - ); - let context = SemanticCacheContext::default(); - - cache.set_cache("key", entry.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key", &context).unwrap(), None); - cache - .async_set_cache("key", entry, context.clone()) - .await - .unwrap(); - assert_eq!(cache.async_get_cache("key", &context).await.unwrap(), None); - assert_eq!(cache.get_ttl(&context), None); -} - -#[rstest] -fn with_embedder_shares_index_state_and_connections(entry: Value, context: SemanticCacheContext) { - let recording = recording([ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]); - recording - .cache - .set_cache("key", entry.clone(), &context) - .unwrap(); - - let prepared = recording - .cache - .with_embedder(PreparedEmbedding(vec![0.1, 0.2, 0.3])); - - assert_eq!(prepared.get_cache("key", &context).unwrap(), Some(entry)); - assert_eq!(recording.text().matches("FT.CREATE").count(), 1); -} - -#[rstest] -fn accessors_report_the_config() { - let recording = recording([]); - - assert_eq!(recording.cache.index_name(), "test"); - assert_eq!(recording.cache.similarity_threshold(), 0.8); - assert_eq!(DEFAULT_INDEX_NAME, "litellm_semantic_cache_index"); -} - -#[rstest] -#[tokio::test] -async fn async_set_and_get_use_shared_document_helpers( - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([ok(), ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]); - let context = SemanticCacheContext { - ttl: Some(Duration::from_millis(1900)), - ..context - }; - - recording - .cache - .async_set_cache("key", entry.clone(), context.clone()) - .await - .unwrap(); - assert_eq!( - recording - .cache - .async_get_cache("key", &context) - .await - .unwrap(), - Some(entry) - ); - - let text = recording.text(); - assert!(text.contains("FT.CREATE")); - assert!(text.contains("HSET")); - assert!(text.contains("EXPIRE")); - assert_eq!( - *recording.calls.lock().unwrap(), - vec![("hello".to_owned(), context.metadata.clone()); 2] - ); -} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs b/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs deleted file mode 100644 index 71081fca92b..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs +++ /dev/null @@ -1,62 +0,0 @@ -//! `overwrite_replaces` does not apply: like Python, every write is a new `:` -//! document, so a second write with the same prompt adds a tie instead of replacing the first. - -mod support; - -use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding}; -use litellm_cache_testing as contract; -use litellm_cache_valkey_semantic::{ - DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig, -}; -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; -use support::FakeSearch; - -type Cache = ValkeySemanticCache, FakeSearch>; - -const PREFIX: &str = "contract:"; - -#[fixture] -fn cache() -> Cache { - ValkeySemanticCache::with_connection( - FakeSearch::default(), - PreparedEmbedding(vec![0.6, 0.8]), - JsonCodec::new(), - ValkeySemanticConfig { - similarity_threshold: 0.9, - index_name: DEFAULT_INDEX_NAME.into(), - }, - ) -} - -#[fixture] -fn context() -> SemanticCacheContext { - SemanticCacheContext { - messages: Some(json!([{"role": "user", "content": "contract prompt"}])), - ..Default::default() - } -} - -#[rstest] -#[tokio::test] -async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) { - contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await; -} - -#[rstest] -#[tokio::test] -async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) { - contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await; -} - -#[rstest] -#[tokio::test] -async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) { - contract::pipeline_writes_every_entry( - &cache, - context, - PREFIX, - vec![json!("a"), json!(2), json!({"c": true})], - ) - .await; -} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs deleted file mode 100644 index 2af05943c98..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs +++ /dev/null @@ -1,352 +0,0 @@ -#![allow(dead_code)] - -use std::{ - collections::{BTreeMap, HashMap, VecDeque}, - sync::{Arc, Mutex}, -}; - -use litellm_cache::{Error, semantic::Embedder}; -use serde_json::Value; - -pub type EmbedCalls = Arc)>>>; - -/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every -/// prompt with its metadata. -pub struct FakeEmbedder { - vectors: HashMap>, - pub calls: EmbedCalls, -} - -impl FakeEmbedder { - pub fn new(vectors: &[(&str, &[f32])]) -> Self { - Self { - vectors: vectors - .iter() - .map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec())) - .collect(), - calls: EmbedCalls::default(), - } - } -} - -impl Embedder for FakeEmbedder { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - self.calls - .lock() - .unwrap() - .push((prompt.to_owned(), metadata.cloned())); - Ok(self - .vectors - .get(prompt) - .cloned() - .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) - } - - async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - self.embed(prompt, metadata) - } -} - -struct FakeIndex { - prefix: Vec, - dims: usize, - vector_field: String, -} - -#[derive(Default)] -struct SearchState { - indexes: HashMap, - hashes: BTreeMap, BTreeMap>>, -} - -/// An in-memory valkey-search speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache -/// sends, with exact cosine KNN over the hashes under an index prefix. -#[derive(Clone, Default)] -pub struct FakeSearch { - state: Arc>, -} - -impl FakeSearch { - fn run(&self, args: Vec>) -> redis::RedisResult { - let mut state = self.state.lock().unwrap(); - let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned(); - match text(0).to_uppercase().as_str() { - "FT.CREATE" => { - let name = text(1); - if state.indexes.contains_key(&name) { - return Err(server_error("Index already exists")); - } - let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes()); - let prefix = args[position("PREFIX").unwrap() + 2].clone(); - let dims = text(position("DIM").unwrap() + 1).parse().unwrap(); - let vector_field = text(position("VECTOR").unwrap() - 1); - state.indexes.insert( - name, - FakeIndex { - prefix, - dims, - vector_field, - }, - ); - Ok(redis::Value::Okay) - } - "FT.INFO" => { - let index = state - .indexes - .get(&text(1)) - .ok_or_else(|| server_error("Unknown index name"))?; - Ok(index_info(index)) - } - "FT.DROPINDEX" => { - state.indexes.remove(&text(1)); - Ok(redis::Value::Okay) - } - "HSET" => { - let hash = state.hashes.entry(args[1].clone()).or_default(); - for pair in args[2..].chunks(2) { - hash.insert( - String::from_utf8_lossy(&pair[0]).into_owned(), - pair[1].clone(), - ); - } - Ok(redis::Value::Int(((args.len() - 2) / 2) as i64)) - } - "EXPIRE" => Ok(redis::Value::Int(i64::from( - state.hashes.contains_key(&args[1]), - ))), - "FT.SEARCH" => { - let index = state - .indexes - .get(&text(1)) - .ok_or_else(|| server_error("no such index"))?; - let query = text(2); - let tag = query_tag(&query); - let params = args.iter().position(|arg| arg == b"PARAMS").unwrap(); - let vector = floats(&args[params + 3]); - let best = state - .hashes - .iter() - .filter(|(key, _)| key.starts_with(&index.prefix)) - .filter(|(_, fields)| { - fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes()) - }) - .filter_map(|(key, fields)| { - let stored = floats(fields.get(&index.vector_field)?); - (stored.len() == index.dims) - .then(|| (key, fields, 1.0 - cosine(&vector, &stored))) - }) - .min_by(|left, right| left.2.total_cmp(&right.2)); - let Some((key, fields, distance)) = best else { - return Ok(redis::Value::Array(vec![redis::Value::Int(0)])); - }; - let mut reply = fields - .iter() - .filter(|(name, _)| **name != index.vector_field) - .flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)]) - .collect::>(); - reply.extend([ - bulk(b"vector_distance"), - bulk(distance.to_string().as_bytes()), - ]); - Ok(redis::Value::Array(vec![ - redis::Value::Int(1), - bulk(key), - redis::Value::Array(reply), - ])) - } - "PING" => Ok(redis::Value::SimpleString("PONG".into())), - _ => Err(server_error("unsupported command")), - } - } -} - -impl redis::ConnectionLike for FakeSearch { - fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { - let mut commands = parse_commands(command); - self.run(commands.remove(0)) - } - - fn req_packed_commands( - &mut self, - commands: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - let replies = parse_commands(commands) - .into_iter() - .map(|args| self.run(args)) - .collect::>>()?; - Ok(replies.into_iter().skip(offset).take(count).collect()) - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} - -pub fn server_error(reply: &str) -> redis::RedisError { - redis::parse_redis_value(format!("-{reply}\r\n").as_bytes()) - .expect("error reply") - .extract_error() - .expect_err("server error") -} - -fn bulk(bytes: &[u8]) -> redis::Value { - redis::Value::BulkString(bytes.to_vec()) -} - -fn index_info(index: &FakeIndex) -> redis::Value { - redis::Value::Array(vec![ - bulk(b"index_name"), - bulk(b"fake"), - bulk(b"attributes"), - redis::Value::Array(vec![ - redis::Value::Array(vec![ - bulk(b"identifier"), - bulk(b"litellm_cache_key"), - bulk(b"type"), - bulk(b"TAG"), - ]), - redis::Value::Array(vec![ - bulk(b"identifier"), - bulk(index.vector_field.as_bytes()), - bulk(b"type"), - bulk(b"VECTOR"), - bulk(b"index"), - redis::Value::Array(vec![ - bulk(b"dimensions"), - redis::Value::Int(index.dims as i64), - ]), - ]), - ]), - ]) -} - -/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed. -fn query_tag(query: &str) -> String { - let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len(); - let mut tag = String::new(); - let mut characters = query[start..].chars(); - while let Some(character) = characters.next() { - match character { - '\\' => tag.extend(characters.next()), - '}' => break, - character => tag.push(character), - } - } - tag -} - -fn floats(bytes: &[u8]) -> Vec { - bytes - .as_chunks::<4>() - .0 - .iter() - .map(|chunk| f32::from_le_bytes(*chunk)) - .collect() -} - -fn cosine(left: &[f32], right: &[f32]) -> f64 { - let dot = left - .iter() - .zip(right) - .map(|(left, right)| f64::from(*left) * f64::from(*right)) - .sum::(); - let norm = |vector: &[f32]| { - vector - .iter() - .map(|value| f64::from(*value).powi(2)) - .sum::() - .sqrt() - }; - dot / (norm(left) * norm(right)) -} - -/// Splits a packed RESP request into each command's arguments. -fn parse_commands(mut bytes: &[u8]) -> Vec>> { - let line = |bytes: &mut &[u8]| { - let end = bytes - .windows(2) - .position(|window| window == b"\r\n") - .unwrap(); - let text = String::from_utf8(bytes[1..end].to_vec()).unwrap(); - *bytes = &bytes[end + 2..]; - text.parse::().unwrap() - }; - let mut commands = Vec::new(); - while !bytes.is_empty() { - let count = line(&mut bytes); - let mut args = Vec::with_capacity(count); - for _ in 0..count { - let length = line(&mut bytes); - args.push(bytes[..length].to_vec()); - bytes = &bytes[length + 2..]; - } - commands.push(args); - } - commands -} - -/// Records every packed request and answers from a script, `OK` once the script runs out. -pub struct RecordingConnection { - requests: Arc>>>, - replies: Mutex>>, -} - -impl RecordingConnection { - pub fn new(replies: impl IntoIterator>) -> Self { - Self { - requests: Arc::default(), - replies: Mutex::new(replies.into_iter().collect()), - } - } - - pub fn requests(&self) -> Arc>>> { - Arc::clone(&self.requests) - } - - fn reply(&self) -> redis::RedisResult { - self.replies - .lock() - .unwrap() - .pop_front() - .unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into()))) - } -} - -impl redis::ConnectionLike for RecordingConnection { - fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { - self.requests.lock().unwrap().push(command.to_vec()); - self.reply() - } - - fn req_packed_commands( - &mut self, - command: &[u8], - _offset: usize, - count: usize, - ) -> redis::RedisResult> { - self.requests.lock().unwrap().push(command.to_vec()); - (0..count).map(|_| self.reply()).collect() - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index 1cd05710366..0133af16b80 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -16,7 +16,7 @@ serde_json.workspace = true thiserror.workspace = true [dev-dependencies] -litellm-auth.workspace = true +litellm-secret-value.workspace = true proptest.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index f65516265c5..8a212223523 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -27,18 +27,42 @@ "api_key", "additional_args" ], + "messages_input": [ + "body", + "kwargs" + ], "post_call": [ "logger", "original_response", "api_key", "additional_args" ], + "record_response": [ + "logger", + "status", + "headers", + "body" + ], + "record_api_window": [ + "logger", + "start", + "end" + ], "defers_async_logging": [ "logger" ], + "claim_async_success_log": [ + "logger" + ], + "mark_async_success_enqueued": [ + "logger" + ], "defer_success": [ "logger", - "pending" + "pending", + "response", + "start", + "end" ], "sync_success_for_async_call": [ "logger", @@ -91,6 +115,13 @@ "logger", "head" ], + "stamp_stream_head": [ + "head", + "logger", + "kwargs", + "start_time", + "end_time" + ], "stream_success": [ "logger", "url_route", @@ -99,7 +130,8 @@ "chunks", "start", "end", - "first_chunk" + "first_chunk", + "stream_teardown" ], "stream_failure": [ "logger", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 493e91bb5ea..f68541c5c7b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -14,7 +14,7 @@ use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::{PyDateTime, PyDict, PyList}, + types::{PyDateTime, PyDict, PyList, PyString}, }; use serde_json::Value; @@ -22,7 +22,7 @@ use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, finalize, is_internal_call, - python::Streaming, + python::{Logging, Streaming}, setup, }; @@ -57,6 +57,8 @@ pub struct LegacyLogging { asynchronous: bool, internal: bool, cache_key: Option, + /// When the provider request went out, until its window is recorded (Messages only). + api_started: Option, } fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult> { @@ -87,6 +89,94 @@ impl LegacyLogging { asynchronous, internal: false, cache_key: None, + api_started: None, + } + } + + /// The kwargs `Logging.update_from_kwargs` sees. Python's Messages handler takes `metadata` as + /// its own argument, so the logger never sees the caller's raw metadata, only the validated + /// request metadata merged in through `litellm_params`. Python updates the logger once, before + /// the call; the refresh once the result is ready leaves out `litellm_metadata` too, so the + /// logged metadata is not rebuilt from what the call wrote into it since. + fn logging_kwargs<'py>(&self, py: Python<'py>, refresh: bool) -> PyResult> { + let kwargs = self.call.kwargs().bind(py).copy()?; + if self.operation != LoggingOperation::Messages { + return Ok(kwargs); + } + let consumed: &[&str] = if refresh { + &["metadata", "litellm_metadata"] + } else { + &["metadata"] + }; + for name in consumed { + if kwargs.contains(name)? { + kwargs.del_item(name)?; + } + } + Ok(kwargs) + } + + /// Python's Messages handler hands the logger its request params, validated `metadata` + /// included, as `litellm_params`, beside the `aanthropic_messages` call flag. + fn logged_request_params( + &self, + context: &RequestContext, + ) -> Option> { + if self.operation != LoggingOperation::Messages { + return None; + } + let mut params = context + .optional_params + .as_object() + .cloned() + .unwrap_or_default(); + params.insert( + "aanthropic_messages".to_owned(), + Value::Bool(self.asynchronous), + ); + Some(params) + } + + // Compared as Python compares them, so a caller's `0` is kept over a typed float's `0.0`. + fn caller_value<'py>( + &self, + py: Python<'py>, + name: &str, + sent: &Value, + ) -> PyResult>> { + match self.call.lookup(py, name)? { + Some(value) if from_py::(&value).is_ok() && value.eq(to_py(py, sent)?)? => { + Ok(Some(value)) + } + _ => Ok(None), + } + } + + // Python's Messages handler logs the caller's own optional params, not their typed copies. + fn as_called(&self, py: Python<'_>, context: &RequestContext) -> PyResult { + let mut context = context.clone(); + if self.operation != LoggingOperation::Messages { + return Ok(context); + } + for (name, sent) in context + .optional_params + .as_object_mut() + .into_iter() + .flatten() + { + if let Some(value) = self.caller_value(py, name, sent)? { + *sent = from_py(&value)?; + } + } + Ok(context) + } + + /// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync + /// client is untimed. + fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> { + match self.api_started.take() { + Some(start) => self.logger()?.record_api_window(py, start, end), + None => Ok(()), } } @@ -102,12 +192,29 @@ impl LegacyLogging { } } - fn input_description(&self) -> &'static str { - match self.operation { + fn logged_input<'py>( + &self, + py: Python<'py>, + body: &Bound<'py, PyDict>, + ) -> PyResult> { + let description = match self.operation { LoggingOperation::Completion => "Chat completions", LoggingOperation::Responses => "Responses", - LoggingOperation::Messages => "Messages", + LoggingOperation::Messages => { + return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py))); + } LoggingOperation::Ocr => "OCR document processing", + }; + Ok(PyString::new(py, description).into_any()) + } + + /// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key. + fn logged_api_key<'a>(&self, context: &'a RequestContext) -> Option<&'a str> { + match self.operation { + LoggingOperation::Messages => Some(""), + LoggingOperation::Completion | LoggingOperation::Responses | LoggingOperation::Ocr => { + context.api_key.as_ref().map(|api_key| api_key.expose()) + } } } @@ -185,30 +292,49 @@ impl LegacyLogging { if !self.asynchronous { return pending().sync(py); } - if !self.internal - && self + if self.internal { + return Ok(()); + } + let defers = logger.defers_async_logging(py); + if logger.claim_async_success_log(py)? { + let enqueues = self .call .kwargs() .bind(py) .get_item("fallbacks")? - .is_none_or(|value| value.is_none()) - { - if logger.defers_async_logging(py) { + .is_none_or(|value| value.is_none()); + if defers { let pending = Py::new( py, PendingLogging { pending: Some(pending()), + enqueues, }, )?; - logger.defer_success(py, pending.bind(py).as_any())?; - } else { + return logger.defer_success( + py, + pending.bind(py).as_any(), + &self.response, + &self.start, + &self.end, + ); + } + if enqueues { pending().asynchronous(py)?; } } + if defers { + return Ok(()); + } logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) } - fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> { + fn stream_success( + &self, + py: Python<'_>, + stream: &DeliveredStream, + teardown: bool, + ) -> PyResult<()> { let logger = self.logger()?; let billing = self.stream_billing().ok_or(HostError::MissingState)?; let billed = Streaming::Success.call( @@ -225,6 +351,7 @@ impl LegacyLogging { &self.start, &self.end, &stream.first_chunk, + teardown, ), ); match billed { @@ -376,18 +503,30 @@ impl LegacyLogging { wire: Box, context: &RequestContext, ) -> PyResult>> { + let context = &self.as_called(py, context)?; let logger = self.logger()?; - logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?; + logger.update_from_kwargs( + py, + &self.logging_kwargs(py, false)?, + &wire, + context, + self.logged_request_params(context).as_ref(), + )?; let body = to_py(py, &wire.body)? .into_bound(py) .cast_into::()?; for (name, sent) in wire.body.as_object().into_iter().flatten() { - if let Some(value) = self.call.lookup(py, name)? - && from_py::(&value).is_ok_and(|caller| caller == *sent) - { + if let Some(value) = self.caller_value(py, name, sent)? { body.set_item(name, value)?; } } + if self.operation == LoggingOperation::Messages { + // Python's Messages handler merges its transformed request body into the call details. + logger + .object(py) + .getattr("model_call_details")? + .call_method1("update", (&body,))?; + } let headers = PyDict::new(py); for (name, value) in &wire.headers { headers.set_item(name, value)?; @@ -397,14 +536,18 @@ impl LegacyLogging { headers: headers.clone().unbind(), context: context.clone(), }); + let input = self.logged_input(py, &body)?; self.logger()?.pre_call( py, - self.input_description(), - context.api_key.as_ref().map(|api_key| api_key.expose()), + &input, + self.logged_api_key(context), &body, &headers, &wire.url, )?; + if self.operation == LoggingOperation::Messages && self.asynchronous { + self.api_started = Some(epoch_seconds()); + } let headers = headers .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) @@ -450,12 +593,16 @@ impl LegacyLogging { .getattr("litellm_params")? .cast_into::()? .copy()?; - params.set_item("custom_llm_provider", &facts.provider.provider)?; + if self.operation != LoggingOperation::Messages { + // Python's Messages handler hands the provider to the logger beside its params, never + // in them, and the proxy's logger carries over to the next fallback hop. + params.set_item("custom_llm_provider", &facts.provider.provider)?; + } crate::python::Logging::Update.call( py, ( &logger, - self.call.kwargs(), + self.logging_kwargs(py, true)?, &facts.provider.model, logger.getattr("optional_params")?, params, @@ -473,8 +620,17 @@ impl LegacyLogging { ResultSource::Provider => None, ResultSource::Cache { key } => Some(key.clone()), }; - details.set_item("cache_hit", self.cache_key.is_some())?; + if self.cache_key.is_some() { + details.set_item("cache_hit", true)?; + } else { + details.call_method1("pop", ("cache_hit", py.None()))?; + } details.set_item("cache_key", self.cache_key.as_deref())?; + if let (LoggingOperation::Messages, ResultSource::Provider, Some(request)) = + (self.operation, &facts.source, &self.request) + { + details.call_method1("update", (request.body.bind(py),))?; + } Ok(HookStep::Ready(())) } @@ -483,6 +639,16 @@ impl LegacyLogging { py: Python<'_>, raw: &RawResponse, ) -> PyResult> { + if self.operation == LoggingOperation::Messages { + // Python's Messages handler never calls `post_call` at receipt: it stores the raw + // response, and logging replays it through `AnthropicConfig.transform_response`. + self.close_api_window(py, epoch_seconds())?; + if let Some(head) = &raw.head { + self.logger()? + .record_response(py, head.status, &head.headers, &raw.body)?; + } + return Ok(HookStep::Ready(())); + } let api_key = self .request .as_ref() @@ -507,7 +673,7 @@ impl LegacyLogging { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); match &self.stream { - Some(stream) => self.stream_success(py, stream)?, + Some(stream) => self.stream_success(py, stream, false)?, None => self.dispatch_success(py)?, } Ok(HookStep::Ready(())) @@ -522,6 +688,9 @@ impl LegacyLogging { ) -> PyResult> { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); + if self.logger.is_some() { + self.close_api_window(py, timing.end_time)?; + } if self.stream.is_some() { return self.stream_failure(py); } @@ -543,6 +712,20 @@ impl LegacyLogging { head.bind(py).set_item("cache_key", key)?; head.bind(py).set_item("cache_hit", true)?; } + self.close_api_window(py, epoch_seconds())?; + if self.operation == LoggingOperation::Messages { + // Python's `@client` stamps the call metadata onto the stream it hands back. + Streaming::StampHead.call( + py, + ( + head.bind(py), + self.logger()?.object(py), + self.call.kwargs(), + &self.start, + datetime(py, epoch_seconds())?, + ), + )?; + } Streaming::Opened.call(py, (self.logger()?.object(py), head.bind(py)))?; self.stream = Some(DeliveredStream { chunks: PyList::empty(py).unbind(), @@ -560,8 +743,28 @@ impl LegacyLogging { } } +impl LegacyLogging { + // A caller that goes away mid-stream cancels the call before it ends: Python's Messages + // stream iterator still bills the chunks it read, straight away rather than deferred. + fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> { + if !self.asynchronous || self.end.is_some() { + return Ok(()); + } + let Some(stream) = self.stream.take() else { + return Ok(()); + }; + self.end = Some(datetime(py, epoch_seconds())?); + self.stream_success(py, &stream, true) + } +} + impl PythonOwned for LegacyLogging { fn close(&mut self, py: Python<'_>) { + if self.logger.is_some() + && let Err(error) = self.bill_abandoned_stream(py) + { + error.write_unraisable(py, None); + } if let Some(logger) = self.logger.take() && let Err(error) = logger.restore_context(py) { @@ -925,11 +1128,11 @@ assert all(value is failure for name, value in logger.calls if name.endswith('_h mod payload_tests { use std::ffi::CStr; - use litellm_auth::SecretValue; use litellm_host::hooks::CallHooks; use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest}; use litellm_host::lifecycle::ExecutionEvent; use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py}; + use litellm_secret_value::SecretValue; use proptest::prelude::*; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -1012,10 +1215,10 @@ logger._provider_reported_cost = None }; let previous = litellm_host::interceptors::ExecutionFacts { provider: facts.provider.clone(), - cost: litellm_host::interceptors::Cost::Reported { - amount: serde_json::Number::from_f64(9.0).unwrap(), + cost: litellm_host::interceptors::Cost::CacheHit, + source: litellm_host::interceptors::ResultSource::Cache { + key: "earlier-hit".into(), }, - source: litellm_host::interceptors::ResultSource::Provider, }; logging.result_ready(py, &previous).unwrap(); local(&locals, "logger") @@ -1033,13 +1236,59 @@ logger._provider_reported_cost = None c" assert logger.model_call_details.get('response_cost') is None, logger.model_call_details assert logger._provider_reported_cost == reported, logger._provider_reported_cost -assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details +assert logger.model_call_details.get('cache_hit') is (True if cache_hit else None), logger.model_call_details assert logger.update['custom_llm_provider'] == 'edenai', logger.update ", ); }); } + #[rstest] + #[case::messages(crate::LoggingOperation::Messages, false)] + #[case::ocr(crate::LoggingOperation::Ocr, true)] + fn result_ready_puts_the_provider_in_the_logged_params_except_for_messages( + #[case] operation: crate::LoggingOperation, + #[case] in_params: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + run( + py, + &locals, + c" +logger.litellm_params = {} +logger.optional_params = {} +logger.model_call_details = {} +logger._provider_reported_cost = None +", + ); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + operation, + ..legacy_call(py, &locals, false) + }; + let facts = litellm_host::interceptors::ExecutionFacts { + provider: litellm_host::interceptors::ProviderIdentity { + model: "native-test".into(), + provider: "anthropic".into(), + }, + cost: litellm_host::interceptors::Cost::Deferred, + source: litellm_host::interceptors::ResultSource::Provider, + }; + logging.result_ready(py, &facts).unwrap(); + locals.set_item("in_params", in_params).unwrap(); + run( + py, + &locals, + c" +assert logger.update['custom_llm_provider'] == 'anthropic', logger.update +assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params, logger.update +", + ); + }); + } + const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; @@ -1117,9 +1366,7 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update let step = logging .before_provider_request(py, Box::new(wire), context) .unwrap(); - let raw = RawResponse { - body: "raw response".into(), - }; + let raw = RawResponse::new("raw response".into()); assert!(matches!( logging .on_event( @@ -1176,7 +1423,10 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update #[rstest] #[case::completion(crate::LoggingOperation::Completion, "Chat completions")] #[case::responses(crate::LoggingOperation::Responses, "Responses")] - #[case::messages(crate::LoggingOperation::Messages, "Messages")] + #[case::messages( + crate::LoggingOperation::Messages, + r#"[{'role': 'user', 'content': '{"pages": [1]}'}]"# + )] #[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")] fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases( #[case] operation: crate::LoggingOperation, @@ -1217,7 +1467,7 @@ prepared = {'pages': replacement} &locals, c" assert logger.pre['complete_input_dict']['pages'] is replacement -assert logger.pre_input == description +assert str(logger.pre_input) == description assert original == [0] ", ); @@ -1773,10 +2023,15 @@ mod terminal_tests { #[case::async_listened( true, c"", - &["async_success_handler", "enqueued", "sync_success_for_async_call"] + &["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call"] )] - #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] + #[case::async_deferred(true, c"logger._defer_async_logging = True", &[])] #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + #[case::async_deferred_with_fallbacks( + true, + c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}", + &[] + )] fn success_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, @@ -1802,12 +2057,88 @@ mod terminal_tests { &locals, c" assert all(value is response for name, value in logger.calls if name.endswith('_handler')) -assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) +assert hasattr(logger, 'enqueue_deferred_logging') == getattr(logger, '_defer_async_logging', False) +assert getattr(logger, 'deferred_success_result', (response,))[0] is response ", ); }); } + #[rstest] + #[case::accepted(c"", &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"])] + #[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + fn a_deferred_success_runs_sync_callbacks_only_once_released( + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"response = object()\nlogger._defer_async_logging = True", + ); + run(py, &locals, script); + let mut logging = logged(py, &locals, true); + succeed(py, &locals, &mut logging); + run( + py, + &locals, + c" +assert logger.names() == [], logger.calls +logger.enqueue_deferred_logging() +logger.enqueue_deferred_logging() +", + ); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + + #[rstest] + #[case::deferred( + c"logger._defer_async_logging = True", + c"logger.enqueue_deferred_logging()", + &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"] + )] + #[case::deferred_and_blocked(c"logger._defer_async_logging = True", c"", &[])] + #[case::immediate( + c"", + c"", + &["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call", "sync_success_for_async_call"] + )] + fn a_fallback_attempt_on_the_same_logger_leaves_the_first_success_claimed( + #[case] script: &CStr, + #[case] release: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"refused = object()\nresponse = refused"); + run(py, &locals, script); + let mut first = logged(py, &locals, true); + succeed(py, &locals, &mut first); + run(py, &locals, c"response = object()"); + let mut fallback = logged(py, &locals, true); + succeed(py, &locals, &mut fallback); + run(py, &locals, release); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c"assert all(value is refused for name, value in logger.calls if name == 'async_success_handler')", + ); + }); + } + #[rstest] fn dropped_observations_preserve_deferred_success_and_response_identity() { Python::initialize(); @@ -1833,10 +2164,10 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy )); assert_eq!(sender.dropped_events(), 1); run(py, &locals, c" -assert logger.names() == ['sync_success_for_async_call'], logger.calls -logger._native_pending_logging.release(True) -logger._native_pending_logging.release(True) -assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls +assert logger.names() == [], logger.calls +logger.enqueue_deferred_logging() +logger.enqueue_deferred_logging() +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls assert logger.calls[0][1] is response assert logger.calls[1][1] is response "); @@ -1869,9 +2200,10 @@ assert logger.calls[1][1] is response py, &locals, c" -assert logger.names() == ['stream_opened', 'stream_success'], logger.calls +assert logger.names() == ['stamp_stream_head', 'stream_opened', 'stream_success'], logger.calls assert logger.calls[0][1] is head -chunks = logger.calls[1][1] +assert logger.calls[1][1] is head +chunks = logger.calls[2][1] assert len(chunks) == 2 assert chunks[0] is first assert chunks[1] is last @@ -1880,6 +2212,45 @@ assert chunks[1] is last }); } + #[rstest] + #[case::cancelled_mid_stream(None, &["stamp_stream_head", "stream_opened", "stream_success", "restore"])] + #[case::after_success(Some(true), &["stamp_stream_head", "stream_opened", "stream_success", "restore"])] + #[case::after_failure(Some(false), &["stamp_stream_head", "stream_opened", "stream_failure", "restore"])] + fn closing_an_unfinished_stream_bills_its_delivered_chunks_once( + #[case] ended: Option, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"first = b'first'\nresponse = None\nfailure = ValueError('provider')\nhead = {}", + ); + let mut logging = LegacyLogging { + operation: crate::LoggingOperation::Messages, + ..logged(py, &locals, true) + }; + logging + .on_stream_open(py, &local(&locals, "head").unbind()) + .unwrap(); + logging + .on_stream_chunk(py, &local(&locals, "first").unbind()) + .unwrap(); + match ended { + Some(true) => drop(succeed(py, &locals, &mut logging)), + Some(false) => drop(fail(py, &locals, &mut logging)), + None => {} + } + litellm_host_python::PythonOwned::close(&mut logging, py); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + #[rstest] #[case::synchronous(false, &["failure_handler"])] #[case::asynchronous(true, &[])] @@ -1908,7 +2279,7 @@ assert chunks[1] is last } #[rstest::rstest] - fn internal_async_calls_skip_the_async_success_fan_out() { + fn internal_async_calls_skip_success_logging() { Python::initialize(); Python::attach(|py| { let locals = namespace(py, c"response = object()"); @@ -1917,11 +2288,7 @@ assert chunks[1] is last ..logged(py, &locals, true) }; succeed(py, &locals, &mut logging); - run( - py, - &locals, - c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", - ); + run(py, &locals, c"assert logger.calls == [], logger.calls"); }); } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index f185f21a9c3..093e86c43dc 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -5,6 +5,7 @@ use litellm_host::interceptors::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; +use serde_json::Value; use crate::logger::PythonLogger; use crate::python::{Logging, Wrapper}; @@ -15,22 +16,37 @@ pub trait LegacyCallbacks { fn update_from_kwargs( &self, py: Python<'_>, - kwargs: &Py, + kwargs: &Bound<'_, PyDict>, wire: &WireRequest, context: &RequestContext, + request_params: Option<&serde_json::Map>, ) -> PyResult<()>; /// `Logging.pre_call`. fn pre_call( &self, py: Python<'_>, - input: &str, + input: &Bound<'_, PyAny>, api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, ) -> PyResult<()>; + /// Records the provider's raw HTTP response on the logger, as Python's Messages handler + /// does with `httpx_response` and `response_headers`. + fn record_response( + &self, + py: Python<'_>, + status: u16, + headers: &[(String, String)], + body: &str, + ) -> PyResult<()>; + + /// Records one provider HTTP exchange's window, as Python's timed `AsyncHTTPHandler.post` + /// does with `llm_api_duration_ms` and the metadata's `llm_api_timing_windows`. + fn record_api_window(&self, py: Python<'_>, start: f64, end: f64) -> PyResult<()>; + /// `Logging.post_call`. fn post_call( &self, @@ -43,7 +59,16 @@ pub trait LegacyCallbacks { fn defers_async_logging(&self, py: Python<'_>) -> bool; - fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>; + fn claim_async_success_log(&self, py: Python<'_>) -> PyResult; + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()>; fn sync_success_for_async_call( &self, @@ -83,12 +108,13 @@ impl LegacyCallbacks for PythonLogger { fn update_from_kwargs( &self, py: Python<'_>, - kwargs: &Py, + kwargs: &Bound<'_, PyDict>, wire: &WireRequest, context: &RequestContext, + request_params: Option<&serde_json::Map>, ) -> PyResult<()> { let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect(); - let redacted_kwargs = redact(py, kwargs.bind(py), &secret_fields)?; + let redacted_kwargs = redact(py, kwargs, &secret_fields)?; let optional_params = redact( py, &to_py(py, &context.optional_params)? @@ -97,23 +123,23 @@ impl LegacyCallbacks for PythonLogger { &secret_fields, )?; let params = PyDict::new(py); - params.set_item( - "litellm_call_id", - kwargs.bind(py).get_item("litellm_call_id")?, - )?; + params.set_item("litellm_call_id", kwargs.get_item("litellm_call_id")?)?; params.set_item("api_base", &wire.url)?; for name in ["logger_fn", "litellm_request_debug"] { - if let Some(value) = kwargs.bind(py).get_item(name)? { + if let Some(value) = kwargs.get_item(name)? { params.set_item(name, value)?; } } for name in custom_pricing_fields(py)? { - if let Some(value) = kwargs.bind(py).get_item(&name)? + if let Some(value) = kwargs.get_item(&name)? && !value.is_none() { params.set_item(name, value)?; } } + for (name, value) in request_params.into_iter().flatten() { + params.set_item(name, to_py(py, value)?)?; + } Logging::Update.call( py, ( @@ -131,7 +157,7 @@ impl LegacyCallbacks for PythonLogger { fn pre_call( &self, py: Python<'_>, - input: &str, + input: &Bound<'_, PyAny>, api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, @@ -145,6 +171,22 @@ impl LegacyCallbacks for PythonLogger { Ok(()) } + fn record_response( + &self, + py: Python<'_>, + status: u16, + headers: &[(String, String)], + body: &str, + ) -> PyResult<()> { + Logging::RecordResponse.call(py, (self.object(py), status, headers.to_vec(), body))?; + Ok(()) + } + + fn record_api_window(&self, py: Python<'_>, start: f64, end: f64) -> PyResult<()> { + Logging::RecordApiWindow.call(py, (self.object(py), start, end))?; + Ok(()) + } + fn post_call( &self, py: Python<'_>, @@ -170,8 +212,21 @@ impl LegacyCallbacks for PythonLogger { .unwrap_or(false) } - fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> { - Logging::DeferSuccess.call(py, (self.object(py), pending))?; + fn claim_async_success_log(&self, py: Python<'_>) -> PyResult { + Logging::ClaimAsyncSuccess + .call(py, (self.object(py),))? + .extract() + } + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + Logging::DeferSuccess.call(py, (self.object(py), pending, response, start, end))?; Ok(()) } @@ -219,6 +274,7 @@ impl LegacyCallbacks for PythonLogger { ) -> PyResult<()> { let coroutine = Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?; + Logging::MarkAsyncSuccessEnqueued.call(py, (self.object(py),))?; let enqueue = Logging::Enqueue.call(py, (&coroutine,)); if enqueue.is_err() && let Err(error) = coroutine.call_method0("close") diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index 292648ba69c..ea72191ddb2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -22,21 +22,35 @@ impl PendingSuccess { self.logger .enqueue_success(py, &self.response, &self.start, &self.end) } + + pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .sync_success_for_async_call(py, &self.response, &self.start, &self.end) + } } #[pyclass] pub(crate) struct PendingLogging { pub(crate) pending: Option, + pub(crate) enqueues: bool, } #[pymethods] impl PendingLogging { - fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { - let pending = slf.borrow_mut().pending.take(); - if let Some(pending) = pending - && success - { - match pending.asynchronous(py) { + fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> { + let (pending, enqueues) = { + let mut slot = slf.borrow_mut(); + (slot.pending.take(), slot.enqueues) + }; + if let Some(pending) = pending { + let released = pending.sync_for_async_call(py).and_then(|()| { + if enqueues { + pending.asynchronous(py) + } else { + Ok(()) + } + }); + match released { Err(error) if error.is_instance_of::(py) => { error.write_unraisable(py, Some(pending.logger.object(py))); } @@ -87,6 +101,7 @@ mod tests { start: py.None(), end: Some(py.None()), }), + enqueues: true, }, ) .unwrap(); @@ -108,7 +123,7 @@ observed = [] def on_enqueue(coroutine): observed.append(marker.get()) - pending.release(True) + pending() logger.on_enqueue = on_enqueue ", @@ -118,28 +133,12 @@ logger.on_enqueue = on_enqueue &locals, c" marker.set('release') -pending.release(True) -pending.release(True) +pending() +pending() assert observed == ['release'], observed -assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls assert logger.calls[0][1] is response -", - ); - }); - } - - #[test] - fn a_blocked_release_drops_the_success_for_good() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -pending.release(False) -pending.release(True) -assert logger.calls == [], logger.calls +assert logger.calls[1][1] is response ", ); }); @@ -168,7 +167,7 @@ logger.on_enqueue = on_enqueue locals .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) .unwrap(); - let released = local(&locals, "pending").call_method1("release", (true,)); + let released = local(&locals, "pending").call0(); match released { Ok(_) => assert!(!propagates), Err(error) => { @@ -181,8 +180,8 @@ logger.on_enqueue = on_enqueue py, &locals, c" -pending.release(True) -assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls +pending() +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued', 'closed'], logger.calls assert unraisable_from(logger) == ([] if propagates else [failure]) ", ); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6f1ee9dfd78..5b6faa3dbc2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -42,10 +42,18 @@ pub(crate) enum Logging { Update, #[strum(serialize = "pre_call")] PreCall, + #[strum(serialize = "messages_input")] + MessagesInput, #[strum(serialize = "post_call")] PostCall, + #[strum(serialize = "record_response")] + RecordResponse, + #[strum(serialize = "record_api_window")] + RecordApiWindow, #[strum(serialize = "defers_async_logging")] DefersAsync, + #[strum(serialize = "claim_async_success_log")] + ClaimAsyncSuccess, #[strum(serialize = "defer_success")] DeferSuccess, #[strum(serialize = "sync_success_for_async_call")] @@ -54,6 +62,8 @@ pub(crate) enum Logging { SubmitSuccess, #[strum(serialize = "async_success_handler")] AsyncSuccessHandler, + #[strum(serialize = "mark_async_success_enqueued")] + MarkAsyncSuccessEnqueued, #[strum(serialize = "enqueue_logging")] Enqueue, #[strum(serialize = "failure_handler")] @@ -77,6 +87,8 @@ pub(crate) enum DeploymentHooks { pub(crate) enum Streaming { #[strum(serialize = "stream_opened")] Opened, + #[strum(serialize = "stamp_stream_head")] + StampHead, #[strum(serialize = "stream_success")] Success, #[strum(serialize = "stream_failure")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 093ce60e510..ee16e0c0059 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -62,11 +62,22 @@ FAKES = { custom_llm_provider=provider, ), 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'messages_input': lambda body, kwargs: [{'role': 'user', 'content': json.dumps(body)}], 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( original_response, api_key, additional_args ), + 'record_response': lambda logger, status, headers, body: logger.model_call_details.update( + response_headers=dict(headers) + ), + 'record_api_window': lambda logger, start, end: logger.model_call_details.update( + llm_api_duration_ms=(end - start) * 1000 + ), 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), - 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'claim_async_success_log': lambda logger: logger.claim_async_success_log(), + 'defer_success': lambda logger, pending, response, start, end: ( + setattr(logger, 'deferred_success_result', (response, start, end)), + setattr(logger, 'enqueue_deferred_logging', pending), + ), 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( response, start, end ), @@ -75,6 +86,7 @@ FAKES = { )(error, ''.join(traceback.format_exception(error)), start, end), 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None), 'enqueue_logging': lambda coroutine: coroutine.enqueue(), 'restore_context': lambda logger: logger.record('restore', None), 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), @@ -85,7 +97,8 @@ FAKES = { ), 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), 'stream_opened': lambda logger, head: logger.record('stream_opened', head), - 'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record( + 'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head), + 'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record( 'stream_success', list(chunks) ), 'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error), @@ -122,6 +135,7 @@ class StubCoroutine: class StubLogger: def __init__(self): self.calls = [] + self.model_call_details = {} self.hooks = {} self.on_enqueue = lambda coroutine: None @@ -149,6 +163,10 @@ class StubLogger: self.record('async_success_handler', response) return StubCoroutine(self) + def claim_async_success_log(self): + claimed, self.claimed = getattr(self, 'claimed', False), True + return not claimed + def handle_sync_success_callbacks_for_async_calls(self, response, start, end): self.record('sync_success_for_async_call', response) diff --git a/litellm-rust/crates/circuit-breaker/AGENTS.md b/litellm-rust/crates/circuit-breaker/AGENTS.md index 52ef0ca7435..852cd72664a 100644 --- a/litellm-rust/crates/circuit-breaker/AGENTS.md +++ b/litellm-rust/crates/circuit-breaker/AGENTS.md @@ -1,13 +1,14 @@ - `fsm` is a generic state machine runtime and knows nothing about breakers. `breaker` is a generic circuit breaker and knows nothing about Redis. Deciding which errors count as failures belongs to the consumer - `Machine::step` stays pure. `&self` is config and time comes in on the event, never read inside the machine. The only clock read is `Admission`'s drop, at the call-site boundary - Transition rules are one `match (state, event)` per machine, with no macro DSL and no `_` arm that would hide a new state from the compiler -- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes, matching Python bumping `generation` only when `_state` changes. Data inside a variant (failure counters, `since`) changes without a transition +- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes. Data inside a variant (failure counters, `since`) changes without a transition - The observer runs under the lock so transitions are seen in generation order. Keep it to logging or a metric - A disabled breaker is `Option::None` at the consumer, not a flag inside the machine. `Config` comes in already resolved, reading `REDIS_CIRCUIT_BREAKER_*` env is the consumer's job -## Parity with Python `RedisCircuitBreaker` (`litellm/caching/redis_cache.py`) -- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against the Python breaker. A behavior change here needs the same check, by driving the Python class with the same steps -- Divergences, each on purpose +## The only Redis breaker +- Python's `RedisCache` admits through this breaker via `NativeRedisBreaker` (`python-bridge/src/redis/breaker.rs`), sharing one per Redis server with the native client through `redis-connect`'s `RedisHosts`, so the Python `RedisCircuitBreaker` it replaced is gone +- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against that Python class and now pins its behavior. A behavior change here changes what every Redis caller in the proxy sees +- Divergences from the old Python class, each on purpose - monotonic `Instant` instead of `time.time()`, so a wall clock jump can't open or close the breaker - a probe `Admission` dropped without an outcome records a hard failure and reopens. Python's probe that dies on `CancelledError` (a `BaseException`) records nothing and leaves the breaker HALF_OPEN for good, and a dropped Rust future is common - `_failure_count` is dropped, it only fed a log line diff --git a/litellm-rust/crates/circuit-breaker/src/breaker.rs b/litellm-rust/crates/circuit-breaker/src/breaker.rs index c9509561b37..cf4961ef099 100644 --- a/litellm-rust/crates/circuit-breaker/src/breaker.rs +++ b/litellm-rust/crates/circuit-breaker/src/breaker.rs @@ -1,5 +1,5 @@ use std::{ - sync::Arc, + sync::{Arc, Weak}, time::{Duration, Instant}, }; @@ -128,6 +128,9 @@ impl Machine for Config { #[derive(Clone)] pub struct Breaker(Arc>); +#[derive(Clone)] +pub struct WeakBreaker(Weak>); + #[derive(Debug, PartialEq, Eq)] pub struct Rejected; @@ -165,6 +168,16 @@ impl Breaker { pub fn state(&self) -> State { self.0.state() } + + pub fn downgrade(&self) -> WeakBreaker { + WeakBreaker(Arc::downgrade(&self.0)) + } +} + +impl WeakBreaker { + pub fn upgrade(&self) -> Option { + self.0.upgrade().map(Breaker) + } } impl Admission { diff --git a/litellm-rust/crates/circuit-breaker/src/lib.rs b/litellm-rust/crates/circuit-breaker/src/lib.rs index e76228122d1..af8df68246e 100644 --- a/litellm-rust/crates/circuit-breaker/src/lib.rs +++ b/litellm-rust/crates/circuit-breaker/src/lib.rs @@ -1,5 +1,7 @@ mod breaker; mod fsm; -pub use breaker::{Admission, Breaker, Config, Event, Failure, Output, Rejected, State}; +pub use breaker::{ + Admission, Breaker, Config, Event, Failure, Output, Rejected, State, WeakBreaker, +}; pub use fsm::{Machine, Observer, Shared, Transition}; diff --git a/litellm-rust/crates/config/AGENTS.md b/litellm-rust/crates/config/AGENTS.md new file mode 100644 index 00000000000..9f6f1db9a67 --- /dev/null +++ b/litellm-rust/crates/config/AGENTS.md @@ -0,0 +1,3 @@ +## Redis settings + +- This crate only deserializes YAML. `RedisSettings` is the raw `coordination_redis` or `cache_params` map and `CacheSettings` tells a `type: redis` cache apart from the rest. Reading those maps the way `litellm/_redis.py` does belongs to `gateway/src/redis`, and turning merged kwargs into a `Resolution` belongs to `redis-resolve` diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 5b0cd22332f..a433aacc1c9 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -6,17 +6,16 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth-types.workspace = true litellm-router-types.workspace = true +litellm-secret-value.workspace = true serde.workspace = true -serde_json.workspace = true serde_yaml_ng = "0.10.0" strum.workspace = true thiserror.workspace = true -tracing.workspace = true +veil.workspace = true [dev-dependencies] litellm-error.workspace = true rstest.workspace = true +serde_json.workspace = true tempfile.workspace = true -tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs index e5219ac992c..300f15380bd 100644 --- a/litellm-rust/crates/config/src/error.rs +++ b/litellm-rust/crates/config/src/error.rs @@ -11,11 +11,4 @@ pub enum Error { path: PathBuf, source: serde_yaml_ng::Error, }, - #[error("Redis host does not yet support configured setting {0}")] - RedisSetting(String), - #[error("invalid Redis setting {setting}: {reason}")] - RedisInvalid { - setting: &'static str, - reason: &'static str, - }, } diff --git a/litellm-rust/crates/config/src/lib.rs b/litellm-rust/crates/config/src/lib.rs index b8af6f11d2e..46657aa91e7 100644 --- a/litellm-rust/crates/config/src/lib.rs +++ b/litellm-rust/crates/config/src/lib.rs @@ -14,10 +14,7 @@ pub use error::Error; pub use litellm_router_types::{LitellmParams, Spelled}; pub use mcp::{McpAuth, McpServer, McpTransport}; pub use model::Model; -pub use redis::{ - CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings, - ResolvedRedis, environment_has_target, -}; +pub use redis::{CacheSettings, RedisSettings}; pub use settings::{ ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings, TracingStoreSettings, diff --git a/litellm-rust/crates/config/src/mcp.rs b/litellm-rust/crates/config/src/mcp.rs index cbe992f9956..00480fa5be1 100644 --- a/litellm-rust/crates/config/src/mcp.rs +++ b/litellm-rust/crates/config/src/mcp.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, fmt}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use crate::Object; diff --git a/litellm-rust/crates/config/src/redis.rs b/litellm-rust/crates/config/src/redis.rs index 5674f993025..d0ba107242c 100644 --- a/litellm-rust/crates/config/src/redis.rs +++ b/litellm-rust/crates/config/src/redis.rs @@ -1,195 +1,21 @@ -use std::{collections::BTreeMap, fmt, time::Duration}; +use std::collections::BTreeMap; -use litellm_auth_types::SecretValue; -use litellm_router_types::Spelled; use serde::{Deserialize, Deserializer, de::Error as _}; +use veil::Redact; -use crate::{AdditionalFields, Error, Object, Value}; +use crate::{Object, Value}; -pub type Lookup<'a> = &'a dyn Fn(&str) -> Option; - -const SUPPORTED: &[&str] = &[ - "url", - "host", - "port", - "username", - "password", - "ssl", - "db", - "socket_timeout", - "startup_nodes", - "namespace", -]; - -const IGNORED: &[&str] = &[ - "max_connections", - "single_connection_client", - "health_check_interval", - "socket_keepalive", - "socket_keepalive_options", - "decode_responses", - "encoding", - "encoding_errors", - "charset", - "errors", - "client_name", - "lib_name", - "lib_version", - "sentinel_password", - "gcp_ssl_ca_certs", - "azure_client_id", - "azure_tenant_id", - "azure_client_secret", - "aws_iam_user_name", - "aws_iam_cache_name", - "aws_iam_region", - "aws_iam_serverless", -]; - -const WARN_IGNORED: &[&str] = &[ - "max_connections", - "health_check_interval", - "socket_keepalive", - "socket_keepalive_options", -]; - -const REDIS_PY_KWARGS: &[&str] = &[ - "aws_iam_auth", - "aws_iam_cache_name", - "aws_iam_region", - "aws_iam_serverless", - "aws_iam_user_name", - "azure_client_id", - "azure_client_secret", - "azure_redis_ad_token", - "azure_tenant_id", - "cache", - "cache_config", - "charset", - "client_name", - "db", - "decode_responses", - "encoding", - "encoding_errors", - "errors", - "event_dispatcher", - "gcp_service_account", - "gcp_ssl_ca_certs", - "health_check_interval", - "host", - "lib_name", - "lib_version", - "max_connections", - "password", - "port", - "protocol", - "redis_connect_func", - "retry_on_error", - "retry_on_timeout", - "single_connection_client", - "socket_connect_timeout", - "socket_keepalive", - "socket_keepalive_options", - "socket_timeout", - "ssl", - "ssl_ca_certs", - "ssl_ca_data", - "ssl_ca_path", - "ssl_cert_reqs", - "ssl_certfile", - "ssl_check_hostname", - "ssl_ciphers", - "ssl_keyfile", - "ssl_min_version", - "ssl_ocsp_context", - "ssl_ocsp_expected_cert", - "ssl_password", - "ssl_validate_ocsp", - "ssl_validate_ocsp_stapled", - "unix_socket_path", - "url", - "username", -]; - -const UNSUPPORTED_EXTRA: &[&str] = &[ - "credential_provider", - "sentinel_nodes", - "retry", - "connection_pool", -]; - -const CACHE_ONLY: &[&str] = &[ - "type", - "mode", - "ttl", - "default_in_memory_ttl", - "default_in_redis_ttl", - "supported_call_types", - "similarity_threshold", - "max_messages", - "_backend", -]; - -#[derive(Clone, Default)] +/// A Redis section as `config.yaml` spells it, before any surface reads it. +#[derive(Clone, Default, Deserialize, Redact)] +#[serde(transparent)] pub struct RedisSettings { - pub url: Option, - pub host: Option>, - pub port: Option>, - pub username: Option, - pub password: Option, - pub ssl: Option>, - pub db: Option>, - pub socket_timeout: Option>, - pub startup_nodes: Option, - pub namespace: Option>, - pub additional_fields: AdditionalFields, - unparsed: DeferredInputs, + #[redact(with = "[REDACTED]")] + fields: BTreeMap, } -impl<'de> Deserialize<'de> for RedisSettings { - fn deserialize>(deserializer: D) -> Result { - let fields = BTreeMap::::deserialize(deserializer)?; - let scalar = |key: &str| fields.get(key).and_then(as_text); - Ok(Self { - url: scalar("url").map(SecretValue::new), - host: scalar("host").map(Spelled::Value), - port: spelled(scalar("port")), - username: scalar("username").map(SecretValue::new), - password: scalar("password").map(SecretValue::new), - ssl: spelled(scalar("ssl")), - db: spelled(scalar("db")), - socket_timeout: spelled(scalar("socket_timeout")), - startup_nodes: fields - .get("startup_nodes") - .filter(|value| !value.is_null()) - .cloned(), - namespace: scalar("namespace").map(Spelled::Value), - additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)), - unparsed: deferred_inputs(&fields), - }) - } -} - -impl fmt::Debug for RedisSettings { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("RedisSettings") - .field("url", &self.url) - .field("host", &self.host) - .field("port", &self.port) - .field("username", &self.username) - .field("password", &self.password) - .field("ssl", &self.ssl) - .field("db", &self.db) - .field("socket_timeout", &self.socket_timeout) - .field( - "startup_nodes", - &self.startup_nodes.as_ref().map(|_| ""), - ) - .field("namespace", &self.namespace) - .field("additional_fields", &self.additional_fields.keys()) - .field("unparsed", &self.unparsed.keys()) - .finish() +impl RedisSettings { + pub fn fields(&self) -> &BTreeMap { + &self.fields } } @@ -216,452 +42,22 @@ impl<'de> Deserialize<'de> for CacheSettings { } } -#[derive(Clone, PartialEq)] -pub struct ResolvedRedis { - pub endpoint: RedisEndpoint, - pub namespace: Option, - pub timeout: Option, -} +#[cfg(test)] +mod tests { + use super::RedisSettings; + use rstest::rstest; -impl fmt::Debug for ResolvedRedis { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("ResolvedRedis") - .field("endpoint", &self.endpoint) - .field("namespace", &self.namespace) - .field("timeout", &self.timeout) - .finish() + #[rstest] + #[case::scalar(serde_json::json!("credential-value"))] + #[case::object(serde_json::json!({"secret": "credential-value"}))] + #[case::array(serde_json::json!(["credential-value"]))] + fn credentials_are_redacted(#[case] password: serde_json::Value) { + let settings: RedisSettings = serde_json::from_value(serde_json::json!({ + "host": "redis", "password": password, + })) + .unwrap(); + let debug = format!("{settings:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("[REDACTED]")); } } - -#[derive(Clone, PartialEq)] -pub enum RedisEndpoint { - Url(SecretValue), - Host(RedisHost), - Cluster(RedisCluster), -} - -impl fmt::Debug for RedisEndpoint { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Url(_) => formatter.debug_tuple("Url").field(&"").finish(), - Self::Host(host) => host.fmt(formatter), - Self::Cluster(cluster) => cluster.fmt(formatter), - } - } -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisHost { - pub host: String, - pub port: u16, - pub username: Option, - pub password: Option, - pub tls: bool, - pub db: i64, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisCluster { - pub nodes: Vec, - pub username: Option, - pub password: Option, - pub tls: bool, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisNode { - pub host: String, - pub port: u16, -} - -#[derive(Clone)] -struct Entry { - value: Value, - source: String, -} - -type Kwargs = BTreeMap; - -impl RedisSettings { - pub fn resolve(&self, lookup: Lookup<'_>) -> Result { - self.resolve_with(lookup, false) - } - - pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result { - self.resolve_with(lookup, true) - } - - fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result { - let block = self.kwargs(cache); - let block: Kwargs = block - .into_iter() - .filter_map(|(key, entry)| { - let value = environment_reference(&entry.value, lookup)?; - Some((key, Entry { value, ..entry })) - }) - .collect(); - let environment = environment_kwargs(lookup); - let names_a_target = ["host", "startup_nodes", "sentinel_nodes"] - .iter() - .any(|key| block.contains_key(*key)); - let drop_environment_url = names_a_target && !block.contains_key("url"); - let merged: Kwargs = environment - .into_iter() - .filter(|(key, _)| !(key == "url" && drop_environment_url)) - .chain(block) - .collect(); - reject_unsupported(&merged, lookup)?; - warn_ignored(&merged); - - let startup_nodes = match merged.get("startup_nodes") { - Some(entry) => Some(cluster_nodes(&entry.value)?), - None => lookup("REDIS_CLUSTER_NODES") - .map(|nodes| cluster_nodes(&Value::String(nodes))) - .transpose()?, - }; - let username = secret_value(&merged, "username"); - let password = secret_value(&merged, "password"); - let tls = bool_value(&merged, "ssl"); - let endpoint = match (startup_nodes, text_value(&merged, "url")) { - (Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster { - nodes, - username, - password, - tls, - }), - (_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)), - _ => RedisEndpoint::Host(RedisHost { - host: text_value(&merged, "host").ok_or(Error::RedisInvalid { - setting: "host", - reason: "either host or url must be specified for redis", - })?, - port: integer(&merged, "port")?.unwrap_or(6379), - username, - password, - tls, - db: integer(&merged, "db")?.unwrap_or(0), - }), - }; - let timeout = merged - .get("socket_timeout") - .and_then(|entry| as_text(&entry.value)) - .map(|seconds| { - seconds - .parse::() - .ok() - .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) - .ok_or(Error::RedisInvalid { - setting: "socket_timeout", - reason: "must be a non-negative number of seconds", - }) - }) - .transpose()?; - Ok(ResolvedRedis { - endpoint, - namespace: text_value(&merged, "namespace"), - timeout, - }) - } - - fn kwargs(&self, cache: bool) -> Kwargs { - let known = [ - ( - "url", - self.url.as_ref().map(|value| value.expose().to_owned()), - ), - ("host", spelled_text(self.host.as_ref())), - ("port", spelled_text(self.port.as_ref())), - ( - "username", - self.username - .as_ref() - .map(|value| value.expose().to_owned()), - ), - ( - "password", - self.password - .as_ref() - .map(|value| value.expose().to_owned()), - ), - ("ssl", spelled_text(self.ssl.as_ref())), - ("db", spelled_text(self.db.as_ref())), - ("socket_timeout", spelled_text(self.socket_timeout.as_ref())), - ("namespace", spelled_text(self.namespace.as_ref())), - ]; - let known = known.into_iter().filter_map(|(key, value)| { - value.map(|value| { - ( - key.to_owned(), - Entry { - value: Value::String(value), - source: key.to_owned(), - }, - ) - }) - }); - let startup_nodes = self.startup_nodes.clone().map(|value| { - ( - "startup_nodes".to_owned(), - Entry { - value, - source: "startup_nodes".to_owned(), - }, - ) - }); - let extras = self.additional_fields.iter().filter_map(|(key, value)| { - if cache && CACHE_ONLY.contains(&key.as_str()) { - return None; - } - let key = match (cache, key.as_str()) { - (true, "redis_startup_nodes") => "startup_nodes".to_owned(), - _ => key.clone(), - }; - Some(( - key.clone(), - Entry { - value: value.clone(), - source: key, - }, - )) - }); - let unparsed = self.unparsed.iter().map(|(key, value)| { - ( - key.clone(), - Entry { - value: value.0.clone(), - source: key.clone(), - }, - ) - }); - extras - .chain(unparsed) - .chain(known) - .chain(startup_nodes) - .collect() - } -} - -pub fn environment_has_target(lookup: Lookup<'_>) -> bool { - [ - "REDIS_HOST", - "REDIS_URL", - "REDIS_CLUSTER_NODES", - "REDIS_SENTINEL_NODES", - ] - .iter() - .any(|name| lookup(name).is_some()) -} - -fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> { - let active_auth = [ - merged - .get("gcp_service_account") - .map(|entry| entry.source.as_str()) - .or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")), - text_value(merged, "azure_redis_ad_token") - .or_else(|| lookup("REDIS_AZURE_AD_TOKEN")) - .filter(|value| value.eq_ignore_ascii_case("true")) - .map(|_| { - merged - .get("azure_redis_ad_token") - .map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str()) - }), - bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()), - merged - .get("sentinel_nodes") - .map(|entry| entry.source.as_str()) - .or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")), - ]; - if let Some(setting) = active_auth.into_iter().flatten().next() { - return Err(Error::RedisSetting(setting.to_owned())); - } - match merged - .iter() - .find(|(key, _)| is_unsupported(key)) - .map(|(_, entry)| entry.source.clone()) - { - Some(setting) => Err(Error::RedisSetting(setting)), - None => Ok(()), - } -} - -fn warn_ignored(merged: &Kwargs) { - for setting in WARN_IGNORED { - if merged.contains_key(*setting) { - tracing::warn!(setting, "Redis setting is accepted but ignored"); - } - } -} - -fn is_unsupported(key: &str) -> bool { - let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key); - let handled = SUPPORTED.contains(&key) - || IGNORED.contains(&key) - || matches!( - key, - "gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth" - ); - known && !handled -} - -fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs { - REDIS_PY_KWARGS - .iter() - .filter_map(|key| { - let source = format!("REDIS_{}", key.to_ascii_uppercase()); - let value = lookup(&source)?; - Some(( - (*key).to_owned(), - Entry { - value: Value::String(value), - source, - }, - )) - }) - .collect() -} - -fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option { - match value { - Value::String(text) => match text.strip_prefix("os.environ/") { - Some(name) => lookup(name).map(Value::String), - None => Some(value.clone()), - }, - Value::Null => None, - _ => Some(value.clone()), - } -} - -fn as_text(value: &Value) -> Option { - match value { - Value::String(text) => Some(text.clone()), - Value::Number(number) => Some(number.to_string()), - Value::Bool(flag) => Some(flag.to_string()), - _ => None, - } -} - -fn text_value(kwargs: &Kwargs, key: &str) -> Option { - kwargs.get(key).and_then(|entry| as_text(&entry.value)) -} - -fn secret_value(kwargs: &Kwargs, key: &str) -> Option { - text_value(kwargs, key).map(SecretValue::new) -} - -fn bool_value(kwargs: &Kwargs, key: &str) -> bool { - match kwargs.get(key).map(|entry| &entry.value) { - Some(Value::Bool(flag)) => *flag, - Some(value) => as_text(value) - .is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")), - None => false, - } -} - -fn integer(kwargs: &Kwargs, key: &'static str) -> Result, Error> { - kwargs - .get(key) - .and_then(|entry| as_text(&entry.value)) - .map(|text| { - text.parse().map_err(|_| Error::RedisInvalid { - setting: key, - reason: "must be an integer", - }) - }) - .transpose() -} - -const INVALID_NODES: Error = Error::RedisInvalid { - setting: "startup_nodes", - reason: "must be a list of {host, port} objects", -}; - -fn cluster_nodes(value: &Value) -> Result, Error> { - match value { - Value::String(json) => { - let nodes: Vec> = - serde_json::from_str(json).map_err(|_| INVALID_NODES)?; - nodes - .iter() - .map(|node| { - cluster_node( - node.get("host").and_then(json_text), - node.get("port").and_then(json_text), - ) - }) - .collect() - } - Value::Sequence(nodes) => nodes - .iter() - .map(|node| { - cluster_node( - node.get("host").and_then(as_text), - node.get("port").and_then(as_text), - ) - }) - .collect(), - _ => Err(INVALID_NODES), - } -} - -fn cluster_node(host: Option, port: Option) -> Result { - match (host, port.and_then(|port| port.parse().ok())) { - (Some(host), Some(port)) => Ok(RedisNode { host, port }), - _ => Err(INVALID_NODES), - } -} - -fn json_text(value: &serde_json::Value) -> Option { - match value { - serde_json::Value::String(text) => Some(text.clone()), - serde_json::Value::Number(number) => Some(number.to_string()), - _ => None, - } -} - -fn spelled_text(value: Option<&Spelled>) -> Option { - value.map(|value| match value { - Spelled::Value(value) => value.to_string(), - Spelled::Text(value) => value.clone(), - }) -} - -fn spelled(text: Option) -> Option> { - text.map(|text| match text.parse() { - Ok(value) => Spelled::Value(value), - Err(_) => Spelled::Text(text), - }) -} - -fn cloned( - fields: &BTreeMap, - keep: impl Fn(&str, &Value) -> bool, -) -> AdditionalFields { - fields - .iter() - .filter(|(key, value)| keep(key, value)) - .map(|(key, value)| (key.clone(), value.clone())) - .collect() -} - -#[derive(Clone)] -struct DeferredInput(Value); - -type DeferredInputs = BTreeMap; - -fn deferred_inputs(fields: &BTreeMap) -> DeferredInputs { - fields - .iter() - .filter(|(key, value)| is_unparsed(key, value)) - .map(|(key, value)| (key.clone(), DeferredInput(value.clone()))) - .collect() -} - -fn is_unparsed(key: &str, value: &Value) -> bool { - SUPPORTED.contains(&key) - && match key { - "startup_nodes" => value.is_null(), - _ => as_text(value).is_none(), - } -} diff --git a/litellm-rust/crates/config/src/settings.rs b/litellm-rust/crates/config/src/settings.rs index 9571c4fa5de..1bfc19983f3 100644 --- a/litellm-rust/crates/config/src/settings.rs +++ b/litellm-rust/crates/config/src/settings.rs @@ -1,6 +1,6 @@ use std::fmt; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use crate::{ diff --git a/litellm-rust/crates/config/tests/config.rs b/litellm-rust/crates/config/tests/config.rs index 3c7d02203e5..fec042e41be 100644 --- a/litellm-rust/crates/config/tests/config.rs +++ b/litellm-rust/crates/config/tests/config.rs @@ -440,3 +440,29 @@ fn mcp_config_redacts_nested_credentials_and_preserves_policy_for_validation() { fn rejects_invalid_typed_mcp_settings(#[case] setting: &str) { assert!(Config::from_yaml(&format!("mcp_servers:\n docs:\n {setting}\n")).is_err()); } + +#[rstest] +fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() { + let config = Config::from_yaml( + "litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid", + ) + .unwrap(); + + assert!(matches!( + config.litellm_settings.cache_params, + Some(litellm_config::CacheSettings::Redis(_)) + )); +} + +#[rstest] +fn non_redis_cache_settings_remain_opaque() { + let config = Config::from_yaml( + "litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true", + ) + .unwrap(); + + assert!(matches!( + config.litellm_settings.cache_params, + Some(litellm_config::CacheSettings::Other(_)) + )); +} diff --git a/litellm-rust/crates/config/tests/redis.rs b/litellm-rust/crates/config/tests/redis.rs deleted file mode 100644 index ab409541c8e..00000000000 --- a/litellm-rust/crates/config/tests/redis.rs +++ /dev/null @@ -1,405 +0,0 @@ -use std::{ - collections::HashMap, - fmt::Write, - sync::{Arc, Mutex}, -}; - -use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis}; -use rstest::rstest; -use tracing::{Event, Subscriber, field::Visit}; -use tracing_subscriber::{Layer, layer::Context, prelude::*}; - -fn resolve( - yaml: &str, - environment: &[(&str, &str)], -) -> Result { - let environment: HashMap = environment - .iter() - .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) - .collect(); - let config = Config::from_yaml(yaml).unwrap(); - config - .general_settings - .coordination_redis - .as_ref() - .unwrap_or(&RedisSettings::default()) - .resolve(&|name| environment.get(name).cloned()) -} - -fn resolve_cache(yaml: &str) -> Result { - let config = Config::from_yaml(yaml).unwrap(); - match config.litellm_settings.cache_params.unwrap() { - CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None), - CacheSettings::Other(_) => panic!("expected a Redis cache"), - } -} - -fn cluster_ports(resolved: &ResolvedRedis) -> Vec { - match &resolved.endpoint { - RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(), - endpoint => panic!("expected a cluster, got {endpoint:?}"), - } -} - -#[rstest] -#[case::host( - "general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5", - &[], - "host:redis:6380:2:true:app:1500" -)] -#[case::environment_url( - "model_list: []", - &[("REDIS_URL", "redis://environment:6379")], - "url:redis://environment:6379" -)] -#[case::explicit_host_drops_environment_url( - "general_settings:\n coordination_redis:\n host: configured", - &[("REDIS_URL", "redis://environment:6379")], - "host:configured:6379:0:false::none" -)] -#[case::yaml_cluster( - "general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"", - &[], - "cluster:node-a:7000,node-b:7001" -)] -#[case::json_environment_cluster( - "model_list: []", - &[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)], - "cluster:node-a:7000,node-b:7001" -)] -#[case::missing_reference_keeps_environment_fallback( - "general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST", - &[("REDIS_HOST", "environment")], - "host:environment:6379:0:false::none" -)] -#[case::configured_values_beat_the_same_environment_values( - "general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3", - &[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")], - "host:configured:6380:3:false::none" -)] -#[case::environment_url_beats_the_environment_host( - "model_list: []", - &[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")], - "url:redis://environment:6379" -)] -#[case::startup_nodes_beat_a_url( - "general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000", - &[], - "cluster:node-a:7000" -)] -#[case::environment_references_resolve_the_host_and_port( - "general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT", - &[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")], - "host:referenced:6400:0:false::none" -)] -#[case::a_null_host_falls_back_to_the_environment( - "general_settings:\n coordination_redis:\n host: null", - &[("REDIS_HOST", "environment")], - "host:environment:6379:0:false::none" -)] -fn resolves_redis_through_the_parsed_config( - #[case] yaml: &str, - #[case] environment: &[(&str, &str)], - #[case] expected: &str, -) { - let resolved = resolve(yaml, environment).unwrap(); - let description = match resolved.endpoint { - RedisEndpoint::Url(url) => format!("url:{}", url.expose()), - RedisEndpoint::Host(host) => format!( - "host:{}:{}:{}:{}:{}:{}", - host.host, - host.port, - host.db, - host.tls, - resolved.namespace.unwrap_or_default(), - resolved - .timeout - .map(|timeout| timeout.as_millis().to_string()) - .unwrap_or_else(|| "none".to_owned()) - ), - RedisEndpoint::Cluster(cluster) => format!( - "cluster:{}", - cluster - .nodes - .iter() - .map(|node| format!("{}:{}", node.host, node.port)) - .collect::>() - .join(",") - ), - }; - - assert_eq!(description, expected); -} - -#[rstest] -#[case::custom_tls_file( - "general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem", - &[], - "ssl_ca_certs" -)] -#[case::gcp_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account", - &[], - "gcp_service_account" -)] -#[case::gcp_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")], - "REDIS_GCP_SERVICE_ACCOUNT" -)] -#[case::azure_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true", - &[], - "azure_redis_ad_token" -)] -#[case::azure_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")], - "REDIS_AZURE_AD_TOKEN" -)] -#[case::aws_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true", - &[], - "aws_iam_auth" -)] -#[case::aws_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")], - "REDIS_AWS_IAM_AUTH" -)] -#[case::sentinel_from_config( - "general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]", - &[], - "sentinel_nodes" -)] -#[case::sentinel_from_environment( - "model_list: []", - &[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")], - "REDIS_SENTINEL_NODES" -)] -fn rejected_settings_name_their_source( - #[case] yaml: &str, - #[case] environment: &[(&str, &str)], - #[case] expected: &str, -) { - let error = resolve(yaml, environment).unwrap_err(); - - assert!(matches!(error, Error::RedisSetting(setting) if setting == expected)); -} - -#[rstest] -fn an_explicit_host_keeps_environment_credentials() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: configured", - &[ - ("REDIS_PASSWORD", "environment-secret"), - ("REDIS_SSL", "True"), - ], - ) - .unwrap(); - - let RedisEndpoint::Host(host) = resolved.endpoint else { - panic!("expected a host endpoint"); - }; - assert_eq!( - host.password.as_ref().map(|password| password.expose()), - Some("environment-secret") - ); - assert!(host.tls); -} - -#[rstest] -#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])] -#[case::a_structured_host_with_an_environment_target( - "general_settings:\n coordination_redis:\n host: []", - &[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")] -)] -fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) { - let error = resolve(yaml, environment).unwrap_err(); - - assert!(matches!( - error, - Error::RedisInvalid { - setting: "host", - .. - } - )); -} - -#[rstest] -fn the_cache_alias_names_cluster_nodes() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000", - ) - .unwrap(); - - assert_eq!(cluster_ports(&resolved), vec![7000]); -} - -#[rstest] -fn canonical_startup_nodes_beat_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101", - ) - .unwrap(); - - assert_eq!(cluster_ports(&resolved), vec![7000]); -} - -#[rstest] -fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100", - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -#[rstest] -fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100", - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -#[rstest] -fn rejected_environment_settings_never_print_the_value() { - let error = resolve( - "model_list: []", - &[ - ("REDIS_HOST", "redis"), - ("REDIS_SSL_CERT_REQS", "private-option-value"), - ], - ) - .unwrap_err(); - let message = error.to_string(); - - assert!(message.contains("REDIS_SSL_CERT_REQS")); - assert!(!message.contains("private-option-value")); -} - -#[rstest] -#[case::negative("-1")] -#[case::not_a_number("slow")] -fn rejects_bad_socket_timeout(#[case] timeout: &str) { - let yaml = format!( - "general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}" - ); - let error = resolve(&yaml, &[]).unwrap_err(); - - assert!(matches!( - error, - Error::RedisInvalid { - setting: "socket_timeout", - .. - } - )); -} - -#[rstest] -fn defaults_socket_timeout_for_the_gateway() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: redis", - &[], - ) - .unwrap(); - - assert_eq!(resolved.timeout, None); -} - -#[rstest] -fn debug_redacts_redis_secrets_and_url() { - let config = Config::from_yaml( - "general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password", - ) - .unwrap(); - let settings = config.general_settings.coordination_redis.unwrap(); - let resolved = settings.resolve(&|_| None).unwrap(); - let debug = format!("{settings:?} {resolved:?}"); - - assert!(!debug.contains("redis://user:password")); - assert!(!debug.contains("credential-user")); - assert!(!debug.contains("credential-password")); - assert!(debug.contains("redacted") || debug.contains("REDACTED")); -} - -#[rstest] -fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() { - let config = Config::from_yaml( - "litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid", - ) - .unwrap(); - - assert!(matches!( - config.litellm_settings.cache_params, - Some(litellm_config::CacheSettings::Redis(_)) - )); -} - -#[rstest] -fn non_redis_cache_settings_remain_opaque() { - let config = Config::from_yaml( - "litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true", - ) - .unwrap(); - - assert!(matches!( - config.litellm_settings.cache_params, - Some(litellm_config::CacheSettings::Other(_)) - )); -} - -#[rstest] -fn legacy_unknown_settings_do_not_change_resolution() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary", - &[], - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -struct WarningCapture(Arc>>); - -impl Layer for WarningCapture { - fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) { - let mut fields = Fields::default(); - event.record(&mut fields); - self.0.lock().unwrap().push(fields.0); - } -} - -#[derive(Default)] -struct Fields(String); - -impl Visit for Fields { - fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { - write!(self.0, "{}={value:?} ", field.name()).unwrap(); - } - - fn record_str(&mut self, field: &tracing::field::Field, value: &str) { - write!(self.0, "{}={value} ", field.name()).unwrap(); - } -} - -#[rstest] -fn ignored_setting_warning_names_the_setting_without_its_value() { - let events = Arc::new(Mutex::new(Vec::new())); - let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone())); - tracing::subscriber::with_default(subscriber, || { - resolve( - "general_settings:\n coordination_redis:\n host: redis\n max_connections: 47", - &[], - ) - .unwrap(); - }); - let warnings = events.lock().unwrap().join(" "); - - assert!(warnings.contains("max_connections")); - assert!(!warnings.contains("47")); -} diff --git a/litellm-rust/crates/core-utils/src/core_helpers.rs b/litellm-rust/crates/core-utils/src/core_helpers.rs index 58ee88d301c..b93f0484bb9 100644 --- a/litellm-rust/crates/core-utils/src/core_helpers.rs +++ b/litellm-rust/crates/core-utils/src/core_helpers.rs @@ -13,6 +13,7 @@ const FINISH_REASONS: &[(&str, &str)] = &[ ("max_tokens", "length"), ("refusal", "content_filter"), ("compaction", "length"), + ("pause_turn", "length"), ("model_context_window_exceeded", "length"), ("guardrail_intervened", "content_filter"), ("content_filtered", "content_filter"), @@ -75,6 +76,7 @@ mod tests { #[case::end_turn("end_turn", "stop")] #[case::stop_sequence("stop_sequence", "stop")] #[case::max_tokens("max_tokens", "length")] + #[case::pause_turn("pause_turn", "length")] #[case::model_context_window_exceeded("model_context_window_exceeded", "length")] #[case::refusal("refusal", "content_filter")] #[case::guardrail_intervened("guardrail_intervened", "content_filter")] diff --git a/litellm-rust/crates/cost/AGENTS.md b/litellm-rust/crates/cost/AGENTS.md new file mode 100644 index 00000000000..0f70dea8921 --- /dev/null +++ b/litellm-rust/crates/cost/AGENTS.md @@ -0,0 +1 @@ +The price vocabulary (meter, band, channel, `When`, `Billing`) is explained in [../model-catalog/AGENTS.md](../model-catalog/AGENTS.md#concepts), the only crate that maps it to JSON keys diff --git a/litellm-rust/crates/cost/src/price.rs b/litellm-rust/crates/cost/src/price.rs index a90bedacb0f..70f6016780c 100644 --- a/litellm-rust/crates/cost/src/price.rs +++ b/litellm-rust/crates/cost/src/price.rs @@ -19,6 +19,7 @@ pub enum ServiceTier { Balanced, } +/// A prompt-size range. `OverPromptTokens(n)` is spelled `_above__tokens` or `_above_k_tokens` #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum Band { #[default] @@ -26,12 +27,15 @@ pub enum Band { OverPromptTokens(NonZeroU64), } +/// The request condition a price applies under, its service tier and prompt band #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct When { pub tier: ServiceTier, pub band: Band, } +/// The route whose price list a call reads, from the `_batches` and `_decisions` suffixes. Not part of `When`, +/// since a missing batch price is half the online rate and a missing decisions price is the full online rate #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum Channel { #[default] @@ -40,6 +44,7 @@ pub enum Channel { Decisions, } +/// The key of one rate, what is counted, through which channel and under which condition #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct Price { pub meter: Meter, @@ -47,6 +52,7 @@ pub struct Price { pub when: When, } +/// Every rate a row lists, its off-peak overlay and the row attributes that say how the rates apply #[derive(Clone, Debug, Default, PartialEq)] pub struct Prices { rates: BTreeMap, @@ -117,6 +123,7 @@ pub enum TierScope { WholeCall, } +/// Rates that replace the row's token rates while a UTC window in the schedule is open #[derive(Clone, Debug, PartialEq)] pub struct OffPeak { schedule: Schedule, @@ -133,6 +140,7 @@ pub struct OffPeakRates { pub reasoning: Option, } +/// Multipliers applied on top of a priced call, never prices themselves #[derive(Clone, Debug, Default, PartialEq)] pub struct Adjustments(BTreeMap); @@ -143,6 +151,8 @@ pub enum Uplift { Residency(Residency), /// Vertex non-global endpoint, from `regional_endpoint_uplift_multiplier` NonGlobalEndpoint, + /// OpenAI FedRAMP endpoint (gov.api.openai.com), from `fedramp_uplift_multiplier` + FedrampEndpoint, /// Anthropic `speed: fast`, from `provider_specific_entry.fast` Speed(Speed), /// Anthropic and Bedrock geo, from `provider_specific_entry.`, lowercased as on main diff --git a/litellm-rust/crates/cost/src/price/meter.rs b/litellm-rust/crates/cost/src/price/meter.rs index 6e68981d011..09fd23e42ed 100644 --- a/litellm-rust/crates/cost/src/price/meter.rs +++ b/litellm-rust/crates/cost/src/price/meter.rs @@ -1,5 +1,6 @@ use std::num::NonZeroU64; +/// What a price counts. A variation only one meter has, such as cache TTL or image size, lives inside its variant #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum Meter { Input(InputUnit), @@ -51,12 +52,14 @@ pub enum OutputUnit { VideoSecond, } +/// `Explicit` is a cache the caller created ahead of time, spelled `explicit_cache_read_*`. `Implicit` is every other read #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum CacheReadMode { Implicit, Explicit, } +/// How long a cache write lives. `OneHour` and `TwelveHours` are spelled `_above_1hr` and `_above_12hr` #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum CacheTtl { Default, diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml index 27f20203843..ba9097e85d8 100644 --- a/litellm-rust/crates/gateway-auth/Cargo.toml +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -6,9 +6,9 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum = { workspace = true, features = ["matched-path"] } axum-login.workspace = true -litellm-auth-types.workspace = true litellm-config.workspace = true litellm-error.workspace = true litellm-secrets.workspace = true diff --git a/litellm-rust/crates/gateway-auth/src/authentication.rs b/litellm-rust/crates/gateway-auth/src/authentication.rs index 4ecfaa375b4..349913619f5 100644 --- a/litellm-rust/crates/gateway-auth/src/authentication.rs +++ b/litellm-rust/crates/gateway-auth/src/authentication.rs @@ -1,8 +1,8 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use axum::http::{Request, request::Parts}; -use litellm_auth_types::SecretValue; use litellm_config::Config; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/gateway-auth/src/http.rs b/litellm-rust/crates/gateway-auth/src/http.rs index ddbf240509a..ad1718161c0 100644 --- a/litellm-rust/crates/gateway-auth/src/http.rs +++ b/litellm-rust/crates/gateway-auth/src/http.rs @@ -4,7 +4,7 @@ use axum::{ middleware::Next, response::Response, }; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{AccessRequest, Auth, AuthenticatedRequest, Credential, Error}; diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 1fe8fa60e0e..0deaf3f7a82 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -1,6 +1,6 @@ use std::{future::Future, pin::Pin, time::SystemTime}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{Error, hash_token}; diff --git a/litellm-rust/crates/gateway-auth/src/ui.rs b/litellm-rust/crates/gateway-auth/src/ui.rs index 3d655bde187..ee4dd670aa7 100644 --- a/litellm-rust/crates/gateway-auth/src/ui.rs +++ b/litellm-rust/crates/gateway-auth/src/ui.rs @@ -3,7 +3,7 @@ use std::{convert::Infallible, sync::Arc}; use axum::extract::FromRequestParts; use axum::http::{header::AUTHORIZATION, request::Parts}; use axum_login::{AuthUser, AuthnBackend, UserId}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/gateway-auth/tests/auth.rs b/litellm-rust/crates/gateway-auth/tests/auth.rs index 781ba385381..54180e30504 100644 --- a/litellm-rust/crates/gateway-auth/tests/auth.rs +++ b/litellm-rust/crates/gateway-auth/tests/auth.rs @@ -8,9 +8,9 @@ use axum::{ routing::get, }; use futures_util::future::BoxFuture; -use litellm_auth_types::SecretValue; use litellm_config::Config; use litellm_gateway_auth::{Auth, RequireMasterKey, hash_token}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use rstest::{fixture, rstest}; use tower::ServiceExt; diff --git a/litellm-rust/crates/gateway-auth/tests/keys.rs b/litellm-rust/crates/gateway-auth/tests/keys.rs index 4d16d3573bc..1f6fc461a5b 100644 --- a/litellm-rust/crates/gateway-auth/tests/keys.rs +++ b/litellm-rust/crates/gateway-auth/tests/keys.rs @@ -1,10 +1,10 @@ use std::time::{Duration, SystemTime}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ Error, keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key}, }; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; struct Lookup { diff --git a/litellm-rust/crates/gateway-auth/tests/ui.rs b/litellm-rust/crates/gateway-auth/tests/ui.rs index 2e5f72d605d..8d465e1198d 100644 --- a/litellm-rust/crates/gateway-auth/tests/ui.rs +++ b/litellm-rust/crates/gateway-auth/tests/ui.rs @@ -1,6 +1,6 @@ use axum_login::{AuthUser, AuthnBackend}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{UiBackend, UiCredentials}; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; #[fixture] diff --git a/litellm-rust/crates/gateway-forward/AGENTS.md b/litellm-rust/crates/gateway-forward/AGENTS.md new file mode 100644 index 00000000000..3d07f9ee342 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/AGENTS.md @@ -0,0 +1,40 @@ +# Forwarding gateway + +`Forward` is a reverse proxy in front of the Python proxy on loopback. Python must see each request as if the client had connected to it directly, because every Python route (inference, MCP, pass-through, UI, WebSocket) is served behind it without Rust knowing which one it is + +- Serve HTTP/1.1 only, as uvicorn does. The supervisor never puts Rust in front of a Python server that speaks HTTP/2 +- Never branch on a route or path. A behavior one endpoint needs is a protocol behavior every endpoint gets +- Forward the method, request target, headers, and body byte for byte within the [URL contract](FORWARDING.md#url-fragments-and-literal-hashes). An unescaped `#` starts a fragment and is omitted with everything after it, while `%23` is preserved. Only drop hop-by-hop headers (RFC 9110 section 7.6.1) and set `x-litellm-rust-peer`, which replaces any client copy +- Stream both bodies. Never buffer, cap, or time out a body, since uploads and provider streams have no upper bound +- A client disconnect closes the upstream connection at every stage, because that is how Python learns to stop a provider stream +- An upstream that dies mid-response fails the client body. Never end it cleanly, or SSE clients treat a truncated stream as complete +- Never invent a Host. Python builds absolute URLs from it, so a missing one must never become the private upstream address +- A request Python could not have received unchanged (HTTP/1.1 without Host, a transfer coding other than chunked, a target without a path) gets the status uvicorn would have given, from Rust +- Pooled upstream connections are retired well before Python's keep-alive timeout, so a request never lands on a connection Python is closing +- Idle client connections close after Python's keep-alive timeout, as uvicorn closed them, so a client that vanishes cannot hold a socket forever. Only the wait for the next request is timed, never an exchange in progress +- WebSocket is the one protocol switch Python serves. Upgrades are relayed as raw bytes after the 101 and drained on shutdown +- Drain on EOF of stdin exactly as on SIGTERM. The supervisor holds the only write end, so EOF means it is gone and nobody is left to stop the gateway. `input_closed` reads on a plain thread, since a read blocked on tokio's blocking pool would hold up the runtime's exit + +## Layout + +Every forwarding decision has one home. `rewrite.rs` owns what Python and the client receive, as pure functions over requests and responses. `Rejection` in `error.rs` is the only place that picks a status for a request Rust answers itself. `settings.rs` parses the environment into typed values, so an `Upstream` is a loopback origin by construction. `server.rs` accepts client connections and owns their timeouts, since `axum::serve` exposes none, and `connections.rs` tracks every connection and relay for shutdown. `forward.rs` and `websocket.rs` only move requests, bodies and upgraded connections and never touch a header themselves. `connector.rs` keeps hyper reading Python's answer after Python stops reading the request + +## Environment + +`Settings` reads only these variables, and only when `LITELLM_FORWARD_UPSTREAM` is set. The Python supervisor in `litellm/proxy/rust_forward/` sets every one of them from the values Python actually runs with, so all of them are required and a missing or invalid one fails startup naming it. The crate never reads Python's own settings (`HOST`, `PORT`, `KEEPALIVE_TIMEOUT`, `GRACEFUL_SHUTDOWN_TIMEOUT`) and keeps no copy of Python's defaults, since Python can get those values from CLI flags the environment never shows + +Only the Python supervisor starts the gateway in forwarding mode, from `litellm` and from `python -m gateway.launch`, on a laptop and in every image. `supervisor/children.py` is the only writer of these variables, so adding, renaming or removing one changes `settings.rs` and `children.rust_spec` in the same change. Without `LITELLM_FORWARD_UPSTREAM` the binary runs as the standalone Rust gateway and this crate does nothing + +| Variable | Meaning | +|---|---| +| `LITELLM_FORWARD_UPSTREAM` | Python's loopback origin, e.g. `http://127.0.0.1:53211` | +| `LITELLM_FORWARD_HOST` | public bind host | +| `LITELLM_FORWARD_PORT` | public bind port | +| `LITELLM_FORWARD_UPSTREAM_KEEPALIVE` | seconds uvicorn keeps an idle connection open. Pooled connections are retired at half of it, and idle client connections are closed after it | +| `LITELLM_FORWARD_SHUTDOWN_DRAIN` | seconds to drain on shutdown, already including Python's graceful timeout and teardown headroom | + +## Tests + +Test through real sockets with `tests/support`, never `oneshot`. `Gateway` runs the forwarder on a listener like production, `recording` is a well-behaved upstream that captures what Python would receive, and `raw_upstream` and `Raw` drive bytes by hand for framing, timing and disconnects a real server would not produce. Waits measured in minutes use paused tokio time instead of real sleeps + +The forwarder sees HTTP/1.1 messages, not routes, so the tests are split into one file per traffic shape uvicorn can emit: `exchange` (plain request and response), `framing` (requests whose framing a parser could read two ways), `slow_head` (minutes before response headers), `stream` (incremental response bodies), `idle_stream` (long-lived streams that go quiet), `upload`, `download`, `websocket`, `client_identity`, `lifecycle` (disconnects, pooling, idle timeouts, shutdown) and `roundtrip` (the forwarding invariant for generated messages). A new route never gets its own test. Add its behavior to the shape it already has, and when it needs a wire behavior none of the files covers, add a new shape file for it diff --git a/litellm-rust/crates/gateway-forward/Cargo.toml b/litellm-rust/crates/gateway-forward/Cargo.toml new file mode 100644 index 00000000000..821ea91d269 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "litellm-gateway-forward" +version = "0.1.0" +edition.workspace = true +rust-version.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +axum.workspace = true +envy = "0.4.2" +hyper = { version = "1.12.0", features = ["http1", "server"] } +hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } +serde.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["io-util", "sync"] } +tokio-util = { version = "0.7", features = ["rt"] } +tower-service = "0.3.3" +tracing.workspace = true + +[dev-dependencies] +bytes.workspace = true +futures-util.workspace = true +http-body-util = "0.1" +httparse = "1" +hyper = { version = "1.12.0", features = ["client"] } +proptest.workspace = true +rstest.workspace = true +tokio = { workspace = true, features = ["io-util", "sync", "test-util"] } +tokio-tungstenite.workspace = true diff --git a/litellm-rust/crates/gateway-forward/FORWARDING.md b/litellm-rust/crates/gateway-forward/FORWARDING.md new file mode 100644 index 00000000000..532f341aba5 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/FORWARDING.md @@ -0,0 +1,31 @@ +# Python forwarding compatibility + +The opt-in forwarding gateway serves HTTP/1.1 and WebSockets in front of Python. Python still owns every route and policy. Its protocol contract and environment settings are in [AGENTS.md](AGENTS.md) + +## URL fragments and literal hashes + +The gateway follows Hyper's `http::Uri` parsing. If a raw client sends an unescaped `#` in the request target, only the path and query before the first `#` reach Python. For example, `/v1/files?purpose=batch#retry&limit=1` becomes `/v1/files?purpose=batch`, and `/gemini/v1beta/files/a#b?key=k` becomes `/gemini/v1beta/files/a` + +This is accepted behavior because `#` starts a URL fragment, which is separated before dereferencing ([RFC 3986 section 3.5](https://www.rfc-editor.org/rfc/rfc3986.html#section-3.5)). HTTP origin-form contains only the path and optional query, not a fragment ([RFC 9112 section 3.2.1](https://www.rfc-editor.org/rfc/rfc9112.html#section-3.2.1)). Browsers and URL-based clients omit the fragment before sending the request, so UI hash routing is unaffected + +A literal hash in path or query data must be encoded as `%23`, which the gateway forwards unchanged. Direct Python with h11 also accepts unescaped hashes as data, but that permissive raw-target behavior is outside this URL contract. Active socket tests cover both [fragment omission](tests/exchange.rs) and [encoded-hash preservation](tests/exchange.rs) + +## Known differences from uvicorn with h11 + +These cases remain executable as ignored socket tests. They are not passing parity checks + +| Input | Direct Python | Through Rust | Regression test | +|---|---|---|---| +| Header after a removed hop-by-hop field | Keeps the remaining header order | `HeaderMap::remove` can move a later header | [dropping_a_hop_by_hop_header_keeps_the_others_in_order](tests/exchange.rs) | +| Repeated header names interleaved with other names | Keeps the interleaving | Groups values by name, keeping their order within that name | [interleaved_repeated_headers_reach_python_interleaved](tests/exchange.rs) | +| Obsolete folded header, such as `x-a: b\r\n c` | Unfolds to `b c` | Rejects with 400 | [a_folded_header_reaches_python_unfolded](tests/exchange.rs) | +| Identical Content-Length list, such as `7, 7` | Accepts length 7 | Rejects with 400 | [a_content_length_list_of_one_value_reaches_python](tests/framing.rs) | +| HTTP/1.0 request with chunked transfer coding | Decodes the body | Rejects with 400 | [an_http10_chunked_upload_reaches_python](tests/framing.rs) | + +Clients can avoid the parser rejections by sending unfolded headers, one Content-Length value, and HTTP/1.1 for chunked uploads. Integrations must not depend on header ordering across different field names + +Run `cargo test -p litellm-gateway-forward --test exchange --test framing -- --ignored` from this directory to reproduce the differences + +## Connection header handling + +The gateway strips standard hop-by-hop fields and fields named by Connection, except routing, representation, credential, forwarding, and WebSocket handshake headers. A client cannot nominate those end-to-end fields for removal. The private `x-litellm-rust-peer` always replaces any client-supplied copy diff --git a/litellm-rust/crates/gateway-forward/src/connections.rs b/litellm-rust/crates/gateway-forward/src/connections.rs new file mode 100644 index 00000000000..10a4867fa51 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/connections.rs @@ -0,0 +1,58 @@ +use tokio_util::{ + sync::{CancellationToken, DropGuard}, + task::TaskTracker, +}; + +/// Every client connection and WebSocket relay the gateway runs, so shutdown can wait for them +/// and then close whatever outlives the drain. +#[derive(Clone, Default)] +pub(crate) struct Connections { + tasks: TaskTracker, + draining: CancellationToken, + closed: CancellationToken, +} + +impl Connections { + /// Runs `work` until it finishes or the gateway closes every connection. + pub fn spawn(&self, work: impl Future + Send + 'static) { + let closed = self.closed.clone(); + self.tasks.spawn(async move { + tokio::select! { + () = closed.cancelled() => {} + () = work => {} + } + }); + } + + pub fn begin_shutdown(&self) { + self.draining.cancel(); + } + + pub async fn draining(&self) { + self.draining.cancelled().await; + } + + pub fn is_draining(&self) -> bool { + self.draining.is_cancelled() + } + + pub async fn until_closed(&self, future: impl Future) -> Option { + tokio::select! { + () = self.closed.cancelled() => None, + output = future => Some(output), + } + } + + pub fn close(&self) { + self.closed.cancel(); + } + + pub fn close_on_drop(&self) -> DropGuard { + self.closed.clone().drop_guard() + } + + pub async fn wait(&self) { + self.tasks.close(); + self.tasks.wait().await; + } +} diff --git a/litellm-rust/crates/gateway-forward/src/connector.rs b/litellm-rust/crates/gateway-forward/src/connector.rs new file mode 100644 index 00000000000..5eb998db0c0 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/connector.rs @@ -0,0 +1,122 @@ +use std::{ + io::{self, IoSlice}, + pin::Pin, + task::{Context, Poll, ready}, +}; + +use axum::http::Uri; +use hyper::rt::{Read, ReadBufCursor, Write}; +use hyper_util::{ + client::legacy::connect::{Connected, Connection, HttpConnector}, + rt::TokioIo, +}; +use tokio::net::TcpStream; + +/// Connects to Python like `HttpConnector`, but keeps hyper reading Python's answer after Python +/// stops reading the request. uvicorn rejects an upload (413, 401) as soon as it reads the head +/// and closes with the body unread, so writing the rest fails, and hyper would fail the request +/// and drop the answer it already received. Like curl and httpcore, the forwarder discards the +/// rest of the request instead and reads on. A connection that never answered still fails, on the +/// read, and no request is ever sent twice. +#[derive(Clone)] +pub(crate) struct Connector(pub HttpConnector); + +impl tower_service::Service for Connector { + type Response = AnswerAfterRefusal; + type Error = >::Error; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll> { + self.0.poll_ready(context) + } + + fn call(&mut self, uri: Uri) -> Self::Future { + let connecting = self.0.call(uri); + Box::pin(async move { + Ok(AnswerAfterRefusal { + stream: connecting.await?, + refused: false, + }) + }) + } +} + +pub(crate) struct AnswerAfterRefusal { + stream: TokioIo, + refused: bool, +} + +impl AnswerAfterRefusal { + fn written(&mut self, written: io::Result, length: usize) -> io::Result { + match written { + Err(error) + if matches!( + error.kind(), + io::ErrorKind::BrokenPipe | io::ErrorKind::ConnectionReset + ) => + { + self.refused = true; + Ok(length) + } + written => written, + } + } +} + +impl Read for AnswerAfterRefusal { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: ReadBufCursor<'_>, + ) -> Poll> { + Pin::new(&mut self.stream).poll_read(context, buffer) + } +} + +impl Write for AnswerAfterRefusal { + fn poll_write( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + if self.refused { + return Poll::Ready(Ok(buffer.len())); + } + let written = ready!(Pin::new(&mut self.stream).poll_write(context, buffer)); + Poll::Ready(self.written(written, buffer.len())) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffers: &[IoSlice<'_>], + ) -> Poll> { + let length = buffers.iter().map(|buffer| buffer.len()).sum(); + if self.refused { + return Poll::Ready(Ok(length)); + } + let written = ready!(Pin::new(&mut self.stream).poll_write_vectored(context, buffers)); + Poll::Ready(self.written(written, length)) + } + + fn is_write_vectored(&self) -> bool { + self.stream.is_write_vectored() + } + + fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if self.refused { + return Poll::Ready(Ok(())); + } + Pin::new(&mut self.stream).poll_flush(context) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_shutdown(context) + } +} + +impl Connection for AnswerAfterRefusal { + fn connected(&self) -> Connected { + self.stream.connected() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/error.rs b/litellm-rust/crates/gateway-forward/src/error.rs new file mode 100644 index 00000000000..7dbdde68a85 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/error.rs @@ -0,0 +1,91 @@ +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, +}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid forwarding upstream URI")] + UpstreamUri(#[source] axum::http::uri::InvalidUri), + #[error( + "LITELLM_FORWARD_UPSTREAM must be a loopback IPv4 http origin without a path, query or credentials" + )] + Upstream, + #[error("{0} must be set in forward mode")] + MissingSetting(String), + #[error("invalid forwarding setting: {0}")] + InvalidSetting(String), + #[error("forwarding gateway I/O failed: {0}")] + Io(#[from] std::io::Error), +} + +/// Why a request got an answer from Rust instead of from Python. +#[derive(Debug, thiserror::Error)] +pub(crate) enum Rejection { + #[error("HTTP/1.1 request without a Host header")] + MissingHost, + #[error("request target has no path to forward")] + NoPath, + #[error("request uses a transfer coding other than chunked")] + UnsupportedTransferCoding, + #[error("client request body failed: {0}")] + ClientBody(#[source] hyper_util::client::legacy::Error), + #[error("python proxy unreachable: {0}")] + Unreachable(#[source] hyper_util::client::legacy::Error), + #[error("python proxy switched protocols without being asked to")] + UnrequestedUpgrade, + #[error("python proxy websocket upgrade failed: {0}")] + UpgradeFailed(#[source] hyper::Error), + #[error("forwarding gateway is shutting down")] + ShuttingDown, + #[error("could not build the request for the python proxy")] + Internal, +} + +impl Rejection { + /// A client body that ends early or breaks its framing fails the upstream request too, but + /// the fault is the client's, not Python's. + pub fn from_upstream(error: hyper_util::client::legacy::Error) -> Self { + let client_body = std::error::Error::source(&error) + .and_then(|source| source.downcast_ref::()) + .is_some_and(hyper::Error::is_user); + match client_body { + true => Self::ClientBody(error), + false => Self::Unreachable(error), + } + } + + fn status(&self) -> StatusCode { + match self { + Self::MissingHost + | Self::NoPath + | Self::UnsupportedTransferCoding + | Self::ClientBody(_) => StatusCode::BAD_REQUEST, + Self::Unreachable(_) | Self::UnrequestedUpgrade | Self::UpgradeFailed(_) => { + StatusCode::BAD_GATEWAY + } + Self::ShuttingDown => StatusCode::SERVICE_UNAVAILABLE, + Self::Internal => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn is_gateway_fault(&self) -> bool { + matches!( + self, + Self::Unreachable(_) + | Self::UnrequestedUpgrade + | Self::UpgradeFailed(_) + | Self::Internal + ) + } +} + +impl IntoResponse for Rejection { + fn into_response(self) -> Response { + match self.is_gateway_fault() { + true => tracing::warn!(error = %self, "request not forwarded"), + false => tracing::debug!(error = %self, "request not forwarded"), + } + self.status().into_response() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/forward.rs b/litellm-rust/crates/gateway-forward/src/forward.rs new file mode 100644 index 00000000000..69fb63ada7d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/forward.rs @@ -0,0 +1,169 @@ +use std::{ + pin::Pin, + task::{Context, Poll}, + time::Duration, +}; + +use axum::{ + body::Body, + extract::Request, + http::{StatusCode, Uri}, + response::{IntoResponse, Response}, +}; +use hyper::body::{Frame, Incoming, SizeHint}; +use hyper_util::{ + client::legacy::{Client, connect::HttpConnector}, + rt::{TokioExecutor, TokioTimer}, +}; +use tokio::net::TcpListener; + +use crate::{ + Rejection, Upstream, connections::Connections, connector::Connector, rewrite, server, + server::ClientSocket, websocket, +}; + +// h11 counts no headers and only caps a head at 16 KiB, so it accepts as many headers as fit in +// that many shortest lines (`a:\r\n`). Python copies every provider header into its response +// too, so hyper's default of 100 would turn ordinary traffic into 431s and 502s. +pub(crate) const MAX_HEADERS: usize = 16 * 1024 / b"a:\r\n".len(); + +#[derive(Clone)] +pub struct Forward { + upstream: Uri, + client: Client, + connections: Connections, + keepalive: Duration, +} + +impl Forward { + /// `upstream_keepalive` is how long the upstream keeps an idle connection open. Pooled + /// connections are dropped well before that, so a request never lands on one it is closing, + /// and idle client connections are closed after it, as uvicorn did. + pub fn new(upstream: &Upstream, upstream_keepalive: Duration) -> Self { + let mut connector = HttpConnector::new(); + connector.set_nodelay(true); + let client = Client::builder(TokioExecutor::new()) + .http1_max_headers(MAX_HEADERS) + .pool_timer(TokioTimer::new()) + .pool_idle_timeout(upstream_keepalive / 2) + .set_host(false) + .build(Connector(connector)); + Self { + upstream: upstream.uri().clone(), + client, + connections: Connections::default(), + keepalive: upstream_keepalive, + } + } + + /// Serves until `shutdown` resolves, then lets open requests and WebSockets finish for up to + /// `drain` before closing whatever is left. Python shuts down alongside and ends long streams + /// itself, so `drain` only backstops a Python that never does. + pub async fn serve( + self, + listener: TcpListener, + shutdown: impl Future, + drain: Duration, + ) { + server::serve(self, listener, shutdown, drain).await; + } + + pub(crate) fn connections(&self) -> &Connections { + &self.connections + } + + pub(crate) fn keepalive(&self) -> Duration { + self.keepalive + } + + pub(crate) async fn handle(&self, client: ClientSocket, request: Request) -> Response { + self.forward(client, request) + .await + .unwrap_or_else(IntoResponse::into_response) + } + + async fn forward( + &self, + client: ClientSocket, + mut request: Request, + ) -> Result { + if let Some(downstream) = websocket::take_handshake(&self.connections, &mut request)? { + return self.forward_websocket(client, request, downstream).await; + } + let request = rewrite::upstream_request(request, &self.upstream, client, false)?; + let response = self + .client + .request(request) + .await + .map_err(Rejection::from_upstream)?; + http_response(response) + } + + async fn forward_websocket( + &self, + client: ClientSocket, + request: Request, + downstream: hyper::upgrade::OnUpgrade, + ) -> Result { + let request = rewrite::upstream_request(request, &self.upstream, client, true)?; + let response = self + .connections + .until_closed(self.client.request(request)) + .await + .ok_or(Rejection::ShuttingDown)? + .map_err(Rejection::from_upstream)?; + if response.status() == StatusCode::SWITCHING_PROTOCOLS { + return websocket::switch(&self.connections, downstream, response).await; + } + http_response(response) + } +} + +fn http_response(response: Response) -> Result { + if response.status() == StatusCode::SWITCHING_PROTOCOLS { + return Err(Rejection::UnrequestedUpgrade); + } + Ok( + rewrite::downstream_response(response) + .map(|body| Body::new(Relayed { body, failed: None })), + ) +} + +/// Python's response body on its way to the client. +struct Relayed { + body: Incoming, + failed: Option, +} + +impl hyper::body::Body for Relayed { + type Data = ::Data; + type Error = hyper::Error; + + // hyper aborts a client connection whose response body fails without flushing the head it + // buffered, so a head Python sent right before dying would never arrive and SDKs would retry + // a request Python already ran. Holding the error for one poll lets hyper flush the head. + fn poll_frame( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + if let Some(error) = self.failed.take() { + return Poll::Ready(Some(Err(error))); + } + match Pin::new(&mut self.body).poll_frame(context) { + Poll::Ready(Some(Err(error))) => { + self.failed = Some(error); + context.waker().wake_by_ref(); + Poll::Pending + } + frame => frame, + } + } + + fn is_end_stream(&self) -> bool { + self.failed.is_none() && self.body.is_end_stream() + } + + fn size_hint(&self) -> SizeHint { + self.body.size_hint() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/input.rs b/litellm-rust/crates/gateway-forward/src/input.rs new file mode 100644 index 00000000000..772d4610683 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/input.rs @@ -0,0 +1,29 @@ +use std::io::{self, Read}; + +use tokio::sync::oneshot; + +/// Resolves once `input` reaches end of file or fails to read. The Python supervisor holds the only +/// write end of the gateway's stdin, so EOF means it is gone and the gateway drains as on SIGTERM. +/// +/// The read runs on a plain thread because a blocked read on tokio's blocking pool would hold up +/// runtime shutdown after a signal had already stopped the gateway. +pub fn input_closed(mut input: impl Read + Send + 'static) -> impl Future + Send { + let (closed, on_closed) = oneshot::channel::<()>(); + let reader = std::thread::Builder::new() + .name("stdin-eof".into()) + .spawn(move || { + let _ = io::copy(&mut input, &mut io::sink()); + let _ = closed.send(()); + }); + async move { + match reader { + Ok(_) => { + let _ = on_closed.await; + } + Err(error) => { + tracing::warn!(%error, "could not watch stdin, so only signals stop the gateway"); + std::future::pending::<()>().await; + } + } + } +} diff --git a/litellm-rust/crates/gateway-forward/src/lib.rs b/litellm-rust/crates/gateway-forward/src/lib.rs new file mode 100644 index 00000000000..35f25d19c3c --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/lib.rs @@ -0,0 +1,31 @@ +mod connections; +mod connector; +mod error; +mod forward; +mod input; +mod rewrite; +mod server; +mod settings; +mod websocket; + +pub use error::Error; +pub(crate) use error::Rejection; +pub use forward::Forward; +pub use input::input_closed; +pub use rewrite::PEER_HEADER; +pub use settings::{Settings, Upstream}; + +/// Binds the public listener and forwards to Python until `shutdown` resolves, then drains for +/// `settings.shutdown_drain`. +pub async fn run( + settings: Settings, + shutdown: impl Future + Send + 'static, +) -> Result<(), Error> { + let forward = Forward::new(&settings.upstream, settings.upstream_keepalive); + let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; + tracing::info!(address = %listener.local_addr()?, "forwarding gateway listening"); + forward + .serve(listener, shutdown, settings.shutdown_drain) + .await; + Ok(()) +} diff --git a/litellm-rust/crates/gateway-forward/src/rewrite.rs b/litellm-rust/crates/gateway-forward/src/rewrite.rs new file mode 100644 index 00000000000..fa22bf19037 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/rewrite.rs @@ -0,0 +1,161 @@ +use axum::{ + body::HttpBody, + extract::Request, + http::{HeaderMap, HeaderValue, Response, Uri, Version, header}, +}; + +use crate::{Rejection, server::ClientSocket}; + +pub const PEER_HEADER: &str = "x-litellm-rust-peer"; + +/// The request Python receives: the client's method, target, headers and body unchanged, sent to +/// `upstream` over HTTP/1.1 without the client connection's own headers, plus the client's TCP +/// peer in `PEER_HEADER`, which replaces any copy the client sent. +pub(crate) fn upstream_request( + mut request: Request, + upstream: &Uri, + client: ClientSocket, + websocket: bool, +) -> Result { + if !only_chunked(request.headers()) { + return Err(Rejection::UnsupportedTransferCoding); + } + let path_and_query = request + .uri() + .path_and_query() + .cloned() + .ok_or(Rejection::NoPath)?; + if let Some(host) = public_host(&request, client)? { + request.headers_mut().insert(header::HOST, host); + } + let mut uri = upstream.clone().into_parts(); + uri.path_and_query = Some(path_and_query); + *request.uri_mut() = Uri::from_parts(uri).map_err(|_| Rejection::Internal)?; + *request.version_mut() = Version::HTTP_11; + let chunked = request.body().size_hint().exact().is_none(); + let headers = request.headers_mut(); + strip_hop_by_hop(headers); + if chunked { + // hyper sends no body for a GET with an unknown length unless told it is chunked. + headers.insert( + header::TRANSFER_ENCODING, + HeaderValue::from_static("chunked"), + ); + } + if websocket { + restore_websocket_handshake(headers); + } + headers.insert( + PEER_HEADER, + HeaderValue::from_str(&client.peer.to_string()).map_err(|_| Rejection::Internal)?, + ); + Ok(request) +} + +// Python builds absolute URLs from Host and, without a usable one, from the address it listens +// on, which behind Rust is the private loopback port. A client that sent no usable Host gets the +// public address it connected to, which is what Python used before. h11 rejects HTTP/1.1 without +// Host. +fn public_host(request: &Request, client: ClientSocket) -> Result, Rejection> { + let host = request.headers().get(header::HOST); + if host.is_some_and(|host| !host.is_empty()) { + return Ok(None); + } + if host.is_none() && request.version() == Version::HTTP_11 { + return Err(Rejection::MissingHost); + } + HeaderValue::from_str(&client.local.to_string()) + .map(Some) + .map_err(|_| Rejection::Internal) +} + +// hyper only decodes chunked framing, so anything else would reach Python still coded but framed +// as a plain body. uvicorn rejects it with the 400 it sends for every h11 parse error. +fn only_chunked(headers: &HeaderMap) -> bool { + headers + .get_all(header::TRANSFER_ENCODING) + .iter() + .flat_map(|value| value.as_bytes().split(|byte| *byte == b',')) + .all(|coding| coding.trim_ascii().eq_ignore_ascii_case(b"chunked")) +} + +/// The response the client receives: Python's, without the upstream connection's own headers. +pub(crate) fn downstream_response(mut response: Response) -> Response { + strip_hop_by_hop(response.headers_mut()); + response +} + +/// Python's 101 as the client receives it, still naming the protocol both sides switch to. +pub(crate) fn switched_response(mut response: Response) -> Response { + strip_hop_by_hop(response.headers_mut()); + restore_websocket_handshake(response.headers_mut()); + response +} + +pub(crate) fn is_websocket_handshake(headers: &HeaderMap) -> bool { + headers + .get(header::UPGRADE) + .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket")) + && headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .any(|token| token.trim().eq_ignore_ascii_case("upgrade")) +} + +const HOP_BY_HOP: [&str; 9] = [ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + "proxy-connection", +]; + +// Per-connection headers (RFC 9110 section 7.6.1) describe the client's connection to Rust, not +// Rust's connection to Python, so a proxy must drop them instead of forwarding them. +fn strip_hop_by_hop(headers: &mut HeaderMap) { + let named: Vec<_> = headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .filter_map(|name| header::HeaderName::from_bytes(name.trim().as_bytes()).ok()) + .filter(|name| !must_forward(name)) + .collect(); + for name in named { + headers.remove(name); + } + for name in HOP_BY_HOP { + headers.remove(name); + } +} + +fn must_forward(name: &header::HeaderName) -> bool { + matches!( + name.as_str(), + "host" + | "authorization" + | "cookie" + | "set-cookie" + | "www-authenticate" + | "authentication-info" + | "api-key" + | "x-api-key" + | "litellm-api-key" + | "forwarded" + | PEER_HEADER + ) || name.as_str().starts_with("content-") + || name.as_str().starts_with("x-forwarded-") + || name.as_str().starts_with("sec-websocket-") +} + +// Stripping drops the handshake's Connection and Upgrade, so each hop gets them back. +fn restore_websocket_handshake(headers: &mut HeaderMap) { + headers.insert(header::CONNECTION, HeaderValue::from_static("upgrade")); + headers.insert(header::UPGRADE, HeaderValue::from_static("websocket")); +} diff --git a/litellm-rust/crates/gateway-forward/src/server.rs b/litellm-rust/crates/gateway-forward/src/server.rs new file mode 100644 index 00000000000..0388642446a --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/server.rs @@ -0,0 +1,116 @@ +use std::{convert::Infallible, io, net::SocketAddr, pin::pin, time::Duration}; + +use axum::{body::Body, extract::Request}; +use hyper::{body::Incoming, server::conn::http1, service::service_fn}; +use hyper_util::rt::{TokioIo, TokioTimer}; +use tokio::net::TcpListener; + +use crate::{Forward, forward::MAX_HEADERS}; + +/// The client's TCP connection to the public listener. +#[derive(Clone, Copy, Debug)] +pub(crate) struct ClientSocket { + pub peer: SocketAddr, + pub local: SocketAddr, +} + +/// Serves each public connection with hyper directly, because `axum::serve` cannot configure +/// the connection timeouts uvicorn applied to the same clients. +pub(crate) async fn serve( + forward: Forward, + listener: TcpListener, + shutdown: impl Future, + drain: Duration, +) { + let connections = forward.connections().clone(); + let _close_on_drop = connections.close_on_drop(); + let keepalive = forward.keepalive(); + let builder = builder(keepalive); + let mut shutdown = pin!(shutdown); + loop { + let (stream, peer) = tokio::select! { + () = &mut shutdown => break, + accepted = listener.accept() => match accepted { + Ok(accepted) => accepted, + Err(error) => { + back_off(error).await; + continue; + } + }, + }; + let local = match stream.local_addr() { + Ok(local) => local, + Err(error) => { + tracing::debug!(%error, "client connection closed before it was served"); + continue; + } + }; + if let Err(error) = stream.set_nodelay(true) { + tracing::warn!(%error, "could not disable Nagle on a client connection"); + } + let socket = ClientSocket { peer, local }; + let forward = forward.clone(); + let service = service_fn(move |request: Request| { + let forward = forward.clone(); + async move { Ok::<_, Infallible>(forward.handle(socket, request.map(Body::new)).await) } + }); + let connection = builder + .serve_connection(TokioIo::new(stream), service) + .with_upgrades(); + let draining = connections.clone(); + connections.spawn(async move { + let mut connection = pin!(connection); + let ended = tokio::select! { + ended = connection.as_mut() => ended, + () = draining.draining() => { + connection.as_mut().graceful_shutdown(); + connection.await + } + }; + if let Err(error) = ended { + tracing::debug!(%error, "client connection ended with an error"); + } + }); + } + drop(listener); + connections.begin_shutdown(); + if tokio::time::timeout(drain, connections.wait()) + .await + .is_err() + { + tracing::warn!( + ?drain, + "closing connections still open after the shutdown drain" + ); + } + connections.close(); + connections.wait().await; +} + +// uvicorn closes a client connection that sits idle for `--keepalive_timeout` between requests. +// Without the same limit, a client that vanishes without closing (NAT expiry, a crashed host) +// would hold its socket for the life of the process until the gateway runs out of descriptors. +// Operators already tune that timeout against their load balancer, so the public side keeps +// Python's value instead of introducing its own. hyper only runs the header timer while a +// connection waits for its next request, so it never cuts a long stream, upload or WebSocket. +fn builder(keepalive: Duration) -> http1::Builder { + let mut builder = http1::Builder::new(); + builder + .max_headers(MAX_HEADERS) + .timer(TokioTimer::new()) + .header_read_timeout(keepalive); + builder +} + +async fn back_off(error: io::Error) { + if matches!( + error.kind(), + io::ErrorKind::ConnectionRefused + | io::ErrorKind::ConnectionAborted + | io::ErrorKind::ConnectionReset + ) { + return; + } + tracing::error!(%error, "could not accept a client connection"); + tokio::time::sleep(Duration::from_secs(1)).await; +} diff --git a/litellm-rust/crates/gateway-forward/src/settings.rs b/litellm-rust/crates/gateway-forward/src/settings.rs new file mode 100644 index 00000000000..de01a96c519 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/settings.rs @@ -0,0 +1,113 @@ +use std::{ffi::OsString, net::SocketAddr, str::FromStr, time::Duration}; + +use axum::http::Uri; +use serde::Deserialize; + +use crate::Error; + +const UPSTREAM: &str = "LITELLM_FORWARD_UPSTREAM"; + +/// The Python proxy's private origin. Only a loopback IPv4 address parses, so nothing can point +/// the forwarder at another host. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Upstream(Uri); + +impl Upstream { + pub(crate) fn uri(&self) -> &Uri { + &self.0 + } +} + +impl FromStr for Upstream { + type Err = Error; + + fn from_str(value: &str) -> Result { + let uri: Uri = value.parse().map_err(Error::UpstreamUri)?; + let address = uri + .authority() + .filter(|_| uri.scheme_str() == Some("http")) + .filter(|authority| !authority.as_str().contains('@')) + .filter(|_| uri.path() == "/" && uri.query().is_none()) + .and_then(|authority| authority.as_str().parse::().ok()); + match address { + Some(SocketAddr::V4(address)) if address.ip().is_loopback() => Ok(Self(uri)), + _ => Err(Error::Upstream), + } + } +} + +#[derive(Deserialize)] +struct Environment { + litellm_forward_upstream: String, + litellm_forward_host: String, + litellm_forward_port: u16, + litellm_forward_upstream_keepalive: f64, + litellm_forward_shutdown_drain: f64, +} + +/// Everything forward mode runs with. The Python supervisor resolves each value from the settings +/// Python actually runs with, so none of them is guessed or defaulted here. +/// Setting `LITELLM_FORWARD_UPSTREAM` requires every other `LITELLM_FORWARD_*` setting below. +#[derive(Debug, Clone, PartialEq)] +pub struct Settings { + pub upstream: Upstream, + pub host: String, + pub port: u16, + pub upstream_keepalive: Duration, + pub shutdown_drain: Duration, +} + +impl Settings { + /// `None` when `LITELLM_FORWARD_UPSTREAM` is unset, which means the gateway is not forwarding. + pub fn from_env() -> Result, Error> { + Self::from_vars(std::env::vars_os()) + } + + pub fn from_vars( + variables: impl IntoIterator, + ) -> Result, Error> { + let variables: Vec<(String, String)> = variables + .into_iter() + .filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?))) + .collect(); + if !variables.iter().any(|(key, _)| key == UPSTREAM) { + return Ok(None); + } + let environment: Environment = envy::from_iter(variables).map_err(|error| match error { + envy::Error::MissingValue(field) => Error::MissingSetting(field.to_uppercase()), + envy::Error::Custom(message) => Error::InvalidSetting(message), + })?; + Ok(Some(Self { + upstream: environment.litellm_forward_upstream.parse()?, + host: environment.litellm_forward_host, + port: environment.litellm_forward_port, + upstream_keepalive: positive_seconds( + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE", + environment.litellm_forward_upstream_keepalive, + )?, + shutdown_drain: seconds( + "LITELLM_FORWARD_SHUTDOWN_DRAIN", + environment.litellm_forward_shutdown_drain, + )?, + })) + } +} + +fn seconds(variable: &str, value: f64) -> Result { + Duration::try_from_secs_f64(value).map_err(|_| { + Error::InvalidSetting(format!( + "{variable} must be a non-negative number of seconds, got {value}" + )) + }) +} + +/// The keep-alive also bounds the wait for a client's first request, so zero would close every +/// client connection before it could send one. +fn positive_seconds(variable: &str, value: f64) -> Result { + match seconds(variable, value)? { + Duration::ZERO => Err(Error::InvalidSetting(format!( + "{variable} must be a positive number of seconds, got {value}" + ))), + duration => Ok(duration), + } +} diff --git a/litellm-rust/crates/gateway-forward/src/websocket.rs b/litellm-rust/crates/gateway-forward/src/websocket.rs new file mode 100644 index 00000000000..84d89c75607 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/websocket.rs @@ -0,0 +1,86 @@ +use axum::{ + body::Body, + extract::Request, + http::{Method, Version}, + response::Response, +}; +use hyper::{ + body::Incoming, + upgrade::{OnUpgrade, Upgraded}, +}; +use hyper_util::rt::TokioIo; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +use crate::{Rejection, connections::Connections, rewrite}; + +pub(crate) fn take_handshake( + connections: &Connections, + request: &mut Request, +) -> Result, Rejection> { + if request.method() != Method::GET + || request.version() != Version::HTTP_11 + || !rewrite::is_websocket_handshake(request.headers()) + { + return Ok(None); + } + if connections.is_draining() { + return Err(Rejection::ShuttingDown); + } + Ok(Some(hyper::upgrade::on(request))) +} + +pub(crate) async fn switch( + connections: &Connections, + downstream: OnUpgrade, + mut response: Response, +) -> Result { + if !rewrite::is_websocket_handshake(response.headers()) { + return Err(Rejection::UnrequestedUpgrade); + } + let upstream = connections + .until_closed(hyper::upgrade::on(&mut response)) + .await + .ok_or(Rejection::ShuttingDown)? + .map_err(Rejection::UpgradeFailed)?; + connections.spawn(async move { + match downstream.await { + Ok(downstream) => relay(downstream, upstream).await, + Err(_) => tracing::debug!("websocket upgrade did not complete"), + } + }); + Ok(rewrite::switched_response(response).map(|_| Body::empty())) +} + +// uvicorn never half-closes a WebSocket, so Python's EOF means the socket is gone and the relay +// ends. Only a client that is gone ends it from the other side: the client's EOF just half-closes +// Python's side, and a write Python refuses means it closed with its last frames still in flight +// to the client. +async fn relay(downstream: Upgraded, upstream: Upgraded) { + let (mut client_read, mut client_write) = tokio::io::split(TokioIo::new(downstream)); + let (mut python_read, mut python_write) = tokio::io::split(TokioIo::new(upstream)); + let to_python = async { + let mut buffer = vec![0; 8192]; + loop { + let read = match client_read.read(&mut buffer).await { + Ok(0) => break, + Ok(read) => read, + Err(_) => return, + }; + if python_write.write_all(&buffer[..read]).await.is_err() { + return std::future::pending().await; + } + } + python_write.shutdown().await.ok(); + std::future::pending().await + }; + let to_client = async { + tokio::io::copy(&mut python_read, &mut client_write) + .await + .ok(); + client_write.shutdown().await.ok(); + }; + tokio::select! { + () = to_python => tracing::debug!("websocket client disconnected"), + () = to_client => {} + } +} diff --git a/litellm-rust/crates/gateway-forward/tests/client_identity.rs b/litellm-rust/crates/gateway-forward/tests/client_identity.rs new file mode 100644 index 00000000000..c484c89dc1d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/client_identity.rs @@ -0,0 +1,223 @@ +mod support; + +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use axum::{body::Body, http::StatusCode, response::Response}; +use litellm_gateway_forward::{Forward, PEER_HEADER}; +use rstest::rstest; +use support::{DRAIN, Gateway, Raw, UVICORN_KEEPALIVE, raw_upstream, recording, within}; +use tokio::net::{TcpListener, TcpStream}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +#[rstest] +#[case::ordinary("")] +#[case::connection_nominated( + "connection: X-Forwarded-For, X-Forwarded-Proto, Forwarded, X-Litellm-Rust-Peer\r\n" +)] +#[tokio::test] +async fn python_sees_forwarding_headers_untouched_and_the_tcp_peer_rust_saw( + #[case] connection: &str, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "GET /sso/callback HTTP/1.1\r\n\ + host: public.example\r\n\ + x-forwarded-for: 192.0.2.1\r\n\ + x-forwarded-for: 192.0.2.2\r\n\ + x-forwarded-proto: https\r\n\ + forwarded: for=192.0.2.1;proto=https\r\n\ + x-litellm-rust-peer: 10.1.2.3:1\r\n{connection}\r\n", + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + + let request = received.recv().await.unwrap(); + let forwarded_for: Vec<_> = request.headers.get_all("x-forwarded-for").iter().collect(); + assert_eq!(forwarded_for, ["192.0.2.1", "192.0.2.2"]); + assert_eq!(request.headers["x-forwarded-proto"], "https"); + assert_eq!(request.headers["forwarded"], "for=192.0.2.1;proto=https"); + let peers: Vec<_> = request.headers.get_all(PEER_HEADER).iter().collect(); + assert_eq!(peers, [raw.local_address().to_string().as_str()]); +} + +/// Python builds redirects and OAuth metadata from Host, and without one from the address it +/// listens on, which behind Rust is its private loopback port. So a request without a usable Host +/// carries the public address the client connected to, except HTTP/1.1, which h11 rejects. +#[rstest] +#[case::http11_without_host("HTTP/1.1", "", false)] +#[case::http10_without_host("HTTP/1.0", "", true)] +#[case::http11_with_an_empty_host("HTTP/1.1", "host:\r\n", true)] +#[tokio::test] +async fn a_request_without_a_usable_host_never_names_the_private_upstream( + #[case] version: &str, + #[case] host: &str, + #[case] forwarded: bool, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!("GET /sso/key/generate {version}\r\n{host}\r\n")) + .await; + + let status = within(raw.response_head()).await.unwrap().status; + let hosts = received.try_recv().ok().map(|request| { + request + .headers + .get_all("host") + .iter() + .map(|value| value.to_str().unwrap().to_owned()) + .collect::>() + }); + let expected = if forwarded { + (StatusCode::OK, Some(vec![gateway.address.to_string()])) + } else { + (StatusCode::BAD_REQUEST, None) + }; + assert_eq!((status, hosts), expected); +} + +#[rstest] +#[case::uppercase_with_port("PUBLIC.Example:4000")] +#[case::ipv6_literal("[::1]:4000")] +#[case::ipv4_with_port("192.0.2.10:8443")] +#[tokio::test] +async fn python_sees_the_host_header_byte_for_byte(#[case] host: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!("GET /ui HTTP/1.1\r\nHost: {host}\r\n\r\n")) + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let hosts: Vec<_> = request.headers.get_all("host").iter().collect(); + assert_eq!(hosts, [host]); +} + +/// Python's IP allowlists and rate limits read the peer exactly as uvicorn would have seen it on +/// the public socket, so a dual-stack listener reports IPv4 clients in their mapped form. +#[rstest] +#[case::ipv6("[::1]:0", "[::1]", "[::1]")] +#[case::ipv4_on_a_dual_stack_listener("[::]:0", "127.0.0.1", "[::ffff:127.0.0.1]")] +#[tokio::test] +async fn ipv6_listeners_report_the_peer_bracketed_with_its_port( + #[case] bind: &str, + #[case] connect: &str, + #[case] peer_ip: &str, +) { + let (upstream, mut received) = recording(ok).await; + let forward = Forward::new( + &format!("http://{upstream}").parse().unwrap(), + UVICORN_KEEPALIVE, + ); + let public = TcpListener::bind(bind).await.unwrap(); + let port = public.local_addr().unwrap().port(); + tokio::spawn(forward.serve(public, std::future::pending(), DRAIN)); + let mut raw = Raw::new( + TcpStream::connect(format!("{connect}:{port}")) + .await + .unwrap(), + ); + + raw.send("GET / HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!( + received.recv().await.unwrap().headers[PEER_HEADER], + format!("{peer_ip}:{}", raw.local_address().port()).as_str() + ); +} + +/// Without the peer header Python would take Rust's loopback hop as the client, which allowlists +/// often trust, so no spelling, repetition or Connection token may remove or replace it. +#[rstest] +#[case::named_as_hop_by_hop( + "connection: x-litellm-rust-peer\r\nx-litellm-rust-peer: 127.0.0.1:1\r\n" +)] +#[case::mixed_case_copies( + "X-LiteLLM-Rust-Peer: 127.0.0.1:1\r\nx-litellm-rust-peer: 192.0.2.7:2\r\n" +)] +#[case::comma_joined("x-litellm-rust-peer: 127.0.0.1:1, 192.0.2.7:2\r\n")] +#[tokio::test] +async fn a_client_can_neither_strip_nor_spoof_the_peer(#[case] spoof: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!( + "GET /key/info HTTP/1.1\r\nhost: public.example\r\n{spoof}\r\n" + )) + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let peers: Vec<_> = request.headers.get_all(PEER_HEADER).iter().collect(); + assert_eq!(peers, [raw.local_address().to_string().as_str()]); +} + +/// Pooled upstream connections carry requests from many clients in turn, so the peer belongs to +/// each request and never to the connection it travels on. +#[rstest] +#[tokio::test] +async fn each_request_carries_its_own_peer_over_a_shared_upstream_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + while let Some(head) = python.request_head().await { + let peer = head.header(PEER_HEADER).unwrap().to_vec(); + python + .send( + [ + format!("HTTP/1.1 200 OK\r\ncontent-length: {}\r\n\r\n", peer.len()) + .as_bytes(), + &peer, + ] + .concat(), + ) + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut clients = [ + gateway.connect().await, + gateway.connect().await, + gateway.connect().await, + ]; + + for index in [0, 1, 2, 2, 1, 0] { + let raw = &mut clients[index]; + raw.send("GET /v1/models HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + let peer = raw.body(head.content_length().unwrap()).await; + assert_eq!(peer, raw.local_address().to_string().into_bytes()); + } + assert_eq!(connections.load(Ordering::SeqCst), 1); +} diff --git a/litellm-rust/crates/gateway-forward/tests/download.rs b/litellm-rust/crates/gateway-forward/tests/download.rs new file mode 100644 index 00000000000..9707d2f53b7 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/download.rs @@ -0,0 +1,192 @@ +mod support; + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{body::Body, extract::Request, http::StatusCode, response::Response}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, recording, within}; +use tokio::sync::{Notify, mpsc}; + +#[rstest] +#[tokio::test] +async fn streams_a_body_without_a_content_length_as_python_produces_it() { + let first: &[u8] = &[0xff, 0xfb, 0x90, 0x64, 0x00, 0x0d, 0x0a]; + let rest = vec![0xa5; 64 * 1024]; + let finish = Arc::new(Notify::new()); + let finished = finish.clone(); + let tail = rest.clone(); + let upstream = raw_upstream(move |mut python| { + let finished = finished.clone(); + let tail = tail.clone(); + async move { + python.request_head().await.unwrap(); + python + .send([&b"HTTP/1.1 200 OK\r\ncontent-type: audio/mpeg\r\ntransfer-encoding: chunked\r\n\r\n7\r\n"[..], first, b"\r\n"].concat()) + .await; + finished.notified().await; + python + .send([format!("{:x}\r\n", tail.len()).as_bytes(), &tail, b"\r\n0\r\n\r\n"].concat()) + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut response = within( + client().request( + Request::post(gateway.url("/v1/audio/speech")) + .body(Body::empty()) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert!(!response.headers().contains_key("content-length")); + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), first); + finish.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, rest); +} + +#[rstest] +#[tokio::test] +async fn range_requests_and_partial_content_pass_unchanged() { + let (upstream, mut received) = recording(|| -> Response { + Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header("content-range", "bytes 2-5/10") + .header("accept-ranges", "bytes") + .body(Body::from("2345")) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/videos/video_1/content HTTP/1.1\r\nhost: x\r\nrange: bytes=2-5\r\nif-range: \"v1\"\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::PARTIAL_CONTENT); + assert_eq!(head.header("content-range"), Some(&b"bytes 2-5/10"[..])); + assert_eq!(raw.body(head.content_length().unwrap()).await, b"2345"); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["range"], "bytes=2-5"); + assert_eq!(request.headers["if-range"], "\"v1\""); +} + +#[rstest] +#[tokio::test] +async fn a_large_download_is_paced_by_the_client_and_keeps_its_content_length() { + const LENGTH: usize = 128 * 1024 * 1024; + let written = Arc::new(AtomicUsize::new(0)); + let writing = written.clone(); + let upstream = raw_upstream(move |mut python| { + let writing = writing.clone(); + async move { + python.request_head().await.unwrap(); + python + .send(format!("HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: {LENGTH}\r\n\r\n")) + .await; + let chunk = [0xa5; 64 * 1024]; + while writing.load(Ordering::SeqCst) < LENGTH { + python.send(chunk).await; + writing.fetch_add(chunk.len(), Ordering::SeqCst); + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/files/file-1/content HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.content_length(), Some(LENGTH)); + assert_eq!(head.header("transfer-encoding"), None); + tokio::time::sleep(Duration::from_secs(1)).await; + assert!( + written.load(Ordering::SeqCst) < LENGTH / 4, + "the forwarder buffered a download the client was not reading" + ); + let mut remaining = LENGTH; + while remaining > 0 { + remaining -= within(raw.body(remaining.min(1 << 20))).await.len(); + } +} + +/// The UI revalidates its static assets, and a 304 must carry exactly the validators Python sent, +/// with no length invented for a body it never has. +#[rstest] +#[tokio::test] +async fn a_conditional_request_gets_python_304_with_its_validators_only() { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + event + .send(head.header("if-none-match").map(<[u8]>::to_vec)) + .unwrap(); + python + .send("HTTP/1.1 304 Not Modified\r\netag: \"5f1-abc\"\r\ncache-control: no-cache\r\nvary: Accept-Encoding\r\ndate: Tue, 03 Nov 2026 10:00:00 GMT\r\n\r\n") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + raw.send("GET /ui/_next/static/chunks/main.js HTTP/1.1\r\nhost: x\r\nif-none-match: \"5f1-abc\"\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::NOT_MODIFIED); + let mut names: Vec<_> = head.headers.iter().map(|(name, _)| name.as_str()).collect(); + names.sort_unstable(); + assert_eq!(names, ["cache-control", "date", "etag", "vary"]); + assert_eq!(head.header("etag"), Some(&b"\"5f1-abc\""[..])); + assert_eq!( + head.header("date"), + Some(&b"Tue, 03 Nov 2026 10:00:00 GMT"[..]) + ); + assert_eq!( + within(events.recv()).await.unwrap().as_deref(), + Some(&b"\"5f1-abc\""[..]) + ); + } +} + +/// A client knows a download was cut short only because fewer bytes than the Content-Length +/// arrived before the connection closed, so the length must survive and the connection must end. +#[rstest] +#[tokio::test] +async fn a_download_python_cuts_short_closes_the_client_connection_before_its_length() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send([&b"HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: 1000\r\n\r\n"[..], &[0xa5; 500]].concat()) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/files/file-1/content HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.content_length(), Some(1000)); + assert_eq!(within(raw.until_closed()).await, [0xa5; 500]); +} diff --git a/litellm-rust/crates/gateway-forward/tests/exchange.rs b/litellm-rust/crates/gateway-forward/tests/exchange.rs new file mode 100644 index 00000000000..8fed7b18453 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/exchange.rs @@ -0,0 +1,855 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::{HeaderValue, StatusCode}, + response::Response, +}; +use litellm_gateway_forward::PEER_HEADER; +use rstest::rstest; +use support::{Gateway, Head, client, raw_upstream, recording, serve, within}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + sync::mpsc, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +/// An upstream that reports each bodiless request head exactly as it arrived on the wire. +async fn heads() -> (std::net::SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = raw_upstream(move |mut python| { + let sender = sender.clone(); + async move { + while let Some(head) = python.request_head().await { + sender.send(head).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + (address, receiver) +} + +#[rstest] +#[tokio::test] +async fn forwards_the_request_as_the_client_sent_it_minus_hop_headers() { + let (upstream, mut received) = recording(|| { + Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header("connection", "x-upstream-hop") + .header("x-upstream-hop", "remove") + .header("x-response", "keep") + .body(Body::from("python error")) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + "PATCH /v1/messages?x=a%2Fb&x=c HTTP/1.1\r\n\ + host: public.example\r\n\ + authorization: Bearer test\r\n\ + connection: x-remove, keep-alive\r\n\ + x-remove: remove\r\n\ + proxy-authorization: remove\r\n\ + te: trailers\r\n\ + content-length: 7\r\n\r\npayload", + ) + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::UNAUTHORIZED); + assert_eq!(head.header("x-response"), Some(&b"keep"[..])); + assert_eq!(head.header("x-upstream-hop"), None); + let body = raw.body(head.content_length().unwrap()).await; + assert_eq!(body, b"python error"); + + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "PATCH"); + assert_eq!(request.uri, "/v1/messages?x=a%2Fb&x=c"); + assert_eq!(request.headers["host"], "public.example"); + assert_eq!(request.headers["authorization"], "Bearer test"); + assert!(!request.headers.contains_key("x-remove")); + assert!(!request.headers.contains_key("proxy-authorization")); + assert!(!request.headers.contains_key("te")); + assert!(!request.headers.contains_key("connection")); + assert_eq!(request.body, "payload"); +} + +#[rstest] +#[case::host("HoSt", "public.example")] +#[case::length("Content-Length", "7")] +#[case::authorization("AUTHORIZATION", "Bearer test")] +#[case::cookie("Cookie", "session=test")] +#[case::api_key("X-Api-Key", "test")] +#[case::azure_key("Api-Key", "test")] +#[case::litellm_key("Litellm-Api-Key", "test")] +#[case::forwarded("Forwarded", "for=192.0.2.1;proto=https")] +#[case::forwarded_for("X-Forwarded-For", "192.0.2.1")] +#[case::forwarded_proto("X-Forwarded-Proto", "https")] +#[case::forwarded_host("X-Forwarded-Host", "public.example")] +#[case::content_type("Content-Type", "application/octet-stream")] +#[case::content_encoding("Content-Encoding", "identity")] +#[tokio::test] +async fn connection_cannot_remove_end_to_end_headers(#[case] name: &str, #[case] value: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let host = if name.eq_ignore_ascii_case("host") { + "" + } else { + "host: public.example\r\n" + }; + let length = if name.eq_ignore_ascii_case("content-length") { + "" + } else { + "content-length: 7\r\n" + }; + raw.send(format!( + "POST / HTTP/1.1\r\n{host}{length}{name}: {value}\r\nconnection: x-remove\r\nconnection: {name}, keep-alive\r\nx-remove: remove\r\n\r\npayload" + )).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = within(received.recv()).await.unwrap(); + assert_eq!(request.headers[name], value); + assert_eq!(request.body, "payload"); + assert!(!request.headers.contains_key("connection")); + assert!(!request.headers.contains_key("x-remove")); +} + +#[rstest] +#[tokio::test] +async fn connection_cannot_remove_response_credentials_or_representation_metadata() { + let (upstream, _received) = recording(|| { + Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header("connection", "Content-Length, Set-Cookie, WWW-Authenticate, Authentication-Info, Content-Type, Content-Encoding, x-remove") + .header("content-length", "2") + .header("set-cookie", "session=; Max-Age=0") + .header("www-authenticate", "Bearer") + .header("authentication-info", "nextnonce=test") + .header("content-type", "text/plain") + .header("content-encoding", "identity") + .header("x-remove", "remove") + .body(Body::from("ok")) + .unwrap() + }).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::UNAUTHORIZED); + assert_eq!(head.content_length(), Some(2)); + assert_eq!(head.header("set-cookie"), Some(&b"session=; Max-Age=0"[..])); + assert_eq!(head.header("www-authenticate"), Some(&b"Bearer"[..])); + assert_eq!( + head.header("authentication-info"), + Some(&b"nextnonce=test"[..]) + ); + assert_eq!(head.header("content-type"), Some(&b"text/plain"[..])); + assert_eq!(head.header("content-encoding"), Some(&b"identity"[..])); + assert_eq!(head.header("x-remove"), None); + assert_eq!(raw.body(2).await, b"ok"); +} + +#[rstest] +#[case::get("GET")] +#[case::post("POST")] +#[case::put("PUT")] +#[case::delete("DELETE")] +#[case::options("OPTIONS")] +#[case::extension("PURGE")] +#[case::webdav("PROPFIND")] +#[case::lowercase("get")] +#[case::punctuated("M-SEARCH")] +#[tokio::test] +async fn forwards_every_method(#[case] method: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let response = client() + .request( + Request::builder() + .method(method) + .uri(gateway.url("/mcp")) + .body(Body::from("body")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, method); + assert_eq!(request.body, "body"); +} + +#[rstest] +#[case::double_slash("//v1//models")] +#[case::encoded_slash("/v1/files/a%2Fb/content")] +#[case::encoded_hash_in_query("/v1/files?purpose=batch%23retry&limit=1")] +#[case::encoded_hash_in_path("/v1/files/report%231.pdf/content")] +#[case::encoded_hash_before_query("/gemini/v1beta/files/a%23b?key=k")] +#[case::encoded_query("/v1/models?q=%E2%9C%93&empty=&flag")] +#[case::utf8_bytes_escaped("/mcp/%F0%9F%A6%80")] +#[case::lowercase_encoded_slash("/v1/files/a%2fb/content")] +#[case::dot_segments("/openai/./v1/../v1/models")] +#[case::encoded_dot_segments("/openai/%2e%2e/v1/%2E/models")] +#[case::matrix_params("/langfuse/api;v=1/public;x=y/traces")] +#[case::plus_and_encoded_space("/gemini/v1beta/models?q=a+b&r=a%20b&s=%2B")] +#[case::empty_query("/v1/models?")] +#[case::question_marks_in_query("/v1/models??a=b?c")] +#[case::colons_in_segments( + "/v1/projects/p/locations/l/publishers/google/models/m:streamGenerateContent?alt=sse" +)] +#[case::sub_delims("/~u/@:!$&'()*+,=/x")] +#[case::raw_bytes_browsers_leave_unescaped("/a{b}|c^d\\e\"f?g={}|^`")] +#[case::invalid_percent_escapes("/a%zz%00b?%")] +#[tokio::test] +async fn forwards_the_request_target_byte_for_byte(#[case] target: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().uri, target); +} + +#[rstest] +#[tokio::test] +async fn forwards_a_query_string_as_long_as_uvicorn_accepts() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let target = format!("/vertex_ai/v1/models?key=%2F{}&end", "a%20b".repeat(3000)); + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().uri, target.as_str()); +} + +#[rstest] +#[case::in_query("/v1/files?purpose=batch#retry&limit=1", "/v1/files?purpose=batch")] +#[case::in_path("/v1/files/report#1.pdf/content", "/v1/files/report")] +#[case::before_query("/gemini/v1beta/files/a#b?key=k", "/gemini/v1beta/files/a")] +#[tokio::test] +async fn forwards_only_the_path_and_query_before_a_fragment( + #[case] target: &str, + #[case] expected: &str, +) { + let (upstream, mut received) = heads().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(within(received.recv()).await.unwrap().path, expected); +} + +#[rstest] +#[tokio::test] +async fn absolute_form_targets_reach_python_as_origin_form() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + "GET http://elsewhere.example/v1/models?x=1 HTTP/1.1\r\nhost: elsewhere.example\r\n\r\n", + ) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.uri, "/v1/models?x=1"); +} + +#[rstest] +#[tokio::test] +async fn relays_a_response_with_more_headers_than_hyper_allows_by_default() { + let upstream = serve(axum::Router::new().fallback(|| async { + let headers = (0..150).map(|index| (format!("llm_provider-x-{index}"), "v")); + let mut response = ok(); + for (name, value) in headers { + response.headers_mut().insert( + axum::http::HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_static(value), + ); + } + response + })) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send("GET /v1/chat/completions HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.header("llm_provider-x-149"), Some(&b"v"[..])); +} + +/// uvicorn's h11 parser has no header count limit, only a 16 KiB cap on the whole head. +#[rstest] +#[tokio::test] +async fn forwards_a_request_with_as_many_headers_as_uvicorn_accepts() { + let (sender, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |python| { + let sender = sender.clone(); + async move { + let mut stream = python.into_stream(); + let mut head = Vec::new(); + while !head.ends_with(b"\r\n\r\n") { + head.push(stream.read_u8().await.unwrap()); + } + sender.send(head).unwrap(); + stream + .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await + .unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let headers: String = (0..1000).map(|index| format!("x-{index}: v\r\n")).collect(); + raw.send(format!( + "GET /v1/chat/completions HTTP/1.1\r\nhost: x\r\n{headers}\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let head = String::from_utf8(within(received.recv()).await.unwrap()).unwrap(); + assert!(head.contains(&headers)); +} + +async fn header_fields_python_receives(fields: &str) -> Vec<(String, Vec)> { + let (upstream, mut received) = heads().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "GET /anthropic/v1/messages HTTP/1.1\r\nhost: x\r\n{fields}\r\n" + )) + .await; + within(raw.response_head()).await.unwrap(); + within(received.recv()) + .await + .unwrap() + .headers + .into_iter() + .filter(|(name, _)| name.starts_with("x-") && name != PEER_HEADER) + .collect() +} + +fn fields(expected: &[(&str, &str)]) -> Vec<(String, Vec)> { + expected + .iter() + .map(|(name, value)| ((*name).to_owned(), value.as_bytes().to_vec())) + .collect() +} + +/// Python reads `scope["headers"]` as the list the client sent, in that order. +#[rstest] +#[ignore = "divergence: HeaderMap::remove can move the headers after a dropped hop-by-hop one"] +#[tokio::test] +async fn dropping_a_hop_by_hop_header_keeps_the_others_in_order() { + let received = header_fields_python_receives( + "x-a: 1\r\nconnection: keep-alive\r\nx-b: 2\r\nx-c: 3\r\nx-d: 4\r\n", + ) + .await; + assert_eq!( + received, + fields(&[("x-a", "1"), ("x-b", "2"), ("x-c", "3"), ("x-d", "4")]) + ); +} + +/// uvicorn speaks only HTTP/1.1, so an HTTP/2 prior-knowledge client never reaches Python through +/// Rust either. hyper closes the connection where h11 answers 400, which no HTTP/2 client reads. +#[rstest] +#[tokio::test] +async fn an_http2_client_never_reaches_python() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00") + .await; + + assert!(within(raw.until_closed()).await.is_empty()); + assert!(received.try_recv().is_err()); +} + +#[rstest] +#[ignore = "divergence: HeaderMap groups every value of a name at its first position"] +#[tokio::test] +async fn interleaved_repeated_headers_reach_python_interleaved() { + let received = header_fields_python_receives("x-a: 1\r\nx-b: 2\r\nx-a: 3\r\n").await; + assert_eq!( + received, + fields(&[("x-a", "1"), ("x-b", "2"), ("x-a", "3")]) + ); +} + +/// h11 trims optional whitespace around a value and keeps everything inside it, including empty +/// values and names using every token character. +#[rstest] +#[tokio::test] +async fn python_sees_header_values_as_h11_would_parse_them() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nx-empty:\r\nx-blank: \r\nx-inner: a \t b \t\r\nX-Token!#$%&'*+.^_`|~: v\r\nx-empty: second\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let values = |name: &str| -> Vec<&[u8]> { + request + .headers + .get_all(name) + .iter() + .map(HeaderValue::as_bytes) + .collect() + }; + assert_eq!(values("x-empty"), [&b""[..], b"second"]); + assert_eq!(values("x-blank"), [&b""[..]]); + assert_eq!(values("x-inner"), [&b"a \t b"[..]]); + assert_eq!(values("x-token!#$%&'*+.^_`|~"), [&b"v"[..]]); +} + +/// RFC 9112 section 5.2 lets a server either reject obs-fold or unfold it, and h11 unfolds. +#[rstest] +#[ignore = "divergence: hyper rejects obs-fold with a 400 where uvicorn unfolds it to one space"] +#[tokio::test] +async fn a_folded_header_reaches_python_unfolded() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nx-a: b\r\n c\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().headers["x-a"], "b c"); +} + +#[rstest] +#[tokio::test] +async fn keeps_repeated_response_headers_and_opaque_header_bytes() { + let (upstream, mut received) = recording(|| { + Response::builder() + .header("set-cookie", "a=1; Path=/") + .header("set-cookie", "b=2; HttpOnly") + .header("x-opaque", HeaderValue::from_bytes(b"caf\xe9").unwrap()) + .body(Body::empty()) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(b"GET / HTTP/1.1\r\nhost: x\r\nx-opaque: caf\xe9\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + let cookies: Vec<_> = head + .headers + .iter() + .filter(|(name, _)| name == "set-cookie") + .map(|(_, value)| value.as_slice()) + .collect(); + assert_eq!(cookies, [&b"a=1; Path=/"[..], b"b=2; HttpOnly"]); + assert_eq!(head.header("x-opaque"), Some(&b"caf\xe9"[..])); + assert_eq!( + received.recv().await.unwrap().headers["x-opaque"].as_bytes(), + b"caf\xe9" + ); +} + +/// An HTTP/1.0 client cannot decode chunked framing, so a streamed answer reaches it delimited by +/// the connection closing, as uvicorn sends it. +#[rstest] +#[tokio::test] +async fn an_http10_client_gets_a_streamed_answer_unframed_and_python_gets_its_body() { + let (upstream, mut received) = recording(|| { + Response::new(Body::from_stream(futures_util::stream::iter( + ["str", "eamed"].map(|piece| Ok::<_, Infallible>(Bytes::from_static(piece.as_bytes()))), + ))) + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/completions HTTP/1.0\r\nhost: x\r\ncontent-length: 7\r\n\r\npayload") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.header("transfer-encoding"), None); + assert_eq!(within(raw.until_closed()).await, b"streamed"); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "POST"); + assert_eq!(request.body, "payload"); +} + +#[rstest] +#[tokio::test] +async fn answers_pipelined_requests_in_order() { + let upstream = serve( + Router::new().fallback(|request: Request| async move { request.uri().path().to_owned() }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /first HTTP/1.1\r\nhost: x\r\n\r\nHEAD /second HTTP/1.1\r\nhost: x\r\n\r\nGET /third HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let first = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(first.content_length().unwrap()).await, b"/first"); + let second = within(raw.response_head()).await.unwrap(); + assert_eq!(second.content_length(), Some(7)); + let third = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(third.content_length().unwrap()).await, b"/third"); +} + +#[rstest] +#[case::no_content(StatusCode::NO_CONTENT)] +#[case::not_modified(StatusCode::NOT_MODIFIED)] +#[tokio::test] +async fn bodiless_statuses_keep_the_connection_reusable(#[case] status: StatusCode) { + let upstream = serve(Router::new().fallback(move |request: Request| async move { + match request.uri().path() { + "/next" => (StatusCode::OK, "next"), + _ => (status, ""), + } + })) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /empty HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + assert_eq!(within(raw.response_head()).await.unwrap().status, status); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +#[rstest] +#[tokio::test] +async fn informational_responses_keep_the_connection_reusable() { + let upstream = raw_upstream(|mut python| async move { + while let Some(head) = python.request_head().await { + match head.path.as_str() { + "/hints" => { + python + .send("HTTP/1.1 103 Early Hints\r\nlink: ; rel=preload\r\n\r\nHTTP/1.1 200 OK\r\ncontent-length: 5\r\n\r\nfinal") + .await + } + _ => { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await + } + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /hints HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let mut hinted = within(raw.response_head()).await.unwrap(); + while hinted.status.is_informational() { + hinted = within(raw.response_head()).await.unwrap(); + } + assert_eq!(hinted.status, StatusCode::OK); + assert_eq!(raw.body(hinted.content_length().unwrap()).await, b"final"); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +#[rstest] +#[case::found( + StatusCode::FOUND, + "https://idp.example/authorize?state=a%2Fb&redirect_uri=https%3A%2F%2Fpublic.example%2Fsso%2Fcallback" +)] +#[case::see_other(StatusCode::SEE_OTHER, "/ui/?login=success")] +#[case::temporary(StatusCode::TEMPORARY_REDIRECT, "./mcp/oauth/authorize?client_id=x")] +#[tokio::test] +async fn redirects_reach_the_client_with_location_and_cookies_unchanged( + #[case] status: StatusCode, + #[case] location: &'static str, +) { + let (upstream, _received) = recording(move || { + Response::builder() + .status(status) + .header("location", location) + .header( + "set-cookie", + "token=a%2Fb; Path=/; HttpOnly; Secure; SameSite=Lax", + ) + .header("set-cookie", "state=; Max-Age=0") + .header( + "set-cookie", + "session=x; Expires=Wed, 04 Nov 2026 10:00:00 GMT; Path=/", + ) + .body(Body::empty()) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /sso/key/generate HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, status); + assert_eq!(head.header("location"), Some(location.as_bytes())); + let cookies: Vec<_> = head + .headers + .iter() + .filter(|(name, _)| name == "set-cookie") + .map(|(_, value)| value.as_slice()) + .collect(); + assert_eq!( + cookies, + [ + &b"token=a%2Fb; Path=/; HttpOnly; Secure; SameSite=Lax"[..], + b"state=; Max-Age=0", + b"session=x; Expires=Wed, 04 Nov 2026 10:00:00 GMT; Path=/", + ] + ); +} + +#[rstest] +#[tokio::test] +async fn non_websocket_upgrades_reach_python_as_plain_requests() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nconnection: upgrade, http2-settings\r\nupgrade: h2c\r\nhttp2-settings: AAMAAABkAAQAoAAAAAIAAAAA\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert!(!request.headers.contains_key("upgrade")); + assert!(!request.headers.contains_key("connection")); + assert!(!request.headers.contains_key("http2-settings")); +} + +#[rstest] +#[tokio::test] +async fn forwards_an_asterisk_form_options_request() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("OPTIONS * HTTP/1.1\r\nhost: x\r\n\r\n").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "OPTIONS"); + assert_eq!(request.uri, "*"); +} + +#[rstest] +#[tokio::test] +async fn an_authority_form_target_is_a_400_python_never_sees() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("CONNECT api.example:443 HTTP/1.1\r\nhost: api.example:443\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + assert!(received.try_recv().is_err()); +} + +#[rstest] +#[tokio::test] +async fn a_client_connection_close_ends_only_the_client_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + while python.request_head().await.is_some() { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..2 { + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nconnection: close\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("connection"), Some(&b"close"[..])); + assert_eq!(raw.body(2).await, b"ok"); + within(raw.closed()).await; + } + assert_eq!(connections.load(Ordering::SeqCst), 1); +} + +#[rstest] +#[tokio::test] +async fn a_python_connection_close_ends_only_the_upstream_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 500 Internal Server Error\r\ncontent-length: 5\r\nconnection: close\r\n\r\nerror") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + for _ in 0..2 { + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(head.header("connection"), None); + assert_eq!(raw.body(5).await, b"error"); + } + assert_eq!(connections.load(Ordering::SeqCst), 2); +} + +/// uvicorn answers HEAD on a streaming route with chunked framing and no body, so the client +/// must not be told the body is empty, and must not wait for one before its next request. +#[rstest] +#[tokio::test] +async fn head_answered_with_chunked_framing_claims_no_length_and_keeps_the_connection() { + let upstream = raw_upstream(|mut python| async move { + while let Some(head) = python.request_head().await { + match head.path.as_str() { + "/next" => { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await + } + _ => { + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n") + .await + } + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("HEAD /mcp HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.content_length(), None); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +const GZIP: &[u8] = &[ + 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x4b, 0x49, 0x2c, 0x49, 0xb4, 0x52, + 0x30, 0xe4, 0xe2, 0x02, 0x00, 0xce, 0xcd, 0x78, 0x54, 0x09, 0x00, 0x00, 0x00, +]; + +#[rstest] +#[case::fixed_length(false)] +#[case::streamed(true)] +#[tokio::test] +async fn content_codings_cross_unchanged_both_ways(#[case] streamed: bool) { + let (upstream, mut received) = recording(move || { + let body = match streamed { + true => Body::from_stream(futures_util::stream::iter( + GZIP.chunks(5) + .map(|piece| Ok::<_, Infallible>(Bytes::from_static(piece))), + )), + false => Body::from(GZIP), + }; + Response::builder() + .header("content-encoding", "gzip") + .header("vary", "Accept-Encoding") + .body(body) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + [ + format!( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-encoding: gzip\r\naccept-encoding: gzip, br;q=0.5\r\ncontent-length: {}\r\n\r\n", + GZIP.len() + ) + .as_bytes(), + GZIP, + ] + .concat(), + ) + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("content-encoding"), Some(&b"gzip"[..])); + assert_eq!(head.header("vary"), Some(&b"Accept-Encoding"[..])); + let body = match head.content_length() { + Some(length) => raw.body(length).await, + None => within(raw.chunked_body()).await, + }; + assert_eq!(body, GZIP); + assert_eq!(head.content_length().is_none(), streamed); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["content-encoding"], "gzip"); + assert_eq!(request.headers["accept-encoding"], "gzip, br;q=0.5"); + assert_eq!(request.body, GZIP); +} + +#[rstest] +#[tokio::test] +async fn python_sees_no_accept_encoding_the_client_did_not_send() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /metrics/ HTTP/1.1\r\nhost: x\r\n\r\n").await; + within(raw.response_head()).await.unwrap(); + assert!( + !received + .recv() + .await + .unwrap() + .headers + .contains_key("accept-encoding") + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/framing.rs b/litellm-rust/crates/gateway-forward/tests/framing.rs new file mode 100644 index 00000000000..32a2c3f2b8d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/framing.rs @@ -0,0 +1,259 @@ +mod support; + +use std::{collections::BTreeSet, net::SocketAddr}; + +use axum::http::StatusCode; +use rstest::rstest; +use support::{Gateway, Head, Raw, raw_upstream, within}; +use tokio::sync::mpsc; + +struct Captured { + head: Head, + body: Vec, +} + +/// A keep-alive upstream that decodes framing like h11 and reports every request it reads. +async fn capturing() -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = raw_upstream(move |mut python: Raw| { + let sender = sender.clone(); + async move { + while let Some(head) = python.request_head().await { + let body = match head.header("transfer-encoding") { + Some(_) => python.chunked_body().await, + None => python.body(content_length(&head)).await, + }; + sender.send(Captured { head, body }).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + (address, receiver) +} + +/// The length h11 reads from every Content-Length field, which may repeat one value. +fn content_length(head: &Head) -> usize { + let lengths: BTreeSet = head + .headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .flat_map(|(_, value)| value.split(|byte| *byte == b',')) + .map(|length| { + std::str::from_utf8(length.trim_ascii()) + .unwrap() + .parse() + .unwrap() + }) + .collect(); + assert!( + lengths.len() < 2, + "python got conflicting lengths {lengths:?}" + ); + lengths.first().copied().unwrap_or(0) +} + +/// RFC 9112 section 6.3: Transfer-Encoding overrides Content-Length. Forwarding both would let +/// Python split the stream at a different place than Rust did. Section 6.1 then closes the +/// connection after the response, so a request smuggled behind it never runs. +#[rstest] +#[tokio::test] +async fn a_request_framed_two_ways_reaches_python_framed_by_transfer_encoding_and_closes() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 3\r\ntransfer-encoding: chunked\r\n\r\n5\r\nhello\r\n0\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + within(raw.until_closed()).await; + + let upload = received.recv().await.unwrap(); + assert_eq!(upload.head.header("content-length"), None); + assert_eq!( + upload.head.header("transfer-encoding"), + Some(&b"chunked"[..]) + ); + assert_eq!(upload.body, b"hello"); + assert!( + received.try_recv().is_err(), + "python ran the smuggled request" + ); +} + +#[rstest] +#[case::conflicting_content_lengths("content-length: 4\r\ncontent-length: 40\r\n")] +#[case::content_length_list("content-length: 4, 40\r\n")] +#[case::signed_content_length("content-length: +4\r\n")] +#[case::chunked_not_last("transfer-encoding: chunked, identity\r\ncontent-length: 4\r\n")] +#[case::unknown_coding("transfer-encoding: xchunked\r\ncontent-length: 4\r\n")] +#[case::space_before_colon("transfer-encoding : chunked\r\ncontent-length: 4\r\n")] +#[case::folded_transfer_encoding("transfer-encoding:\r\n chunked\r\ncontent-length: 4\r\n")] +#[case::bare_cr("x-a: b\rx-b: c\r\ncontent-length: 4\r\n")] +#[case::nul_byte("x-a: b\x00c\r\ncontent-length: 4\r\n")] +#[tokio::test] +async fn ambiguous_framing_never_reaches_python(#[case] headers: &str) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\nbodyGET /smuggled HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + within(raw.closed()).await; + assert!(received.try_recv().is_err(), "python received a request"); +} + +#[rstest] +#[case::bare_lf("4\nbody\n0\n\n")] +#[case::invalid_size("zz\r\nbody\r\n0\r\n\r\n")] +#[case::size_overflow("10000000000000004\r\nbody\r\n0\r\n\r\n")] +#[tokio::test] +async fn a_malformed_request_body_is_a_400_not_a_502(#[case] chunks: &str) { + let upstream = raw_upstream(|mut python| async move { python.closed().await }).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{chunks}" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); +} + +/// hyper decodes only the chunked layer, so any other coding would reach Python still coded but +/// framed as a plain body. uvicorn turns every h11 parse error into a 400, ignoring the 501 h11 +/// suggests for codings. +#[rstest] +#[case::listed_before_chunked("transfer-encoding: gzip, chunked\r\n")] +#[case::split_across_fields("transfer-encoding: gzip\r\ntransfer-encoding: chunked\r\n")] +#[tokio::test] +async fn a_transfer_coding_other_than_chunked_is_the_400_uvicorn_sends_python_never_sees( + #[case] headers: &str, +) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\n4\r\nbody\r\n0\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + assert!( + received.try_recv().is_err(), + "python received a still-coded body" + ); +} + +async fn assert_python_reads_the_body_and_the_next_request(headers: &str) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\npayloadGET /next HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + for _ in 0..2 { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + } + assert_eq!(received.recv().await.unwrap().body, b"payload"); + assert_eq!(received.recv().await.unwrap().head.path, "/next"); +} + +#[rstest] +#[case::repeated_fields("content-length: 7\r\ncontent-length: 7\r\n")] +#[case::leading_zeros("content-length: 007\r\n")] +#[tokio::test] +async fn a_content_length_h11_reads_one_way_reaches_python_framed_the_same(#[case] headers: &str) { + assert_python_reads_the_body_and_the_next_request(headers).await; +} + +/// h11 accepts a Content-Length list that repeats one value, as RFC 9110 section 8.6 allows. +#[rstest] +#[ignore = "divergence: hyper's server parses each Content-Length field as one number and 400s a list"] +#[tokio::test] +async fn a_content_length_list_of_one_value_reaches_python() { + assert_python_reads_the_body_and_the_next_request("content-length: 7, 7\r\n").await; +} + +#[rstest] +#[case::leading_zeros("0004\r\nbody\r\n0\r\n\r\n", "body")] +#[case::uppercase_hex("A\r\nbody\r\nbody\r\n0\r\n\r\n", "body\r\nbody")] +#[case::whitespace_after_size("4 \r\nbody\r\n0\r\n\r\n", "body")] +#[case::extension_and_trailer("4;a=\"b;c\"\r\nbody\r\n0\r\nx-sum: 1\r\n\r\n", "body")] +#[tokio::test] +async fn chunk_framing_h11_accepts_reaches_python_decoded( + #[case] chunks: &str, + #[case] body: &str, +) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{chunks}GET /next HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + for _ in 0..2 { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + } + assert_eq!(received.recv().await.unwrap().body, body.as_bytes()); + assert_eq!(received.recv().await.unwrap().head.path, "/next"); +} + +/// h11 reads a head whose lines end in a bare LF, so Python gets the same request through Rust. +#[rstest] +#[tokio::test] +async fn a_head_with_bare_lf_line_endings_reaches_python() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.1\nhost: x\nx-a: b\ncontent-length: 4\n\nbody") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let upload = received.recv().await.unwrap(); + assert_eq!(upload.head.header("x-a"), Some(&b"b"[..])); + assert_eq!(upload.body, b"body"); +} + +/// h11 decodes a chunked body from an HTTP/1.0 client, and RFC 9112 section 6.1 only asks the +/// server to close the connection afterwards. +#[rstest] +#[ignore = "divergence: hyper rejects Transfer-Encoding on HTTP/1.0 with a 400 that uvicorn never sends"] +#[tokio::test] +async fn an_http10_chunked_upload_reaches_python() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.0\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n4\r\nbody\r\n0\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, b"body"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/idle_stream.rs b/litellm-rust/crates/gateway-forward/tests/idle_stream.rs new file mode 100644 index 00000000000..a07c0466481 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/idle_stream.rs @@ -0,0 +1,102 @@ +mod support; + +use std::time::Duration; + +use axum::{ + body::Body, + extract::Request, + http::StatusCode, + response::{IntoResponse, Response}, +}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, recording, within}; +use tokio::{sync::mpsc, time::Instant}; + +const SSE_HEAD: &str = "HTTP/1.1 200 OK\r\n\ + date: Tue, 03 Nov 2026 00:00:00 GMT\r\n\ + server: uvicorn\r\n\ + cache-control: no-cache, no-transform\r\n\ + connection: keep-alive\r\n\ + content-type: text/event-stream\r\n\ + mcp-session-id: session-1\r\n\ + x-accel-buffering: no\r\n\ + transfer-encoding: chunked\r\n\r\n"; + +const PING: &str = ": ping - 2026-11-03 00:20:00.000000+00:00\r\n\r\n"; + +const MESSAGE: &str = "event: message\r\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\r\n\r\n"; + +fn chunk(data: &str) -> String { + format!("{:x}\r\n{data}\r\n", data.len()) +} + +/// An MCP client holds a GET stream open for server-initiated messages, which may not come for +/// many minutes, so neither hop may close a stream for being idle. +#[rstest] +#[tokio::test(start_paused = true)] +async fn an_idle_mcp_stream_stays_open_in_both_directions() { + let idle = Duration::from_secs(20 * 60); + let (session, mut sessions) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let session = session.clone(); + async move { + let head = python.request_head().await.unwrap(); + session + .send(head.header("mcp-session-id").map(<[u8]>::to_vec)) + .unwrap(); + python.send(SSE_HEAD).await; + tokio::time::sleep(idle).await; + python.send(chunk(PING)).await; + python.send(chunk(MESSAGE)).await; + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let started = Instant::now(); + + let response = client() + .request( + Request::get(gateway.url("/github/mcp")) + .header("accept", "text/event-stream") + .header("mcp-session-id", "session-1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.headers()["mcp-session-id"], "session-1"); + assert_eq!(response.headers()["x-accel-buffering"], "no"); + assert!(!response.headers().contains_key("connection")); + assert_eq!( + sessions.recv().await.unwrap().as_deref(), + Some(&b"session-1"[..]) + ); + let body = response.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, [PING, MESSAGE].concat()); + assert!(started.elapsed() >= idle); +} + +#[rstest] +#[tokio::test] +async fn ends_a_session_with_a_bodiless_delete() { + let (upstream, mut received) = recording(|| -> Response { + ([("mcp-session-id", "session-1")], StatusCode::NO_CONTENT).into_response() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("DELETE /mcp HTTP/1.1\r\nhost: x\r\nmcp-session-id: session-1\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::NO_CONTENT); + assert_eq!(head.header("mcp-session-id"), Some(&b"session-1"[..])); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "DELETE"); + assert_eq!(request.headers["mcp-session-id"], "session-1"); + assert!(!request.headers.contains_key("transfer-encoding")); + assert!(request.body.is_empty()); +} diff --git a/litellm-rust/crates/gateway-forward/tests/lifecycle.rs b/litellm-rust/crates/gateway-forward/tests/lifecycle.rs new file mode 100644 index 00000000000..97907738ad6 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/lifecycle.rs @@ -0,0 +1,524 @@ +mod support; + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use axum::{Router, body::Body, extract::Request, http::StatusCode, response::Response}; +use http_body_util::BodyExt; +use litellm_gateway_forward::input_closed; +use rstest::rstest; +use support::{ + DRAIN, Gateway, Raw, UVICORN_KEEPALIVE, client, raw_upstream, recording, serve, serve_until, + unreachable_upstream, within, +}; +use tokio::{ + net::{TcpListener, TcpStream}, + sync::{Notify, mpsc}, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +/// Python stops a provider stream when it sees the client go away, so a client disconnect must +/// close Rust's connection to Python at every stage of the exchange. +#[rstest] +#[case::waiting_for_headers( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n", + false +)] +#[case::mid_stream( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n", + true +)] +#[case::mid_upload( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 1000000\r\n\r\npartial", + false +)] +#[case::waiting_for_websocket_handshake( + "GET /v1/realtime HTTP/1.1\r\nhost: x\r\nconnection: Upgrade\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: dGhlIHNhbXBsZSBub25jZQ==\r\n\r\n", + false +)] +#[tokio::test] +async fn client_disconnect_closes_the_upstream_connection( + #[case] request: &'static str, + #[case] headers_sent: bool, +) { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + python.request_head().await.unwrap(); + if headers_sent { + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n6\r\ndata: \r\n") + .await; + } + event.send("received").unwrap(); + python.closed().await; + event.send("closed").unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(request).await; + assert_eq!(within(events.recv()).await, Some("received")); + if headers_sent { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + } + drop(raw); + assert_eq!(within(events.recv()).await, Some("closed")); +} + +/// Python treats EOF on its socket as the client leaving, so Rust must not half-close the +/// upstream connection once it has sent the request body. +#[rstest] +#[tokio::test] +async fn python_sees_no_eof_while_its_response_is_pending() { + let (verdict, mut verdicts) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let verdict = verdict.clone(); + async move { + let head = python.request_head().await.unwrap(); + python.body(head.content_length().unwrap()).await; + let eof = tokio::time::timeout(Duration::from_millis(200), python.closed()) + .await + .is_ok(); + verdict.send(eof).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + + let response = within( + client().request( + Request::post(gateway.url("/mcp")) + .body(Body::from("{}")) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(verdicts.recv().await, Some(false)); +} + +/// uvicorn closes a connection that sat idle past its keep-alive timeout. A request Rust sends on +/// that connection as uvicorn closes it is lost, so Rust must retire idle connections first. +#[rstest] +#[case::reused_while_fresh(Duration::from_millis(0))] +#[case::past_the_pool_timeout(Duration::from_millis(250))] +#[case::past_python_keepalive(Duration::from_millis(450))] +#[tokio::test] +async fn never_reuses_a_connection_python_is_closing_for_idleness(#[case] idle: Duration) { + let keepalive = Duration::from_millis(400); + let upstream = raw_upstream(move |mut python: Raw| async move { + let mut idle_since = None::; + while python.request_head().await.is_some() { + if idle_since.is_some_and(|since| since.elapsed() >= keepalive) { + return; + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + idle_since = Some(Instant::now()); + } + }) + .await; + let gateway = Gateway::with_keepalive(upstream, keepalive).await; + let client = client(); + let first = within(client.request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + first.into_body().collect().await.unwrap(); + tokio::time::sleep(idle).await; + let second = + within(client.request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(second.status(), StatusCode::OK); +} + +#[rstest] +#[case::health("/health/liveliness")] +#[case::messages("/v1/messages")] +#[tokio::test] +async fn unreachable_python_is_a_502(#[case] path: &str) { + let gateway = Gateway::start(unreachable_upstream()).await; + let response = + within(client().request(Request::get(gateway.url(path)).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); +} + +#[rstest] +#[case::bad_status_line(b"HTTP/1.1 OK\r\ncontent-length: 0\r\n\r\n".to_vec())] +#[case::header_without_colon(b"HTTP/1.1 200 OK\r\nbroken header\r\ncontent-length: 0\r\n\r\n".to_vec())] +#[case::not_http(b"SSH-2.0-OpenSSH_9.6\r\n".to_vec())] +#[case::conflicting_lengths( + b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\ncontent-length: 3\r\n\r\nok".to_vec() +)] +#[case::oversized_head( + [&b"HTTP/1.1 200 OK\r\nx-big: "[..], &vec![b'a'; 1024 * 1024], b"\r\n\r\n"].concat() +)] +#[case::closed_without_reply(Vec::new())] +#[tokio::test] +async fn an_upstream_that_breaks_http_is_a_502_and_the_client_connection_survives( + #[case] reply: Vec, +) { + let upstream = raw_upstream(move |mut python: Raw| { + let reply = reply.clone(); + async move { + python.request_head().await.unwrap(); + if reply.is_empty() { + return; + } + python.send(reply).await; + std::future::pending::<()>().await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::BAD_GATEWAY); + assert_eq!(head.content_length(), Some(0)); + } +} + +#[rstest] +#[tokio::test] +async fn python_restarting_is_a_502_until_it_listens_again() { + let address = unreachable_upstream(); + let gateway = Gateway::start(address).await; + let get = || client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()); + assert_eq!( + within(get()).await.unwrap().status(), + StatusCode::BAD_GATEWAY + ); + + let listener = TcpListener::bind(address).await.unwrap(); + let router = Router::new().fallback(|| async { "up" }); + tokio::spawn(async move { axum::serve(listener, router).await }); + + assert_eq!(within(get()).await.unwrap().status(), StatusCode::OK); +} + +/// A request Python already read may have run, so it is answered with a 502 and never replayed. +#[rstest] +#[tokio::test] +async fn a_request_python_read_before_closing_is_a_502_and_never_replayed() { + let (event, mut seen) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + python.body(head.content_length().unwrap()).await; + event.send(head.path.clone()).unwrap(); + if head.path == "/v1/chat/completions" { + return; + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /warm HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + let warm = within(raw.response_head()).await.unwrap(); + raw.body(warm.content_length().unwrap()).await; + + raw.send("POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 2\r\n\r\n{}") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::BAD_GATEWAY); + let paths: Vec<_> = std::iter::from_fn(|| seen.try_recv().ok()).collect(); + assert_eq!(paths, ["/warm", "/v1/chat/completions"]); +} + +#[rstest] +#[tokio::test] +async fn shutdown_lets_an_in_flight_stream_finish() { + let release = Arc::new(Notify::new()); + let released = release.clone(); + let upstream = raw_upstream(move |mut python: Raw| { + let released = released.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n") + .await; + released.notified().await; + python.send("4\r\nlast\r\n0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(!server.is_finished()); + release.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, "firstlast"); + within(server).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn the_supervisor_closing_stdin_drains_like_a_stop_signal() { + let release = Arc::new(Notify::new()); + let released = release.clone(); + let upstream = raw_upstream(move |mut python: Raw| { + let released = released.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n") + .await; + released.notified().await; + python.send("4\r\nlast\r\n0\r\n\r\n").await; + } + }) + .await; + let (stdin, supervisor) = std::io::pipe().unwrap(); + let (address, server) = serve_until(upstream, UVICORN_KEEPALIVE, input_closed(stdin)).await; + let response = within( + client().request( + Request::get(format!("http://{address}/")) + .body(Body::empty()) + .unwrap(), + ), + ) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(TcpStream::connect(address).await.is_ok()); + + drop(supervisor); + + within(async { + while TcpStream::connect(address).await.is_ok() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await; + assert!(!server.is_finished()); + release.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, "firstlast"); + within(server).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn shutdown_stops_waiting_for_a_stream_still_open_after_the_drain() { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n") + .await; + std::future::pending::<()>().await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + let started = std::time::Instant::now(); + + within(gateway.shutdown()).await.unwrap(); + + assert!(started.elapsed() >= DRAIN); + drop(response); +} + +#[rstest] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_keep_alive_connections_each_get_their_own_responses() { + let upstream = serve(Router::new().fallback(|request: Request| async move { + let index: usize = request.uri().query().unwrap().parse().unwrap(); + Response::builder() + .header("etag", format!("\"{index}\"")) + .body(Body::from(index.to_string().repeat(1 + index * 1000))) + .unwrap() + })) + .await; + let gateway = Gateway::start(upstream).await; + let connections = 6; + + let tasks: Vec<_> = (0..connections) + .map(|first| { + let address = gateway.address; + tokio::spawn(async move { + let mut raw = Raw::new(TcpStream::connect(address).await.unwrap()); + for index in (first..50).step_by(connections) { + raw.send(format!( + "GET /ui/chunk.js?{index} HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("etag"), Some(format!("\"{index}\"").as_bytes())); + let body = raw.body(head.content_length().unwrap()).await; + assert_eq!( + body, + index.to_string().repeat(1 + index * 1000).into_bytes() + ); + } + }) + }) + .collect(); + + for task in tasks { + within(task).await.unwrap(); + } +} + +// Python's `--keepalive_timeout` governed idle client connections before Rust took the public +// port, and operators tune it against their load balancer's idle timeout, so it still does. +// Without it, a client that vanishes without closing its connection holds a socket forever and +// the gateway eventually runs out of file descriptors. +const IDLE: Duration = Duration::from_millis(400); + +#[rstest] +#[tokio::test] +async fn closes_a_client_connection_idle_past_python_keepalive() { + let (upstream, _received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + tokio::time::sleep(IDLE / 2).await; + raw.send("GET /health/liveliness HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + raw.body(head.content_length().unwrap()).await; + } + let idle_since = Instant::now(); + within(raw.closed()).await; + + assert!(idle_since.elapsed() >= IDLE * 3 / 4); +} + +/// uvicorn would hold a connection that never sends a request, which lets a client park sockets +/// for free, so the same limit applies before the first request too. +#[rstest] +#[tokio::test] +async fn closes_a_client_connection_that_never_sends_a_request() { + let (upstream, _received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + within(raw.closed()).await; +} + +/// The limit only covers the wait for the next request. A response that goes quiet for longer +/// belongs to Python, which decides how long it may take. +#[rstest] +#[tokio::test] +async fn keeps_a_connection_open_while_a_response_goes_quiet() { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + tokio::time::sleep(IDLE * 3).await; + python.send("4\r\ndone\r\n0\r\n\r\n").await; + }) + .await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + raw.send("GET /v1/chat/completions HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + within(raw.response_head()).await.unwrap(); + assert_eq!(within(raw.chunked_body()).await, b"done"); +} + +/// Likewise for an upload that pauses longer than the limit before it finishes. +#[rstest] +#[tokio::test] +async fn keeps_a_connection_open_while_an_upload_goes_quiet() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + raw.send("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 8\r\n\r\nhalf") + .await; + tokio::time::sleep(IDLE * 3).await; + raw.send("done").await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "halfdone"); +} + +/// uvicorn answers a request it rejects (413, 401) as soon as it reads the head and then closes, +/// leaving the upload unread. The client must get that answer, not a 502 or a reset. +#[rstest] +#[case::sized("content-length: 16777216\r\n", vec![b'x'; 256 * 1024])] +#[case::chunked( + "transfer-encoding: chunked\r\n", + [&b"40000\r\n"[..], &[b'x'; 256 * 1024], b"\r\n"].concat() +)] +#[tokio::test] +async fn python_rejecting_an_upload_unread_and_closing_still_answers_the_client( + #[case] framing: &'static str, + #[case] partial_upload: Vec, +) { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 413 Payload Too Large\r\ncontent-type: application/json\r\ncontent-length: 26\r\nconnection: close\r\n\r\n{\"error\":\"file too large\"}") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{framing}\r\n" + )) + .await; + raw.send(partial_upload).await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!( + raw.body(head.content_length().unwrap()).await, + b"{\"error\":\"file too large\"}" + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/roundtrip.rs b/litellm-rust/crates/gateway-forward/tests/roundtrip.rs new file mode 100644 index 00000000000..c220c7fe8e9 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/roundtrip.rs @@ -0,0 +1,360 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{Arc, Mutex}, +}; + +use axum::{ + body::{Body, Bytes}, + http::{HeaderMap, HeaderName, HeaderValue, Method, Request, StatusCode}, + response::Response, +}; +use http_body_util::BodyExt; +use litellm_gateway_forward::PEER_HEADER; +use proptest::{prelude::*, test_runner::TestRunner}; +use support::{Gateway, client, recording, within}; + +#[derive(Clone, Debug)] +struct Message { + headers: Vec<(String, Vec>)>, + body: Vec, + chunked: bool, +} + +#[derive(Clone, Debug)] +struct Exchange { + method: Method, + target: String, + request: Message, + status: StatusCode, + response: Message, +} + +/// Any value h11 reads back unchanged: visible ASCII, obs-text and interior whitespace. +fn header_value() -> impl Strategy> { + prop::collection::vec( + prop_oneof![8 => 0x21u8..=0x7e, 2 => 0x80u8..=0xff, 1 => Just(b' '), 1 => Just(b'\t')], + 0..24, + ) + .prop_map(|value| value.trim_ascii().to_vec()) +} + +fn method() -> impl Strategy { + prop_oneof![ + Just(Method::GET), + Just(Method::POST), + Just(Method::PUT), + Just(Method::PATCH), + Just(Method::DELETE), + Just(Method::OPTIONS), + "[A-Za-z0-9!#$%&'*+.^_`|~-]{1,16}" + .prop_filter("HEAD and CONNECT frame differently", |method| { + !["HEAD", "CONNECT"].contains(&method.as_str()) + }) + .prop_map(|method| Method::from_bytes(method.as_bytes()).unwrap()), + ] +} + +fn message() -> impl Strategy { + ( + prop::collection::btree_map( + "x-[a-z0-9!#$%&'*+.^_`|~-]{1,12}", + prop::collection::vec(header_value(), 1..4), + 0..12, + ), + prop::collection::vec(any::(), 0..32 * 1024), + any::(), + ) + .prop_map(|(headers, body, chunked)| Message { + headers: headers.into_iter().collect(), + body, + chunked, + }) +} + +fn target() -> impl Strategy { + let segment = prop_oneof![ + "[A-Za-z0-9._~!$&'()*+,;=:@-]{0,8}", + "%[0-9A-Fa-f]{2}", + Just(".".to_owned()), + Just("..".to_owned()), + Just(String::new()), + ]; + ( + prop::collection::vec(segment, 0..6), + prop::option::of("[A-Za-z0-9=&._~%/?:@!$'()*+,;-]{0,40}"), + ) + .prop_map(|(segments, query)| { + let path = format!("/{}", segments.join("/")); + match query { + Some(query) => format!("{path}?{query}"), + None => path, + } + }) +} + +fn exchange() -> impl Strategy { + let status = (200u16..600) + .prop_filter("bodiless status", |code| ![204, 205, 304].contains(code)) + .prop_map(|code| StatusCode::from_u16(code).unwrap()); + (method(), target(), message(), status, message()).prop_map( + |(method, target, request, status, response)| Exchange { + method, + target, + request, + status, + response, + }, + ) +} + +fn body(message: &Message) -> Body { + if !message.chunked { + return Body::from(message.body.clone()); + } + let chunks: Vec> = message + .body + .chunks(1000) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect(); + Body::from_stream(futures_util::stream::iter(chunks)) +} + +fn apply(headers: &mut HeaderMap, message: &Message) { + if message.chunked { + headers.insert("transfer-encoding", HeaderValue::from_static("chunked")); + } + for (name, values) in &message.headers { + for value in values { + headers.append( + HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_bytes(value).unwrap(), + ); + } + } +} + +fn assert_carried(headers: &HeaderMap, message: &Message, framing_and_hop: &[&str]) { + for (name, values) in &message.headers { + let carried: Vec<&[u8]> = headers + .get_all(name) + .iter() + .map(HeaderValue::as_bytes) + .collect(); + let sent: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect(); + assert_eq!(carried, sent, "header {name}"); + } + for name in headers.keys() { + assert!( + framing_and_hop.contains(&name.as_str()) + || message + .headers + .iter() + .any(|(sent, _)| sent == name.as_str()), + "unexpected header {name}" + ); + } +} + +/// Any request a client can send and any response Python can give cross the forwarder unchanged, +/// apart from the framing and peer headers the forwarder owns. +#[test] +fn every_exchange_crosses_the_forwarder_unchanged() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let current = Arc::new(Mutex::new(None::)); + let responding = current.clone(); + let (gateway, received) = runtime.block_on(async { + let (upstream, received) = recording(move || { + let exchange = responding.lock().unwrap().clone().unwrap(); + let mut response = Response::new(body(&exchange.response)); + *response.status_mut() = exchange.status; + apply(response.headers_mut(), &exchange.response); + response + }) + .await; + ( + Gateway::start(upstream).await, + tokio::sync::Mutex::new(received), + ) + }); + let client = client(); + TestRunner::new(ProptestConfig::with_cases(256)) + .run(&exchange(), |exchange| { + *current.lock().unwrap() = Some(exchange.clone()); + runtime.block_on(async { + while received.lock().await.try_recv().is_ok() {} + let mut request = Request::builder() + .method(exchange.method.clone()) + .uri(gateway.url(&exchange.target)) + .body(body(&exchange.request)) + .unwrap(); + apply(request.headers_mut(), &exchange.request); + let response = support::within(client.request(request)).await.unwrap(); + + let python = received.lock().await.recv().await.unwrap(); + assert_eq!(python.method, exchange.method); + assert_eq!(python.uri, exchange.target.as_str()); + assert_eq!(python.body, exchange.request.body); + assert_carried( + &python.headers, + &exchange.request, + &["host", "content-length", "transfer-encoding", PEER_HEADER], + ); + assert_eq!(python.headers.get_all(PEER_HEADER).iter().count(), 1); + + assert_eq!(response.status(), exchange.status); + let headers = response.headers().clone(); + assert_carried( + &headers, + &exchange.response, + &["date", "content-length", "transfer-encoding"], + ); + let body = support::within(response.into_body().collect()) + .await + .unwrap(); + assert_eq!(body.to_bytes(), exchange.response.body); + }); + Ok(()) + }) + .unwrap(); +} + +#[derive(Clone, Debug)] +enum Framing { + Unframed, + Length, + Chunked(usize), +} + +#[derive(Clone, Debug)] +struct WireRequest { + method: Method, + target: String, + headers: Vec<(String, Vec)>, + body: Vec, + framing: Framing, +} + +fn wire_request() -> impl Strategy { + let framing = prop_oneof![ + Just(Framing::Unframed), + Just(Framing::Length), + (1usize..2048).prop_map(Framing::Chunked), + ]; + ( + method(), + target(), + prop::collection::vec( + ("[xX]-[A-Za-z0-9!#$%&'*+.^_`|~-]{1,12}", header_value()), + 0..24, + ), + prop::collection::vec(any::(), 0..8 * 1024), + framing, + ) + .prop_map(|(method, target, headers, body, framing)| WireRequest { + method, + target, + headers, + body: match framing { + Framing::Unframed => Vec::new(), + _ => body, + }, + framing, + }) +} + +fn encode(request: &WireRequest) -> Vec { + let framing = match request.framing { + Framing::Unframed => String::new(), + Framing::Length => format!("Content-Length: {}\r\n", request.body.len()), + Framing::Chunked(_) => "Transfer-Encoding: chunked\r\n".to_owned(), + }; + let fields = request + .headers + .iter() + .flat_map(|(name, value)| [name.as_bytes(), b": ", value, b"\r\n"].concat()); + let body = match request.framing { + Framing::Chunked(size) => request + .body + .chunks(size) + .flat_map(|chunk| { + [format!("{:x}\r\n", chunk.len()).as_bytes(), chunk, b"\r\n"].concat() + }) + .chain(*b"0\r\n\r\n") + .collect(), + _ => request.body.clone(), + }; + [ + format!( + "{} {} HTTP/1.1\r\nHost: x\r\n", + request.method, request.target + ) + .into_bytes(), + fields.collect(), + framing.into_bytes(), + b"\r\n".to_vec(), + body, + ] + .concat() +} + +fn by_name<'a>(fields: impl Iterator) -> Vec<(String, &'a [u8])> { + let mut fields: Vec<_> = fields + .map(|(name, value)| (name.to_ascii_lowercase(), value)) + .collect(); + fields.sort_by(|(left, _), (right, _)| left.cmp(right)); + fields +} + +/// Pipelined requests in any framing, written in arbitrary pieces, reach Python one by one with +/// the bytes h11 would have read from the client directly. +#[test] +fn pipelined_wire_requests_reach_python_as_sent() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let (gateway, received) = runtime.block_on(async { + let (upstream, received) = recording(|| Response::new(Body::from("ok"))).await; + ( + Gateway::start(upstream).await, + tokio::sync::Mutex::new(received), + ) + }); + let requests = (prop::collection::vec(wire_request(), 1..5), 1usize..4096); + TestRunner::new(ProptestConfig::with_cases(256)) + .run(&requests, |(requests, piece)| { + runtime.block_on(async { + let mut received = received.lock().await; + while received.try_recv().is_ok() {} + let mut raw = gateway.connect().await; + let wire: Vec = requests.iter().flat_map(encode).collect(); + for bytes in wire.chunks(piece) { + raw.send(bytes).await; + } + for request in &requests { + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK, "{request:?}"); + assert_eq!(raw.body(2).await, b"ok"); + let python = within(received.recv()).await.unwrap(); + assert_eq!(python.method, request.method); + assert_eq!(python.uri, request.target.as_str()); + assert_eq!(python.body, request.body); + assert_eq!(python.headers["host"], "x"); + let forwarded = python.headers.iter().filter(|(name, _)| { + !["host", "content-length", "transfer-encoding", PEER_HEADER] + .contains(&name.as_str()) + }); + assert_eq!( + by_name(forwarded.map(|(name, value)| (name.as_str(), value.as_bytes()))), + by_name( + request + .headers + .iter() + .map(|(name, value)| (name.as_str(), value.as_slice())) + ) + ); + } + }); + Ok(()) + }) + .unwrap(); +} diff --git a/litellm-rust/crates/gateway-forward/tests/settings.rs b/litellm-rust/crates/gateway-forward/tests/settings.rs new file mode 100644 index 00000000000..ce0d95e723f --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/settings.rs @@ -0,0 +1,110 @@ +use std::{ffi::OsString, time::Duration}; + +use litellm_gateway_forward::{Error, Settings}; +use rstest::rstest; + +const COMPLETE: [(&str, &str); 5] = [ + ("LITELLM_FORWARD_UPSTREAM", "http://127.0.0.1:53211"), + ("LITELLM_FORWARD_HOST", "0.0.0.0"), + ("LITELLM_FORWARD_PORT", "4000"), + ("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "2"), + ("LITELLM_FORWARD_SHUTDOWN_DRAIN", "40"), +]; + +fn parse(variables: &[(&str, &str)]) -> Result, Error> { + Settings::from_vars( + variables + .iter() + .map(|(key, value)| (OsString::from(key), OsString::from(value))), + ) +} + +fn without(variable: &str) -> Vec<(&'static str, &'static str)> { + COMPLETE + .into_iter() + .filter(|(key, _)| *key != variable) + .collect() +} + +fn with<'a>(variable: &'a str, value: &'a str) -> Vec<(&'a str, &'a str)> { + without(variable) + .into_iter() + .chain([(variable, value)]) + .collect() +} + +#[rstest] +fn reads_every_variable() { + assert_eq!( + parse(&COMPLETE).unwrap(), + Some(Settings { + upstream: "http://127.0.0.1:53211".parse().unwrap(), + host: "0.0.0.0".into(), + port: 4000, + upstream_keepalive: Duration::from_secs(2), + shutdown_drain: Duration::from_secs(40), + }) + ); +} + +#[rstest] +fn without_an_upstream_the_gateway_is_not_forwarding() { + assert_eq!(parse(&without("LITELLM_FORWARD_UPSTREAM")).unwrap(), None); +} + +#[rstest] +#[case::host("LITELLM_FORWARD_HOST")] +#[case::port("LITELLM_FORWARD_PORT")] +#[case::upstream_keepalive("LITELLM_FORWARD_UPSTREAM_KEEPALIVE")] +#[case::shutdown_drain("LITELLM_FORWARD_SHUTDOWN_DRAIN")] +fn a_missing_variable_fails_naming_it(#[case] variable: &str) { + let error = parse(&without(variable)).unwrap_err(); + + assert!(error.to_string().contains(variable), "{error}"); +} + +#[rstest] +#[case::port_overflow("LITELLM_FORWARD_PORT", "65536")] +#[case::port_text("LITELLM_FORWARD_PORT", "http")] +#[case::keepalive_text("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "soon")] +#[case::keepalive_negative("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "-1")] +#[case::keepalive_zero("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "0")] +#[case::drain_text("LITELLM_FORWARD_SHUTDOWN_DRAIN", "soon")] +#[case::drain_negative("LITELLM_FORWARD_SHUTDOWN_DRAIN", "-1")] +#[case::drain_infinite("LITELLM_FORWARD_SHUTDOWN_DRAIN", "inf")] +fn an_invalid_variable_fails_naming_it(#[case] variable: &str, #[case] value: &str) { + let error = parse(&with(variable, value)).unwrap_err(); + + assert!(error.to_string().contains(variable), "{error}"); +} + +#[rstest] +fn fractional_drain_seconds_are_used_as_given() { + let settings = parse(&with("LITELLM_FORWARD_SHUTDOWN_DRAIN", "12.5")) + .unwrap() + .unwrap(); + + assert_eq!(settings.shutdown_drain, Duration::from_millis(12_500)); +} + +#[rstest] +#[case::keepalive("KEEPALIVE_TIMEOUT")] +#[case::graceful_shutdown("GRACEFUL_SHUTDOWN_TIMEOUT")] +#[case::host("HOST")] +#[case::port("PORT")] +fn python_settings_are_never_read(#[case] variable: &str) { + let mut variables = COMPLETE.to_vec(); + variables.push((variable, "9")); + + assert_eq!(parse(&variables).unwrap(), parse(&COMPLETE).unwrap()); +} + +#[rstest] +fn a_non_loopback_upstream_fails_startup() { + let error = parse(&with("LITELLM_FORWARD_UPSTREAM", "http://192.0.2.1:4000")).unwrap_err(); + + assert!( + error.to_string().contains("LITELLM_FORWARD_UPSTREAM"), + "{error}" + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/slow_head.rs b/litellm-rust/crates/gateway-forward/tests/slow_head.rs new file mode 100644 index 00000000000..764733aa524 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/slow_head.rs @@ -0,0 +1,41 @@ +mod support; + +use std::time::Duration; + +use axum::{body::Body, extract::Request, http::StatusCode}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream}; +use tokio::time::Instant; + +/// Reasoning models with `stream=false`, image and video generation, and model health checks send +/// no response headers until the provider finishes, which can take many minutes. +#[rstest] +#[tokio::test(start_paused = true)] +async fn waits_for_response_headers_as_long_as_python_takes() { + let thinking = Duration::from_secs(30 * 60); + let upstream = raw_upstream(move |mut python| async move { + python.request_head().await.unwrap(); + tokio::time::sleep(thinking).await; + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 5\r\n\r\nimage") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let started = Instant::now(); + + let response = client() + .request( + Request::post(gateway.url("/v1/images/generations")) + .body(Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert!(started.elapsed() >= thinking); + let body = response.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, "image"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/stream.rs b/litellm-rust/crates/gateway-forward/tests/stream.rs new file mode 100644 index 00000000000..eaa12f20b77 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/stream.rs @@ -0,0 +1,313 @@ +mod support; + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{body::Body, extract::Request, http::StatusCode}; +use http_body_util::BodyExt; +use hyper::body::Incoming; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, within}; +use tokio::{ + io::AsyncWriteExt, + sync::{Notify, mpsc}, +}; + +fn chunk(bytes: &[u8]) -> Vec { + [format!("{:x}\r\n", bytes.len()).as_bytes(), bytes, b"\r\n"].concat() +} + +const EVENTS: [&str; 3] = ["data: 1\n\n", "data: 2\n\n", "data: [DONE]\n\n"]; + +#[rstest] +#[tokio::test] +async fn each_event_reaches_the_client_before_python_produces_the_next( + #[values(true, false)] chunked: bool, +) { + let produce = Arc::new(Notify::new()); + let next = produce.clone(); + let upstream = raw_upstream(move |mut python| { + let next = next.clone(); + async move { + python.request_head().await.unwrap(); + let framing = match chunked { + true => "transfer-encoding: chunked".to_owned(), + false => format!("content-length: {}", EVENTS.concat().len()), + }; + python + .send(format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n{framing}\r\n\r\n" + )) + .await; + for event in EVENTS { + match chunked { + true => python.send(chunk(event.as_bytes())).await, + false => python.send(event).await, + } + next.notified().await; + } + if chunked { + python.send("0\r\n\r\n").await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut response = within( + client().request( + Request::post(gateway.url("/v1/chat/completions")) + .body(Body::from(r#"{"stream":true}"#)) + .unwrap(), + ), + ) + .await + .unwrap(); + + for event in EVENTS { + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), event); + produce.notify_one(); + } + let rest = within(response.into_body().collect()).await.unwrap(); + assert!(rest.to_bytes().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn relays_a_binary_framed_stream_unchanged() { + let message: Vec = [ + &[ + 0x00, 0x00, 0x00, 0x3d, 0x00, 0x00, 0x00, 0x20, 0x8a, 0x0d, 0x0a, 0xff, + ][..], + b"\x0b:event-type\x07\x00\x05chunk\x0d:content-type\x07\x00\x10application/json", + b"0\r\n\r\n{\"bytes\":\"eyJ0eXBlIjoiY29udGVudCJ9\"}", + &[0xc3, 0x28, 0x00, 0x0d], + ] + .concat(); + let sent = message.clone(); + let upstream = raw_upstream(move |mut python| { + let sent = sent.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: application/vnd.amazon.eventstream\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + for piece in sent.chunks(7) { + python.send(chunk(piece)).await; + } + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = within( + client().request( + Request::post(gateway.url("/bedrock/model/m/converse-stream")) + .body(Body::empty()) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert_eq!( + response.headers()["content-type"], + "application/vnd.amazon.eventstream" + ); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, message); +} + +#[rstest] +#[tokio::test] +async fn relays_a_body_delimited_by_the_upstream_closing() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\nconnection: close\r\n\r\nuntil eof") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()); + let body = within(response).await.unwrap().into_body().collect(); + assert_eq!(within(body).await.unwrap().to_bytes(), "until eof"); +} + +async fn drain(mut body: Incoming) -> (Vec, Result<(), hyper::Error>) { + let mut delivered = Vec::new(); + loop { + match within(body.frame()).await { + None => return (delivered, Ok(())), + Some(Err(error)) => return (delivered, Err(error)), + Some(Ok(frame)) => delivered.extend_from_slice(&frame.into_data().unwrap()), + } + } +} + +#[rstest] +#[case::after_the_head("transfer-encoding: chunked", "", "")] +#[case::between_events("transfer-encoding: chunked", "6\r\ndata: \r\n", "data: ")] +#[case::mid_event("transfer-encoding: chunked", "a\r\ndata:", "data:")] +#[case::short_of_its_content_length("content-length: 100", "data: 1\n\n", "data: 1\n\n")] +#[tokio::test] +async fn an_upstream_dying_mid_stream_fails_the_client_body_instead_of_ending_it( + #[case] framing: &'static str, + #[case] sent: &'static str, + #[case] delivered: &'static str, +) { + let upstream = raw_upstream(move |mut python| async move { + python.request_head().await.unwrap(); + python + .send(format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n{framing}\r\n\r\n{sent}" + )) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let (received, ended) = drain(response.into_body()).await; + assert_eq!(received, delivered.as_bytes()); + assert!(ended.is_err()); +} + +#[rstest] +#[tokio::test] +async fn response_headers_reach_the_client_before_python_sends_any_body() { + let release = Arc::new(Notify::new()); + let released = release.clone(); + let upstream = raw_upstream(move |mut python| { + let released = released.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream; charset=utf-8\r\nx-accel-buffering: no\r\nx-litellm-call-id: call-1\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + released.notified().await; + python.send(chunk(b"data: 1\n\n")).await; + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + + let response = + within(client().request(Request::post(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + + assert_eq!(response.headers()["x-litellm-call-id"], "call-1"); + assert_eq!(response.headers()["x-accel-buffering"], "no"); + release.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, "data: 1\n\n"); +} + +/// A connection closed without any response looks retryable to client SDKs, which would re-run a +/// request Python already ran. Python dying right after its head races the relay, hence the loop. +#[rstest] +#[tokio::test] +async fn a_head_python_sent_right_before_dying_still_reaches_the_client() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..20 { + let mut raw = gateway.connect().await; + raw.send("POST / HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + let head = within(raw.response_head()).await; + assert_eq!(head.map(|head| head.status), Some(StatusCode::OK)); + } +} + +#[rstest] +#[tokio::test] +async fn a_client_that_stops_reading_stalls_python_instead_of_being_buffered_for() { + let written = Arc::new(AtomicUsize::new(0)); + let counted = written.clone(); + let (failed, mut write_failed) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let counted = counted.clone(); + let failed = failed.clone(); + async move { + python.request_head().await.unwrap(); + let mut python = python.into_stream(); + python + .write_all(b"HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n") + .await + .unwrap(); + let piece = chunk(&[0x5a; 64 * 1024]); + while python.write_all(&piece).await.is_ok() { + counted.fetch_add(piece.len(), Ordering::SeqCst); + } + failed.send(()).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST / HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + within(raw.response_head()).await.unwrap(); + + tokio::time::sleep(Duration::from_secs(1)).await; + let stalled = written.load(Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(250)).await; + + assert_eq!(written.load(Ordering::SeqCst), stalled); + assert!(stalled < 64 * 1024 * 1024, "{stalled} bytes in flight"); + drop(raw); + within(write_failed.recv()).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn an_http10_client_receives_a_streamed_body_delimited_by_close() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send([ + &b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n"[..], + &chunk(b"data: 1\n"), + &chunk(b"data: 2\n"), + b"0\r\n\r\n", + ] + .concat()) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/stream HTTP/1.0\r\nhost: x\r\n\r\n").await; + + let response = String::from_utf8(within(raw.until_closed()).await).unwrap(); + let (head, body) = response.split_once("\r\n\r\n").unwrap(); + assert!(head.starts_with("HTTP/1.0 200 OK"), "{head}"); + assert!( + !head.to_ascii_lowercase().contains("transfer-encoding"), + "{head}" + ); + assert_eq!(body, "data: 1\ndata: 2\n"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/support/mod.rs b/litellm-rust/crates/gateway-forward/tests/support/mod.rs new file mode 100644 index 00000000000..7d92ca4a14a --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/support/mod.rs @@ -0,0 +1,320 @@ +#![allow( + dead_code, + reason = "each test binary uses a different subset of the harness" +)] + +use std::{future::Future, net::SocketAddr, time::Duration}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::{HeaderMap, Method, StatusCode, Uri}, +}; +use http_body_util::BodyExt; +use hyper_util::{ + client::legacy::{Client, connect::HttpConnector}, + rt::TokioExecutor, +}; +use litellm_gateway_forward::Forward; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::{mpsc, oneshot}, + task::JoinHandle, +}; + +pub const WAIT: Duration = Duration::from_secs(5); +pub const UVICORN_KEEPALIVE: Duration = Duration::from_secs(5); +pub const DRAIN: Duration = Duration::from_secs(1); + +pub async fn within(future: impl Future) -> T { + tokio::time::timeout(WAIT, future).await.expect("timed out") +} + +/// A forwarder in front of `upstream` that drains once `shutdown` resolves, as `run` does. +pub async fn serve_until( + upstream: SocketAddr, + keepalive: Duration, + shutdown: impl Future + Send + 'static, +) -> (SocketAddr, JoinHandle<()>) { + let forward = Forward::new(&format!("http://{upstream}").parse().unwrap(), keepalive); + let public = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = public.local_addr().unwrap(); + ( + address, + tokio::spawn(forward.serve(public, shutdown, DRAIN)), + ) +} + +/// A running forwarder in front of `upstream`, listening on a real socket like production. +pub struct Gateway { + pub address: SocketAddr, + stop: oneshot::Sender<()>, + server: JoinHandle<()>, +} + +impl Gateway { + pub async fn start(upstream: SocketAddr) -> Self { + Self::with_keepalive(upstream, UVICORN_KEEPALIVE).await + } + + pub async fn with_keepalive(upstream: SocketAddr, keepalive: Duration) -> Self { + let (stop, stopped) = oneshot::channel::<()>(); + let (address, server) = serve_until(upstream, keepalive, async { + stopped.await.ok(); + }) + .await; + Self { + address, + stop, + server, + } + } + + pub fn url(&self, path_and_query: &str) -> Uri { + format!("http://{}{path_and_query}", self.address) + .parse() + .unwrap() + } + + pub async fn connect(&self) -> Raw { + Raw::new(TcpStream::connect(self.address).await.unwrap()) + } + + /// Starts graceful shutdown and returns the server task, which ends once in-flight work drains + /// or `DRAIN` passes. + pub fn shutdown(self) -> JoinHandle<()> { + self.stop.send(()).unwrap(); + self.server + } +} + +pub fn client() -> Client { + Client::builder(TokioExecutor::new()).build_http() +} + +/// What the upstream received, after the forwarder rewrote it. +pub struct Received { + pub method: Method, + pub uri: Uri, + pub headers: HeaderMap, + pub body: Bytes, +} + +/// An axum upstream that records every request it fully reads and answers with `respond`. +pub async fn recording( + respond: impl Fn() -> axum::response::Response + Clone + Send + Sync + 'static, +) -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = serve(Router::new().fallback(move |request: Request| { + let sender = sender.clone(); + let respond = respond.clone(); + async move { + let (parts, body) = request.into_parts(); + let body = body.collect().await.unwrap().to_bytes(); + sender + .send(Received { + method: parts.method, + uri: parts.uri, + headers: parts.headers, + body, + }) + .unwrap(); + respond() + } + })) + .await; + (address, receiver) +} + +pub async fn serve(router: Router) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + address +} + +/// An upstream that hands each accepted connection to `handle` as raw bytes, for behavior a +/// well-behaved HTTP server would never produce. +pub async fn raw_upstream(handle: F) -> SocketAddr +where + F: Fn(Raw) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, +{ + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + tokio::spawn(handle(Raw::new(stream))); + } + }); + address +} + +pub fn unreachable_upstream() -> SocketAddr { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap() +} + +pub struct Head { + pub status: StatusCode, + pub headers: Vec<(String, Vec)>, + pub path: String, +} + +impl Head { + pub fn header(&self, name: &str) -> Option<&[u8]> { + self.headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_slice()) + } + + pub fn content_length(&self) -> Option { + self.header("content-length") + .map(|value| std::str::from_utf8(value).unwrap().parse().unwrap()) + } +} + +/// One TCP connection driven byte by byte, so tests control framing, timing and disconnects. +/// Bytes read past a message head stay buffered for the next read. +pub struct Raw { + stream: TcpStream, + buffer: Vec, +} + +impl Raw { + pub fn new(stream: TcpStream) -> Self { + Self { + stream, + buffer: Vec::new(), + } + } + + pub fn into_stream(self) -> TcpStream { + assert!(self.buffer.is_empty(), "buffered bytes would be lost"); + self.stream + } + + pub fn local_address(&self) -> SocketAddr { + self.stream.local_addr().unwrap() + } + + pub async fn send(&mut self, bytes: impl AsRef<[u8]>) { + self.stream.write_all(bytes.as_ref()).await.unwrap(); + } + + /// Reads one request head, or `None` when the peer closed the connection first. + pub async fn request_head(&mut self) -> Option { + self.head(|bytes| { + let mut headers = [httparse::EMPTY_HEADER; 256]; + let mut request = httparse::Request::new(&mut headers); + let httparse::Status::Complete(length) = request.parse(bytes).unwrap() else { + return None; + }; + Some(( + length, + Head { + status: StatusCode::OK, + headers: owned(request.headers), + path: request.path.unwrap().to_owned(), + }, + )) + }) + .await + } + + pub async fn response_head(&mut self) -> Option { + self.head(|bytes| { + let mut headers = [httparse::EMPTY_HEADER; 256]; + let mut response = httparse::Response::new(&mut headers); + let httparse::Status::Complete(length) = response.parse(bytes).unwrap() else { + return None; + }; + Some(( + length, + Head { + status: StatusCode::from_u16(response.code.unwrap()).unwrap(), + headers: owned(response.headers), + path: String::new(), + }, + )) + }) + .await + } + + async fn head(&mut self, parse: impl Fn(&[u8]) -> Option<(usize, Head)>) -> Option { + loop { + if let Some((length, head)) = parse(&self.buffer) { + self.buffer.drain(..length); + return Some(head); + } + if self.fill().await == 0 { + return None; + } + } + } + + /// Reads exactly `length` body bytes, using what is already buffered first. + pub async fn body(&mut self, length: usize) -> Vec { + while self.buffer.len() < length { + assert_ne!(self.fill().await, 0, "closed mid-body"); + } + self.buffer.drain(..length).collect() + } + + /// Reads a whole chunked body and returns its decoded bytes. + pub async fn chunked_body(&mut self) -> Vec { + let mut decoded = Vec::new(); + loop { + let size_line = self.line().await; + let size = + usize::from_str_radix(size_line.split(';').next().unwrap().trim(), 16).unwrap(); + if size == 0 { + while !self.line().await.is_empty() {} + return decoded; + } + decoded.extend(self.body(size).await); + assert_eq!(self.body(2).await, b"\r\n"); + } + } + + async fn line(&mut self) -> String { + loop { + if let Some(end) = self.buffer.windows(2).position(|window| window == b"\r\n") { + let line = String::from_utf8(self.buffer[..end].to_vec()).unwrap(); + self.buffer.drain(..end + 2); + return line; + } + assert_ne!(self.fill().await, 0, "closed mid-line"); + } + } + + async fn fill(&mut self) -> usize { + let mut chunk = [0; 8192]; + let read = self.stream.read(&mut chunk).await.unwrap_or(0); + self.buffer.extend_from_slice(&chunk[..read]); + read + } + + /// Everything the peer sends until it closes, including bytes already buffered. + pub async fn until_closed(&mut self) -> Vec { + while self.fill().await != 0 {} + std::mem::take(&mut self.buffer) + } + + /// Resolves once the peer closes the connection, discarding anything it still sends. + pub async fn closed(&mut self) { + let mut chunk = [0; 8192]; + while self.stream.read(&mut chunk).await.unwrap_or(0) != 0 {} + } +} + +fn owned(headers: &[httparse::Header<'_>]) -> Vec<(String, Vec)> { + headers + .iter() + .map(|header| (header.name.to_owned(), header.value.to_vec())) + .collect() +} diff --git a/litellm-rust/crates/gateway-forward/tests/upload.rs b/litellm-rust/crates/gateway-forward/tests/upload.rs new file mode 100644 index 00000000000..3dc865c9c86 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/upload.rs @@ -0,0 +1,339 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::StatusCode, + response::Response, +}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, Raw, client, raw_upstream, recording, serve, within}; +use tokio::{ + io::AsyncWriteExt, + sync::{Notify, mpsc}, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +#[rstest] +#[tokio::test] +async fn streams_upload_and_response_incrementally_without_a_body_limit() { + let upstream = serve( + Router::new().fallback(|request: Request| async { Response::new(request.into_body()) }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let (sender, receiver) = mpsc::channel::>(1); + let stream = futures_util::stream::unfold(receiver, |mut receiver| async { + receiver.recv().await.map(|chunk| (chunk, receiver)) + }); + sender.send(Ok(Bytes::from_static(b"first"))).await.unwrap(); + let mut response = within( + client().request( + Request::post(gateway.url("/v1/files")) + .body(Body::from_stream(stream)) + .unwrap(), + ), + ) + .await + .unwrap(); + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), "first"); + let large = Bytes::from(vec![b'x'; 8 * 1024 * 1024]); + sender.send(Ok(large.clone())).await.unwrap(); + drop(sender); + let rest = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(rest, large); +} + +/// hyper sends the 100 as soon as the request is forwarded, where uvicorn waits for the app to +/// read the body, so an upload Python rejects on headers alone is still sent in full. +#[rstest] +#[tokio::test] +async fn answers_expect_continue_so_large_uploads_start() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/audio/transcriptions HTTP/1.1\r\nhost: x\r\nexpect: 100-continue\r\ncontent-length: 4\r\n\r\n") + .await; + let interim = within(raw.response_head()).await.unwrap(); + assert_eq!(interim.status, StatusCode::CONTINUE); + raw.send("data").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "data"); +} + +/// curl sends `expect: 100-continue` for any body over 1 MiB and waits for the 100 before sending +/// the multipart body, which must reach Python intact even where it looks like framing. +#[rstest] +#[tokio::test] +async fn expect_continue_delivers_a_large_multipart_upload_byte_for_byte() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let boundary = "------------------------d74496d66958873e"; + let audio: Vec = (0..3 * 1024 * 1024) + .map(|index: u32| (index.wrapping_mul(2_654_435_761) >> 24) as u8) + .chain(*b"\r\n0\r\n\r\nPOST /smuggled HTTP/1.1\r\n\r\n--") + .collect(); + let body = [ + format!("--{boundary}\r\ncontent-disposition: form-data; name=\"model\"\r\n\r\nwhisper-1\r\n--{boundary}\r\ncontent-disposition: form-data; name=\"file\"; filename=\"a b.wav\"\r\ncontent-type: audio/wav\r\n\r\n").into_bytes(), + audio, + format!("\r\n--{boundary}--\r\n").into_bytes(), + ] + .concat(); + let content_type = format!("multipart/form-data; boundary={boundary}"); + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/audio/transcriptions HTTP/1.1\r\nhost: x\r\ncontent-type: {content_type}\r\nexpect: 100-continue\r\ncontent-length: {}\r\n\r\n", + body.len() + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::CONTINUE + ); + raw.send(&body).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["content-type"], content_type.as_str()); + assert_eq!(request.body, body); +} + +#[rstest] +#[case::get("GET")] +#[case::delete("DELETE")] +#[case::options("OPTIONS")] +#[case::post("POST")] +#[tokio::test] +async fn forwards_a_chunked_body_on_any_method(#[case] method: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "{method} /v1/search HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n4\r\nbody\r\n0\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "body"); +} + +#[rstest] +#[case::no_framing("", "", &[], "")] +#[case::empty_content_length("content-length: 0\r\n", "", &[("content-length", "0")], "")] +#[case::empty_chunked( + "transfer-encoding: chunked\r\n", + "0\r\n\r\n", + &[("transfer-encoding", "chunked")], + "" +)] +#[case::sized( + "content-length: 7\r\n", + "payload", + &[("content-length", "7")], + "payload" +)] +#[case::chunk_extensions_and_trailers( + "transfer-encoding: chunked\r\n", + "4;name=\"v\"\r\nbody\r\n0\r\nx-checksum: 1\r\n\r\n", + &[("transfer-encoding", "chunked")], + "body" +)] +#[tokio::test] +async fn a_body_reaches_python_with_the_framing_the_client_chose( + #[values("POST", "GET")] method: &str, + #[case] framing: &str, + #[case] wire_body: &str, + #[case] python_framing: &[(&str, &str)], + #[case] python_body: &str, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "{method} /v1/files HTTP/1.1\r\nhost: x\r\n{framing}\r\n{wire_body}" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + + let request = received.recv().await.unwrap(); + let forwarded: Vec<(&str, &str)> = ["content-length", "transfer-encoding"] + .into_iter() + .flat_map(|name| { + request + .headers + .get_all(name) + .iter() + .map(move |value| (name, value.to_str().unwrap())) + }) + .collect(); + assert_eq!(request.method, method); + assert_eq!(forwarded, python_framing); + assert_eq!(request.body, python_body); +} + +#[rstest] +#[tokio::test] +async fn a_large_upload_is_paced_by_python_reading_it() { + const LENGTH: usize = 128 * 1024 * 1024; + let start_reading = Arc::new(Notify::new()); + let reading = start_reading.clone(); + let upstream = raw_upstream(move |mut python| { + let reading = reading.clone(); + async move { + let head = python.request_head().await.unwrap(); + reading.notified().await; + let mut remaining = head.content_length().unwrap(); + while remaining > 0 { + remaining -= python.body(remaining.min(1 << 20)).await.len(); + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (read, mut write) = gateway.connect().await.into_stream().into_split(); + let sent = Arc::new(AtomicUsize::new(0)); + let sending = sent.clone(); + let upload = tokio::spawn(async move { + write + .write_all( + format!("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: {LENGTH}\r\n\r\n") + .as_bytes(), + ) + .await + .unwrap(); + let chunk = [0x5a; 64 * 1024]; + while sending.load(Ordering::SeqCst) < LENGTH { + write.write_all(&chunk).await.unwrap(); + sending.fetch_add(chunk.len(), Ordering::SeqCst); + } + write + }); + + tokio::time::sleep(Duration::from_secs(1)).await; + assert!( + sent.load(Ordering::SeqCst) < LENGTH / 4, + "the forwarder buffered an upload Python was not reading" + ); + start_reading.notify_one(); + let write = tokio::time::timeout(Duration::from_secs(60), upload) + .await + .unwrap() + .unwrap(); + let mut raw = Raw::new(read.reunite(write).unwrap()); + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); +} + +#[rstest] +#[case::between_chunks("5\r\nhello\r\n")] +#[case::mid_chunk("100\r\nhello")] +#[tokio::test] +async fn a_client_abandoning_a_chunked_upload_never_ends_python_body_cleanly( + #[case] sent: &'static str, +) { + let (forwarded, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let forwarded = forwarded.clone(); + async move { + python.request_head().await.unwrap(); + forwarded.send(python.until_closed().await).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{sent}" + )) + .await; + tokio::time::sleep(Duration::from_millis(100)).await; + drop(raw); + + let body = within(received.recv()).await.unwrap(); + assert!(body.starts_with(b"5\r\nhello\r\n")); + assert!( + !body.ends_with(b"0\r\n\r\n"), + "python saw a truncated upload end cleanly" + ); +} + +#[rstest] +#[tokio::test] +async fn an_early_answer_during_an_upload_keeps_both_connections_in_sync() { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + let length = head.content_length(); + event.send(head.path.into_bytes()).unwrap(); + let Some(length) = length else { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await; + continue; + }; + python + .send("HTTP/1.1 413 Payload Too Large\r\ncontent-length: 2\r\n\r\n{}") + .await; + event.send(python.body(length).await).unwrap(); + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let upload = b"GET /smuggled HTTP/1.1\r\nhost: x\r\n\r\n".repeat(1024); + let (first_half, second_half) = upload.split_at(upload.len() / 2); + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: {}\r\n\r\n", + upload.len() + )) + .await; + raw.send(first_half).await; + + let rejected = within(raw.response_head()).await.unwrap(); + assert_eq!(rejected.status, StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!(raw.body(2).await, b"{}"); + raw.send(second_half).await; + raw.send("GET /next HTTP/1.1\r\nhost: x\r\n\r\n").await; + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); + + assert_eq!(within(events.recv()).await.unwrap(), b"/v1/files"); + assert_eq!(within(events.recv()).await.unwrap(), upload); + assert_eq!(within(events.recv()).await.unwrap(), b"/next"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/upstream.rs b/litellm-rust/crates/gateway-forward/tests/upstream.rs new file mode 100644 index 00000000000..25e4768a54e --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/upstream.rs @@ -0,0 +1,21 @@ +use litellm_gateway_forward::Upstream; +use rstest::rstest; + +#[rstest] +#[case::https("https://127.0.0.1:4000")] +#[case::public("http://192.0.2.1:4000")] +#[case::hostname("http://localhost:4000")] +#[case::ipv6("http://[::1]:4000")] +#[case::path("http://127.0.0.1:4000/api")] +#[case::query("http://127.0.0.1:4000/?x=y")] +#[case::credentials("http://user@127.0.0.1:4000")] +fn rejects_anything_but_a_loopback_http_origin(#[case] upstream: &str) { + assert!(upstream.parse::().is_err()); +} + +#[rstest] +#[case::standard("http://127.0.0.1:4000")] +#[case::loopback_range("http://127.0.0.2:4000/")] +fn accepts_a_loopback_http_origin(#[case] upstream: &str) { + assert!(upstream.parse::().is_ok()); +} diff --git a/litellm-rust/crates/gateway-forward/tests/websocket.rs b/litellm-rust/crates/gateway-forward/tests/websocket.rs new file mode 100644 index 00000000000..ff64ddf38d3 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/websocket.rs @@ -0,0 +1,719 @@ +mod support; + +use std::{ + net::SocketAddr, + sync::Arc, + time::{Duration, Instant}, +}; + +use axum::{ + Router, + http::{HeaderValue, StatusCode}, + response::IntoResponse, +}; +use futures_util::{SinkExt, StreamExt, future::join_all}; +use litellm_gateway_forward::PEER_HEADER; +use rstest::rstest; +use support::{DRAIN, Gateway, Head, Raw, raw_upstream, serve, within}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, + sync::{Barrier, mpsc}, +}; +use tokio_tungstenite::{ + WebSocketStream, accept_hdr_async, client_async, + tungstenite::{ + Error, Message, + client::IntoClientRequest, + handshake::{ + derive_accept_key, + server::{ErrorResponse, Request as Handshake, Response as Accepted}, + }, + protocol::{CloseFrame, frame::coding::CloseCode}, + }, +}; + +const KEY: &str = "dGhlIHNhbXBsZSBub25jZQ=="; + +async fn connect( + address: SocketAddr, + connection: &'static str, +) -> Result, Box> { + let mut request = format!("ws://{address}/v1/realtime?model=m") + .into_client_request() + .unwrap(); + let headers = request.headers_mut(); + headers.insert("authorization", HeaderValue::from_static("Bearer test")); + headers.insert( + "sec-websocket-protocol", + HeaderValue::from_static("realtime"), + ); + headers.insert("connection", HeaderValue::from_static(connection)); + let stream = TcpStream::connect(address).await.unwrap(); + let (socket, response) = client_async(request, stream).await.map_err(Box::new)?; + assert_eq!(response.headers()["sec-websocket-protocol"], "realtime"); + Ok(socket) +} + +/// Greets, echoes every message, and closes with 4000 when told to "bye". +#[expect( + clippy::result_large_err, + reason = "tungstenite's handshake callback returns ErrorResponse" +)] +async fn echo() -> SocketAddr { + raw_upstream(|python| async move { + let callback = |request: &Handshake, mut response: Accepted| { + assert_eq!(request.uri(), "/v1/realtime?model=m"); + assert_eq!(request.headers()["authorization"], "Bearer test"); + assert!(request.headers().contains_key(PEER_HEADER)); + response.headers_mut().insert( + "sec-websocket-protocol", + HeaderValue::from_static("realtime"), + ); + Ok::<_, ErrorResponse>(response) + }; + let mut socket = accept_hdr_async(python.into_stream(), callback) + .await + .unwrap(); + socket.send(Message::text("greeting")).await.unwrap(); + while let Some(Ok(message)) = socket.next().await { + if message == Message::text("bye") { + socket + .close(Some(CloseFrame { + code: CloseCode::Library(4000), + reason: "done".into(), + })) + .await + .unwrap(); + continue; + } + if message.is_close() { + break; + } + socket.send(message).await.unwrap(); + } + }) + .await +} + +fn handshake(target: &str, headers: &str) -> String { + format!( + "GET {target} HTTP/1.1\r\nhost: public.example\r\nconnection: Upgrade\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: {KEY}\r\n{headers}\r\n" + ) +} + +/// The 101 uvicorn writes for `handshake`, with `headers` added. +fn switching(handshake: &Head, headers: &str) -> String { + let accept = derive_accept_key(handshake.header("sec-websocket-key").unwrap()); + format!( + "HTTP/1.1 101 Switching Protocols\r\nupgrade: websocket\r\nconnection: Upgrade\r\nsec-websocket-accept: {accept}\r\n{headers}\r\n" + ) +} + +fn values(head: &Head, name: &str) -> Vec { + head.headers + .iter() + .filter(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| String::from_utf8(value.clone()).unwrap()) + .collect() +} + +/// Accepts with `headers` in the 101, hands the handshake head to the test, then echoes bytes. +async fn byte_echo(headers: &'static str) -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, handshakes) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let sender = sender.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, headers)).await; + sender.send(handshake).unwrap(); + echo_bytes(python).await; + } + }) + .await; + (upstream, handshakes) +} + +async fn echo_bytes(python: Raw) { + let mut stream = python.into_stream(); + let (mut read, mut write) = stream.split(); + tokio::io::copy(&mut read, &mut write).await.ok(); +} + +async fn assert_echoes(socket: &mut TcpStream, message: &[u8]) { + socket.write_all(message).await.unwrap(); + let mut echoed = vec![0; message.len()]; + within(socket.read_exact(&mut echoed)).await.unwrap(); + assert_eq!(echoed, message); +} + +async fn open(gateway: &Gateway, handshake: String) -> (TcpStream, Head) { + let mut raw = gateway.connect().await; + raw.send(handshake).await; + let switched = within(raw.response_head()).await.unwrap(); + assert_eq!(switched.status, StatusCode::SWITCHING_PROTOCOLS); + (raw.into_stream(), switched) +} + +#[rstest] +#[tokio::test] +async fn relays_messages_both_ways( + #[values( + "Upgrade", + "keep-alive, Upgrade", + "upgrade", + "Upgrade, Sec-WebSocket-Key, Sec-WebSocket-Version" + )] + connection: &'static str, + #[values( + Message::text("hello"), + Message::binary(vec![0, 1, 2]), + Message::binary(vec![7; 4 * 1024 * 1024]) + )] + message: Message, +) { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, connection)).await.unwrap(); + assert_eq!( + socket.next().await.unwrap().unwrap(), + Message::text("greeting") + ); + socket.send(message.clone()).await.unwrap(); + assert_eq!(within(socket.next()).await.unwrap().unwrap(), message); + socket.close(None).await.unwrap(); + while socket.next().await.is_some() {} + drop(socket); + within(gateway.shutdown()).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn relays_a_close_python_initiates() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + socket.send(Message::text("bye")).await.unwrap(); + let Message::Close(Some(frame)) = within(socket.next()).await.unwrap().unwrap() else { + panic!("expected python's close frame"); + }; + assert_eq!(frame.code, CloseCode::Library(4000)); + assert_eq!(frame.reason, "done"); +} + +#[rstest] +#[tokio::test] +async fn passes_an_upstream_rejection_through() { + let upstream = serve( + Router::new() + .fallback(|| async { (StatusCode::FORBIDDEN, "origin not allowed").into_response() }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let error = within(connect(gateway.address, "Upgrade")) + .await + .unwrap_err(); + let Error::Http(response) = *error else { + panic!("expected the upstream rejection, got {error}"); + }; + assert_eq!(response.status(), StatusCode::FORBIDDEN); + assert_eq!(response.body().as_deref(), Some(&b"origin not allowed"[..])); +} + +/// The idle timeout only covers a connection waiting for its next HTTP request, and a realtime +/// session may sit silent between turns. +#[rstest] +#[tokio::test] +async fn keeps_a_quiet_socket_open_past_the_keepalive() { + let gateway = Gateway::with_keepalive(echo().await, Duration::from_millis(400)).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + tokio::time::sleep(Duration::from_millis(1200)).await; + socket.send(Message::text("still here")).await.unwrap(); + assert_eq!( + within(socket.next()).await.unwrap().unwrap(), + Message::text("still here") + ); +} + +/// Browsers carry the API key in a subprotocol, Python accepts with the first one it was offered, +/// and uvicorn negotiates permessage-deflate, so every handshake header must cross unchanged. +#[rstest] +#[tokio::test] +async fn handshake_headers_cross_unchanged_both_ways() { + let (upstream, mut handshakes) = byte_echo( + "sec-websocket-protocol: realtime\r\nsec-websocket-extensions: permessage-deflate; server_max_window_bits=12; client_max_window_bits=12\r\nset-cookie: a=1\r\nset-cookie: b=2\r\nx-litellm-call-id: abc\r\n", + ) + .await; + let gateway = Gateway::start(upstream).await; + let target = "/v1/realtime?model=gpt-realtime&api-key=sk%2Fa%2Bb%3D&dup=1&dup=2"; + let (_socket, switched) = open( + &gateway, + handshake( + target, + "sec-websocket-protocol: realtime, openai-insecure-api-key.sk-websocket-test, openai-beta.realtime-v1\r\nsec-websocket-protocol: second-line\r\nsec-websocket-extensions: permessage-deflate; client_max_window_bits\r\ncookie: token=abc\r\n", + ), + ) + .await; + let handshake = within(handshakes.recv()).await.unwrap(); + + assert_eq!(handshake.path, target); + assert_eq!( + values(&handshake, "sec-websocket-protocol"), + [ + "realtime, openai-insecure-api-key.sk-websocket-test, openai-beta.realtime-v1", + "second-line" + ] + ); + assert_eq!( + values(&handshake, "sec-websocket-extensions"), + ["permessage-deflate; client_max_window_bits"] + ); + assert_eq!(values(&handshake, "cookie"), ["token=abc"]); + assert_eq!(values(&switched, "sec-websocket-protocol"), ["realtime"]); + assert_eq!( + values(&switched, "sec-websocket-extensions"), + ["permessage-deflate; server_max_window_bits=12; client_max_window_bits=12"] + ); + assert_eq!(values(&switched, "set-cookie"), ["a=1", "b=2"]); + assert_eq!(values(&switched, "x-litellm-call-id"), ["abc"]); +} + +#[rstest] +#[tokio::test] +async fn invents_no_subprotocol_or_extension() { + let (upstream, mut handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (_socket, switched) = open(&gateway, handshake("/v1/realtime", "")).await; + let handshake = within(handshakes.recv()).await.unwrap(); + for name in ["sec-websocket-protocol", "sec-websocket-extensions"] { + assert!(handshake.header(name).is_none()); + assert!(switched.header(name).is_none()); + } +} + +#[rstest] +#[case::capitalized_upgrade("Upgrade: WebSocket\r\nConnection: keep-alive, Upgrade\r\n")] +#[case::uppercase("UPGRADE: WEBSOCKET\r\nCONNECTION: UPGRADE\r\n")] +#[case::split_connection("Upgrade: websocket\r\nConnection: keep-alive\r\nConnection: Upgrade\r\n")] +#[tokio::test] +async fn accepts_any_spelling_of_the_upgrade_headers(#[case] upgrade: &str) { + let (upstream, _handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open( + &gateway, + format!( + "GET /v1/realtime HTTP/1.1\r\nhost: x\r\n{upgrade}sec-websocket-version: 13\r\nsec-websocket-key: {KEY}\r\n\r\n" + ), + ) + .await; + socket.write_all(b"ping").await.unwrap(); + let mut echoed = [0; 4]; + within(socket.read_exact(&mut echoed)).await.unwrap(); + assert_eq!(&echoed, b"ping"); +} + +/// Compressed frames set RSV1 and every client frame is masked, so the socket is relayed as bytes, +/// never re-framed. +#[rstest] +#[tokio::test] +async fn relays_socket_bytes_exactly() { + let (upstream, _handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let sent: Vec = (0..3_000_017u32) + .map(|index| index.wrapping_mul(2_654_435_761).to_be_bytes()[1]) + .collect(); + let (mut read, mut write) = socket.into_split(); + let writing = sent.clone(); + let writer = tokio::spawn(async move { + for chunk in writing.chunks(7919) { + write.write_all(chunk).await.unwrap(); + } + write + }); + let mut echoed = vec![0; sent.len()]; + within(read.read_exact(&mut echoed)).await.unwrap(); + assert!(echoed == sent); + drop(writer.await.unwrap()); +} + +/// uvicorn writes the 101, an error event and a close frame back to back and then closes, so they +/// share a segment with the FIN, and a client may send its first frame with the handshake. +#[rstest] +#[tokio::test] +async fn keeps_bytes_that_arrive_with_the_handshake() { + let (early, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let early = early.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + let mut segment = switching(&handshake, "").into_bytes(); + segment.extend(b"\x81\x05error\x88\x08\x03\xf0denied"); + python.send(segment).await; + early.send(within(python.body(6)).await).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(handshake("/v1/realtime", "") + "early!").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::SWITCHING_PROTOCOLS + ); + assert_eq!(within(received.recv()).await.unwrap(), b"early!"); + assert_eq!( + within(raw.until_closed()).await, + b"\x81\x05error\x88\x08\x03\xf0denied" + ); +} + +/// Python frees a realtime session's budget and parallel slot when its socket closes, so a client +/// that disappears without a close frame must still close Python's socket. +#[rstest] +#[case::fin(false)] +#[case::reset(true)] +#[tokio::test] +async fn a_vanished_client_closes_the_python_socket(#[case] reset: bool) { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + event.send("open").unwrap(); + python.closed().await; + event.send("closed").unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_eq!(within(events.recv()).await, Some("open")); + if reset { + #[expect(deprecated, reason = "a zero linger is how a test sends a TCP reset")] + socket.set_linger(Some(Duration::ZERO)).unwrap(); + } + drop(socket); + assert_eq!(within(events.recv()).await, Some("closed")); +} + +#[rstest] +#[case::fin(false)] +#[case::reset(true)] +#[tokio::test] +async fn a_vanished_python_ends_the_client_socket(#[case] reset: bool) { + let upstream = raw_upstream(move |mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + let socket = python.into_stream(); + if reset { + #[expect(deprecated, reason = "a zero linger is how a test sends a TCP reset")] + socket.set_linger(Some(Duration::ZERO)).unwrap(); + } + drop(socket); + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let mut rest = Vec::new(); + assert!(matches!( + within(socket.read_to_end(&mut rest)).await, + Ok(0) | Err(_) + )); +} + +#[rstest] +#[tokio::test] +async fn shutdown_closes_open_sockets_after_drain() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + let started = Instant::now(); + within(gateway.shutdown()).await.unwrap(); + assert!(started.elapsed() >= DRAIN); + assert!(!matches!(socket.next().await, Some(Ok(Message::Text(_))))); +} + +#[rstest] +#[tokio::test] +async fn shutdown_relays_the_close_python_sends_while_draining() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + socket.send(Message::text("bye")).await.unwrap(); + let Message::Close(Some(frame)) = within(socket.next()).await.unwrap().unwrap() else { + panic!("expected python's close frame"); + }; + assert_eq!(frame.code, CloseCode::Library(4000)); + while socket.next().await.is_some() {} + drop(socket); + let started = Instant::now(); + within(server).await.unwrap(); + assert!(started.elapsed() < DRAIN / 2); +} + +#[rstest] +#[tokio::test] +async fn refuses_new_sockets_while_draining() { + let gateway = Gateway::start(echo().await).await; + let mut open = within(connect(gateway.address, "Upgrade")).await.unwrap(); + open.next().await.unwrap().unwrap(); + let address = gateway.address; + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!( + TcpStream::connect(address).await.is_err() || connect(address, "Upgrade").await.is_err() + ); + open.send(Message::text("still open")).await.unwrap(); + assert_eq!( + within(open.next()).await.unwrap().unwrap(), + Message::text("still open") + ); + drop(open); + within(server).await.unwrap(); +} + +/// uvicorn never half-closes: once Python closes the socket nothing more can cross it, so the +/// relay must end instead of waiting for the client to close its side too. +#[rstest] +#[tokio::test] +async fn relay_ends_when_python_closes_while_the_client_holds_its_socket() { + let upstream = raw_upstream(move |mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_eq!(within(socket.read(&mut [0; 1])).await.unwrap(), 0); + let started = Instant::now(); + within(gateway.shutdown()).await.unwrap(); + assert!(started.elapsed() < DRAIN / 2); + drop(socket); +} + +/// uvicorn answers a handshake Python closes before accepting with a 403, and one Python denies +/// through `websocket.http.response` with any status, headers and body, then drops the socket. The +/// client gets that response whole and keeps a connection that still reaches Python. +#[rstest] +#[case::closed_before_accept("403 Forbidden", None, "")] +#[case::denial_response( + "401 Unauthorized", + Some("Bearer realm=\"litellm\""), + "{\"error\":\"invalid api key\"}" +)] +#[tokio::test] +async fn a_denied_handshake_reaches_the_client_whole( + #[case] status: &'static str, + #[case] authenticate: Option<&'static str>, + #[case] body: &'static str, +) { + let upstream = raw_upstream(move |mut python: Raw| async move { + python.request_head().await.unwrap(); + let authenticate = authenticate + .map(|value| format!("www-authenticate: {value}\r\n")) + .unwrap_or_default(); + python + .send(format!( + "HTTP/1.1 {status}\r\ndate: Tue, 03 Nov 2026 08:00:00 GMT\r\nconnection: close\r\ncontent-length: {}\r\ncontent-type: text/plain; charset=utf-8\r\n{authenticate}\r\n{body}", + body.len() + )) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + for _ in 0..2 { + raw.send(handshake("/v1/realtime", "")).await; + let denied = within(raw.response_head()).await.unwrap(); + assert_eq!(denied.status.as_str(), &status[..3]); + assert_eq!( + values(&denied, "www-authenticate"), + Vec::from_iter(authenticate.map(String::from)) + ); + assert!(denied.header("sec-websocket-accept").is_none()); + let length = denied.content_length().unwrap(); + assert_eq!(within(raw.body(length)).await, body.as_bytes()); + } +} + +/// Only Python's EOF ends a relay, so a client that half-closes after its close frame still gets +/// everything Python sends before closing its own side. +#[rstest] +#[tokio::test] +async fn a_client_half_close_still_receives_what_python_sends_after_it() { + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + let received = python.until_closed().await; + python.send(received).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let close = b"\x88\x82\x01\x02\x03\x04\x02\xea"; + socket.write_all(close).await.unwrap(); + socket.shutdown().await.unwrap(); + let mut replied = Vec::new(); + within(socket.read_to_end(&mut replied)).await.unwrap(); + assert_eq!(replied, close); +} + +/// Python authenticates every handshake before accepting it, so sessions opened together must all +/// reach Python before any of them is accepted, and each must stay on its own socket. +#[rstest] +#[tokio::test] +async fn opens_many_sockets_at_once_without_crossing_them() { + const SOCKETS: usize = 64; + let everyone = Arc::new(Barrier::new(SOCKETS)); + let upstream = raw_upstream(move |mut python: Raw| { + let everyone = everyone.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + everyone.wait().await; + python.send(switching(&handshake, "")).await; + echo_bytes(python).await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let gateway = &gateway; + within(join_all((0..SOCKETS).map(|index| async move { + let (mut socket, _) = open(gateway, handshake("/v1/realtime", "")).await; + assert_echoes(&mut socket, format!("socket {index}").as_bytes()).await; + }))) + .await; +} + +/// Python resolves the key, budget and model before it accepts, and nothing bounds how long that +/// takes, so the idle timeout must not cut a handshake waiting on its 101. +#[rstest] +#[tokio::test] +async fn waits_for_a_101_slower_than_the_keepalive() { + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + tokio::time::sleep(Duration::from_millis(1200)).await; + python.send(switching(&handshake, "")).await; + echo_bytes(python).await; + }) + .await; + let gateway = Gateway::with_keepalive(upstream, Duration::from_millis(400)).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_echoes(&mut socket, b"late").await; +} + +/// A 101 the client did not ask for, or to a protocol other than the one it asked for, would hand +/// the client a socket it cannot speak, so it is a 502 and Python's socket is closed. +#[rstest] +#[case::plain_request("GET /v1/models HTTP/1.1\r\nhost: x\r\n\r\n".to_owned(), "websocket")] +#[case::other_protocol(handshake("/v1/realtime", ""), "h2c")] +#[tokio::test] +async fn an_unrequested_switch_is_a_502_that_closes_python( + #[case] request: String, + #[case] protocol: &'static str, +) { + let (closed, mut closes) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let closed = closed.clone(); + async move { + python.request_head().await.unwrap(); + python + .send(format!( + "HTTP/1.1 101 Switching Protocols\r\nconnection: upgrade\r\nupgrade: {protocol}\r\n\r\n" + )) + .await; + python.closed().await; + closed.send(()).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(request).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_GATEWAY + ); + within(closes.recv()).await.unwrap(); +} + +/// A realtime session streams audio for as long as the client reads it, so a client that stops +/// reading must stall Python's writes instead of piling them up in the gateway. +#[rstest] +#[tokio::test] +async fn a_client_that_stops_reading_stalls_python() { + let (sent, mut totals) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let sent = sent.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + python.body(2).await; + let mut stream = python.into_stream(); + let chunk = vec![0x82; 64 * 1024]; + let mut total = 0; + while let Ok(Ok(())) = + tokio::time::timeout(Duration::from_millis(300), stream.write_all(&chunk)).await + { + total += chunk.len(); + } + sent.send(total).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + socket.write_all(b"go").await.unwrap(); + let total = within(totals.recv()).await.unwrap(); + assert!( + total < 64 << 20, + "python wrote {total} bytes into a stalled socket" + ); +} + +/// Python closes a session that is still streaming audio by sending an error event and a close +/// frame and dropping its socket with the client's next frames unread, which resets it. Those last +/// bytes reached the client when it talked to uvicorn directly, so they must reach it through the +/// relay too, even though writing the client's frames to Python now fails. +#[rstest] +#[tokio::test] +async fn python_closing_mid_stream_still_delivers_its_last_frames() { + const LAST: &[u8] = b"\x81\x05error\x88\x02\x03\xf0"; + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + python.body(1).await; + tokio::time::sleep(Duration::from_millis(20)).await; + python.send(LAST).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..20 { + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let (mut read, mut write) = socket.split(); + let streaming = async { + while write.write_all(b"\x82\x80\0\0\0\0").await.is_ok() { + tokio::time::sleep(Duration::from_micros(200)).await; + } + std::future::pending().await + }; + let mut received = Vec::new(); + let receiving = async { + let mut chunk = [0; 4096]; + while let Ok(read @ 1..) = read.read(&mut chunk).await { + received.extend_from_slice(&chunk[..read]); + } + }; + within(async { + tokio::select! { + () = streaming => {} + () = receiving => {} + } + }) + .await; + assert_eq!(received, LAST); + } +} diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 034c3a9826e..69ca561e4b5 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-cache-response.workspace = true axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs index f200fd7b141..b8e6ef4ff4f 100644 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -2,10 +2,10 @@ use litellm_gateway_auth::AuthenticatedRequest; use std::sync::Arc; use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse}; -use litellm_auth::SecretValue; use litellm_inference_ocr::types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput}; use litellm_llms::base_llm::ocr::transformation::decode_request_value; use litellm_llms_types::formats::ocr::OcrDocument; +use litellm_secret_value::SecretValue; use serde_json::Value; use crate::{ diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 737dc0c830c..185cbc8594d 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -13,7 +13,9 @@ use futures_util::future::BoxFuture; use litellm_gateway_inference::{Deployment, Gateway, router}; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; use litellm_inference::resources::CoreResources; -use litellm_secrets::{SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::source::SecretSource; + use serde_json::Value; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; diff --git a/litellm-rust/crates/gateway-management/Cargo.toml b/litellm-rust/crates/gateway-management/Cargo.toml index 502760f8fbe..1ff5dd5319d 100644 --- a/litellm-rust/crates/gateway-management/Cargo.toml +++ b/litellm-rust/crates/gateway-management/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true base64.workspace = true -litellm-auth-types.workspace = true litellm-error.workspace = true litellm-gateway-auth.workspace = true rand.workspace = true diff --git a/litellm-rust/crates/gateway-management/src/keys.rs b/litellm-rust/crates/gateway-management/src/keys.rs index 45d915b7b3d..feeb40b2403 100644 --- a/litellm-rust/crates/gateway-management/src/keys.rs +++ b/litellm-rust/crates/gateway-management/src/keys.rs @@ -1,8 +1,8 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::keys::{KeyHash, KeyStatus}; +use litellm_secret_value::SecretValue; use rand::{RngCore, rngs::OsRng}; use crate::Error; diff --git a/litellm-rust/crates/gateway-management/tests/keys.rs b/litellm-rust/crates/gateway-management/tests/keys.rs index bf0221cf1be..8a970bb407c 100644 --- a/litellm-rust/crates/gateway-management/tests/keys.rs +++ b/litellm-rust/crates/gateway-management/tests/keys.rs @@ -7,7 +7,6 @@ use std::{ }; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ hash_token, keys::{KeyHash, KeyLookup, KeyStatus, verify_key}, @@ -16,6 +15,7 @@ use litellm_gateway_management::{ Error, keys::{Keys, Revocation}, }; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; use support::{MemoryStore, Unavailable}; diff --git a/litellm-rust/crates/gateway-mcp/Cargo.toml b/litellm-rust/crates/gateway-mcp/Cargo.toml index bed48bad58d..65077eec0ea 100644 --- a/litellm-rust/crates/gateway-mcp/Cargo.toml +++ b/litellm-rust/crates/gateway-mcp/Cargo.toml @@ -7,11 +7,11 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum = { workspace = true, features = ["json", "query"] } futures-util.workspace = true litellm-config.workspace = true litellm-secrets.workspace = true -litellm-auth-types.workspace = true litellm-error.workspace = true base64.workspace = true http.workspace = true diff --git a/litellm-rust/crates/gateway-mcp/src/configured.rs b/litellm-rust/crates/gateway-mcp/src/configured.rs index cc7a52fc029..33cbd243530 100644 --- a/litellm-rust/crates/gateway-mcp/src/configured.rs +++ b/litellm-rust/crates/gateway-mcp/src/configured.rs @@ -8,8 +8,8 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use futures_util::future::try_join_all; use http::{HeaderName, HeaderValue}; -use litellm_auth_types::SecretValue; use litellm_config::{McpAuth, McpServer, McpTransport}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use rmcp::{ RoleClient, ServiceExt, diff --git a/litellm-rust/crates/gateway-ui/Cargo.toml b/litellm-rust/crates/gateway-ui/Cargo.toml index 1c65665be72..a6104455e96 100644 --- a/litellm-rust/crates/gateway-ui/Cargo.toml +++ b/litellm-rust/crates/gateway-ui/Cargo.toml @@ -23,10 +23,10 @@ tower-http = { version = "0.6.11", features = ["fs", "set-header"] } tower-sessions.workspace = true [dev-dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true serde_json.workspace = true tower = { version = "0.5", features = ["util"] } -litellm-auth-types.workspace = true rstest.workspace = true tempfile.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/gateway-ui/tests/routes.rs b/litellm-rust/crates/gateway-ui/tests/routes.rs index 95246f33a35..86b7db63ed9 100644 --- a/litellm-rust/crates/gateway-ui/tests/routes.rs +++ b/litellm-rust/crates/gateway-ui/tests/routes.rs @@ -16,8 +16,8 @@ use tower_sessions::{ session::{Id, Record}, }; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::UiBackend; +use litellm_secret_value::SecretValue; struct App { router: Router, diff --git a/litellm-rust/crates/gateway/AGENTS.md b/litellm-rust/crates/gateway/AGENTS.md index a33bd19cc80..16714cd9a1f 100644 --- a/litellm-rust/crates/gateway/AGENTS.md +++ b/litellm-rust/crates/gateway/AGENTS.md @@ -5,8 +5,11 @@ - Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate ## Redis -- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored +- `build_redis` builds the Redis clients once at boot (`src/redis/mod.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored - Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands -- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped +- `src/redis/layers.rs` is the Rust copy of `_merge_redis_kwargs` in `litellm/_redis.py`, which only this host needs because the Python host merges in Python. It puts the configured kwargs over their `REDIS_*` values and fills the `get_secret` fallbacks from the `SETTINGS` rows in `redis-resolve`, under the same conditions Python reads them, then hands the result to `litellm_redis_resolve::interpret`. Values come from a `SecretSnapshot` filled once at boot, read like `get_secret` (secret manager, then environment, with `true` and `false` read as bools), and the only straight environment read is Azure's `REDIS_USERNAME` +- A surface is a `ConfiguredRedis` constructor: `coordination` for `general_settings.coordination_redis`, which rejects a malformed `port` like the proxy's `CoordinationRedisParams` pydantic model, `cache_params` for `litellm_settings.cache_params`, which drops the cache-only keys and reads the `redis_startup_nodes` alias, and `redis_cache_kwargs` for a `RedisCache.redis_kwargs` dict in the contract tests. A new surface is a new constructor, never a flag +- Nulls depend on the surface. `coordination_redis` drops every null like `model_dump(exclude_none=True)`, `cache_params` keeps a null as an explicit `None` that hides the `REDIS_*` value unless `Cache` or `RedisCache` names that parameter, and `RedisCache.redis_kwargs` keeps every null. A missing `os.environ/` reference is a null, as `load_config` resolves it +- A setting the Rust client cannot honor warns and is dropped, and only a security check it cannot enforce fails boot, per `redis-resolve/AGENTS.md` - Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python - `/health/readiness` never touches Redis, matching the proxy diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 4d831aa2314..a14df9b9c53 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -6,22 +6,28 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum.workspace = true envy = "0.4.2" http-body-util = "0.1" litellm-inference.workspace = true litellm-gateway-inference.workspace = true +litellm-gateway-forward.workspace = true litellm-gateway-mcp.workspace = true tokio-util = "0.7" thiserror.workspace = true futures-util.workspace = true litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true -litellm-auth-types.workspace = true litellm-config.workspace = true +litellm-core-utils.workspace = true litellm-error.workspace = true litellm-circuit-breaker.workspace = true litellm-redis-client.workspace = true +litellm-redis-connect.workspace = true +litellm-redis-resolve.workspace = true +litellm-redis-types.workspace = true +rustls.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets.workspace = true @@ -31,10 +37,12 @@ serde.workspace = true tracing.workspace = true tokio = { workspace = true, features = ["signal"] } uuid.workspace = true +veil.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] litellm-redis-testing.workspace = true +litellm-http = { workspace = true, features = ["test-support"] } base64.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index e2f646c5448..c267e79353d 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -12,13 +12,12 @@ pub enum Error { Mcp(#[from] litellm_gateway_mcp::ConnectError), #[error(transparent)] Config(#[from] litellm_config::Error), - #[error("invalid Redis setting {setting}: {reason}")] - RedisInvalid { - setting: &'static str, - reason: &'static str, - }, #[error(transparent)] - RedisConnect(#[from] litellm_redis_client::ConnectError), + Redis(#[from] litellm_redis_resolve::Error), + #[error("Redis setting {setting} is not plain data")] + RedisNotPlainData { setting: String }, + #[error(transparent)] + RedisConnect(#[from] litellm_redis_connect::Error), #[error("MCP requires working gateway authentication")] Auth(#[from] litellm_gateway_auth::Error), } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 12a9bf9110a..0a0191b0499 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -68,15 +68,29 @@ pub fn build_inference(config: &Config) -> Result, Error> { Ok(Arc::new(Gateway::new(resources, http, secrets, router)?)) } -pub async fn build_redis(config: &Config) -> Result { - let environment = secrets::environment_values(&config.environment_variables)?; - let lookup = |name: &str| { +pub fn install_crypto_provider() { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } +} + +pub async fn build_redis( + config: &Config, + pool: &HttpClientPool, + http: &litellm_http::HttpClientConfig, +) -> Result { + let environment = Arc::new(secrets::environment_values(&config.environment_variables)?); + let lookup: litellm_redis_connect::SharedEnvironment = Arc::new(move |name: &str| { environment .get(name) .map(|value| value.expose().to_owned()) .or_else(|| std::env::var(name).ok()) + }); + let resources = redis::Resources { + http: pool.client(http, ClientVariant::Provider)?, + lookup: lookup.clone(), }; - redis::connect(config, &lookup).await + redis::connect(config, &*lookup, &resources).await } pub async fn build_mcp( diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 47c1fe83f50..176bc68d247 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -10,7 +10,7 @@ use serde_json::json; mod settings; -use settings::{Settings, UiSettings}; +use settings::{LogSettings, Settings, UiSettings}; struct StderrSink { level: Level, @@ -41,12 +41,27 @@ impl Sink for StderrSink { #[tokio::main] async fn main() -> Result<(), Box> { + litellm_gateway::install_crypto_provider(); + if let Some(forward) = litellm_gateway_forward::Settings::from_env()? { + let log: LogSettings = settings::from_iter(std::env::vars_os())?; + Logger::new(StderrSink { + level: log.log_level(), + }) + .install_global()?; + let stdin_closed = litellm_gateway_forward::input_closed(std::io::stdin()); + litellm_gateway_forward::run(forward, async { + tokio::select! { () = shutdown_signal() => (), () = stdin_closed => () } + }) + .await?; + return Ok(()); + } let settings: Settings = settings::from_iter(std::env::vars_os())?; let level = settings.log_level(); Logger::new(StderrSink { level }).install_global()?; let config = Config::load(settings.litellm_config)?; let inference = litellm_gateway::build_inference(&config)?; - let _redis = litellm_gateway::build_redis(&config).await?; + let _redis = + litellm_gateway::build_redis(&config, &inference.resources.pool, &inference.http).await?; let ui = match std::env::var_os("LITELLM_UI_PATH") { Some(directory) => { let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?; @@ -126,7 +141,20 @@ async fn shutdown_signal() { } } } - #[cfg(not(unix))] + // Windows support exists for local CLI runs only, the proxy is never deployed on Windows. The supervisor + // stops the gateway with CTRL_BREAK_EVENT, which has to drain the same way SIGTERM does. + #[cfg(windows)] + { + match tokio::signal::windows::ctrl_break() { + Ok(mut ctrl_break) => { + tokio::select! { _ = tokio::signal::ctrl_c() => (), _ = ctrl_break.recv() => () } + } + Err(_) => { + let _ = tokio::signal::ctrl_c().await; + } + } + } + #[cfg(not(any(unix, windows)))] { let _ = tokio::signal::ctrl_c().await; } diff --git a/litellm-rust/crates/gateway/src/redis/configured.rs b/litellm-rust/crates/gateway/src/redis/configured.rs new file mode 100644 index 00000000000..93297d0e57f --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/configured.rs @@ -0,0 +1,144 @@ +use std::collections::BTreeMap; + +use litellm_config::{RedisSettings, Value}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_resolve::{Expected, Kwargs}; +use litellm_redis_types::{ProviderHandle, Resolution}; +use veil::Redact; + +use super::{layers, secrets::SecretSnapshot}; +use crate::Error; + +const CACHE_ONLY: &[&str] = &[ + "type", + "mode", + "ttl", + "default_in_memory_ttl", + "default_in_redis_ttl", + "supported_call_types", + "similarity_threshold", + "max_messages", + "_backend", +]; + +/// The redis-py kwargs one surface hands `_get_redis_client_logic`, before `REDIS_*` and secrets. +#[derive(Clone, Redact)] +pub(crate) struct ConfiguredRedis { + #[redact(with = "[REDACTED]")] + pub(super) fields: BTreeMap, + pub(super) nulls: Nulls, + pub(super) credential_provider: Option, +} + +/// Whether a null setting reaches redis-py as an explicit `None`, which outranks its `REDIS_*` value. +#[derive(Clone, Copy, Debug)] +pub(super) enum Nulls { + /// `coordination_redis` reaches redis-py through `model_dump(exclude_none=True)`. + Dropped, + /// `Cache` drops a null only for the parameters it names and passes any other key on. + KeptUnlessNamedByCache, + /// `RedisCache.redis_kwargs` already holds only the nulls redis-py receives. + #[cfg(test)] + Kept, +} + +impl ConfiguredRedis { + /// `general_settings.coordination_redis`, which the proxy validates as a pydantic model first. + /// An `os.environ/` port is only known after merging, where it is read like any other port. + pub(crate) fn coordination( + settings: &RedisSettings, + ) -> Result { + if let Some(port) = settings.fields().get("port").and_then(text) + && !port.starts_with("os.environ/") + && port.parse::().is_err() + { + return Err(litellm_redis_resolve::Error::Invalid { + setting: "port".to_owned(), + expected: Expected::Integer, + }); + } + Ok(Self { + fields: settings.fields().clone(), + nulls: Nulls::Dropped, + credential_provider: None, + }) + } + + /// `litellm_settings.cache_params` with `type: redis`, as `litellm.Cache` hands it to `RedisCache`. + pub(crate) fn cache_params(settings: &RedisSettings) -> Self { + let fields = settings.fields(); + let alias = fields + .get("redis_startup_nodes") + .filter(|_| !fields.contains_key("startup_nodes")) + .map(|nodes| ("startup_nodes".to_owned(), nodes.clone())); + let fields = fields + .iter() + .filter(|(key, _)| { + key.as_str() != "redis_startup_nodes" && !CACHE_ONLY.contains(&key.as_str()) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .chain(alias) + .collect(); + Self { + fields, + nulls: Nulls::KeptUnlessNamedByCache, + credential_provider: None, + } + } + + /// `RedisCache.redis_kwargs`, the dict every Python surface ends in. + #[cfg(test)] + pub(crate) fn redis_cache_kwargs(fields: BTreeMap) -> Self { + Self { + fields, + nulls: Nulls::Kept, + credential_provider: None, + } + } + + /// The handle the host holds for the `credential_provider` object in these settings. + #[cfg(test)] + pub(crate) fn with_credential_provider(self, handle: ProviderHandle) -> Self { + Self { + credential_provider: Some(handle), + ..self + } + } + + /// No configured settings, so only `REDIS_*` and the other environment fallbacks apply. + pub(crate) fn environment_only() -> Self { + Self { + fields: BTreeMap::new(), + nulls: Nulls::Dropped, + credential_provider: None, + } + } + + /// `_merge_redis_kwargs` in `litellm/_redis.py`: the kwargs over their `REDIS_*` values with every + /// `get_secret` fallback filled. + pub(crate) fn merged(&self, secrets: &SecretSnapshot) -> Result { + layers::merged(self, secrets) + } + + /// The environment is only for what Python reads straight from `os.environ`, `REDIS_USERNAME` for Azure. + pub(crate) fn resolve( + &self, + secrets: &SecretSnapshot, + environment: &dyn Lookup, + ) -> Result { + Ok(litellm_redis_resolve::interpret( + &self.merged(secrets)?, + self.credential_provider, + environment.get("REDIS_USERNAME").as_deref(), + )?) + } +} + +pub(super) fn text(value: &Value) -> Option { + match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} diff --git a/litellm-rust/crates/gateway/src/redis/layers.rs b/litellm-rust/crates/gateway/src/redis/layers.rs new file mode 100644 index 00000000000..1cf68b2cd4f --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/layers.rs @@ -0,0 +1,226 @@ +use litellm_config::Value; +use litellm_core_utils::settings::Lookup; +use litellm_redis_resolve::{Entry, Kwargs, SETTINGS, is_present, setting}; + +use super::configured::{ConfiguredRedis, Nulls}; +use crate::Error; + +const NAMED_BY_CACHE: &[&str] = &[ + "host", + "port", + "password", + "socket_timeout", + "startup_nodes", + "namespace", + "gcp_service_account", + "gcp_ssl_ca_certs", +]; + +/// How `_merge_redis_kwargs` reads a fallback: `get_secret` turns `true` and `false` into a bool, and +/// `get_secret_str` treats that bool as unset. +#[derive(Clone, Copy)] +enum Read { + Secret, + SecretStr, +} + +const TARGETS: &[(&str, Read)] = &[ + ("startup_nodes", Read::Secret), + ("sentinel_nodes", Read::Secret), + ("sentinel_password", Read::SecretStr), + ("service_name", Read::Secret), +]; + +const IDENTITIES: &[(&str, Read)] = &[ + ("gcp_service_account", Read::SecretStr), + ("gcp_ssl_ca_certs", Read::SecretStr), + ("azure_redis_ad_token", Read::Secret), +]; + +const AZURE: &[(&str, Read)] = &[ + ("azure_client_id", Read::SecretStr), + ("azure_tenant_id", Read::SecretStr), + ("azure_client_secret", Read::SecretStr), +]; + +const AWS: &[(&str, Read)] = &[("aws_iam_region", Read::SecretStr)]; + +/// The configured settings over their `REDIS_*` values, `{**environment_kwargs, **env_overrides}` +/// in `_merge_redis_kwargs`, with the `get_secret` fallbacks it reads filled in. A null left in the result is +/// an explicit `None` that hid its environment value. +pub(super) fn merged(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result { + let block: Kwargs = configured + .fields + .iter() + .map(|(key, value)| (key, environment_reference(value, lookup))) + .filter(|(key, value)| !value.is_null() || keeps_null(configured.nulls, key)) + .map(|(key, value)| { + let value = serde_json::to_value(&value).map_err(|_| Error::RedisNotPlainData { + setting: key.clone(), + })?; + let entry = Entry { + value, + source: key.clone(), + }; + Ok((key.clone(), entry)) + }) + .collect::>()?; + let names_a_target = ["host", "startup_nodes", "sentinel_nodes"] + .iter() + .any(|key| is_set(&block, key)); + let drop_environment_url = names_a_target && !is_set(&block, "url"); + let layered: Kwargs = environment(lookup) + .filter(|(key, _)| !(key == "url" && drop_environment_url)) + .chain(block) + .collect(); + let fallbacks = fallbacks(&layered, configured.credential_provider.is_some(), lookup); + Ok(layered.into_iter().chain(fallbacks).collect()) +} + +// The fallbacks are read under the same conditions as _merge_redis_kwargs, so a secret Python never asks +// for is never filled in +fn fallbacks(layered: &Kwargs, has_provider: bool, lookup: &dyn Lookup) -> Kwargs { + let fill = |kwargs: &[(&str, Read)]| -> Kwargs { + kwargs + .iter() + .filter_map(|(kwarg, read)| fallback(layered, kwarg, *read, lookup)) + .collect() + }; + let targets = fill(TARGETS); + if has_provider || is_set(layered, "credential_provider") { + return targets; + } + let identities = fill(IDENTITIES); + let value = |kwarg: &str| { + layered + .get(kwarg) + .filter(|entry| is_present(&entry.value)) + .or_else(|| identities.get(kwarg)) + .map(|entry| &entry.value) + }; + let gcp = value("gcp_service_account").is_some(); + let azure = value("azure_redis_ad_token") + .is_some_and(|token| text(token).is_some_and(|token| token.eq_ignore_ascii_case("true"))); + let aws = layered + .get("aws_iam_auth") + .is_some_and(|entry| coerces_to_true(&entry.value)); + let azure_settings = match azure && !gcp { + true => fill(AZURE), + false => Kwargs::new(), + }; + let aws_settings = match aws && !gcp && !azure { + true => fill(AWS), + false => Kwargs::new(), + }; + targets + .into_iter() + .chain(identities) + .chain(azure_settings) + .chain(aws_settings) + .collect() +} + +fn fallback( + layered: &Kwargs, + kwarg: &str, + read: Read, + lookup: &dyn Lookup, +) -> Option<(String, Entry)> { + if layered + .get(kwarg) + .is_some_and(|entry| is_present(&entry.value)) + { + return None; + } + let names = setting(kwarg).map_or(&[][..], |setting| setting.fallback); + names.iter().find_map(|name| { + let value = lookup + .get(name) + .filter(|value| !value.is_empty()) + .map(secret)?; + match (read, &value) { + (Read::SecretStr, serde_json::Value::Bool(_)) => None, + _ => { + let entry = Entry { + value, + source: (*name).to_owned(), + }; + Some((kwarg.to_owned(), entry)) + } + } + }) +} + +// get_secret reads an environment value of true or false, in any case and padding, as a bool +fn secret_flag(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "true" => Some(true), + "false" => Some(false), + _ => None, + } +} + +fn secret(value: String) -> serde_json::Value { + secret_flag(&value).map_or(serde_json::Value::String(value), serde_json::Value::Bool) +} + +// _coerces_to_true: text is read like _str_to_bool, anything else by its truthiness +fn coerces_to_true(value: &serde_json::Value) -> bool { + match value { + serde_json::Value::String(text) => { + matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes") + } + serde_json::Value::Null => false, + serde_json::Value::Bool(flag) => *flag, + serde_json::Value::Number(number) => number.as_f64().is_some_and(|number| number != 0.0), + serde_json::Value::Array(items) => !items.is_empty(), + serde_json::Value::Object(entries) => !entries.is_empty(), + } +} + +fn text(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::String(text) => Some(text.clone()), + serde_json::Value::Number(number) => Some(number.to_string()), + serde_json::Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} + +fn keeps_null(nulls: Nulls, key: &str) -> bool { + match nulls { + Nulls::Dropped => false, + Nulls::KeptUnlessNamedByCache => !NAMED_BY_CACHE.contains(&key), + #[cfg(test)] + Nulls::Kept => true, + } +} + +fn environment(lookup: &dyn Lookup) -> impl Iterator + '_ { + SETTINGS.iter().filter_map(|setting| { + let source = setting.environment_name()?; + let value = lookup.get(&source)?; + let entry = Entry { + value: secret(value), + source, + }; + Some((setting.kwarg.to_owned(), entry)) + }) +} + +// load_config resolves os.environ/ to None when unset, so a missing reference is a null like any other +fn environment_reference(value: &Value, lookup: &dyn Lookup) -> Value { + match value { + Value::String(text) => match text.strip_prefix("os.environ/") { + Some(name) => lookup.get(name).map_or(Value::Null, |value| { + secret_flag(&value).map_or(Value::String(value), Value::Bool) + }), + None => value.clone(), + }, + _ => value.clone(), + } +} + +fn is_set(kwargs: &Kwargs, key: &str) -> bool { + kwargs.get(key).is_some_and(|entry| !entry.value.is_null()) +} diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis/mod.rs similarity index 60% rename from litellm-rust/crates/gateway/src/redis.rs rename to litellm-rust/crates/gateway/src/redis/mod.rs index 197f40e7988..91428a825f7 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis/mod.rs @@ -1,19 +1,31 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; -use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State}; -use litellm_config::{ - CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value, - environment_has_target, -}; -use litellm_redis_client::{ - ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient, - RedisConfig, -}; +use litellm_circuit_breaker::Config as BreakerConfig; +use litellm_config::{CacheSettings, Config, RedisSettings, Value}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{RedisClient, RedisEvents}; +use litellm_redis_connect::{HostServices, RedisHosts, SharedEnvironment, TracingEvents}; +use litellm_redis_resolve::Expected; +use litellm_redis_types::Resolution; use crate::Error; +mod configured; +mod layers; +mod secrets; +#[cfg(test)] +mod surface_tests; + +use configured::ConfiguredRedis; +use secrets::SecretSnapshot; + +pub(crate) struct Resources { + pub(crate) http: litellm_http::Client, + pub(crate) lookup: SharedEnvironment, +} + const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2); -type Params = ResolvedRedis; +type Params = Resolution; pub struct RedisClients { pub coordination: Option, @@ -34,12 +46,15 @@ enum Target { Environment(Params), } -pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result { +pub(crate) fn plan(config: &Config, lookup: &dyn Lookup) -> Result { let response_cache = redis_cache_params(config) - .map(|settings| settings.resolve_cache(lookup)) + .map(|settings| resolve(&ConfiguredRedis::cache_params(settings), lookup)) .transpose()?; let coordination = match &config.general_settings.coordination_redis { - Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)), + Some(settings) => Some(Target::Configured(resolve( + &ConfiguredRedis::coordination(settings)?, + lookup, + )?)), None => match (&response_cache, config.litellm_settings.cache == Some(true)) { (Some(cache), _) => Some(Target::Configured(cache.clone())), (None, true) => environment(lookup)?.map(Target::Configured), @@ -57,12 +72,30 @@ pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result { }) } -pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result { +pub(crate) async fn connect( + config: &Config, + lookup: &dyn Lookup, + resources: &Resources, +) -> Result { let plan = plan(config, lookup)?; let breaker = breaker_config(lookup)?; + let events: Arc = Arc::new(TracingEvents::new(lookup)); + let hosts = RedisHosts::default(); + let client = |params: Params| -> Result { + Ok(hosts.client( + params, + HostServices { + events: Some(events.clone()), + http: resources.http.clone(), + environment: resources.lookup.clone(), + credential_provider: None, + }, + breaker, + )?) + }; let coordination = match plan.coordination { - Some(Target::Configured(params)) => Some(client(params, breaker)?), - Some(Target::Environment(params)) => match client(params, breaker) { + Some(Target::Configured(params)) => Some(client(params)?), + Some(Target::Environment(params)) => match client(params) { Ok(client) => reachable(client).await, Err(_) => { tracing::warn!( @@ -75,13 +108,10 @@ pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result Some(coordination.clone()), - (None, Some(params)) => Some(client(params, breaker)?), + (None, Some(params)) => Some(client(params)?), (None, None) => None, }; - let response_cache = plan - .response_cache - .map(|params| client(params, breaker)) - .transpose()?; + let response_cache = plan.response_cache.map(client).transpose()?; Ok(RedisClients { coordination, transaction_buffer, @@ -98,17 +128,38 @@ fn redis_cache_params(config: &Config) -> Option<&RedisSettings> { } } -fn environment(lookup: Lookup<'_>) -> Result, Error> { +fn resolve(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result { + let secrets = SecretSnapshot::read(configured.secret_names(), lookup); + let resolution = configured.resolve(&secrets, lookup)?; + for ignored in &resolution.ignored { + tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored: {}", + ignored.setting, + ignored.consequence + ); + } + Ok(resolution) +} + +fn environment(lookup: &dyn Lookup) -> Result, Error> { match environment_has_target(lookup) { - true => RedisSettings::default() - .resolve(lookup) - .map(Some) - .map_err(Into::into), + true => Ok(Some(resolve(&ConfiguredRedis::environment_only(), lookup)?)), false => Ok(None), } } -fn optional_environment(lookup: Lookup<'_>) -> Option { +fn environment_has_target(lookup: &dyn Lookup) -> bool { + [ + "REDIS_HOST", + "REDIS_URL", + "REDIS_CLUSTER_NODES", + "REDIS_SENTINEL_NODES", + ] + .iter() + .any(|name| lookup.get(name).is_some()) +} + +fn optional_environment(lookup: &dyn Lookup) -> Option { match environment(lookup) { Ok(params) => params, Err(error) => { @@ -149,122 +200,101 @@ async fn reachable(client: RedisClient) -> Option { } } -fn client(params: Params, breaker: Option) -> Result { - tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured"); - let endpoint = match params.endpoint { - RedisEndpoint::Url(url) => Endpoint::Url(url.expose().to_owned()), - RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint { - host: host.host, - port: host.port, - username: host.username.map(|value| value.expose().to_owned()), - password: host.password.map(|value| value.expose().to_owned()), - tls: host.tls, - db: host.db, - }), - RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { - nodes: cluster - .nodes - .into_iter() - .map(|node| ClusterNode { - host: node.host, - port: node.port, - }) - .collect(), - username: cluster.username.map(|value| value.expose().to_owned()), - password: cluster.password.map(|value| value.expose().to_owned()), - tls: cluster.tls, - }), - }; - Ok(RedisClient::new(RedisConfig { - endpoint, - namespace: params.namespace.map(Namespace::new).unwrap_or_default(), - timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT), - breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))), - })?) +fn parse_seconds(text: &str) -> Option { + text.parse::() + .ok() + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) } -fn log_transition(transition: &litellm_circuit_breaker::Transition) { - match transition.to { - State::Open { .. } => tracing::warn!( - generation = transition.generation, - "redis circuit breaker opened, fast-failing redis calls" - ), - State::HalfOpen => tracing::info!( - generation = transition.generation, - "redis circuit breaker half open, probing redis" - ), - State::Closed { .. } => tracing::info!( - generation = transition.generation, - "redis circuit breaker closed, redis recovered" - ), - } -} - -fn breaker_config(lookup: Lookup<'_>) -> Result, Error> { - let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED") +fn breaker_config(lookup: &dyn Lookup) -> Result, Error> { + let enabled = lookup + .get("REDIS_CIRCUIT_BREAKER_ENABLED") .is_none_or(|value| value.eq_ignore_ascii_case("true")); if !enabled { return Ok(None); } let defaults = BreakerConfig::default(); - let seconds = |name: &'static str, default: Duration| -> Result { - lookup(name) - .map(|value| { - value - .parse::() - .ok() - .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) - .ok_or(Error::RedisInvalid { - setting: name, - reason: "must be a non-negative number of seconds", - }) - }) - .unwrap_or(Ok(default)) - }; - let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD") + Ok(Some(BreakerConfig { + failure_threshold: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", + Expected::Integer, + |text| text.parse().ok(), + )? + .unwrap_or(defaults.failure_threshold), + recovery_timeout: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", + Expected::Seconds, + parse_seconds, + )? + .unwrap_or(defaults.recovery_timeout), + timeout_min_duration: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", + Expected::Seconds, + parse_seconds, + )? + .unwrap_or(defaults.timeout_min_duration), + })) +} + +fn environment_value( + lookup: &dyn Lookup, + name: &str, + expected: Expected, + parse: impl Fn(&str) -> Option, +) -> Result, litellm_redis_resolve::Error> { + lookup + .get(name) .map(|value| { - value.parse().map_err(|_| Error::RedisInvalid { - setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", - reason: "must be an integer", + parse(&value).ok_or_else(|| litellm_redis_resolve::Error::Invalid { + setting: name.to_owned(), + expected, }) }) - .unwrap_or(Ok(defaults.failure_threshold))?; - Ok(Some(BreakerConfig { - failure_threshold, - recovery_timeout: seconds( - "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", - defaults.recovery_timeout, - )?, - timeout_min_duration: seconds( - "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", - defaults.timeout_min_duration, - )?, - })) + .transpose() } #[cfg(test)] mod tests { + use litellm_core_utils::settings::Lookup; use std::{collections::HashMap, time::Duration}; - use litellm_config::{Config, RedisEndpoint, RedisHost}; + use litellm_config::Config; + use litellm_redis_client::{Standalone, StaticAuth, commands}; use litellm_redis_testing::RedisServer; + use litellm_redis_types::{Auth, HealthCheck}; use rstest::rstest; - use super::{Params, Plan, Target, connect, plan}; + use std::sync::Arc; + + use super::{Params, Plan, Resources, Target, plan}; use crate::Error; + async fn connect(config: &Config, lookup: &dyn Lookup) -> Result { + let resources = Resources { + http: litellm_http::Client::for_test(), + lookup: Arc::new(|_: &str| None), + }; + super::connect(config, lookup, &resources).await + } + fn host(name: &str, port: u16) -> Params { Params { - endpoint: RedisEndpoint::Host(RedisHost { + target: litellm_redis_client::Target::Standalone(Standalone { host: name.to_owned(), port, - username: None, - password: None, - tls: false, + tls: None, db: 0, }), + auth: Auth::Static(StaticAuth::default()), namespace: None, timeout: None, + connect_timeout: None, + health_check: HealthCheck::Unset, + keepalive: true, + ignored: Vec::new(), } } @@ -273,7 +303,7 @@ mod tests { .iter() .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) .collect(); - plan(&Config::from_yaml(yaml).unwrap(), &|name| { + plan(&Config::from_yaml(yaml).unwrap(), &|name: &str| { environment.get(name).cloned() }) } @@ -350,36 +380,44 @@ mod tests { port_of(&server) )) .unwrap(); - let clients = connect(&config, &|_| None).await.unwrap(); + let clients = connect(&config, &|_: &str| None).await.unwrap(); let coordination = clients.coordination.unwrap(); let buffer = clients.transaction_buffer.unwrap(); let cache = clients.response_cache.unwrap(); - let unnamespaced = connect(&coordination_config(&server), &|_| None) + let unnamespaced = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap() .coordination .unwrap(); coordination - .set("rate-limit", "1", Duration::from_secs(60)) + .send(commands::set("rate-limit", "1", Duration::from_secs(60))) .await .unwrap(); buffer - .set("spend", "2", Duration::from_secs(60)) + .send(commands::set("spend", "2", Duration::from_secs(60))) .await .unwrap(); assert_eq!( - cache.get::("rate-limit").await.unwrap().as_deref(), + cache + .send(commands::get::("rate-limit")) + .await + .unwrap() + .as_deref(), Some("1") ); assert_eq!( - cache.get::("spend").await.unwrap().as_deref(), + cache + .send(commands::get::("spend")) + .await + .unwrap() + .as_deref(), Some("2") ); assert_eq!( unnamespaced - .get::("responses:rate-limit") + .send(commands::get::("responses:rate-limit")) .await .unwrap() .as_deref(), @@ -387,7 +425,7 @@ mod tests { ); assert_eq!( unnamespaced - .get::("responses:spend") + .send(commands::get::("responses:spend")) .await .unwrap() .as_deref(), @@ -395,16 +433,21 @@ mod tests { ); assert!( unnamespaced - .get::("rate-limit") + .send(commands::get::("rate-limit")) + .await + .unwrap() + .is_none() + ); + assert!( + unnamespaced + .send(commands::get::("spend")) .await .unwrap() .is_none() ); - assert!(unnamespaced.get::("spend").await.unwrap().is_none()); } #[rstest] - #[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])] #[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])] fn redis_migration_regression_malformed_optional_environment_is_absent( #[case] environment: &[(&str, &str)], @@ -413,54 +456,45 @@ mod tests { } #[rstest] - fn redis_migration_regression_explicit_malformed_setting_still_fails() { + #[case::optional_environment("model_list: []", &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])] + #[case::transaction_buffer_environment( + "general_settings:\n use_redis_transaction_buffer: true", + &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")] + )] + #[case::response_cache( + "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: env\n port: not-a-port", + &[] + )] + fn a_malformed_port_outside_the_coordination_block_falls_back_to_6379_like_redis_py( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + ) { + let plan = resolved(yaml, environment).unwrap(); + + let target = match (plan.coordination, plan.transaction_buffer_fallback) { + (Some(Target::Configured(params) | Target::Environment(params)), _) + | (None, Some(params)) => params, + (None, None) => panic!("expected a Redis target"), + }; + assert_eq!(target.target, host("env", 6379).target); + } + + #[rstest] + #[case::port("port: not-a-port", "port")] + #[case::ssl("ssl: maybe", "ssl")] + fn a_malformed_coordination_block_fails_like_the_pydantic_model( + #[case] setting: &str, + #[case] expected: &str, + ) { let error = resolved( - "general_settings:\n coordination_redis:\n host: coord\n port: not-a-port", + &format!("general_settings:\n coordination_redis:\n host: coord\n {setting}"), &[], ) .unwrap_err(); - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); - } - - #[rstest] - fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() { - let error = resolved( - "general_settings:\n use_redis_transaction_buffer: true", - &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")], - ) - .unwrap_err(); - - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); - } - - #[rstest] - fn redis_migration_regression_invalid_response_cache_still_fails() { - let error = resolved( - "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port", - &[], - ) - .unwrap_err(); - - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); + assert!( + matches!(error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == expected) + ); } #[rstest] @@ -468,7 +502,7 @@ mod tests { async fn redis_migration_regression_invalid_optional_environment_url_is_absent() { let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| { + let clients = connect(&config, &|name: &str| { (name == "REDIS_URL").then(|| "not a redis URL".to_owned()) }) .await @@ -506,7 +540,7 @@ mod tests { let server = RedisServer::start().await.unwrap(); server.stop().await.unwrap(); - let clients = connect(&coordination_config(&server), &|_| None) + let clients = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap(); let coordination = clients.coordination.unwrap(); @@ -521,12 +555,12 @@ mod tests { #[tokio::test] async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() { let server = RedisServer::start().await.unwrap(); - let clients = connect(&coordination_config(&server), &|_| None) + let clients = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap(); let coordination = clients.coordination.unwrap(); coordination - .set("written", "1", Duration::from_secs(60)) + .send(commands::set("written", "1", Duration::from_secs(60))) .await .unwrap(); @@ -537,6 +571,30 @@ mod tests { assert!(buffer.ping().await.is_ok()); } + #[rstest] + #[tokio::test] + async fn an_environment_only_tls_redis_is_reached_with_its_ca() { + crate::install_crypto_provider(); + let server = litellm_redis_testing::RedisTlsServer::start(false) + .await + .unwrap(); + let ca = server.files().ca.display().to_string(); + let port = server.port().to_string(); + let config = Config::from_yaml("model_list: []").unwrap(); + + let clients = connect(&config, &|name: &str| match name { + "REDIS_HOST" => Some(server.host().to_owned()), + "REDIS_PORT" => Some(port.clone()), + "REDIS_SSL" => Some("true".to_owned()), + "REDIS_SSL_CA_CERTS" => Some(ca.clone()), + _ => None, + }) + .await + .unwrap(); + + assert!(clients.coordination.is_some()); + } + #[rstest] #[tokio::test] async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() { @@ -545,7 +603,7 @@ mod tests { server.stop().await.unwrap(); let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| match name { + let clients = connect(&config, &|name: &str| match name { "REDIS_HOST" => Some("127.0.0.1".to_owned()), "REDIS_PORT" => Some(port.clone()), _ => None, @@ -564,7 +622,7 @@ mod tests { let port = port_of(&server); let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| match name { + let clients = connect(&config, &|name: &str| match name { "REDIS_HOST" => Some("127.0.0.1".to_owned()), "REDIS_PORT" => Some(port.clone()), _ => None, diff --git a/litellm-rust/crates/gateway/src/redis/secrets.rs b/litellm-rust/crates/gateway/src/redis/secrets.rs new file mode 100644 index 00000000000..a969497f4c3 --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/secrets.rs @@ -0,0 +1,61 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; + +use litellm_config::Value; +use litellm_redis_resolve::SETTINGS; + +use super::configured::ConfiguredRedis; + +#[derive(Clone, Debug, Default)] +pub(crate) struct SecretSnapshot(BTreeMap>); + +impl SecretSnapshot { + pub(crate) fn read(names: BTreeSet, lookup: &dyn Lookup) -> Self { + Self( + names + .into_iter() + .map(|name| { + let value = lookup.get(&name).map(SecretValue::new); + (name, value) + }) + .collect(), + ) + } +} + +impl Lookup for SecretSnapshot { + fn get(&self, name: &str) -> Option { + match self.0.get(name) { + Some(value) => value.as_ref().map(|value| value.expose().to_owned()), + None => { + debug_assert!( + false, + "{name} is missing from ConfiguredRedis::secret_names" + ); + tracing::error!( + name, + "a Redis secret was read without being in the snapshot" + ); + None + } + } + } +} + +impl ConfiguredRedis { + pub(crate) fn secret_names(&self) -> BTreeSet { + let references = self.fields.values().filter_map(|value| match value { + Value::String(text) => text.strip_prefix("os.environ/").map(str::to_owned), + _ => None, + }); + let declared = SETTINGS.iter().flat_map(|setting| { + setting + .environment_name() + .into_iter() + .chain(setting.fallback.iter().map(|name| (*name).to_owned())) + }); + declared.chain(references).collect() + } +} diff --git a/litellm-rust/crates/gateway/src/redis/surface_tests.rs b/litellm-rust/crates/gateway/src/redis/surface_tests.rs new file mode 100644 index 00000000000..9ee881e2e04 --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/surface_tests.rs @@ -0,0 +1,964 @@ +use std::collections::{BTreeMap, HashMap}; + +use litellm_config::{CacheSettings, Config, Value}; +use litellm_redis_resolve::Expected; +use litellm_redis_types::{ + Auth, AwsIamAuth, AzureAdAuth, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions, +}; +use litellm_secret_value::SecretValue; +use rstest::rstest; +use serde_json::json; + +use super::{configured::ConfiguredRedis, secrets::SecretSnapshot}; +use crate::Error; + +fn resolution( + configured: &ConfiguredRedis, + environment: &HashMap, +) -> Result { + let lookup = |name: &str| environment.get(name).cloned(); + configured.resolve( + &SecretSnapshot::read(configured.secret_names(), &lookup), + &lookup, + ) +} + +const PROVIDER: ProviderHandle = ProviderHandle::new(7); + +fn redis_cache_kwargs(kwargs: serde_json::Value) -> ConfiguredRedis { + let has_provider = kwargs + .get("credential_provider") + .is_some_and(|provider| !provider.is_null()); + let configured = ConfiguredRedis::redis_cache_kwargs( + serde_json::from_value::>(kwargs).unwrap(), + ); + match has_provider { + true => configured.with_credential_provider(PROVIDER), + false => configured, + } +} + +fn coordination_resolution(yaml: &str, environment: &[(&str, &str)]) -> Result { + let environment: HashMap = environment + .iter() + .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) + .collect(); + let config = Config::from_yaml(yaml).unwrap(); + let configured = config + .general_settings + .coordination_redis + .unwrap_or_default(); + let configured = ConfiguredRedis::coordination(&configured)?; + resolution(&configured, &environment) +} + +fn resolve(yaml: &str, environment: &[(&str, &str)]) -> Result { + coordination_resolution(yaml, environment) +} + +fn resolve_cache(yaml: &str) -> Result { + resolve_cache_in(yaml, &[]) +} + +fn resolve_cache_in(yaml: &str, environment: &[(&str, &str)]) -> Result { + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + let config = Config::from_yaml(yaml).unwrap(); + let configured = match config.litellm_settings.cache_params.unwrap() { + CacheSettings::Redis(settings) => ConfiguredRedis::cache_params(&settings), + CacheSettings::Other(_) => panic!("expected a Redis cache"), + }; + configured.resolve( + &SecretSnapshot::read(configured.secret_names(), &lookup), + &lookup, + ) +} + +fn cluster_ports(resolved: &Resolution) -> Vec { + match &resolved.target { + Target::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(), + target => panic!("expected a cluster, got {target:?}"), + } +} + +#[rstest] +#[case::host( + "general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5", + &[], + "host:redis:6380:2:true:app:1500" +)] +#[case::environment_references_resolve_the_host_and_port( + "general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT", + &[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")], + "host:referenced:6400:0:false::none" +)] +#[case::a_null_host_falls_back_to_the_environment( + "general_settings:\n coordination_redis:\n host: null", + &[("REDIS_HOST", "environment")], + "host:environment:6379:0:false::none" +)] +fn resolves_redis_through_the_parsed_config( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: &str, +) { + let resolved = resolve(yaml, environment).unwrap(); + let description = match resolved.target { + Target::Unix(socket) => format!("unix:{}", socket.path.display()), + Target::Url(url) => format!("url:{}", url.url.expose()), + Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name), + Target::Standalone(host) => format!( + "host:{}:{}:{}:{}:{}:{}", + host.host, + host.port, + host.db, + host.tls.is_some(), + resolved.namespace.unwrap_or_default(), + resolved + .timeout + .map(|timeout| timeout.as_millis().to_string()) + .unwrap_or_else(|| "none".to_owned()) + ), + Target::Cluster(cluster) => format!( + "cluster:{}", + cluster + .nodes + .iter() + .map(|node| format!("{}:{}", node.host, node.port)) + .collect::>() + .join(",") + ), + }; + + assert_eq!(description, expected); +} + +#[rstest] +#[case::tls_ciphers( + "general_settings:\n coordination_redis:\n host: redis\n ssl_ciphers: HIGH", + &[], + "ssl_ciphers" +)] +#[case::ssl_ciphers_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_SSL_CIPHERS", "HIGH")], + "REDIS_SSL_CIPHERS" +)] +#[case::connection_pool( + "general_settings:\n coordination_redis:\n host: redis\n connection_pool: shared", + &[], + "connection_pool" +)] +fn unsupported_settings_resolve_and_are_named_by_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: &str, +) { + let resolution = coordination_resolution(yaml, environment).unwrap(); + + assert_eq!( + resolution.target, + Target::Standalone(litellm_redis_types::Standalone { + host: "redis".into(), + port: 6379, + db: 0, + tls: None, + }) + ); + let [ignored] = resolution.ignored.as_slice() else { + panic!("expected one ignored setting, got {:?}", resolution.ignored); + }; + assert_eq!(ignored.setting, expected); + assert!(!ignored.consequence.is_empty()); +} + +#[rstest] +#[case::kwarg( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_validate_ocsp: true", + &[], + "ssl_validate_ocsp" +)] +#[case::stapled_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_SSL", "true"), ("REDIS_SSL_VALIDATE_OCSP_STAPLED", "yes")], + "REDIS_SSL_VALIDATE_OCSP_STAPLED" +)] +#[case::rediss_url_query( + "general_settings:\n coordination_redis:\n url: rediss://cache:6380/0?ssl_validate_ocsp=true", + &[], + "url ?ssl_validate_ocsp" +)] +fn an_ocsp_check_fails_instead_of_connecting_without_it( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] setting: &str, +) { + let error = coordination_resolution(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Unenforceable { setting: named, .. }) if named == setting), + "{error:?}" + ); +} + +#[rstest] +#[case::kwarg(" host: redis\n ssl: true\n ssl_validate_ocsp: false")] +#[case::rediss_url_query(" url: rediss://cache:6380/0?ssl_validate_ocsp=false")] +fn an_ocsp_flag_set_to_false_asks_for_nothing(#[case] block: &str) { + let yaml = format!("general_settings:\n coordination_redis:\n{block}"); + + let resolution = coordination_resolution(&yaml, &[]).unwrap(); + + assert!(resolution.ignored.is_empty(), "{:?}", resolution.ignored); +} + +fn static_auth(username: Option<&str>, password: Option<&str>) -> Auth { + Auth::Static(StaticAuth { + username: username.map(litellm_secret_value::SecretValue::new), + password: password.map(litellm_secret_value::SecretValue::new), + }) +} + +#[rstest] +fn an_explicit_host_keeps_environment_credentials() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: configured", + &[ + ("REDIS_PASSWORD", "environment-secret"), + ("REDIS_SSL", "True"), + ], + ) + .unwrap(); + + assert_eq!(resolved.auth, static_auth(None, Some("environment-secret"))); + assert!(matches!(resolved.target, Target::Standalone(host) if host.tls.is_some())); +} + +fn azure(username: Option<&str>) -> Auth { + Auth::AzureAd(AzureAdAuth { + client_id: Some("client".to_owned()), + tenant_id: None, + client_secret: None, + username: username.map(str::to_owned), + }) +} + +fn aws(serverless: bool) -> Auth { + Auth::AwsIam(AwsIamAuth { + user_name: "iam-user".to_owned(), + cache_name: "cache-1".to_owned(), + region: "us-east-1".to_owned(), + serverless, + }) +} + +#[rstest] +#[case::static_credentials( + "general_settings:\n coordination_redis:\n host: redis\n username: user\n password: pass", + &[], + static_auth(Some("user"), Some("pass")) +)] +#[case::gcp_beats_azure_and_aws( + "general_settings:\n coordination_redis:\n host: redis\n password: pass\n azure_redis_ad_token: \"True\"\n aws_iam_auth: true", + &[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com")], + Auth::GcpIam { service_account: "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com".to_owned() } +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + "general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client\n aws_iam_auth: true", + &[("REDIS_USERNAME", "environment-user")], + azure(Some("environment-user")) +)] +#[case::azure_without_an_environment_username_sends_only_the_token( + "general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client", + &[], + azure(None) +)] +#[case::azure_needs_the_literal_true( + "general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: \"1\"", + &[], + static_auth(None, None) +)] +#[case::aws_with_a_region_from_the_environment( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n password: dropped\n aws_iam_auth: \"yes\"\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1", + &[("AWS_DEFAULT_REGION", "us-east-1")], + aws(false) +)] +#[case::aws_serverless_over_a_rediss_url( + "general_settings:\n coordination_redis:\n url: rediss://cache-1:6379\n aws_iam_auth: true\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1\n aws_iam_region: us-east-1\n aws_iam_serverless: true", + &[], + aws(true) +)] +fn auth_follows_python_precedence( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: Auth, +) { + assert_eq!(resolve(yaml, environment).unwrap().auth, expected); +} + +#[rstest] +#[case::coordination( + "general_settings:\n coordination_redis:\n host: redis\n credential_provider: host-object", + &[] +)] +#[case::without_a_target( + "general_settings:\n coordination_redis:\n credential_provider: host-object", + &[("REDIS_HOST", "redis")] +)] +fn a_credential_provider_without_a_host_handle_fails( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!( + error, + Error::Redis(litellm_redis_resolve::Error::ProviderWithoutHandle) + ), + "{error}" + ); +} + +#[rstest] +#[case::without_tls( + "general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r", + "requires TLS" +)] +#[case::a_plain_url_ignores_ssl( + "general_settings:\n coordination_redis:\n url: redis://cache:6379\n ssl: true\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r", + "requires TLS" +)] +#[case::missing_settings( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n aws_iam_auth: true\n aws_iam_cache_name: c", + "requires: aws_iam_user_name, aws_iam_region" +)] +fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) { + let error = resolve(yaml, &[]).unwrap_err(); + + assert!(error.to_string().contains(message), "{error}"); +} + +#[rstest] +#[case::defaults_verify_the_certificate_but_not_the_hostname( + " ssl: true", + &[], + Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false }) +)] +#[case::every_option( + " ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_password: key-secret\n ssl_cert_reqs: none\n ssl_check_hostname: true", + &[], + Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), key_password: Some(SecretValue::new("key-secret")), verify_certificate: false, check_hostname: true }) +)] +#[case::options_from_the_environment( + " ssl: true", + &[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional"), ("REDIS_SSL_PASSWORD", "env-key-secret")], + Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: Some(SecretValue::new("env-key-secret")), verify_certificate: true, check_hostname: false }) +)] +#[case::gcp_ca_certs_replace_the_configured_ones( + " ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem", + &[], + Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false }) +)] +#[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)] +fn tls_options_resolve_like_redis_py( + #[case] block: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: Option, +) { + let yaml = format!("general_settings:\n coordination_redis:\n host: redis\n{block}"); + + let resolved = resolve(&yaml, environment).unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(host) if host.tls == expected)); +} + +#[rstest] +fn an_unknown_cert_requirement_fails() { + let error = resolve( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_cert_reqs: strict", + &[], + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::CertRequirement }) if setting == "ssl_cert_reqs" + )); +} + +#[rstest] +#[case::from_config( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[sentinel-a, 26379], [sentinel-b, \"26380\"]]\n service_name: primary\n sentinel_password: sentinel-secret\n db: 2", + &[] +)] +#[case::from_the_environment( + "model_list: []", + &[ + ("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379], ["sentinel-b", "26380"]]"#), + ("REDIS_SERVICE_NAME", "primary"), + ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), + ("REDIS_DB", "2"), + ] +)] +fn sentinel_nodes_resolve_to_a_sentinel_endpoint( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], +) { + let Target::Sentinel(sentinel) = resolve(yaml, environment).unwrap().target else { + panic!("expected a sentinel endpoint"); + }; + + assert_eq!( + sentinel.monitors, + vec![ + Node { + host: "sentinel-a".to_owned(), + port: 26379 + }, + Node { + host: "sentinel-b".to_owned(), + port: 26380 + }, + ] + ); + assert_eq!(sentinel.service_name, "primary"); + assert_eq!( + sentinel.monitor_auth, + Some(StaticAuth { + username: None, + password: Some(litellm_secret_value::SecretValue::new("sentinel-secret")), + }) + ); + assert_eq!(sentinel.db, 2); +} + +#[rstest] +#[case::a_url_beats_sentinels( + "general_settings:\n coordination_redis:\n url: redis://single:6379\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + "url:redis://single:6379" +)] +fn sentinel_precedence(#[case] yaml: &str, #[case] expected: &str) { + let resolved = resolve(yaml, &[]).unwrap(); + let description = match resolved.target { + Target::Url(url) => format!("url:{}", url.url.expose()), + Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name), + other => panic!("unexpected {other:?}"), + }; + + assert_eq!(description, expected); +} + +#[rstest] +fn sentinels_without_a_service_name_fail() { + let error = resolve( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]", + &[], + ) + .unwrap_err(); + + assert!( + matches!( + error, + Error::Redis(litellm_redis_resolve::Error::NoSentinelService) + ), + "{error:?}" + ); +} + +#[rstest] +#[case::configured_host( + "general_settings:\n coordination_redis:\n host: ignored\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + &[], + "host" +)] +#[case::environment_port( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + &[("REDIS_PORT", "6380")], + "REDIS_PORT" +)] +fn a_host_or_port_cannot_reach_sentinel( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::SentinelWithHost { setting }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::not_pairs( + "general_settings:\n coordination_redis:\n sentinel_nodes: [{host: s, port: 26379}]\n service_name: primary", + &[], + "sentinel_nodes" +)] +#[case::from_the_environment( + "model_list: []", + &[("REDIS_SENTINEL_NODES", "not-json"), ("REDIS_SERVICE_NAME", "primary")], + "REDIS_SENTINEL_NODES" +)] +fn invalid_sentinel_nodes_name_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::SentinelNodes }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::from_config("general_settings:\n coordination_redis:\n startup_nodes: [{host: a}]", &[], "startup_nodes")] +#[case::from_the_environment("model_list: []", &[("REDIS_CLUSTER_NODES", "not-json")], "REDIS_CLUSTER_NODES")] +#[case::out_of_range_port_from_the_environment("model_list: []", &[("REDIS_HOST", "redis"), ("REDIS_PORT", "70000")], "REDIS_PORT")] +fn invalid_values_name_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::falls_back_to_the_socket_timeout(" socket_timeout: 1.5", Some(1500), None)] +#[case::explicit_connect_timeout( + " socket_timeout: 1.5\n socket_connect_timeout: 0.25", + Some(1500), + Some(250) +)] +fn connect_timeout_is_resolved_separately( + #[case] block: &str, + #[case] timeout_ms: Option, + #[case] connect_ms: Option, +) { + let resolved = resolve( + &format!("general_settings:\n coordination_redis:\n host: redis\n{block}"), + &[], + ) + .unwrap(); + + assert_eq!( + resolved.timeout.map(|timeout| timeout.as_millis()), + timeout_ms + ); + assert_eq!( + resolved.connect_timeout.map(|timeout| timeout.as_millis()), + connect_ms + ); +} + +#[rstest] +#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])] +#[case::a_structured_host_with_an_environment_target( + "general_settings:\n coordination_redis:\n host: []", + &[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")] +)] +fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(error, Error::Redis(litellm_redis_resolve::Error::NoTarget)), + "{error:?}" + ); +} + +#[rstest] +fn the_cache_alias_names_cluster_nodes() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000", + ) + .unwrap(); + + assert_eq!(cluster_ports(&resolved), vec![7000]); +} + +#[rstest] +fn canonical_startup_nodes_beat_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101", + ) + .unwrap(); + + assert_eq!(cluster_ports(&resolved), vec![7000]); +} + +#[rstest] +fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100", + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100", + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +fn unsupported_environment_settings_never_carry_the_value() { + let resolution = coordination_resolution( + "model_list: []", + &[ + ("REDIS_HOST", "redis"), + ("REDIS_SSL_CIPHERS", "private-option-value"), + ], + ) + .unwrap(); + let ignored = format!("{:?}", resolution.ignored); + + assert!(ignored.contains("REDIS_SSL_CIPHERS"), "{ignored}"); + assert!(!ignored.contains("private-option-value"), "{ignored}"); +} + +#[rstest] +#[case::negative("-1")] +#[case::infinite("inf")] +fn rejects_bad_socket_timeout(#[case] timeout: &str) { + let yaml = format!( + "general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}" + ); + let error = resolve(&yaml, &[]).unwrap_err(); + + assert!(matches!( + error, + Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::Seconds }) if setting == "socket_timeout" + )); +} + +#[rstest] +fn debug_redacts_redis_secrets_and_url() { + let config = Config::from_yaml( + "general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password\n sentinel_kwargs: {password: credential-passthrough}", + ) + .unwrap(); + let settings = config.general_settings.coordination_redis.unwrap(); + let configured = ConfiguredRedis::coordination(&settings).unwrap(); + let resolved = resolution(&configured, &HashMap::new()).unwrap(); + let debug = format!("{settings:?} {configured:?} {resolved:?}"); + + assert!(!debug.contains("redis://user:password")); + assert!(!debug.contains("credential-user")); + assert!(!debug.contains("credential-password")); + assert!(!debug.contains("credential-passthrough")); + assert!(debug.contains("redacted") || debug.contains("REDACTED")); +} + +#[rstest] +fn legacy_unknown_settings_do_not_change_resolution() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary", + &[], + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +#[case::configured( + "general_settings:\n coordination_redis:\n host: redis\n max_connections: 47", + &[], + "max_connections" +)] +#[case::from_the_environment( + "general_settings:\n coordination_redis:\n host: redis", + &[("REDIS_MAX_CONNECTIONS", "47")], + "REDIS_MAX_CONNECTIONS" +)] +fn ignored_settings_are_named_by_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let resolution = coordination_resolution(yaml, environment).unwrap(); + + let settings: Vec<&str> = resolution + .ignored + .iter() + .map(|ignored| ignored.setting.as_str()) + .collect(); + assert_eq!(settings, vec![source]); +} + +fn snapshot(values: &[(&str, &str)]) -> SecretSnapshot { + let values: HashMap<&str, &str> = values.iter().copied().collect(); + SecretSnapshot::read( + ConfiguredRedis::environment_only().secret_names(), + &|name: &str| values.get(name).map(|value| (*value).to_owned()), + ) +} + +#[rstest] +fn secret_manager_values_reach_the_resolver_through_the_snapshot() { + let settings = redis_cache_kwargs(json!({"host": "os.environ/COORD_HOST"})); + let secrets = SecretSnapshot::read(settings.secret_names(), &|name: &str| match name { + "COORD_HOST" => Some("from-the-manager".to_owned()), + "REDIS_PASSWORD" => Some("manager-secret".to_owned()), + _ => None, + }); + + let resolved = settings.resolve(&secrets, &|_: &str| None).unwrap(); + + assert!( + matches!(&resolved.target, Target::Standalone(host) if host.host == "from-the-manager") + ); + assert_eq!(resolved.auth, static_auth(None, Some("manager-secret"))); +} + +#[rstest] +#[case::from_the_environment(&[], &[("REDIS_USERNAME", "environment-user")], Some("environment-user"))] +#[case::never_from_the_secret_manager(&[("REDIS_USERNAME", "manager-user")], &[], None)] +fn the_azure_username_comes_from_the_process_environment_like_python( + #[case] secrets: &[(&str, &str)], + #[case] environment: &[(&str, &str)], + #[case] expected: Option<&str>, +) { + let settings = redis_cache_kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"})); + let environment: HashMap<&str, &str> = environment.iter().copied().collect(); + + let resolved = settings + .resolve(&snapshot(secrets), &|name: &str| { + environment.get(name).map(|value| (*value).to_owned()) + }) + .unwrap(); + + assert!( + matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected), + "{:?}", + resolved.auth + ); +} + +#[rstest] +#[case::coordination_drops_the_null( + "general_settings:\n coordination_redis:\n host: redis\n ssl: null", + false, + true +)] +#[case::cache_passes_an_unnamed_null_on( + "litellm_settings:\n cache_params:\n type: redis\n host: redis\n ssl: null", + true, + false +)] +#[case::cache_drops_a_null_it_names( + "litellm_settings:\n cache_params:\n type: redis\n host: null\n ssl: true", + true, + true +)] +fn a_null_overrides_the_environment_only_where_python_passes_it_on( + #[case] yaml: &str, + #[case] cache: bool, + #[case] tls: bool, +) { + let environment = [("REDIS_HOST", "environment"), ("REDIS_SSL", "true")]; + let resolved = match cache { + true => resolve_cache_in(yaml, &environment), + false => resolve(yaml, &environment), + } + .unwrap(); + + let Target::Standalone(host) = resolved.target else { + panic!("expected a standalone target, got {:?}", resolved.target); + }; + assert_eq!(host.tls.is_some(), tls); +} + +#[rstest] +#[case::empty_cluster_nodes_use_the_environment( + json!({"startup_nodes": []}), + &[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)], + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], +)] +#[case::empty_sentinel_nodes_use_the_environment( + json!({"sentinel_nodes": [], "service_name": "primary"}), + &[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)], + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES")], +)] +#[case::empty_service_name_uses_the_environment( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}), + &[("REDIS_SERVICE_NAME", "primary")], + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[("service_name", "REDIS_SERVICE_NAME")], +)] +#[case::environment_host_port_and_password( + json!({}), + &[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")], + json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}), + &[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")], +)] +#[case::padded_environment_numbers_parse_like_python_int( + json!({}), + &[("REDIS_HOST", "redis"), ("REDIS_PORT", "6_380"), ("REDIS_DB", "3\n")], + json!({"port": "6_380", "host": "redis", "db": "3\n"}), + &[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")], +)] +#[case::socket_timeout_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SOCKET_TIMEOUT", "0.75")], + json!({"socket_timeout": "0.75", "host": "redis"}), + &[("socket_timeout", "REDIS_SOCKET_TIMEOUT")], +)] +#[case::environment_url( + json!({}), + &[("REDIS_URL", "redis://environment:6379/1")], + json!({"url": "redis://environment:6379/1"}), + &[("url", "REDIS_URL")], +)] +#[case::environment_url_beats_the_environment_host( + json!({}), + &[("REDIS_HOST", "ignored"), ("REDIS_URL", "redis://from-url:6379")], + json!({"host": "ignored", "url": "redis://from-url:6379"}), + &[("host", "REDIS_HOST"), ("url", "REDIS_URL")], +)] +#[case::an_explicit_host_drops_the_environment_url( + json!({"host": "configured"}), + &[("REDIS_URL", "redis://environment:6379")], + json!({"host": "configured"}), + &[], +)] +#[case::environment_cluster_nodes( + json!({}), + &[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)], + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], +)] +#[case::sentinel_from_the_environment( + json!({}), + &[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")], + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")], +)] +#[case::an_explicit_credential_provider_beats_cloud_auth( + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[("REDIS_AZURE_AD_TOKEN", "true")], + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[], +)] +#[case::gcp_beats_azure_and_aws( + json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")], + json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")], +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[("REDIS_USERNAME", "environment-user")], + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[], +)] +#[case::an_empty_environment_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}), + &[("REDIS_USERNAME", "")], + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}), + &[("username", "REDIS_USERNAME")], +)] +#[case::aws_iam_over_tls_with_the_default_region( + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}), + &[("AWS_DEFAULT_REGION", "us-east-1")], + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}), + &[("aws_iam_region", "AWS_DEFAULT_REGION")], +)] +#[case::aws_region_beats_the_default_region( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}), + &[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")], + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}), + &[("aws_iam_region", "AWS_REGION")], +)] +#[case::a_null_cert_requirement_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[("REDIS_SSL_CERT_REQS", "required")], + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[], +)] +#[case::a_cert_requirement_spelled_like_python_none_fails( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CERT_REQS", "None")], + json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], +)] +#[case::tls_options_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")], + json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], +)] +#[case::environment_references_resolve( + json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}), + &[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")], + json!({"host": "referenced", "password": "referenced-secret"}), + &[], +)] +#[case::malformed_port_from_the_environment( + json!({}), + &[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")], + json!({"port": "abc", "host": "redis"}), + &[("host", "REDIS_HOST"), ("port", "REDIS_PORT")], +)] +#[case::retry_on_error_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_RETRY_ON_ERROR", "true")], + json!({"retry_on_error": true, "host": "redis"}), + &[("retry_on_error", "REDIS_RETRY_ON_ERROR")], +)] +#[case::ssl_ca_path_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")], + json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")], +)] +fn merges_like_python_redis( + #[case] kwargs: serde_json::Value, + #[case] environment: &[(&str, &str)], + #[case] merged: serde_json::Value, + #[case] sources: &[(&str, &str)], +) { + let configured = redis_cache_kwargs(kwargs); + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + + let kwargs = configured + .merged(&SecretSnapshot::read(configured.secret_names(), &lookup)) + .unwrap(); + + let values: serde_json::Map = kwargs + .iter() + .map(|(kwarg, entry)| (kwarg.clone(), entry.value.clone())) + .collect(); + let renamed: Vec<(&str, &str)> = kwargs + .iter() + .filter(|(kwarg, entry)| *kwarg != &entry.source) + .map(|(kwarg, entry)| (kwarg.as_str(), entry.source.as_str())) + .collect(); + assert_eq!(serde_json::Value::Object(values), merged); + assert_eq!(renamed, sources); +} diff --git a/litellm-rust/crates/gateway/src/secrets.rs b/litellm-rust/crates/gateway/src/secrets.rs index bcd293fe1e9..aafb0abf2e5 100644 --- a/litellm-rust/crates/gateway/src/secrets.rs +++ b/litellm-rust/crates/gateway/src/secrets.rs @@ -1,8 +1,8 @@ use std::{collections::BTreeMap, sync::Arc}; use futures_util::future::BoxFuture; -use litellm_auth_types::SecretValue; use litellm_config::{Object, Value}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use crate::Error; diff --git a/litellm-rust/crates/gateway/src/settings.rs b/litellm-rust/crates/gateway/src/settings.rs index 2dd8872065f..c2c8bc30934 100644 --- a/litellm-rust/crates/gateway/src/settings.rs +++ b/litellm-rust/crates/gateway/src/settings.rs @@ -1,9 +1,26 @@ use std::ffi::OsString; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use litellm_tracing::Level; use serde::{Deserialize, de::DeserializeOwned}; +#[derive(Deserialize)] +pub(super) struct LogSettings { + rust_log: Option, +} + +impl LogSettings { + pub fn log_level(&self) -> Level { + log_level(self.rust_log.as_deref()) + } +} + +fn log_level(rust_log: Option<&str>) -> Level { + rust_log + .and_then(|value| value.parse().ok()) + .unwrap_or(Level::INFO) +} + #[derive(Deserialize)] #[serde(default)] pub(super) struct Settings { @@ -26,10 +43,7 @@ impl Default for Settings { impl Settings { pub fn log_level(&self) -> Level { - self.rust_log - .as_deref() - .and_then(|value| value.parse().ok()) - .unwrap_or(Level::INFO) + log_level(self.rust_log.as_deref()) } } diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index 3d12efd730e..a33c9d7c60e 100644 --- a/litellm-rust/crates/gateway/tests/mcp.rs +++ b/litellm-rust/crates/gateway/tests/mcp.rs @@ -424,7 +424,7 @@ impl litellm_secrets::source::SecretSource for NoSecrets { _: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async { Ok(None) }) } diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index c876c1cbcbb..7de7a97c870 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -159,7 +159,7 @@ async fn mounts_ui_without_exposing_credentials_or_authorizing_inference(inferen std::fs::write(assets.path().join("index.html"), "dashboard").unwrap(); let backend = litellm_gateway_auth::UiBackend::new( "admin".into(), - litellm_auth_types::SecretValue::new("ui-password"), + litellm_secret_value::SecretValue::new("ui-password"), ) .unwrap(); let ui = litellm_gateway_ui::router( diff --git a/litellm-rust/crates/host-http/tests/serve.rs b/litellm-rust/crates/host-http/tests/serve.rs index 2ba39ad4cd2..fd1a8940b33 100644 --- a/litellm-rust/crates/host-http/tests/serve.rs +++ b/litellm-rust/crates/host-http/tests/serve.rs @@ -215,9 +215,7 @@ async fn projection_custom_operations_and_hooks_feed_the_http_response(intercept ) .await?; route_hooks - .after_provider_response(RawResponse { - body: wire.url.clone(), - }) + .after_provider_response(RawResponse::new(wire.url.clone())) .await?; Ok(CallOutput::Complete(Bytes::from(wire.url))) }, @@ -501,9 +499,7 @@ async fn hook_rejection_stops_execution_and_is_reported_once( move |_, _, route_hooks, _observations| async move { if event { route_hooks - .after_provider_response(RawResponse { - body: "response".into(), - }) + .after_provider_response(RawResponse::new("response".into())) .await?; } else { route_hooks @@ -637,9 +633,7 @@ async fn unary_calls_use_into_response_after_hooks_and_before_success(intercepto ) .await?; route_hooks - .after_provider_response(RawResponse { - body: wire.url.clone(), - }) + .after_provider_response(RawResponse::new(wire.url.clone())) .await?; Ok(CallOutput::Complete(json!({"url": wire.url}))) }, @@ -689,7 +683,7 @@ async fn unary_failure_preserves_the_error_without_converting( None, |_, _, route_hooks, _observations| async move { route_hooks - .after_provider_response(RawResponse { body: "raw".into() }) + .after_provider_response(RawResponse::new("raw".into())) .await?; Err(TestError::Provider) }, diff --git a/litellm-rust/crates/host-native/tests/driver.rs b/litellm-rust/crates/host-native/tests/driver.rs index c41c1e5a43b..44c30234328 100644 --- a/litellm-rust/crates/host-native/tests/driver.rs +++ b/litellm-rust/crates/host-native/tests/driver.rs @@ -170,9 +170,7 @@ fn dispatching_call() -> TestMachine { ) .await?; route_hooks - .after_provider_response(RawResponse { - body: wire.url.clone(), - }) + .after_provider_response(RawResponse::new(wire.url.clone())) .await?; Ok(CallOutput::Complete(wire.url)) }, @@ -660,9 +658,7 @@ async fn response_interception_waits_and_can_reject_after_observation(#[case] re "input", Some(sender), |_, _, interceptors, observers| async move { - let raw = RawResponse { - body: "provider response".into(), - }; + let raw = RawResponse::new("provider response".into()); observers.unwrap().emit(CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, )); diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index e36b7d70c28..fe6173dd100 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned { ) -> PyResult>; /// What the stream carries at hand-off, as the caller's stream receives it. + /// The unary response after the deployment transform, before success hooks run. + fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> { + Ok(()) + } + fn encode_stream_head( &mut self, py: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index dbf843a5295..052156cbc5d 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -399,8 +399,16 @@ where Ok(step) => step, Err(error) => return self.failure(error, FailureOrigin::Host).await, }; - match self.settle(step).await { - Ok(response) => self.succeeded(response).await, + let response = match self.settle(step).await { + Ok(response) => response, + Err(error) => return self.failure(error, FailureOrigin::Host).await, + }; + let delivered = { + let mut parts = self.parts(); + Python::attach(|py| parts.binding.deliver_response(py, response.bind(py))) + }; + match delivered { + Ok(()) => self.succeeded(response).await, Err(error) => self.failure(error, FailureOrigin::Host).await, } } diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8e37ab02669..a9ebae427d9 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -36,8 +36,8 @@ pub use marshal::{ pub use owned::PythonOwned; pub use runtime::{ ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value, - ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync, - run_sync_value, runtime_started, + ready_future, reserve_process_for_forking, run_async, run_async_value, run_isolated, run_sync, + run_sync_value, runtime_handle, runtime_started, }; pub use services::PythonHostCalls; diff --git a/litellm-rust/crates/host-python/src/runtime.rs b/litellm-rust/crates/host-python/src/runtime.rs index e8872c9340a..2a8fc598fc4 100644 --- a/litellm-rust/crates/host-python/src/runtime.rs +++ b/litellm-rust/crates/host-python/src/runtime.rs @@ -68,6 +68,23 @@ fn runtime() -> PyResult<&'static Runtime> { Ok(pyo3_async_runtimes::tokio::get_runtime()) } +/// Runs `future` on a current-thread runtime built and shut down inside the call, its blocking pool joined before +/// it returns. No thread outlives the call, so unlike every other door this one never claims the process in the fork +/// gate: a process that only used it can still fork. The future must not hold resources tied to another runtime +pub fn run_isolated(future: F) -> std::io::Result { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let output = runtime.block_on(future); + runtime.shutdown_timeout(Duration::from_secs(5)); + Ok(output) +} + +/// The native runtime, for work that must outlive the call that schedules it. +pub fn runtime_handle() -> PyResult { + Ok(runtime()?.handle().clone()) +} + #[expect(clippy::disallowed_methods, reason = "this is the gated door")] fn future_into_py(py: Python<'_>, future: F) -> PyResult> where @@ -813,4 +830,21 @@ asyncio.run(exercise()) .expect("async route contract should hold"); }); } + + #[rstest] + fn an_isolated_run_joins_blocking_work_it_left_running_before_returning() { + let finished = Arc::new(AtomicBool::new(false)); + let flag = Arc::clone(&finished); + let output = run_isolated(async move { + drop(tokio::task::spawn_blocking(move || { + thread::sleep(Duration::from_millis(50)); + flag.store(true, Ordering::SeqCst); + })); + 7 + }) + .unwrap(); + + assert_eq!(output, 7); + assert!(finished.load(Ordering::SeqCst)); + } } diff --git a/litellm-rust/crates/host-python/tests/isolated_run.rs b/litellm-rust/crates/host-python/tests/isolated_run.rs new file mode 100644 index 00000000000..7b8603229a7 --- /dev/null +++ b/litellm-rust/crates/host-python/tests/isolated_run.rs @@ -0,0 +1,9 @@ +use litellm_host_python::{reserve_process_for_forking, run_isolated, runtime_started}; + +#[test] +fn an_isolated_run_leaves_the_process_free_to_reserve_for_forking() { + run_isolated(async { tokio::task::spawn_blocking(|| ()).await.unwrap() }).unwrap(); + + assert!(!runtime_started()); + assert!(reserve_process_for_forking().is_ok()); +} diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index 003fce74309..15e864c0718 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -467,7 +467,7 @@ pub fn success_machine() -> MachineFactory { .before_provider_request(wire(), context()) .await?; host.interceptors - .after_provider_response(RawResponse { body: "raw".into() }) + .after_provider_response(RawResponse::new("raw".into())) .await?; Ok(format!("{projected}|{signed}|{}", wire.url)) }) diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index d3354fa8622..e0ccab4e0ad 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true -litellm-auth.workspace = true litellm-coroutine.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index 3cff781792f..9398ea4bad6 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -22,12 +22,34 @@ pub struct RequestContext { /// Optional-param names that carry credentials and must be redacted when logged. pub secret_fields: Vec, /// The credential the route resolved for the provider call. - pub api_key: Option, + pub api_key: Option, } #[derive(Clone, Debug, PartialEq, Eq)] pub struct RawResponse { pub body: String, + /// The provider's HTTP status and headers, when the route forwards them to logging. + pub head: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ResponseHead { + pub status: u16, + pub headers: Vec<(String, String)>, +} + +impl RawResponse { + pub fn new(body: String) -> Self { + Self { body, head: None } + } + + #[must_use] + pub fn with_head(self, head: ResponseHead) -> Self { + Self { + head: Some(head), + ..self + } + } } #[derive(Clone, Debug, PartialEq, Eq)] @@ -186,7 +208,7 @@ mod tests { .await?; Interceptors::after_provider_response( &channel.interceptors, - RawResponse { body: "raw".into() }, + RawResponse::new("raw".into()), ) .await?; Ok((sent, ())) diff --git a/litellm-rust/crates/host/tests/observation.rs b/litellm-rust/crates/host/tests/observation.rs index 22872047d81..fba85b63bea 100644 --- a/litellm-rust/crates/host/tests/observation.rs +++ b/litellm-rust/crates/host/tests/observation.rs @@ -74,9 +74,7 @@ fn snapshots_do_not_retain_runtime_objects(capacity: NonZeroUsize, #[case] faile #[rstest] fn provider_snapshots_own_the_response_body(capacity: NonZeroUsize) { - let mut raw = RawResponse { - body: "provider response".into(), - }; + let mut raw = RawResponse::new("provider response".into()); let event: CallEvent<(), (), &RawResponse> = CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw: &raw }); let expected = raw.clone(); diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 29c8b988c0b..9505aacf44b 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -16,6 +16,15 @@ pub enum TlsSource { #[error(transparent)] pub struct Error(Arc); +// A streamed body the connection ended early, after the response head arrived; `expected` is +// the declared `content-length`, `None` for a chunked body. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BodyCut { + pub received: u64, + pub expected: Option, + pub reset: bool, +} + #[derive(Debug, thiserror::Error)] pub(crate) enum Repr { #[error("could not read {}", path.display())] @@ -56,10 +65,28 @@ pub(crate) enum Repr { Connect(#[source] reqwest::Error), #[error("upstream network error")] Network(#[source] reqwest::Error), + #[error("upstream network error")] + BodyCut { + #[source] + source: reqwest::Error, + cut: BodyCut, + }, #[error("upstream request timed out")] Timeout(#[source] reqwest::Error), + /// No bytes arrived within the read timeout, which the caller enforces per read; `body` says + /// the response head had already arrived. + #[error("upstream request timed out")] + ReadTimeout { + #[source] + elapsed: Elapsed, + body: bool, + }, #[error("upstream request failed with status {status}: {body}")] - Status { status: u16, body: UpstreamBody }, + Status { + status: u16, + body: UpstreamBody, + headers: Vec<(String, String)>, + }, #[error("invalid WebSocket URL")] WebSocketUrl(#[source] url::ParseError), @@ -136,13 +163,90 @@ impl Repr { impl Error { pub fn upstream(status: u16, body: String) -> Self { + Self::upstream_with_headers(status, body, Vec::new()) + } + + /// An upstream failure that keeps the provider's response headers (`retry-after`, + /// `request-id`) for the caller. + pub fn upstream_with_headers( + status: u16, + body: String, + headers: Vec<(String, String)>, + ) -> Self { Repr::Status { status, body: UpstreamBody(body), + headers, } .into() } + /// The provider's response headers on an upstream failure; empty for every other error. + pub fn upstream_headers(&self) -> &[(String, String)] { + match &*self.0 { + Repr::Status { headers, .. } => headers, + _ => &[], + } + } + + /// A timeout reached before the connection was established, as opposed to while reading. + pub fn is_connect_timeout(&self) -> bool { + matches!(&*self.0, Repr::Timeout(error) if error.is_connect()) + } + + /// A read the caller timed itself (httpx's per-read timeout) ran out. + pub fn read_timeout(elapsed: Elapsed) -> Self { + Repr::ReadTimeout { + elapsed, + body: false, + } + .into() + } + + /// A read of the response body the caller timed itself ran out, after the head arrived. + pub fn body_read_timeout(elapsed: Elapsed) -> Self { + Repr::ReadTimeout { + elapsed, + body: true, + } + .into() + } + + pub fn is_body_read_timeout(&self) -> bool { + matches!(&*self.0, Repr::ReadTimeout { body: true, .. }) + } + + pub fn body_read(error: reqwest::Error, received: u64, expected: Option) -> Self { + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(&error); + let kind = std::iter::from_fn(|| { + let current = source?; + source = current.source(); + Some(current) + }) + .find_map(|cause| cause.downcast_ref::().map(io::Error::kind)); + let reset = match kind { + Some(io::ErrorKind::UnexpectedEof) => false, + Some(io::ErrorKind::ConnectionReset) => true, + _ => return error.into(), + }; + Repr::BodyCut { + source: error.without_url(), + cut: BodyCut { + received, + expected, + reset, + }, + } + .into() + } + + pub fn body_cut(&self) -> Option { + match &*self.0 { + Repr::BodyCut { cut, .. } => Some(*cut), + _ => None, + } + } + pub fn computed_header(name: impl Into) -> Self { Repr::ComputedHeader(name.into()).into() } @@ -168,7 +272,9 @@ impl Error { pub fn dispatched(&self) -> bool { match &*self.0 { Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::Handshake(_) | Repr::HandshakeTimeout(_) @@ -226,7 +332,9 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::Handshake(_) | Repr::HandshakeTimeout(_) @@ -258,7 +366,9 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -294,7 +404,9 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -332,7 +444,9 @@ impl Error { | Repr::UnsignableHeader(_) | Repr::Signature(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -356,7 +470,10 @@ impl Error { pub fn is_timeout(&self) -> bool { match &*self.0 { - Repr::Timeout(_) | Repr::HandshakeTimeout(_) | Repr::DownloadTimeout(_) => true, + Repr::Timeout(_) + | Repr::ReadTimeout { .. } + | Repr::HandshakeTimeout(_) + | Repr::DownloadTimeout(_) => true, Repr::Read { .. } | Repr::InvalidPem { .. } | Repr::Client(_) @@ -371,6 +488,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -392,7 +510,7 @@ impl Error { pub fn upstream_response(&self) -> Option<(u16, &str)> { match &*self.0 { - Repr::Status { status, body } => Some((*status, &body.0)), + Repr::Status { status, body, .. } => Some((*status, &body.0)), Repr::Read { .. } | Repr::InvalidPem { .. } | Repr::Client(_) @@ -407,7 +525,9 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) | Repr::WebSocketSchemeRewrite(_) @@ -431,7 +551,7 @@ impl Error { pub fn status(&self) -> Option { match &*self.0 { Repr::Status { status, .. } | Repr::DownloadStatus(status) => Some(*status), - Repr::Timeout(_) | Repr::DownloadTimeout(_) => Some(408), + Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::DownloadTimeout(_) => Some(408), Repr::Read { .. } | Repr::InvalidPem { .. } | Repr::Client(_) @@ -446,6 +566,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) | Repr::WebSocketSchemeRewrite(_) @@ -481,7 +602,9 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) + | Repr::ReadTimeout { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -674,4 +797,53 @@ mod tests { assert_eq!(error.upstream_response(), Some((429, body.as_str()))); assert_eq!(error.status(), Some(429)); } + + #[rstest] + #[case::chunked("transfer-encoding: chunked\r\n\r\n5\r\nhello\r\n", None)] + #[case::content_length("content-length: 100\r\n\r\nhello", Some(100))] + #[tokio::test] + async fn a_connection_that_ends_mid_body_is_a_body_cut( + #[case] tail: &'static str, + #[case] expected: Option, + ) { + use std::io::{Read as _, Write as _}; + + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("address"); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept"); + let mut request = [0_u8; 1024]; + let _ = stream.read(&mut request).expect("request"); + stream + .write_all(format!("HTTP/1.1 200 OK\r\n{tail}").as_bytes()) + .expect("response"); + }); + let mut response = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}/")) + .send() + .await + .expect("head"); + server.join().expect("server"); + let mut received = 0_u64; + let error = loop { + match response.chunk().await { + Ok(Some(chunk)) => received += chunk.len() as u64, + Ok(None) => panic!("the body must end with an error"), + Err(error) => break crate::Error::body_read(error, received, expected), + } + }; + assert_eq!( + error.body_cut(), + Some(super::BodyCut { + received: 5, + expected, + reset: false, + }) + ); + assert!(error.dispatched()); + assert_eq!(error.to_string(), "upstream network error"); + } } diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index 02129c62083..1aa11843b46 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -18,7 +18,7 @@ pub mod websocket; pub use client::Client; pub use config::{ClientIdentity, HttpClientConfig, Resolution, Verify}; -pub use error::{Error, TlsSource}; +pub use error::{BodyCut, Error, TlsSource}; pub use pool::{ClientVariant, HttpClientPool}; pub use proxy::EnvironmentProxies; pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive}; diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 3bdc63453ef..3b20bbda607 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -55,36 +55,13 @@ impl HttpClientPool { config: &HttpClientConfig, variant: ClientVariant, ) -> Result { - let effective = match variant { - ClientVariant::Media => HttpClientConfig { - client_certificate: None, - proxies: EnvironmentProxies::default(), - ..config.clone() - }, - ClientVariant::UnpinnedMedia => HttpClientConfig { - client_certificate: None, - ..config.clone() - }, - ClientVariant::WebSocket => HttpClientConfig { - http2: false, - verify: match &config.verify { - Verify::BuiltInRoots => Verify::NativeRoots, - verify => verify.clone(), - }, - ..config.clone() - }, - ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), - }; - let key = (effective, variant); + let key = (effective(config, variant), variant); if let Some(pooled) = self.lock().get(&key) && pooled.built_at.elapsed() < self.ttl { return Ok(Client::new(pooled.client.clone())); } - let client = self - .apply(variant, reqwest::ClientBuilder::try_from(&key.0)?) - .build() - .map_err(Repr::Client)?; + let client = self.build(&key.0, variant)?; self.lock().insert( key, PooledClient { @@ -95,6 +72,26 @@ impl HttpClientPool { Ok(Client::new(client)) } + /// A client nobody else holds, for a caller that drives it on a runtime of its own and drops both together + pub fn unpooled( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + self.build(&effective(config, variant), variant) + .map(Client::new) + } + + fn build( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + self.apply(variant, reqwest::ClientBuilder::try_from(config)?) + .build() + .map_err(|error| Error::from(Repr::Client(error))) + } + fn lock(&self) -> MutexGuard<'_, Clients> { self.clients.lock().unwrap_or_else(PoisonError::into_inner) } @@ -116,6 +113,29 @@ impl HttpClientPool { } } +fn effective(config: &HttpClientConfig, variant: ClientVariant) -> HttpClientConfig { + match variant { + ClientVariant::Media => HttpClientConfig { + client_certificate: None, + proxies: EnvironmentProxies::default(), + ..config.clone() + }, + ClientVariant::UnpinnedMedia => HttpClientConfig { + client_certificate: None, + ..config.clone() + }, + ClientVariant::WebSocket => HttpClientConfig { + http2: false, + verify: match &config.verify { + Verify::BuiltInRoots => Verify::NativeRoots, + verify => verify.clone(), + }, + ..config.clone() + }, + ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), + } +} + #[cfg(test)] mod tests { use std::{ diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index 8a779d44976..bcde3c8d30f 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-cache-response.workspace = true litellm-core-utils.workspace = true diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index 1c724a47956..cdcc3e911b9 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -58,7 +58,7 @@ pub(super) async fn execute( custom_llm_provider: provider_name.to_owned(), optional_params: Value::Object(call.optional_params), secret_fields: Vec::new(), - api_key: call.api_key.map(litellm_auth::SecretValue::new), + api_key: call.api_key.map(litellm_secret_value::SecretValue::new), }; let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?; let identity = litellm_host::interceptors::ProviderIdentity { @@ -108,7 +108,7 @@ pub(super) async fn execute( if !status.is_success() { return Err(litellm_http::Error::upstream(status.as_u16(), text).into()); } - let raw = RawResponse { body: text.clone() }; + let raw = RawResponse::new(text.clone()); if let Some(observers) = observers { observers.emit(litellm_host::lifecycle::CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 4b3e32f04bc..dd229d3d657 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -140,7 +140,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -154,7 +154,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index 694d9b4de31..0c5be52550d 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true bytes.workspace = true futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } @@ -23,6 +24,7 @@ litellm-tracing.workspace = true reqwest.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } +tokio.workspace = true tracing.workspace = true [dev-dependencies] @@ -33,5 +35,4 @@ litellm-host-native.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-inference-testing.workspace = true rstest.workspace = true -tokio.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index a57eca1c6a9..d0f109c1e99 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -1,5 +1,5 @@ use litellm_core_utils::get_llm_provider_logic::LlmProviders; -use litellm_http::request::string_headers as shared_string_headers; +use litellm_llms::ErrorDetail; use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG; use litellm_llms::{ anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, @@ -19,8 +19,6 @@ use serde_json::{Map, Value}; use super::Error; -const HEADER_CONTEXT: &str = "messages"; - #[derive(Clone, Copy)] pub(crate) struct MessagesProvider { provider: LlmProviders, @@ -35,6 +33,14 @@ impl MessagesProvider { pub(crate) fn config(self) -> &'static dyn BaseMessagesConfig { self.config } + + /// Python's `MESSAGES_SERVICE_TIER_PROVIDERS`: everyone else has `service_tier` dropped. + pub(crate) fn accepts_service_tier(self) -> bool { + matches!( + self.provider, + LlmProviders::Anthropic | LlmProviders::AzureAi | LlmProviders::BedrockMantle + ) + } } /// Python's `get_provider_anthropic_messages_config`: Vertex AI serves only its Claude @@ -67,7 +73,52 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option>, ) -> Result, Error> { - shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from) + extra_headers + .unwrap_or_default() + .into_iter() + .map(|(name, value)| match value { + Value::String(value) => Ok((name, value)), + value => Err(Error::invalid_request(ErrorDetail::invalid( + "extra_headers", + NonStringHeader { + actual: python_type_name(&value), + name, + }, + ))), + }) + .collect() +} + +/// A header value the caller set to something other than a string, named by its Python type +/// since the caller wrote it in Python. +#[derive(Debug)] +struct NonStringHeader { + name: String, + actual: &'static str, +} + +impl std::fmt::Display for NonStringHeader { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "header {} must be a string, got {}", + self.name, self.actual + ) + } +} + +impl std::error::Error for NonStringHeader {} + +fn python_type_name(value: &Value) -> &'static str { + match value { + Value::Null => "NoneType", + Value::Bool(_) => "bool", + Value::Number(number) if number.is_f64() => "float", + Value::Number(_) => "int", + Value::String(_) => "str", + Value::Array(_) => "list", + Value::Object(_) => "dict", + } } #[cfg(test)] @@ -77,7 +128,6 @@ mod tests { use rstest::rstest; use super::{messages_provider, string_headers}; - use crate::Error; use litellm_core_utils::get_llm_provider_logic::LlmProviders; #[rstest] @@ -105,17 +155,94 @@ mod tests { assert!(messages_provider(provider, model).is_none()); } - #[test] - fn string_headers_rejects_non_string_values() { - let headers = json!({"x-count": 3}).as_object().unwrap().clone(); - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + #[rstest] + #[case::int(json!(12345), "int")] + #[case::float(json!(1.5), "float")] + #[case::bool(json!(true), "bool")] + #[case::null(json!(null), "NoneType")] + #[case::list(json!([1]), "list")] + #[case::dict(json!({"a": 1}), "dict")] + fn non_string_header_is_a_bad_request_naming_the_header_and_its_type( + #[case] value: serde_json::Value, + #[case] python_type: &str, + ) { + let headers = json!({"x-kept": "v", "x-count": value}) + .as_object() + .unwrap() + .clone(); + let error = string_headers(Some(headers)).expect_err("non-string header refused"); + assert!(error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); + let rejection = error.rejected_request().expect("a rejected request"); assert_eq!( - err, - Error::Headers(litellm_http::request::HeaderError { - context: "messages", - name: "x-count".to_string(), - actual: "number", - }) + litellm_error::Report::new(rejection).to_string(), + format!("invalid extra_headers: header x-count must be a string, got {python_type}") ); } } + +/// Python's `_merge_messages_header_sources`: later sources win per header name, except +/// `anthropic-beta` (any casing), whose values from every source are unioned, deduplicated and +/// sorted into one lowercase `anthropic-beta` entry. +pub fn merge_messages_header_sources( + sources: impl IntoIterator>, +) -> Map { + let mut merged = Map::new(); + let mut betas = std::collections::BTreeSet::new(); + for (name, value) in sources.into_iter().flatten() { + if !name.eq_ignore_ascii_case("anthropic-beta") { + merged.insert(name, value); + continue; + } + let values: Vec<&str> = match &value { + Value::String(text) => vec![text.as_str()], + Value::Array(entries) => entries.iter().filter_map(Value::as_str).collect(), + _ => Vec::new(), + }; + betas.extend( + values + .iter() + .flat_map(|value| value.split(',')) + .map(str::trim) + .filter(|beta| !beta.is_empty()) + .map(str::to_owned), + ); + } + if !betas.is_empty() { + merged.insert( + "anthropic-beta".to_owned(), + Value::String(betas.into_iter().collect::>().join(",")), + ); + } + merged +} + +#[cfg(test)] +mod merge_tests { + use serde_json::{Map, Value, json}; + + use super::merge_messages_header_sources; + + fn map(value: Value) -> Map { + value.as_object().cloned().unwrap() + } + + #[test] + fn later_sources_win_except_anthropic_beta_which_is_unioned_and_sorted() { + let merged = merge_messages_header_sources([ + map(json!({"x-a": "1", "anthropic-beta": "b-2, a-1"})), + map(json!({"x-a": "2", "Anthropic-Beta": ["c-3", 7, "a-1"]})), + map(json!({"anthropic-beta": 5})), + ]); + assert_eq!( + Value::Object(merged), + json!({"x-a": "2", "anthropic-beta": "a-1,b-2,c-3"}) + ); + } + + #[test] + fn no_beta_values_leave_no_beta_header() { + let merged = merge_messages_header_sources([map(json!({"anthropic-beta": ""}))]); + assert!(merged.is_empty()); + } +} diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 9d0bf09197e..c7a7d93bf98 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -2,15 +2,17 @@ use std::time::Duration; use bytes::Bytes; use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream}; -use litellm_auth::{AuthServices, SecretValue}; +use litellm_auth::AuthServices; +use litellm_secret_value::SecretValue; + use litellm_core_utils::{ dot_notation_indexing::delete_nested_value, get_provider_specific_headers::get_provider_specific_headers, settings::Lookup, }; use litellm_host::interceptors::{ - Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, WireRequest, + Cost, ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext, + ResponseHead, ResultSource, WireRequest, }; -use litellm_http::request::with_default_headers; use litellm_llms::base_llm::{ auth::{Authenticated, ValidatedEnvironment, resolve_auth}, messages::{ @@ -19,14 +21,16 @@ use litellm_llms::base_llm::{ transformation::BaseMessagesConfig, }, }; -use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse}; +use litellm_llms_types::formats::messages::{ + MessagesOptionalParams, MessagesRequest, MessagesResponse, +}; use litellm_secrets::source::Secrets; use litellm_tracing::ByteChunk; -use serde_json::Value; +use serde_json::{Map, Value}; use super::{ - Error, MessagesCall, MessagesCallResponse, MessagesRoute, - common_utils::{MessagesProvider, string_headers}, + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, + common_utils::{MessagesProvider, merge_messages_header_sources, string_headers}, constants::MESSAGES_TIMEOUT_SECS, route::Messages, types::invalid_request, @@ -81,7 +85,10 @@ pub(super) async fn execute( ), }; let reported = match &output { - MessagesCallResponse::Complete(response) => provider.config().reported_cost(response), + MessagesCallResponse::Complete(response) => match response { + MessagesBody::Message(message) => provider.config().reported_cost(message), + MessagesBody::Unchecked(_) => None, + }, MessagesCallResponse::Stream { .. } => None, }; context @@ -124,14 +131,19 @@ async fn provider_call( MessagesRequest { model, ..body }, shaping.settings.reasoning_auto_summary, )?; - let trimmed = + let transform_context = + MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params); + let without_listed = without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?; + let trimmed = requested_optional_params(without_listed, provider, &transform_context)?; + let requested_params = serde_json::to_value(&trimmed.params).map_err(serialize_failure)?; let scoped = get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str()); - let forwarded = string_headers(Some( - extra_headers.into_iter().flatten().chain(scoped).collect(), - ))?; + let forwarded = string_headers(Some(merge_messages_header_sources([ + extra_headers.unwrap_or_default(), + scoped, + ])))?; let validated = config.validate_environment( forwarded, api_key.as_deref(), @@ -142,16 +154,13 @@ async fn provider_call( )?; let body = config.request_body( - config.transform_anthropic_messages_request( - trimmed, - &MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params), - )?, + config.transform_anthropic_messages_request(trimmed, &transform_context)?, &validated.headers, shaping.settings.drop_params, )?; let environment = ValidatedEnvironment { headers: config.request_headers( - with_default_headers(validated.headers, config.default_headers()), + with_python_default_headers(validated.headers, config.default_headers()), &body, ), auth: validated.auth, @@ -168,7 +177,7 @@ async fn provider_call( let request_context = RequestContext { model: body.model.clone(), custom_llm_provider: provider.as_str().to_string(), - optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?, + optional_params: requested_params, secret_fields: Vec::new(), api_key: api_key.map(SecretValue::new), }; @@ -182,7 +191,7 @@ async fn provider_call( .before_provider_request( WireRequest { url, - headers: authenticated.headers, + headers: joined_header_values(authenticated.headers), body: config.wire_body(serde_json::to_value(&body).map_err(serialize_failure)?), }, request_context, @@ -211,6 +220,64 @@ async fn provider_call( }) } +/// Python sends its header dict through `httpx.Headers`, which folds names that differ only in +/// case into one lowercase header whose values are joined with `", "`. +fn joined_header_values(headers: Vec<(String, String)>) -> Vec<(String, String)> { + let mut joined: Vec<(String, String)> = Vec::with_capacity(headers.len()); + let mut positions: std::collections::HashMap = + std::collections::HashMap::with_capacity(headers.len()); + for (name, value) in headers { + match positions.entry(name.to_ascii_lowercase()) { + std::collections::hash_map::Entry::Occupied(position) => { + let present = &mut joined[*position.get()].1; + present.push_str(", "); + present.push_str(&value); + } + std::collections::hash_map::Entry::Vacant(position) => { + position.insert(joined.len()); + joined.push((name, value)); + } + } + } + joined +} + +/// Python's `finalize_messages_headers` adds a default only when the exact lowercase name is +/// absent, so a caller's `Anthropic-Version` is sent alongside the default, not instead of it. +fn with_python_default_headers( + headers: Vec<(String, String)>, + defaults: &[(&str, &str)], +) -> Vec<(String, String)> { + let missing: Vec<(String, String)> = defaults + .iter() + .filter(|(name, _)| !headers.iter().any(|(present, _)| present == name)) + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect(); + headers.into_iter().chain(missing).collect() +} + +/// Python's `get_requested_anthropic_messages_optional_param`, which gates speed, service tier +/// and sampling params before `additional_drop_params` and the environment check run. +fn requested_optional_params( + request: MessagesRequest, + provider: MessagesProvider, + context: &MessagesTransformContext, +) -> Result { + let request = litellm_llms::anthropic::messages::transformation::drop_unsupported_params( + request, context, + )?; + if provider.accepts_service_tier() { + return Ok(request); + } + Ok(MessagesRequest { + params: MessagesOptionalParams { + service_tier: None, + ..request.params + }, + ..request + }) +} + fn without_additional_drop_params( request: MessagesRequest, paths: &[String], @@ -243,6 +310,7 @@ async fn call_provider( } = request; let provider_name = provider.as_str(); log_request_body(provider_name, stream, &wire.body); + let timeout = timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)); let response = send( http, Authenticated { @@ -252,10 +320,11 @@ async fn call_provider( &wire.url, &wire.body, timeout, + stream, ) .await?; if !response.status().is_success() { - return Err(provider_error(response).await); + return Err(provider_error(response, timeout).await); } let config = provider.config(); if stream { @@ -263,13 +332,25 @@ async fn call_provider( response, config.stream_decoder(), provider_name, + timeout, )); } + let head = ResponseHead { + status: response.status().as_u16(), + headers: litellm_http::request::header_pairs(response.headers()), + }; + let status = head.status; let text = response.text().await.map_err(litellm_http::Error::from)?; log_response_body(&text); - context.response_received(&text).await?; - decode_response(config, &identity.model, &text) - .map(|message| MessagesCallResponse::Complete(Box::new(message))) + context + .raw_response_received(RawResponse::new(text.clone()).with_head(head)) + .await?; + if serde_json::from_str::(&text).is_err() { + // Python raises `AnthropicError(status_code=, message=)` for a + // body it cannot parse, and the public mapping classifies it by that status. + return Err(litellm_http::Error::upstream(status, text).into()); + } + decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete) } fn serialize_failure(err: serde_json::Error) -> Error { @@ -279,30 +360,60 @@ fn serialize_failure(err: serde_json::Error) -> Error { )) } +/// Python's httpx timeout bounds each phase and each read, never the whole exchange. A complete +/// response keeps the client's overall deadline; a stream is bounded until its head arrives and +/// then per read (see [`timed_chunks`]), so a long stream that keeps sending is not cut off. async fn send( http: &litellm_http::Client, authenticated: Authenticated, url: &str, body: &Value, - timeout: Option, + timeout: Duration, + stream: bool, ) -> Result { - let request = outbound_request( - authenticated, - url, - body, - Some(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), - )?; - Ok(litellm_inference::outbound::send(request, http).await?) + let request = outbound_request(authenticated, url, body, (!stream).then_some(timeout))?; + if !stream { + return Ok(litellm_inference::outbound::send(request, http).await?); + } + tokio::time::timeout(timeout, litellm_inference::outbound::send(request, http)) + .await + .map_err(litellm_http::Error::read_timeout)? + .map_err(Error::from) } -async fn provider_error(response: reqwest::Response) -> Error { +/// The stream's body, each read bounded by `timeout` as httpx bounds each read. +fn timed_chunks( + response: reqwest::Response, + timeout: Duration, +) -> impl futures_util::Stream> + Send + 'static { + let expected = response.content_length(); + futures_util::stream::try_unfold( + (response, 0_u64), + move |(mut response, received)| async move { + let chunk = tokio::time::timeout(timeout, response.chunk()) + .await + .map_err(litellm_http::Error::body_read_timeout)? + .map_err(|error| litellm_http::Error::body_read(error, received, expected))?; + Ok(chunk.map(|chunk| { + let received = received + chunk.len() as u64; + (chunk, (response, received)) + })) + }, + ) +} + +/// The provider's error, its body read under `timeout`: a stream has no overall deadline once +/// its head arrives, so an error body that stalls would otherwise hold the call forever. +async fn provider_error(response: reqwest::Response, timeout: Duration) -> Error { let status = response.status().as_u16(); - match response.text().await { - Ok(text) => { + let headers = litellm_http::request::header_pairs(response.headers()); + match tokio::time::timeout(timeout, response.text()).await { + Ok(Ok(text)) => { log_error_body(status, &text); - litellm_http::Error::upstream(status, text).into() + litellm_http::Error::upstream_with_headers(status, text, headers).into() } - Err(error) => litellm_http::Error::from(error).into(), + Ok(Err(error)) => litellm_http::Error::from(error).into(), + Err(elapsed) => litellm_http::Error::read_timeout(elapsed).into(), } } @@ -310,15 +421,23 @@ fn decode_response( config: &dyn BaseMessagesConfig, model: &str, text: &str, -) -> Result { - let response = serde_json::from_str(text).map_err(|err| { - Error::invalid_response(litellm_llms::ErrorDetail::invalid( - "messages response JSON", - err, - )) - })?; +) -> Result { + let response = match serde_json::from_str::(text) { + Ok(response) => response, + Err(err) => { + let Ok(mut body) = serde_json::from_str::>(text) else { + return Err(Error::invalid_response(litellm_llms::ErrorDetail::invalid( + "messages response JSON", + err, + ))); + }; + body.entry("stop_sequence").or_insert(Value::Null); + return Ok(MessagesBody::Unchecked(body)); + } + }; config .transform_anthropic_messages_response(model, response) + .map(|message| MessagesBody::Message(Box::new(message))) .map_err(Error::from) .map_err(Error::after_send) } @@ -327,22 +446,22 @@ fn streaming_response( response: reqwest::Response, decoder: Option, provider: &'static str, + timeout: Duration, ) -> MessagesCallResponse { let headers = response .headers() .iter() .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) .collect(); + let chunks = timed_chunks(response, timeout) + .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)); let chunks = match decoder { - None => futures_util::stream::try_unfold(response, move |mut response| async move { - let chunk = response.chunk().await.map_err(litellm_http::Error::from)?; - Ok(chunk.map(|chunk| { - log_chunk(provider, "provider_response", &chunk); - (chunk, response) - })) - }) - .boxed(), - Some(decode) => decoded_chunks(response, decode, provider), + None => chunks.map_err(Error::from).boxed(), + Some(decode) => decoded_chunks( + chunks.map_err(std::io::Error::other).boxed(), + decode, + provider, + ), }; MessagesCallResponse::Stream { head: super::route::MessagesStreamHead { headers }, @@ -351,15 +470,10 @@ fn streaming_response( } fn decoded_chunks( - response: reqwest::Response, + bytes: ByteStream, decode: StreamDecoder, provider: &'static str, ) -> BoxStream<'static, Result> { - let bytes: ByteStream = response - .bytes_stream() - .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)) - .map_err(|error| std::io::Error::other(litellm_http::Error::from(error))) - .boxed(); futures_util::stream::try_unfold(decode(bytes), move |mut events| async move { let Some(event) = events.try_next().await? else { return Ok(None); @@ -962,9 +1076,12 @@ mod tests { .send() .await .unwrap(); - let MessagesCallResponse::Stream { mut chunks, .. } = - streaming_response(response, Some(streaming::event_stream), "test") - else { + let MessagesCallResponse::Stream { mut chunks, .. } = streaming_response( + response, + Some(streaming::event_stream), + "test", + Duration::from_secs(5), + ) else { panic!("a streaming response returns chunks"); }; diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index bfefb9c63b0..25ebeaacb49 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -14,10 +14,11 @@ use std::sync::Arc; use crate::common_utils::messages_provider; +pub use common_utils::merge_messages_header_sources; pub use litellm_inference::Error; pub use types::{ - MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, - messages_body, + MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, + litellm_params, messages_body, }; #[derive(Clone)] diff --git a/litellm-rust/crates/inference-messages/src/route.rs b/litellm-rust/crates/inference-messages/src/route.rs index 7ff886689da..5f06d16d49f 100644 --- a/litellm-rust/crates/inference-messages/src/route.rs +++ b/litellm-rust/crates/inference-messages/src/route.rs @@ -5,11 +5,10 @@ use litellm_host::{ call::{HostedCompletion, HostedMachine, hosted_call}, protocol::Protocol, }; -use litellm_llms_types::formats::messages::MessagesResponse; -use super::{Error, MessagesCall}; +use super::{Error, MessagesBody, MessagesCall}; -pub type MessagesOutput = HostedCompletion>; +pub type MessagesOutput = HostedCompletion; /// The upstream response as the caller sees it at stream hand-off, before any chunk. pub struct MessagesStreamHead { @@ -19,7 +18,7 @@ pub struct MessagesStreamHead { pub struct Messages; impl Protocol for Messages { - type Response = Box; + type Response = MessagesBody; type Error = Error; type Request = MessagesCall; type HostCall = Infallible; diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index ae346259559..772ea6f7292 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -42,7 +42,16 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error { } pub type MessagesCallResponse = - CallOutput, super::route::MessagesStreamHead, Bytes, Error>; + CallOutput; + +// Python's `transform_anthropic_messages_response` returns any JSON object the provider sends +// with a 200, message or not, so one that is not a message passes through as sent. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MessagesBody { + Message(Box), + Unchecked(Map), +} #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct MessagesShaping { @@ -144,6 +153,7 @@ mod tests { supports_reasoning: true, supports_adaptive_thinking: true, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: true, supports_sampling_params: false, diff --git a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs index f91d5c79f7c..5207f4e9d58 100644 --- a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs +++ b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs @@ -1,7 +1,9 @@ use std::{sync::Arc, time::Duration}; use futures_util::TryStreamExt; -use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping}; +use litellm_inference_messages::{ + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesShaping, +}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse}; use rstest::{fixture, rstest}; @@ -74,7 +76,10 @@ async fn run(call: MessagesCall) -> Result { async fn run_message(call: MessagesCall) -> MessagesResponse { match run(call).await.unwrap() { - MessagesCallResponse::Complete(message) => *message, + MessagesCallResponse::Complete(message) => match message { + MessagesBody::Message(message) => *message, + MessagesBody::Unchecked(_) => panic!("expected a message"), + }, MessagesCallResponse::Stream { .. } => panic!("expected a completed message"), } } diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 480b3c4f3b7..787501a8fed 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -50,42 +50,47 @@ async fn messages_cache_identity_includes_provider_native_parameters( #[case] original: Value, #[case] changed: Value, ) { - use litellm_inference_messages::route::Messages; + use litellm_inference_messages::{MessagesBody, route::Messages}; use litellm_llms_types::formats::messages::MessagesResponse; let calls = AtomicUsize::new(0); for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] { - let response = - execute_unary::( - CacheRequest::from_wire( - ProviderIdentity { - model: "test".into(), - provider: "anthropic".into(), - }, - Some(&WireRequest { - url: "https://example.test/v1/messages".into(), - headers: vec![], - body: json!({ - "model":"test", "messages":[{"role":"user","content":"hello"}], - "max_tokens":32, (field):value - }), - }), - ), - Some(cache.clone()), - Some(CacheOptions::new(CacheScope::Shared)), - &(), - None, - || async { - let call = calls.fetch_add(1, Ordering::SeqCst); - Ok(Box::new(serde_json::from_value::(json!({ - "id":call.to_string(), "type":"message", "role":"assistant", "model":"test", - "content":[{"type":"text","text":format!("answer {call}")}], - "stop_reason":"end_turn", "stop_sequence":null - })).unwrap())) + let response = execute_unary::( + CacheRequest::from_wire( + ProviderIdentity { + model: "test".into(), + provider: "anthropic".into(), }, - ) - .await - .unwrap(); + Some(&WireRequest { + url: "https://example.test/v1/messages".into(), + headers: vec![], + body: json!({ + "model":"test", "messages":[{"role":"user","content":"hello"}], + "max_tokens":32, (field):value + }), + }), + ), + Some(cache.clone()), + Some(CacheOptions::new(CacheScope::Shared)), + &(), + None, + || async { + let call = calls.fetch_add(1, Ordering::SeqCst); + Ok(MessagesBody::Message(Box::new( + serde_json::from_value::(json!({ + "id":call.to_string(), "type":"message", "role":"assistant", "model":"test", + "content":[{"type":"text","text":format!("answer {call}")}], + "stop_reason":"end_turn", "stop_sequence":null + })) + .unwrap(), + ))) + }, + ) + .await + .unwrap(); + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!(response.id, expected_call.to_string()); assert_eq!( response.content[0]["text"], @@ -107,7 +112,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -121,7 +126,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } @@ -253,7 +258,7 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ #[case] cost: Value, #[case] expected: Option, ) { - use litellm_inference_messages::{MessagesCall, MessagesCallResponse}; + use litellm_inference_messages::{MessagesBody, MessagesCall, MessagesCallResponse}; let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test", "content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null, "usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost}); @@ -277,6 +282,9 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ panic!("expected Messages responses"); }; assert_eq!(first, second); + let MessagesBody::Message(first) = first else { + panic!("expected a message"); + }; assert_eq!(first.extra.get("cost"), Some(&cost)); assert_eq!(support::received(&upstream).await.len(), 1); let facts = hooks.facts.lock().unwrap(); diff --git a/litellm-rust/crates/inference-messages/tests/live/support.rs b/litellm-rust/crates/inference-messages/tests/live/support.rs index 81e295e34ad..52684928f1f 100644 --- a/litellm-rust/crates/inference-messages/tests/live/support.rs +++ b/litellm-rust/crates/inference-messages/tests/live/support.rs @@ -4,7 +4,9 @@ use futures_util::{FutureExt, TryStreamExt}; use litellm_host::interceptors::{ ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest, }; -use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute}; +use litellm_inference_messages::{ + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, +}; use litellm_inference_testing::live::{LiveResources, model, within_deadline}; use litellm_llms_types::formats::messages::MessagesResponse; use serde_json::{Value, json}; @@ -111,7 +113,10 @@ pub(super) fn call(provider: &str, payload: Value) -> MessagesCall { pub(super) fn complete(response: MessagesCallResponse) -> Box { match response { - MessagesCallResponse::Complete(message) => message, + MessagesCallResponse::Complete(message) => match message { + MessagesBody::Message(message) => message, + MessagesBody::Unchecked(_) => panic!("a complete request returned no message"), + }, MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"), } } diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index c3d06027c32..063e9d18738 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -6,9 +6,11 @@ use litellm_host::{ interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest}, lifecycle::CallEvent, }; -use litellm_inference_messages::{MessagesCallResponse, route::Messages}; +use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; -use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities; +use litellm_llms::base_llm::messages::context::{ + MessagesModelCapabilities as AnthropicModelCapabilities, SupportedEffortTiers, +}; use rstest::rstest; use super::*; @@ -338,7 +340,10 @@ async fn response_mode_follows_the_intercepted_request( } MessagesCallResponse::Complete(message) => { assert_eq!(expected_stream, Some(false)); - assert_eq!(*message, serde_json::from_value(message_body()).unwrap()); + assert_eq!( + message, + MessagesBody::Message(serde_json::from_value(message_body()).unwrap()) + ); } } Ok::<_, Error>(()) @@ -419,7 +424,7 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure( assert!(host.raw_responses().is_empty()); } -/// Python logs `optional_params` as what it is about to send, so a dropped param must +/// Python logs `optional_params` after its requested-param gate, so a dropped param must /// not resurface in callbacks. #[rstest] #[tokio::test] @@ -450,3 +455,82 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); assert_eq!(optional_params, json!({"max_tokens": 16})); } + +#[rstest] +#[tokio::test] +async fn the_request_context_leaves_out_a_listed_param_the_model_refuses(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + shaping: MessagesShaping { + settings: MessagesSettings { + additional_drop_params: vec!["temperature".into()], + ..MessagesSettings::default() + }, + capabilities: AnthropicModelCapabilities { + supports_sampling_params: false, + ..AnthropicModelCapabilities::default() + }, + }, + ..with_fields(call, json!({"temperature": 0.2})) + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); + assert_eq!(optional_params, json!({"max_tokens": 16})); +} + +#[rstest] +#[case::reasoning_effort_on_an_adaptive_model( + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_output_config: true, + effort_tiers: SupportedEffortTiers { low: true, ..SupportedEffortTiers::default() }, + ..AnthropicModelCapabilities::default() + }, + json!({"reasoning_effort": "low"}) +)] +#[case::adaptive_payload_on_a_legacy_model( + AnthropicModelCapabilities { + supports_reasoning: true, + ..AnthropicModelCapabilities::default() + }, + json!({"max_tokens": 3000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) +)] +#[tokio::test] +async fn the_request_context_carries_the_requested_reasoning_not_its_translation( + call: MessagesCall, + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, +) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + shaping: MessagesShaping { + capabilities, + ..MessagesShaping::default() + }, + ..with_fields(call, fields.clone()) + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); + assert_eq!( + optional_params, + Value::Object( + object(json!({"max_tokens": 16})) + .into_iter() + .chain(object(fields)) + .collect() + ) + ); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index a01e571ec75..fe203287766 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -5,7 +5,7 @@ use std::{ use litellm_http::{HttpSettings, Resolution}; use litellm_inference_messages::{ - Error, MessagesCall, MessagesSettings, MessagesShaping, + Error, MessagesBody, MessagesCall, MessagesSettings, MessagesShaping, route::{Messages, MessagesMachine, MessagesOutput}, }; use litellm_inference_testing::RecordingSecrets; @@ -119,7 +119,10 @@ async fn run(call: MessagesCall) -> Result { async fn run_message(call: MessagesCall) -> MessagesResponse { match run(call).await.expect("messages call succeeds") { - MessagesOutput::Complete(message) => *message, + MessagesOutput::Complete(message) => match message { + MessagesBody::Message(message) => *message, + MessagesBody::Unchecked(_) => panic!("a call returned no message"), + }, MessagesOutput::StreamEnded | MessagesOutput::Detached => { panic!("a non-streaming call returned a stream") } diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index e53bdb8d252..3784eb6cc68 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -440,7 +440,12 @@ async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: M #[case::deepseek("deepseek")] #[case::minimax("minimax")] #[tokio::test] -async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) { +/// Python adds a default unless that exact lowercase name is present, then folds the casings +/// into one header, so a capitalised caller header is sent beside the default. +async fn capitalised_caller_protocol_headers_are_joined_with_the_defaults( + call: MessagesCall, + #[case] provider: &str, +) { let upstream = upstream([message_response()]).await; run_message(MessagesCall { @@ -457,10 +462,13 @@ async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[cas .await; let request = only_request(&upstream).await; - assert_eq!(request.header_values("anthropic-version"), ["2024-01-01"]); + assert_eq!( + request.header_values("anthropic-version"), + ["2024-01-01, 2023-06-01"] + ); assert_eq!( request.header_values("content-type"), - ["application/json; charset=utf-8"] + ["application/json; charset=utf-8, application/json"] ); } @@ -522,6 +530,30 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i assert_eq!(sent["max_tokens"], 16); } +#[rstest] +#[tokio::test] +async fn additional_drop_params_clear_a_param_the_model_does_not_support(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities: sampling_removed(), + settings: MessagesSettings { + additional_drop_params: vec!["top_k".into()], + ..MessagesSettings::default() + }, + }, + ..with_fields(call, json!({"top_k": 5})) + }) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!(sent.get("top_k"), None); + assert_eq!(sent["max_tokens"], 16); +} + #[rstest] #[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))] #[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))] diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index 7aff1286858..f978da5e995 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -2,7 +2,7 @@ use litellm_host::{ interceptors::{ExecutionFacts, ResultSource}, lifecycle::ExecutionEvent, }; -use litellm_inference_messages::{MessagesCallResponse, messages_body}; +use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body}; use litellm_inference_testing::{ RecordingSecrets, http_config, no_secrets, provider_http, resources, }; @@ -48,6 +48,9 @@ async fn calls_defer_execution_until_polled( let MessagesCallResponse::Complete(response) = future.await.unwrap() else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!( response.content, message_body()["content"].as_array().unwrap().as_slice() @@ -225,14 +228,57 @@ async fn an_upstream_error_keeps_its_status_and_body(call: MessagesCall, #[case] } #[rstest] -#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))] -#[case::not_a_message(json_response(json!({"unexpected": true})))] #[tokio::test] -async fn an_unreadable_success_body_is_an_invalid_response( +async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: MessagesCall) { + let upstream = upstream([ResponseTemplate::new(200).set_body_string("not json")]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .expect_err("an unparseable body fails"); + + assert_eq!(error.upstream_response(), Some((200, "not json"))); +} + +#[rstest] +#[case::error_object( + json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}), + json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}, + "stop_sequence": null}), +)] +#[case::message_missing_id( + json!({"type": "message", "content": [], "stop_sequence": "###"}), + json!({"type": "message", "content": [], "stop_sequence": "###"}), +)] +#[tokio::test] +async fn a_success_body_that_is_not_a_message_passes_through( call: MessagesCall, - #[case] response: ResponseTemplate, + #[case] body: Value, + #[case] expected: Value, ) { - let upstream = upstream([response]).await; + let upstream = upstream([json_response(body)]).await; + + let response = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .expect("any JSON object is returned"); + + let MessagesOutput::Complete(response) = response else { + panic!("expected a completed response"); + }; + assert_eq!(serde_json::to_value(response).unwrap(), expected); +} + +#[rstest] +#[tokio::test] +async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) { + let upstream = upstream([json_response(json!(["unexpected"]))]).await; let error = run(MessagesCall { api_key: Some("sk".into()), @@ -309,6 +355,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes let MessagesCallResponse::Complete(message) = response else { panic!("a non-streaming request returns a message"); }; + let MessagesBody::Message(message) = message else { + panic!("expected a message"); + }; assert_eq!(message.id, "msg_1"); let sent = only_request(&upstream).await; assert_eq!(sent.header("x-api-key"), Some("sk-ant")); @@ -401,6 +450,9 @@ async fn route_uses_injected_dependencies_and_optional_cache( else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!( response.content, message_body()["content"].as_array().unwrap().as_slice() @@ -473,6 +525,9 @@ async fn cache_overrides_preserve_the_routes_isolated_scope(call: MessagesCall) else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!(response.id, expected["id"].as_str().unwrap()); } assert_eq!(received(&upstream).await.len(), 2); diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index 4a8a061f764..99ba08495aa 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -432,6 +432,49 @@ async fn dropping_the_sdk_stream_closes_the_unfinished_upstream( .unwrap(); } +/// Answers with an error head whose promised body never arrives. +async fn stalling_error_upstream() -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let connection = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all(b"HTTP/1.1 429 Too Many Requests\r\ncontent-type: application/json\r\ncontent-length: 100\r\n\r\n") + .await + .unwrap(); + let _ = socket.read_to_end(&mut Vec::new()).await; + }); + (base, connection) +} + +#[rstest] +#[tokio::test] +async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) { + let (base, connection) = stalling_error_upstream().await; + let host = RecordingStreamHost::new( + MessagesCall { + timeout: Some(Duration::from_millis(300)), + ..streaming(call, base) + }, + usize::MAX, + ); + + let error = tokio::time::timeout(Duration::from_secs(5), stream_through(&host)) + .await + .expect("the stalled error body gives up within the timeout") + .expect_err("a stalled error body fails the call"); + + assert!(error.dispatched(), "{error:?}"); + assert_eq!(error.upstream_response(), None); + assert!(host.seen.into_inner().unwrap().is_empty()); + tokio::time::timeout(Duration::from_secs(5), connection) + .await + .expect("timing out closes the upstream connection") + .unwrap(); +} + #[rstest] #[tokio::test] async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesCall) { diff --git a/litellm-rust/crates/inference-ocr/Cargo.toml b/litellm-rust/crates/inference-ocr/Cargo.toml index bbbdaddb050..ad53126b9fe 100644 --- a/litellm-rust/crates/inference-ocr/Cargo.toml +++ b/litellm-rust/crates/inference-ocr/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true base64.workspace = true bytes.workspace = true futures-util.workspace = true diff --git a/litellm-rust/crates/inference-ocr/src/handler.rs b/litellm-rust/crates/inference-ocr/src/handler.rs index 426d4297f28..be23af60c69 100644 --- a/litellm-rust/crates/inference-ocr/src/handler.rs +++ b/litellm-rust/crates/inference-ocr/src/handler.rs @@ -91,9 +91,7 @@ impl> CallHooks for OcrCallHooks<'_, H> { } fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - let raw = RawResponse { - body: String::from_utf8_lossy(body).into_owned(), - }; + let raw = RawResponse::new(String::from_utf8_lossy(body).into_owned()); if let Some(observers) = self.observers { observers.emit(litellm_host::lifecycle::CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, @@ -128,7 +126,7 @@ mod tests { decode_request(OcrWireRequest { model: model.into(), document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_key: Some(litellm_secret_value::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/inference-ocr/src/provider_config.rs b/litellm-rust/crates/inference-ocr/src/provider_config.rs index 19e50af0dd1..e76d291254b 100644 --- a/litellm-rust/crates/inference-ocr/src/provider_config.rs +++ b/litellm-rust/crates/inference-ocr/src/provider_config.rs @@ -1,4 +1,6 @@ -use litellm_auth::{InputSource, SecretValue, Sourced}; +use litellm_auth::{InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::get_llm_provider_logic::{ CustomLlmProvider, LlmProviders, get_custom_llm_provider, }; @@ -407,7 +409,7 @@ mod tests { fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("explicit-key"), + litellm_secret_value::SecretValue::new("explicit-key"), InputSource::Deployment, )), api_base: Some(Sourced::new( @@ -415,7 +417,7 @@ mod tests { InputSource::Deployment, )), dynamic_api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("dynamic-key"), + litellm_secret_value::SecretValue::new("dynamic-key"), InputSource::Environment, )), dynamic_api_base: Some(Sourced::new( @@ -455,7 +457,7 @@ mod tests { ) { let dynamic_key = dynamic_value.map(|value| { Sourced::new( - litellm_auth::SecretValue::new(value), + litellm_secret_value::SecretValue::new(value), InputSource::Environment, ) }); @@ -463,7 +465,7 @@ mod tests { dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("explicit-key"), + litellm_secret_value::SecretValue::new("explicit-key"), InputSource::Deployment, )), api_base: Some(Sourced::new( @@ -502,14 +504,14 @@ mod tests { OcrCredentialInputs { api_key: explicit_key.map(|value| { Sourced::new( - litellm_auth::SecretValue::new(value), + litellm_secret_value::SecretValue::new(value), InputSource::Deployment, ) }), api_base: explicit_base .map(|value| Sourced::new(value.into(), InputSource::Deployment)), dynamic_api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("dynamic-key"), + litellm_secret_value::SecretValue::new("dynamic-key"), InputSource::Environment, )), dynamic_api_base: Some(Sourced::new( diff --git a/litellm-rust/crates/inference-ocr/src/types.rs b/litellm-rust/crates/inference-ocr/src/types.rs index 98c609a158f..1b2dbb52d77 100644 --- a/litellm-rust/crates/inference-ocr/src/types.rs +++ b/litellm-rust/crates/inference-ocr/src/types.rs @@ -1,7 +1,9 @@ use std::{collections::BTreeMap, path::PathBuf, time::Duration}; use bytes::Bytes; -use litellm_auth::{InputSource, SecretValue, TokenProviderHandle}; +use litellm_auth::{InputSource, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::call_arguments::CallArguments; use litellm_llms::base_llm::ocr::{ error::Error, diff --git a/litellm-rust/crates/inference-ocr/src/wire.rs b/litellm-rust/crates/inference-ocr/src/wire.rs index 56d22be6523..45a2865823f 100644 --- a/litellm-rust/crates/inference-ocr/src/wire.rs +++ b/litellm-rust/crates/inference-ocr/src/wire.rs @@ -1,6 +1,8 @@ use std::{collections::BTreeMap, time::Duration}; -use litellm_auth::{InputSource, SecretValue}; +use litellm_auth::InputSource; +use litellm_secret_value::SecretValue; + use litellm_llms::base_llm::ocr::{error::Error, transformation::decode_request_value}; use litellm_llms_types::formats::ocr::OcrDocument; use serde::Deserialize; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs index 6f591af73e0..65c7f309b82 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs @@ -3,9 +3,9 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, }; -use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, -}; +use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use rstest::rstest; use super::*; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs index 75a91092daf..9f59bb9a2a2 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs @@ -239,7 +239,7 @@ impl litellm_host_native::services::HostCallHandler for CallerTokenHost { OcrOp::AcquireAzureAdToken(reply) => { self.trace.lock().unwrap().push("token".into()); reply.send(litellm_auth::ResolvedCredential::Static( - litellm_auth::SecretValue::new("caller-token"), + litellm_secret_value::SecretValue::new("caller-token"), )); Ok(()) } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs index 5a48bba9e1a..6871d1ca1f6 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs @@ -80,7 +80,7 @@ fn wire(model: &str, base: &str, document: Value, options: Value) -> OcrWireRequ OcrWireRequest { model: model.into(), document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_key: Some(litellm_secret_value::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/inference-ocr/tests/resources.rs b/litellm-rust/crates/inference-ocr/tests/resources.rs index d618647e89b..3132a837eaf 100644 --- a/litellm-rust/crates/inference-ocr/tests/resources.rs +++ b/litellm-rust/crates/inference-ocr/tests/resources.rs @@ -7,7 +7,7 @@ use std::sync::{ use litellm_auth::AuthServices; use litellm_auth_gcp::{ - CredentialSource, VertexAuth, VertexAuthFuture, VertexProviderLoader, VertexTokenSource, + CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth, }; use litellm_http::{HttpSettings, Resolution}; use litellm_inference::resources::CoreResources; @@ -20,28 +20,28 @@ use support::{ReceivedRequest, json_response, upstream}; struct TokenSource(String); -impl VertexTokenSource for TokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { +impl GcpTokenSource for TokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { Box::pin(async { Ok(self.0.clone()) }) } - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async { Ok(self.0.clone()) }) + fn token(&self) -> GcpAuthFuture<'_, litellm_secret_value::SecretValue> { + Box::pin(async { Ok(litellm_secret_value::SecretValue::new(self.0.clone())) }) } } #[derive(Default)] struct Loader(AtomicUsize); -impl VertexProviderLoader for Loader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { +impl GcpProviderLoader for Loader { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc> { Box::pin(async move { self.0.fetch_add(1, Ordering::SeqCst); let identity = match source { CredentialSource::Trusted(secret) => secret.expose().to_string(), other => panic!("unexpected credential source: {other:?}"), }; - Ok(Arc::new(TokenSource(identity)) as Arc) + Ok(Arc::new(TokenSource(identity)) as Arc) }) } } diff --git a/litellm-rust/crates/inference-responses/Cargo.toml b/litellm-rust/crates/inference-responses/Cargo.toml index 083179c50c8..f90651634c5 100644 --- a/litellm-rust/crates/inference-responses/Cargo.toml +++ b/litellm-rust/crates/inference-responses/Cargo.toml @@ -25,6 +25,7 @@ tokio-tungstenite.workspace = true tracing.workspace = true [dev-dependencies] +litellm-secret-value.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache.workspace = true litellm-cache-memory.workspace = true diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index 19092a63a14..f81a51ad192 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -114,7 +114,7 @@ pub(super) async fn execute( }); } let body = response.text().await.map_err(litellm_http::Error::from)?; - let raw = RawResponse { body: body.clone() }; + let raw = RawResponse::new(body.clone()); if let Some(observers) = observers { observers.emit(litellm_host::lifecycle::CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, diff --git a/litellm-rust/crates/inference-responses/tests/caching.rs b/litellm-rust/crates/inference-responses/tests/caching.rs index 462461695c3..56e25fbcb98 100644 --- a/litellm-rust/crates/inference-responses/tests/caching.rs +++ b/litellm-rust/crates/inference-responses/tests/caching.rs @@ -98,7 +98,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -112,7 +112,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-testing/Cargo.toml b/litellm-rust/crates/inference-testing/Cargo.toml index 748b1df0d3b..5f690532d9c 100644 --- a/litellm-rust/crates/inference-testing/Cargo.toml +++ b/litellm-rust/crates/inference-testing/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true publish = false [dependencies] +litellm-secret-value.workspace = true litellm-auth.workspace = true futures-util.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/inference-testing/src/lib.rs b/litellm-rust/crates/inference-testing/src/lib.rs index 2e20d5a9e91..f330e02e99a 100644 --- a/litellm-rust/crates/inference-testing/src/lib.rs +++ b/litellm-rust/crates/inference-testing/src/lib.rs @@ -5,7 +5,8 @@ use litellm_http::{ ClientVariant, HttpClientConfig, HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver, }; -use litellm_secrets::{SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::source::SecretSource; use litellm_inference::resources::CoreResources; diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index 963aa7b90c5..ef2635d24a1 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -42,9 +42,11 @@ impl<'a, I: Interceptors> CallContext<'a, I> { } pub async fn response_received(&self, body: &str) -> Result<(), Error> { - let raw = RawResponse { - body: body.to_owned(), - }; + self.raw_response_received(RawResponse::new(body.to_owned())) + .await + } + + pub async fn raw_response_received(&self, raw: RawResponse) -> Result<(), Error> { if let Some(observers) = &self.observers { observers.emit(CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }, diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index a8e6fb4e1b3..4135424d6ef 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -146,6 +146,42 @@ impl Error { } } + /// The provider's response headers on an upstream failure; empty for every other error. + pub fn upstream_headers(&self) -> &[(String, String)] { + match self { + Self::Llm(error) => error.upstream_headers(), + Self::Http(error) => error.upstream_headers(), + _ => &[], + } + } + + /// A read of a streamed body timed out after the response head arrived. + pub fn body_read_timeout(&self) -> bool { + match self { + Self::Llm(error) => error.is_body_read_timeout(), + Self::Http(error) => error.is_body_read_timeout(), + _ => false, + } + } + + pub fn body_cut(&self) -> Option { + match self { + Self::Llm(error) => error.body_cut(), + Self::Http(error) => error.body_cut(), + _ => None, + } + } + + /// `Some(connect)` when the call timed out, where `connect` says the connection was never + /// established; `None` for every other error. + pub fn timeout(&self) -> Option { + match self { + Self::Llm(error) => error.is_timeout(), + Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()), + _ => None, + } + } + pub fn upstream_response(&self) -> Option<(u16, &str)> { match self { Self::Llm(error) => error.upstream_response(), diff --git a/litellm-rust/crates/llms-types/src/formats/messages/request.rs b/litellm-rust/crates/llms-types/src/formats/messages/request.rs index 9dbfa068025..d30d21e74e7 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/request.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/request.rs @@ -264,7 +264,12 @@ pub enum ThinkingDisplay { #[macro_rules_attribute::apply(crate::wire_type)] #[derive(Default)] pub struct EnabledThinking { - #[serde(default, skip_serializing_if = "Option::is_none")] + /// An explicit `null` stays `Unrecognized(Null)`: Python rejects it, but defaults a missing key to 0. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "crate::recognized::present" + )] pub budget_tokens: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub display: Option>, diff --git a/litellm-rust/crates/llms-types/src/formats/messages/response.rs b/litellm-rust/crates/llms-types/src/formats/messages/response.rs index 9b0c5d31b3a..276a30a6cb4 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/response.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/response.rs @@ -10,14 +10,29 @@ pub struct MessagesResponse { pub content: Vec, pub stop_reason: Option, pub stop_sequence: Option, - #[serde(skip_serializing_if = "Option::is_none")] + /// `Some(Value::Null)` when the provider sent an explicit `null`, which Python keeps. + #[serde( + default, + deserialize_with = "present_value", + skip_serializing_if = "Option::is_none" + )] pub usage: Option, - #[serde(skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "present_value", + skip_serializing_if = "Option::is_none" + )] pub container: Option, #[serde(flatten)] pub extra: Map, } +fn present_value<'de, D: serde::Deserializer<'de>>( + deserializer: D, +) -> Result, D::Error> { + >::deserialize(deserializer).map(Some) +} + #[cfg(test)] mod tests { use super::*; @@ -73,4 +88,21 @@ mod tests { assert_eq!(body.get("usage").cloned(), usage); assert_eq!(body.get("container").cloned(), container); } + + #[test] + fn explicit_null_usage_and_container_survive_a_round_trip() { + let body = json!({"id": "msg_1", "type": "message", "role": "assistant", "model": "m", + "content": [], "stop_reason": null, "stop_sequence": null, + "usage": null, "container": null}); + let parsed: MessagesResponse = serde_json::from_value(body.clone()).expect("parses"); + assert_eq!(serde_json::to_value(parsed).expect("serializable"), body); + } + + #[test] + fn absent_usage_and_container_stay_absent() { + let body = json!({"id": "msg_1", "type": "message", "role": "assistant", "model": "m", + "content": [], "stop_reason": null, "stop_sequence": null}); + let parsed: MessagesResponse = serde_json::from_value(body.clone()).expect("parses"); + assert_eq!(serde_json::to_value(parsed).expect("serializable"), body); + } } diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index 80095cac36c..46c14a21c59 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -124,12 +124,86 @@ pub enum AnthropicBeta { ToolSearchTool20251019, #[strum(serialize = "inline-tools-2026-09-15")] InlineTools20260915, + #[strum(serialize = "afk-mode-2026-01-31")] + AfkMode20260131, + #[strum(serialize = "agent-memory-2026-07-22")] + AgentMemory20260722, + #[strum(serialize = "analytics-usage-cost-2026-08-24")] + AnalyticsUsageCost20260824, + #[strum(serialize = "cache-diagnosis-2026-04-07")] + CacheDiagnosis20260407, + #[strum(serialize = "ce-plugins-2026-09-01")] + CePlugins20260901, + #[strum(serialize = "ce-user-management-2026-07-13")] + CeUserManagement20260713, + #[strum(serialize = "code-execution-2025-05-22")] + CodeExecution20250522, + #[strum(serialize = "computer-use-2024-10-22")] + ComputerUse20241022, + #[strum(serialize = "dev-full-thinking-2025-05-14")] + DevFullThinking20250514, + #[strum(serialize = "dreaming-2026-04-21")] + Dreaming20260421, + #[strum(serialize = "experimental-claude-tag-2026-10-21")] + ExperimentalClaudeTag20261021, + #[strum(serialize = "extended-cache-ttl-2025-04-11")] + ExtendedCacheTtl20250411, + #[strum(serialize = "fallback-credit-2026-06-01")] + FallbackCredit20260601, + #[strum(serialize = "fallback-credit-2026-07-01")] + FallbackCredit20260701, + #[strum(serialize = "managed-agents-2026-04-01")] + ManagedAgents20260401, + #[strum(serialize = "mcp-client-2026-09-15")] + McpClient20260915, + #[strum(serialize = "mcp-tunnels-2026-06-22")] + McpTunnels20260622, + #[strum(serialize = "message-batches-2024-09-24")] + MessageBatches20240924, + #[strum(serialize = "message-threads-2026-08-12")] + MessageThreads20260812, + #[strum(serialize = "mid-conversation-system-clear-at-2026-08-21")] + MidConversationSystemClearAt20260821, + #[strum(serialize = "model-context-window-exceeded-2025-08-26")] + ModelContextWindowExceeded20250826, + #[strum(serialize = "output-300k-2026-03-24")] + Output300k20260324, + #[strum(serialize = "pdfs-2024-09-25")] + Pdfs20240925, + #[strum(serialize = "prompt-caching-2024-07-31")] + PromptCaching20240731, + #[strum(serialize = "redact-thinking-2026-02-12")] + RedactThinking20260212, + #[strum(serialize = "server-side-fallback-2026-06-01")] + ServerSideFallback20260601, + #[strum(serialize = "server-side-fallback-2026-07-01")] + ServerSideFallback20260701, + #[strum(serialize = "spend-limit-reads-2026-09-26")] + SpendLimitReads20260926, + #[strum(serialize = "structured-outputs-2025-12-15")] + StructuredOutputs20251215, + #[strum(serialize = "task-budgets-2026-03-13")] + TaskBudgets20260313, + #[strum(serialize = "telemetry-destinations-2026-08-11")] + TelemetryDestinations20260811, + #[strum(serialize = "thinking-token-count-2026-05-13")] + ThinkingTokenCount20260513, + #[strum(serialize = "token-counting-2024-11-01")] + TokenCounting20241101, + #[strum(serialize = "user-profiles-2026-03-24")] + UserProfiles20260324, + #[strum(serialize = "user-profiles-2026-08-18")] + UserProfiles20260818, + #[strum(serialize = "user-profiles-2026-09-04")] + UserProfiles20260904, + #[strum(serialize = "webhooks-2026-10-07")] + Webhooks20261007, #[strum(default, transparent)] Other(String), } impl AnthropicBeta { - pub const KNOWN: [Self; 40] = [ + pub const KNOWN: [Self; 77] = [ Self::Oauth20250420, Self::WebFetch20250910, Self::WebSearch20250305, @@ -170,6 +244,43 @@ impl AnthropicBeta { Self::ToolExamples20251029, Self::ToolSearchTool20251019, Self::InlineTools20260915, + Self::AfkMode20260131, + Self::AgentMemory20260722, + Self::AnalyticsUsageCost20260824, + Self::CacheDiagnosis20260407, + Self::CePlugins20260901, + Self::CeUserManagement20260713, + Self::CodeExecution20250522, + Self::ComputerUse20241022, + Self::DevFullThinking20250514, + Self::Dreaming20260421, + Self::ExperimentalClaudeTag20261021, + Self::ExtendedCacheTtl20250411, + Self::FallbackCredit20260601, + Self::FallbackCredit20260701, + Self::ManagedAgents20260401, + Self::McpClient20260915, + Self::McpTunnels20260622, + Self::MessageBatches20240924, + Self::MessageThreads20260812, + Self::MidConversationSystemClearAt20260821, + Self::ModelContextWindowExceeded20250826, + Self::Output300k20260324, + Self::Pdfs20240925, + Self::PromptCaching20240731, + Self::RedactThinking20260212, + Self::ServerSideFallback20260601, + Self::ServerSideFallback20260701, + Self::SpendLimitReads20260926, + Self::StructuredOutputs20251215, + Self::TaskBudgets20260313, + Self::TelemetryDestinations20260811, + Self::ThinkingTokenCount20260513, + Self::TokenCounting20241101, + Self::UserProfiles20260324, + Self::UserProfiles20260818, + Self::UserProfiles20260904, + Self::Webhooks20261007, ]; pub fn as_str(&self) -> &str { @@ -195,11 +306,10 @@ impl AnthropicBeta { fn rejected_by(&self) -> &'static [BetaProvider] { match self { - Self::AdvancedToolUse20251120 => { + Self::AdvancedToolUse20251120 | Self::ToolSearchTool20251019 => { &[BetaProvider::BedrockConverse, BetaProvider::Databricks] } - Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse], - Self::AdvisorTool20260301 | Self::Compact20260904 => &[ + Self::AdvisorTool20260301 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, BetaProvider::Bedrock, @@ -207,88 +317,99 @@ impl AnthropicBeta { BetaProvider::VertexAi, BetaProvider::Databricks, ], - Self::Bash20241022 + Self::AfkMode20260131 + | Self::CePlugins20260901 + | Self::CeUserManagement20260713 + | Self::Compact20260904 + | Self::Dreaming20260421 + | Self::ExperimentalClaudeTag20261021 + | Self::ExtendedCacheTtl20250411 + | Self::FallbackCredit20260601 + | Self::FallbackCredit20260701 + | Self::InlineTools20260915 + | Self::ManagedAgents20260401 + | Self::McpTunnels20260622 + | Self::MidConversationSystemClearAt20260821 + | Self::ModelContextWindowExceeded20250826 + | Self::Oauth20250420 + | Self::Pdfs20240925 + | Self::PromptCaching20240731 + | Self::PromptCachingScope20260105 + | Self::RedactThinking20260212 + | Self::ServerSideFallback20260601 + | Self::SpendLimitReads20260926 + | Self::StructuredOutputs20251215 + | Self::TaskBudgets20260313 + | Self::TelemetryDestinations20260811 + | Self::ThinkingTokenCount20260513 + | Self::TokenCounting20241101 + | Self::UserProfiles20260324 + | Self::UserProfiles20260818 + | Self::UserProfiles20260904 => &[ + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::BedrockMantle, + BetaProvider::Databricks, + ], + Self::AgentMemory20260722 + | Self::Bash20241022 | Self::Bash20250124 + | Self::DevFullThinking20250514 | Self::McpServers20251204 | Self::StructuredOutput20240301 | Self::TextEditor20241022 | Self::TextEditor20250124 => BetaProvider::VARIANTS, - Self::ClaudeCode20250219 | Self::ToolExamples20251029 => &[ - BetaProvider::Anthropic, - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::CodeExecution20250825 + Self::AnalyticsUsageCost20260824 + | Self::CacheDiagnosis20260407 + | Self::CodeExecution20250522 + | Self::CodeExecution20250825 + | Self::ComputerUse20241022 + | Self::FastMode20260201 | Self::FilesApi20250414 | Self::McpClient20250404 | Self::McpClient20251120 - | Self::PromptCachingScope20260105 + | Self::McpClient20260915 + | Self::MessageBatches20240924 + | Self::MessageThreads20260812 + | Self::Output300k20260324 + | Self::ServerSideFallback20260701 | Self::Skills20251002 - | Self::WebFetch20250910 => &[ + | Self::WebFetch20250910 + | Self::Webhooks20261007 => &[ BetaProvider::BedrockConverse, BetaProvider::Bedrock, BetaProvider::BedrockMantle, BetaProvider::VertexAi, BetaProvider::Databricks, ], - Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse], + Self::ClaudeCode20250219 + | Self::PerTurnControl20260701 + | Self::ToolExamples20251029 => &[ + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::Databricks, + ], + Self::Compact20260112 + | Self::ContextManagement20250627 + | Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 + | Self::Effort20251124 + | Self::FineGrainedToolStreaming20250514 | Self::ThinkingBindingControls20260801 | Self::ThinkingDisplayUpdates20260818 => &[], - Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], - Self::Effort20251124 => &[BetaProvider::VertexAi], - Self::FastMode20260201 | Self::Oauth20250420 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::BedrockMantle, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::FineGrainedToolStreaming20250514 => { - &[BetaProvider::AzureAi, BetaProvider::VertexAi] - } - Self::InterleavedThinking20250514 | Self::WebSearch20250305 => { - &[BetaProvider::BedrockConverse, BetaProvider::Bedrock] - } + Self::InterleavedThinking20250514 + | Self::Output128k20250219 + | Self::TokenEfficientTools20250219 + | Self::WebSearch20250305 => &[BetaProvider::BedrockConverse, BetaProvider::Bedrock], Self::MidConversationOutputConfig20260701 | Self::MidConversationToolChanges20260701 => &[ - BetaProvider::AzureAi, BetaProvider::BedrockConverse, BetaProvider::BedrockMantle, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::VertexAi, - ], - Self::PerTurnControl20260701 => &[ - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::Databricks, - ], - Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi], - Self::ToolSearchTool20251019 => &[ - BetaProvider::Anthropic, - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Databricks, - ], - Self::InlineTools20260915 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::BedrockMantle, BetaProvider::Databricks, ], + Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock], Self::Other(_) => BetaProvider::VARIANTS, } } @@ -441,7 +562,44 @@ mod tests { AnthropicBeta::TokenEfficientTools20250219, AnthropicBeta::ToolExamples20251029, AnthropicBeta::ToolSearchTool20251019, - AnthropicBeta::InlineTools20260915 + AnthropicBeta::InlineTools20260915, + AnthropicBeta::AfkMode20260131, + AnthropicBeta::AgentMemory20260722, + AnthropicBeta::AnalyticsUsageCost20260824, + AnthropicBeta::CacheDiagnosis20260407, + AnthropicBeta::CePlugins20260901, + AnthropicBeta::CeUserManagement20260713, + AnthropicBeta::CodeExecution20250522, + AnthropicBeta::ComputerUse20241022, + AnthropicBeta::DevFullThinking20250514, + AnthropicBeta::Dreaming20260421, + AnthropicBeta::ExperimentalClaudeTag20261021, + AnthropicBeta::ExtendedCacheTtl20250411, + AnthropicBeta::FallbackCredit20260601, + AnthropicBeta::FallbackCredit20260701, + AnthropicBeta::ManagedAgents20260401, + AnthropicBeta::McpClient20260915, + AnthropicBeta::McpTunnels20260622, + AnthropicBeta::MessageBatches20240924, + AnthropicBeta::MessageThreads20260812, + AnthropicBeta::MidConversationSystemClearAt20260821, + AnthropicBeta::ModelContextWindowExceeded20250826, + AnthropicBeta::Output300k20260324, + AnthropicBeta::Pdfs20240925, + AnthropicBeta::PromptCaching20240731, + AnthropicBeta::RedactThinking20260212, + AnthropicBeta::ServerSideFallback20260601, + AnthropicBeta::ServerSideFallback20260701, + AnthropicBeta::SpendLimitReads20260926, + AnthropicBeta::StructuredOutputs20251215, + AnthropicBeta::TaskBudgets20260313, + AnthropicBeta::TelemetryDestinations20260811, + AnthropicBeta::ThinkingTokenCount20260513, + AnthropicBeta::TokenCounting20241101, + AnthropicBeta::UserProfiles20260324, + AnthropicBeta::UserProfiles20260818, + AnthropicBeta::UserProfiles20260904, + AnthropicBeta::Webhooks20261007 )] beta: AnthropicBeta, ) { @@ -483,7 +641,7 @@ mod tests { let config = beta_headers_config(); let config_columns: Vec = config .keys() - .filter(|column| column.as_str() != "description") + .filter(|column| column.as_str() != "description" && !column.ends_with("_header_only")) .cloned() .collect(); let beta_provider_columns: Vec = BetaProvider::VARIANTS @@ -502,12 +660,16 @@ mod tests { let config = beta_headers_config(); for provider in BetaProvider::VARIANTS { for beta in &AnthropicBeta::KNOWN { - let expected = config - .get(provider.as_ref()) - .and_then(serde_json::Value::as_object) - .and_then(|column| column.get(beta.as_str())) - .and_then(serde_json::Value::as_str) - .map(str::to_string); + let cell = |column: &str| { + config + .get(column) + .and_then(serde_json::Value::as_object) + .and_then(|values| values.get(beta.as_str())) + .and_then(serde_json::Value::as_str) + .map(str::to_string) + }; + let expected = + cell(&format!("{provider}_header_only")).or_else(|| cell(provider.as_ref())); let actual = beta.on(*provider).map(|name| name.to_string()); assert_eq!(actual, expected, "provider {provider}, beta {beta}"); if matches!(beta, AnthropicBeta::AdvancedToolUse20251120) @@ -537,7 +699,7 @@ mod tests { BetaProvider::Anthropic, Some(AnthropicBeta::Oauth20250420) )] - #[case::rejected("effort-2025-11-24", BetaProvider::VertexAi, None)] + #[case::rejected("skills-2025-10-02", BetaProvider::VertexAi, None)] #[case::unknown("example-beta-2099-01-01", BetaProvider::Anthropic, None)] fn on_resolves_known_spellings_held_as_other( #[case] raw: &str, @@ -566,12 +728,11 @@ mod tests { #[test] fn unknown_values_are_kept_verbatim() { - let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap(); - assert_eq!( - parsed, - AnthropicBeta::Other("claude-code-20250219".to_string()) + let parsed: AnthropicBeta = "example-beta-2099-01-01".parse().unwrap(); + assert!( + matches!(&parsed, AnthropicBeta::Other(value) if value == "example-beta-2099-01-01") ); - assert_eq!(parsed.to_string(), "claude-code-20250219"); + assert_eq!(parsed.to_string(), "example-beta-2099-01-01"); } #[test] diff --git a/litellm-rust/crates/llms-types/src/recognized.rs b/litellm-rust/crates/llms-types/src/recognized.rs index 58e8dd7a2e7..0545389afbc 100644 --- a/litellm-rust/crates/llms-types/src/recognized.rs +++ b/litellm-rust/crates/llms-types/src/recognized.rs @@ -7,6 +7,16 @@ pub enum Recognized { Unrecognized(Value), } +/// `deserialize_with` for an optional `Recognized` field whose explicit `null` must stay +/// distinguishable from an absent key: `null` becomes `Some(Unrecognized(Null))`. +pub fn present<'de, D, T>(deserializer: D) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, + Recognized: serde::Deserialize<'de>, +{ + as serde::Deserialize<'de>>::deserialize(deserializer).map(Some) +} + impl Recognized { pub fn known(&self) -> Option<&T> { match self { diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 962f61ef4ec..79ee17e8b97 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true test-support = ["litellm-http/test-support"] [dependencies] +litellm-secret-value.workspace = true litellm-llms-types.workspace = true litellm-core-utils.workspace = true litellm-error.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs index a0843621043..7842b478dce 100644 --- a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs +++ b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs @@ -96,7 +96,10 @@ mod tests { } #[rstest] - #[case::azure_ai(BetaProvider::AzureAi, "fast-mode-2026-02-01,example-beta-2099-01-01")] + #[case::azure_ai( + BetaProvider::AzureAi, + "advisor-tool-2026-03-01,example-beta-2099-01-01" + )] #[case::anthropic(BetaProvider::Anthropic, "bash_20241022")] #[case::vertex_ai(BetaProvider::VertexAi, " , ")] fn filter_removes_header_when_nothing_survives( diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index eb796a900f5..164dbce1264 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,4 +1,3 @@ -use litellm_auth::SecretValue; use litellm_core_utils::decline::Decline; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, @@ -10,6 +9,7 @@ use litellm_llms_types::{ }, providers::anthropic::DEFAULT_HEADERS, }; +use litellm_secret_value::SecretValue; use serde::Deserialize; use serde_json::{Map, Value, json}; diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 7991a57353e..0854820358b 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,5 +1,7 @@ use crate::base_llm::messages::context::MessagesModelCapabilities; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{ has_header, header_value, header_values, with_header, without_headers, @@ -152,11 +154,42 @@ pub fn complete_anthropic_url( ) -> String { let api_base = resolve_anthropic_api_base(api_base, env_lookup); - let api_base = api_base.trim_end_matches('/'); - if api_base.ends_with(MESSAGES_PATH) { + ensure_api_route(&api_base, MESSAGES_PATH) +} + +/// Python's `urlsplit(url).path`, for the URLs an `api_base` holds. +fn url_path(url: &str) -> &str { + let before_fragment = url.split('#').next().unwrap_or_default(); + let before_query = before_fragment.split('?').next().unwrap_or_default(); + match before_query.split_once("://") { + Some((_, rest)) => rest.find('/').map_or("", |slash| &rest[slash..]), + None => before_query, + } +} + +/// Python's `_strip_trailing_slashes_keeping_the_scheme`. +fn strip_trailing_slashes(api_base: &str) -> String { + let stripped = api_base.trim_end_matches('/'); + if stripped.ends_with(':') && !stripped.contains('/') && !stripped.contains('?') { return api_base.to_string(); } - format!("{api_base}{MESSAGES_PATH}") + match stripped.split_once('?') { + Some((before, query)) if url_path(stripped).ends_with('/') => { + format!("{}?{query}", before.trim_end_matches('/')) + } + _ => stripped.to_string(), + } +} + +/// Python's `ensure_api_route`: an `api_base` whose path already ends in the route is kept as +/// it is, query string included; anything else gets the route appended to the whole string. +fn ensure_api_route(api_base: &str, route: &str) -> String { + let base = strip_trailing_slashes(api_base); + let route = route.trim_start_matches('/'); + if url_path(&base).ends_with(&format!("/{route}")) { + return base; + } + format!("{}/{route}", strip_trailing_slashes(&base)) } pub fn existing_betas(headers: &[(String, String)]) -> BetaSet { @@ -1880,6 +1913,7 @@ mod tests { supports_reasoning: false, supports_adaptive_thinking: false, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: false, supports_sampling_params: true, @@ -1920,4 +1954,24 @@ mod tests { expected ); } + + #[rstest] + #[case::bare_host("https://api.example.com", "https://api.example.com/v1/messages")] + #[case::trailing_slash("https://api.example.com/", "https://api.example.com/v1/messages")] + #[case::route_present( + "https://api.example.com/v1/messages/", + "https://api.example.com/v1/messages" + )] + #[case::route_and_query( + "https://api.example.com/v1/messages?beta=true", + "https://api.example.com/v1/messages?beta=true" + )] + #[case::query_without_route( + "https://api.example.com/?x=1", + "https://api.example.com?x=1/v1/messages" + )] + #[case::scheme_only("https://", "https:///v1/messages")] + fn the_messages_url_follows_python_ensure_api_route(#[case] api_base: &str, #[case] url: &str) { + assert_eq!(complete_anthropic_url(Some(api_base), &|_| None), url); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs index df1ae13fd28..0ddbc915154 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs @@ -2,8 +2,8 @@ use litellm_llms_types::{ formats::{ chat_completions::ReasoningEffort, messages::{ - EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig, ThinkingConfig, - ThinkingDisplay, + AdaptiveThinking, DisabledThinking, EffortLevel, MessagesOptionalParams, + MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay, }, }, recognized::Recognized, @@ -36,16 +36,17 @@ fn budget_for_effort(budgets: &ThinkingBudgets, effort: ReasoningEffort) -> Opti fn effort_for_budget( budgets: &ThinkingBudgets, - budget_tokens: u64, + budget_tokens: i64, capabilities: &MessagesModelCapabilities, ) -> EffortLevel { - if budget_tokens >= budgets.xhigh && capabilities.effort_tiers.xhigh { + let reaches = |tier: u64| i128::from(budget_tokens) >= i128::from(tier); + if reaches(budgets.xhigh) && capabilities.effort_tiers.xhigh { return EffortLevel::Xhigh; } - if budget_tokens >= budgets.high { + if reaches(budgets.high) { return EffortLevel::High; } - if budget_tokens >= budgets.medium { + if reaches(budgets.medium) { return EffortLevel::Medium; } EffortLevel::Low @@ -171,7 +172,11 @@ fn translate_reasoning_effort( ) else { return Ok(MessagesRequest { params: MessagesOptionalParams { - thinking: None, + thinking: (context.capabilities.thinking_on_by_default + && !context.capabilities.thinking_always_on) + .then(|| { + Recognized::Known(ThinkingConfig::Disabled(DisabledThinking::default())) + }), output_config: None, ..request.params }, @@ -231,32 +236,129 @@ fn drop_disabled_thinking(request: MessagesRequest, context: &ThinkingContext) - } } +/// `_numeric_budget_tokens`: Python reads any int, float (truncated) or int-parseable string +/// as a budget, and nothing else. +fn numeric_budget_tokens(budget: &Recognized) -> Option { + let value = match budget { + Recognized::Known(budget) => return Some(i64::try_from(*budget).unwrap_or(i64::MAX)), + Recognized::Unrecognized(value) => value, + }; + match value { + Value::Number(number) => number + .as_i64() + .or_else(|| number.as_f64().map(|float| float.trunc() as i64)), + Value::String(text) => python_int(text), + Value::Null | Value::Bool(_) | Value::Array(_) | Value::Object(_) => None, + } +} + +/// Python's `int(str)`: surrounding whitespace, a sign, and single underscores between digits. +fn python_int(text: &str) -> Option { + let trimmed = text.trim(); + let digits = trimmed.trim_start_matches(['+', '-']); + let well_formed = trimmed.len() - digits.len() <= 1 + && !digits.is_empty() + && !digits.starts_with('_') + && !digits.ends_with('_') + && !digits.contains("__") + && digits.chars().all(|ch| ch.is_ascii_digit() || ch == '_'); + if !well_formed { + return None; + } + let cleaned: String = trimmed.chars().filter(|ch| *ch != '_').collect(); + cleaned.parse::().ok().or_else(|| { + Some(if cleaned.starts_with('-') { + i64::MIN + } else { + i64::MAX + }) + }) +} + +fn python_type_name(value: &Value) -> &'static str { + match value { + Value::Null => "NoneType", + Value::Bool(_) => "bool", + Value::Number(number) if number.is_f64() => "float", + Value::Number(_) => "int", + Value::String(_) => "str", + Value::Array(_) => "list", + Value::Object(_) => "dict", + } +} + +fn kept_display( + display: Option<&Recognized>, + context: &ThinkingContext, +) -> Option> { + match display? { + Recognized::Known(ThinkingDisplay::Summarized | ThinkingDisplay::Omitted) => { + display.cloned() + } + Recognized::Known(ThinkingDisplay::Updates) if context.keeps_display_updates => { + display.cloned() + } + Recognized::Known(ThinkingDisplay::Updates) | Recognized::Unrecognized(_) => None, + } +} + +/// Keys Python's `carried_thinking_keys` keeps on the rebuilt adaptive config: the adaptive +/// keys always, and any other key unless `drop_params` is set. +fn carried_thinking_keys( + extra: &serde_json::Map, + context: &ThinkingContext, +) -> serde_json::Map { + extra + .iter() + .filter(|(key, _)| key.as_str() == "block_binding" || !context.drop_params) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} + +/// `AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model_silently`. fn translate_legacy_thinking_for_adaptive_model( request: MessagesRequest, context: &ThinkingContext, -) -> MessagesRequest { +) -> Result { let capabilities = &context.capabilities; if !capabilities.supports_adaptive_thinking || capabilities.supports_legacy_thinking { - return request; + return Ok(request); } let Some(ThinkingConfig::Enabled(enabled)) = known_thinking(&request) else { - return request; + return Ok(request); }; - let budget = enabled - .budget_tokens - .as_ref() - .and_then(Recognized::known) - .copied() - .unwrap_or(0); - let level = effort_for_budget(&context.budgets, budget, capabilities); - MessagesRequest { + let budget = match &enabled.budget_tokens { + None => Some(0), + Some(budget) => numeric_budget_tokens(budget), + }; + if budget.is_none() && !context.drop_params { + let actual = match &enabled.budget_tokens { + Some(Recognized::Unrecognized(value)) => python_type_name(value), + Some(Recognized::Known(_)) | None => "int", + }; + return Err(Error::InvalidRequest( + crate::ErrorDetail::NonNumericThinkingBudget { actual }, + )); + } + let thinking = ThinkingConfig::Adaptive(AdaptiveThinking { + display: kept_display(enabled.display.as_ref(), context), + extra: carried_thinking_keys(&enabled.extra, context), + }); + let output_config = match budget { + Some(budget) => with_default_effort( + request.params.output_config, + effort_for_budget(&context.budgets, budget, capabilities), + ), + None => request.params.output_config, + }; + Ok(MessagesRequest { params: MessagesOptionalParams { - thinking: Some(Recognized::Known(ThinkingConfig::adaptive(None))), - output_config: with_default_effort(request.params.output_config, level), + thinking: Some(Recognized::Known(thinking)), + output_config, ..request.params }, ..request - } + }) } fn translate_adaptive_effort_for_non_adaptive_model( @@ -337,7 +439,7 @@ pub fn translate_thinking( ) -> Result { let request = translate_reasoning_effort(request, context)?; let request = drop_disabled_thinking(request, context); - let request = translate_legacy_thinking_for_adaptive_model(request, context); + let request = translate_legacy_thinking_for_adaptive_model(request, context)?; let request = translate_adaptive_effort_for_non_adaptive_model(request, context)?; Ok(drop_incompatible_temperature_for_thinking(request, context)) } @@ -363,6 +465,7 @@ mod tests { ThinkingContext { capabilities, budgets: ThinkingBudgets::default(), + ..ThinkingContext::default() } } @@ -649,6 +752,33 @@ mod tests { ); } + #[rstest] + #[case::default_thinking(true, false, Some(serde_json::json!({"type": "disabled"}))) ] + #[case::legacy(false, false, None)] + #[case::always_on(true, true, None)] + fn none_disables_only_optional_default_thinking( + #[case] thinking_on_by_default: bool, + #[case] thinking_always_on: bool, + #[case] thinking: Option, + ) { + let capabilities = MessagesModelCapabilities { + thinking_on_by_default, + thinking_always_on, + ..opus_4_7() + }; + let expected = match thinking { + Some(thinking) => serde_json::json!({"max_tokens": 1024, "thinking": thinking}), + None => serde_json::json!({"max_tokens": 1024}), + }; + assert_eq!( + translate( + capabilities, + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "none"}) + ), + Ok(request(expected)) + ); + } + #[rstest] #[case::adaptive_model(opus_4_7())] #[case::effort_capable_model(opus_4_5())] @@ -1239,10 +1369,102 @@ mod tests { let context = ThinkingContext { capabilities, budgets: overridden_budgets(overrides), + ..ThinkingContext::default() }; assert_eq!( translate_thinking(request(input), &context), Ok(request(expected)) ); } + + use serde_json::json; + + fn legacy_context(drop_params: bool, keeps_display_updates: bool) -> ThinkingContext { + ThinkingContext { + drop_params, + keeps_display_updates, + ..context(opus_4_7()) + } + } + + #[rstest] + #[case::float_budget(json!(8192.0), "xhigh")] + #[case::float_below_a_tier(json!(1023.9), "low")] + #[case::string_budget(json!(" 16_000 "), "xhigh")] + #[case::negative_budget(json!(-5), "low")] + fn legacy_budgets_read_like_python_numbers(#[case] budget: Value, #[case] effort: &str) { + let translated = translate_thinking( + request( + json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": budget}}), + ), + &legacy_context(false, false), + ); + assert_eq!( + translated, + Ok(request(json!({ + "max_tokens": 64000, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": effort}, + }))) + ); + } + + #[rstest] + #[case::text(json!("lots"), "str")] + #[case::list(json!([1]), "list")] + #[case::null(Value::Null, "NoneType")] + #[case::boolean(json!(true), "bool")] + fn a_non_numeric_legacy_budget_is_rejected_unless_dropped( + #[case] budget: Value, + #[case] actual: &'static str, + ) { + let fields = + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": budget}}); + assert_eq!( + translate_thinking(request(fields.clone()), &legacy_context(false, false)), + Err(Error::InvalidRequest( + crate::ErrorDetail::NonNumericThinkingBudget { actual } + )) + ); + assert_eq!( + translate_thinking(request(fields), &legacy_context(true, false)), + Ok(request( + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + )) + ); + } + + #[rstest] + #[case::summarized("summarized", false, Some("summarized"))] + #[case::updates_kept("updates", true, Some("updates"))] + #[case::updates_dropped("updates", false, None)] + #[case::unknown("verbose", true, None)] + fn legacy_display_survives_only_where_python_keeps_it( + #[case] display: &str, + #[case] keeps_display_updates: bool, + #[case] kept: Option<&str>, + ) { + let translated = translate_thinking( + request(json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 2048, "display": display}})), + &legacy_context(false, keeps_display_updates), + ) + .unwrap(); + let thinking = serde_json::to_value(translated.params.thinking).unwrap(); + assert_eq!(thinking.get("display").and_then(Value::as_str), kept); + } + + #[rstest] + #[case::kept(false, json!({"type": "adaptive", "block_binding": "strict", "note": "x"}))] + #[case::dropped(true, json!({"type": "adaptive", "block_binding": "strict"}))] + fn legacy_extra_keys_follow_drop_params(#[case] drop_params: bool, #[case] expected: Value) { + let translated = translate_thinking( + request(json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 2048, "block_binding": "strict", "note": "x"}})), + &legacy_context(drop_params, false), + ) + .unwrap(); + assert_eq!( + serde_json::to_value(translated.params.thinking).unwrap(), + expected + ); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 6657cbf262f..18789e5b2be 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -1,7 +1,8 @@ use litellm_auth::CredentialPlacement; use litellm_llms_types::{ formats::messages::{ - ContextEdit, ContextManagement, Message, MessagesOptionalParams, MessagesRequest, Speed, + ContentBlock, ContentBlockType, ContextEdit, ContextManagement, Message, + MessagesOptionalParams, MessagesRequest, Speed, ThinkingConfig, ThinkingDisplay, }, providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS}, recognized::Recognized, @@ -55,7 +56,7 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - transform_messages_request(request, context) + transform_messages_request(request, &context.keeping_display_updates()) } fn secret_names(&self) -> &'static [&'static str] { @@ -203,6 +204,12 @@ fn feature_betas(request: &MessagesRequest) -> BetaSet { .then_some(AnthropicBeta::FastMode20260201), messages_carry_output_config(&request.messages) .then_some(AnthropicBeta::PerTurnControl20260701), + uses_thinking_display_updates(params) + .then_some(AnthropicBeta::ThinkingDisplayUpdates20260818), + uses_thinking_block_binding(params) + .then_some(AnthropicBeta::ThinkingBindingControls20260801), + messages_carry_tool_changes(&request.messages) + .then_some(AnthropicBeta::MidConversationToolChanges20260701), has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301), is_tool_search_used(tools).then_some(AnthropicBeta::AdvancedToolUse20251120), ] @@ -253,6 +260,55 @@ fn messages_carry_output_config(messages: &[Message]) -> bool { .any(|message| message.extra.contains_key("output_config")) } +fn uses_thinking_display_updates(params: &MessagesOptionalParams) -> bool { + let display = match params.thinking.as_ref().and_then(Recognized::known) { + Some(ThinkingConfig::Enabled(thinking)) => thinking.display.as_ref(), + Some(ThinkingConfig::Adaptive(thinking)) => thinking.display.as_ref(), + _ => None, + }; + display.and_then(Recognized::known) == Some(&ThinkingDisplay::Updates) +} + +/// `AnthropicModelInfo.thinking_betas`: any thinking object carrying `block_binding`. +fn uses_thinking_block_binding(params: &MessagesOptionalParams) -> bool { + match params.thinking.as_ref() { + Some(Recognized::Known(ThinkingConfig::Enabled(thinking))) => { + thinking.extra.contains_key("block_binding") + } + Some(Recognized::Known(ThinkingConfig::Adaptive(thinking))) => { + thinking.extra.contains_key("block_binding") + } + Some(Recognized::Known(ThinkingConfig::Disabled(thinking))) => { + thinking.extra.contains_key("block_binding") + } + Some(Recognized::Unrecognized(Value::Object(thinking))) => { + thinking.contains_key("block_binding") + } + Some(Recognized::Unrecognized(_)) | None => false, + } +} + +fn is_tool_reference_change(block: &ContentBlock) -> bool { + matches!( + &block.block_type, + Some(ContentBlockType::Other(block_type)) + if block_type == "tool_addition" || block_type == "tool_removal" + ) && block + .extra + .get("tool") + .and_then(|tool| tool.get("type")) + .and_then(Value::as_str) + == Some("tool_reference") +} + +fn messages_carry_tool_changes(messages: &[Message]) -> bool { + messages + .iter() + .filter(|message| message.role == "system") + .flat_map(Message::blocks) + .any(is_tool_reference_change) +} + fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error { Error::InvalidRequest(crate::ErrorDetail::UnsupportedParameter { model: model.into(), @@ -262,7 +318,8 @@ fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error }) } -fn drop_unsupported_params( +/// The speed and sampling gates of Python's `get_requested_anthropic_messages_optional_param`. +pub fn drop_unsupported_params( request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { @@ -803,6 +860,8 @@ mod tests { low: 2000, ..ThinkingBudgets::default() }, + drop_params: true, + keeps_display_updates: false, }, drop_params: true, } @@ -1071,6 +1130,41 @@ mod tests { json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}), &[] )] + #[case::adaptive_display_updates( + json!({"thinking": {"type": "adaptive", "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::enabled_display_updates( + json!({"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::summarized_display(json!({"thinking": {"type": "adaptive", "display": "summarized"}}), &[])] + #[case::disabled_thinking(json!({"thinking": {"type": "disabled"}}), &[])] + #[case::system_tool_addition( + json!({"messages": [ + {"role": "user", "content": "Use the ping tool"}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] + #[case::system_tool_removal( + json!({"messages": [ + {"role": "system", "content": [{"type": "tool_removal", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] + #[case::user_tool_addition( + json!({"messages": [ + {"role": "user", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ]}), + &[] + )] + #[case::tool_addition_without_a_tool_reference( + json!({"messages": [ + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_definition"}}]}, + ]}), + &[] + )] fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { assert_eq!(feature_betas(&request(fields)), betas(expected)); } @@ -1093,6 +1187,33 @@ mod tests { ); } + #[rstest] + #[case::vertex_ai( + BetaProvider::VertexAi, + &["mid-conversation-tool-changes-2026-07-01", "thinking-display-updates-2026-08-18"] + )] + #[case::azure_ai( + BetaProvider::AzureAi, + &["mid-conversation-tool-changes-2026-07-01", "thinking-display-updates-2026-08-18"] + )] + #[case::bedrock_converse(BetaProvider::BedrockConverse, &["thinking-display-updates-2026-08-18"])] + #[case::databricks(BetaProvider::Databricks, &["thinking-display-updates-2026-08-18"])] + fn tool_change_and_display_betas_reach_the_providers_that_accept_them( + #[case] provider: BetaProvider, + #[case] expected: &[&str], + ) { + let fields = json!({ + "thinking": {"type": "adaptive", "display": "updates"}, + "messages": [ + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ], + }); + assert_eq!( + provider_feature_betas(&request(fields), provider), + betas(expected) + ); + } + #[rstest] fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() { let client = headers(&[("anthropic-beta", "claude-code-20250219")]); diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index bf934fb4d4c..d13cfa81a94 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -1,10 +1,10 @@ -use litellm_auth::SecretValue; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; +use litellm_secret_value::SecretValue; use crate::{ Error, @@ -68,7 +68,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { .capabilities .supports_mid_conversation_system, )), - context, + &context.keeping_display_updates(), ) } @@ -545,6 +545,11 @@ mod tests { &[("anthropic-beta", "web-search-2025-03-05")], &[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")] )] + #[case::compaction_needs_the_compact_2026_09_04_beta( + json!({"compaction": {"type": "summarize"}}), + &[], + &[("x-api-key", "k"), ("anthropic-beta", "compact-2026-09-04")] + )] #[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])] #[case::safeguards_need_the_dangerous_tool_use_beta( json!({"safeguards": [{"type": "dangerous_tool_use"}]}), diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 36dfe2c890e..870306b1a83 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -1250,7 +1250,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 2f7957fe56e..25f93846967 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -211,7 +211,7 @@ mod tests { #[fixture] fn connection() -> OcrConnection { OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_base: Some("https://example.com".into()), ..Default::default() } @@ -312,7 +312,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/base_llm/auth.rs b/litellm-rust/crates/llms/src/base_llm/auth.rs index de271a269d3..ab2cf00614e 100644 --- a/litellm-rust/crates/llms/src/base_llm/auth.rs +++ b/litellm-rust/crates/llms/src/base_llm/auth.rs @@ -5,7 +5,9 @@ //! that split: `validate_environment` shapes the forwarded headers and names the credential //! as an [`AuthScheme`], and [`resolve_auth`] turns the scheme into headers and a signer. -use litellm_auth::{AuthServices, CredentialPlacement, SecretValue, TokenProviderHandle}; +use litellm_auth::{AuthServices, CredentialPlacement, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{AwsCredentialSource, SigV4Signer}; use litellm_auth_gcp::VertexConfig; use litellm_http::request::with_header; diff --git a/litellm-rust/crates/llms/src/base_llm/messages/context.rs b/litellm-rust/crates/llms/src/base_llm/messages/context.rs index e59a63813b0..3d8b0f41d8c 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/context.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/context.rs @@ -32,6 +32,8 @@ pub struct MessagesModelCapabilities { #[serde(default)] pub thinking_always_on: bool, #[serde(default)] + pub thinking_on_by_default: bool, + #[serde(default)] pub supports_legacy_thinking: bool, #[serde(default)] pub supports_output_config: bool, @@ -55,6 +57,7 @@ impl Default for MessagesModelCapabilities { supports_reasoning: false, supports_adaptive_thinking: false, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: false, supports_sampling_params: true, @@ -110,6 +113,11 @@ impl ThinkingBudgets { pub struct ThinkingContext { pub capabilities: MessagesModelCapabilities, pub budgets: ThinkingBudgets, + /// `drop_params` as Python's thinking translation reads it. + pub drop_params: bool, + /// The provider accepts `display: "updates"` on thinking translated from the legacy shape + /// (Python's `LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS`). + pub keeps_display_updates: bool, } #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub struct MessagesTransformContext { @@ -130,8 +138,22 @@ impl MessagesTransformContext { thinking: ThinkingContext { capabilities, budgets: ThinkingBudgets::from_lookup(env), + drop_params, + keeps_display_updates: false, }, drop_params, } } + + /// The same context for a provider in Python's `LEGACY_THINKING_DISPLAY_UPDATES_PROVIDERS`. + #[must_use] + pub fn keeping_display_updates(&self) -> Self { + Self { + thinking: ThinkingContext { + keeps_display_updates: true, + ..self.thinking + }, + ..*self + } + } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index c90fd99809e..9574e99d357 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -5,7 +5,9 @@ use serde::{ use serde_json::{Map, Value}; use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; -use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; +use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; use litellm_http::outbound::RequestSigner; use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index 976c9161dd2..6cf0167a30f 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -1,5 +1,7 @@ use litellm_auth::AwsParams; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 79c30df9b19..5b89345a1c5 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -1,6 +1,8 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region}; use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{header_value, with_header, without_headers}; diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 3ec60f12021..1d380c917f6 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -7,7 +7,9 @@ use crate::{ }; use futures_util::StreamExt; use litellm_auth::AwsParams; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{ diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs index ae13704f4a0..f139b93e259 100644 --- a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -1,6 +1,8 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::AwsCredentialSource; use litellm_auth_aws::bedrock_model_id_and_region; use litellm_core_utils::settings::resolve_non_empty; diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index 76ff27b56fc..5bad06adda2 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -735,7 +735,7 @@ mod tests { assert!(matches!( CohereParseConfig.resolve_headers( &OcrConnection { - api_key: Some(litellm_auth::SecretValue::new(" ")), + api_key: Some(litellm_secret_value::SecretValue::new(" ")), ..Default::default() }, &|_| None, diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 168c3e63920..0dbcea01439 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool}, diff --git a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs index 2c3a7eaef1d..d5944484522 100644 --- a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::formats::messages::{ MessagesOptionalParams, MessagesRequest, MessagesResponse, diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index b3d1421083a..2975ceee3cb 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -92,6 +92,31 @@ impl Error { } } + pub fn upstream_headers(&self) -> &[(String, String)] { + match self { + Self::Http(error) => error.upstream_headers(), + _ => &[], + } + } + + pub fn is_body_read_timeout(&self) -> bool { + matches!(self, Self::Http(error) if error.is_body_read_timeout()) + } + + pub fn body_cut(&self) -> Option { + match self { + Self::Http(error) => error.body_cut(), + _ => None, + } + } + + pub fn is_timeout(&self) -> Option { + match self { + Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()), + _ => None, + } + } + pub fn upstream_response(&self) -> Option<(u16, &str)> { match self { Self::Http(error) => error.upstream_response(), @@ -158,6 +183,10 @@ pub enum ErrorDetail { value: String, hint: String, }, + #[error( + "thinking.budget_tokens must be a number, got {actual}. To drop it, set `litellm.drop_params=True`, or `drop_params: true` under `litellm_settings` on the proxy" + )] + NonNumericThinkingBudget { actual: &'static str }, #[error( "Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({{LLM_PROVIDER}}_API_KEY) or via params" )] diff --git a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs index 012a657fceb..dafe5fbd16b 100644 --- a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index fa4e2254e86..679aa65f7fd 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -219,7 +219,7 @@ mod tests { #[default(vec![])] extra_headers: Vec<(String, String)>, ) -> OcrConnection { OcrConnection { - api_key: api_key.map(litellm_auth::SecretValue::new), + api_key: api_key.map(litellm_secret_value::SecretValue::new), extra_headers, ..OcrConnection::default() } diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index cf74871894b..deb163daa07 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -3,7 +3,8 @@ use litellm_llms_types::formats::responses::{ }; use serde_json::{Map, Value}; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; use crate::{ Error, ErrorDetail, diff --git a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs index e998f8785cb..83fb8ac3fa2 100644 --- a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs @@ -4,7 +4,9 @@ //! the `max_completion_tokens` -> `max_tokens` rename and the usage //! `*_tokens` null-to-zero sanitize. -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::core_helpers::unix_now; use litellm_llms_types::formats::chat_completions::{ ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, diff --git a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs index 5b4e4f77f2a..eeb3d1cb7f8 100644 --- a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{has_bearer_auth, without_headers}; use litellm_llms_types::{ diff --git a/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs b/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs index 900fba9e916..76155ba55e5 100644 --- a/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::formats::responses::ResponsesApiResponse; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index bf0ddc1c7a8..beda6899d04 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -690,7 +690,7 @@ mod tests { #[test] fn explicit_key_precedes_environment_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("passed-key")), + api_key: Some(litellm_secret_value::SecretValue::new("passed-key")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); @@ -700,7 +700,7 @@ mod tests { #[test] fn blank_explicit_key_uses_environment_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new(" ")), + api_key: Some(litellm_secret_value::SecretValue::new(" ")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index a70336748ad..621a09f4fb0 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index 46e5a68dcb8..038617c2e10 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -77,7 +77,7 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - let request = transform_messages_request(request, context)?; + let request = transform_messages_request(request, &context.keeping_display_updates())?; let request = strip_cache_control_scope(normalize_system_role_messages( request, context @@ -798,6 +798,18 @@ mod tests { ]}}), &["compact-2026-01-12", "context-management-2025-06-27"] )] + #[case::compaction(json!({"compaction": {"type": "summarize"}}), &["compact-2026-09-04"])] + #[case::thinking_display_updates( + json!({"thinking": {"type": "adaptive", "display": "updates"}}), + &["thinking-display-updates-2026-08-18"] + )] + #[case::tool_addition( + json!({"messages": [ + {"role": "user", "content": "Use the ping tool"}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]} + ]}), + &["mid-conversation-tool-changes-2026-07-01"] + )] #[case::advisor_tool_is_not_a_vertex_beta( json!({"tools": [{"type": "advisor_20260301", "name": "advisor"}]}), &[] diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index 77881cd2e78..b902967d9dc 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -137,7 +137,7 @@ impl VertexAiOcrConfig { connection .api_key .as_ref() - .map(litellm_auth::SecretValue::expose), + .map(litellm_secret_value::SecretValue::expose), config, &|name: &str| connection.secret(name), ) diff --git a/litellm-rust/crates/model-catalog-refresh/AGENTS.md b/litellm-rust/crates/model-catalog-refresh/AGENTS.md new file mode 100644 index 00000000000..96228716abf --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/AGENTS.md @@ -0,0 +1,9 @@ +# Model catalog refresh + +- `fetch` fetches one document over http(s) or `s3://` and never parses it +- Only 429, 500, 502, 503, 504 and transport errors are retried. S3 errors are final +- `Failed.retryable` means only the attempt limit ended the fetch +- A server's Retry-After overrides the backoff, capped at `max_delay` +- `Start::After(failed)` resumes the backoff schedule where another caller's first attempt left off +- Error reasons use the URL with userinfo, query and fragment stripped, since presigned URLs carry credentials +- Bodies over `MAX_BODY_BYTES` are rejected diff --git a/litellm-rust/crates/model-catalog-refresh/Cargo.toml b/litellm-rust/crates/model-catalog-refresh/Cargo.toml index 27219042c6c..76f4a9fb969 100644 --- a/litellm-rust/crates/model-catalog-refresh/Cargo.toml +++ b/litellm-rust/crates/model-catalog-refresh/Cargo.toml @@ -6,14 +6,19 @@ license.workspace = true repository.workspace = true [dependencies] +backon.workspace = true +bytes.workspace = true +litellm-error.workspace = true litellm-http.workspace = true -litellm-model-catalog.workspace = true +litellm-s3.workspace = true +litellm-tracing.workspace = true +reqwest.workspace = true thiserror.workspace = true +tokio = { workspace = true, features = ["time"] } url.workspace = true [dev-dependencies] -litellm-error.workspace = true +aws-credential-types = "1.3.0" litellm-http = { workspace = true, features = ["test-support"] } -reqwest.workspace = true rstest.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } diff --git a/litellm-rust/crates/model-catalog-refresh/src/attempt.rs b/litellm-rust/crates/model-catalog-refresh/src/attempt.rs new file mode 100644 index 00000000000..45d80e36a18 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/attempt.rs @@ -0,0 +1,153 @@ +use std::time::Duration; + +use bytes::Bytes; + +use litellm_http::Client; +use litellm_s3::S3Uri; +use litellm_s3::sdk::config::retry::RetryConfig; +use litellm_s3::sdk::config::timeout::TimeoutConfig; +use reqwest::header::{ETAG, HeaderMap, RETRY_AFTER}; +use url::Url; + +use crate::{Error, RETRYABLE_STATUSES, S3Access}; + +pub(crate) const MAX_BODY_BYTES: u64 = 100 * 1024 * 1024; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Fetched { + pub body: Vec, + pub etag: Option, +} + +#[derive(Debug)] +pub(crate) enum Cause { + Status(u16), + Transport(Error), +} + +#[derive(Debug)] +pub(crate) struct Rejected { + pub(crate) cause: Cause, + pub(crate) retryable: bool, + pub(crate) retry_after: Option, +} + +impl Rejected { + pub(crate) fn fatal(error: Error) -> Self { + Self { + cause: Cause::Transport(error), + retryable: false, + retry_after: None, + } + } +} + +/// Any status below 400 is the document +pub(crate) async fn http( + client: &Client, + url: Url, + timeout: Duration, +) -> Result { + let transport = |error: litellm_http::Error| Rejected { + cause: Cause::Transport(Error::from(error)), + retryable: true, + retry_after: None, + }; + let mut response = tokio::time::timeout(timeout, client.get(url).send()) + .await + .map_err(litellm_http::Error::read_timeout) + .and_then(|sent| sent.map_err(litellm_http::Error::from)) + .map_err(transport)?; + let status = response.status().as_u16(); + let etag = etag(response.headers()); + let retry_after = response + .headers() + .get(RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.trim().parse().ok()) + .map(Duration::from_secs); + if status >= 400 { + return Err(Rejected { + cause: Cause::Status(status), + retryable: RETRYABLE_STATUSES.contains(&status), + retry_after, + }); + } + if response + .content_length() + .is_some_and(|length| length > MAX_BODY_BYTES) + { + return Err(Rejected::fatal(Error::body_too_large(MAX_BODY_BYTES))); + } + let mut body = Vec::new(); + while let Some(chunk) = tokio::time::timeout(timeout, response.chunk()) + .await + .map_err(litellm_http::Error::body_read_timeout) + .and_then(|read| read.map_err(litellm_http::Error::from)) + .map_err(transport)? + { + append(&mut body, &chunk).map_err(Rejected::fatal)?; + } + Ok(Fetched { body, etag }) +} + +/// S3 failures are final: the SDK already classifies its own errors and a missing or denied object stays so +pub(crate) async fn s3( + client: &Client, + s3: &S3Access, + object: &S3Uri, + timeout: Duration, +) -> Result { + let fetched = async { + let config = match s3 { + S3Access::Environment { region } => { + litellm_s3::environment_config(client.clone(), region.clone()).await + } + S3Access::Config(config) => config.to_builder(), + } + .retry_config(RetryConfig::disabled()) + .timeout_config( + TimeoutConfig::builder() + .operation_attempt_timeout(timeout) + .build(), + ) + .build(); + let output = litellm_s3::sdk::Client::from_conf(config) + .get_object() + .bucket(&object.bucket) + .key(&object.key) + .send() + .await?; + if output + .content_length() + .is_some_and(|length| length.unsigned_abs() > MAX_BODY_BYTES) + { + return Err(Error::body_too_large(MAX_BODY_BYTES)); + } + let etag = output.e_tag().map(str::to_owned); + let mut stream = output.body; + let mut body = Vec::new(); + while let Some(chunk) = tokio::time::timeout(timeout, stream.try_next()).await?? { + append(&mut body, &chunk)?; + } + Ok::<_, Error>(Fetched { body, etag }) + }; + fetched.await.map_err(Rejected::fatal) +} + +fn append(body: &mut Vec, chunk: &Bytes) -> Result<(), Error> { + if (body.len() + chunk.len()) as u64 > MAX_BODY_BYTES { + return Err(Error::body_too_large(MAX_BODY_BYTES)); + } + body.extend_from_slice(chunk); + Ok(()) +} + +fn etag(headers: &HeaderMap) -> Option { + let values: Vec<&str> = headers + .get_all(ETAG) + .iter() + .filter_map(|value| value.to_str().ok()) + .collect(); + (!values.is_empty()).then(|| values.join(", ")) +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/error.rs b/litellm-rust/crates/model-catalog-refresh/src/error.rs index 656f32370cb..c79daeb4777 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/error.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/error.rs @@ -1,6 +1,9 @@ -use std::io; - +use litellm_s3::sdk::config::http::HttpResponse; +use litellm_s3::sdk::error::SdkError; +use litellm_s3::sdk::operation::get_object::GetObjectError; +use litellm_s3::sdk::primitives::ByteStreamError; use thiserror::Error; +use tokio::time::error::Elapsed; #[derive(Debug, Error)] #[error(transparent)] @@ -10,18 +13,48 @@ pub struct Error(Box); enum Repr { #[error(transparent)] Http(#[from] litellm_http::Error), - #[error("could not read the catalog")] - Read(#[source] io::Error), #[error("invalid catalog URL")] InvalidUrl(#[source] url::ParseError), #[error("unsupported catalog URL scheme {scheme}")] UnsupportedScheme { scheme: String }, + #[error("invalid S3 URL")] + S3Uri(#[source] litellm_s3::Error), + #[error("S3 GetObject failed")] + S3GetObject(#[source] Box>), + #[error("S3 GetObject body read failed")] + S3Body(#[source] ByteStreamError), + #[error("S3 GetObject timed out")] + S3Timeout(#[source] Elapsed), + #[error("response body is over the {limit} byte limit")] + BodyTooLarge { limit: u64 }, } impl Error { pub fn unsupported_scheme(scheme: String) -> Self { Self(Box::new(Repr::UnsupportedScheme { scheme })) } + + pub(crate) fn body_too_large(limit: u64) -> Self { + Self(Box::new(Repr::BodyTooLarge { limit })) + } +} + +impl From> for Error { + fn from(cause: SdkError) -> Self { + Self(Box::new(Repr::S3GetObject(Box::new(cause)))) + } +} + +impl From for Error { + fn from(cause: ByteStreamError) -> Self { + Self(Box::new(Repr::S3Body(cause))) + } +} + +impl From for Error { + fn from(cause: Elapsed) -> Self { + Self(Box::new(Repr::S3Timeout(cause))) + } } impl From for Error { @@ -30,9 +63,9 @@ impl From for Error { } } -impl From for Error { - fn from(cause: io::Error) -> Self { - Self(Box::new(Repr::Read(cause))) +impl From for Error { + fn from(cause: litellm_s3::Error) -> Self { + Self(Box::new(Repr::S3Uri(cause))) } } @@ -44,29 +77,18 @@ impl From for Error { #[cfg(test)] mod tests { + use std::io; + use litellm_error::{Report, find}; use rstest::rstest; use super::*; - use crate::{DoneReason, Reply, Step}; #[rstest] - fn exhausted_transport_keeps_its_typed_cause() { - let reply = Reply::Transport { - error: io::Error::from(io::ErrorKind::ConnectionReset).into(), - }; - let Reply::Transport { error } = reply else { - panic!("transport failure") - }; - let step = Step::Done { - reason: DoneReason::Exhausted(error), - }; - let Step::Done { - reason: DoneReason::Exhausted(error), - } = step - else { - panic!("exhausted attempts") - }; + fn an_s3_failure_keeps_its_typed_cause_once() { + let error = Error::from(SdkError::::timeout_error( + io::Error::from(io::ErrorKind::ConnectionReset), + )); let cause = find::(&error).expect("typed I/O cause"); assert_eq!(cause.kind(), io::ErrorKind::ConnectionReset); assert_eq!( diff --git a/litellm-rust/crates/model-catalog-refresh/src/fetch.rs b/litellm-rust/crates/model-catalog-refresh/src/fetch.rs new file mode 100644 index 00000000000..bd58827b546 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/fetch.rs @@ -0,0 +1,122 @@ +use std::borrow::Cow; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::Duration; + +use backon::Retryable; +use litellm_error::Report; +use litellm_http::Client; +use url::Url; + +use crate::attempt::{self, Cause, Fetched, Rejected}; +use crate::{RetryPolicy, Source}; + +/// How an `s3://` fetch reaches S3 +#[derive(Clone, Debug)] +pub enum S3Access { + /// What an AWS SDK finds in the environment and shared config files, loaded on each `s3://` attempt, with + /// `region` winning over the region it finds + Environment { region: Option }, + /// A config the caller built with its region, endpoint and credentials + Config(Box), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Failed { + pub reason: String, + /// Another attempt could succeed. Only the attempt limit ended the fetch + pub retryable: bool, + pub retry_after: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Start { + Fresh, + /// Continues a fetch whose first attempt failed retryably elsewhere, waiting before its next attempt + After(Failed), +} + +pub async fn fetch( + client: &Client, + s3: &S3Access, + url: &str, + policy: RetryPolicy, + start: Start, +) -> Result { + let shown = redacted(url); + let source = + Source::parse(url).map_err(|error| Failed::of(&shown, &Rejected::fatal(error), 1))?; + let mut backoff = policy.backoff(); + let made = AtomicU32::new(0); + if let Start::After(failed) = start { + let Some(wait) = backoff.next() else { + return Err(failed); + }; + tokio::time::sleep(policy.delay(wait, failed.retry_after)).await; + made.store(1, Ordering::Relaxed); + } + let attempt = || { + made.fetch_add(1, Ordering::Relaxed); + async { + match &source { + Source::Http(http) => attempt::http(client, http.clone(), policy.timeout).await, + Source::S3(object) => attempt::s3(client, s3, object, policy.timeout).await, + } + } + }; + attempt + .retry(backoff) + .sleep(tokio::time::sleep) + .when(|rejected| rejected.retryable) + .adjust(|rejected, wait| wait.map(|wait| policy.delay(wait, rejected.retry_after))) + .notify(|rejected, wait| { + litellm_tracing::warn!( + "LiteLLM: model cost map fetch attempt {}/{} failed ({}); retrying in {:.1}s", + made.load(Ordering::Relaxed), + policy.max_attempts, + reason(&shown, &rejected.cause), + wait.as_secs_f64(), + ); + }) + .await + .map_err(|rejected| Failed::of(&shown, &rejected, made.load(Ordering::Relaxed))) +} + +impl Failed { + fn of(url: &str, rejected: &Rejected, attempts: u32) -> Self { + let reason = reason(url, &rejected.cause); + Self { + reason: match rejected.retryable && attempts > 1 { + true => format!("{reason} (after {attempts} attempts)"), + false => reason, + }, + retryable: rejected.retryable, + retry_after: rejected.retry_after, + } + } +} + +fn reason(url: &str, cause: &Cause) -> String { + match cause { + Cause::Status(status) => format!("HTTP {status} from {url}"), + Cause::Transport(error) => format!("error fetching {url}: {}", Report::new(error)), + } +} + +/// Drops userinfo, query and fragment, where presigned URLs and tokens carry credentials +fn redacted(url: &str) -> Cow<'_, str> { + let Ok(mut parsed) = Url::parse(url) else { + return Cow::Borrowed(url); + }; + if parsed.username().is_empty() + && parsed.password().is_none() + && parsed.query().is_none() + && parsed.fragment().is_none() + { + return Cow::Borrowed(url); + } + let _ = parsed.set_username(""); + let _ = parsed.set_password(None); + parsed.set_query(None); + parsed.set_fragment(None); + Cow::Owned(parsed.into()) +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/lib.rs b/litellm-rust/crates/model-catalog-refresh/src/lib.rs index 006c4887fad..1742a669483 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/lib.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/lib.rs @@ -1,9 +1,11 @@ +mod attempt; mod error; +mod fetch; mod policy; mod source; -mod step; +pub use attempt::Fetched; pub use error::Error; +pub use fetch::*; pub use policy::*; pub use source::*; -pub use step::*; diff --git a/litellm-rust/crates/model-catalog-refresh/src/policy.rs b/litellm-rust/crates/model-catalog-refresh/src/policy.rs index bc32e40408e..1cdb737ecd8 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/policy.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/policy.rs @@ -1,20 +1,76 @@ use std::time::Duration; -use litellm_model_catalog::IntegrityLimits; +use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder}; -/// Main's defaults, 3 attempts, 30 s max wait, 5 s timeout, backoff base 2.0, jitter 0.0 to 1.0 s -#[derive(Clone, Copy, Debug, PartialEq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct RetryPolicy { pub max_attempts: u32, - pub max_wait: Duration, pub timeout: Duration, - pub backoff_base_seconds: f64, - pub jitter_min_seconds: f64, - pub jitter_max_seconds: f64, - /// Main's 50 model floor and 0.5 shrink ratio - pub limits: IntegrityLimits, + pub min_delay: Duration, + pub max_delay: Duration, } +pub const RETRYABLE_STATUSES: [u16; 5] = [429, 500, 502, 503, 504]; + /// Main's default `LITELLM_MODEL_COST_MAP_URL` pub const DEFAULT_REGISTRY_URL: &str = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"; + +impl RetryPolicy { + pub(crate) fn backoff(&self) -> ExponentialBackoff { + ExponentialBuilder::new() + .with_min_delay(self.min_delay) + .with_max_delay(self.max_delay) + .with_max_times(self.max_attempts.saturating_sub(1) as usize) + .with_jitter() + .build() + } + + pub(crate) fn delay(&self, backoff: Duration, retry_after: Option) -> Duration { + retry_after.map_or(backoff, |server| server.min(self.max_delay)) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + const POLICY: RetryPolicy = RetryPolicy { + max_attempts: 3, + timeout: Duration::from_secs(5), + min_delay: Duration::from_secs(1), + max_delay: Duration::from_secs(30), + }; + + #[rstest] + #[case::backoff_without_retry_after(None, 4)] + #[case::retry_after_wins(Some(7), 7)] + #[case::retry_after_zero(Some(0), 0)] + #[case::retry_after_capped(Some(120), 30)] + fn the_server_retry_after_overrides_the_backoff_up_to_max_delay( + #[case] retry_after: Option, + #[case] seconds: u64, + ) { + assert_eq!( + POLICY.delay(Duration::from_secs(4), retry_after.map(Duration::from_secs)), + Duration::from_secs(seconds) + ); + } + + #[rstest] + #[case::one_attempt(1, 0)] + #[case::three_attempts(3, 2)] + #[case::zero_attempts(0, 0)] + fn the_backoff_allows_one_retry_fewer_than_the_attempts( + #[case] max_attempts: u32, + #[case] retries: usize, + ) { + let policy = RetryPolicy { + max_attempts, + ..POLICY + }; + assert_eq!(policy.backoff().count(), retries); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/source.rs b/litellm-rust/crates/model-catalog-refresh/src/source.rs index 713e0386941..28849ba2841 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/source.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/source.rs @@ -1,17 +1,59 @@ -use std::path::PathBuf; - +use litellm_s3::S3Uri; use url::Url; +use crate::Error; + /// Where a fetched catalog comes from #[derive(Clone, Debug, PartialEq, Eq)] pub enum Source { Http(Url), - File(PathBuf), - S3(S3Object), + S3(S3Uri), } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct S3Object { - pub bucket: String, - pub key: String, +impl Source { + pub fn parse(url: &str) -> Result { + let parsed = Url::parse(url)?; + match parsed.scheme() { + "http" | "https" => Ok(Self::Http(parsed)), + "s3" => Ok(Self::S3(S3Uri::try_from(url)?)), + scheme => Err(Error::unsupported_scheme(scheme.to_owned())), + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::https("https://example.com/cost.json")] + #[case::http("http://127.0.0.1:8000/cost.json")] + fn http_urls_are_fetched_over_http(#[case] url: &str) { + assert_eq!( + Source::parse(url).unwrap(), + Source::Http(Url::parse(url).unwrap()) + ); + } + + #[rstest] + fn s3_urls_name_a_bucket_and_key() { + assert_eq!( + Source::parse("s3://bucket/prices/cost.json").unwrap(), + Source::S3(S3Uri { + bucket: "bucket".into(), + key: "prices/cost.json".into(), + }) + ); + } + + #[rstest] + #[case::s3_without_a_key("s3://bucket/")] + #[case::unsupported("ftp://example.com/cost.json")] + #[case::file("file:///tmp/cost.json")] + #[case::not_a_url("not a url")] + fn other_urls_are_refused(#[case] url: &str) { + assert!(Source::parse(url).is_err()); + } } diff --git a/litellm-rust/crates/model-catalog-refresh/src/step.rs b/litellm-rust/crates/model-catalog-refresh/src/step.rs deleted file mode 100644 index 822d1807916..00000000000 --- a/litellm-rust/crates/model-catalog-refresh/src/step.rs +++ /dev/null @@ -1,66 +0,0 @@ -use std::time::Duration; - -use litellm_model_catalog::Catalog; - -use crate::Error; -use crate::source::Source; - -/// One attempt. Where to fetch from is the `Source`, so the request carries no URL -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct FetchRequest { - pub if_none_match: Option, - pub timeout: Duration, -} - -#[derive(Debug)] -pub enum Reply { - /// An HTTP, S3 or file answer. The body is read only on 200 - Response { - status: u16, - etag: Option, - retry_after: Option, - body: Option>, - }, - /// No answer reached us - Transport { error: Error }, -} - -/// What the fetch state machine asks its driver to do next -#[derive(Debug)] -pub enum Step { - Fetch(FetchRequest), - Sleep(Duration), - /// A parsed document to install, with its rejected row errors - Adopt { - catalog: Box, - installed: Installed, - errors: Vec, - }, - /// 304, a final status, or retries exhausted - Done { - reason: DoneReason, - }, -} - -#[derive(Debug)] -pub enum DoneReason { - NotModified, - Status(u16), - Exhausted(Error), -} - -/// What is installed now, the input to the next revalidation -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct Installed { - pub origin: Origin, - pub etag: Option, - /// `git_blob_id` of the installed bytes - pub revision: String, - pub model_count: usize, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum Origin { - Bundled, - Remote { source: Source }, -} diff --git a/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs b/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs new file mode 100644 index 00000000000..2a4ec996bd4 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs @@ -0,0 +1,541 @@ +use std::io::{BufRead, BufReader, Write}; +use std::net::TcpListener; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use aws_credential_types::Credentials; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::{ProvideCredentials, future}; +use litellm_http::Client; +use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start, fetch}; +use rstest::rstest; + +const POLICY: RetryPolicy = RetryPolicy { + max_attempts: 3, + timeout: Duration::from_secs(5), + min_delay: Duration::from_millis(10), + max_delay: Duration::from_millis(50), +}; + +const OK: &str = + "HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}"; +const UNAVAILABLE: &str = "HTTP/1.1 503 Service Unavailable\r\nRetry-After: 0\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"; +const NOT_FOUND: &str = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"; + +const SLOW_BODY: &[u8] = b"{ }"; + +#[derive(Clone, Copy)] +enum Answer { + Raw(&'static str), + Stall, + StallBody, + SlowBody, + Oversized, +} + +/// A scripted upstream: the n-th request gets the n-th answer, repeating the last, and every request head is kept +struct Upstream { + base: String, + heads: Arc>>, +} + +impl Upstream { + fn serve(answers: Vec) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let heads = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&heads); + std::thread::spawn(move || { + for (index, stream) in listener.incoming().enumerate() { + let answer = answers[index.min(answers.len() - 1)]; + let recorded = Arc::clone(&recorded); + std::thread::spawn(move || { + let mut stream = stream.unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut head = String::new(); + loop { + let mut line = String::new(); + if reader.read_line(&mut line).unwrap() <= 2 { + break; + } + head.push_str(&line); + } + recorded.lock().unwrap().push(head); + match answer { + Answer::Raw(raw) => stream.write_all(raw.as_bytes()).unwrap(), + Answer::Stall => std::thread::sleep(Duration::from_secs(5)), + Answer::Oversized => { + let chunk = vec![b' '; 1024 * 1024]; + let written = stream + .write_all(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n") + .and_then(|()| { + (0..=100).try_for_each(|_| stream.write_all(&chunk)) + }); + drop(written); + } + Answer::StallBody => { + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{") + .unwrap(); + std::thread::sleep(Duration::from_secs(5)); + } + Answer::SlowBody => { + stream.set_nodelay(true).unwrap(); + stream + .write_all( + b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 8\r\nConnection: close\r\n\r\n", + ) + .unwrap(); + for byte in SLOW_BODY { + std::thread::sleep(Duration::from_millis(60)); + stream.write_all(&[*byte]).unwrap(); + } + } + } + }); + } + }); + Self { base, heads } + } + + fn url(&self) -> String { + format!("{}/cost.json", self.base) + } + + fn heads(&self) -> Vec { + self.heads.lock().unwrap().clone() + } +} + +fn header<'a>(head: &'a str, name: &str) -> Option<&'a str> { + head.lines().find_map(|line| { + let (key, value) = line.split_once(':')?; + key.eq_ignore_ascii_case(name).then(|| value.trim()) + }) +} + +#[derive(Clone, Debug)] +struct NoCredentials; + +impl ProvideCredentials for NoCredentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::ready(Err(CredentialsError::not_loaded("no credentials"))) + } +} + +fn keys() -> Credentials { + Credentials::new( + "AKIDMODELCOST", + "model-cost-secret", + None, + None, + "model-cost-test", + ) +} + +async fn fetched(url: &str, policy: RetryPolicy, start: Start) -> Result { + fetch( + &Client::for_test(), + &S3Access::Environment { region: None }, + url, + policy, + start, + ) + .await +} + +#[rstest] +#[case::every_etag_joined( + "HTTP/1.1 200 OK\r\nETag: \"a\"\r\nETag: \"b\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}", + Some("\"a\", \"b\"") +)] +#[case::no_etag( + "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}", + None +)] +#[tokio::test] +async fn a_200_returns_the_body_and_its_etags( + #[case] response: &'static str, + #[case] etag: Option<&str>, +) { + let upstream = Upstream::serve(vec![Answer::Raw(response)]); + assert_eq!( + fetched(&upstream.url(), POLICY, Start::Fresh).await, + Ok(Fetched { + body: b"{}".to_vec(), + etag: etag.map(str::to_owned), + }) + ); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test(flavor = "multi_thread")] +async fn a_fetch_runs_as_a_task_on_a_multi_thread_runtime() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let url = upstream.url(); + let fetched = tokio::spawn(async move { + fetch( + &Client::for_test(), + &S3Access::Environment { region: None }, + &url, + POLICY, + Start::Fresh, + ) + .await + }) + .await + .unwrap(); + assert!(fetched.is_ok()); +} + +#[rstest] +#[tokio::test] +async fn a_redirect_is_not_followed_and_its_body_is_the_document() { + let upstream = Upstream::serve(vec![ + Answer::Raw( + "HTTP/1.1 301 Moved Permanently\r\nLocation: /elsewhere.json\r\nContent-Length: 7\r\nConnection: close\r\n\r\n{\"a\":1}", + ), + Answer::Raw(OK), + ]); + assert_eq!( + fetched(&upstream.url(), POLICY, Start::Fresh).await, + Ok(Fetched { + body: b"{\"a\":1}".to_vec(), + etag: None, + }) + ); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_retry_after_shorter_than_the_backoff_is_honored() { + let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE), Answer::Raw(OK)]); + let policy = RetryPolicy { + min_delay: Duration::from_secs(10), + max_delay: Duration::from_secs(10), + ..POLICY + }; + let fetched = tokio::time::timeout( + Duration::from_secs(2), + fetched(&upstream.url(), policy, Start::Fresh), + ) + .await + .expect("Retry-After: 0 replaces the 10s backoff"); + assert!(fetched.is_ok(), "{fetched:?}"); + assert_eq!(upstream.heads().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn retryable_failures_stop_at_the_attempt_limit() { + let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]); + let url = upstream.url(); + assert_eq!( + fetched(&url, POLICY, Start::Fresh).await, + Err(Failed { + reason: format!("HTTP 503 from {url} (after 3 attempts)"), + retryable: true, + retry_after: Some(Duration::ZERO), + }) + ); + assert_eq!(upstream.heads().len(), 3); +} + +#[rstest] +#[tokio::test] +async fn a_404_fails_on_the_first_attempt() { + let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]); + let url = upstream.url(); + assert_eq!( + fetched(&url, POLICY, Start::Fresh).await, + Err(Failed { + reason: format!("HTTP 404 from {url}"), + retryable: false, + retry_after: None, + }) + ); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_stalled_response_times_out_and_is_retried() { + let upstream = Upstream::serve(vec![Answer::Stall]); + let url = upstream.url(); + let policy = RetryPolicy { + max_attempts: 2, + timeout: Duration::from_millis(100), + ..POLICY + }; + let failed = fetched(&url, policy, Start::Fresh).await.unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")) + && failed.reason.ends_with(" (after 2 attempts)"), + "{}", + failed.reason + ); + assert!(failed.retryable); + assert_eq!(upstream.heads().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn a_resumed_fetch_makes_only_the_attempts_left() { + let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]); + let url = upstream.url(); + let first = Failed { + reason: "first".into(), + retryable: true, + retry_after: Some(Duration::ZERO), + }; + let failed = fetched(&url, POLICY, Start::After(first)) + .await + .unwrap_err(); + assert_eq!( + failed.reason, + format!("HTTP 503 from {url} (after 3 attempts)") + ); + assert_eq!(upstream.heads().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn a_resumed_fetch_with_no_attempts_left_returns_the_failure_it_resumed() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let first = Failed { + reason: "first".into(), + retryable: true, + retry_after: None, + }; + let policy = RetryPolicy { + max_attempts: 1, + ..POLICY + }; + assert_eq!( + fetched(&upstream.url(), policy, Start::After(first.clone())).await, + Err(first) + ); + assert!(upstream.heads().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn a_failure_reason_drops_url_credentials() { + let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]); + let address = upstream.base.trim_start_matches("http://"); + let url = format!("http://user:hunter2@{address}/cost.json?X-Amz-Signature=sig-secret#frag"); + let failed = fetched(&url, POLICY, Start::Fresh).await.unwrap_err(); + assert_eq!(failed.reason, format!("HTTP 404 from {}", upstream.url())); +} + +#[rstest] +#[case::unsupported_scheme("ftp://example.com/cost.json")] +#[case::not_a_url("not a url")] +#[case::s3_without_a_key("s3://prices")] +#[tokio::test] +async fn an_unfetchable_url_fails_without_a_request(#[case] url: &str) { + let failed = fetched(url, POLICY, Start::Fresh).await.unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")), + "{}", + failed.reason + ); + assert!(!failed.retryable); +} + +fn s3(upstream: &Upstream, credentials: impl ProvideCredentials + 'static) -> S3Access { + S3Access::Config(Box::new( + litellm_s3::config( + "us-west-2".into(), + Some(&upstream.base), + Client::for_test(), + credentials, + ) + .build(), + )) +} + +async fn fetched_from_s3( + url: &str, + upstream: &Upstream, + credentials: impl ProvideCredentials + 'static, +) -> Result { + fetch( + &Client::for_test(), + &s3(upstream, credentials), + url, + POLICY, + Start::Fresh, + ) + .await +} + +#[rstest] +#[tokio::test] +async fn an_s3_fetch_sends_a_signed_get_object_to_the_endpoint() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + assert_eq!( + fetched_from_s3("s3://prices/maps/cost.json", &upstream, keys()).await, + Ok(Fetched { + body: b"{}".to_vec(), + etag: Some("\"v1\"".into()), + }) + ); + let heads = upstream.heads(); + let [head] = heads.as_slice() else { + panic!("expected one request, got {heads:?}"); + }; + assert!(head.starts_with("GET /prices/maps/cost.json"), "{head}"); + let authorization = header(head, "authorization").unwrap_or_default(); + assert!( + authorization.starts_with("AWS4-HMAC-SHA256 Credential=AKIDMODELCOST/") + && authorization.contains("/us-west-2/s3/aws4_request"), + "{authorization}" + ); +} + +#[rstest] +#[tokio::test] +async fn an_s3_body_that_stalls_after_the_headers_times_out() { + let upstream = Upstream::serve(vec![Answer::StallBody]); + let url = "s3://prices/cost.json"; + let policy = RetryPolicy { + timeout: Duration::from_millis(200), + ..POLICY + }; + let failed = tokio::time::timeout( + Duration::from_secs(2), + fetch( + &Client::for_test(), + &s3(&upstream, keys()), + url, + policy, + Start::Fresh, + ), + ) + .await + .expect("the attempt timeout bounds the body read") + .unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")), + "{}", + failed.reason + ); + assert!(!failed.retryable); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[case::http(|upstream: &Upstream| (S3Access::Environment { region: None }, upstream.url()))] +#[case::s3(|upstream: &Upstream| (s3(upstream, keys()), "s3://prices/cost.json".to_owned()))] +#[tokio::test] +async fn a_body_that_keeps_arriving_finishes_after_the_timeout( + #[case] source: fn(&Upstream) -> (S3Access, String), +) { + let upstream = Upstream::serve(vec![Answer::SlowBody]); + let (access, url) = source(&upstream); + let policy = RetryPolicy { + timeout: Duration::from_millis(300), + ..POLICY + }; + assert_eq!( + fetch(&Client::for_test(), &access, &url, policy, Start::Fresh).await, + Ok(Fetched { + body: SLOW_BODY.to_vec(), + etag: Some("\"v1\"".into()), + }) + ); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[case::missing_key(NOT_FOUND)] +#[case::denied("HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")] +#[tokio::test] +async fn an_s3_error_fails_without_a_retry(#[case] response: &'static str) { + let upstream = Upstream::serve(vec![Answer::Raw(response), Answer::Raw(OK)]); + let url = "s3://prices/cost.json"; + let failed = fetched_from_s3(url, &upstream, keys()).await.unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")), + "{}", + failed.reason + ); + assert!(!failed.retryable); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn an_s3_fetch_without_credentials_fails_without_a_request() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let url = "s3://prices/cost.json"; + let failed = fetched_from_s3(url, &upstream, NoCredentials) + .await + .unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")), + "{}", + failed.reason + ); + assert!(!failed.retryable); + assert!(upstream.heads().is_empty()); +} + +const DECLARED_OVERSIZED: &str = + "HTTP/1.1 200 OK\r\nContent-Length: 104857601\r\nConnection: close\r\n\r\n"; + +#[rstest] +#[case::http_declared(Answer::Raw(DECLARED_OVERSIZED), false)] +#[case::http_streamed(Answer::Oversized, false)] +#[case::s3_declared(Answer::Raw(DECLARED_OVERSIZED), true)] +#[case::s3_streamed(Answer::Oversized, true)] +#[tokio::test] +async fn a_body_over_100_mib_fails_without_a_retry(#[case] answer: Answer, #[case] s3: bool) { + let upstream = Upstream::serve(vec![answer, Answer::Raw(OK)]); + let failed = match s3 { + true => fetched_from_s3("s3://prices/cost.json", &upstream, keys()).await, + false => fetched(&upstream.url(), POLICY, Start::Fresh).await, + } + .unwrap_err(); + assert!( + failed + .reason + .contains("response body is over the 104857600 byte limit"), + "{}", + failed.reason + ); + assert!(!failed.retryable); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_body_of_exactly_100_mib_is_fetched() { + let upstream = Upstream::serve(vec![Answer::Raw(Box::leak( + format!( + "HTTP/1.1 200 OK\r\nContent-Length: 104857600\r\nConnection: close\r\n\r\n{}", + " ".repeat(100 * 1024 * 1024) + ) + .into_boxed_str(), + ))]); + let fetched = fetched(&upstream.url(), POLICY, Start::Fresh) + .await + .unwrap(); + assert_eq!(fetched.body.len(), 100 * 1024 * 1024); +} diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index ff24fb81143..a9c473f2ca3 100644 --- a/litellm-rust/crates/model-catalog/AGENTS.md +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -1,3 +1,41 @@ +## Concepts + +Every price in a row lowers to one `Price { meter, channel, when }` and its `UnitPrice`. These keys come from the `FlatKey::parse` tests + +- `input_cost_per_token` + - meter `Input(Token(Text))`, channel `Online`, tier `Standard`, band `Base` +- `output_cost_per_token_above_272k_tokens_priority` + - meter `Output(Token(Text))`, channel `Online` + - `_priority` sets tier `Priority` + - `_above_272k_tokens` sets band `OverPromptTokens(272000)` +- `cache_creation_input_token_cost_above_1hr_above_200k_tokens` + - meter `CacheWrite { Text, OneHour }`, where `_above_1hr` is the cache TTL, not a band + - `_above_200k_tokens` sets band `OverPromptTokens(200000)` + - channel `Online`, tier `Standard` +- `ocr_cost_per_page_batches` + - meter `Feature(OcrPage)`, tier `Standard`, band `Base` + - `_batches` sets channel `Batch` + +- Meter is what gets counted, such as text input tokens, reasoning tokens, an image of a given size or a web search + - A variation only one meter has stays inside it (cache TTL, image size, resolution, `_above_Ns_interval`, `max_results_range`) +- `UnitPrice` is USD for one unit of the meter, such as one token, image, page or second +- Band is a range of prompt sizes, starting where `_above__tokens` or `_above_k_tokens` says + - `Prices::band` picks one band from prompt tokens alone, the highest edge passed wins, and `Band::Base` is below every edge + - The whole call moves, so every token meter, output included, bills all its tokens at that band's rate + - `Billing::band_edge` says whether a prompt exactly at the edge is in the band (exclusive by default, inclusive for xAI) + - A band only another tier or the batch channel prices never moves the call +- `When` is the request's condition, its service tier and band + - A missing price falls back from `(tier, band)` to `(standard, band)`, `(tier, base)`, then `(standard, base)` +- `Channel` is the route whose price list the call reads, and each falls back differently + - `Online` falls back across `When` as above + - `Batch` (`_batches`) uses bands but no service tier, and a missing price is half the online rate + - `Decisions` (`_decisions`, OpenAI `/v1/decisions`) overrides single online fields, and a missing price is the full online rate +- Off-peak and `Adjustments` are never prices + - Off-peak swaps in other token rates while a UTC window is open + - `Adjustments` multiplies a priced call for region, residency, geo or `speed: fast` +- `Billing` on `Prices` holds row attributes that change how its prices apply (text unit, band edge, cache-read fallback, wall clock, web search unit, tier scope) + - `lower.rs` derives it from the row and its `litellm_provider`, so `litellm-cost` holds no provider list + ## Layout `wire/` reads one entry as written. `EntryWire::read` makes one pass over the entry's keys and puts every non-null key in exactly one place, which is an attribute, a flat price parsed by `FlatKey::parse`, a price block, a `supports_*` capability or `unknown`. A known key with a bad value becomes an `Error` that names the key, and the rest of the entry is still read @@ -6,17 +44,22 @@ `catalog.rs` reads the document, indexes aliases and runs the checks across entries. A bad entry keeps its raw fields and its errors are listed in `Loaded.errors`. Deciding what to do with them is the caller's job. `Catalog::written` gives back every top-level key and value as the document wrote them, in document order. `Catalog::served` gives back every model and alias as Python serves it, without `sample_spec` and `fallback_generalizations` -`Billing` on `Prices` holds the row attributes pricing reads (text unit, band edge, cache-read fallback, wall clock, web search unit, tier scope). `lower.rs` derives it from the row and its `litellm_provider`, so `litellm-cost` holds no provider list - `store/` is the catalog's in-memory store and the only place that indexes model names. Its rules are in [src/store/AGENTS.md](src/store/AGENTS.md) -`layers.rs`, `resolve/`, `model_info.rs` and `register.rs` hold types only. The functions over them land with their first caller +`cost_map.rs` holds a host's inputs, the installed document and the `register_model` overlay in `register.rs`. It serves the document in the key order and alias sharing Python's `litellm.model_cost` has, without the overlay, since where a registration lands is a lookup in the published table. It never holds the live table, which belongs to the host + +- `CostMap` is immutable behind an `Arc`. A new document is a new `CostMap` +- `Registrations` is generic over the host's value type, because Python registrations can hold `datetime` and other values JSON cannot. The host keeps its own objects (`Py` in python-bridge) and does the merging + - Entry order is latest-registration order. `landed_on` records which registered keys landed on each row key +- `write_pass` stops at the first `place` that returns `None` and reports `Pass::Swapped`. Retrying is the host's job + +`resolve/` and `model_info.rs` hold types only. The functions over them land with their first caller Catalog errors use one opaque `Error` with typed causes and JSON Pointer accessors (`path`, `key`, `field`). `Error::rejects` alone decides whether an error fails the document, rejects its row or neither, and `ErrorKind` is the stable name hosts report it by ## Prices -The price vocabulary lives in `litellm-cost`, and only this crate knows JSON key names. Every price is keyed by `Meter`, `Channel` and `When`. A variation only one meter has stays in the meter (cache TTL, image size, resolution, `_above_Ns_interval`, `max_results_range`). Service tier and prompt band make up `When`. Batch and decisions are a `Channel` because Python bills them with their own rules. Off-peak is an overlay on `Prices`, and regional uplifts are `Adjustments`, never prices +The price vocabulary lives in `litellm-cost`, and only this crate knows JSON key names A new `N`, image size or resolution in an existing key shape parses with no Rust change. A new unit, token kind, feature or tier needs a variant in `litellm-cost` plus its spelling in `wire/flat_key.rs` and `FLAT_KEY_PATTERN` in `schema.rs` diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index 63b9ad14768..8cd8ee5760b 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -11,6 +11,7 @@ schema = ["dep:schemars", "litellm-cost/schema"] [dependencies] litellm-cost = { workspace = true, features = ["serde"] } litellm-error.workspace = true +litellm-tracing.workspace = true indexmap = { version = "2.14.0", features = ["serde"] } jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] } @@ -19,6 +20,7 @@ schemars = { workspace = true, optional = true } serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } serde_with.workspace = true +sha1.workspace = true strum.workspace = true thiserror.workspace = true time.workspace = true diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index 98af785b17e..6c11a39b0eb 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -12,7 +12,8 @@ use crate::store::{Rows, Store}; use crate::wire::{EntryWire, Fields}; const SAMPLE_SPEC: &str = "sample_spec"; -const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; +/// The top-level key of the rules block, which is never a model row +pub const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; #[derive(Clone, Copy, Debug, PartialEq)] pub struct IntegrityLimits { @@ -111,6 +112,11 @@ impl Rows for CatalogRows { } impl IntegrityLimits { + pub const UNCHECKED: Self = Self { + min_model_count: 0, + min_reference_ratio: 0.0, + }; + pub fn new(min_model_count: usize, min_reference_ratio: f64) -> Option { (0.0..=1.0).contains(&min_reference_ratio).then_some(Self { min_model_count, @@ -142,6 +148,7 @@ impl Catalog { let mut sections = Vec::new(); let mut entries = IndexMap::with_capacity(document.len()); let mut alias_lists = Vec::new(); + let mut pointer_rows = Vec::new(); for (index, (name, value)) in document.into_iter().enumerate() { let fields = match value { Value::Object(fields) if !is_section(&name) => fields, @@ -157,16 +164,34 @@ impl Catalog { continue; } }; - let (model, aliases, row_problems) = read_row(&fields); + let row = read_row(&fields); + let readable = row.problems.is_empty(); problems.extend( - row_problems + row.problems .into_iter() .map(|problem| problem.prepend(&name)), ); - if let Some(aliases) = aliases { + if let Some(target) = row.alias_of { + if readable { + pointer_rows.push((name.clone(), target)); + } + sections.push(Section { + index, + key: name, + value: Value::Object(fields), + }); + continue; + } + if let Some(aliases) = row.aliases { alias_lists.push((name.clone(), aliases)); } - entries.insert(name, ModelEntry { fields, model }); + entries.insert( + name, + ModelEntry { + fields, + model: row.model, + }, + ); } if entries.is_empty() { return Err(Error::new("", Repr::Empty)); @@ -190,6 +215,21 @@ impl Catalog { } } } + for (alias, target) in pointer_rows { + if !entries.contains_key(&target) { + problems.push(Error::from_key(&alias, Repr::AliasTargetMissing { target })); + } else if let Some(taken_by) = aliases.get(&alias) { + problems.push(Error::from_key( + &alias, + Repr::AliasTaken { + alias: alias.clone(), + taken_by: taken_by.clone(), + }, + )); + } else { + aliases.insert(alias, target); + } + } let fallback_generalizations = sections .iter() @@ -242,8 +282,8 @@ impl Catalog { } pub fn row_errors(model: &str, fields: &Fields) -> Vec { - let (_, _, errors) = read_row(fields); - errors + read_row(fields) + .problems .into_iter() .map(|error| error.prepend(model)) .collect() @@ -255,7 +295,13 @@ impl Catalog { (0..self.rows().entries.len() + self.sections.len()).map_while(move |index| match sections .next_if(|section| section.index == index) { - Some(section) => Some((section.key.as_str(), Written::Section(§ion.value))), + Some(section) => Some(( + section.key.as_str(), + match §ion.value { + Value::Object(fields) if !is_section(§ion.key) => Written::Entry(fields), + value => Written::Section(value), + }, + )), None => entries .next() .map(|(name, entry)| (name.as_str(), Written::Entry(&entry.fields))), @@ -378,6 +424,20 @@ impl Catalog { }) } + /// The row under an exact model or alias key, an alias reading its owner's row + pub fn get(&self, key: &str) -> Option<&Fields> { + self.rows().fields(key) + } + + /// A top-level value that is not a model row, such as `sample_spec` + pub fn section(&self, key: &str) -> Option<&Value> { + self.sections + .iter() + .filter(|section| is_section(§ion.key)) + .find(|section| section.key == key) + .map(|section| §ion.value) + } + pub fn model_count(&self) -> usize { self.rows().entries.len() } @@ -406,16 +466,41 @@ impl Catalog { } } -fn read_row(fields: &Fields) -> (Option, Option>, Vec) { +struct RowRead { + model: Option, + aliases: Option>, + alias_of: Option, + problems: Vec, +} + +fn read_row(fields: &Fields) -> RowRead { let wire = EntryWire::read(fields); + if let Some(target) = wire.attributes.alias_of.clone() { + let other_fields = fields + .keys() + .filter(|key| *key != "alias_of") + .map(|key| Error::from_key(key, Repr::AliasRowField)); + return RowRead { + model: None, + aliases: None, + alias_of: Some(target), + problems: wire.problems.into_iter().chain(other_fields).collect(), + }; + } let aliases = wire.attributes.aliases.clone(); let timezone_error = unknown_timezone(&wire); - match Model::try_from(wire) { - Ok(model) => (Some(model), aliases, timezone_error.into_iter().collect()), + let (model, problems) = match Model::try_from(wire) { + Ok(model) => (Some(model), timezone_error.into_iter().collect()), Err(mut errors) => { errors.extend(timezone_error); - (None, aliases, errors) + (None, errors) } + }; + RowRead { + model, + aliases, + alias_of: None, + problems, } } diff --git a/litellm-rust/crates/model-catalog/src/cost_map.rs b/litellm-rust/crates/model-catalog/src/cost_map.rs new file mode 100644 index 00000000000..bec417acae7 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/cost_map.rs @@ -0,0 +1,130 @@ +use std::sync::{Arc, LazyLock}; + +use indexmap::IndexMap; +use serde_json::Value; +use sha1::{Digest, Sha1}; + +use crate::bundled::BUNDLED_REGISTRY; +use crate::catalog::{Catalog, FALLBACK_GENERALIZATIONS, IntegrityLimits}; +use crate::error::Error; + +/// Which document is installed: the `git_blob_id` of its bytes and the ETag it was served with +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Provenance { + pub revision: String, + pub etag: Option, +} + +/// An installed cost map document. A host builds its live table from it +#[derive(Clone, Debug)] +pub struct CostMap { + base: Arc, + provenance: Provenance, +} + +#[derive(Debug)] +struct Base { + catalog: Catalog, + /// Each owner's names, the owner first, for owners with at least one alias + groups: IndexMap>, +} + +impl Base { + fn new(catalog: Catalog) -> Self { + let groups = catalog.aliases().iter().fold( + IndexMap::>::new(), + |mut groups, (alias, owner)| { + groups + .entry(owner.clone()) + .or_insert_with(|| vec![owner.clone()]) + .push(alias.clone()); + groups + }, + ); + Self { catalog, groups } + } +} + +static BUNDLED: LazyLock = LazyLock::new(|| CostMap { + base: Arc::new(Base::new( + Catalog::parse(BUNDLED_REGISTRY) + .expect("the bundled registry parses") + .value, + )), + provenance: Provenance { + revision: git_blob_id(BUNDLED_REGISTRY), + etag: None, + }, +}); + +impl CostMap { + pub fn bundled() -> Self { + BUNDLED.clone() + } + + /// Reads a fetched document. One that does not parse, or that has fewer models than `limits` allow against the + /// bundled one, is refused. Every alias problem it has is logged + pub fn read(body: &[u8], etag: Option, limits: IntegrityLimits) -> Result { + let loaded = Catalog::parse(body)?; + loaded + .value + .check_size(BUNDLED.base.catalog.model_count(), limits)?; + for problem in loaded.errors.iter().filter(|error| error.is_alias()) { + litellm_tracing::warn!( + "LiteLLM: model cost map alias problem: {}", + litellm_error::Report::new(problem) + ); + } + Ok(Self { + base: Arc::new(Base::new(loaded.value)), + provenance: Provenance { + revision: git_blob_id(body), + etag, + }, + }) + } + + pub fn catalog(&self) -> &Catalog { + &self.base.catalog + } + + pub fn provenance(&self) -> &Provenance { + &self.provenance + } + + /// Every name the row under `key` is served under, the owner first. Empty for a row without aliases + pub fn names(&self, key: &str) -> &[String] { + let owner = self + .catalog() + .aliases() + .get(key) + .map_or(key, String::as_str); + self.base + .groups + .get(owner) + .map(Vec::as_slice) + .unwrap_or_default() + } + + /// Each owner with its names, the owner first, in the order the document declared the aliases + pub fn alias_groups(&self) -> impl Iterator { + self.base + .groups + .iter() + .map(|(owner, names)| (owner.as_str(), names.as_slice())) + } + + /// The `fallback_generalizations` block as the document wrote it + pub fn fallback_generalizations(&self) -> Option<&Value> { + self.catalog().section(FALLBACK_GENERALIZATIONS) + } +} + +/// `git hash-object` of the bytes, which main reports as `source_revision` +pub fn git_blob_id(body: &[u8]) -> String { + let digest = Sha1::new() + .chain_update(format!("blob {}\0", body.len())) + .chain_update(body) + .finalize(); + format!("{digest:x}") +} diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index 89e27c05f46..a4b54b5a551 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -67,6 +67,12 @@ pub(crate) enum Repr { #[error("alias {alias:?} already belongs to {taken_by}")] #[strum_discriminants(strum(serialize = "alias_taken"))] AliasTaken { alias: String, taken_by: String }, + #[error("alias_of {target:?} is not a model name")] + #[strum_discriminants(strum(serialize = "alias_target_missing"))] + AliasTargetMissing { target: String }, + #[error("an alias_of row carries no other field")] + #[strum_discriminants(strum(serialize = "alias_row_field"))] + AliasRowField, #[error("{first} and {second} price the same meter under the same conditions")] #[strum_discriminants(strum(serialize = "duplicate_price"))] DuplicatePrice { first: String, second: String }, @@ -141,6 +147,8 @@ impl Error { | Repr::UnknownTimezone { .. } | Repr::AliasShadowsModel { .. } | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField | Repr::DuplicatePrice { .. } | Repr::SuffixInsideBlock { .. } | Repr::TierRange @@ -159,6 +167,74 @@ impl Error { } } +impl Error { + /// The document parsed but holds fewer models than the integrity limits trust + pub fn is_integrity(&self) -> bool { + match self.repr() { + Repr::Empty | Repr::BelowMinimum { .. } | Repr::Shrunk { .. } => true, + Repr::Json(_) + | Repr::NotAnObject + | Repr::Invalid(_) + | Repr::UnknownBlockKey + | Repr::UnknownTimezone { .. } + | Repr::AliasShadowsModel { .. } + | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField + | Repr::DuplicatePrice { .. } + | Repr::SuffixInsideBlock { .. } + | Repr::TierRange + | Repr::TierReversed + | Repr::TierGap { .. } + | Repr::MixedTierBasis + | Repr::NotPricedByResults + | Repr::BandsFromTwoSources + | Repr::NotAnOffPeakRate + | Repr::NoOffPeakPrices + | Repr::NoOffPeakHours + | Repr::EmptyList + | Repr::InvalidDate { .. } + | Repr::MissingProvider + | Repr::UnreadField + | Repr::NoFallbackRules => false, + } + } + + /// A problem with a row's aliases, which the served map skips + pub fn is_alias(&self) -> bool { + match self.repr() { + Repr::AliasShadowsModel { .. } + | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField => true, + Repr::Invalid(_) => matches!(self.field().as_deref(), Some("aliases" | "alias_of")), + Repr::Json(_) + | Repr::Empty + | Repr::BelowMinimum { .. } + | Repr::Shrunk { .. } + | Repr::NotAnObject + | Repr::UnknownBlockKey + | Repr::UnknownTimezone { .. } + | Repr::DuplicatePrice { .. } + | Repr::SuffixInsideBlock { .. } + | Repr::TierRange + | Repr::TierReversed + | Repr::TierGap { .. } + | Repr::MixedTierBasis + | Repr::NotPricedByResults + | Repr::BandsFromTwoSources + | Repr::NotAnOffPeakRate + | Repr::NoOffPeakPrices + | Repr::NoOffPeakHours + | Repr::EmptyList + | Repr::InvalidDate { .. } + | Repr::MissingProvider + | Repr::UnreadField + | Repr::NoFallbackRules => false, + } + } +} + #[derive(Debug)] pub struct Loaded { pub value: T, @@ -327,6 +403,8 @@ mod tests { #[case::unknown_timezone(Repr::UnknownTimezone { zone: "invalid".into() }, Rejects::Row)] #[case::alias_shadows_model(Repr::AliasShadowsModel { alias: "model".into() }, Rejects::Row)] #[case::alias_taken(Repr::AliasTaken { alias: "alias".into(), taken_by: "owner".into() }, Rejects::Row)] + #[case::alias_target_missing(Repr::AliasTargetMissing { target: "owner".into() }, Rejects::Row)] + #[case::alias_row_field(Repr::AliasRowField, Rejects::Row)] #[case::duplicate_price(Repr::DuplicatePrice { first: "a".into(), second: "b".into() }, Rejects::Row)] #[case::suffix_inside_block(Repr::SuffixInsideBlock { block: "tiered_pricing" }, Rejects::Row)] #[case::tier_range(Repr::TierRange, Rejects::Row)] diff --git a/litellm-rust/crates/model-catalog/src/layers.rs b/litellm-rust/crates/model-catalog/src/layers.rs deleted file mode 100644 index 881f0f6af9b..00000000000 --- a/litellm-rust/crates/model-catalog/src/layers.rs +++ /dev/null @@ -1,32 +0,0 @@ -use std::sync::Arc; - -use indexmap::IndexMap; -use serde_json::Value; - -use crate::catalog::Catalog; -use crate::register::Registrations; -use crate::wire::Fields; - -/// One field of a fetched row over the layer below, RFC 7386. An absent field inherits, JSON `null` is `Remove` -#[derive(Clone, Debug, PartialEq)] -pub enum Patch { - Set(Value), - Remove, -} - -/// The four layers served rows are composed from, lowest first -#[derive(Clone, Debug)] -#[expect( - dead_code, - reason = "02 and 07 add the functions that read these layers" -)] -pub struct Layers { - /// 0, the wheel's embedded map - bundled: Arc, - /// 1, the last installed fetch, rejected rows left out - fetched: Option>, - /// 2, `register(persist=False)`, cleared by every install and refilled by the router replay - transient: IndexMap, - /// 3, `register(persist=True)`, applied over every install - persisted: Registrations, -} diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index fe419dde054..9b869a078db 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -17,6 +17,9 @@ //! failing the document. //! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known //! families. They are data in the file, not entries. +//! - `CostMap`: an installed document with its provenance, which a host builds its live table from. +//! `Registrations` is the ordered `register_model` overlay over the host's own values, and +//! `write_pass` writes registrations into a host's table through `RegistrationTable`. //! - Validation: `Catalog::check_size` checks a catalog against `IntegrityLimits` (a remote copy //! that lost most of its models is refused), and `Catalog::row_errors` validates one row. //! @@ -37,9 +40,9 @@ macro_rules_attribute::attribute_alias! { mod bundled; mod catalog; +mod cost_map; mod error; mod fallback; -mod layers; mod lower; mod model; mod model_info; @@ -51,9 +54,9 @@ pub mod wire; pub use bundled::*; pub use catalog::*; +pub use cost_map::*; pub use error::*; pub use fallback::*; -pub use layers::*; pub use model::*; pub use model_info::*; pub use register::*; diff --git a/litellm-rust/crates/model-catalog/src/lower.rs b/litellm-rust/crates/model-catalog/src/lower.rs index 4c3fc38a01a..7acfc781431 100644 --- a/litellm-rust/crates/model-catalog/src/lower.rs +++ b/litellm-rust/crates/model-catalog/src/lower.rs @@ -104,6 +104,7 @@ impl From<&EntryWire> for Adjustments { (Uplift::NonGlobalEndpoint, wire.regional_endpoint_uplift), (Uplift::Residency(Residency::Eu), wire.residency_uplift_eu), (Uplift::Residency(Residency::Us), wire.residency_uplift_us), + (Uplift::FedrampEndpoint, wire.fedramp_uplift), ] .into_iter() .filter_map(|(uplift, factor)| Some((uplift, factor?))), @@ -746,7 +747,8 @@ mod tests { fn uplifts_are_adjustments_not_prices() { let wire = wire(json!({ "regional_endpoint_uplift_multiplier": 1.1, - "regional_processing_uplift_multiplier_eu": "1.2" + "regional_processing_uplift_multiplier_eu": "1.2", + "fedramp_uplift_multiplier": 1.1 })); let adjustments = Adjustments::from(&wire); @@ -764,6 +766,12 @@ mod tests { Some(1.2) ); assert_eq!(adjustments.get(&Uplift::Residency(Residency::Us)), None); + assert_eq!( + adjustments + .get(&Uplift::FedrampEndpoint) + .map(Multiplier::get), + Some(1.1) + ); } #[rstest] diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index 8f7549c1c6d..fcd0c027323 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -28,6 +28,7 @@ pub struct ModelInfoFields { pub cache_read_input_audio_token_cost_flex: Option, pub cache_read_input_audio_token_cost_priority: Option, pub cache_read_input_audio_token_cost_batches: Option, + pub cache_read_input_image_token_cost_batches: Option, pub explicit_cache_read_input_token_cost: Option, pub prompt_cache_min_tokens: Option, pub cache_read_input_token_cost_above_200k_tokens: Option, @@ -96,6 +97,7 @@ pub struct ModelInfoFields { pub output_cost_per_token_ultrafast: Option, pub regional_processing_uplift_multiplier_eu: Option, pub regional_processing_uplift_multiplier_us: Option, + pub fedramp_uplift_multiplier: Option, pub regional_endpoint_uplift_multiplier: Option, pub output_cost_per_audio_token: Option, pub output_cost_per_character: Option, @@ -186,6 +188,9 @@ pub struct ModelInfoFields { pub supports_image_size: Option, pub supported_audio_formats: Option, pub token_count_framing: Option, + pub token_count_tool_framing: Option, + pub token_count_responses_framing: Option, + pub token_count_responses_tool_framing: Option, pub token_count_model: Option, pub token_count_tokenizer: Option, pub vertex_ai_audio_api: Option, diff --git a/litellm-rust/crates/model-catalog/src/register.rs b/litellm-rust/crates/model-catalog/src/register.rs index ed34082ae57..7aaf2ed3978 100644 --- a/litellm-rust/crates/model-catalog/src/register.rs +++ b/litellm-rust/crates/model-catalog/src/register.rs @@ -1,26 +1,252 @@ -use std::collections::{BTreeSet, HashMap}; +use std::collections::BTreeMap; use indexmap::IndexMap; -use crate::wire::Fields; - -/// Where one `register` call writes and which row it inherits from -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RegistrationTarget { - /// Always the caller's key, never a case variant of it - pub write: String, - /// The first candidate that is an exact key whose row passes `provider_matches` - pub inherit_from: Option, +/// `register_model(persist_across_reloads=True)` calls, replayed above every installed document +#[derive(Clone, Debug, PartialEq)] +pub struct Registrations { + /// Each registered key's merged value, ordered by its latest registration + entries: IndexMap, + /// The registered keys whose write landed on each row key, in landing order + landed_on: IndexMap>, } -/// Python's inheritance candidates in main's order, `get_model_info`'s then `_resolve_builtin_model_cost_entry`'s -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RegistrationCandidates(pub Vec); - -/// Layer 3, and the fields `register` last wrote, which litellm owns again -#[derive(Clone, Debug, Default)] -#[expect(dead_code, reason = "07 adds the functions that read registrations")] -pub struct Registrations { - persisted: IndexMap, - registered_fields: HashMap>, +impl Default for Registrations { + fn default() -> Self { + Self { + entries: IndexMap::new(), + landed_on: IndexMap::new(), + } + } +} + +/// One registration written to a host's table: the key it was registered under and the row key it landed on +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Landed { + pub registered: String, + pub written: String, +} + +impl Registrations { + /// Keeps `value` as the key's registration, already merged with `get(key)` by the host, and makes it the latest + pub fn register(&mut self, key: &str, value: V) { + self.entries.shift_remove(key); + self.entries.insert(key.to_owned(), value); + } + + pub fn get(&self, key: &str) -> Option<&V> { + self.entries.get(key) + } + + pub fn record(&mut self, landed: &[Landed]) { + for Landed { + registered, + written, + } in landed + { + let names = self.landed_on.entry(written.clone()).or_default(); + if !names.contains(registered) { + names.push(registered.clone()); + } + } + } + + /// Every registration, oldest first + pub fn entries(&self) -> impl Iterator { + self.entries + .iter() + .map(|(key, value)| (key.as_str(), value)) + } + + /// The registrations that apply to the row under `key`, oldest first: `key`'s own and every one that landed on it + pub fn over<'a>(&'a self, key: &str) -> Vec<&'a V> { + let landed = self.landed_on.get(key).into_iter().flatten(); + let by_order: BTreeMap = landed + .map(String::as_str) + .chain([key]) + .filter_map(|name| self.entries.get_full(name)) + .map(|(order, _, value)| (order, value)) + .collect(); + by_order.into_values().collect() + } + + pub fn map(&self, mut convert: impl FnMut(&V) -> W) -> Registrations { + Registrations { + entries: self + .entries + .iter() + .map(|(key, value)| (key.clone(), convert(value))) + .collect(), + landed_on: self.landed_on.clone(), + } + } +} + +/// A host's live table that registrations are written into +pub trait RegistrationTable { + type Registration; + type Error; + + /// Writes `registration` where it lands, `None` when this table is no longer the one the host publishes + fn place( + &mut self, + key: &str, + registration: &Self::Registration, + ) -> Result, Self::Error>; +} + +/// Where a registration was written and whether a built-in row matched it +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Placed { + pub key: String, + pub matched_builtin: bool, +} + +#[derive(Debug, PartialEq)] +pub enum Pass<'a, R> { + /// Every entry was written. `unmatched` are those no built-in row matched + Finished { + landed: Vec, + unmatched: Vec<(&'a str, &'a R)>, + }, + /// The table stopped being the published one before every entry was written + Swapped { landed: Vec }, +} + +impl Pass<'_, R> { + pub fn landed(&self) -> &[Landed] { + match self { + Pass::Finished { landed, .. } | Pass::Swapped { landed } => landed, + } + } +} + +/// Writes `entries` into `table` in order +pub fn write_pass<'a, T: RegistrationTable>( + table: &mut T, + entries: &'a [(String, T::Registration)], +) -> Result, T::Error> { + let mut landed = Vec::with_capacity(entries.len()); + let mut unmatched = Vec::new(); + for (key, registration) in entries { + let Some(written) = table.place(key, registration)? else { + return Ok(Pass::Swapped { landed }); + }; + if !written.matched_builtin { + unmatched.push((key.as_str(), registration)); + } + landed.push(Landed { + registered: key.clone(), + written: written.key, + }); + } + Ok(Pass::Finished { landed, unmatched }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn landed(registered: &str, written: &str) -> Landed { + Landed { + registered: registered.to_owned(), + written: written.to_owned(), + } + } + + #[rstest] + fn reregistering_a_key_makes_it_the_latest() { + let mut registrations = Registrations::default(); + registrations.register("a", 1); + registrations.register("b", 2); + registrations.register("a", 3); + assert_eq!( + registrations.entries().collect::>(), + [("b", &2), ("a", &3)] + ); + } + + #[rstest] + fn over_applies_names_landed_on_the_row_in_registration_order() { + let mut registrations = Registrations::default(); + registrations.register("openai/o3-pro", 1); + registrations.register("o3-pro", 2); + registrations.register("other", 3); + registrations.record(&[ + landed("openai/o3-pro", "o3-pro"), + landed("o3-pro", "o3-pro"), + ]); + registrations.register("openai/o3-pro", 4); + assert_eq!(registrations.over("o3-pro"), [&2, &4]); + } + + #[rstest] + fn over_applies_a_keys_own_registration_that_landed_on_another_row() { + let mut registrations = Registrations::default(); + registrations.register("openai/o3-pro", 1); + registrations.record(&[landed("openai/o3-pro", "o3-pro")]); + assert_eq!( + ( + registrations.over("openai/o3-pro"), + registrations.over("o3-pro") + ), + (vec![&1], vec![&1]) + ); + } + + /// A table published until `swaps_after` writes, whose rows match a built-in only when the key starts with `b` + struct Table { + writes: Vec<(String, i32)>, + swaps_after: usize, + } + + impl RegistrationTable for Table { + type Registration = i32; + type Error = (); + + fn place(&mut self, key: &str, registration: &i32) -> Result, ()> { + if self.writes.len() >= self.swaps_after { + return Ok(None); + } + let written = key.to_uppercase(); + self.writes.push((written.clone(), *registration)); + Ok(Some(Placed { + matched_builtin: written.starts_with('B'), + key: written, + })) + } + } + + #[rstest] + #[case::published_throughout( + usize::MAX, + Pass::Finished { landed: vec![landed("a", "A"), landed("b", "B")], unmatched: vec![("a", &1)] }, + &[("A", 1), ("B", 2)], + )] + #[case::swapped_after_the_first_write( + 1, + Pass::Swapped { landed: vec![landed("a", "A")] }, + &[("A", 1)], + )] + fn a_pass_writes_entries_in_order_until_the_table_is_swapped( + #[case] swaps_after: usize, + #[case] expected: Pass<'static, i32>, + #[case] writes: &[(&str, i32)], + ) { + let entries = [("a".to_owned(), 1), ("b".to_owned(), 2)]; + let mut table = Table { + writes: Vec::new(), + swaps_after, + }; + assert_eq!(write_pass(&mut table, &entries), Ok(expected)); + assert_eq!( + table.writes, + writes + .iter() + .map(|(key, value)| ((*key).to_owned(), *value)) + .collect::>() + ); + } } diff --git a/litellm-rust/crates/model-catalog/src/schema.rs b/litellm-rust/crates/model-catalog/src/schema.rs index 66f9933a3e3..60c755fe94f 100644 --- a/litellm-rust/crates/model-catalog/src/schema.rs +++ b/litellm-rust/crates/model-catalog/src/schema.rs @@ -29,6 +29,8 @@ struct PriceBlocks { regional_processing_uplift_multiplier_eu: Option, /// Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%). regional_processing_uplift_multiplier_us: Option, + /// Multiplier applied to all token costs when served from the OpenAI FedRAMP endpoint (e.g. 1.10 = +10%). + fedramp_uplift_multiplier: Option, /// USD cost per web search query, keyed by search context size. search_context_cost_per_query: Option, /// Context-length or result-count tiered rates; each tier's costs apply within its range. @@ -102,9 +104,18 @@ pub fn registry_json_schema() -> Schema { "title".into(), "LiteLLM model prices and context window registry".into(), ); + definitions.insert( + "aliasRow".into(), + json!({ + "type": "object", + "required": ["alias_of"], + "properties": {"alias_of": {"type": "string"}}, + "additionalProperties": false, + }), + ); root.insert( "additionalProperties".into(), - json!({"$ref": "#/$defs/modelEntry"}), + json!({"anyOf": [{"$ref": "#/$defs/modelEntry"}, {"$ref": "#/$defs/aliasRow"}]}), ); root.insert("$defs".into(), Value::Object(definitions)); Schema::from(root.clone()) @@ -151,6 +162,7 @@ fn decorate_model_entry(schema: &mut Value) { let rates = json!({FLAT_KEY_PATTERN: {"type": "number", "minimum": 0}}); let object = schema.as_object_mut().unwrap(); object.insert("required".into(), json!(["litellm_provider"])); + object.insert("not".into(), json!({"required": ["alias_of"]})); object.insert("additionalProperties".into(), Value::Bool(true)); object.insert( "patternProperties".into(), diff --git a/litellm-rust/crates/model-catalog/src/store/AGENTS.md b/litellm-rust/crates/model-catalog/src/store/AGENTS.md index f1dacd2e217..580798da895 100644 --- a/litellm-rust/crates/model-catalog/src/store/AGENTS.md +++ b/litellm-rust/crates/model-catalog/src/store/AGENTS.md @@ -24,7 +24,7 @@ | Source | Stamp | Blind spot | |---|---|---| | `Catalog`'s `CatalogRows` | `()`, every change goes through `Store::edit`, which drops the index | none | -| python-bridge's bound `litellm.model_cost` | the dict object and its `len` | a delete plus an add between two reads, or a `litellm_provider` rewrite on an existing key | +| python-bridge's bound `litellm.model_cost` | the dict object, its `len` and its last key | deletes plus adds between two reads that put the old last key back last, or a `litellm_provider` rewrite on an existing key | A source with a blind spot calls `Store::invalidate` when it knows keys changed. Until then a key added in the blind spot is not found, and a key moved to another provider stays listed under the old one. Removed keys are still never returned diff --git a/litellm-rust/crates/model-catalog/src/vocabulary.rs b/litellm-rust/crates/model-catalog/src/vocabulary.rs index 2b7ac8bd1ab..6bf2cc55224 100644 --- a/litellm-rust/crates/model-catalog/src/vocabulary.rs +++ b/litellm-rust/crates/model-catalog/src/vocabulary.rs @@ -93,6 +93,12 @@ pub enum Endpoint { Realtime, #[serde(rename = "/v1/realtime/transcription_sessions")] RealtimeTranscriptionSessions, + #[serde(rename = "/v1/realtime/translations")] + RealtimeTranslations, + #[serde(rename = "/v1/realtime/translations/calls")] + RealtimeTranslationCalls, + #[serde(rename = "/v1/realtime/translations/client_secrets")] + RealtimeTranslationClientSecrets, #[serde(rename = "/v1/videos")] Videos, #[serde(rename = "/v1/ocr")] diff --git a/litellm-rust/crates/model-catalog/src/wire/attributes.rs b/litellm-rust/crates/model-catalog/src/wire/attributes.rs index 93c10a81524..52cd2e8525f 100644 --- a/litellm-rust/crates/model-catalog/src/wire/attributes.rs +++ b/litellm-rust/crates/model-catalog/src/wire/attributes.rs @@ -11,7 +11,7 @@ use crate::error::{DateCause, Error, Repr}; use crate::vocabulary::{ AudioFormat, EffortLevel, Endpoint, InputModality, Mode, OutputModality, VertexAiAudioApi, }; -use crate::wire::lax::Count; +use crate::wire::lax::{Count, NumericCount}; macro_rules! attributes { ($($(#[$doc:meta])* $name:ident: $ty:ty $(as $via:ty)?),* $(,)?) => { @@ -47,6 +47,8 @@ macro_rules! attributes { } attributes! { + /// Model this row is another name for; such a row carries no other key. + alias_of: String, /// Other names this entry is served under. aliases: Vec, audio_transcription_config: String, @@ -93,6 +95,10 @@ attributes! { search_free_requests: SearchFreeRequests, /// URL of the provider pricing/model page this entry was taken from. source: String, + /// Model the provider recommends migrating to before `deprecation_date`, as a catalog key when one exists. + successor_model: String, + /// URL of the provider page or feed `deprecation_date` was read from. + deprecation_source: String, /// Audio container formats the model can return. supported_audio_formats: Vec, /// OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions. @@ -104,9 +110,14 @@ attributes! { /// Cloud regions the model is available in ('global' or region ids). supported_regions: Vec, thinking_always_on: bool, - token_count_framing: u64 as Count, + /// Whether the model thinks when the request carries no thinking param. + thinking_on_by_default: bool, + token_count_framing: u64 as NumericCount, token_count_model: TokenCountModel, + token_count_responses_framing: u64 as NumericCount, + token_count_responses_tool_framing: u64 as NumericCount, token_count_tokenizer: TokenCountTokenizer, + token_count_tool_framing: u64 as NumericCount, /// Provider default tokens-per-minute limit. tpm: u64 as Count, use_openai_responses_path: bool, @@ -298,10 +309,28 @@ mod tests { #[case::known_mode("mode", json!("chat"), true)] #[case::unknown_mode("mode", json!("chatt"), false)] #[case::negative_dbu("input_dbu_cost_per_token", json!(-1), false)] + #[case::successor_model("successor_model", json!("gpt-6-luna"), true)] + #[case::successor_model_not_a_name("successor_model", json!(["gpt-6-luna"]), false)] + #[case::deprecation_source("deprecation_source", json!("https://developers.openai.com/api/docs/deprecations"), true)] + #[case::deprecation_source_not_a_url_string("deprecation_source", json!(7), false)] + #[case::thinking_on_by_default("thinking_on_by_default", json!(true), true)] + #[case::thinking_on_by_default_not_a_flag("thinking_on_by_default", json!("yes"), false)] #[case::reasoning_contexts("reasoning_contexts", json!(["reasoning", "non_reasoning"]), true)] #[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)] #[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)] #[case::reasoning_mode_not_a_string("reasoning_modes", json!([1]), false)] + #[case::framing("token_count_framing", json!(12), true)] + #[case::framing_written_as_text("token_count_framing", json!("12"), false)] + #[case::framing_written_as_whole_float("token_count_framing", json!(12.0), true)] + #[case::fractional_tool_framing("token_count_tool_framing", json!(1.5), false)] + #[case::tool_framing("token_count_tool_framing", json!(214), true)] + #[case::tool_framing_written_as_text("token_count_tool_framing", json!("214"), false)] + #[case::tool_framing_as_a_flag("token_count_tool_framing", json!(true), false)] + #[case::negative_tool_framing("token_count_tool_framing", json!(-1), false)] + #[case::responses_framing("token_count_responses_framing", json!(38), true)] + #[case::responses_framing_as_a_flag("token_count_responses_framing", json!(false), false)] + #[case::responses_tool_framing("token_count_responses_tool_framing", json!(267), true)] + #[case::responses_tool_framing_written_as_text("token_count_responses_tool_framing", json!("267"), false)] fn each_attribute_reads_its_own_type( #[case] key: &str, #[case] value: Value, diff --git a/litellm-rust/crates/model-catalog/src/wire/entry.rs b/litellm-rust/crates/model-catalog/src/wire/entry.rs index 9c155c6496a..d0bdb394720 100644 --- a/litellm-rust/crates/model-catalog/src/wire/entry.rs +++ b/litellm-rust/crates/model-catalog/src/wire/entry.rs @@ -24,6 +24,7 @@ pub struct EntryWire { pub regional_endpoint_uplift: Option, pub residency_uplift_eu: Option, pub residency_uplift_us: Option, + pub fedramp_uplift: Option, pub unknown: Fields, pub problems: Vec, } @@ -57,6 +58,7 @@ impl EntryWire { "regional_processing_uplift_multiplier_us" => { self.residency_uplift_us = reader.value(path(), value) } + "fedramp_uplift_multiplier" => self.fedramp_uplift = reader.value(path(), value), _ => self.read_open_key(reader, key, value), } } diff --git a/litellm-rust/crates/model-catalog/src/wire/lax.rs b/litellm-rust/crates/model-catalog/src/wire/lax.rs index d514a6a82c1..d98a601e8c1 100644 --- a/litellm-rust/crates/model-catalog/src/wire/lax.rs +++ b/litellm-rust/crates/model-catalog/src/wire/lax.rs @@ -8,6 +8,8 @@ pub(crate) struct LaxI64; pub(crate) struct Count; +pub(crate) struct NumericCount; + impl<'de> DeserializeAs<'de, i64> for LaxI64 { fn deserialize_as>(deserializer: D) -> Result { deserializer.deserialize_any(Self) @@ -52,6 +54,34 @@ impl<'de> DeserializeAs<'de, u64> for Count { } } +impl<'de> DeserializeAs<'de, u64> for NumericCount { + fn deserialize_as>(deserializer: D) -> Result { + deserializer.deserialize_any(Self) + } +} + +impl Visitor<'_> for NumericCount { + type Value = u64; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a whole number of zero or more") + } + + fn visit_u64(self, value: u64) -> Result { + Ok(value) + } + + fn visit_i64(self, value: i64) -> Result { + u64::try_from(value).map_err(|_| E::invalid_value(de::Unexpected::Signed(value), &self)) + } + + fn visit_f64(self, value: f64) -> Result { + integral_float(value) + .and_then(|whole| u64::try_from(whole).ok()) + .ok_or_else(|| E::invalid_value(de::Unexpected::Float(value), &self)) + } +} + fn integer_string(value: &str) -> Option { let integer = match value.split_once('.') { Some((integer, fraction)) => { @@ -99,6 +129,12 @@ mod tests { .map(DeserializeAsWrap::into_inner) } + fn numeric_count(value: Value) -> Option { + DeserializeAsWrap::::deserialize(value) + .ok() + .map(DeserializeAsWrap::into_inner) + } + use serde::Deserialize; #[rstest] @@ -125,4 +161,23 @@ mod tests { ) { assert_eq!(count(written), expected); } + + #[rstest] + #[case::integer(json!(214), Some(214))] + #[case::zero(json!(0), Some(0))] + #[case::whole_float(json!(214.0), Some(214))] + #[case::negative_zero(json!(-0.0), Some(0))] + #[case::above_i64(json!(u64::MAX), Some(u64::MAX))] + #[case::negative(json!(-1), None)] + #[case::negative_whole_float(json!(-1.0), None)] + #[case::fraction(json!(1.5), None)] + #[case::string(json!("214"), None)] + #[case::boolean(json!(true), None)] + #[case::list(json!([214]), None)] + fn numeric_counts_accept_only_whole_json_numbers( + #[case] written: Value, + #[case] expected: Option, + ) { + assert_eq!(numeric_count(written), expected); + } } diff --git a/litellm-rust/crates/model-catalog/tests/catalog.rs b/litellm-rust/crates/model-catalog/tests/catalog.rs index 31f907689d3..679c67412c8 100644 --- a/litellm-rust/crates/model-catalog/tests/catalog.rs +++ b/litellm-rust/crates/model-catalog/tests/catalog.rs @@ -139,6 +139,18 @@ fn a_bad_entry_is_reported_by_model_and_key_and_kept_raw() { json!({"First": {"litellm_provider": "t", "aliases": ["shared"]}, "Second": {"litellm_provider": "t", "aliases": ["shared"]}}), r#"Second: alias "shared" already belongs to First"# )] +#[case::alias_of_target_missing( + json!({"Real": {"litellm_provider": "t"}, "ghost-latest": {"alias_of": "Ghost"}}), + r#"ghost-latest: alias_of "Ghost" is not a model name"# +)] +#[case::alias_of_row_with_other_fields( + json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real", "mode": "chat"}}), + "real-latest.mode: an alias_of row carries no other field" +)] +#[case::alias_of_claimed_by_a_list( + json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real"}, "Other": {"litellm_provider": "t", "aliases": ["real-latest"]}}), + r#"real-latest: alias "real-latest" already belongs to Other"# +)] #[case::aliases_not_a_list( json!({"a": {"litellm_provider": "t", "aliases": "bad"}}), "a.aliases: invalid value: invalid type: string \"bad\", expected a sequence" @@ -399,3 +411,70 @@ fn a_bad_value_keeps_its_json_error_as_the_source_and_names_its_key_and_field() assert_eq!(error.key().as_deref(), Some("open/router")); assert_eq!(error.field().as_deref(), Some("input_cost_per_token")); } + +#[rstest] +fn an_alias_of_row_is_served_as_its_target_and_written_as_itself() { + let catalog = catalog(json!({ + "Real": {"litellm_provider": "t", "input_cost_per_token": 1e-6}, + "real-latest": {"alias_of": "Real"}, + "After": {"litellm_provider": "t"} + })); + let found = catalog.lookup("REAL-LATEST").unwrap(); + assert_eq!( + (found.matched_key, found.canonical_key), + ("real-latest", "Real") + ); + assert_eq!( + found.entry.field("input_cost_per_token"), + Some(&json!(1e-6)) + ); + assert_eq!(catalog.model_count(), 2); + assert_eq!( + catalog.served().map(|(key, _)| key).collect::>(), + ["Real", "After", "real-latest"] + ); + assert_eq!( + written(&catalog), + r#"{"Real":{"litellm_provider":"t","input_cost_per_token":1e-6},"real-latest":{"alias_of":"Real"},"After":{"litellm_provider":"t"}}"# + ); +} + +#[rstest] +fn a_dangling_alias_of_row_is_rejected_and_dropped_from_the_written_document() { + let parsed = loaded(json!({ + "Real": {"litellm_provider": "t"}, + "ghost-latest": {"alias_of": "Ghost"} + })); + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(removed.len(), 1); + assert!(catalog.lookup("ghost-latest").is_none()); + assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#); +} + +#[rstest] +fn an_alias_of_row_with_other_fields_is_never_served() { + let parsed = loaded(json!({ + "Real": {"litellm_provider": "t"}, + "real-latest": {"alias_of": "Real", "mode": "chat"} + })); + assert!(parsed.value.lookup("real-latest").is_none()); + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(removed.len(), 1); + assert!(catalog.lookup("real-latest").is_none()); + assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#); +} + +#[rstest] +fn a_pointer_row_that_repeats_a_listed_alias_is_dropped_and_the_model_is_kept() { + let parsed = loaded(json!({ + "Real": {"litellm_provider": "t", "aliases": ["real-latest"]}, + "real-latest": {"alias_of": "Real"} + })); + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(removed.len(), 1); + assert_eq!(catalog.lookup("real-latest").unwrap().canonical_key, "Real"); + assert_eq!( + written(&catalog), + r#"{"Real":{"litellm_provider":"t","aliases":["real-latest"]}}"# + ); +} diff --git a/litellm-rust/crates/model-catalog/tests/cost_map.rs b/litellm-rust/crates/model-catalog/tests/cost_map.rs new file mode 100644 index 00000000000..04a04170df6 --- /dev/null +++ b/litellm-rust/crates/model-catalog/tests/cost_map.rs @@ -0,0 +1,53 @@ +use litellm_model_catalog::{ + BUNDLED_REGISTRY, CostMap, ErrorKind, IntegrityLimits, Provenance, git_blob_id, +}; +use rstest::rstest; +use serde_json::json; + +fn any_size() -> IntegrityLimits { + IntegrityLimits::new(0, 0.0).unwrap() +} + +#[rstest] +#[case::not_json(b"{".as_slice(), false, ErrorKind::Json)] +#[case::empty(b"{}".as_slice(), true, ErrorKind::Empty)] +fn a_refused_document_says_whether_it_failed_the_integrity_limits( + #[case] body: &[u8], + #[case] integrity: bool, + #[case] kind: ErrorKind, +) { + let error = CostMap::read(body, None, any_size()).unwrap_err(); + assert_eq!((error.is_integrity(), error.kind()), (integrity, kind)); +} + +#[rstest] +fn provenance_is_the_read_bytes_and_etag() { + let body = br#"{"m":{"litellm_provider":"p"}}"#; + let cost_map = CostMap::read(body, Some("\"v1\"".into()), any_size()).unwrap(); + assert_eq!( + cost_map.provenance(), + &Provenance { + revision: "eaa747f210a1527d5ea7138be0f768c47578ca68".to_owned(), + etag: Some("\"v1\"".into()), + } + ); + assert_eq!( + CostMap::bundled().provenance().revision, + git_blob_id(BUNDLED_REGISTRY) + ); +} + +#[rstest] +fn every_name_of_an_aliased_row_lists_its_owner_first( + #[values("owner", "short", "owner-latest")] key: &str, +) { + let body = json!({ + "owner": {"litellm_provider": "p", "aliases": ["short"]}, + "owner-latest": {"alias_of": "owner"}, + "plain": {"litellm_provider": "p"}, + }) + .to_string(); + let cost_map = CostMap::read(body.as_bytes(), None, any_size()).unwrap(); + assert_eq!(cost_map.names(key), ["owner", "short", "owner-latest"]); + assert!(cost_map.names("plain").is_empty()); +} diff --git a/litellm-rust/crates/model-catalog/tests/schema.rs b/litellm-rust/crates/model-catalog/tests/schema.rs index 464cf069cac..5fe5832cb58 100644 --- a/litellm-rust/crates/model-catalog/tests/schema.rs +++ b/litellm-rust/crates/model-catalog/tests/schema.rs @@ -35,6 +35,7 @@ fn generated_registry_schema_validates_checked_in_catalog() { #[case(json!({"example": {"litellm_provider": "test"}}))] #[case(json!({"example": {"litellm_provider": "test", "future_field": true}}))] #[case(json!({"sample_spec": {"litellm_provider": "placeholder"}}))] +#[case(json!({"example": {"litellm_provider": "test"}, "example-latest": {"alias_of": "example"}}))] fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) { assert!(registry_validator().is_valid(&document)); } @@ -54,6 +55,9 @@ fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) #[case::empty_windows(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}))] #[case::invalid_weekday(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}))] #[case::invalid_aliases(json!({"litellm_provider": "test", "aliases": "wrong"}))] +#[case::alias_of_on_a_model(json!({"litellm_provider": "test", "alias_of": "other"}))] +#[case::alias_row_with_other_fields(json!({"alias_of": "other", "mode": "chat"}))] +#[case::alias_of_not_a_string(json!({"alias_of": 4}))] #[case::non_object_model(json!(4))] fn generated_registry_schema_rejects_invalid_entries(#[case] entry: Value) { assert!(!registry_validator().is_valid(&json!({"example": entry}))); diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 26858b2afe2..b18a1e2ca19 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -81,6 +81,7 @@ GIL handling to `litellm-host-python`. ## Bridge Shape +- Reach Python only through `litellm/rust_bridge/host/` and the `litellm.rust_bridge.execution` driver. Never import another LiteLLM module or read a `litellm` global. `litellm/rust_bridge/host/AGENTS.md` lists what may be imported directly - Prefer one stable method per top-level LiteLLM route, for example `messages(...)`, calling the matching `litellm-inference-` entrypoint. - Do not add one exported PyO3 function per provider helper unless there is a diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index a4e1c57ab1e..48206d242e4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,17 +19,26 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +rustls.workspace = true +litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true litellm-spend-clickhouse.workspace = true litellm-storage-clickhouse.workspace = true litellm-host.workspace = true +async-trait.workspace = true bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true litellm-error.workspace = true -litellm-cache-memory.workspace = true -litellm-cache-redis.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-client.workspace = true +litellm-redis-connect.workspace = true +litellm-redis-store.workspace = true +litellm-redis-resolve.workspace = true +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true +tracing.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true @@ -48,14 +57,16 @@ litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "has litellm-secrets-types.workspace = true litellm-llms-types.workspace = true litellm-router-types.workspace = true +litellm-router-core.workspace = true litellm-host-python.workspace = true litellm-model-catalog.workspace = true +litellm-model-catalog-refresh.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } -pyo3.workspace = true +indexmap = "2.14.0" +pyo3 = { workspace = true, features = ["indexmap"] } prost.workspace = true pyo3-async-runtimes.workspace = true reqwest.workspace = true -redis = { workspace = true, features = ["tls-rustls"] } serde_json.workspace = true strum.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md index 41c82187373..e79562b8762 100644 --- a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md @@ -1,7 +1,7 @@ # Cache boundary -This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates +This folder owns how Rust inference reaches the selected Python cache. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates -`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters +`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission and inference protocol composition -`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports +`python/` delegates operations to the selected Python cache without discovering configuration. Routes use only the parent module's exports diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 87dba3efcb1..fe9e06b3d23 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,32 +1,17 @@ -mod native; mod python; mod selection; -pub(crate) use native::NativeCacheHandle; pub(crate) use python::{CacheCall, PythonCache}; -pub(crate) use selection::{Cached, Selection, configure, configured_native}; +pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled}; -use litellm_cache::Error; -use pyo3::{ - exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, - prelude::*, -}; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; + +const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache"; -/// A cache call the bridge rejects itself, before or after the cache ran. #[derive(Debug, thiserror::Error)] pub(crate) enum HandleError { - #[error("recreate the v2 cache after fork")] - Forked, - #[error("cache limits must be positive")] - Limits, - #[error("namespace and a positive cache limit are required")] - RedisLimits, - #[error("cache durations must be finite and positive")] - Duration(#[source] Option), - #[error("unsupported v2 cache control: {0}")] - Control(String), - #[error("native inference needs a v2 cache, the catalog keeps Python caches on Python")] - NotNative, + #[error("this inference route does not support configured caches")] + Unsupported, #[error("cache operation without configured cache")] Unconfigured, #[error("cache arguments unavailable")] @@ -39,66 +24,10 @@ impl From for PyErr { fn from(error: HandleError) -> Self { let message = litellm_error::Report::new(&error).to_string(); match error { - HandleError::Limits - | HandleError::RedisLimits - | HandleError::Duration(_) - | HandleError::Control(_) => PyValueError::new_err(message), - HandleError::Forked - | HandleError::NotNative + HandleError::Unsupported | HandleError::Unconfigured | HandleError::MissingArguments | HandleError::NothingPending => PyRuntimeError::new_err(message), } } } - -fn cache_error(error: Error) -> PyErr { - let message = litellm_error::Report::new(&error).to_string(); - match error { - Error::InvalidEntry(_) => PyValueError::new_err(message), - Error::UnsupportedOperation => PyNotImplementedError::new_err(message), - Error::Unavailable(_) | Error::UnscopedFlush | Error::MissingPrompt => { - PyRuntimeError::new_err(message) - } - } -} - -#[cfg(test)] -mod tests { - use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, - prelude::*, - }; - use rstest::rstest; - - use super::{Error, cache_error}; - - fn refused() -> std::io::Error { - std::io::Error::other("connection refused") - } - - #[rstest] - #[case::unavailable( - Error::unavailable(refused()), - false, - "cache is unavailable: connection refused" - )] - #[case::invalid_entry( - Error::invalid_entry(refused()), - true, - "invalid cache entry: connection refused" - )] - fn python_sees_the_class_and_the_whole_chain( - #[case] error: Error, - #[case] caller_fault: bool, - #[case] message: &str, - ) { - Python::initialize(); - Python::attach(|py| { - let failure = cache_error(error); - assert_eq!(failure.is_instance_of::(py), caller_fault); - assert_eq!(failure.is_instance_of::(py), !caller_fault); - assert_eq!(failure.value(py).to_string(), message); - }); - } -} diff --git a/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md deleted file mode 100644 index 09a34d188cd..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md +++ /dev/null @@ -1,9 +0,0 @@ -# Native cache bindings - -This directory exposes the experimental V2 native cache handles (`NativeCacheHandle`) to Python. Cache algorithms and storage protocols remain in their cache crates - -Accept the cache object selected by the parent module. Do not read global `litellm.cache`, decide route admission, or select the Python cache adapter here - -Keep Python-facing cache classes and method signatures stable when reorganizing modules. Native internals stay private to this directory unless the shared cache boundary or Python module registration needs them - -Verify changes with the V2 cache tests using a freshly built extension. Test cache behavior, not module paths or file structure diff --git a/litellm-rust/crates/python-bridge/src/cache/native/mod.rs b/litellm-rust/crates/python-bridge/src/cache/native/mod.rs deleted file mode 100644 index da1460f90ed..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/native/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub(super) mod v2; - -pub(crate) use v2::NativeCacheHandle; diff --git a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs deleted file mode 100644 index c54cb641843..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs +++ /dev/null @@ -1,336 +0,0 @@ -use crate::cache::{HandleError, cache_error}; -use std::{ - sync::Arc, - time::{Duration, SystemTime, UNIX_EPOCH}, -}; - -use litellm_cache::{DeleteCache, DisconnectCache, PingCache}; -use litellm_host_python::{from_py, release_gil, to_py}; -use serde_json::Value; - -use litellm_cache_memory::InMemoryCache; -use litellm_cache_redis::{RedisCache, RedisTopology}; -use litellm_cache_response::{ - CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec, - ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, -}; -use pyo3::{prelude::*, types::PyDict}; - -#[pyclass( - frozen, - name = "NativeCacheHandle", - module = "litellm.rust_bridge._native" -)] -pub(crate) struct NativeCacheHandle { - service: Arc, - backend: Arc, - storage: Storage, - pid: u32, -} - -#[derive(Clone)] -enum Storage { - Memory(Arc>), - Redis(Arc>), -} - -impl NativeCacheHandle { - fn check_process(&self) -> PyResult<()> { - if self.pid != std::process::id() { - return Err(HandleError::Forked.into()); - } - Ok(()) - } -} - -fn request(key: String, ttl: Option) -> PyResult { - let mut request: ResponseCacheRequest = ResponseCacheRequest::new(CacheKeyInput { - preset: Some(key), - ..Default::default() - }); - request.context.ttl = ttl.map(duration).transpose()?; - Ok(request) -} - -#[pymethods] -impl NativeCacheHandle { - #[staticmethod] - #[pyo3(signature = (*, ttl=600.0, capacity=200, max_entry_bytes=4194304))] - fn memory(ttl: f64, capacity: usize, max_entry_bytes: usize) -> PyResult { - let ttl = duration(ttl)?; - if capacity == 0 || max_entry_bytes == 0 { - return Err(HandleError::Limits.into()); - } - let storage = Arc::new(InMemoryCache::new(Some(capacity), Some(ttl))); - let backend = Arc::new(ResponseCache::new(storage.clone()).with_config( - ResponseCacheConfig { - namespace: "sdk".into(), - max_entry_bytes, - }, - )); - Ok(Self { - service: backend.clone(), - backend, - storage: Storage::Memory(storage), - pid: std::process::id(), - }) - } - - #[staticmethod] - #[pyo3(signature = (url, *, namespace, ttl=600.0, max_entry_bytes=4194304))] - fn redis( - py: Python<'_>, - url: &str, - namespace: String, - ttl: f64, - max_entry_bytes: usize, - ) -> PyResult { - let ttl = duration(ttl)?; - if namespace.is_empty() || max_entry_bytes == 0 { - return Err(HandleError::RedisLimits.into()); - } - let storage = Arc::new( - release_gil(py, || { - RedisCache::connect( - url, - &RedisTopology::Standalone, - Some(ttl), - ResponseCacheCodec, - ) - }) - .map_err(cache_error)? - .with_namespace(Some(namespace.clone())), - ); - let backend = Arc::new(ResponseCache::new(storage.clone()).with_config( - ResponseCacheConfig { - namespace, - max_entry_bytes, - }, - )); - Ok(Self { - service: backend.clone(), - backend, - storage: Storage::Redis(storage), - pid: std::process::id(), - }) - } - fn get(&self, py: Python<'_>, key: String) -> PyResult> { - self.check_process()?; - let request = request(key, None)?; - let value = - release_gil(py, || self.backend.lookup(&request, now())).map_err(cache_error)?; - to_py(py, &value) - } - - #[pyo3(signature = (key, value, *, ttl=None))] - fn set( - &self, - py: Python<'_>, - key: String, - value: &Bound<'_, PyAny>, - ttl: Option, - ) -> PyResult<()> { - self.check_process()?; - let request = request(key, ttl)?; - let value: Value = from_py(value)?; - release_gil(py, || self.backend.store(&request, value, now())).map_err(cache_error) - } - - fn async_get<'py>(&self, py: Python<'py>, key: String) -> PyResult> { - self.check_process()?; - let request = request(key, None)?; - let backend = self.backend.clone(); - crate::execution::run_async( - py, - async move { backend.async_lookup(&request, now()).await }, - cache_error, - ) - } - - #[pyo3(signature = (key, value, *, ttl=None))] - fn async_set<'py>( - &self, - py: Python<'py>, - key: String, - value: &Bound<'_, PyAny>, - ttl: Option, - ) -> PyResult> { - self.check_process()?; - let request = request(key, ttl)?; - let value: Value = from_py(value)?; - let backend = self.backend.clone(); - crate::execution::run_async( - py, - async move { backend.async_store(&request, value, now()).await }, - cache_error, - ) - } - - #[pyo3(signature = (entries, *, ttl=None))] - fn async_set_many<'py>( - &self, - py: Python<'py>, - entries: &Bound<'_, PyAny>, - ttl: Option, - ) -> PyResult> { - self.check_process()?; - let entries: Vec<(String, Value)> = from_py(entries)?; - let entries = entries - .into_iter() - .map(|(key, value)| Ok((request(key, ttl)?, value))) - .collect::>>()?; - let backend = self.backend.clone(); - crate::execution::run_async( - py, - async move { backend.async_store_batch(entries, now()).await }, - cache_error, - ) - } - - fn flush(&self, py: Python<'_>) -> PyResult> { - self.check_process()?; - let backend = self.backend.clone(); - crate::execution::run_sync(py, async move { backend.async_flush().await }, cache_error) - } - - fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - let backend = self.backend.clone(); - crate::execution::run_async(py, async move { backend.async_flush().await }, cache_error) - } - - fn ping<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - let storage = self.storage.clone(); - crate::execution::run_async( - py, - async move { - match storage { - Storage::Memory(_) => Ok(true), - Storage::Redis(cache) => cache.ping().await, - } - }, - cache_error, - ) - } - - fn disconnect<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - let storage = self.storage.clone(); - crate::execution::run_async( - py, - async move { - match storage { - Storage::Memory(cache) => cache.disconnect().await, - Storage::Redis(cache) => cache.disconnect().await, - } - }, - cache_error, - ) - } - - fn delete<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { - self.check_process()?; - let storage = self.storage.clone(); - crate::execution::run_async( - py, - async move { - for key in keys { - match &storage { - Storage::Memory(cache) => cache.async_delete_cache(&key).await?, - Storage::Redis(cache) => cache.async_delete_cache(&key).await?, - } - } - Ok(()) - }, - cache_error, - ) - } -} - -fn duration(seconds: f64) -> Result { - match Duration::try_from_secs_f64(seconds) { - Ok(duration) if !duration.is_zero() => Ok(duration), - Ok(_) => Err(HandleError::Duration(None)), - Err(error) => Err(HandleError::Duration(Some(error))), - } -} - -pub(in crate::cache) fn native_handle<'py>( - configured: &Bound<'py, PyAny>, -) -> PyResult>> { - Ok(configured - .getattr_opt("cache")? - .map(|backend| backend.getattr_opt("native_handle")) - .transpose()? - .flatten() - .filter(|handle| handle.is_instance_of::())) -} - -pub(in crate::cache) fn configured( - configured: &Bound<'_, PyAny>, - kwargs: &Bound<'_, PyDict>, -) -> PyResult<( - Option>, - litellm_cache_response::CacheOptions, -)> { - let handle = native_handle(configured)?.ok_or(HandleError::NotNative)?; - let cache = handle.extract::>()?; - cache.check_process()?; - let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none()); - let controls = controls - .as_ref() - .map(|value| value.cast::()) - .transpose()?; - if let Some(controls) = controls { - for name in controls.keys() { - let name = name.extract::()?; - if !matches!( - name.as_str(), - "no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache" - ) { - return Err(HandleError::Control(name).into()); - } - } - } - let boolean = |name: &str| -> PyResult { - controls - .map(|values| values.get_item(name)) - .transpose()? - .flatten() - .map(|value| value.extract()) - .transpose() - .map(|value| value.unwrap_or(false)) - }; - let seconds = |name: &str| -> PyResult> { - controls - .map(|values| values.get_item(name)) - .transpose()? - .flatten() - .map(|value| Ok(duration(value.extract()?)?)) - .transpose() - }; - Ok(( - Some(cache.service.clone()), - litellm_cache_response::CacheOptions { - policy: litellm_cache_response::CachePolicy { - caching: kwargs - .get_item("caching")? - .filter(|value| !value.is_none()) - .map(|value| value.extract()) - .transpose()?, - no_cache: boolean("no-cache")?, - no_store: boolean("no-store")?, - ttl: seconds("ttl")?, - max_age: seconds("s-max-age")?.or(seconds("s-maxage")?), - }, - scope: litellm_cache_response::CacheScope::Shared, - }, - )) -} - -fn now() -> Duration { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() -} diff --git a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md index 81a0997f02a..3b516251f43 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md @@ -4,6 +4,6 @@ This directory lets Rust inference use a selected Python cache. `service.rs` imp Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details -Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions +Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. The one exception is the asynchronous unary messages store: Python's `_async_set_cache` writes the returned response from `create_cache_write_task` after the router has stamped its `_hidden_params`, so `host.rs` defers that store until the deployment transform has returned the response the caller receives and hands that response to `litellm.rust_bridge.host.cache.store_response`, which schedules the same write task. Lookups, sync stores and stream stores stay in the caller's task. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs index 82b7d881df6..a8cbd0ce607 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/host.rs @@ -9,7 +9,7 @@ use pyo3::{ use serde_json::Value; use super::service::CacheCall; -use crate::cache::HandleError; +use crate::cache::{CACHE_HOST_MODULE, HandleError}; enum Pending { Lookup(Reply, Error>>), @@ -21,6 +21,7 @@ pub(crate) struct PythonCache { arguments: Option>, pending: Option, asynchronous: bool, + deferred: bool, } impl PythonCache { @@ -30,6 +31,7 @@ impl PythonCache { arguments: None, pending: None, asynchronous, + deferred: false, } } @@ -64,7 +66,14 @@ impl PythonCache { None, ) } - CacheCall::Store { value, reply } => { + CacheCall::Store { + unary: true, reply, .. + } if self.asynchronous => { + self.deferred = true; + reply.send(Ok(())); + return Ok(None); + } + CacheCall::Store { value, reply, .. } => { self.pending = Some(Pending::Store(reply)); ( if self.asynchronous { @@ -89,6 +98,19 @@ impl PythonCache { self.resume(py, result) } + pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> { + if !std::mem::take(&mut self.deferred) { + return Ok(()); + } + let (Some(cache), Some(arguments)) = (self.cache.as_ref(), self.arguments.as_ref()) else { + return Err(HandleError::Unconfigured.into()); + }; + py.import(CACHE_HOST_MODULE)? + .getattr("store_response")? + .call((cache, response), Some(&arguments.bind(py).copy()?))?; + Ok(()) + } + pub fn resume( &mut self, py: Python<'_>, @@ -124,6 +146,7 @@ impl PythonCache { pub fn close(&mut self) { self.pending = None; + self.deferred = false; self.cache = None; self.arguments = None; } diff --git a/litellm-rust/crates/python-bridge/src/cache/python/service.rs b/litellm-rust/crates/python-bridge/src/cache/python/service.rs index 38872b85d30..018c8647043 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/service.rs @@ -45,6 +45,7 @@ pub(crate) enum CacheCall { }, Store { value: Value, + unary: bool, reply: Reply>, }, } @@ -101,8 +102,13 @@ where ) -> Pin> + Send + 'a>> { Box::pin(async move { let value = to_python(value)?; + let unary = value.get(STREAM_EVENTS_KEY).is_none(); self.services - .call(|reply| CacheCall::Store { value, reply }) + .call(|reply| CacheCall::Store { + value, + unary, + reply, + }) .await .map_err(Error::unavailable)? }) diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index 00fc7cd822f..78c1606ea40 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -1,15 +1,10 @@ -use super::{native, python}; -use litellm_cache_response::{ - CacheOptions, CachePolicy, CacheScope, ResponseCacheService, ScopedCache, -}; +use super::{CACHE_HOST_MODULE, HandleError, python}; +use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache}; use litellm_host::{ machine::{HostServices, MachineFault}, protocol::Protocol, }; use pyo3::{prelude::*, types::PyDict}; -use std::sync::Arc; - -const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache"; pub(crate) struct Cached

(std::marker::PhantomData

); @@ -24,7 +19,6 @@ impl Protocol for Cached

{ enum Backend { Disabled, - Native(Arc), Python { namespace: String }, } @@ -43,7 +37,6 @@ impl Selection { { let service = match self.backend { Backend::Disabled => None, - Backend::Native(service) => Some(service), Backend::Python { namespace } => Some(python::service(services, namespace)), }; ( @@ -76,37 +69,21 @@ fn selected_cache<'py>( Ok(Some(configured)) } -pub(crate) fn configured_native( +pub(crate) fn ensure_cache_disabled( py: Python<'_>, kwargs: &Bound<'_, PyDict>, call_type: &str, -) -> PyResult<( - Option>, - litellm_cache_response::CacheOptions, -)> { +) -> PyResult<()> { let Some(configured) = selected_cache(py, kwargs, call_type)? else { - return Ok(( - None, - litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared), - )); + return Ok(()); }; - native_configuration(&configured, kwargs) -} - -fn native_configuration( - configured: &Bound<'_, PyAny>, - kwargs: &Bound<'_, PyDict>, -) -> PyResult<(Option>, CacheOptions)> { - if !configured + if configured .call_method("should_use_cache", (), Some(kwargs))? .extract::()? { - return Ok(( - None, - litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared), - )); + return Err(HandleError::Unsupported.into()); } - native::v2::configured(configured, kwargs) + Ok(()) } pub(crate) fn configure( @@ -122,13 +99,6 @@ pub(crate) fn configure( options: CacheOptions::new(CacheScope::Shared), }); }; - if native::v2::native_handle(&cache)?.is_some() { - let (native, options) = native_configuration(&cache, arguments)?; - return Ok(Selection { - backend: native.map_or(Backend::Disabled, Backend::Native), - options, - }); - } let enabled = cache .call_method("should_use_cache", (), Some(arguments))? .extract::()?; diff --git a/litellm-rust/crates/python-bridge/src/credentials.rs b/litellm-rust/crates/python-bridge/src/credentials.rs index a4bb064aae0..adc52a7e77a 100644 --- a/litellm-rust/crates/python-bridge/src/credentials.rs +++ b/litellm-rust/crates/python-bridge/src/credentials.rs @@ -2,7 +2,9 @@ //! keyword arguments and acquired on the host's own thread when the call asks for one. use crate::callable::wrap_failure; -use litellm_auth::{ResolvedCredential, SecretValue}; +use litellm_auth::ResolvedCredential; +use litellm_secret_value::SecretValue; + use pyo3::{ exceptions::PyTypeError, gc::{PyTraverseError, PyVisit}, diff --git a/litellm-rust/crates/python-bridge/src/db.rs b/litellm-rust/crates/python-bridge/src/db.rs new file mode 100644 index 00000000000..03fb30b986f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/db.rs @@ -0,0 +1,19 @@ +use litellm_host_python::ready_future; +use pyo3::{prelude::*, types::PyNone}; + +#[pyclass(frozen)] +pub(crate) struct NativeDatabase {} + +// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first +// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens +#[pymethods] +impl NativeDatabase { + #[staticmethod] + fn connect(py: Python<'_>) -> PyResult> { + ready_future(py, Bound::new(py, NativeDatabase {})?.as_any()) + } + + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + ready_future(py, PyNone::get(py).as_any()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index b42c08b7e78..db5696d52ce 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -23,6 +23,19 @@ pyo3::create_exception!( ); const MESSAGES_REQUEST_ERROR: &str = "messages_request_error"; +/// The rejection is a parameter the model does not support, which Python raises as +/// `UnsupportedParamsError`. +const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param"; +/// The phase label (`Connection` or `Read`) of a timed-out call, for Python's timeout text, or +/// `Body` for a streamed body read that timed out after the head arrived. +const MESSAGES_TIMEOUT: &str = "messages_timeout"; +// `(received, expected, reset)` of a streamed body the connection ended early. +const MESSAGES_BODY_CUT: &str = "messages_body_cut"; +/// No credential reached the provider, which Python raises as `AuthenticationError`. +const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential"; +/// Python refuses the request only once its logger knows the resolved model and provider (the +/// provider config's request transform), so the failure logs them. +const MESSAGES_LATE_REJECTION: &str = "messages_late_rejection"; pub(crate) fn upstream_error( py: Python<'_>, @@ -78,17 +91,60 @@ pub(crate) fn route_failure( return Ok(raised); } if let Some((status, body)) = error.upstream_response() { - return public(upstream_error(py, status, body, &HeaderMap::new())?); + let headers = error + .upstream_headers() + .iter() + .filter_map(|(name, value)| { + Some(( + reqwest::header::HeaderName::from_bytes(name.as_bytes()).ok()?, + reqwest::header::HeaderValue::from_str(value).ok()?, + )) + }) + .collect::(); + return public(upstream_error(py, status, body, &headers)?); } let native = match (marker, error.rejected_request()) { (RequestMarker::Messages, Some(rejection)) => { let native = PyValueError::new_err(Report::new(rejection).to_string()); - native.value(py).setattr(MESSAGES_REQUEST_ERROR, true)?; + let value = native.value(py); + value.setattr(MESSAGES_REQUEST_ERROR, true)?; + if matches!( + rejection.downcast_ref::(), + Some( + litellm_llms::ErrorDetail::UnsupportedParameter { .. } + | litellm_llms::ErrorDetail::NonNumericThinkingBudget { .. } + ) + ) { + value.setattr(MESSAGES_UNSUPPORTED_PARAM, true)?; + } + if matches!( + rejection.downcast_ref::(), + Some(litellm_llms::ErrorDetail::NonNumericThinkingBudget { .. }) + ) { + value.setattr(MESSAGES_LATE_REJECTION, true)?; + } native } - (RequestMarker::Messages, None) | (RequestMarker::Unmarked, _) => { - by_fault(error.is_request(), &error) + (RequestMarker::Messages, None) => { + let native = by_fault(error.is_request(), &error); + let value = native.value(py); + if let Some(connect) = error.timeout() { + let phase = match (connect, error.body_read_timeout()) { + (true, _) => "Connection", + (false, true) => "Body", + (false, false) => "Read", + }; + value.setattr(MESSAGES_TIMEOUT, phase)?; + } + if let Some(cut) = error.body_cut() { + value.setattr(MESSAGES_BODY_CUT, (cut.received, cut.expected, cut.reset))?; + } + if error.is_missing_credential() { + value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?; + } + native } + (RequestMarker::Unmarked, _) => by_fault(error.is_request(), &error), }; public(native) } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index b4e20a8f750..57c58c3f7cd 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -108,13 +108,38 @@ pub(crate) fn call_config( Some(snapshot) => configured(&snapshot)?, None => HttpSettingsLayer::default(), }, - ]) - .without_missing_files(&|path: &Path| path.exists()); - let resolution = Resolution::from(&settings); + ]); + Ok(resolved(py, settings)) +} + +pub(crate) fn settings_config( + py: Python<'_>, + use_http_settings: bool, +) -> PyResult { + let snapshot = if use_http_settings { + PythonSettings::Http.read_or_unset(py)? + } else { + None + }; + let layer = match snapshot { + Some(snapshot) => configured(&snapshot)?, + None => HttpSettingsLayer::default(), + }; + let settings = HttpSettings::from_layers([ + HttpSettingsLayer::from_environment(&ProcessEnvironment), + layer, + ]); + Ok(resolved(py, settings)) +} + +/// The client config `settings` resolve to, warning once per process about each setting the client cannot apply +fn resolved(py: Python<'_>, settings: HttpSettings) -> HttpClientConfig { + let resolution = + Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists())); for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) { crate::logger::capture(py).scope(|| litellm_tracing::warn!("{unsupported}")); } - Ok(resolution.config) + resolution.config } pub(crate) fn provider_client( diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 683e90d02a2..7edf13035c1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -2,15 +2,19 @@ mod cache; mod callable; mod coercion; mod credentials; +mod db; mod diagnostics; mod errors; mod execution; mod http; mod logger; mod marshal; +mod model_cost; mod model_registry; mod preflight; mod python_settings; +mod redis; +mod router; mod routes; mod secrets; mod tokenizer; @@ -21,6 +25,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; @@ -31,8 +37,41 @@ mod _native { #[pymodule_export] use crate::logger::NativeDiagnosticProcessor; #[pymodule_export] + use crate::model_cost::{ + NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched, NativeInstalled, + NativeModelCostMap, NativeModelCostMapState, NativeProvenance, NativeRegistered, + NativeRejected, NativeTableSwapped, + }; + #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states}; + #[pymodule_export] + use crate::redis::client::NativeRedisClient; + #[pymodule_export] + use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes}; + #[pymodule_export] + use crate::redis::store::rate_limits::{ + NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired, + PyGuardedReply, PyWindowCount, + }; + #[pymodule_export] + use crate::redis::store::request::{NativeRequestRedis, NativeSend}; + #[pymodule_export] + use crate::redis::store::router_usage::NativeRouterUsage; + #[pymodule_export] + use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins; + #[pymodule_export] + use crate::redis::store::spend_counters::{ + NativeSpendCounters, PyPostCallSettlement, PySpendRead, + }; + #[pymodule_export] + use crate::redis::store::{ + NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns, + }; + #[pymodule_export] + use crate::router::{RouterCall, RouterCoreHandle}; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -60,12 +99,9 @@ mod _native { #[pymodule_init] fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { + let _ = rustls::crypto::ring::default_provider().install_default(); let py = module.py(); let dict = module.dict(); - dict.set_item( - "NativeCacheHandle", - py.get_type::(), - )?; dict.set_item( "_SecretManagerRuntime", py.get_type::(), @@ -89,7 +125,6 @@ mod tests { Python::initialize(); Python::attach(|py| { let mut expected = vec![ - "NativeCacheHandle", "RustBridgeDeclined", "RustUpstreamError", "ForkedAfterNativeRuntimeStarted", @@ -110,12 +145,49 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", + "NativeRedisClient", + "NativeRedisBreaker", + "NativeAdmission", + "NativeRequestRedis", + "NativeSend", + "NativeAuthCache", + "NativeRateLimits", + "WindowCount", + "AtomicApplied", + "AtomicOverLimit", + "GuardedReply", + "GaugesAcquired", + "GaugeOverLimit", + "NativePending", + "NativeRouterCooldowns", + "NativeRouterUsage", + "NativeSensitiveRoutePins", + "NativeRedisStoreError", + "NativeSpendCounters", + "PostCallSettlement", + "SpendRead", + "LocalSizes", + "native_redis_breaker_states", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", + "NativeDatabase", "trace_encode_error", "bundled_model_cost_map", + "NativeModelCostMap", + "NativeModelCostMapState", + "NativeInstalled", + "NativeRejected", + "NativeRegistered", + "NativeTableSwapped", + "NativeProvenance", + "NativeCostMapSource", + "NativeFetchPolicy", + "NativeFetched", + "NativeFetchFailed", "TokenCounter", "Tokenizer", + "RouterCoreHandle", + "RouterCall", "gil_stats", "process_state_started", "reserve_process_for_forking", @@ -136,4 +208,18 @@ mod tests { assert_eq!(public_names, expected); }); } + #[rstest::rstest] + fn module_import_installs_ring_as_the_process_crypto_provider() { + Python::initialize(); + Python::attach(|py| { + native_module(py); + }); + + let installed = rustls::crypto::CryptoProvider::get_default() + .expect("importing _native should install a process crypto provider"); + assert_eq!( + format!("{installed:?}"), + format!("{:?}", rustls::crypto::ring::default_provider()) + ); + } } diff --git a/litellm-rust/crates/python-bridge/src/logger/mod.rs b/litellm-rust/crates/python-bridge/src/logger/mod.rs index 8f40d94f458..5032b418404 100644 --- a/litellm-rust/crates/python-bridge/src/logger/mod.rs +++ b/litellm-rust/crates/python-bridge/src/logger/mod.rs @@ -7,7 +7,7 @@ use litellm_tracing::{DiagnosticInput, Level, Logger, Metadata, Policy, Processo use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; -const MODULE: &str = "litellm.rust_bridge.logger"; +const MODULE: &str = "litellm.rust_bridge.host.logger"; type NativeDiagnosticOutput = (String, Option, Option, Vec, bool); #[pyclass] diff --git a/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md b/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md new file mode 100644 index 00000000000..4ba23aac1ce --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md @@ -0,0 +1,16 @@ +# Model cost bridge + +- `NativeModelCostMap` holds inputs only, the installed `CostMap`, kept `Registrations` and the name index. Python's `litellm.model_cost` is the only copy of the live rows +- Every call that reads or writes rows takes the dict as an argument. Never store rows or build a second table here +- `LiveDict` holds a reference to that same dict and reads it under the GIL on every query, so a Python write is seen on the next read + - It needs the GIL, so Rust code running without the GIL cannot use it + - Its stamp is the dict's address, length and last key. [model-catalog/src/store/AGENTS.md](../../../model-catalog/src/store/AGENTS.md) says what that misses +- `host.rs` is the only code that calls Python, by name through `MODULE`. Calls happen with the GIL held and no `NativeModelCostMap` lock held + - `resolve` reaches `find_key` on this same object, so holding a lock across a host call deadlocks + - Test `a_host_lookup_that_reaches_find_key_on_the_same_map_finishes` guards this +- Rows and registrations cross as Python objects, never JSON, because router `model_info` holds `datetime` +- `table::build` serves the document as main's `_expand_model_aliases` did, so all names of a model share one row object made by `host.entry` +- `install` touches no state until `CostMap::read` succeeds. `NativeRejected` means nothing changed +- `register` keeps a registration only after its pass finished, so an exception in a host call keeps nothing +- `rows=None` writes each entry into whichever table is published at that moment and never reports a swap +- `fetch_blocking` runs on a throwaway runtime with an unpooled client so an importing process can still fork. `fetch` is for workers after the fork diff --git a/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs b/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs new file mode 100644 index 00000000000..dff6296cf91 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs @@ -0,0 +1,238 @@ +use std::time::Duration; + +use litellm_http::{ClientVariant, HttpClientConfig}; +use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +/// A cost map URL and how to connect to it. An `s3://` URL reaches S3 as an AWS SDK finds it in the environment, +/// with LiteLLM's `AWS_REGION_NAME` winning over the region it finds +#[pyclass( + frozen, + name = "NativeCostMapSource", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeCostMapSource { + url: String, + http: HttpClientConfig, +} + +#[pymethods] +impl NativeCostMapSource { + #[new] + #[pyo3(signature = (url, *, use_http_settings))] + fn new(py: Python<'_>, url: String, use_http_settings: bool) -> PyResult { + Ok(Self { + url, + http: crate::http::settings_config(py, use_http_settings)?, + }) + } + + /// Fetches on the shared runtime with a pooled client + fn fetch<'py>( + &self, + py: Python<'py>, + policy: PyRef<'_, NativeFetchPolicy>, + ) -> PyResult> { + let client = crate::http::pool() + .client(&self.http, ClientVariant::Provider) + .map_err(crate::http::client_error)?; + let url = self.url.clone(); + let policy = policy.0; + crate::execution::run_async_value(py, async move { + Ok(Outcome( + litellm_model_catalog_refresh::fetch(&client, &s3(), &url, policy, Start::Fresh) + .await, + )) + }) + } + + /// Fetches on a throwaway runtime with a client of its own. `resume_from` continues a fetch whose first attempt + /// failed retryably + #[pyo3(signature = (policy, resume_from = None))] + fn fetch_blocking( + &self, + py: Python<'_>, + policy: PyRef<'_, NativeFetchPolicy>, + resume_from: Option>, + ) -> PyResult> { + let start = resume_from.map_or(Start::Fresh, |failed| Start::After(failed.0.clone())); + let (config, url, policy) = (self.http.clone(), self.url.clone(), policy.0); + let logger = crate::logger::capture(py); + let fetched = litellm_host_python::release_gil(py, move || { + logger.scope(|| { + let client = crate::http::pool() + .unpooled(&config, ClientVariant::Provider) + .map_err(crate::http::client_error)?; + litellm_host_python::run_isolated(litellm_model_catalog_refresh::fetch( + &client, + &s3(), + &url, + policy, + start, + )) + .map_err(PyErr::from) + }) + })?; + Outcome(fetched).into_pyobject(py).map(Bound::unbind) + } +} + +fn s3() -> S3Access { + S3Access::Environment { + region: std::env::var("AWS_REGION_NAME").ok(), + } +} + +#[pyclass( + frozen, + name = "NativeFetchPolicy", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeFetchPolicy(RetryPolicy); + +#[pymethods] +impl NativeFetchPolicy { + #[new] + /// Retries wait an exponential backoff with jitter from `min_delay_seconds` up to `max_wait_seconds`, or the + /// server's `Retry-After` up to that cap + #[pyo3(signature = (*, max_attempts, max_wait_seconds, timeout_seconds, min_delay_seconds))] + fn new( + max_attempts: u32, + max_wait_seconds: f64, + timeout_seconds: f64, + min_delay_seconds: f64, + ) -> PyResult { + if max_attempts == 0 { + return Err(PyValueError::new_err("max_attempts must be at least 1")); + } + let seconds = |name: &str, value: f64| { + Duration::try_from_secs_f64(value).map_err(|_| { + PyValueError::new_err(format!("{name} must be a non-negative number of seconds")) + }) + }; + let min_delay = seconds("min_delay_seconds", min_delay_seconds)?; + let max_delay = seconds("max_wait_seconds", max_wait_seconds)?; + if min_delay > max_delay { + return Err(PyValueError::new_err( + "min_delay_seconds must not exceed max_wait_seconds", + )); + } + Ok(Self(RetryPolicy { + max_attempts, + timeout: seconds("timeout_seconds", timeout_seconds)?, + min_delay, + max_delay, + })) + } +} + +#[pyclass(frozen, name = "NativeFetched", module = "litellm.rust_bridge._native")] +pub(crate) struct NativeFetched { + #[pyo3(get)] + body: Py, + #[pyo3(get)] + etag: Option, +} + +#[pyclass( + frozen, + skip_from_py_object, + name = "NativeFetchFailed", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeFetchFailed(Failed); + +#[pymethods] +impl NativeFetchFailed { + #[getter] + fn reason(&self) -> &str { + &self.0.reason + } + + /// Another attempt could succeed. Only the attempt limit ended the fetch + #[getter] + fn retryable(&self) -> bool { + self.0.retryable + } + + #[getter] + fn retry_after(&self) -> Option { + self.0.retry_after.map(|wait| wait.as_secs_f64()) + } +} + +struct Outcome(Result); + +impl<'py> IntoPyObject<'py> for Outcome { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + match self.0 { + Ok(fetched) => Bound::new( + py, + NativeFetched { + body: PyBytes::new(py, &fetched.body).unbind(), + etag: fetched.etag, + }, + ) + .map(Bound::into_any), + Err(failed) => Bound::new(py, NativeFetchFailed(failed)).map(Bound::into_any), + } + } +} + +#[cfg(test)] +mod tests { + use std::io::{BufRead, BufReader, Write}; + use std::net::TcpListener; + + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::equal_bounds(2.0, 2.0, true)] + #[case::min_below_max(2.0, 30.0, true)] + #[case::min_above_max(31.0, 30.0, false)] + fn a_policy_accepts_a_min_delay_only_up_to_the_max_wait( + #[case] min_delay_seconds: f64, + #[case] max_wait_seconds: f64, + #[case] accepted: bool, + ) { + Python::initialize(); + let policy = NativeFetchPolicy::new(3, max_wait_seconds, 5.0, min_delay_seconds); + assert_eq!(policy.is_ok(), accepted); + } + + #[rstest] + fn a_blocking_fetch_hands_python_the_document_without_starting_the_shared_runtime() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/cost.json", listener.local_addr().unwrap()); + std::thread::spawn(move || { + let mut stream = listener.incoming().next().unwrap().unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut line = String::new(); + while reader.read_line(&mut line).unwrap() > 2 { + line.clear(); + } + stream + .write_all(b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}") + .unwrap(); + }); + Python::initialize(); + Python::attach(|py| { + let source = NativeCostMapSource::new(py, url, false).unwrap(); + let policy = + Bound::new(py, NativeFetchPolicy::new(3, 30.0, 5.0, 2.0).unwrap()).unwrap(); + let outcome = source.fetch_blocking(py, policy.borrow(), None).unwrap(); + let fetched = outcome.bind(py).cast::().unwrap().get(); + assert_eq!(fetched.body.bind(py).as_bytes(), b"{}"); + assert_eq!(fetched.etag.as_deref(), Some("\"v1\"")); + }); + assert!(!litellm_host_python::runtime_started()); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/host.rs b/litellm-rust/crates/python-bridge/src/model_cost/host.rs new file mode 100644 index 00000000000..25fbbefbb73 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/host.rs @@ -0,0 +1,81 @@ +//! Calls into `litellm.rust_bridge.host.model_cost` + +use pyo3::call::PyCallArgs; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +pub(crate) const MODULE: &str = "litellm.rust_bridge.host.model_cost"; + +/// Where a registration lands in the published table and the row it starts from +#[derive(FromPyObject)] +pub(crate) struct Landing<'py> { + #[pyo3(attribute)] + pub(crate) key: String, + #[pyo3(attribute)] + pub(crate) base: Bound<'py, PyAny>, + #[pyo3(attribute)] + pub(crate) matched_builtin: bool, +} + +fn call<'py>( + py: Python<'py>, + name: &str, + args: impl PyCallArgs<'py>, +) -> PyResult> { + py.import(MODULE)?.getattr(name)?.call1(args) +} + +/// The row object every name in `names` shares, the owner first +pub(crate) fn entry<'py>( + fields: Bound<'py, PyAny>, + names: &[String], +) -> PyResult> { + let py = fields.py(); + call(py, "entry", (fields, PyTuple::new(py, names)?)) +} + +/// Where `registration` lands in `rows`, `None` when `rows` is no longer the published table. `rows` of `None` is the +/// published table +pub(crate) fn resolve<'py>( + rows: Option<&Bound<'py, PyDict>>, + key: &str, + registration: &Bound<'py, PyAny>, + custom_llm_provider: Option<&str>, +) -> PyResult>> { + call( + registration.py(), + "resolve", + (rows, key, registration, custom_llm_provider), + )? + .extract() +} + +/// The row `registration` writes over `current` +pub(crate) fn merged<'py>( + current: Option<&Bound<'py, PyAny>>, + landing: &Landing<'py>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + call( + registration.py(), + "merged", + (current, &landing.base, registration), + ) +} + +/// The registration kept for a key: `registration` merged into the one kept before it +pub(crate) fn kept<'py>( + previous: Option<&Bound<'py, PyDict>>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + Ok(call(registration.py(), "kept", (previous, registration))?.cast_into::()?) +} + +pub(crate) fn published(py: Python<'_>) -> PyResult> { + Ok(call(py, "published", ())?.cast_into::()?) +} + +/// Writes `written` into `rows`, the published table, through the host, which tracks the table's version +pub(crate) fn write(rows: &Bound<'_, PyDict>, written: &Bound<'_, PyDict>) -> PyResult<()> { + call(rows.py(), "write", (rows, written)).map(drop) +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/live.rs b/litellm-rust/crates/python-bridge/src/model_cost/live.rs new file mode 100644 index 00000000000..0b7f071127d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/live.rs @@ -0,0 +1,176 @@ +use std::borrow::Cow; + +use litellm_model_catalog::Rows; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +/// The dict Python last handed in as `litellm.model_cost`, read under the GIL on every query +pub(crate) struct LiveDict(Py); + +/// One row exactly as the dict holds it +pub(crate) struct LiveRow(Py); + +impl Clone for LiveRow { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } +} + +impl LiveDict { + pub(crate) fn new(dict: &Bound<'_, PyDict>) -> Self { + Self(dict.clone().unbind()) + } + + pub(crate) fn is(&self, dict: &Bound<'_, PyDict>) -> bool { + self.0.bind(dict.py()).is(dict) + } +} + +impl Rows for LiveDict { + type Row = LiveRow; + /// The dict's address, length and last key + type Stamp = (usize, usize, Option); + + fn fold(name: &str) -> String { + name.to_lowercase() + } + + fn stamp(&self) -> Self::Stamp { + Python::attach(|py| { + let rows = self.0.bind(py); + let last = rows + .call_method0(intern!(py, "__reversed__")) + .and_then(|keys| keys.try_iter()) + .ok() + .and_then(|mut keys| keys.next()) + .and_then(Result::ok) + .and_then(|key| key.extract::().ok()); + (self.0.as_ptr() as usize, rows.len(), last) + }) + } + + fn contains(&self, key: &str) -> bool { + Python::attach(|py| self.0.bind(py).contains(key).unwrap_or(false)) + } + + fn row(&self, key: &str) -> Option> { + Python::attach(|py| { + let row = self.0.bind(py).get_item(key).ok().flatten()?; + Some(Cow::Owned(LiveRow(row.unbind()))) + }) + } + + /// Groups each key by its row's `litellm_provider` when that is a string + fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) { + Python::attach(|py| { + for (key, row) in self.0.bind(py).iter() { + let Ok(key) = key.extract::() else { + continue; + }; + let provider = row + .cast::() + .ok() + .and_then(|row| row.get_item("litellm_provider").ok().flatten()) + .and_then(|provider| provider.extract::().ok()); + visit(&key, provider.as_deref()); + } + }); + } +} + +#[cfg(test)] +mod tests { + use litellm_model_catalog::Store; + use pyo3::ffi::c_str; + use rstest::rstest; + + use super::*; + + fn dict<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + py.eval(source, None, None) + .unwrap() + .cast_into::() + .unwrap() + } + + #[rstest] + fn a_key_added_to_the_dict_is_found_on_the_next_query() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}")); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.find_key("GPT-4O"), Some("gpt-4o".to_owned())); + rows.set_item("Claude-X", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("claude-x"), Some("Claude-X".to_owned())); + }); + } + + #[rstest] + fn a_removed_key_is_never_returned() { + Python::initialize(); + Python::attach(|py| { + let rows = dict( + py, + c_str!("{'gpt-4o': {'litellm_provider': 'openai'}, 'o3': {}}"), + ); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.keys_in("openai"), ["gpt-4o"]); + rows.del_item("gpt-4o").unwrap(); + rows.set_item("o4", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("GPT-4O"), None); + assert!(store.keys_in("openai").is_empty()); + }); + } + + #[rstest] + fn a_key_swapped_in_for_a_removed_one_is_found() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {}, 'Lit-Gone': {}}")); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.find_key("lit-gone"), Some("Lit-Gone".to_owned())); + rows.del_item("Lit-Gone").unwrap(); + rows.set_item("Lit-New", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("lit-new"), Some("Lit-New".to_owned())); + assert_eq!(store.find_key("lit-gone"), None); + }); + } + + #[rstest] + fn two_keys_folding_alike_resolve_to_the_later_one() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'Model': {}, 'MODEL': {}}")); + assert_eq!( + Store::new(LiveDict::new(&rows)).find_key("model"), + Some("MODEL".to_owned()) + ); + }); + } + + #[rstest] + fn rows_are_the_objects_the_dict_holds() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}")); + let store = Store::new(LiveDict::new(&rows)); + let row = store.get("gpt-4o").unwrap().into_owned(); + assert!(row.0.bind(py).is(rows.get_item("gpt-4o").unwrap().unwrap())); + }); + } + + #[rstest] + fn providers_group_only_string_providers() { + Python::initialize(); + Python::attach(|py| { + let rows = dict( + py, + c_str!( + "{'a': {'litellm_provider': 'p'}, 'b': {'litellm_provider': None}, 'c': 1, 1: {'litellm_provider': 'p'}}" + ), + ); + assert_eq!(Store::new(LiveDict::new(&rows)).keys_in("p"), ["a"]); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/mod.rs b/litellm-rust/crates/python-bridge/src/model_cost/mod.rs new file mode 100644 index 00000000000..99451033122 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/mod.rs @@ -0,0 +1,497 @@ +mod fetch; +mod host; +mod live; +mod table; + +use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard}; + +use litellm_host_python::to_py; +use litellm_model_catalog::{ + CostMap, IntegrityLimits, Landed, Pass, Placed, Provenance, RegistrationTable, Registrations, + Store, write_pass, +}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +pub(crate) use fetch::{NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched}; +use live::LiveDict; + +/// The installed document and the `register_model` calls kept above it, each as the dict Python merged +struct Inputs { + cost_map: CostMap, + registrations: Registrations>, +} + +impl Inputs { + fn clone_ref(&self, py: Python<'_>) -> Self { + Self { + cost_map: self.cost_map.clone(), + registrations: self.registrations.map(|kept| kept.clone_ref(py)), + } + } +} + +/// The cost map's inputs and the name index over `litellm.model_cost` +#[pyclass( + frozen, + name = "NativeModelCostMap", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeModelCostMap { + inputs: RwLock, + limits: IntegrityLimits, + index: RwLock>>>, +} + +#[pyclass( + frozen, + name = "NativeInstalled", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeInstalled { + /// The new live table, the document without registrations + #[pyo3(get)] + rows: Py, + /// How many keys the document itself serves, aliases included and registrations left out + #[pyo3(get)] + model_count: usize, + /// The document's `fallback_generalizations` block as written, `None` when it has none + #[pyo3(get)] + fallback_generalizations: Py, +} + +#[pyclass( + frozen, + name = "NativeRejected", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeRejected { + #[pyo3(get)] + reason: String, + /// The document parsed but has fewer models than the integrity limits allow + #[pyo3(get)] + integrity: bool, +} + +/// A finished registration pass, with the registrations no built-in row matched, by registered key +#[pyclass( + frozen, + name = "NativeRegistered", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeRegistered { + #[pyo3(get)] + unmatched: Py, +} + +/// A registration pass that stopped because the table it was handed is no longer the published one +#[pyclass( + frozen, + name = "NativeTableSwapped", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeTableSwapped; + +/// The inputs at one moment, for a caller that has to put them back, such as test isolation +#[pyclass( + frozen, + name = "NativeModelCostMapState", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeModelCostMapState(Inputs); + +#[pyclass( + frozen, + name = "NativeProvenance", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeProvenance { + #[pyo3(get)] + revision: String, + #[pyo3(get)] + etag: Option, +} + +#[pymethods] +impl NativeModelCostMap { + /// Starts from the bundled document. A fetched document with fewer than `min_model_count` models, or fewer than + /// `max_shrink_ratio` times the bundled count, is rejected + #[staticmethod] + fn bundled(min_model_count: i64, max_shrink_ratio: f64) -> PyResult { + let min_model_count = usize::try_from(min_model_count).map_err(|_| { + PyValueError::new_err(format!( + "min_model_count must be 0 or more, got {min_model_count}" + )) + })?; + let limits = IntegrityLimits::new(min_model_count, max_shrink_ratio).ok_or_else(|| { + PyValueError::new_err(format!( + "max_shrink_ratio must be between 0 and 1, got {max_shrink_ratio}" + )) + })?; + Ok(Self { + inputs: RwLock::new(Inputs { + cost_map: CostMap::bundled(), + registrations: Registrations::default(), + }), + limits, + index: RwLock::new(None), + }) + } + + /// The live table built from a fetched document, or why it was rejected, in which case nothing changed + #[pyo3(signature = (body, etag = None))] + fn install(&self, py: Python<'_>, body: &[u8], etag: Option) -> PyResult> { + match crate::logger::capture(py).scope(|| CostMap::read(body, etag, self.limits)) { + Ok(cost_map) => self.install_document(py, cost_map), + Err(error) => rejected(py, &error), + } + } + + /// A caller's document as a table, read without size limits and installed nowhere + #[staticmethod] + fn finalize(py: Python<'_>, body: &[u8]) -> PyResult> { + match crate::logger::capture(py) + .scope(|| CostMap::read(body, None, IntegrityLimits::UNCHECKED)) + { + Ok(cost_map) => Ok(Py::new(py, built(py, &cost_map)?)?.into_any()), + Err(error) => rejected(py, &error), + } + } + + /// The live table built from the bundled document, which is never size checked + fn install_bundled(&self, py: Python<'_>) -> PyResult> { + self.install_document(py, CostMap::bundled()) + } + + /// Writes each entry into `rows`, or into the published table when `rows` is `None` + #[pyo3(signature = (rows, entries, persist, custom_llm_provider = None))] + fn register( + &self, + py: Python<'_>, + rows: Option<&Bound<'_, PyDict>>, + entries: &Bound<'_, PyDict>, + persist: bool, + custom_llm_provider: Option<&str>, + ) -> PyResult> { + let entries: Vec<(String, Bound<'_, PyAny>)> = entries + .iter() + .map(|(key, registration)| Ok((key.str()?.to_string(), registration))) + .collect::>()?; + let kept = match persist { + true => self.merged_with_kept(&entries)?, + false => Vec::new(), + }; + let pass = self.write_pass(rows, &entries, custom_llm_provider)?; + if persist { + self.keep(kept, pass.landed()); + } + into_py(py, pass) + } + + /// Writes every kept registration into `rows` again, oldest first + fn reapply(&self, rows: &Bound<'_, PyDict>) -> PyResult> { + let py = rows.py(); + let entries: Vec<(String, Bound<'_, PyAny>)> = self + .read_inputs() + .registrations + .entries() + .map(|(key, kept)| (key.to_owned(), kept.bind(py).clone().into_any())) + .collect(); + let pass = self.write_pass(Some(rows), &entries, None)?; + self.write_inputs().registrations.record(pass.landed()); + into_py(py, pass) + } + + /// The installed row under `key` with the kept registrations that apply to it, without anything the router wrote + fn registered_row<'py>( + &self, + py: Python<'py>, + key: &str, + ) -> PyResult>> { + let (cost_map, applied) = { + let inputs = self.read_inputs(); + let applied: Vec> = inputs + .registrations + .over(key) + .into_iter() + .map(|kept| kept.bind(py).clone()) + .collect(); + (inputs.cost_map.clone(), applied) + }; + table::registered_row(py, &cost_map, key, &applied) + } + + fn registered_keys(&self) -> Vec { + let inputs = self.read_inputs(); + table::registered_keys( + &inputs.cost_map, + inputs.registrations.entries().map(|(key, _)| key), + ) + } + + /// The installed document as a table, without any registration + fn loaded<'py>(&self, py: Python<'py>) -> PyResult> { + let cost_map = self.read_inputs().cost_map.clone(); + table::build(py, &cost_map) + } + + /// The key itself when `rows` has it, otherwise the key that differs from `name` only in case + fn find_key(&self, rows: &Bound<'_, PyDict>, name: &str) -> Option { + self.index_over(rows).find_key(name) + } + + /// Keys whose row's `litellm_provider` is exactly `provider`, in row order + fn keys_in(&self, rows: &Bound<'_, PyDict>, provider: &str) -> Vec { + self.index_over(rows).keys_in(provider) + } + + /// Drops the name index, for a caller that changed keys in a way the dict's length does not show + fn invalidate(&self) { + if let Some(index) = self + .index + .read() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + { + index.invalidate(); + } + } + + fn snapshot(&self, py: Python<'_>) -> NativeModelCostMapState { + NativeModelCostMapState(self.read_inputs().clone_ref(py)) + } + + /// Puts back the document, registrations and provenance a `snapshot` saw + fn restore(&self, py: Python<'_>, state: PyRef<'_, NativeModelCostMapState>) { + *self.write_inputs() = state.0.clone_ref(py); + self.invalidate(); + } + + #[getter] + fn provenance(&self) -> NativeProvenance { + let Provenance { revision, etag } = self.read_inputs().cost_map.provenance().clone(); + NativeProvenance { revision, etag } + } +} + +impl NativeModelCostMap { + fn read_inputs(&self) -> RwLockReadGuard<'_, Inputs> { + self.inputs.read().unwrap_or_else(PoisonError::into_inner) + } + + fn write_inputs(&self) -> RwLockWriteGuard<'_, Inputs> { + self.inputs.write().unwrap_or_else(PoisonError::into_inner) + } + + fn install_document(&self, py: Python<'_>, cost_map: CostMap) -> PyResult> { + let installed = built(py, &cost_map)?; + self.write_inputs().cost_map = cost_map; + *self.index.write().unwrap_or_else(PoisonError::into_inner) = + Some(Arc::new(Store::new(LiveDict::new(installed.rows.bind(py))))); + Ok(Py::new(py, installed)?.into_any()) + } + + /// Each entry merged into the registration kept for its key + fn merged_with_kept<'py>( + &self, + entries: &[(String, Bound<'py, PyAny>)], + ) -> PyResult)>> { + entries + .iter() + .map(|(key, registration)| { + let previous = self + .read_inputs() + .registrations + .get(key) + .map(|kept| kept.bind(registration.py()).clone()); + Ok((key.clone(), host::kept(previous.as_ref(), registration)?)) + }) + .collect() + } + + fn keep(&self, kept: Vec<(String, Bound<'_, PyDict>)>, landed: &[Landed]) { + let mut inputs = self.write_inputs(); + for (key, registration) in kept { + inputs.registrations.register(&key, registration.unbind()); + } + inputs.registrations.record(landed); + } + + fn write_pass<'a, 'py>( + &self, + rows: Option<&Bound<'py, PyDict>>, + entries: &'a [(String, Bound<'py, PyAny>)], + custom_llm_provider: Option<&str>, + ) -> PyResult>> { + let mut table = PublishedTable { + map: self, + rows, + custom_llm_provider, + }; + write_pass(&mut table, entries) + } + + /// The index over `rows`, rebuilt when `rows` is not the dict it was built over + fn index_over(&self, rows: &Bound<'_, PyDict>) -> Arc> { + let current = self + .index + .read() + .unwrap_or_else(PoisonError::into_inner) + .clone(); + if let Some(index) = current.filter(|index| index.rows().is(rows)) { + return index; + } + let index = Arc::new(Store::new(LiveDict::new(rows))); + *self.index.write().unwrap_or_else(PoisonError::into_inner) = Some(Arc::clone(&index)); + index + } +} + +/// The dict Python publishes as `litellm.model_cost`, written through the host +struct PublishedTable<'a, 'py> { + map: &'a NativeModelCostMap, + rows: Option<&'a Bound<'py, PyDict>>, + custom_llm_provider: Option<&'a str>, +} + +impl<'py> RegistrationTable for PublishedTable<'_, 'py> { + type Registration = Bound<'py, PyAny>; + type Error = PyErr; + + fn place(&mut self, key: &str, registration: &Self::Registration) -> PyResult> { + let Some(landing) = host::resolve(self.rows, key, registration, self.custom_llm_provider)? + else { + return Ok(None); + }; + let rows = match self.rows { + Some(rows) => rows.clone(), + None => host::published(registration.py())?, + }; + table::write(&rows, &landing, registration)?; + self.map.invalidate(); + Ok(Some(Placed { + key: landing.key, + matched_builtin: landing.matched_builtin, + })) + } +} + +fn built(py: Python<'_>, cost_map: &CostMap) -> PyResult { + let rows = table::build(py, cost_map)?; + let fallback_generalizations = match cost_map.fallback_generalizations() { + Some(block) => to_py(py, block)?, + None => py.None(), + }; + Ok(NativeInstalled { + model_count: rows.len(), + rows: rows.unbind(), + fallback_generalizations, + }) +} + +fn rejected(py: Python<'_>, error: &litellm_model_catalog::Error) -> PyResult> { + let rejected = NativeRejected { + reason: litellm_error::Report::new(error).to_string(), + integrity: error.is_integrity(), + }; + Ok(Py::new(py, rejected)?.into_any()) +} + +fn into_py<'py>(py: Python<'py>, pass: Pass<'_, Bound<'py, PyAny>>) -> PyResult> { + match pass { + Pass::Swapped { .. } => Ok(Py::new(py, NativeTableSwapped)?.into_any()), + Pass::Finished { unmatched, .. } => { + let table = PyDict::new(py); + for (key, registration) in unmatched { + table.set_item(key, registration)?; + } + let registered = NativeRegistered { + unmatched: table.unbind(), + }; + Ok(Py::new(py, registered)?.into_any()) + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::mpsc; + use std::time::Duration; + + use pyo3::ffi::c_str; + use rstest::rstest; + + use super::*; + + const HOST: &std::ffi::CStr = c_str!( + r#" +from types import SimpleNamespace + +class Row(dict): + names = () + +found = [] + +def resolve(rows, key, registration, custom_llm_provider): + found.append(cost_map.find_key(rows, key)) + return SimpleNamespace(key=key, base={}, matched_builtin=True) + +def merged(current, base, registration): + return Row({**(current or {}), **base, **registration}) + +def write(rows, written): + rows.update(written) +"# + ); + + #[rstest] + fn a_host_lookup_that_reaches_find_key_on_the_same_map_finishes() { + Python::initialize(); + let (sender, receiver) = mpsc::channel(); + std::thread::spawn(move || { + let outcome = Python::attach(|py| -> PyResult<(Vec>, f64)> { + let cost_map = Py::new(py, NativeModelCostMap::bundled(0, 0.0)?)?; + let module_name = std::ffi::CString::new(host::MODULE).unwrap(); + let host = + PyModule::from_code(py, HOST, c_str!("host_model_cost.py"), &module_name)?; + host.setattr("cost_map", &cost_map)?; + let modules = py.import("sys")?.getattr("modules")?; + let packages = ["litellm", "litellm.rust_bridge", "litellm.rust_bridge.host"]; + for package in packages { + modules.set_item(package, PyModule::new(py, package)?)?; + } + modules.set_item(host::MODULE, &host)?; + let rows = PyDict::new(py); + rows.set_item("Lit-Reentry", PyDict::new(py))?; + let entry = PyDict::new(py); + entry.set_item("input_cost_per_token", 1.5)?; + let entries = PyDict::new(py); + entries.set_item("lit-reentry", entry)?; + let registered = cost_map + .bind(py) + .call_method1("register", (&rows, &entries, false)); + for name in [host::MODULE].into_iter().chain(packages) { + modules.del_item(name)?; + } + registered?; + let price = rows + .get_item("lit-reentry")? + .expect("the registration landed") + .get_item("input_cost_per_token")? + .extract()?; + Ok((host.getattr("found")?.extract()?, price)) + }); + sender + .send(outcome.map_err(|error| error.to_string())) + .unwrap(); + }); + + let (found, price) = receiver + .recv_timeout(Duration::from_secs(30)) + .expect("register finished instead of deadlocking on its own lock") + .unwrap(); + assert_eq!(found, [Some("Lit-Reentry".to_owned())]); + assert!((price - 1.5).abs() < f64::EPSILON); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/table.rs b/litellm-rust/crates/python-bridge/src/model_cost/table.rs new file mode 100644 index 00000000000..67b5f581b49 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/table.rs @@ -0,0 +1,166 @@ +//! Builds `litellm.model_cost` from an installed document + +use std::borrow::Cow; +use std::collections::HashMap; + +use indexmap::IndexSet; +use litellm_host_python::to_py; +use litellm_model_catalog::wire::Fields; +use litellm_model_catalog::{CostMap, FALLBACK_GENERALIZATIONS, Written}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::Value; + +use super::host::{self, Landing}; + +const ALIASES: &str = "aliases"; + +enum Served<'a> { + Model, + AliasOf(&'a str), + /// A top-level value that is not a model row, such as `sample_spec` + Section(&'a Value), +} + +/// Every served key in dict order. An `alias_of` row repeats among the aliases and keeps its document position +fn served(cost_map: &CostMap) -> impl Iterator)> { + let catalog = cost_map.catalog(); + let documented = catalog + .written() + .filter(|(key, _)| *key != FALLBACK_GENERALIZATIONS) + .filter_map(|(key, written)| { + let served = match written { + Written::Section(value) => Served::Section(value), + Written::Entry(_) => match catalog.aliases().get(key) { + Some(owner) => Served::AliasOf(owner), + None => catalog.get(key).map(|_| Served::Model)?, + }, + }; + Some((key, served)) + }); + let listed = cost_map.alias_groups().flat_map(|(owner, names)| { + names + .iter() + .skip(1) + .map(move |alias| (alias.as_str(), Served::AliasOf(owner))) + }); + documented.chain(listed) +} + +/// The installed document as a new dict +pub(crate) fn build<'py>(py: Python<'py>, cost_map: &CostMap) -> PyResult> { + let table = PyDict::new(py); + let mut rows = HashMap::new(); + for (key, served) in served(cost_map) { + let value = match served { + Served::Model => row(py, cost_map, key, &mut rows)?, + Served::AliasOf(owner) => row(py, cost_map, owner, &mut rows)?, + Served::Section(value) => to_py(py, value)?.into_bound(py), + }; + table.set_item(key, value)?; + } + Ok(table) +} + +fn row<'a, 'py>( + py: Python<'py>, + cost_map: &'a CostMap, + owner: &'a str, + rows: &mut HashMap<&'a str, Bound<'py, PyAny>>, +) -> PyResult> { + if let Some(row) = rows.get(owner) { + return Ok(row.clone()); + } + let fields = cost_map + .catalog() + .get(owner) + .expect("a served model or alias owner is a catalog row"); + let fields = to_py(py, without_alias_list(fields).as_ref())?.into_bound(py); + let row = match cost_map.names(owner) { + [] => fields, + names => host::entry(fields, names)?, + }; + rows.insert(owner, row.clone()); + Ok(row) +} + +fn without_alias_list(fields: &Fields) -> Cow<'_, Fields> { + if fields.get(ALIASES).is_none_or(Value::is_null) { + return Cow::Borrowed(fields); + } + Cow::Owned( + fields + .iter() + .filter(|(name, _)| *name != ALIASES) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + ) +} + +/// The installed row under `key` with `applied` written over it in order, `None` when neither has one +pub(crate) fn registered_row<'py>( + py: Python<'py>, + cost_map: &CostMap, + key: &str, + applied: &[Bound<'py, PyDict>], +) -> PyResult>> { + let catalog = cost_map.catalog(); + let installed = catalog + .get(key) + .or_else(|| match catalog.section(key) { + Some(Value::Object(fields)) if key != FALLBACK_GENERALIZATIONS => Some(fields), + _ => None, + }) + .map(without_alias_list); + let row = match (installed, applied) { + (None, []) => return Ok(None), + (Some(fields), _) => to_py(py, fields.as_ref())? + .into_bound(py) + .cast_into::()?, + (None, _) => PyDict::new(py), + }; + for registration in applied { + row.update(registration.as_mapping())?; + } + Ok(Some(row)) +} + +/// The keys `registered_row` answers for, installed keys first +pub(crate) fn registered_keys<'a>( + cost_map: &'a CostMap, + registered: impl Iterator, +) -> Vec { + let served: IndexSet<&str> = served(cost_map).map(|(key, _)| key).collect(); + let only_registered: Vec<&str> = registered.filter(|key| !served.contains(key)).collect(); + served + .into_iter() + .chain(only_registered) + .map(str::to_owned) + .collect() +} + +/// Writes `registration` over the row under `landing.key`, and under every name that shares that row object +pub(crate) fn write<'py>( + rows: &Bound<'py, PyDict>, + landing: &Landing<'py>, + registration: &Bound<'py, PyAny>, +) -> PyResult<()> { + let current = rows.get_item(&landing.key)?; + let row = host::merged(current.as_ref(), landing, registration)?; + let names: Vec = row.getattr("names")?.extract()?; + let sharing: Vec<&String> = names + .iter() + .filter(|name| { + let served = rows.get_item(name.as_str()).ok().flatten(); + served + .zip(current.as_ref()) + .is_some_and(|(served, current)| served.is(current)) + }) + .collect(); + let written = PyDict::new(rows.py()); + written.set_item(&landing.key, &row)?; + sharing + .into_iter() + .try_for_each(|name| written.set_item(name, &row))?; + host::write(rows, &written) +} diff --git a/litellm-rust/crates/python-bridge/src/preflight.rs b/litellm-rust/crates/python-bridge/src/preflight.rs index bc692498880..2a07723b94c 100644 --- a/litellm-rust/crates/python-bridge/src/preflight.rs +++ b/litellm-rust/crates/python-bridge/src/preflight.rs @@ -7,7 +7,7 @@ use pyo3::{ }; use strum::{IntoStaticStr, VariantArray}; -const MODULE: &str = "litellm.rust_bridge.preflight"; +const MODULE: &str = "litellm.rust_bridge.host.preflight"; /// The litellm globals the preflight still reads through Python. `preflight_contract.json` /// pins each function's parameters on both sides. @@ -115,7 +115,7 @@ mod tests { /// tests that install it run one at a time. static PREFLIGHT_MODULE: Mutex<()> = Mutex::new(()); - /// A fresh stand-in for `litellm.rust_bridge.preflight` that records every call, then + /// A fresh stand-in for `litellm.rust_bridge.host.preflight` that records every call, then /// `script` run against it with the module bound as `preflight`. fn preflight_stubs<'py>(py: Python<'py>, script: &std::ffi::CStr) -> Bound<'py, PyDict> { let locals = PyDict::new(py); @@ -124,15 +124,15 @@ mod tests { import sys import types -for name in ('litellm', 'litellm.rust_bridge'): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'): sys.modules.setdefault(name, types.ModuleType(name)) -preflight = types.ModuleType('litellm.rust_bridge.preflight') +preflight = types.ModuleType('litellm.rust_bridge.host.preflight') preflight.warnings = [] preflight.checked = [] preflight.credential_list = lambda: [] preflight.warn_unknown_credential = lambda name, loaded: preflight.warnings.append((name, loaded)) preflight.check_limits = lambda kwargs: preflight.checked.append(kwargs) -sys.modules['litellm.rust_bridge.preflight'] = preflight +sys.modules['litellm.rust_bridge.host.preflight'] = preflight ", Some(&locals), Some(&locals), diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index eacec6859b5..1cb1813049c 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -13,6 +13,8 @@ pub(crate) enum PythonSettings { UrlPolicy, #[strum(serialize = "provider_defaults")] ProviderDefaults, + #[strum(serialize = "router_tiers")] + RouterTiers, #[strum(serialize = "secret_manager")] SecretManager, #[strum(serialize = "secret_manager_binding")] @@ -24,10 +26,15 @@ pub(crate) struct Snapshot<'py> { value: Bound<'py, PyAny>, } -impl Snapshot<'_> { +impl<'py> Snapshot<'py> { pub(crate) fn read(&self, spec: &FieldSpec) -> Result { spec.read(&self.value, self.group.into()) } + + pub(crate) fn attribute(&self, name: &str) -> PyResult>> { + let value = self.value.getattr(name)?; + Ok((!value.is_none()).then_some(value)) + } } impl PythonSettings { diff --git a/litellm-rust/crates/python-bridge/src/redis/breaker.rs b/litellm-rust/crates/python-bridge/src/redis/breaker.rs new file mode 100644 index 00000000000..5fdac76cb02 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/breaker.rs @@ -0,0 +1,134 @@ +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::{Duration, Instant}, +}; + +use litellm_circuit_breaker::{Admission, Breaker, Config, Failure}; +use litellm_redis_client::RedisEvents; +use litellm_redis_types::Endpoint; +use pyo3::{exceptions::PyValueError, prelude::*}; + +use crate::redis::{ + client::{RedisSpec, resolved_endpoint}, + events::{ServiceLoggingEvents, state_label}, + hosts, +}; + +/// `litellm.redis.store_settings.RedisBreakerSettings`, filled from `litellm.constants`. +#[derive(FromPyObject)] +pub(crate) struct BreakerSettings { + failure_threshold: u32, + recovery_timeout: f64, + timeout_min_duration: f64, +} + +impl BreakerSettings { + pub(crate) fn config(&self) -> PyResult { + let seconds = |seconds: f64| { + Duration::try_from_secs_f64(seconds) + .map_err(|error| PyValueError::new_err(error.to_string())) + }; + Ok(Config { + failure_threshold: self.failure_threshold, + recovery_timeout: seconds(self.recovery_timeout)?, + timeout_min_duration: seconds(self.timeout_min_duration)?, + }) + } +} + +/// The one health signal for a Redis, shared by its `RedisCache` and its native client, so a Redis the +/// cache found wedged is fast-failed by every store too. +#[pyclass(frozen, name = "NativeRedisBreaker")] +pub(crate) struct NativeRedisBreaker(pub(crate) Breaker); + +impl NativeRedisBreaker { + fn build( + py: Python<'_>, + settings: &BreakerSettings, + endpoint: Option, + ) -> PyResult { + let config = settings.config()?; + let runtime = litellm_host_python::runtime_handle()?; + let events: Arc = + Arc::new(ServiceLoggingEvents::transitions_only(py, runtime)?); + Ok(Self(py.detach(|| { + hosts().breaker(endpoint, config, Some(events)) + }))) + } +} + +#[pymethods] +impl NativeRedisBreaker { + #[new] + fn new(py: Python<'_>, settings: BreakerSettings) -> PyResult { + Self::build(py, &settings, None) + } + + /// The breaker of the Redis server `redis_kwargs` names, which every cache and client on that server shares + /// whatever its db, credentials or namespace. The first caller's settings win. + #[staticmethod] + fn shared( + py: Python<'_>, + redis_kwargs: RedisSpec, + settings: BreakerSettings, + ) -> PyResult { + let endpoint = resolved_endpoint(&redis_kwargs); + Self::build(py, &settings, endpoint) + } + + fn admit(&self) -> Option { + self.0 + .admit(Instant::now()) + .ok() + .map(|admission| NativeAdmission(Mutex::new(Some(admission)))) + } + + fn failed(&self, timeout: bool) { + self.0.failed(failure(timeout), Instant::now()); + } + + #[getter] + fn state(&self) -> &'static str { + state_label(&self.0.state()) + } +} + +/// One admitted call. Settling it twice records nothing more, and dropping an unsettled probe reopens the breaker. +#[pyclass(frozen, name = "NativeAdmission")] +pub(crate) struct NativeAdmission(Mutex>); + +impl NativeAdmission { + fn take(&self) -> Option { + self.0.lock().unwrap_or_else(PoisonError::into_inner).take() + } +} + +#[pymethods] +impl NativeAdmission { + fn succeeded(&self) { + if let Some(admission) = self.take() { + admission.succeeded(); + } + } + + fn failed(&self, timeout: bool) { + if let Some(admission) = self.take() { + admission.failed(failure(timeout), Instant::now()); + } + } +} + +fn failure(timeout: bool) -> Failure { + match timeout { + true => Failure::Timeout, + false => Failure::Hard, + } +} + +#[pyfunction] +pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> { + py.detach(|| hosts().breakers()) + .iter() + .map(|breaker| state_label(&breaker.state())) + .collect() +} diff --git a/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs b/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs new file mode 100644 index 00000000000..6995981a03e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs @@ -0,0 +1,47 @@ +use std::sync::Arc; + +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_redis_connect::{CredentialSource, CredentialStage, Error, Expiring, SourceFuture}; +use litellm_secret_value::SecretValue; +use pyo3::prelude::*; + +const CREDENTIALS_MODULE: &str = "litellm.rust_bridge.host.redis_credentials"; + +/// A redis-py `CredentialProvider` the SDK caller passed in `cache_kwargs`, asked through the host +/// module the way redis-py's async connection asks it. +pub(crate) struct PythonCredentialSource { + provider: Arc>, + context: PythonContext, +} + +impl PythonCredentialSource { + pub(crate) fn new(provider: Py, context: PythonContext) -> Self { + Self { + provider: Arc::new(provider), + context, + } + } +} + +impl CredentialSource for PythonCredentialSource { + fn fetch(&self) -> SourceFuture<'_> { + let provider = Arc::clone(&self.provider); + let context = self.context.clone(); + Box::pin(async move { + let (username, password) = attach_blocking(context, move |py| { + py.import(CREDENTIALS_MODULE)? + .getattr("credentials")? + .call1((provider.bind(py),))? + .extract::<(Option, String)>() + }) + .await + .and_then(|credentials| credentials) + .map_err(Error::credential(CredentialStage::Provider))?; + Ok(Expiring { + username, + password: SecretValue::new(password), + expires_at: None, + }) + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs new file mode 100644 index 00000000000..5e5ab7e6f5a --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs @@ -0,0 +1,294 @@ +mod credentials; + +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use litellm_circuit_breaker::Config as BreakerConfig; +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; +use litellm_host_python::PythonContext; +use litellm_python_compat::number::{self, SecondsError}; +use litellm_redis_client::{RedisClient, RedisEvents}; +use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment}; +use litellm_redis_resolve::{Kwargs, PythonException}; +use litellm_redis_types::{Endpoint, ProviderHandle, Resolution}; +use pyo3::{ + exceptions::{PyOverflowError, PyTypeError, PyValueError}, + prelude::*, + types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple}, +}; + +use self::credentials::PythonCredentialSource; +use crate::redis::{breaker::BreakerSettings, events::ServiceLoggingEvents, hosts}; + +pyo3::import_exception!(redis.exceptions, RedisClusterException); + +/// The kwargs `_merge_redis_kwargs` returns for a `RedisCache`, which the resolver reads the way +/// redis-py's constructors do. A key whose value is a Python object, such as `retry`, reaches the +/// resolver by name only, so it can name the setting it cannot honor. A `credential_provider` object +/// is also kept, so the connector can ask it for credentials, and reaches the resolver as a handle. +pub(crate) struct RedisSpec { + kwargs: Kwargs, + credential_provider: Option>, +} + +impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec { + type Error = PyErr; + + fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult { + let py = kwargs.py(); + let source = PyDict::new(py); + let mapping = kwargs.cast::()?; + source.update(&mapping)?; + let plain = PyDict::new(py); + let mut host_objects = Vec::new(); + for (name, value) in source.iter() { + match value.is_none() || is_plain(&value) { + true => plain.set_item(&name, value)?, + false => host_objects.push(name.extract::()?), + } + } + let credential_provider = source + .get_item("credential_provider")? + .filter(|provider| !provider.is_none() && !is_plain(provider)) + .map(Bound::unbind); + let fields: BTreeMap = + litellm_host_python::from_py_argument(plain.as_any())?; + let objects = host_objects + .into_iter() + .map(|name| (name, serde_json::Value::Bool(true))); + Ok(Self { + kwargs: litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)), + credential_provider, + }) + } +} + +/// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built. +#[derive(FromPyObject)] +pub(crate) struct StoreSettings { + redis_kwargs: RedisSpec, + namespace: Option, + default_ttl: f64, + breaker: Option, +} + +fn is_plain(value: &Bound<'_, PyAny>) -> bool { + value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() +} + +/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool +/// and its breaker. A config redis-py refuses building its client raises the same class here. One it only fails on +/// once a command connects, where `RedisCache` logs every failed call, leaves the family on memory with an error. +fn shared_client( + py: Python<'_>, + spec: RedisSpec, + namespace: Option, + breaker: Option, +) -> PyResult> { + let environment: SharedEnvironment = Arc::new(ProcessEnvironment); + let logger = crate::logger::capture(py); + let resolved = match logger.scope(|| resolve_spec(&spec)) { + Ok(resolution) => { + logger.scope(|| { + for ignored in &resolution.ignored { + litellm_tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored: {}", + ignored.setting, + ignored.consequence + ); + } + }); + Resolution { + namespace, + ..resolution + } + } + Err(error) => return unresolved(py, &error), + }; + let runtime = litellm_host_python::runtime_handle()?; + let events: Arc = Arc::new(ServiceLoggingEvents::new(py, runtime.clone())?); + let host = HostServices { + events: Some(events), + http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?, + environment, + credential_provider: spec + .credential_provider + .as_ref() + .map(|provider| -> PyResult> { + Ok(Arc::new(PythonCredentialSource::new( + provider.clone_ref(py), + PythonContext::capture(py)?, + ))) + }) + .transpose()?, + }; + let connected = py.detach(|| { + let _runtime = runtime.enter(); + hosts().client(resolved, host, breaker) + }); + match connected { + Ok(client) => Ok(Some(client)), + Err(error) => Ok(unconnected(py, &error)), + } +} + +/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store +/// never sees the redis-py kwargs or the namespace. +#[pyclass(frozen, name = "NativeRedisClient")] +pub(crate) struct NativeRedisClient { + pub(crate) client: RedisClient, + /// `RedisCache.default_ttl`, the TTL a store's write takes when its caller gives none. + pub(crate) default_ttl: Duration, +} + +#[pymethods] +impl NativeRedisClient { + /// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory. + #[staticmethod] + fn connect(py: Python<'_>, settings: StoreSettings) -> PyResult> { + let StoreSettings { + redis_kwargs, + namespace, + default_ttl, + breaker, + } = settings; + let default_ttl = number::seconds(default_ttl).map_err(seconds_error)?; + let breaker = breaker.as_ref().map(BreakerSettings::config).transpose()?; + Ok( + shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self { + client, + default_ttl, + }), + ) + } + + #[getter] + fn default_ttl_seconds(&self) -> f64 { + self.default_ttl.as_secs_f64() + } +} + +/// The Redis server `spec` names, whatever its db, credentials or namespace, or none when the Rust client cannot reach it. +pub(crate) fn resolved_endpoint(spec: &RedisSpec) -> Option { + resolve_spec(spec) + .ok() + .and_then(|resolution| resolution.target.endpoint()) +} + +// A live registry entry's client keeps its provider alive, so no live entry shares its address +fn resolve_spec(spec: &RedisSpec) -> Result { + let handle = spec + .credential_provider + .as_ref() + .map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64)); + litellm_redis_resolve::interpret( + &spec.kwargs, + handle, + ProcessEnvironment.get("REDIS_USERNAME").as_deref(), + ) +} + +fn unresolved( + py: Python<'_>, + error: &litellm_redis_resolve::Error, +) -> PyResult> { + let message = || litellm_error::Report::new(error).to_string(); + match error.raised_building_the_python_client() { + Some(PythonException::ValueError) => Err(PyValueError::new_err(message())), + Some(PythonException::TypeError) => Err(PyTypeError::new_err(message())), + Some(PythonException::RedisClusterException) => { + Err(RedisClusterException::new_err(message())) + } + None => Ok(memory_only(py, error)), + } +} + +// redis-py reads TLS files and fetches credentials only once a command connects, so none of these fails building +// its client +fn unconnected(py: Python<'_>, error: &litellm_redis_connect::Error) -> Option { + match error { + litellm_redis_connect::Error::Connect(_) + | litellm_redis_connect::Error::NoHostCredentialProvider + | litellm_redis_connect::Error::ProcessEnvironmentRequired { .. } + | litellm_redis_connect::Error::Credential { .. } + | litellm_redis_connect::Error::CredentialTimeout { .. } + | litellm_redis_connect::Error::ExpiredCredential + | litellm_redis_connect::Error::MissingCredential { .. } => memory_only(py, error), + } +} + +fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option { + let report = litellm_error::Report::new(error).to_string(); + crate::logger::capture(py).scope(|| { + litellm_tracing::error!( + "Redis is not usable from the Rust stores, their state stays in memory: {report}" + ) + }); + None +} + +fn seconds_error(error: SecondsError) -> PyErr { + match error { + SecondsError::Overflow => PyOverflowError::new_err(error.to_string()), + SecondsError::NotANumber => PyValueError::new_err(error.to_string()), + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn kwargs(fields: serde_json::Value) -> Kwargs { + let fields: BTreeMap = serde_json::from_value(fields).unwrap(); + litellm_redis_resolve::kwargs(fields) + } + + #[rstest] + #[case::no_target(json!({"port": 6379}), Some("ValueError"))] + #[case::unenforceable(json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), Some("TypeError"))] + #[case::fails_only_once_a_command_connects( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": "None"}), + None + )] + fn a_config_raises_where_redis_py_raises_building_its_client( + #[case] fields: serde_json::Value, + #[case] raised: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let error = litellm_redis_resolve::interpret(&kwargs(fields), None, None).unwrap_err(); + + let outcome = unresolved(py, &error) + .map(|client| client.is_some()) + .map_err(|error| error.get_type(py).name().unwrap().to_string()); + + assert_eq!( + outcome, + raised.map_or(Ok(false), |class| Err(class.to_owned())) + ); + }); + } + + #[rstest] + #[case::overflow(f64::INFINITY, "OverflowError")] + #[case::not_a_number(f64::NAN, "ValueError")] + fn a_default_ttl_python_cannot_make_an_int_raises_like_int( + #[case] seconds: f64, + #[case] raised: &str, + ) { + Python::initialize(); + Python::attach(|py| { + let error = number::seconds(seconds).map_err(seconds_error).unwrap_err(); + + assert_eq!(error.get_type(py).name().unwrap().to_string(), raised); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/events.rs b/litellm-rust/crates/python-bridge/src/redis/events.rs new file mode 100644 index 00000000000..a281d78d5a6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/events.rs @@ -0,0 +1,403 @@ +use std::{ + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_redis_client::{ + CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged, +}; +use pyo3::{prelude::*, types::PyModule}; +use tokio::runtime::Handle; + +/// The asyncio loop a Python caller ran on. `RedisCache` scheduled its service events there with +/// `asyncio.create_task`, so the native events go there too instead of a new loop per event. +#[derive(Clone)] +pub(crate) struct EventLoop(Arc>); + +impl EventLoop { + pub(crate) fn running(py: Python<'_>) -> Option { + pyo3_async_runtimes::get_running_loop(py) + .ok() + .map(|event_loop| Self(Arc::new(event_loop.unbind()))) + } +} + +/// The Python caller of one store operation. Its context carries `with_service_target` and the +/// ambient OTel span, and its `call_type` is the label `RedisCache` reports, so dashboards keep +/// their series. +pub(crate) struct RedisCall { + pub(crate) context: PythonContext, + pub(crate) event_loop: Option, + pub(crate) call_type: &'static str, +} + +impl RedisCall { + pub(crate) fn tag(self) -> CallTag { + Arc::new(self) + } + + pub(crate) async fn scope(self, operation: F) -> F::Output { + tagged(self.tag(), operation).await + } +} + +/// Reports every native Redis call and breaker transition to `ServiceLogging`. +pub(crate) struct ServiceLoggingEvents { + record_call: Arc>, + record_breaker: Arc>, + context: PythonContext, + /// The loop the client was built on, for calls whose caller had none running. + event_loop: Option, + runtime: Handle, +} + +impl ServiceLoggingEvents { + const MODULE: &str = "litellm.rust_bridge.host.redis_events"; + + pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult { + Self::from_module( + py, + &py.import(Self::MODULE)?, + runtime, + EventLoop::running(py), + ) + } + + /// For a breaker, whose transitions never run on the caller's loop, so building it never looks one up. + pub(crate) fn transitions_only(py: Python<'_>, runtime: Handle) -> PyResult { + Self::from_module(py, &py.import(Self::MODULE)?, runtime, None) + } + + fn from_module( + py: Python<'_>, + module: &Bound<'_, PyModule>, + runtime: Handle, + event_loop: Option, + ) -> PyResult { + Ok(Self { + record_call: Arc::new(module.getattr("record_redis_call")?.unbind()), + record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()), + context: PythonContext::capture(py)?, + event_loop, + runtime, + }) + } + + fn report(&self, context: PythonContext, record: Arc>, args: ReportArgs) { + self.runtime.spawn(async move { + let reported = attach_blocking(context, move |py| match args { + ReportArgs::Call(call) => record + .bind(py) + .call1(( + call.event_loop + .as_ref() + .map(|event_loop| event_loop.0.bind(py)), + call.call_type, + call.start_time, + call.end_time, + call.failure, + )) + .map(drop), + ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop), + }) + .await + .and_then(|reported| reported); + if let Err(error) = reported { + tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed"); + } + }); + } +} + +/// One call as `record_redis_call` takes it, its times in `time.time()` seconds. +struct CallReport { + event_loop: Option, + call_type: &'static str, + start_time: f64, + end_time: f64, + failure: Option<&'static str>, +} + +enum ReportArgs { + Call(CallReport), + Breaker(&'static str), +} + +/// The wall-clock start and end of a call that just finished after `duration`. +fn wall_clock_span(duration: Duration) -> (f64, f64) { + let end = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default(); + ( + end.saturating_sub(duration).as_secs_f64(), + end.as_secs_f64(), + ) +} + +fn operation_call_type(operation: Operation) -> &'static str { + match operation { + Operation::Command(name) => name, + Operation::Script => "eval", + Operation::Pipeline { .. } => "pipeline", + } +} + +fn failure_label(outcome: Outcome) -> Option<&'static str> { + match outcome { + Outcome::Succeeded => None, + Outcome::Failed(FailureKind::Timeout) => Some("timeout"), + Outcome::Failed(FailureKind::Transport) => Some("transport"), + Outcome::Failed(FailureKind::Server) => Some("server"), + Outcome::Failed(FailureKind::Decode) => Some("decode"), + Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"), + } +} + +pub(crate) fn state_label(state: &State) -> &'static str { + match state { + State::Closed { .. } => "closed", + State::Open { .. } => "open", + State::HalfOpen => "half_open", + } +} + +impl RedisEvents for ServiceLoggingEvents { + fn call(&self, event: &CallEvent) { + let scoped = current_tag().and_then(|tag| tag.downcast::().ok()); + let (context, event_loop, call_type) = match scoped { + Some(call) => ( + call.context.clone(), + call.event_loop.clone().or_else(|| self.event_loop.clone()), + call.call_type, + ), + None => ( + self.context.clone(), + self.event_loop.clone(), + operation_call_type(event.operation), + ), + }; + let (start_time, end_time) = wall_clock_span(event.duration); + self.report( + context, + self.record_call.clone(), + ReportArgs::Call(CallReport { + event_loop, + call_type, + start_time, + end_time, + failure: failure_label(event.outcome), + }), + ); + } + + fn breaker(&self, transition: &Transition) { + self.report( + self.context.clone(), + self.record_breaker.clone(), + ReportArgs::Breaker(state_label(&transition.to)), + ); + } +} + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use litellm_host_python::PythonContext; + use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; + use pyo3::{ + prelude::*, + types::{PyCFunction, PyModule, PyString}, + }; + use rstest::rstest; + use tokio::runtime::Handle; + + use super::{EventLoop, RedisCall, ServiceLoggingEvents}; + + const RECORDER: &str = " +import contextvars, time +target = contextvars.ContextVar('target', default=None) +seen = [] +def record_redis_call(event_loop, call_type, start_time, end_time, failure): + just_now = abs(time.time() - end_time) < 60 + seen.append((event_loop, call_type, round(end_time - start_time, 3), just_now, failure, target.get())) +def record_redis_breaker(state): + seen.append(('breaker', state)) +"; + + fn stand_in_loop(py: Python<'_>, name: &str) -> EventLoop { + EventLoop(Arc::new(PyString::new(py, name).into_any().unbind())) + } + + fn recorder(name: &str, client_loop: Option<&str>) -> (ServiceLoggingEvents, Py) { + Python::initialize(); + Python::attach(|py| { + let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap(); + let module = + PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name) + .unwrap(); + let events = ServiceLoggingEvents::from_module( + py, + &module, + Handle::current(), + client_loop.map(|client_loop| stand_in_loop(py, client_loop)), + ) + .unwrap(); + (events, module.unbind()) + }) + } + + async fn seen(module: &Py, count: usize) -> String { + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + loop { + let seen: (usize, String) = Python::attach(|py| { + let seen = module.bind(py).getattr("seen").unwrap(); + (seen.len().unwrap(), seen.repr().unwrap().to_string()) + }); + if seen.0 >= count || tokio::time::Instant::now() > deadline { + return seen.1; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + } + + fn scoped_call(namespace: &Py, caller_loop: Option<&str>) -> RedisCall { + Python::attach(|py| { + namespace + .bind(py) + .getattr("target") + .unwrap() + .call_method1("set", ("spend_counter",)) + .unwrap(); + RedisCall { + context: PythonContext::capture(py).unwrap(), + event_loop: caller_loop.map(|caller_loop| stand_in_loop(py, caller_loop)), + call_type: "async_increment", + } + }) + } + + #[rstest] + #[case::on_the_callers_loop(Some("caller loop"), "caller loop")] + #[case::on_the_clients_loop_when_the_caller_had_none(None, "client loop")] + #[tokio::test(flavor = "multi_thread")] + async fn a_scoped_call_reports_the_store_labels_and_its_times_in_the_callers_context( + #[case] caller_loop: Option<&str>, + #[case] expected_loop: &str, + ) { + let (events, namespace) = recorder( + expected_loop.replace(' ', "_").as_str(), + Some("client loop"), + ); + let call = scoped_call(&namespace, caller_loop); + let event = CallEvent { + operation: Operation::Command("INCRBYFLOAT"), + duration: Duration::from_millis(250), + outcome: Outcome::Failed(FailureKind::Timeout), + }; + + call.scope(async { events.call(&event) }).await; + + assert_eq!( + seen(&namespace, 1).await, + format!( + "[('{expected_loop}', 'async_increment', 0.25, True, 'timeout', 'spend_counter')]" + ) + ); + } + + #[rstest] + #[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")] + #[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")] + #[tokio::test(flavor = "multi_thread")] + async fn breaker_transitions_report_python_state_names( + #[case] to: litellm_circuit_breaker::State, + #[case] expected: &str, + ) { + let (events, namespace) = recorder(expected, None); + + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to, + generation: 1, + }); + + assert_eq!( + seen(&namespace, 1).await, + format!("[('breaker', '{expected}')]") + ); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn a_transition_on_a_thread_without_a_runtime_is_still_reported() { + let (events, namespace) = recorder("plain_thread", None); + + std::thread::spawn(move || { + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to: litellm_circuit_breaker::State::HalfOpen, + generation: 1, + }); + }) + .join() + .unwrap(); + + assert_eq!(seen(&namespace, 1).await, "[('breaker', 'half_open')]"); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn an_unscoped_call_falls_back_to_the_command_name_and_the_clients_loop() { + let (events, namespace) = recorder("unscoped", None); + + events.call(&CallEvent { + operation: Operation::Pipeline { ops: 3 }, + duration: Duration::from_millis(5), + outcome: Outcome::Succeeded, + }); + + assert_eq!( + seen(&namespace, 1).await, + "[(None, 'pipeline', 0.005, True, None, None)]" + ); + } + + #[rstest] + fn running_is_the_loop_the_caller_runs_on_and_none_outside_one() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::from_code( + py, + c"import asyncio\nasync def probe_inside(probe):\n return probe() is asyncio.get_running_loop()\n", + c"event_loop_probe", + c"event_loop_probe", + ) + .unwrap(); + let probe = PyCFunction::new_closure(py, None, None, |_, _| { + Python::attach(|py| -> PyResult>> { + Ok(EventLoop::running(py).map(|event_loop| event_loop.0.clone_ref(py))) + }) + }) + .unwrap(); + let coroutine = module + .getattr("probe_inside") + .unwrap() + .call1((probe,)) + .unwrap(); + + let inside: bool = py + .import("asyncio") + .unwrap() + .call_method1("run", (coroutine,)) + .unwrap() + .extract() + .unwrap(); + + assert!(inside); + assert!(EventLoop::running(py).is_none()); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/mod.rs b/litellm-rust/crates/python-bridge/src/redis/mod.rs new file mode 100644 index 00000000000..808f930dd20 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/mod.rs @@ -0,0 +1,15 @@ +use std::sync::LazyLock; + +use litellm_redis_connect::RedisHosts; + +pub(crate) mod breaker; +pub(crate) mod client; +mod events; +pub(crate) mod store; + +static HOSTS: LazyLock = LazyLock::new(RedisHosts::default); + +/// The Redis servers this process talks to, shared by every `RedisCache` and native store in it. +pub(crate) fn hosts() -> &'static RedisHosts { + &HOSTS +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs new file mode 100644 index 00000000000..8d7788b136e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs @@ -0,0 +1,265 @@ +use std::sync::Arc; + +use litellm_python_compat::Value; +use litellm_redis_store::auth_cache::{AuthCache, Settings}; +use pyo3::prelude::*; + +use super::{ + NativePending, PyNumber, kind_error, pending, redis_async, redis_sync, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, + value::PyValue, +}; +use crate::redis::client::NativeRedisClient; + +fn values(values: Vec>) -> Vec> { + values.into_iter().map(|value| value.map(PyValue)).collect() +} + +/// `UserApiKeyCache` with `enable_redis_auth_cache` on. The sync methods back `DualCache`'s sync +/// API and block with the GIL released, like redis-py's sync client. +#[pyclass(frozen, name = "NativeAuthCache")] +pub(crate) struct NativeAuthCache(Arc); + +#[pymethods] +impl NativeAuthCache { + #[new] + #[pyo3(signature = (client, *, memory_entries, memory_ttl_seconds = None))] + fn new( + client: Option>, + memory_entries: usize, + memory_ttl_seconds: Option, + ) -> PyResult { + Ok(Self(Arc::new(AuthCache::new( + client.map(|client| remote(&client)), + Settings { + memory_entries, + memory_ttl: ttl(memory_ttl_seconds)?, + }, + )))) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + /// `request` answers a key its pre-call reads already found absent without another `GET`. + #[pyo3(signature = (key, request = None))] + fn get<'py>( + &self, + py: Python<'py>, + key: String, + request: Option<&NativeRequestRedis>, + ) -> PyResult> { + let store = self.0.clone(); + let request = request.map(|request| request.0.clone()); + redis_async(py, "async_get_cache", kind_error, async move { + store + .get(request.as_ref(), &key) + .await + .map(|value| value.map(PyValue)) + }) + } + + fn get_sync(&self, py: Python<'_>, key: String) -> PyResult> { + let store = self.0.clone(); + redis_sync(py, "get_cache", kind_error, async move { + store.get(None, &key).await.map(|value| value.map(PyValue)) + }) + } + + /// `update_cache_ttl` and `update_in_memory_max_size`: both partitions keep what they hold. + #[pyo3(signature = (*, memory_entries, memory_ttl_seconds = None))] + fn reconfigure(&self, memory_entries: usize, memory_ttl_seconds: Option) -> PyResult<()> { + self.0.reconfigure(Settings { + memory_entries, + memory_ttl: ttl(memory_ttl_seconds)?, + }); + Ok(()) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set<'py>( + &self, + py: Python<'py>, + key: String, + value: PyValue, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.set(&key, value.0, ttl).await + }) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set_sync( + &self, + py: Python<'_>, + key: String, + value: PyValue, + ttl_seconds: Option, + ) -> PyResult<()> { + let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; + redis_sync(py, "set_cache", kind_error, async move { + store.set_str(&key, value.0, ttl).await + }) + } + + /// Memory now and the `SET` on the batch `send` picks. + #[pyo3(signature = (key, value, ttl_seconds, send))] + fn declare_set( + &self, + key: String, + value: PyValue, + ttl_seconds: Option, + send: &NativeSend, + ) -> PyResult { + let written = self + .0 + .declare_set(send.send(), &key, value.0, ttl(ttl_seconds)?) + .map_err(kind_error)?; + Ok(pending( + "async_set_cache", + kind_error, + written, + std::convert::identity, + )) + } + + fn declare_delete(&self, key: String, send: &NativeSend) -> PyResult { + let deleted = self + .0 + .declare_delete(send.send(), &key) + .map_err(kind_error)?; + Ok(pending( + "async_delete_cache", + kind_error, + deleted, + std::convert::identity, + )) + } + + #[pyo3(signature = (entries, ttl_seconds = None))] + fn set_many<'py>( + &self, + py: Python<'py>, + entries: Vec<(String, PyValue)>, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let entries = entries + .into_iter() + .map(|(key, value)| (key, value.0)) + .collect(); + let ttl = ttl(ttl_seconds)?; + redis_async(py, "async_set_cache_pipeline", kind_error, async move { + store.set_many(entries, ttl).await + }) + } + + /// On a request batch it returns before anything is sent, and `Now` is a pipeline of its own. + fn write_back<'py>( + &self, + py: Python<'py>, + entries: Vec<(String, PyValue, Option)>, + send: &NativeSend, + ) -> PyResult> { + let entries = entries + .into_iter() + .map(|(key, value, seconds)| Ok((key, value.0, ttl(seconds)?))) + .collect::>()?; + let written = self + .0 + .write_back(send.send(), entries) + .map_err(kind_error)?; + redis_async( + py, + "async_set_cache_pipeline_with_ttls", + kind_error, + written, + ) + } + + fn delete<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_delete_cache", kind_error, async move { + store.delete(&keys).await + }) + } + + fn delete_sync(&self, py: Python<'_>, keys: Vec) -> PyResult<()> { + let store = self.0.clone(); + redis_sync(py, "delete_cache", kind_error, async move { + store.delete(&keys).await + }) + } + + /// One `MGET` of keys the caller found missing in memory, on the batch `send` picks; `None` + /// without Redis. + fn get_many(&self, keys: Vec, send: &NativeSend) -> Option { + let read = self.0.get_many(send.send(), &keys)?; + Some(pending("async_batch_get_cache", kind_error, read, values)) + } + + fn increment_if_present<'py>( + &self, + py: Python<'py>, + key: String, + amount: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_increment_if_present", kind_error, async move { + store.increment_if_present(&key, amount.0).await + }) + } + + fn get_local(&self, key: &str) -> PyResult> { + Ok(self.0.get_local(key).map_err(kind_error)?.map(PyValue)) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set_local(&self, key: &str, value: PyValue, ttl_seconds: Option) -> PyResult<()> { + self.0 + .set_local(key, value.0, ttl(ttl_seconds)?) + .map_err(kind_error) + } + + fn delete_local(&self, key: &str) -> PyResult<()> { + self.0.delete_local(key).map_err(kind_error) + } + + fn flush_local(&self) -> PyResult<()> { + self.0.flush_local().map_err(kind_error) + } + + fn len_local(&self) -> PyResult { + let (objects, key_objects) = self.0.len_local().map_err(kind_error)?; + Ok(PyLocalSizes { + objects, + key_objects, + }) + } +} + +/// Entries the management-object and the user-key partitions hold in memory. +#[pyclass(frozen, eq, get_all, name = "LocalSizes")] +#[derive(PartialEq)] +pub(crate) struct PyLocalSizes { + objects: usize, + key_objects: usize, +} + +#[pymethods] +impl PyLocalSizes { + #[new] + fn new(objects: usize, key_objects: usize) -> Self { + Self { + objects, + key_objects, + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs new file mode 100644 index 00000000000..d9c3aabb814 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -0,0 +1,271 @@ +pub(crate) mod auth_cache; +mod policy; +pub(crate) mod rate_limits; +pub(crate) mod request; +pub(crate) mod router_cooldowns; +pub(crate) mod router_usage; +pub(crate) mod sensitive_route_pins; +pub(crate) mod spend_counters; +mod value; + +use value::PyValue; + +use std::{ + pin::Pin, + sync::{Mutex, PoisonError}, + time::Duration, +}; + +use litellm_host_python::PythonContext; +use litellm_python_compat::number::{SecondsError, seconds}; +use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote}; +use pyo3::{ + IntoPyObjectExt, + exceptions::{PyOverflowError, PyRuntimeError, PyTypeError, PyValueError}, + prelude::*, + types::{PyBool, PyFloat, PyInt}, +}; + +use crate::execution::{run_async_value, run_sync_value}; +use crate::redis::{ + client::NativeRedisClient, + events::{EventLoop, RedisCall}, +}; + +/// The labels `RedisCache` reports for this operation, in the caller's context so its +/// `service_target` reaches the service event, dispatched on the loop the caller runs on. +fn call(py: Python<'_>, call_type: &'static str) -> PyResult { + Ok(RedisCall { + context: PythonContext::capture(py)?, + event_loop: EventLoop::running(py), + call_type, + }) +} + +/// Runs one Redis operation as a Python awaitable, reported to `ServiceLogging` as `call_type`. +fn redis_async<'py, T, E: 'static>( + py: Python<'py>, + call_type: &'static str, + error: fn(E) -> PyErr, + operation: impl Future> + Send + 'static, +) -> PyResult> +where + T: for<'a> IntoPyObject<'a> + Send + 'static, +{ + let scope = call(py, call_type)?; + run_async_value( + py, + async move { scope.scope(operation).await.map_err(error) }, + ) +} + +/// Runs one Redis operation with the GIL released, for the router's sync paths. +fn redis_sync( + py: Python<'_>, + call_type: &'static str, + error: fn(E) -> PyErr, + operation: impl Future> + Send + 'static, +) -> PyResult +where + T: Send + 'static, +{ + let scope = call(py, call_type)?; + run_sync_value( + py, + async move { scope.scope(operation).await.map_err(error) }, + ) +} + +/// A TTL in seconds, raising what `int(seconds)` raises for infinity and NaN. +fn duration(value: f64) -> PyResult { + seconds(value).map_err(|error| match error { + SecondsError::Overflow => PyOverflowError::new_err(error.to_string()), + SecondsError::NotANumber => PyValueError::new_err(error.to_string()), + }) +} + +fn ttl(value: Option) -> PyResult> { + value.map(duration).transpose() +} + +fn remote(client: &NativeRedisClient) -> Remote { + Remote { + client: client.client.clone(), + default_ttl: client.default_ttl, + } +} + +type Resolve = Pin>> + Send>>; + +/// A Redis call declared on a request's batch. `resolve` flushes that batch if it has not gone +/// out and answers the call, once. +#[pyclass(frozen, name = "NativePending")] +pub(crate) struct NativePending { + resolve: Mutex>, + call_type: &'static str, +} + +impl NativePending { + pub(crate) fn new( + call_type: &'static str, + error: fn(E) -> PyErr, + reply: impl Future> + Send + 'static, + ) -> Self + where + T: for<'a> IntoPyObject<'a> + Send + 'static, + { + let resolve = async move { + let reply = reply.await.map_err(error)?; + Python::attach(|py| reply.into_py_any(py)) + }; + Self { + resolve: Mutex::new(Some(Box::pin(resolve))), + call_type, + } + } +} + +#[pymethods] +impl NativePending { + fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { + let resolve = self + .resolve + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .ok_or_else(|| PyRuntimeError::new_err("this Redis call was already resolved"))?; + redis_async(py, self.call_type, std::convert::identity, resolve) + } +} + +pyo3::create_exception!( + _native, + NativeRedisStoreError, + pyo3::exceptions::PyException, + "A native Redis store operation failed. Args are (kind, message), kind spelled as litellm-redis-store's ErrorKind." +); + +/// What every native store raises, its kind telling a caller which failures it may drop. +fn kind_error(error: Error) -> PyErr { + let kind: &'static str = error.kind().into(); + NativeRedisStoreError::new_err((kind, litellm_error::Report::new(&error).to_string())) +} + +/// A store op declared on the batch the caller's `NativeSend` picked, answered as `reply` shapes it. +fn pending( + call_type: &'static str, + error: fn(Error) -> PyErr, + declared: Pending, + reply: fn(T) -> R, +) -> NativePending +where + T: Send + 'static, + R: for<'a> IntoPyObject<'a> + Send + 'static, +{ + NativePending::new(call_type, error, async move { declared.await.map(reply) }) +} + +/// A Python `int` or `float` across the bridge, keeping which one it is. +#[derive(Clone, Copy, PartialEq)] +pub(crate) struct PyNumber(Number); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + if value.is_instance_of::() && !value.is_instance_of::() { + return Ok(Self(Number::Int(value.extract()?))); + } + if value.is_instance_of::() { + return Ok(Self(Number::Float(value.extract()?))); + } + Err(PyTypeError::new_err("expected an int or a float")) + } +} + +impl<'py> IntoPyObject<'py> for PyNumber { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Ok(match self.0 { + Number::Int(value) => value.into_pyobject(py)?.into_any(), + Number::Float(value) => value.into_pyobject(py)?.into_any(), + }) + } +} + +/// A store enum read from its wire spelling, so the bridge names no variant. +pub(crate) struct Spelled(pub(crate) T); + +impl<'a, 'py, T: std::str::FromStr> FromPyObject<'a, 'py> for Spelled { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let spelling: String = value.extract()?; + spelling + .parse() + .map(Self) + .map_err(|_| PyValueError::new_err(format!("unknown spelling {spelling:?}"))) + } +} + +/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue` +/// so the bridge never names a field. A shape the store refuses is `ValueError`. +pub(crate) struct PyData(pub(crate) T); + +impl<'a, 'py, T: FromValue> FromPyObject<'a, 'py> for PyData { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let PyValue(value) = value.extract()?; + T::from_value(&value) + .map(Self) + .map_err(|error| PyValueError::new_err(error.to_string())) + } +} + +impl<'py, T: IntoValue> IntoPyObject<'py> for PyData { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + PyValue(self.0.into_value()).into_pyobject(py) + } +} + +fn unwrap_all(values: Vec>) -> Vec { + values.into_iter().map(|value| value.0).collect() +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pyo3::prelude::*; + use rstest::rstest; + + use super::ttl; + + #[rstest] + #[case::infinity(f64::INFINITY, "OverflowError")] + #[case::nan(f64::NAN, "ValueError")] + fn a_ttl_int_refuses_raises_what_int_raises(#[case] seconds: f64, #[case] raised: &str) { + Python::initialize(); + Python::attach(|py| { + let error = ttl(Some(seconds)).unwrap_err(); + + assert_eq!(error.get_type(py).name().unwrap().to_string(), raised); + }); + } + + #[rstest] + #[case::fraction(Some(1.5), Some(Duration::from_millis(1500)))] + #[case::negative_clamps_to_zero(Some(-1.0), Some(Duration::ZERO))] + #[case::absent(None, None)] + fn a_finite_ttl_converts(#[case] seconds: Option, #[case] expected: Option) { + assert_eq!(ttl(seconds).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/policy.rs b/litellm-rust/crates/python-bridge/src/redis/store/policy.rs new file mode 100644 index 00000000000..ffdf7b97a2e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/policy.rs @@ -0,0 +1,151 @@ +use std::time::Duration; + +use litellm_redis_store::TierPolicy; +use pyo3::prelude::*; + +use crate::{ + coercion::{Field, FieldSpec, ProjectionError}, + python_settings::{PythonSettings, Snapshot}, +}; + +const COOLDOWN_REFRESH: FieldSpec = + FieldSpec::new("cooldown_redis_read_interval", seconds); +const USAGE_REFRESH: FieldSpec = FieldSpec::new("usage_redis_read_interval", seconds); +const TRACKED_READS: FieldSpec = FieldSpec::new("max_redis_batch_cache_size", count); +const MEMORY_TTL: FieldSpec> = + FieldSpec::new("default_in_memory_ttl", optional_seconds); + +/// The tier values this Python process runs with, from its environment and `litellm` globals. +#[derive(Debug)] +pub(crate) struct RouterTierPolicies { + pub(crate) cooldowns: TierPolicy, + pub(crate) usage: TierPolicy, +} + +pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult { + Ok(project(&PythonSettings::RouterTiers.read(py)?)?) +} + +fn project(snapshot: &Snapshot<'_>) -> Result { + let tracked_reads = snapshot.read(&TRACKED_READS)?; + let default_in_memory_ttl = snapshot.read(&MEMORY_TTL)?; + Ok(RouterTierPolicies { + cooldowns: TierPolicy { + refresh_after: snapshot.read(&COOLDOWN_REFRESH)?, + tracked_reads, + default_in_memory_ttl, + }, + usage: TierPolicy { + refresh_after: snapshot.read(&USAGE_REFRESH)?, + tracked_reads, + default_in_memory_ttl, + }, + }) +} + +fn seconds(field: &Field<'_>) -> Result { + match field.value().extract::() { + Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => { + Ok(Duration::from_secs_f64(seconds)) + } + _ => Err(field.invalid("a non-negative number of seconds")), + } +} + +fn optional_seconds(field: &Field<'_>) -> Result, ProjectionError> { + match field.value().is_none() { + true => Ok(None), + false => seconds(field).map(Some), + } +} + +fn count(field: &Field<'_>) -> Result { + field + .value() + .extract::() + .map_err(|_| field.invalid("a non-negative integer")) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pyo3::{prelude::*, types::PyDict}; + + use super::project; + use crate::python_settings::PythonSettings; + + fn projected(py: Python<'_>, fields: &str) -> Result { + let locals = PyDict::new(py); + let source = std::ffi::CString::new(format!( + "import types\nsettings = types.SimpleNamespace({fields})" + )) + .unwrap(); + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + let snapshot = + PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap()); + project(&snapshot).map_err(|error| PyErr::from(error).to_string()) + } + + #[test] + fn the_host_values_are_each_familys_refresh_and_tracked_reads() { + Python::initialize(); + Python::attach(|py| { + let policies = projected( + py, + "cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42, default_in_memory_ttl=None", + ) + .unwrap(); + assert_eq!( + policies.cooldowns.refresh_after, + Duration::from_millis(2500) + ); + assert_eq!(policies.usage.refresh_after, Duration::from_secs(30)); + assert_eq!(policies.cooldowns.tracked_reads, 42); + assert_eq!(policies.usage.tracked_reads, 42); + assert_eq!(policies.cooldowns.default_in_memory_ttl, None); + }); + } + + #[test] + fn a_set_default_in_memory_ttl_reaches_both_familys_backfill() { + Python::initialize(); + Python::attach(|py| { + let policies = projected( + py, + "cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=7.5", + ) + .unwrap(); + assert_eq!( + policies.cooldowns.default_in_memory_ttl, + Some(Duration::from_millis(7500)) + ); + assert_eq!( + policies.usage.default_in_memory_ttl, + Some(Duration::from_millis(7500)) + ); + }); + } + + #[test] + fn a_value_python_would_not_run_with_is_a_configuration_error() { + Python::initialize(); + Python::attach(|py| { + let negative = projected( + py, + "cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=None", + ) + .unwrap_err(); + let text = projected( + py, + "cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many', default_in_memory_ttl=None", + ) + .unwrap_err(); + assert!( + negative.contains("cooldown_redis_read_interval"), + "{negative}" + ); + assert!(text.contains("max_redis_batch_cache_size"), "{text}"); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs new file mode 100644 index 00000000000..ca95e1530f4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs @@ -0,0 +1,577 @@ +use std::{sync::Arc, time::Duration}; + +use indexmap::IndexMap; +use litellm_redis_store::{ + Pending, + rate_limits::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge, GuardedIncrement, + GuardedReply, Increment, Kind, Local, RateLimits, Scope, WindowCount, + }, +}; +use pyo3::{exceptions::PyValueError, prelude::*}; + +use super::{ + NativePending, PyNumber, Spelled, kind_error, redis_async, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, +}; +use crate::redis::client::NativeRedisClient; + +/// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`; +/// every Redis call is a future. +#[pyclass(frozen, name = "NativeRateLimits")] +pub(crate) struct NativeRateLimits(Arc); + +/// An `int` or `float` counter, a `str` window start, or a `{slot_id: acquired_at}` registry. +#[derive(FromPyObject)] +enum PyLocal { + Number(PyNumber), + Text(String), + Slots(IndexMap), +} + +impl From for Local { + fn from(local: PyLocal) -> Self { + match local { + PyLocal::Number(number) => Self::Number(number.0), + PyLocal::Text(text) => Self::Text(text), + PyLocal::Slots(slots) => Self::Slots(slots.into_iter().collect()), + } + } +} + +/// What the memory tier holds, back as the Python value it was written as. +struct LocalValue(Local); + +impl<'py> IntoPyObject<'py> for LocalValue { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Ok(match self.0 { + Local::Number(number) => PyNumber(number).into_pyobject(py)?, + Local::Text(text) => text.into_pyobject(py)?.into_any(), + Local::Slots(slots) => slots + .into_iter() + .collect::>() + .into_pyobject(py)? + .into_any(), + }) + } +} + +/// `litellm.redis.stores.rate_limits.RateLimitScope`. +#[derive(FromPyObject)] +struct PyScope { + descriptor_key: String, + descriptor_value: String, +} + +impl From for Scope { + fn from(scope: PyScope) -> Self { + Self { + descriptor: scope.descriptor_key, + value: scope.descriptor_value, + } + } +} + +/// `RateLimitCounter`. +#[derive(FromPyObject)] +struct PyCounter { + scope: PyScope, + kind: Spelled, +} + +impl From for Counter { + fn from(counter: PyCounter) -> Self { + Self { + scope: counter.scope.into(), + kind: counter.kind.0, + } + } +} + +/// `RateLimitWindow`, `RateLimitCounter` or `RateLimitGauge`, told apart by the key suffix each names. +#[derive(FromPyObject)] +struct PyEntry { + suffix: String, + scope: PyScope, +} + +impl TryFrom for Entry { + type Error = PyErr; + + fn try_from(entry: PyEntry) -> PyResult { + Self::new(&entry.suffix, entry.scope.into()) + .map_err(|_| PyValueError::new_err(format!("unknown key suffix {:?}", entry.suffix))) + } +} + +/// `CounterIncrement`. +#[derive(FromPyObject)] +struct PyIncrement { + counter: PyCounter, + amount: PyNumber, + ttl: Option, +} + +impl From for Increment { + fn from(increment: PyIncrement) -> Self { + Self { + counter: increment.counter.into(), + amount: increment.amount.0, + ttl: increment.ttl.map(Duration::from_secs), + } + } +} + +/// `AtomicCounter`. +#[derive(FromPyObject)] +struct PyAtomicCounter { + counter: PyCounter, + limit: i64, + increment: i64, + ttl: i64, + window_size: i64, +} + +impl From for AtomicCounter { + fn from(counter: PyAtomicCounter) -> Self { + Self { + counter: counter.counter.into(), + limit: counter.limit, + increment: counter.increment, + ttl_seconds: counter.ttl, + window_size_seconds: counter.window_size, + } + } +} + +fn all>(records: Vec

) -> Vec { + records.into_iter().map(T::from).collect() +} + +#[pyclass(frozen, eq, get_all, name = "WindowCount")] +#[derive(PartialEq)] +pub(crate) struct PyWindowCount { + window_start: String, + count: i64, +} + +#[pymethods] +impl PyWindowCount { + #[new] + fn new(window_start: String, count: i64) -> Self { + Self { + window_start, + count, + } + } +} + +/// Per counter, the count after the increment and the start of the window it landed in. +#[pyclass(frozen, eq, get_all, name = "AtomicApplied")] +#[derive(PartialEq)] +pub(crate) struct PyAtomicApplied { + counters: Vec<(i64, i64)>, +} + +#[pymethods] +impl PyAtomicApplied { + #[new] + fn new(counters: Vec<(i64, i64)>) -> Self { + Self { counters } + } +} + +#[pyclass(frozen, eq, get_all, name = "AtomicOverLimit")] +#[derive(PartialEq)] +pub(crate) struct PyAtomicOverLimit { + index: usize, + current: i64, + limit: i64, +} + +#[pymethods] +impl PyAtomicOverLimit { + #[new] + fn new(index: usize, current: i64, limit: i64) -> Self { + Self { + index, + current, + limit, + } + } +} + +#[derive(IntoPyObject)] +enum PyAtomicReply { + Applied(PyAtomicApplied), + OverLimit(PyAtomicOverLimit), +} + +#[pyclass(frozen, eq, get_all, name = "GuardedReply")] +#[derive(PartialEq)] +pub(crate) struct PyGuardedReply { + applied: bool, + count: i64, +} + +#[pymethods] +impl PyGuardedReply { + #[new] + fn new(applied: bool, count: i64) -> Self { + Self { applied, count } + } +} + +#[pyclass(frozen, eq, get_all, name = "GaugesAcquired")] +#[derive(PartialEq)] +pub(crate) struct PyGaugesAcquired { + in_flight: Vec, +} + +#[pymethods] +impl PyGaugesAcquired { + #[new] + fn new(in_flight: Vec) -> Self { + Self { in_flight } + } +} + +#[pyclass(frozen, eq, get_all, name = "GaugeOverLimit")] +#[derive(PartialEq)] +pub(crate) struct PyGaugeOverLimit { + index: usize, + in_flight: i64, + limit: i64, +} + +#[pymethods] +impl PyGaugeOverLimit { + #[new] + fn new(index: usize, in_flight: i64, limit: i64) -> Self { + Self { + index, + in_flight, + limit, + } + } +} + +#[derive(IntoPyObject)] +enum PyAcquireReply { + Acquired(PyGaugesAcquired), + OverLimit(PyGaugeOverLimit), +} + +fn window_counts(counts: Vec) -> Vec { + counts + .into_iter() + .map( + |WindowCount { + window_start, + count, + }| PyWindowCount { + window_start, + count, + }, + ) + .collect() +} + +fn atomic(reply: AtomicReply) -> PyAtomicReply { + match reply { + AtomicReply::Applied(applied) => PyAtomicReply::Applied(PyAtomicApplied { + counters: applied + .into_iter() + .map( + |Applied { + count, + window_start, + }| (count, window_start), + ) + .collect(), + }), + AtomicReply::OverLimit { + index, + current, + limit, + } => PyAtomicReply::OverLimit(PyAtomicOverLimit { + index, + current, + limit, + }), + } +} + +fn guarded(GuardedReply { applied, count }: GuardedReply) -> PyGuardedReply { + PyGuardedReply { applied, count } +} + +fn acquire(reply: AcquireReply) -> PyAcquireReply { + match reply { + AcquireReply::Acquired(in_flight) => { + PyAcquireReply::Acquired(PyGaugesAcquired { in_flight }) + } + AcquireReply::OverLimit { + index, + in_flight, + limit, + } => PyAcquireReply::OverLimit(PyGaugeOverLimit { + index, + in_flight, + limit, + }), + } +} + +/// One script call declared on the batch the caller's `NativeSend` picked; `None` without Redis. +fn script( + call_type: &'static str, + declared: Option>, + reply: fn(T) -> R, +) -> Option +where + T: Send + 'static, + R: for<'a> IntoPyObject<'a> + Send + 'static, +{ + Some(super::pending(call_type, kind_error, declared?, reply)) +} + +#[pymethods] +impl NativeRateLimits { + #[new] + #[pyo3(signature = (client, backfill_ttl_seconds = None))] + fn new( + client: Option>, + backfill_ttl_seconds: Option, + ) -> PyResult { + Ok(Self(Arc::new(RateLimits::new( + client.map(|client| remote(&client)), + ttl(backfill_ttl_seconds)?, + )))) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn slot_groups(&self, scopes: Vec, force: bool) -> Vec> { + self.0.slot_groups(&all(scopes), force) + } + + fn get_local(&self, entry: PyEntry) -> PyResult> { + Ok(self + .0 + .get_local(&entry.try_into()?) + .map_err(kind_error)? + .map(LocalValue)) + } + + #[pyo3(signature = (entry, value, ttl_seconds))] + fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option) -> PyResult<()> { + self.0 + .set_local(&entry.try_into()?, value.into(), ttl(ttl_seconds)?) + .map_err(kind_error) + } + + /// `request` answers a key its pre-call reads already found absent without another `GET`. + #[pyo3(signature = (entry, *, backfill_ttl_seconds = None, request = None))] + fn get<'py>( + &self, + py: Python<'py>, + entry: PyEntry, + backfill_ttl_seconds: Option, + request: Option<&NativeRequestRedis>, + ) -> PyResult> { + let store = self.0.clone(); + let entry = Entry::try_from(entry)?; + let backfill_ttl = ttl(backfill_ttl_seconds)?; + let request = request.map(|request| request.0.clone()); + redis_async(py, "async_get_cache", kind_error, async move { + store + .get(request.as_ref(), &entry, backfill_ttl) + .await + .map(|local| local.map(LocalValue)) + }) + } + + fn sliding_window( + &self, + counters: Vec, + now: i64, + window_size: i64, + send: &NativeSend, + ) -> Option { + script( + "batch_rate_limiter_script", + self.0 + .sliding_window(send.send(), &all(counters), now, window_size), + window_counts, + ) + } + + fn read_counters(&self, counters: Vec, send: &NativeSend) -> Option { + script( + "batch_counter_read_script", + self.0.read_counters(send.send(), &all(counters)), + std::convert::identity, + ) + } + + fn check_and_increment( + &self, + counters: Vec, + send: &NativeSend, + ) -> Option { + script( + "check_and_increment_by_n_script", + self.0.check_and_increment(send.send(), &all(counters)), + atomic, + ) + } + + fn guarded_increment( + &self, + counter: PyCounter, + expected_window_start: String, + amount: PyNumber, + ttl_seconds: Option, + send: &NativeSend, + ) -> Option { + let increment = GuardedIncrement { + counter: counter.into(), + expected_window_start, + amount: amount.0, + ttl: ttl_seconds.map(Duration::from_secs), + }; + script( + "window_guarded_token_increment_script", + self.0.guarded_increment(send.send(), &increment), + guarded, + ) + } + + fn parallel_acquire( + &self, + gauges: Vec<(PyScope, i64)>, + slot_ttl_seconds: u64, + slot_id: &str, + send: &NativeSend, + ) -> Option { + let gauges: Vec = gauges + .into_iter() + .map(|(scope, limit)| Gauge { + scope: scope.into(), + limit, + }) + .collect(); + script( + "parallel_acquire_script", + self.0.parallel_acquire( + send.send(), + &gauges, + Duration::from_secs(slot_ttl_seconds), + slot_id, + ), + acquire, + ) + } + + fn parallel_release( + &self, + gauges: Vec, + slot_id: &str, + send: &NativeSend, + ) -> Option { + script( + "parallel_release_script", + self.0.parallel_release(send.send(), &all(gauges), slot_id), + std::convert::identity, + ) + } + + fn parallel_count( + &self, + gauges: Vec, + slot_ttl_seconds: u64, + send: &NativeSend, + ) -> Option { + script( + "parallel_count_script", + self.0.parallel_count( + send.send(), + &all(gauges), + Duration::from_secs(slot_ttl_seconds), + ), + std::convert::identity, + ) + } + + fn token_increment( + &self, + increments: Vec, + send: &NativeSend, + ) -> Option { + script( + "token_increment_script", + self.0.token_increment(send.send(), &all(increments)), + std::convert::identity, + ) + } + + /// `False` without Redis. + fn token_increment_post_call( + &self, + request: &NativeRequestRedis, + increments: Vec, + ) -> bool { + self.0 + .token_increment_post_call(&request.0, all(increments)) + } + + /// `False` without Redis. + fn parallel_release_post_call( + &self, + request: &NativeRequestRedis, + gauges: Vec, + slot_id: &str, + ) -> bool { + self.0 + .parallel_release_post_call(&request.0, &all(gauges), slot_id) + } + + fn increment<'py>( + &self, + py: Python<'py>, + increment: PyIncrement, + ) -> PyResult> { + let store = self.0.clone(); + let increment = Increment::from(increment); + redis_async(py, "async_increment", kind_error, async move { + Ok(store.increment(&increment).await.map(PyNumber)) + }) + } + + fn increment_many<'py>( + &self, + py: Python<'py>, + increments: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let increments: Vec = all(increments); + redis_async(py, "async_increment_pipeline", kind_error, async move { + store.increment_many(&increments).await; + Ok(()) + }) + } + + fn increment_many_post_call(&self, request: &NativeRequestRedis, increments: Vec) { + self.0 + .increment_many_post_call(&request.0, &all(increments)); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/request.rs b/litellm-rust/crates/python-bridge/src/redis/store/request.rs new file mode 100644 index 00000000000..3436b26a16c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/request.rs @@ -0,0 +1,60 @@ +use litellm_redis_store::{RequestRedis, Send}; +use pyo3::prelude::*; + +use super::{call, duration}; +use crate::execution::run_async_value; + +/// One request's native Redis batches. Python builds one per request scope and passes it to each +/// request-aware store call. +#[pyclass(frozen, name = "NativeRequestRedis")] +pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis); + +#[pymethods] +impl NativeRequestRedis { + #[new] + fn new(py: Python<'_>, post_call_deadline_seconds: f64) -> PyResult { + let runtime = litellm_host_python::runtime_handle()?; + let deadline = duration(post_call_deadline_seconds)?; + let deadline_tag = call(py, "post_call_redis_batch")?.tag(); + Ok(Self(crate::logger::capture(py).scope(|| { + RequestRedis::with_deadline_tag(runtime, deadline, Some(deadline_tag)) + }))) + } + + fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "request_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_pre()).await; + Ok(()) + }) + } + + fn flush_post<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "post_call_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_post()).await; + Ok(()) + }) + } +} + +/// The Python face of `Send`: which batch a store op rides. Each variant carries exactly what it +/// needs, so there is no phase string and no phase without a request. +#[pyclass(frozen, name = "NativeSend")] +pub(crate) enum NativeSend { + Now(), + PreCall(Py), + PostCall(Py), +} + +impl NativeSend { + pub(crate) fn send(&self) -> Send<'_> { + match self { + Self::Now() => Send::Now, + Self::PreCall(request) => Send::PreCall(&request.get().0), + Self::PostCall(request) => Send::PostCall(&request.get().0), + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs new file mode 100644 index 00000000000..5d695ca7090 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs @@ -0,0 +1,182 @@ +use std::{collections::HashMap, sync::Arc}; + +use litellm_redis_store::{ + Send, + router_cooldowns::{Cooldown, RouterCooldowns}, +}; +use pyo3::{prelude::*, types::PyTuple}; + +use super::{ + NativePending, PyData, PyNumber, duration, kind_error, pending, policy::router_tier_policies, + redis_async, redis_sync, remote, request::NativeSend, +}; +use crate::redis::client::NativeRedisClient; + +/// One router's cooldown store. The router's sync completion path calls `put`, `active_sync`, +/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py; the +/// `async_` twins return awaitables for async callers. +#[pyclass(frozen, name = "NativeRouterCooldowns")] +pub(crate) struct NativeRouterCooldowns(Arc); + +/// `[(model_id, cooldown), ...]` the way `RedisCache` hands the router its batch read. +struct PyActive(Vec<(String, Cooldown)>); + +impl<'py> IntoPyObject<'py> for PyActive { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + active_to_py(py, self.0)? + .into_pyobject(py) + .map(Bound::into_any) + } +} + +fn active_to_py( + py: Python<'_>, + active: Vec<(String, Cooldown)>, +) -> PyResult>> { + active + .into_iter() + .map(|(model_id, cooldown)| { + PyTuple::new( + py, + [ + model_id.into_pyobject(py)?.into_any(), + PyData(cooldown).into_pyobject(py)?, + ], + ) + }) + .collect() +} + +#[pymethods] +impl NativeRouterCooldowns { + #[new] + fn new(py: Python<'_>, client: Option>) -> PyResult { + let remote = client.map(|client| remote(&client)); + let policy = router_tier_policies(py)?.cooldowns; + Ok(Self(Arc::new(RouterCooldowns::new(remote, policy)))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool { + self.0.attach(remote(&client)) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn put( + &self, + py: Python<'_>, + model_id: String, + cooldown: PyData, + ttl_seconds: f64, + ) -> PyResult<()> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_sync(py, "set_cache", kind_error, async move { + store.put_str(&model_id, &cooldown.0, ttl).await + }) + } + + fn increment_allowed_fails( + &self, + py: Python<'_>, + model_id: String, + suffix: Option, + ttl_seconds: f64, + ) -> PyResult { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_sync(py, "increment_cache", kind_error, async move { + store + .increment_allowed_fails_int(&model_id, suffix.as_deref(), ttl) + .await + }) + } + + fn async_put<'py>( + &self, + py: Python<'py>, + model_id: String, + cooldown: PyData, + ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.put(&model_id, &cooldown.0, ttl).await + }) + } + + fn async_increment_allowed_fails<'py>( + &self, + py: Python<'py>, + model_id: String, + suffix: Option, + ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_async(py, "async_increment", kind_error, async move { + store + .increment_allowed_fails(&model_id, suffix.as_deref(), ttl) + .await + }) + } + + /// The deployments still cooling down, read on the batch `send` picks. + fn active(&self, model_ids: Vec, send: &NativeSend) -> PyResult { + let active = self.0.active(send.send(), &model_ids).map_err(kind_error)?; + Ok(pending( + "async_batch_get_cache", + kind_error, + active, + PyActive, + )) + } + + /// `active` on its own with the GIL released, for the router's sync completion path. + fn active_sync<'py>( + &self, + py: Python<'py>, + model_ids: Vec, + ) -> PyResult>> { + let active = self.0.active(Send::Now, &model_ids).map_err(kind_error)?; + let active = redis_sync(py, "batch_get_cache", kind_error, active)?; + active_to_py(py, active) + } + + fn minimum(&self, py: Python<'_>, model_ids: Vec) -> PyResult> { + let store = self.0.clone(); + redis_sync(py, "batch_get_cache", kind_error, async move { + store + .minimum(&model_ids) + .await + .map(|minimum| minimum.map(PyNumber)) + }) + } + + fn delete_local(&self, model_id: &str) -> PyResult<()> { + self.0.delete_local(model_id).map_err(kind_error) + } + + fn flush_local_counters(&self) -> PyResult<()> { + self.0.flush_local_counters().map_err(kind_error) + } + + fn snapshot(&self) -> PyResult>> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, cooldown)| (key, PyData(cooldown))) + .collect()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs new file mode 100644 index 00000000000..48565028de3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -0,0 +1,182 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage}; +use pyo3::prelude::*; + +use super::{ + PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote, +}; +use crate::redis::client::NativeRedisClient; + +/// `litellm.redis.stores.router_usage.UsageCounter`. +#[derive(FromPyObject)] +struct PyCounter { + deployment_id: String, + model: Option, + minute: String, + kind: Spelled, +} + +impl From for Counter { + fn from(counter: PyCounter) -> Self { + Self { + deployment_id: counter.deployment_id, + model: counter.model, + minute: counter.minute, + kind: counter.kind.0, + } + } +} + +fn counters(counters: Vec) -> Vec { + counters.into_iter().map(Counter::from).collect() +} + +/// One router's global TPM/RPM counters. The `_sync` methods block with the GIL released, for the +/// Rust router's sync calls, which have no event loop to await on. +#[pyclass(frozen, name = "NativeRouterUsage")] +pub(crate) struct NativeRouterUsage(Arc); + +#[pymethods] +impl NativeRouterUsage { + #[new] + fn new(py: Python<'_>, client: Option>) -> PyResult { + let remote = client.map(|client| remote(&client)); + let policy = router_tier_policies(py)?.usage; + Ok(Self(Arc::new(RouterUsage::new(remote, policy)))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool { + self.0.attach(remote(&client)) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn increment<'py>( + &self, + py: Python<'py>, + counter: PyCounter, + value: PyNumber, + ttl_seconds: u64, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_async(py, "async_increment", kind_error, async move { + Ok(store + .increment(&counter, value.0, Duration::from_secs(ttl_seconds)) + .await + .map(PyNumber)) + }) + } + + /// `_increment_deployment`'s `async_increment_pipeline`, always on its own round trip. + fn increment_deployment<'py>( + &self, + py: Python<'py>, + deployment_id: String, + model: Option, + minute: String, + tokens: PyNumber, + requests: i64, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_increment_pipeline", kind_error, async move { + store + .increment_deployment( + &deployment_id, + model.as_deref(), + &minute, + tokens.0, + requests, + ) + .await + }) + } + + fn get_many<'py>( + &self, + py: Python<'py>, + counters: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let counters = self::counters(counters); + redis_async(py, "async_batch_get_cache", kind_error, async move { + let read = store.get_many(&counters).await?; + Ok(read + .into_iter() + .map(|count| count.map(PyNumber)) + .collect::>()) + }) + } + + fn increment_sync( + &self, + py: Python<'_>, + counter: PyCounter, + value: PyNumber, + ttl_seconds: u64, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_sync(py, "increment_cache", kind_error, async move { + Ok(store + .increment(&counter, value.0, Duration::from_secs(ttl_seconds)) + .await + .map(PyNumber)) + }) + } + + fn increment_deployment_sync( + &self, + py: Python<'_>, + deployment_id: String, + model: Option, + minute: String, + tokens: PyNumber, + requests: i64, + ) -> PyResult<()> { + let store = self.0.clone(); + redis_sync(py, "increment_pipeline", kind_error, async move { + store + .increment_deployment( + &deployment_id, + model.as_deref(), + &minute, + tokens.0, + requests, + ) + .await + }) + } + + fn get_many_sync( + &self, + py: Python<'_>, + counters: Vec, + ) -> PyResult>> { + let store = self.0.clone(); + let counters = self::counters(counters); + redis_sync(py, "batch_get_cache", kind_error, async move { + let read = store.get_many(&counters).await?; + Ok(read.into_iter().map(|count| count.map(PyNumber)).collect()) + }) + } + + fn flush_local_counters(&self) -> PyResult<()> { + self.0.flush_local_counters().map_err(kind_error) + } + + fn snapshot(&self) -> PyResult> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, count)| (key, PyNumber(count))) + .collect()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/sensitive_route_pins.rs b/litellm-rust/crates/python-bridge/src/redis/store/sensitive_route_pins.rs new file mode 100644 index 00000000000..8b1b82945fb --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/sensitive_route_pins.rs @@ -0,0 +1,49 @@ +use std::sync::Arc; + +use litellm_redis_store::sensitive_route_pins::SensitiveRoutePins; +use pyo3::prelude::*; + +use super::{duration, kind_error, redis_async, remote}; +use crate::redis::client::NativeRedisClient; + +#[pyclass(frozen, name = "NativeSensitiveRoutePins")] +pub(crate) struct NativeSensitiveRoutePins(Arc); + +#[pymethods] +impl NativeSensitiveRoutePins { + #[new] + fn new(client: Option>) -> Self { + Self(Arc::new(SensitiveRoutePins::new( + client.map(|client| remote(&client)), + ))) + } + + fn get<'py>( + &self, + py: Python<'py>, + tenant: String, + session_id: String, + default_ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let default_ttl = duration(default_ttl_seconds)?; + redis_async(py, "async_get_cache", kind_error, async move { + store.get(&tenant, &session_id, default_ttl).await + }) + } + + fn set<'py>( + &self, + py: Python<'py>, + tenant: String, + session_id: String, + model: String, + ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.set(&tenant, &session_id, &model, ttl).await + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs new file mode 100644 index 00000000000..9da6a0e00e5 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -0,0 +1,277 @@ +use std::{collections::HashMap, sync::Arc}; + +use litellm_redis_store::{ + Number, + spend_counters::{LostSeed, PostCallSettlement, Reading, SpendCounter, SpendCounters}, +}; +use pyo3::prelude::*; + +use super::{ + NativePending, PyData, PyNumber, duration, kind_error, pending, redis_async, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, unwrap_all, +}; +use crate::redis::client::NativeRedisClient; + +/// A counter's value and whether Redis answered it, a clean miss included. +#[pyclass(frozen, eq, get_all, name = "SpendRead")] +#[derive(PartialEq)] +pub(crate) struct PySpendRead { + value: Option, + from_redis: bool, +} + +#[pymethods] +impl PySpendRead { + #[new] + fn new(value: Option, from_redis: bool) -> Self { + Self { value, from_redis } + } +} + +/// Whether every increment one `increment_many_post_call` declared has its reply. +#[pyclass(frozen, name = "PostCallSettlement")] +pub(crate) struct PyPostCallSettlement(PostCallSettlement); + +#[pymethods] +impl PyPostCallSettlement { + #[getter] + fn settled(&self) -> bool { + self.0.is_settled() + } +} + +fn increments(increments: Vec<(PyData, PyNumber)>) -> Vec<(SpendCounter, Number)> { + increments + .into_iter() + .map(|(counter, amount)| (counter.0, amount.0)) + .collect() +} + +fn numbers(values: Option>>) -> Option>> { + values.map(|values| { + values + .into_iter() + .map(|value| value.map(PyNumber)) + .collect() + }) +} + +/// The proxy's spend counters, Redis first with this worker's memory copy behind it. +#[pyclass(frozen, name = "NativeSpendCounters")] +pub(crate) struct NativeSpendCounters(Arc); + +#[pymethods] +impl NativeSpendCounters { + #[new] + fn new(client: Option>) -> Self { + Self(Arc::new(SpendCounters::new( + client.map(|client| remote(&client)), + ))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool { + self.0.attach(remote(&client)) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + #[getter] + fn redis_ttl_seconds(&self) -> Option { + self.0.redis_ttl().map(|ttl| ttl.as_secs_f64()) + } + + /// The least TTL a Redis write gives a committed spend twin. + fn set_committed_ttl(&self, ttl_seconds: f64) -> PyResult<()> { + self.0.set_committed_ttl(duration(ttl_seconds)?); + Ok(()) + } + + /// `("redis" | "memory", value)`: where the answer came from, a Redis miss included. + fn read<'py>( + &self, + py: Python<'py>, + counter: PyData, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_get_cache", kind_error, async move { + let reading = store.read(&counter.0).await?; + Ok(PySpendRead { + value: reading.value().map(PyNumber), + from_redis: matches!(reading, Reading::Redis(_)), + }) + }) + } + + /// One `MGET` on the batch `send` picks; `None` without Redis. It resolves to `None` when the + /// read failed, so each caller reads on its own. + fn read_many( + &self, + counters: Vec>, + send: &NativeSend, + ) -> Option { + let read = self.0.read_many(send.send(), &unwrap_all(counters))?; + Some(pending("async_batch_get_cache", kind_error, read, numbers)) + } + + fn increment<'py>( + &self, + py: Python<'py>, + counter: PyData, + amount: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_increment", kind_error, async move { + store.increment(&counter.0, amount.0).await.map(PyNumber) + }) + } + + fn increment_many<'py>( + &self, + py: Python<'py>, + increments: Vec<(PyData, PyNumber)>, + ) -> PyResult> { + let store = self.0.clone(); + let increments = self::increments(increments); + redis_async(py, "async_increment_pipeline", kind_error, async move { + let totals = store.increment_many(&increments).await?; + Ok(totals.into_iter().map(PyNumber).collect::>()) + }) + } + + /// Declares the increments on the request's post-call batch; `None` when there is no Redis. + fn increment_many_post_call( + &self, + request: &NativeRequestRedis, + increments: Vec<(PyData, PyNumber)>, + ) -> Option { + self.0 + .increment_many_post_call(&request.0, &self::increments(increments)) + .map(PyPostCallSettlement) + } + + fn set_max<'py>( + &self, + py: Python<'py>, + counter: PyData, + spend: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_set_max", kind_error, async move { + store + .set_max(&counter.0, spend.0) + .await + .map(|stored| stored.map(PyNumber)) + }) + } + + #[pyo3(signature = (counter, spend, *, increment_if_lost))] + fn seed<'py>( + &self, + py: Python<'py>, + counter: PyData, + spend: PyNumber, + increment_if_lost: bool, + ) -> PyResult> { + let store = self.0.clone(); + let lost = if increment_if_lost { + LostSeed::Increment + } else { + LostSeed::Answer + }; + redis_async(py, "async_set_cache", kind_error, async move { + store.seed(&counter.0, spend.0, lost).await.map(PyNumber) + }) + } + + #[pyo3(signature = (counter, spend, ttl_seconds = None))] + fn set<'py>( + &self, + py: Python<'py>, + counter: PyData, + spend: PyNumber, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.set(&counter.0, spend.0, ttl).await + }) + } + + fn delete<'py>( + &self, + py: Python<'py>, + counters: Vec>, + ) -> PyResult> { + let store = self.0.clone(); + let counters = unwrap_all(counters); + redis_async(py, "async_delete_cache", kind_error, async move { + store.delete(&counters).await + }) + } + + fn refresh_ttl<'py>( + &self, + py: Python<'py>, + counter: PyData, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_refresh_ttl", kind_error, async move { + store.refresh_ttl(&counter.0).await + }) + } + + fn get_local(&self, counter: PyData) -> PyResult> { + Ok(self + .0 + .get_local(&counter.0) + .map_err(kind_error)? + .map(PyNumber)) + } + + #[pyo3(signature = (counter, spend, ttl_seconds = None))] + fn set_local( + &self, + counter: PyData, + spend: PyNumber, + ttl_seconds: Option, + ) -> PyResult<()> { + self.0 + .set_local(&counter.0, spend.0, ttl(ttl_seconds)?) + .map_err(kind_error) + } + + fn delete_local(&self, counter: PyData) -> PyResult<()> { + self.0.delete_local(&counter.0).map_err(kind_error) + } + + fn increment_local( + &self, + counter: PyData, + amount: PyNumber, + ) -> PyResult { + self.0 + .increment_local(&counter.0, amount.0) + .map(PyNumber) + .map_err(kind_error) + } + + fn flush_local(&self) -> PyResult<()> { + self.0.flush_local().map_err(kind_error) + } + + fn snapshot(&self) -> PyResult> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, spend)| (key, PyNumber(spend))) + .collect()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/value.rs b/litellm-rust/crates/python-bridge/src/redis/store/value.rs new file mode 100644 index 00000000000..c86c9e9bb20 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/value.rs @@ -0,0 +1,125 @@ +use litellm_python_compat::{BigInt, MAX_DEPTH, Value}; +use pyo3::{ + exceptions::{PyTypeError, PyValueError}, + prelude::*, + types::{ + PyBool, PyBytes, PyComplex, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PySet, PyString, + PyTuple, + }, +}; + +/// A plain Python value (what `json.loads` or `ast.literal_eval` returns) across the bridge. +/// Anything else, a live object included, is refused with `TypeError`. +pub(crate) struct PyValue(pub(crate) Value); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyValue { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + to_value(&value, 0).map(Self) + } +} + +impl<'py> IntoPyObject<'py> for PyValue { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + to_python(py, self.0) + } +} + +fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult { + if depth > MAX_DEPTH { + return Err(PyValueError::new_err("value is nested too deeply")); + } + let nested = |items: Bound<'_, PyAny>| -> PyResult> { + items + .try_iter()? + .map(|item| to_value(&item?, depth + 1)) + .collect() + }; + if value.is_none() { + return Ok(Value::None); + } + if let Ok(flag) = value.cast::() { + return Ok(Value::Bool(flag.is_true())); + } + if value.is_instance_of::() { + return Ok(Value::Int(match value.extract::() { + Ok(small) => small.into(), + Err(_) => value + .str()? + .to_str()? + .parse::() + .map_err(|error| PyValueError::new_err(error.to_string()))?, + })); + } + if let Ok(number) = value.cast::() { + return Ok(Value::Float(number.value())); + } + if let Ok(text) = value.cast::() { + return Ok(Value::Str(text.to_str()?.to_owned())); + } + if let Ok(bytes) = value.cast::() { + return Ok(Value::Bytes(bytes.as_bytes().to_vec())); + } + if let Ok(complex) = value.cast::() { + return Ok(Value::Complex { + re: complex.real(), + im: complex.imag(), + }); + } + if let Ok(entries) = value.cast::() { + return entries + .iter() + .map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?))) + .collect::>() + .map(Value::Dict); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::List); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::Tuple); + } + if value.is_instance_of::() || value.is_instance_of::() { + return nested(value.clone()).map(Value::Set); + } + Err(PyTypeError::new_err(format!( + "{} is not a plain data value", + value.get_type().name()? + ))) +} + +fn to_python(py: Python<'_>, value: Value) -> PyResult> { + let all = |items: Vec| -> PyResult>> { + items.into_iter().map(|item| to_python(py, item)).collect() + }; + Ok(match value { + Value::None => py.None().into_bound(py), + Value::Bool(flag) => PyBool::new(py, flag).to_owned().into_any(), + Value::Int(number) => match i64::try_from(&number) { + Ok(small) => small.into_pyobject(py)?.into_any(), + Err(_) => py + .get_type::() + .call1((number.to_string(),))? + .into_any(), + }, + Value::Float(number) => number.into_pyobject(py)?.into_any(), + Value::Complex { re, im } => PyComplex::from_doubles(py, re, im).into_any(), + Value::Str(text) => text.into_pyobject(py)?.into_any(), + Value::Bytes(bytes) => PyBytes::new(py, &bytes).into_any(), + Value::Tuple(items) => PyTuple::new(py, all(items)?)?.into_any(), + Value::List(items) => PyList::new(py, all(items)?)?.into_any(), + Value::Set(items) => PySet::new(py, all(items)?)?.into_any(), + Value::Dict(entries) => { + let dict = PyDict::new(py); + for (key, item) in entries { + dict.set_item(to_python(py, key)?, to_python(py, item)?)?; + } + dict.into_any() + } + }) +} diff --git a/litellm-rust/crates/python-bridge/src/router/config.rs b/litellm-rust/crates/python-bridge/src/router/config.rs new file mode 100644 index 00000000000..5e099b8841c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/config.rs @@ -0,0 +1,414 @@ +//! The documents `RouterCoreHandle` reads: the router's static and dynamic config, a call's +//! request and an `apply` change set. A config key router-core does not read declines the +//! handle, so Python keeps the router on its own core. + +use std::collections::BTreeMap; + +use litellm_router_core::select::PassThrough; +use litellm_router_core::{ + ChangeSet, ConfigError, DeploymentSpec, OutcomeCallbacks, RouterConfig, RouterDynamicConfig, + RouterRequest, + config::{Fallbacks, RetryPolicy}, + host::UncountedRetries, + mock::MockFailure, + registry::IndexOverrides, +}; +use serde::{Deserialize, de::DeserializeOwned}; +use serde_json::{Map, Value}; + +/// The `RouterConfig` keys router-core reads. +const STATIC_KEYS: &[&str] = &[ + "model_list", + "routing_strategy", + "routing_groups", + "routing_group_members", + "plugins", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_policy", + "model_group_retry_policy", + "retry_after", + "timeout", + "request_timeout", + "drop_params", + "enable_pre_call_checks", + "routing_plugins", + "callback_filters", + "enable_health_check_routing", + "health_check_staleness_threshold", + "background_health_check_model_groups", + "enable_weighted_failover", + "fallback_on_unsupported_params", + "provider_default_deployment_ids", +]; + +/// The `RouterDynamicConfig` keys router-core reads. +const DYNAMIC_KEYS: &[&str] = &[ + "version", + "model_alias_map", + "model_group_alias", + "fallbacks", + "default_fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "allowed_fails_is_default", + "fallback_keys", + "default_deployment", + "default_deployment_raises", +]; + +/// Router arguments that do not change routing, accepted in either config and dropped. +const IGNORED_KEYS: &[&str] = &[ + "set_verbose", + "debug_level", + "alerting_config", + "client_ttl", +]; + +/// Fallback lists Python treats as unset when empty (`context_window_fallbacks or None`). +const EMPTY_MEANS_UNSET: &[&str] = &["context_window_fallbacks", "content_policy_fallbacks"]; + +/// Which document a failure is about. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)] +pub(crate) enum Document { + #[strum(serialize = "router config")] + Static, + #[strum(serialize = "dynamic router config")] + Dynamic, + #[strum(serialize = "router request")] + Request, + #[strum(serialize = "change set")] + Changes, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Error { + #[error("the {document} must be a JSON object")] + NotObject { document: Document }, + #[error("router-core does not read `{key}` in the {document}")] + UnknownKey { document: Document, key: String }, + #[error("the {document} does not have the shape router-core reads")] + Shape { + document: Document, + #[source] + source: serde_json::Error, + }, + #[error("the request names no model")] + MissingModel, + #[error("router-core rejected the router config")] + Core(#[from] ConfigError), +} + +impl Error { + /// Whether router-core cannot serve this config, so Python keeps its own core. + pub(crate) fn declined(&self) -> bool { + match self { + Self::UnknownKey { .. } | Self::Core(_) => true, + Self::Shape { document, .. } => { + matches!(document, Document::Static | Document::Dynamic) + } + Self::NotObject { .. } | Self::MissingModel => false, + } + } +} + +fn object(document: Document, value: Value) -> Result, Error> { + match value { + Value::Object(object) => Ok(object), + _ => Err(Error::NotObject { document }), + } +} + +fn shaped(document: Document, object: Map) -> Result { + serde_json::from_value(Value::Object(object)) + .map_err(|source| Error::Shape { document, source }) +} + +/// Drop unset and ignored keys, then reject any key outside `known`. +fn routing_keys( + document: Document, + value: Value, + known: &[&str], +) -> Result, Error> { + let mut object = object(document, value)?; + object.retain(|key, value| { + let empty_unset = EMPTY_MEANS_UNSET.contains(&key.as_str()) + && value.as_array().is_some_and(Vec::is_empty); + !(value.is_null() || empty_unset || IGNORED_KEYS.contains(&key.as_str())) + }); + if let Some(key) = object.keys().find(|key| !known.contains(&key.as_str())) { + return Err(Error::UnknownKey { + document, + key: key.clone(), + }); + } + Ok(object) +} + +pub(crate) fn parse_router_config(value: Value) -> Result { + let object = routing_keys(Document::Static, value, STATIC_KEYS)?; + shaped(Document::Static, object) +} + +pub(crate) fn parse_dynamic_config(value: Value) -> Result { + let object = routing_keys(Document::Dynamic, value, DYNAMIC_KEYS)?; + shaped(Document::Dynamic, object) +} + +/// A call's request as Python spells it. +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestSpec { + model: Option, + team_id: Option, + listing_team_id: Option, + proxy_admin: Option, + allowed_models: Option>, + admin_model: Option, + specific_deployment: Option, + num_retries: Option, + uncounted_retries: Option, + disable_fallbacks: Option, + router_weights: Option, + pinned_deployment_id: Option, + routing_strategy: Option, + filter_tags: Option, + /// `listed` or `healthy`: which pass-through pick the request makes. + pass_through: Option, + target_order: Option, + excluded_ids: Option>, + retry_skipped_ids: Option>, + pre_routes: Option, + has_image: Option, + endpoint: Option, + web_search: Option, + route_handler: Option, + decisions_media: Option, + batch_create: Option, + batch_upload: Option, + embedding_batch: Option, + wildcard_request: Option, + unrouted: Option, + passthrough_on_no_deployment: Option, + skip_weighted_failover: Option, + same_group_fallbacks_only: Option, + model_group_retry_policy: Option>, + retry_policy_raises: Option>, + #[serde(default, deserialize_with = "present")] + fallbacks: Option>, + #[serde(default, deserialize_with = "present")] + context_window_fallbacks: Option>, + #[serde(default, deserialize_with = "present")] + content_policy_fallbacks: Option>, + max_fallbacks: Option, + mock_failure: Option, + cooldown_time: Option, + health_check_probe: Option, + outcome_callbacks: Option, + body: Option, +} + +/// A key the request passed, even as `None`: `Some(None)`. +fn present<'de, D, T>(deserializer: D) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, + T: Deserialize<'de>, +{ + Option::::deserialize(deserializer).map(Some) +} + +/// Read a request; `model` replaces the request's own when given. +pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result { + let spec: RequestSpec = shaped(Document::Request, object(Document::Request, value)?)?; + let model = model + .map(str::to_owned) + .or(spec.model) + .ok_or(Error::MissingModel)?; + Ok(RouterRequest { + model, + body: spec.body.unwrap_or_default(), + team_id: spec.team_id, + listing_team_id: spec.listing_team_id, + proxy_admin: spec.proxy_admin.unwrap_or(false), + allowed_models: spec.allowed_models.unwrap_or_default(), + admin_model: spec.admin_model, + specific_deployment: spec.specific_deployment.unwrap_or(false), + num_retries: spec.num_retries, + uncounted_retries: spec.uncounted_retries, + disable_fallbacks: spec.disable_fallbacks.unwrap_or(false), + router_weights: spec.router_weights, + pinned_deployment_id: spec.pinned_deployment_id, + routing_strategy: spec.routing_strategy, + filter_tags: spec.filter_tags.unwrap_or(false), + pass_through: match spec.pass_through.as_deref() { + None => None, + Some("healthy") => Some(PassThrough::Healthy), + Some(_) => Some(PassThrough::Listed), + }, + target_order: spec.target_order, + excluded_ids: spec.excluded_ids.unwrap_or_default(), + retry_skipped_ids: spec.retry_skipped_ids.unwrap_or_default(), + pre_routes: spec.pre_routes.unwrap_or(false), + has_image: spec.has_image.unwrap_or(false), + endpoint: spec.endpoint, + web_search: spec.web_search.unwrap_or(false), + route_handler: spec.route_handler, + decisions_media: spec.decisions_media.unwrap_or(false), + batch_create: spec.batch_create.unwrap_or(false), + batch_upload: spec.batch_upload.unwrap_or(false), + embedding_batch: spec.embedding_batch.unwrap_or(false), + wildcard_request: spec.wildcard_request, + unrouted: spec.unrouted.unwrap_or(false), + passthrough_on_no_deployment: spec.passthrough_on_no_deployment.unwrap_or(false), + skip_weighted_failover: spec.skip_weighted_failover.unwrap_or(false), + same_group_fallbacks_only: spec.same_group_fallbacks_only.unwrap_or(false), + model_group_retry_policy: spec.model_group_retry_policy, + retry_policy_raises: spec.retry_policy_raises.unwrap_or_default(), + fallbacks: spec.fallbacks, + context_window_fallbacks: spec.context_window_fallbacks, + content_policy_fallbacks: spec.content_policy_fallbacks, + max_fallbacks: spec.max_fallbacks, + mock_failure: spec.mock_failure, + cooldown_time: spec.cooldown_time, + health_check_probe: spec.health_check_probe.unwrap_or(false), + outcome_callbacks: spec.outcome_callbacks.unwrap_or_default(), + payload: None, + }) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct ChangesSpec { + /// Every deployment, replacing the current list before the upserts and deletes. + #[serde(default)] + replace: Option>, + /// Deployments added at the end, beside any that share their id. + #[serde(default)] + append: Vec, + #[serde(default)] + upsert: Vec, + #[serde(default)] + delete: Vec, + /// A new dynamic document (the same keys `RouterCoreHandle(static, dynamic)` takes). + #[serde(default)] + settings: Option, + /// The registry indexes the host wrote (`IndexOverrides`); `{}` drops them. + #[serde(default)] + index: Option, + /// The deployment at the id map's position for this id, removed with every index shifted. + #[serde(default)] + remove: Option, + /// A write to the indexes the host hands out (`HostIndexes`), field by field. + #[serde(default)] + host_index: Option, +} + +pub(crate) fn parse_changes(value: Value) -> Result { + let spec: ChangesSpec = shaped(Document::Changes, object(Document::Changes, value)?)?; + Ok(ChangeSet { + replace: spec.replace, + appends: spec.append, + upserts: spec.upsert, + deletes: spec.delete, + settings: spec.settings.map(parse_dynamic_config).transpose()?, + index: spec.index, + remove: spec.remove, + host_index: spec.host_index, + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::unknown_key(json!({"routing_strategy_args": {"ttl": 1}}), "routing_strategy_args")] + #[case::dynamic_key_in_static(json!({"num_retries": 2}), "num_retries")] + fn an_unread_static_key_declines_naming_the_key(#[case] json: Value, #[case] key: &str) { + let error = parse_router_config(json).unwrap_err(); + assert!(error.declined()); + assert!(matches!(&error, Error::UnknownKey { key: named, .. } if named == key)); + } + + #[rstest] + #[case::null_unknown(json!({"routing_strategy_args": null}))] + #[case::ignored(json!({"set_verbose": true, "debug_level": "INFO"}))] + #[case::pre_call_checks_off(json!({"enable_pre_call_checks": false}))] + #[case::pre_call_checks_on(json!({"enable_pre_call_checks": true}))] + fn unset_and_ignored_keys_are_accepted(#[case] json: Value) { + parse_router_config(json).unwrap(); + } + + #[rstest] + #[case::not_object(json!([]), false)] + #[case::bad_shape(json!({"retry_after": "soon"}), true)] + fn a_config_that_is_not_an_object_is_a_value_error_and_a_wrong_shape_declines( + #[case] json: Value, + #[case] declined: bool, + ) { + assert_eq!(parse_router_config(json).unwrap_err().declined(), declined); + } + + #[rstest] + fn python_router_defaults_read_as_unset() { + let dynamic = parse_dynamic_config(json!({ + "num_retries": null, "fallbacks": [], "context_window_fallbacks": [], + "content_policy_fallbacks": [], "model_group_alias": {"gpt": "gpt-4"}, + })) + .unwrap(); + assert_eq!(dynamic.num_retries, None); + assert!(dynamic.context_window_fallbacks.is_none()); + assert!(dynamic.content_policy_fallbacks.is_none()); + assert_eq!(dynamic.model_alias_map["gpt"], "gpt-4"); + } + + #[rstest] + #[case::own_model(json!({"model": "gpt", "num_retries": 1}), None, "gpt")] + #[case::override_model(json!({"model": "gpt"}), Some("other"), "other")] + fn a_request_reads_its_model( + #[case] json: Value, + #[case] model: Option<&str>, + #[case] expected: &str, + ) { + assert_eq!(parse_request(json, model).unwrap().model, expected); + } + + #[rstest] + #[case::negative("negative", UncountedRetries::Negative)] + #[case::not_positive("not_positive", UncountedRetries::NotPositive)] + #[case::raises("raises", UncountedRetries::Raises)] + fn a_request_reads_a_num_retries_that_is_no_count( + #[case] spelled: &str, + #[case] expected: UncountedRetries, + ) { + let request = + parse_request(json!({"model": "gpt", "uncounted_retries": spelled}), None).unwrap(); + assert_eq!(request.uncounted_retries, Some(expected)); + assert_eq!(request.num_retries, None); + } + + #[rstest] + #[case::unknown_field(json!({"model": "gpt", "temperature": 1}))] + #[case::no_model(json!({"team_id": "t"}))] + fn a_bad_request_is_a_value_error(#[case] json: Value) { + assert!(!parse_request(json, None).unwrap_err().declined()); + } + + #[rstest] + fn a_change_set_reads_upserts_and_deletes() { + let changes = parse_changes(json!({ + "upsert": [{"model_name": "m", "litellm_params": {"model": "openai/x"}}], + "delete": ["old"], + })) + .unwrap(); + assert_eq!(changes.upserts.len(), 1); + assert_eq!(changes.deletes, ["old"]); + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs new file mode 100644 index 00000000000..04d19e70edc --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -0,0 +1,1699 @@ +//! The host behind `RouterCoreHandle`. Rust never awaits Python: each host operation the call +//! loop needs becomes a [`HostCall`] parked in the call's [`Exchange`], and the call future stays +//! pending until Python runs it and hands the [`HostCallResponse`] back through [`Driver::resume`]. +//! +//! The driver polls the call future with a no-op waker while the call's exchange is bound to the +//! current thread, so a host future finds its exchange without the router passing it through. +//! +//! The routing cache is a host operation too: [`StepCache`] parks every [`RouterCache`] op as a +//! `cache` step, and Python runs it against the router's own `DualCache` (memory plus Redis), so +//! Python and Rust read and write the same cooldowns, fail counters and usage. + +use std::{ + cell::RefCell, + future::{Future, poll_fn}, + pin::Pin, + sync::{Arc, Mutex, MutexGuard, PoisonError}, + task::{Context, Poll, Waker}, + time::Duration, +}; + +use async_trait::async_trait; +use litellm_router_core::{ + Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Deployment, + Error, Host, Output, Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, + RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, SelectStage, + ThreadRandomness, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, + precall::PreCallFacts, +}; +use pyo3::{PyResult, Python, types::PyAnyMethods}; +use serde_json::{Map, Value}; + +/// The `kind` of a step dict. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] +pub(crate) enum StepKind { + #[strum(serialize = "attempt")] + Attempt, + #[strum(serialize = "sleep")] + Sleep, + #[strum(serialize = "allow_fallback")] + AllowFallback, + #[strum(serialize = "refuse_fallbacks")] + RefuseFallbacks, + #[strum(serialize = "select")] + Select, + #[strum(serialize = "pre_route")] + PreRoute, + #[strum(serialize = "cache")] + Cache, + #[strum(serialize = "secrets")] + Secrets, + #[strum(serialize = "validate")] + Validate, + #[strum(serialize = "provider")] + Provider, + #[strum(serialize = "added")] + Added, + #[strum(serialize = "model_info")] + ModelInfo, + #[strum(serialize = "done")] + Done, + #[strum(serialize = "failed")] + Failed, +} + +/// Which selection stage Python is asked to run. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum Stage { + /// Pre-call checks and sub-routers: return the deployments that remain. + Plugins, + /// A strategy router-core does not implement: return the candidates with the pick first. + /// When candidates share an id, `records` holds each candidate's deployment, in order, and the + /// host answers with the pick's position, so a pick splits over every entry as Python's does. + Strategy { + strategy: String, + records: Vec, + }, + /// Tag routing: return the deployments that remain. + Tags, + /// `enable_pre_call_checks`: return the facts the filter decides by. + PreCall, + /// `_deployment_declares_no_vision`: return the candidates that declare no vision support. + TextOnly, + /// `_deployment_serves_route`: return the candidates that cannot serve `handler`'s call. + RouteIncapable { handler: String }, + /// `_decisions_deployment_refuses_media`: return the candidates that refuse the call's media. + MediaRefusing, + /// `deployments_admitting_embedding_batch`: return the candidates the embedding pick keeps. + EmbeddingBatch, + /// The override strategy's pre-call check on a deployment named by id. + OverrideCheck(String), + /// The routing plugins: run them when `run`, then return the candidate models they left. + RoutingPlugins { run: bool }, + /// The host's `CustomLogger` deployment filters: return the ids that remain, or a failure. + /// They read the ids this request already failed on, as Python's request kwargs carry them. + CallbackFilters { + excluded_ids: Vec, + skipped_ids: Vec, + target_order: Option, + }, + /// `f"{provider}/{model}"` for a name no wildcard pattern matches as spelled: return it as + /// the one entry of `models`. + ProviderRequest, +} + +/// The `op` of a `cache` step: one per [`RouterCache`] method. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] +pub(crate) enum CacheOp { + /// `async_get_cache(key)`: answered with the stored JSON value, or `None`. + #[strum(serialize = "get")] + Get, + /// `async_set_cache(key, value, ttl)`: answered with `None`. + #[strum(serialize = "set")] + Set, + /// `async_increment(key, value, ttl)`: answered with the new number. The TTL applies when + /// the key is created. + #[strum(serialize = "increment")] + Increment, + /// `async_increment_usage(keys, value, ttl)`: `value` is one amount per key, added as + /// `_increment_deployment_usage` adds them. Answered with `None`. + #[strum(serialize = "increment_usage")] + IncrementUsage, + /// `async_batch_get_cache(keys)`: answered with one JSON value or `None` per key, in order. + #[strum(serialize = "batch_get")] + BatchGet, + /// `async_delete_cache(key)`: answered with `None`. + #[strum(serialize = "delete")] + Delete, +} + +/// Python's outcome of a `cache` step. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum CacheOutcome { + /// `get`: the stored value (`None` when missing). + Value(Option), + /// `batch_get`: one entry per key. + Values(Vec>), + /// `increment`: the counter after the increment. + Number(f64), + /// `set`, `delete` and `increment_usage`. + Stored, +} + +impl CacheOutcome { + /// The op this outcome answers (`Stored` answers both `set` and `delete`). + fn answers(&self, op: CacheOp) -> bool { + matches!( + (self, op), + (Self::Value(_), CacheOp::Get) + | (Self::Values(_), CacheOp::BatchGet) + | (Self::Number(_), CacheOp::Increment) + | ( + Self::Stored, + CacheOp::Set | CacheOp::Delete | CacheOp::IncrementUsage + ) + ) + } +} + +/// A fallback chain's dict targets, which can carry an api key. +#[derive(Clone, PartialEq, veil::Redact)] +pub(crate) struct DictTargets(#[redact(with = "[REDACTED]")] pub(crate) Vec>); + +/// What a call is blocked on, for Python to run. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum HostCall { + Attempt { + deployment_id: String, + /// The picked deployment as configured: `model_name`, `litellm_params`, `model_info` + /// (with its `id`). The host builds the request from it; it never looks the deployment up + /// in a Python router. + deployment: Value, + /// The name the attempt resolved: the requested name or the fallback target. + model: String, + model_group: String, + /// The group a pre-routing hook selected for `model`, when it differs. + routed: Option, + /// The `order` level of a same-group order fallback. + target_order: Option, + /// Every deployment `model` resolved to for this caller, the picked one included. + candidate_ids: Vec, + retry: u32, + fallback_depth: u32, + fallback_level: u32, + timeout: Option, + /// Request params the fallback walk's dict targets merged in. + params: Option>, + /// The walk level whose request metadata the attempt carries. + scope: u32, + max_retries: u32, + model_group_size: usize, + }, + Sleep { + seconds: f64, + }, + AllowFallback { + target: String, + }, + RefuseFallbacks { + /// The walk level whose chain the targets are. + level: u32, + targets: DictTargets, + }, + Select { + stage: Stage, + model_group: String, + /// The name the request resolved from, which strategies key their state by. + name: String, + /// Whose selector a strategy reads: `default`, `group` (with `group`) or `override`. + scope: &'static str, + group: Option, + /// The deployments the name resolved to, before health filters (the tag stage reads them). + pool: Vec, + candidates: Vec, + }, + /// Pre-routing before a pick resolves `model`: answered with the model group to resolve + /// instead (a string), a custom strategy's deployment document (a mapping), or `None`. + PreRoute { + model: String, + fallback_depth: u32, + retry: u32, + }, + /// One [`RouterCache`] op against the host's cache. `key` is set for every op but + /// `batch_get`, which sets `keys`; `value` is set for `set` (the JSON to store) and + /// `increment` (the amount); `ttl` is in seconds. + Cache { + op: CacheOp, + key: Option, + keys: Option>, + value: Option, + ttl: Option, + /// The op reads or writes this worker's in-memory tier only (`local_only=True`). + local: bool, + }, + /// The deployment build's `os.environ/` references: answered with one value per ref. + Secrets { + index: usize, + refs: Vec, + }, + /// The deployment build's `validate`, `provider` and `added` steps for one deployment. + Build { + kind: StepKind, + path: BuildPath, + index: usize, + entry: Value, + declared_id: Option, + duplicate_id: bool, + /// The provider the `provider` step resolved (`added` only). + provider: Option, + }, + /// One lookup of a model-info query ([`ModelInfoLookup`], as its `lookup` field and + /// arguments): answered as the lookup documents. + ModelInfo { + lookup: Value, + }, +} + +impl HostCall { + pub(crate) fn kind(&self) -> StepKind { + match self { + Self::Attempt { .. } => StepKind::Attempt, + Self::Sleep { .. } => StepKind::Sleep, + Self::AllowFallback { .. } => StepKind::AllowFallback, + Self::RefuseFallbacks { .. } => StepKind::RefuseFallbacks, + Self::Select { .. } => StepKind::Select, + Self::PreRoute { .. } => StepKind::PreRoute, + Self::Cache { .. } => StepKind::Cache, + Self::Secrets { .. } => StepKind::Secrets, + Self::Build { kind, .. } => *kind, + Self::ModelInfo { .. } => StepKind::ModelInfo, + } + } +} + +/// Python's response to a [`HostCall`]. +#[derive(Debug)] +pub(crate) enum HostCallResponse { + /// The attempt's total tokens, or its failure. + Attempt(Result), + /// A successful attempt whose response the provider blocked, or a stream whose outcome the + /// host reports when it ends (`deferred`), with its tokens. + Served { + total_tokens: u64, + blocked: Option, + deferred: bool, + tracks_usage: Option, + }, + Slept, + AllowFallback(bool), + /// Whether the chain a `refuse_fallbacks` step named is refused. + RefuseFallbacks(bool), + /// The candidate ids that remain, pick first. + Select(Vec), + /// A strategy's pick among candidates sharing an id: positions in the candidates, pick first. + Positions(Vec), + /// The model group to resolve instead, a custom strategy's deployment, or `None`. + PreRoute(Option), + /// The facts a `pre_call` select step asked for. + PreCall(PreCallFacts), + /// The candidate models a `routing_plugins` select step reported. + Models(Option>), + /// The ids a `callback_filters` select step left with the one a filter prefers, or the error + /// a filter raised. + Filtered(Result<(Vec, Option), AttemptError>), + /// What a pick's select or pre-route step, or a fallback check, raised on the host instead of + /// answering: the walk takes it as the try's failure. + Raised(AttemptError), + Cache(CacheOutcome), + /// One value per ref of a `secrets` step. + Secrets(Vec), + Validated(Built), + Provided(Provided), + Added(Built), + /// A `model_info` lookup's response, as the host spelled it. + ModelInfo(Value), +} + +impl HostCallResponse { + fn kind(&self) -> StepKind { + match self { + Self::Attempt(_) | Self::Served { .. } => StepKind::Attempt, + // Answers whichever step raised it (`check`). + Self::Raised(_) => StepKind::Select, + Self::Slept => StepKind::Sleep, + Self::AllowFallback(_) => StepKind::AllowFallback, + Self::RefuseFallbacks(_) => StepKind::RefuseFallbacks, + Self::Select(_) + | Self::Positions(_) + | Self::PreCall(_) + | Self::Models(_) + | Self::Filtered(_) => StepKind::Select, + Self::PreRoute(_) => StepKind::PreRoute, + Self::Cache(_) => StepKind::Cache, + Self::Secrets(_) => StepKind::Secrets, + Self::Validated(_) => StepKind::Validate, + Self::Provided(_) => StepKind::Provider, + Self::Added(_) => StepKind::Added, + Self::ModelInfo(_) => StepKind::ModelInfo, + } + } + + fn check(&self, call: &HostCall) -> Result<(), DriveError> { + if matches!(self, Self::Raised(_)) + && matches!( + call.kind(), + StepKind::Select | StepKind::PreRoute | StepKind::AllowFallback + ) + { + return Ok(()); + } + if self.kind() != call.kind() { + return Err(DriveError::WrongAnswer { + expected: call.kind(), + got: self.kind(), + }); + } + if let ( + Self::Select(ids) | Self::Filtered(Ok((ids, _))), + HostCall::Select { candidates, .. }, + ) = (self, call) + && let Some(id) = { + let known: std::collections::HashSet<&str> = + candidates.iter().map(String::as_str).collect(); + ids.iter().find(|id| !known.contains(id.as_str())) + } + { + return Err(DriveError::UnknownCandidate { id: id.clone() }); + } + if let (Self::Positions(positions), HostCall::Select { candidates, .. }) = (self, call) + && let Some(position) = positions.iter().find(|p| **p >= candidates.len()) + { + return Err(DriveError::UnknownCandidate { + id: format!("#{position}"), + }); + } + if let (Self::Secrets(values), HostCall::Secrets { refs, .. }) = (self, call) + && values.len() != refs.len() + { + return Err(DriveError::SecretsLength { + expected: refs.len(), + got: values.len(), + }); + } + if let (Self::Cache(outcome), HostCall::Cache { op, keys, .. }) = (self, call) { + if !outcome.answers(*op) { + return Err(DriveError::WrongCacheOutcome { op: *op }); + } + if let (CacheOutcome::Values(values), Some(keys)) = (outcome, keys) + && values.len() != keys.len() + { + return Err(DriveError::BatchLength { + expected: keys.len(), + got: values.len(), + }); + } + } + Ok(()) + } +} + +/// How a driven future ended well. +#[derive(Debug)] +pub(crate) enum Finish { + /// `start`: the call's response. + Call(RouterResponse), + /// `select`: the deployment picked and the model group it serves. + Selected { + deployment_id: String, + model_group: String, + /// The deployment as served: a priority group's member carries its priority as `order`. + deployment: Value, + }, + /// `healthy`: the deployments a pick would choose among, as served. + Healthy { + /// A deployment named by its id, which the listing returns alone rather than as a list. + single: bool, + deployment_ids: Vec, + deployments: Vec, + }, + /// `report_failure`: the cooldown applied, if the deployment was cooled down. + Reported { cooldown: Option }, + /// `report_success` / `count_success`: the success counter's key, if one was counted. + Counted { key: Option }, + /// `build`: the deployments it installed, in order. + Built(Vec), + /// `router_model_info`: the merged model info (`None` when the lookup answered `None`). + RouterModelInfo(Option), + /// `deployment_model_info`: the merged model info, if any. + DeploymentModelInfo(Option>), + /// A model-info query that ends with one JSON answer (`model_group_info`, + /// `configured_token_limits`, `model_listing_info`). + Value(Value), +} + +/// One step of a call: what it is blocked on, or how it ended. +#[derive(Debug)] +pub(crate) enum Step { + HostCall(HostCall), + Done(Finish), + Failed(Error), +} + +/// A misuse of the step protocol by the caller. +#[derive(Debug, thiserror::Error)] +pub(crate) enum DriveError { + #[error("the router call already finished")] + Finished, + #[error("the router call has not made a host call yet, so the outcome must be None")] + UnexpectedResponse, + #[error("the router call is waiting for an outcome of its `{expected}` step")] + MissingResponse { expected: StepKind }, + #[error("the router call is waiting for an outcome of its `{expected}` step, not `{got}`")] + WrongAnswer { expected: StepKind, got: StepKind }, + #[error("deployment `{id}` was not one of the select step's candidates")] + UnknownCandidate { id: String }, + #[error("the outcome does not answer a `{op}` cache step")] + WrongCacheOutcome { op: CacheOp }, + #[error("a `batch_get` cache step for {expected} keys was answered with {got} values")] + BatchLength { expected: usize, got: usize }, + #[error("a `secrets` step for {expected} references was answered with {got} values")] + SecretsLength { expected: usize, got: usize }, + #[error("the router call is pending on something other than a step")] + Stalled, +} + +impl DriveError { + /// Whether the caller passed a bad outcome, rather than driving a call that cannot go on. + pub(crate) fn is_request(&self) -> bool { + match self { + Self::UnexpectedResponse + | Self::MissingResponse { .. } + | Self::WrongAnswer { .. } + | Self::UnknownCandidate { .. } + | Self::WrongCacheOutcome { .. } + | Self::BatchLength { .. } + | Self::SecretsLength { .. } => true, + Self::Finished | Self::Stalled => false, + } + } +} + +/// The host call a router call is blocked on and the response that releases it. +#[derive(Default)] +struct Exchange { + call: Option, + response: Option, + /// Events the call reported since the last step, handed to Python with the next one. + events: Vec, +} + +type Shared = Arc>; + +fn lock(exchange: &Mutex) -> MutexGuard<'_, Exchange> { + exchange.lock().unwrap_or_else(PoisonError::into_inner) +} + +thread_local! { + static CURRENT: RefCell> = const { RefCell::new(None) }; +} + +/// Binds an exchange to this thread for one poll and restores the previous one after. +struct Bound(Option); + +impl Bound { + fn new(exchange: &Shared) -> Self { + Self(CURRENT.replace(Some(exchange.clone()))) + } +} + +impl Drop for Bound { + fn drop(&mut self) { + CURRENT.set(self.0.take()); + } +} + +/// Park the host call `call` in the bound exchange and wait for its response. +/// `response`, or the failure a step answered with instead ([`HostCallResponse::Raised`]). +fn answered(response: HostCallResponse) -> Result { + match response { + HostCallResponse::Raised(failure) => Err(failure), + response => Ok(response), + } +} + +async fn call_host(call: HostCall) -> HostCallResponse { + let mut call = Some(call); + poll_fn(move |_| { + CURRENT.with_borrow(|current| { + let Some(exchange) = current else { + return Poll::Pending; + }; + let mut exchange = lock(exchange); + if let Some(call) = call.take() { + exchange.call = Some(call); + return Poll::Pending; + } + exchange.response.take().map_or(Poll::Pending, Poll::Ready) + }) + }) + .await +} + +const CHECKED: &str = "the driver checks each response against its host call"; + +fn ttl_seconds(ttl: Option) -> Option { + ttl.map(|ttl| ttl.as_secs_f64()) +} + +/// An increment amount as Python's counters store it: an int when it is whole. +fn amount_value(amount: f64) -> Value { + if amount.fract() == 0.0 && amount.abs() < i64::MAX as f64 { + Value::from(amount as i64) + } else { + Value::from(amount) + } +} + +/// The host's cache: every op is a `cache` step Python runs against the router's `DualCache`. +pub(crate) struct StepCache; + +async fn cache_step( + op: CacheOp, + key: Option<&str>, + keys: Option<&[String]>, + value: Option, + ttl: Option, +) -> CacheOutcome { + call_cache(op, key, keys, value, ttl, false).await +} + +async fn call_cache( + op: CacheOp, + key: Option<&str>, + keys: Option<&[String]>, + value: Option, + ttl: Option, + local: bool, +) -> CacheOutcome { + let call = HostCall::Cache { + op, + key: key.map(str::to_owned), + keys: keys.map(<[String]>::to_vec), + value, + ttl: ttl_seconds(ttl), + local, + }; + let HostCallResponse::Cache(outcome) = call_host(call).await else { + unreachable!("{CHECKED}"); + }; + outcome +} + +#[async_trait] +impl RouterCache for StepCache { + async fn async_get_cache(&self, key: &str) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = + cache_step(CacheOp::Get, Some(key), None, None, None).await + else { + unreachable!("{CHECKED}"); + }; + Ok(value) + } + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), litellm_cache::Error> { + cache_step(CacheOp::Set, Some(key), None, Some(value), ttl).await; + Ok(()) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + let value = Some(amount_value(amount)); + let CacheOutcome::Number(next) = + cache_step(CacheOp::Increment, Some(key), None, value, ttl).await + else { + unreachable!("{CHECKED}"); + }; + Ok(next) + } + + async fn async_batch_get_cache( + &self, + keys: &[String], + ) -> Result>, litellm_cache::Error> { + let CacheOutcome::Values(values) = + cache_step(CacheOp::BatchGet, None, Some(keys), None, None).await + else { + unreachable!("{CHECKED}"); + }; + Ok(values) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), litellm_cache::Error> { + cache_step(CacheOp::Delete, Some(key), None, None, None).await; + Ok(()) + } + + async fn async_increment_usage( + &self, + increments: &[(String, f64)], + ttl: Option, + ) -> Result<(), litellm_cache::Error> { + let keys: Vec = increments.iter().map(|(key, _)| key.clone()).collect(); + let amounts = increments + .iter() + .map(|(_, amount)| amount_value(*amount)) + .collect(); + let value = Some(Value::Array(amounts)); + cache_step(CacheOp::IncrementUsage, None, Some(&keys), value, ttl).await; + Ok(()) + } + + async fn async_get_local(&self, key: &str) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = + call_cache(CacheOp::Get, Some(key), None, None, None, true).await + else { + unreachable!("{CHECKED}"); + }; + Ok(value) + } + + async fn async_increment_local( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + let value = Some(amount_value(amount)); + let CacheOutcome::Number(next) = + call_cache(CacheOp::Increment, Some(key), None, value, ttl, true).await + else { + unreachable!("{CHECKED}"); + }; + Ok(next) + } +} + +/// A deployment as Python spells a `model_list` entry, with its id in `model_info`. +/// Whether two candidates share an id, which a pick answered by id can't tell apart. +fn shares_ids(candidates: &[Arc]) -> bool { + let mut seen = std::collections::HashSet::new(); + !candidates.iter().all(|d| seen.insert(d.id.as_str())) +} + +/// The candidates a host answer names, in its order. An id names every candidate that carries +/// it, so deployments sharing an id stay together through a filter, as Python's list keeps them. +fn answered_by_id(candidates: &[Arc], ids: &[String]) -> Vec> { + let mut positions: std::collections::HashMap<&str, Vec> = + std::collections::HashMap::new(); + for (position, deployment) in candidates.iter().enumerate() { + positions + .entry(deployment.id.as_str()) + .or_default() + .push(position); + } + let mut kept = Vec::new(); + for id in ids { + if let Some(found) = positions.remove(id.as_str()) { + kept.extend( + found + .into_iter() + .map(|position| candidates[position].clone()), + ); + } + } + kept +} + +pub(crate) fn deployment_record(deployment: &Deployment) -> Value { + deployment.record() +} + +const RANDOMNESS_MODULE: &str = "litellm.rust_bridge.host.randomness"; + +/// The process's `random` module: a Python router draws its shuffle and retry jitter there, so +/// a seeded process picks the same deployments on either core. The bare module functions share +/// one generator, so these calls interleave with every other draw the process makes. +pub(crate) struct PythonRandomness; + +impl Randomness for PythonRandomness { + fn next_f64(&self) -> f64 { + Python::attach(|py| -> PyResult { + py.import(RANDOMNESS_MODULE)? + .call_method0("next_float")? + .extract() + }) + .unwrap_or_else(|_| ThreadRandomness.next_f64()) + } + + fn index_below(&self, len: usize) -> usize { + // `random.choice` is `seq[self._randbelow(len(seq))]`, and so is `randrange(len)`. + Python::attach(|py| -> PyResult { + py.import(RANDOMNESS_MODULE)? + .call_method1("index_below", (len,))? + .extract() + }) + .unwrap_or_else(|_| ThreadRandomness.index_below(len)) + } +} + +/// The [`Host`] every `RouterCoreHandle` router runs on. +pub(crate) struct StepHost; + +#[async_trait] +impl Host for StepHost { + async fn invoke(&self, attempt: &Attempt) -> Result { + let call = HostCall::Attempt { + deployment_id: attempt.deployment.id.clone(), + deployment: deployment_record(&attempt.deployment), + model: attempt.name.clone(), + model_group: attempt.model_group.clone(), + routed: attempt.routed.clone(), + target_order: attempt.target_order, + candidate_ids: attempt.candidate_ids.clone(), + retry: attempt.retry, + fallback_depth: attempt.fallback_depth, + fallback_level: attempt.fallback_level, + timeout: attempt.timeout.map(|timeout| timeout.as_secs_f64()), + params: attempt + .params + .as_ref() + .map(|params| Box::new(Value::Object(Map::clone(params)))), + scope: attempt.scope, + max_retries: attempt.max_retries, + model_group_size: attempt.model_group_size, + }; + let (result, blocked, deferred, tracks_usage) = match call_host(call).await { + HostCallResponse::Attempt(result) => (result, None, false, None), + HostCallResponse::Served { + total_tokens, + blocked, + deferred, + tracks_usage, + } => (Ok(total_tokens), blocked, deferred, tracks_usage), + _ => unreachable!("{CHECKED}"), + }; + result.map(|total_tokens| Output { + value: Payload(Arc::new(())), + total_tokens, + blocked, + deferred, + tracks_usage, + }) + } + + async fn select( + &self, + stage: &SelectStage, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + let (scope, group) = ctx.scope.wire(); + let call = HostCall::Select { + stage: match stage { + SelectStage::Plugins => Stage::Plugins, + SelectStage::Strategy(strategy) => Stage::Strategy { + strategy: strategy.to_string(), + records: if shares_ids(&candidates) { + candidates.iter().map(|d| deployment_record(d)).collect() + } else { + Vec::new() + }, + }, + SelectStage::Tags => Stage::Tags, + SelectStage::TextOnly => Stage::TextOnly, + SelectStage::RouteIncapable { handler } => Stage::RouteIncapable { + handler: handler.clone(), + }, + SelectStage::MediaRefusing => Stage::MediaRefusing, + SelectStage::EmbeddingBatch => Stage::EmbeddingBatch, + SelectStage::OverrideCheck(strategy) => Stage::OverrideCheck(strategy.to_string()), + }, + pool: ctx.pool.iter().map(|d| d.id.clone()).collect(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + match answered(call_host(call).await)? { + HostCallResponse::Select(ids) => Ok(answered_by_id(&candidates, &ids)), + HostCallResponse::Positions(positions) => Ok(positions + .iter() + .map(|position| candidates[*position].clone()) + .collect()), + _ => unreachable!("{CHECKED}"), + } + } + + async fn pre_route(&self, pre: &PreRoute<'_>) -> Result, Error> { + let call = HostCall::PreRoute { + model: pre.name.to_owned(), + fallback_depth: pre.fallback_depth, + retry: pre.retry, + }; + let HostCallResponse::PreRoute(routed) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(routed) + } + + async fn routing_plugin_models( + &self, + ctx: &RoutingContext<'_>, + run: bool, + ) -> Result>, Error> { + let (scope, group) = ctx.scope.wire(); + let call = HostCall::Select { + stage: Stage::RoutingPlugins { run }, + pool: Vec::new(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: Vec::new(), + }; + let HostCallResponse::Models(models) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(models) + } + + async fn provider_request(&self, ctx: &RoutingContext<'_>) -> Result, Error> { + let (scope, group) = ctx.scope.wire(); + let call = HostCall::Select { + stage: Stage::ProviderRequest, + pool: Vec::new(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: Vec::new(), + }; + let HostCallResponse::Models(models) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(models.into_iter().flatten().next()) + } + + async fn pre_call_facts( + &self, + ctx: &RoutingContext<'_>, + candidates: &[Arc], + ) -> Result { + let (scope, group) = ctx.scope.wire(); + let call = HostCall::Select { + stage: Stage::PreCall, + pool: ctx.pool.iter().map(|d| d.id.clone()).collect(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let HostCallResponse::PreCall(facts) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(facts) + } + + async fn callback_filters( + &self, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result { + let (scope, group) = ctx.scope.wire(); + let joined = |own: &[String], walked: &[String]| [own, walked].concat(); + let call = HostCall::Select { + stage: Stage::CallbackFilters { + excluded_ids: joined(&ctx.request.excluded_ids, ctx.excluded_ids), + skipped_ids: joined(&ctx.request.retry_skipped_ids, ctx.skipped_ids), + target_order: ctx.target_order, + }, + pool: Vec::new(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let HostCallResponse::Filtered(result) = call_host(call).await else { + unreachable!("{CHECKED}"); + }; + let (ids, preferred) = result.map_err(Error::Attempt)?; + Ok(CallbackFiltered { + kept: answered_by_id(&candidates, &ids), + preferred, + }) + } + + async fn allow_fallback( + &self, + _request: &RouterRequest, + target: &str, + ) -> Result { + let call = HostCall::AllowFallback { + target: target.to_owned(), + }; + let HostCallResponse::AllowFallback(allowed) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(allowed) + } + + async fn refuses_fallbacks(&self, level: u32, targets: &[Map]) -> bool { + let call = HostCall::RefuseFallbacks { + level, + targets: DictTargets(targets.to_vec()), + }; + let HostCallResponse::RefuseFallbacks(refused) = call_host(call).await else { + unreachable!("{CHECKED}"); + }; + refused + } + + /// A zero sleep is asked too: Python's retry loop awaits `asyncio.sleep(0)` between sibling + /// retries, which yields to the event loop. + async fn sleep(&self, duration: Duration) { + let call = HostCall::Sleep { + seconds: duration.as_secs_f64(), + }; + let HostCallResponse::Slept = call_host(call).await else { + unreachable!("{CHECKED}"); + }; + } + + /// Kept for Python only when it acts on it: an exhausted retry loop is stamped on the + /// error Python raises, a fallback hop tells it which target a failed walk ended on, and a + /// cooldown or a pick is reported to the host's Prometheus and service loggers (a pick also + /// joins the router cache's per-deployment count, a failed walk the router's error log). + fn observe(&self, event: &RouterEvent) { + if !matches!( + event, + RouterEvent::RetriesExhausted { .. } + | RouterEvent::Retry { .. } + | RouterEvent::Fallback { .. } + | RouterEvent::FallingBack { .. } + | RouterEvent::OrderChainFailed { .. } + | RouterEvent::FallbacksFailed { .. } + | RouterEvent::LevelFailed { .. } + | RouterEvent::PickFailed { .. } + | RouterEvent::RetriesOpened { .. } + | RouterEvent::RetryBudget { .. } + | RouterEvent::MockFailed { .. } + | RouterEvent::RetryLogged { .. } + | RouterEvent::RetryPolicyRaised { .. } + | RouterEvent::RetryCountRaised { .. } + | RouterEvent::FallbackDone { .. } + | RouterEvent::Blocked { .. } + | RouterEvent::CooldownSet { .. } + | RouterEvent::Selected { .. } + | RouterEvent::PickStarted { .. } + ) { + return; + } + CURRENT.with_borrow(|current| { + if let Some(exchange) = current { + lock(exchange).events.push(event.clone()); + } + }); + } + + fn cache(&self) -> Arc { + Arc::new(StepCache) + } + + async fn secrets(&self, index: usize, refs: &[SecretRef]) -> Result, Error> { + let HostCallResponse::Secrets(values) = call_host(HostCall::Secrets { + index, + refs: refs.to_vec(), + }) + .await + else { + unreachable!("{CHECKED}"); + }; + Ok(values) + } + + async fn validate(&self, build: &BuildAsk<'_>) -> Result { + let HostCallResponse::Validated(built) = + call_host(build_host_call(StepKind::Validate, build, None)).await + else { + unreachable!("{CHECKED}"); + }; + Ok(built) + } + + async fn provider(&self, build: &BuildAsk<'_>) -> Result { + let HostCallResponse::Provided(provided) = + call_host(build_host_call(StepKind::Provider, build, None)).await + else { + unreachable!("{CHECKED}"); + }; + Ok(provided) + } + + async fn added(&self, build: &BuildAsk<'_>, provider: Option<&str>) -> Result { + let HostCallResponse::Added(built) = + call_host(build_host_call(StepKind::Added, build, provider)).await + else { + unreachable!("{CHECKED}"); + }; + Ok(built) + } + + async fn model_info(&self, lookup: &ModelInfoLookup) -> Result { + let lookup = serde_json::to_value(lookup).unwrap_or(Value::Null); + let HostCallResponse::ModelInfo(response) = call_host(HostCall::ModelInfo { lookup }).await + else { + unreachable!("{CHECKED}"); + }; + Ok(response) + } +} + +fn build_host_call(kind: StepKind, build: &BuildAsk<'_>, provider: Option<&str>) -> HostCall { + HostCall::Build { + kind, + path: build.path, + index: build.index, + entry: Value::Object(build.entry.clone()), + declared_id: build.declared_id.map(str::to_owned), + duplicate_id: build.duplicate_id, + provider: provider.map(str::to_owned), + } +} + +type CallFuture = Pin> + Send>>; + +/// Poll `future` once with `exchange` bound. +fn poll_bound(future: Pin<&mut F>, exchange: &Shared) -> Poll { + let _bound = Bound::new(exchange); + future.poll(&mut Context::from_waker(Waker::noop())) +} + +/// Drives one router call a step at a time. +pub(crate) struct Driver { + future: Option, + exchange: Shared, + pending_call: Option, +} + +impl Driver { + pub(crate) fn new(call: impl Future> + Send + 'static) -> Self { + Self { + future: Some(Box::pin(call)), + exchange: Shared::default(), + pending_call: None, + } + } + + /// The host call the router call is blocked on, if it made one. + pub(crate) fn pending_call(&self) -> Option<&HostCall> { + self.pending_call.as_ref() + } + + /// The events reported since the last call. + pub(crate) fn take_events(&self) -> Vec { + std::mem::take(&mut lock(&self.exchange).events) + } + + /// Hand back the response to the pending host call (`None` on the first call) and run the call to + /// its next step or its end. + pub(crate) fn resume( + &mut self, + response: Option, + ) -> Result { + let Some(future) = self.future.as_mut() else { + return Err(DriveError::Finished); + }; + match (&self.pending_call, response) { + (None, None) => {} + (None, Some(_)) => return Err(DriveError::UnexpectedResponse), + (Some(call), None) => { + return Err(DriveError::MissingResponse { + expected: call.kind(), + }); + } + (Some(call), Some(response)) => { + response.check(call)?; + lock(&self.exchange).response = Some(response); + } + } + self.pending_call = None; + match poll_bound(future.as_mut(), &self.exchange) { + Poll::Ready(result) => { + self.future = None; + Ok(match result { + Ok(finish) => Step::Done(finish), + Err(error) => Step::Failed(error), + }) + } + Poll::Pending => { + let Some(call) = lock(&self.exchange).call.take() else { + self.future = None; + return Err(DriveError::Stalled); + }; + self.pending_call = Some(call.clone()); + Ok(Step::HostCall(call)) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::HashMap, + time::{SystemTime, UNIX_EPOCH}, + }; + + use litellm_router_core::{ + FailureKind, Op, Router, RouterConfig, RouterDynamicConfig, SeededRandomness, + }; + use rstest::{fixture, rstest}; + use serde_json::{Value, json}; + + use super::*; + + fn router(config: Value, dynamic: Value) -> Router { + let config: RouterConfig = serde_json::from_value(config).unwrap(); + let dynamic: RouterDynamicConfig = serde_json::from_value(dynamic).unwrap(); + Router::builder(config, dynamic, Arc::new(StepHost)) + .randomness(Arc::new(SeededRandomness::new(7))) + .build() + .unwrap() + } + + fn deployment(group: &str, id: &str) -> Value { + json!({"model_name": group, "litellm_params": {"model": format!("openai/{id}")}, "model_info": {"id": id}}) + } + + fn start(router: &Router, model: &str) -> Driver { + let router = router.clone(); + let request = RouterRequest { + model: model.to_owned(), + ..RouterRequest::default() + }; + Driver::new(async move { + router + .call(Op::from("acompletion"), request) + .await + .map(Finish::Call) + }) + } + + fn rate_limited() -> HostCallResponse { + HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::RateLimit, Some(429)))) + } + + /// A dict-backed stand-in for Python's `DualCache` that answers `cache` steps and keeps + /// every one it answered. + #[derive(Default)] + struct Store { + entries: HashMap, + calls: Vec, + } + + impl Store { + fn respond(&mut self, call: &HostCall) -> HostCallResponse { + let HostCall::Cache { + op, + key, + keys, + value, + .. + } = call + else { + panic!("not a cache step: {call:?}"); + }; + self.calls.push(call.clone()); + let key = key.clone().unwrap_or_default(); + HostCallResponse::Cache(match op { + CacheOp::Get => CacheOutcome::Value(self.entries.get(&key).cloned()), + CacheOp::Set => { + self.entries.insert(key, value.clone().unwrap()); + CacheOutcome::Stored + } + CacheOp::Increment => { + let current = self.entries.get(&key).and_then(Value::as_f64); + let next = + current.unwrap_or(0.0) + value.as_ref().and_then(Value::as_f64).unwrap(); + self.entries.insert(key, amount_value(next)); + CacheOutcome::Number(next) + } + CacheOp::IncrementUsage => { + let amounts = value.as_ref().and_then(Value::as_array).unwrap(); + for (key, amount) in keys.iter().flatten().zip(amounts) { + let current = self.entries.get(key).and_then(Value::as_f64); + let next = current.unwrap_or(0.0) + amount.as_f64().unwrap(); + self.entries.insert(key.clone(), amount_value(next)); + } + CacheOutcome::Stored + } + CacheOp::BatchGet => CacheOutcome::Values( + keys.iter() + .flatten() + .map(|key| self.entries.get(key).cloned()) + .collect(), + ), + CacheOp::Delete => { + self.entries.remove(&key); + CacheOutcome::Stored + } + }) + } + + /// Resume with `response`, answering every cache step on the way, up to the next other + /// step. + fn resume(&mut self, driver: &mut Driver, response: Option) -> Step { + let mut step = driver.resume(response).unwrap(); + while let Step::HostCall(call @ HostCall::Cache { .. }) = &step { + let response = self.respond(call); + step = driver.resume(Some(response)).unwrap(); + } + step + } + } + + #[fixture] + fn pair() -> Router { + router( + json!({"model_list": [deployment("primary", "a1"), deployment("primary", "a2")]}), + json!({}), + ) + } + + #[rstest] + fn a_first_attempt_success_is_done(pair: Router) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + let Step::HostCall(HostCall::Attempt { + deployment_id, + model_group, + retry: 0, + fallback_depth: 0, + .. + }) = store.resume(&mut driver, None) + else { + panic!("the call should start with an attempt"); + }; + assert_eq!(model_group, "primary"); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(HostCallResponse::Attempt(Ok(12)))) + else { + panic!("a successful attempt should finish the call"); + }; + assert_eq!(response.deployment_id, deployment_id); + assert_eq!(response.output.total_tokens, 12); + assert_eq!( + (response.attempted_retries, response.attempted_fallbacks), + (0, 0) + ); + assert!(matches!(driver.resume(None), Err(DriveError::Finished))); + } + + #[rstest] + fn cooldown_reads_are_cache_steps(pair: Router) { + let mut driver = start(&pair, "primary"); + let step = driver.resume(None).unwrap(); + let Step::HostCall(HostCall::Cache { + op: CacheOp::BatchGet, + key: None, + keys: Some(keys), + value: None, + ttl: None, + local: false, + }) = step + else { + panic!("the call should read cooldowns first, got {step:?}"); + }; + assert_eq!(keys, ["deployment:a1:cooldown", "deployment:a2:cooldown"]); + } + + #[rstest] + #[case::lone_deployment_backs_off(vec![deployment("primary", "a1")], true)] + #[case::sibling_retries_at_once( + vec![deployment("primary", "a1"), deployment("primary", "a2")], + false + )] + fn a_rate_limit_retries_after_the_backoff( + #[case] model_list: Vec, + #[case] backs_off: bool, + ) { + let router = router(json!({"model_list": model_list}), json!({"num_retries": 1})); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + assert_eq!( + store.resume(&mut driver, None).unwrap_kind(), + StepKind::Attempt + ); + let step = store.resume(&mut driver, Some(rate_limited())); + let Step::HostCall(HostCall::Sleep { seconds }) = step else { + panic!("a retry should sleep first, got {step:?}"); + }; + assert_eq!(seconds > 0.0, backs_off, "slept {seconds}s"); + let step = store.resume(&mut driver, Some(HostCallResponse::Slept)); + assert!( + matches!(step, Step::HostCall(HostCall::Attempt { retry: 1, .. })), + "expected the retry, got {step:?}" + ); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(HostCallResponse::Attempt(Ok(1)))) + else { + panic!("the retry should finish the call"); + }; + assert_eq!(response.attempted_retries, 1); + } + + #[rstest] + fn a_cooldown_python_wrote_is_read_back_by_the_next_call() { + let pair_without_retries = router( + json!({"model_list": [deployment("primary", "a1"), deployment("primary", "a2")]}), + json!({"num_retries": 0}), + ); + let mut store = Store::default(); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs_f64(); + store.entries.insert( + "deployment:a1:cooldown".into(), + json!({"exception_received": "slow down", "status_code": "429", "timestamp": now, "cooldown_time": 60}), + ); + for _ in 0..8 { + let mut call = start(&pair_without_retries, "primary"); + let Step::HostCall(HostCall::Attempt { deployment_id, .. }) = + store.resume(&mut call, None) + else { + panic!("the call should attempt the healthy deployment"); + }; + assert_eq!(deployment_id, "a2", "the cooled-down deployment was picked"); + } + } + + /// The host only runs the attempt; the core records its outcome (cooldowns, fail counters, + /// usage) through the host's cache. + #[rstest] + #[case::failure(rate_limited())] + #[case::success(HostCallResponse::Attempt(Ok(7)))] + fn attempt_outcomes_are_recorded_by_the_core(#[case] outcome: HostCallResponse) { + let tracked = |id: &str| { + json!({ + "model_name": "primary", + "litellm_params": {"model": format!("openai/{id}"), "tpm": 100, "cooldown_time": 7}, + "model_info": {"id": id, "allowed_fails": 3}, + }) + }; + let router = router( + json!({"model_list": [tracked("a1"), tracked("a2")]}), + json!({"num_retries": 0}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let step = store.resume(&mut driver, Some(outcome)); + assert!( + matches!(step, Step::Done(_) | Step::Failed(_)), + "got {step:?}" + ); + let writes: Vec<&HostCall> = store + .calls + .iter() + .filter(|call| { + matches!(call, HostCall::Cache { op, .. } if !matches!(op, CacheOp::Get | CacheOp::BatchGet)) + }) + .collect(); + assert!(!writes.is_empty(), "the core wrote no outcome"); + } + + #[rstest] + fn a_failed_group_falls_back_once_python_allows_it() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": ["backup"]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + assert_eq!( + store.resume(&mut driver, None).unwrap_kind(), + StepKind::Attempt + ); + let bad_request = + HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + let step = store.resume(&mut driver, Some(bad_request)); + assert!( + matches!(&step, Step::HostCall(HostCall::AllowFallback { target }) if target == "backup"), + "expected the fallback check, got {step:?}" + ); + let step = store.resume(&mut driver, Some(HostCallResponse::AllowFallback(true))); + assert!( + matches!(&step, Step::HostCall(HostCall::Attempt { deployment_id, fallback_depth: 1, .. }) if deployment_id == "b1"), + "expected the fallback attempt, got {step:?}" + ); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(HostCallResponse::Attempt(Ok(3)))) + else { + panic!("the fallback should finish the call"); + }; + assert_eq!(response.model_group, "backup"); + assert_eq!(response.attempted_fallbacks, 1); + } + + #[rstest] + fn a_refused_fallback_fails_with_the_attempt_error() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": ["backup"]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let bad_request = + HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + store.resume(&mut driver, Some(bad_request)); + let Step::Failed(error) = + store.resume(&mut driver, Some(HostCallResponse::AllowFallback(false))) + else { + panic!("a refused fallback should fail the call"); + }; + assert_eq!( + error.attempt().map(|a| a.kind), + Some(FailureKind::BadRequest) + ); + } + + #[rstest] + #[case::none_while_waiting(None, StepKind::Attempt)] + #[case::wrong_kind(Some(HostCallResponse::Slept), StepKind::Attempt)] + #[case::cache_outcome(Some(HostCallResponse::Cache(CacheOutcome::Stored)), StepKind::Attempt)] + fn a_mismatched_outcome_is_rejected_and_the_step_kept( + pair: Router, + #[case] response: Option, + #[case] expected: StepKind, + ) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + store.resume(&mut driver, None); + let error = driver.resume(response).unwrap_err(); + assert!(error.is_request(), "{error}"); + assert_eq!(driver.pending_call().map(HostCall::kind), Some(expected)); + assert!(matches!( + store.resume(&mut driver, Some(HostCallResponse::Attempt(Ok(1)))), + Step::Done(_) + )); + } + + #[rstest] + #[case::wrong_op(HostCallResponse::Cache(CacheOutcome::Stored))] + #[case::number_for_batch(HostCallResponse::Cache(CacheOutcome::Number(1.0)))] + #[case::short_batch(HostCallResponse::Cache(CacheOutcome::Values(vec![None])))] + #[case::not_a_cache_outcome(HostCallResponse::Slept)] + fn a_mismatched_cache_outcome_is_rejected_and_the_step_kept( + pair: Router, + #[case] response: HostCallResponse, + ) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + let step = driver.resume(None).unwrap(); + let Step::HostCall( + call @ HostCall::Cache { + op: CacheOp::BatchGet, + .. + }, + ) = &step + else { + panic!("expected the cooldown read, got {step:?}"); + }; + let error = driver.resume(Some(response)).unwrap_err(); + assert!(error.is_request(), "{error}"); + assert_eq!(driver.pending_call(), Some(call)); + let response = store.respond(call); + assert_eq!( + store.resume(&mut driver, Some(response)).unwrap_kind(), + StepKind::Attempt + ); + } + + #[rstest] + fn a_strategy_python_runs_asks_for_a_selection_and_keeps_the_pick() { + let router = router( + json!({ + "model_list": [deployment("primary", "a1"), deployment("primary", "a2")], + "routing_strategy": "least-busy", + }), + json!({}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + let step = store.resume(&mut driver, None); + let Step::HostCall(HostCall::Select { + stage: Stage::Strategy { strategy, records }, + candidates, + .. + }) = step + else { + panic!("a host strategy should ask for a selection, got {step:?}"); + }; + assert_eq!(strategy, "least-busy"); + assert_eq!(candidates, ["a1", "a2"]); + assert!(records.is_empty(), "distinct ids send no records"); + let step = store.resume( + &mut driver, + Some(HostCallResponse::Select(vec!["a2".to_owned()])), + ); + assert!( + matches!(&step, Step::HostCall(HostCall::Attempt { deployment_id, .. }) if deployment_id == "a2"), + "expected the picked deployment, got {step:?}" + ); + } + + #[rstest] + fn an_answer_by_id_keeps_its_order_and_every_candidate_sharing_an_id() { + let candidates: Vec> = ["a", "b", "a", "c"] + .into_iter() + .map(|id| { + let spec: litellm_router_core::DeploymentSpec = + serde_json::from_value(deployment("g", id)).unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + }) + .collect(); + let ids = ["c", "a", "x", "a"].map(String::from); + + let kept = answered_by_id(&candidates, &ids); + + let positions: Vec = kept + .iter() + .map(|d| candidates.iter().position(|c| Arc::ptr_eq(c, d)).unwrap()) + .collect(); + assert_eq!(positions, [3, 0, 2]); + } + #[rstest] + fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() { + let mut second = deployment("primary", "a1"); + second["litellm_params"]["api_base"] = json!("https://second.example"); + let router = router( + json!({ + "model_list": [deployment("primary", "a1"), second], + "routing_strategy": "least-busy", + }), + json!({}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + let step = store.resume(&mut driver, None); + let Step::HostCall(HostCall::Select { + stage: Stage::Strategy { records, .. }, + .. + }) = step + else { + panic!("a host strategy should ask for a selection, got {step:?}"); + }; + assert_eq!( + records.len(), + 2, + "each candidate sharing an id sends its record" + ); + let step = store.resume(&mut driver, Some(HostCallResponse::Positions(vec![1]))); + assert!( + matches!(&step, Step::HostCall(HostCall::Attempt { deployment, .. }) + if deployment["litellm_params"]["api_base"] == "https://second.example"), + "expected the second entry, got {step:?}" + ); + } + + #[rstest] + fn a_selection_is_driven_through_cache_steps(pair: Router) { + let request = RouterRequest { + model: "primary".to_owned(), + ..RouterRequest::default() + }; + let mut store = Store::default(); + let mut driver = Driver::new(async move { + pair.select(&Op::from("acompletion"), &request) + .await + .map(|deployment| Finish::Selected { + deployment_id: deployment.id.clone(), + model_group: deployment.model_name.clone(), + deployment: deployment_record(&deployment), + }) + }); + let Step::Done(Finish::Selected { + deployment_id, + model_group, + .. + }) = store.resume(&mut driver, None) + else { + panic!("a selection should finish with the pick"); + }; + assert!(["a1", "a2"].contains(&deployment_id.as_str())); + assert_eq!(model_group, "primary"); + assert!(matches!( + store.calls.as_slice(), + [HostCall::Cache { + op: CacheOp::BatchGet, + .. + }] + )); + } + + #[rstest] + fn dict_targets_keep_their_api_key_out_of_debug() { + let target = json!({"model": "b", "api_key": "credential-value"}); + let Value::Object(target) = target else { + unreachable!("a JSON object"); + }; + let shown = format!("{:?}", DictTargets(vec![target])); + assert!(!shown.contains("credential-value")); + assert!(shown.contains("[REDACTED]")); + } + + #[rstest] + fn a_chain_with_a_dict_target_asks_python_before_any_hop() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": [{"model": "backup", "api_key": "k"}]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let bad_request = + HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + let step = store.resume(&mut driver, Some(bad_request)); + assert!( + matches!(&step, Step::HostCall(HostCall::RefuseFallbacks { level: 0, targets }) if targets.0.len() == 1), + "expected the chain's check, got {step:?}" + ); + let Step::Failed(error) = + store.resume(&mut driver, Some(HostCallResponse::RefuseFallbacks(true))) + else { + panic!("a refused chain should fail the call"); + }; + assert_eq!( + error.attempt().map(|a| a.kind), + Some(FailureKind::BadRequest) + ); + } + + trait StepKindOf { + fn unwrap_kind(self) -> StepKind; + } + + impl StepKindOf for Step { + fn unwrap_kind(self) -> StepKind { + match self { + Step::HostCall(call) => call.kind(), + Step::Done(_) => StepKind::Done, + Step::Failed(_) => StepKind::Failed, + } + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs new file mode 100644 index 00000000000..6b58d63a714 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -0,0 +1,2087 @@ +//! `RouterCoreHandle`: router-core behind a step-driven handle. Python drives each call through +//! [`RouterCall::next`], running every attempt, sleep, fallback check, selection and cache op the +//! call loop asks for and handing the outcome back; Rust never awaits Python. `select` is driven +//! the same way, since reading cooldowns is a cache step. + +mod config; +mod host; +mod retries; +mod stream; + +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::{Duration, SystemTime}, +}; + +use litellm_error::Report; +use litellm_host_python::{from_py_argument, json_loads, to_py}; +use litellm_router_core::lookup::{ListScope, Lookup, Row}; +use litellm_router_core::{ + AttemptError, Blocked, BuildPath, Built, Deployment, DeploymentSpec, Error, FailureKind, + FailureReport, Op, OutcomeCallbacks, PreRouted, Provided, RetryCause, Router, RouterEvent, + SecretAt, SuccessReport, + build::Install, + health::{UsageMetric, current_minute, usage_key}, + registry::{DeploymentFilter, ResolveKind, ResolveScope}, + strategy::routing_context, +}; +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::{PyBool, PyDict, PyList, PyString, PyTuple}, +}; +use serde_json::Value; + +use self::host::{ + CacheOp, CacheOutcome, DriveError, Driver, Finish, HostCall, HostCallResponse, + PythonRandomness, Stage, Step, StepHost, StepKind, +}; +use crate::errors::RustBridgeDeclined; + +/// The `error_kind` of a `failed` step. +#[derive(Clone, Copy, Debug, strum::IntoStaticStr)] +enum ErrorKind { + #[strum(serialize = "config")] + Config, + #[strum(serialize = "unknown_model")] + UnknownModel, + #[strum(serialize = "no_healthy_deployments")] + NoHealthyDeployments, + #[strum(serialize = "attempt")] + Attempt, + #[strum(serialize = "cache")] + Cache, + #[strum(serialize = "retired")] + Retired, + #[strum(serialize = "paused")] + Paused, + #[strum(serialize = "not_pass_through")] + NotPassThrough, + #[strum(serialize = "no_pass_through_deployments")] + NoPassThroughDeployments, + #[strum(serialize = "only_markers")] + OnlyMarkers, + #[strum(serialize = "endpoint_not_supported")] + EndpointNotSupported, + #[strum(serialize = "ptu_reserved")] + PtuReserved, + #[strum(serialize = "context_window_exceeded")] + ContextWindowExceeded, + #[strum(serialize = "unsupported_params")] + UnsupportedParams, + #[strum(serialize = "routing_plugins_left_nothing")] + RoutingPluginsLeftNothing, + #[strum(serialize = "routing_plugins_need_async")] + RoutingPluginsNeedAsync, + #[strum(serialize = "multiple_teams")] + MultipleTeams, + #[strum(serialize = "reserved")] + Reserved, + #[strum(serialize = "invalid_router_weights")] + InvalidRouterWeights, + #[strum(serialize = "non_numeric_weight")] + NonNumericWeight, + #[strum(serialize = "missing_routing_group")] + MissingRoutingGroup, + #[strum(serialize = "default_deployment_raises")] + DefaultDeploymentRaises, + #[strum(serialize = "index_out_of_range")] + IndexOutOfRange, + #[strum(serialize = "model_info")] + ModelInfo, +} + +impl From<&Error> for ErrorKind { + fn from(error: &Error) -> Self { + match error { + Error::Config(_) => Self::Config, + Error::UnknownModel { .. } => Self::UnknownModel, + Error::NoHealthyDeployments { .. } => Self::NoHealthyDeployments, + Error::Attempt(_) => Self::Attempt, + Error::Cache(_) => Self::Cache, + Error::Retired => Self::Retired, + Error::Paused { .. } => Self::Paused, + Error::NotPassThrough { .. } => Self::NotPassThrough, + Error::NoPassThroughDeployments { .. } => Self::NoPassThroughDeployments, + Error::OnlyMarkers { .. } => Self::OnlyMarkers, + Error::EndpointNotSupported { .. } => Self::EndpointNotSupported, + Error::PtuReserved { .. } => Self::PtuReserved, + Error::ContextWindowExceeded { .. } => Self::ContextWindowExceeded, + Error::UnsupportedParams { .. } => Self::UnsupportedParams, + Error::RoutingPluginsLeftNothing { .. } => Self::RoutingPluginsLeftNothing, + Error::RoutingPluginsNeedAsync => Self::RoutingPluginsNeedAsync, + Error::MultipleTeams { .. } => Self::MultipleTeams, + Error::Reserved { .. } => Self::Reserved, + Error::InvalidRouterWeights => Self::InvalidRouterWeights, + Error::NonNumericWeight { .. } => Self::NonNumericWeight, + Error::MissingRoutingGroup { .. } => Self::MissingRoutingGroup, + Error::DefaultDeploymentRaises { .. } => Self::DefaultDeploymentRaises, + Error::IndexOutOfRange { .. } => Self::IndexOutOfRange, + Error::ModelInfo { .. } => Self::ModelInfo, + } + } +} + +fn config_error_to_pyerr(error: config::Error) -> PyErr { + let message = Report::new(&error).to_string(); + if error.declined() { + RustBridgeDeclined::new_err(message) + } else { + PyValueError::new_err(message) + } +} + +fn drive_error_to_pyerr(error: DriveError) -> PyErr { + let message = Report::new(&error).to_string(); + if error.is_request() { + PyValueError::new_err(message) + } else { + PyRuntimeError::new_err(message) + } +} + +/// A document Python passes as a mapping or as its JSON text. +fn document(value: &Bound<'_, PyAny>) -> PyResult { + match value.cast::() { + Ok(text) => { + let parsed = json_loads(value.py(), text.to_str()?.as_bytes())?; + from_py_argument(parsed.bind(value.py())) + } + Err(_) => from_py_argument(value), + } +} + +/// The `OutcomeCallbacks` a reported outcome came through; none is a report no callback makes. +fn outcome_callbacks(value: Option<&Bound<'_, PyAny>>) -> PyResult { + let Some(value) = value else { + return Ok(OutcomeCallbacks::default()); + }; + serde_json::from_value(document(value)?) + .map_err(|error| PyValueError::new_err(error.to_string())) +} + +/// A router-core router whose host operations Python runs. +#[pyclass(frozen)] +pub(crate) struct RouterCoreHandle { + router: Router, +} + +#[pymethods] +impl RouterCoreHandle { + #[new] + fn new(config: &Bound<'_, PyAny>, dynamic: &Bound<'_, PyAny>) -> PyResult { + let (config, dynamic) = (document(config)?, document(dynamic)?); + let build = || { + let config = config::parse_router_config(config)?; + let dynamic = config::parse_dynamic_config(dynamic)?; + Ok(Router::builder(config, dynamic, Arc::new(StepHost)) + .randomness(Arc::new(PythonRandomness)) + .build()?) + }; + build() + .map(|router| Self { router }) + .map_err(config_error_to_pyerr) + } + + /// Start a call; nothing runs until the first `next(None)`. + fn start(&self, op: &str, request: &Bound<'_, PyAny>) -> PyResult { + let op = Op::new(op); + let request = + config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router.call(op, request).await.map(Finish::Call) + })) + } + + /// Pick the deployment `model` routes to now, without calling it. Driven like `start`: its + /// steps are the cache reads and selections the pick needs, and it ends with `done` + /// (`deployment_id`, `model_group`) or `failed`. + fn select(&self, model: &str, request: &Bound<'_, PyAny>) -> PyResult { + let request = config::parse_request(document(request)?, Some(model)) + .map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let deployment = router.select(&Op::from("acompletion"), &request).await?; + Ok(Finish::Selected { + deployment_id: deployment.id.clone(), + model_group: deployment.model_name.clone(), + deployment: host::deployment_record(&deployment), + }) + })) + } + + /// List the deployments a pick for `model` would choose among, driven like `select`. + fn healthy(&self, model: &str, request: &Bound<'_, PyAny>) -> PyResult { + let request = config::parse_request(document(request)?, Some(model)) + .map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let candidates = router.healthy(&Op::from("acompletion"), &request).await?; + // A deployment named by its id is listed alone, as a dict (`_core_async_get_healthy_deployments`). + let single = candidates.kind.is_single() && candidates.healthy.len() == 1; + let deployments = candidates.healthy; + Ok(Finish::Healthy { + single, + deployment_ids: deployments.iter().map(|d| d.id.clone()).collect(), + deployments: deployments + .iter() + .map(|d| host::deployment_record(d)) + .collect(), + }) + })) + } + + /// Apply `{"upsert": [deployments], "delete": [ids], "settings": {dynamic document}}` and + /// return the new registry version. + fn apply(&self, changes: &Bound<'_, PyAny>) -> PyResult { + let changes = config::parse_changes(document(changes)?).map_err(config_error_to_pyerr)?; + self.router + .apply(changes) + .map_err(|error| config_error_to_pyerr(error.into())) + } + + /// Build `entries` (model_list entries) through the host, driven like `start`: its steps + /// are `secrets`, `validate`, `provider` and `added`, and it ends with `done` (`built`: one + /// `{"index", "id", "provider"}` per deployment installed). `path` is `create` + /// (`set_model_list`) or `add` (`add_deployment`); `install` is `replace`, `append` or + /// `upsert`. + #[pyo3(signature = (entries, path, resolve_secrets, install))] + fn build( + &self, + entries: &Bound<'_, PyAny>, + path: &str, + resolve_secrets: bool, + install: &str, + ) -> PyResult { + let entries: Vec> = + serde_json::from_value(document(entries)?) + .map_err(|error| PyValueError::new_err(format!("bad model_list: {error}")))?; + let path: BuildPath = path + .parse() + .map_err(|_| PyValueError::new_err(format!("`{path}` is not a build path")))?; + let install: Install = install + .parse() + .map_err(|_| PyValueError::new_err(format!("`{install}` is not a build install")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .build(path, entries, resolve_secrets, install) + .await + .map(Finish::Built) + })) + } + + /// `get_router_model_info` for the deployment `query` describes (a + /// `RouterModelInfoQuery` document), driven like `start`: its steps are `model_info` + /// lookups, and it ends with `done` (`info`: the merged map or None, `user_keys`: the keys + /// whose values are the deployment's own). + fn router_model_info(&self, query: &Bound<'_, PyAny>) -> PyResult { + let query: litellm_router_core::model_info::RouterModelInfoQuery = + serde_json::from_value(document(query)?) + .map_err(|error| PyValueError::new_err(format!("bad model-info query: {error}")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .router_model_info(&query) + .await + .map(Finish::RouterModelInfo) + })) + } + + /// `get_deployment_model_info` for deployment `model_id` serving `model_name`, driven like + /// `router_model_info`; it ends with `done` (`info`: the merged map or None). + fn deployment_model_info( + &self, + model_id: &str, + model_name: &Bound<'_, PyAny>, + ) -> PyResult { + let model_id = model_id.to_owned(); + let model_name = from_py_argument::(model_name)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .deployment_model_info(&model_id, &model_name) + .await + .map(Finish::DeploymentModelInfo) + })) + } + + /// `_set_model_group_info` for the deployments `query` names (a `GroupInfoQuery` document), + /// driven like `router_model_info`; it ends with `done` (`value`: None when no deployment + /// matched, else `{"assignments": {field: value}}`, what the merge set on the group the host + /// built from its first deployment). + fn model_group_info(&self, query: &Bound<'_, PyAny>) -> PyResult { + let query: litellm_router_core::group_info::GroupInfoQuery = + serde_json::from_value(document(query)?) + .map_err(|error| PyValueError::new_err(format!("bad group-info query: {error}")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let info = router.model_group_info(&query).await?; + Ok(Finish::Value(info.map_or( + Value::Null, + |info| serde_json::json!({"assignments": info.assignments}), + ))) + })) + } + + /// `get_configured_token_limits` for the deployment with `id` and its `configured` model_info, + /// driven like `router_model_info`; it ends with `done` (`value`: `[max_input, max_output]`). + fn configured_token_limits( + &self, + id: &Bound<'_, PyAny>, + configured: &Bound<'_, PyAny>, + ) -> PyResult { + let id = from_py_argument::(id)?; + let Value::Object(configured) = document(configured)? else { + return Err(PyValueError::new_err( + "the configured model_info must be a dict", + )); + }; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let (input, output) = router.configured_token_limits(id, &configured).await?; + Ok(Finish::Value(serde_json::json!([input, output]))) + })) + } + + /// `get_model_listing_info` over the `deployments` a listed name names, driven like + /// `router_model_info`; it ends with `done` (`value`: `{"cost_map_keys", "max_input_tokens", + /// "max_output_tokens"}`). + fn model_listing_info(&self, deployments: &Bound<'_, PyAny>) -> PyResult { + let deployments: Vec = from_py_argument(deployments)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let listing = router.model_listing_info(&deployments).await?; + Ok(Finish::Value(serde_json::json!({ + "cost_map_keys": listing.cost_map_keys, + "max_input_tokens": listing.max_input_tokens, + "max_output_tokens": listing.max_output_tokens, + }))) + })) + } + + /// `get_configured_mode` / `get_configured_display_name`: `value` when it is a string that is + /// not blank. + #[staticmethod] + fn configured_text(value: &Bound<'_, PyAny>) -> PyResult> { + let value = from_py_argument::(value).ok(); + Ok(litellm_router_core::group_info::configured_text( + value.as_ref(), + )) + } + + /// `get_configured_service_tiers`: each of `model_infos`' `service_tiers`, unvalidated. + #[staticmethod] + fn configured_service_tiers( + py: Python<'_>, + model_infos: &Bound<'_, PyAny>, + ) -> PyResult>> { + let model_infos: Vec> = from_py_argument(model_infos)?; + litellm_router_core::group_info::configured_service_tiers(&model_infos) + .iter() + .map(|tiers| to_py(py, tiers)) + .collect() + } + + /// The routing-cache key deployment `id` serving litellm `model` counts `metric` (`tpm`, + /// `rpm`, `itpm` or `otpm`) under in `minute` (`HH-MM`). + fn usage_key(&self, id: &str, model: &str, metric: &str, minute: &str) -> PyResult { + let metric: UsageMetric = metric + .parse() + .map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?; + Ok(usage_key(id, model, metric, minute)) + } + + /// Keep the providers `other`'s build resolved for the deployments this router serves too. + fn adopt_providers(&self, other: &Self) { + self.router.adopt_providers(&other.router); + } + + fn registry_version(&self) -> u64 { + self.router.registry().version() + } + + /// The version the deployments before the newest appends last changed at: a list read at + /// or after it is current once it reads `deployments_from(len(list))`. + fn registry_rewritten(&self) -> u64 { + self.router.registry().rewritten() + } + + /// Record a failure the router did not see itself (a health check, an interceptor, a + /// logging callback), driven like `start`: its steps are cache ops, and it ends with `done` + /// (`cooled`, `cooldown_time`). `counts_request` also counts the failed request against the deployment's + /// per-minute requests, as an async call's failure callback does. `callbacks` is the logged + /// request's `OutcomeCallbacks`. + #[pyo3(signature = (id, failure, deployment_cooldown=false, count_only=false, requested_group=None, cooldown_time=None, uncounted=false, counts_request=false, usage_model=None, callbacks=None))] + #[allow(clippy::too_many_arguments)] + fn report_failure( + &self, + id: &str, + failure: &Bound<'_, PyDict>, + deployment_cooldown: bool, + count_only: bool, + requested_group: Option, + cooldown_time: Option, + uncounted: bool, + counts_request: bool, + usage_model: Option, + callbacks: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let failure = match attempt_outcome(failure)? { + Err(failure) => failure, + Ok(_) => { + return Err(PyValueError::new_err( + "a reported failure needs `ok: False`", + )); + } + }; + let router = self.router.clone(); + let id = id.to_owned(); + let report = FailureReport { + deployment_cooldown, + cooldown_time, + count_only, + requested_group, + uncounted, + counts_request, + usage_model, + callbacks: outcome_callbacks(callbacks)?, + }; + Ok(RouterCall::new(async move { + router + .report_failure(&id, &failure, report) + .await + .map(|cooldown| Finish::Reported { cooldown }) + })) + } + + /// Record a success a logging callback saw, driven like `report_failure`; it ends with + /// `done` (`key`: the tokens counter's key when usage was added, else None). + /// `usage_model` is the model its usage is keyed by, and `tracks_usage` the host's answer to + /// whether it counts any (see `Output::tracks_usage`). `callbacks`: see `report_failure`. + #[pyo3(signature = (id, total_tokens, requests, usage_model=None, tracks_usage=None, callbacks=None))] + fn report_success( + &self, + id: &str, + total_tokens: f64, + requests: f64, + usage_model: Option, + tracks_usage: Option, + callbacks: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let router = self.router.clone(); + let id = id.to_owned(); + let report = SuccessReport { + tracks_usage, + usage_model, + callbacks: outcome_callbacks(callbacks)?, + }; + Ok(RouterCall::new(async move { + router + .report_success(&id, total_tokens, requests, report) + .await + .map(|key| Finish::Counted { key }) + })) + } + + /// Count a success against `id`, configured or not; ends with `done` (`key`). `callbacks`: + /// see `report_failure`. + #[pyo3(signature = (id, callbacks=None))] + fn count_success( + &self, + id: &str, + callbacks: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let router = self.router.clone(); + let id = id.to_owned(); + let callbacks = outcome_callbacks(callbacks)?; + Ok(RouterCall::new(async move { + router + .count_success(&id, callbacks) + .await + .map(|key| Finish::Counted { key }) + })) + } + + /// The usage keys of this minute for every deployment of `model_group`: + /// `[(id, tpm_key, rpm_key)]`, or None when the group has no deployment. + fn usage_keys(&self, model_group: &str) -> Option> { + let minute = current_minute(SystemTime::now()); + let group = self.router.registry().group(model_group)?; + Some( + group + .iter() + .map(|deployment| { + let model = &deployment.params.model; + ( + deployment.id.clone(), + usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), + usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), + ) + }) + .collect(), + ) + } + + /// Every configured deployment, as `model_list` entries with their ids, in config order. + fn deployments<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + let records = registry + .deployments(&DeploymentFilter::default()) + .iter() + .map(|deployment| to_py(py, &host::deployment_record(deployment))) + .collect::>>()?; + PyList::new(py, records) + } + + /// The deployments from position `start` on, as `deployments()` lists them. + fn deployments_from<'py>(&self, py: Python<'py>, start: usize) -> PyResult> { + let registry = self.router.registry(); + let records = registry + .deployments_from(start) + .iter() + .map(|deployment| to_py(py, &host::deployment_record(deployment))) + .collect::>>()?; + PyList::new(py, records) + } + + /// The deployment with this id, or None. + fn deployment(&self, py: Python<'_>, id: &str) -> PyResult>> { + self.router + .registry() + .deployment(id) + .map(|deployment| to_py(py, &host::deployment_record(&deployment))) + .transpose() + } + + // -- registry lookups (router-core `lookup.rs`): rows the host reads its stored entries at -- + + /// `_core_get_all_deployments`, as rows. + #[pyo3(signature = (model_name, team_id=None, rename=None))] + fn lookup_all<'py>( + &self, + py: Python<'py>, + model_name: &str, + team_id: Option<&str>, + rename: Option<&str>, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).all(model_name, team_id, rename)) + } + + /// `_team_deployments_across_teams`: the rows, and how many teams own them. + fn lookup_across_teams<'py>( + &self, + py: Python<'py>, + model: &str, + ) -> PyResult<(Bound<'py, PyAny>, usize)> { + let registry = self.router.registry(); + let (rows, teams) = Lookup::new(®istry).across_teams(model); + Ok((rows_to_py(py, &rows)?, teams)) + } + + /// `_core_config_deployments`, as rows. + fn lookup_config<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).config()) + } + + /// `_core_get_model_ids`: the ids, or the row the read raises at. + #[pyo3(signature = (model_name=None, exclude_team_models=false))] + fn lookup_model_ids<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + exclude_team_models: bool, + ) -> PyResult<(Vec, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).model_ids(model_name, exclude_team_models) { + Ok(ids) => Ok((ids, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// `_core_get_model_list`, as rows. + #[pyo3(signature = (model_name, team_id, scope))] + fn lookup_model_list<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + rows_to_py( + py, + &Lookup::new(®istry).model_list(model_name, team_id, &scope), + ) + } + + /// `_core_get_model_names`, or the row the read raises at. + fn lookup_model_names<'py>( + &self, + py: Python<'py>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(Vec, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).model_names(team_id, &scope) { + Ok(names) => Ok((names, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// `_core_get_model_access_groups`: `(group, names)` pairs in the order the groups appear, + /// or the row the read raises at. + fn lookup_access_groups<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + access_group: Option<&str>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(AccessGroups, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).access_groups(model_name, access_group, team_id, &scope) { + Ok(groups) => Ok((groups, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// Whether the whole list lists any row, and the name of its first row whose litellm model + /// is `model_id` (or ends in `/model_id` or `:model_id`); or the row the read raises at. + fn lookup_name_serving<'py>( + &self, + py: Python<'py>, + model_id: &str, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(bool, Option, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).name_serving(model_id, &scope) { + Ok((listed, name)) => Ok((listed, name, None)), + Err(raw) => Ok((false, None, Some(row_to_py(py, &raw)?))), + } + } + + /// Whether `model_names` (the group names) holds `name`. + fn has_model_name(&self, name: &str) -> bool { + self.router + .registry() + .host_indexes() + .model_names + .contains(&Value::String(name.to_owned())) + } + + /// `_core_team_model_has_alternatives`, or the row the read raises at. + fn lookup_team_model_has_alternatives<'py>( + &self, + py: Python<'py>, + id: &str, + ) -> PyResult<(bool, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).team_model_has_alternatives(id) { + Ok(found) => Ok((found, None)), + Err(raw) => Ok((false, Some(row_to_py(py, &raw)?))), + } + } + + /// The position of the group's first deployment the team may use, or the row the read + /// raises at. + fn lookup_group_usable_by_team<'py>( + &self, + py: Python<'py>, + model_group: &str, + team_id: Option<&str>, + ) -> PyResult<(Option, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).group_deployment_usable_by_team(model_group, team_id) { + Ok(position) => Ok((position, None)), + Err(raw) => Ok((None, Some(row_to_py(py, &raw)?))), + } + } + + /// The position of the team public name's first deployment, or the row the read raises at. + fn lookup_team_public_name<'py>( + &self, + py: Python<'py>, + model_id: &str, + team_id: &str, + ) -> PyResult<(Option, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).team_public_name_deployment(model_id, team_id) { + Ok(position) => Ok((position, None)), + Err(raw) => Ok((None, Some(row_to_py(py, &raw)?))), + } + } + + /// `_routable_deployments`, as rows. + fn lookup_routable<'py>( + &self, + py: Python<'py>, + model_group: &str, + team_id: Option<&str>, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).routable(model_group, team_id)) + } + + /// `_core_wildcard_route_deployments`: the team router's routes, else the global router's. + fn lookup_wildcard_route<'py>( + &self, + py: Python<'py>, + model: &str, + team_id: Option<&str>, + requests: Vec, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py( + py, + &Lookup::new(®istry).wildcard_route(model, team_id, &requests), + ) + } + + /// `_core_get_candidate_model_ids_for_route` for the resolved name: the members' ids when it + /// names a callable routing group (`members`), else for a name that is no model group the + /// team's or the wildcard route's, else the group's; or the row the read raises at. + fn lookup_candidate_ids<'py>( + &self, + py: Python<'py>, + resolved: &str, + team_id: Option<&str>, + members: Option>, + requests: Vec, + ) -> PyResult<(Vec, Option>)> { + let registry = self.router.registry(); + let lookup = Lookup::new(®istry); + match lookup.candidate_ids(resolved, team_id, members.as_deref(), &requests) { + Ok(ids) => Ok((ids, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// Whether the id map holds `id` (`_has_model_id`). + fn has_model_id(&self, id: &str) -> bool { + self.router.registry().host_indexes().has_id(id) + } + + /// What the id map holds for `id`, as written (a position, or what a caller wrote), or None. + fn id_position<'py>(&self, py: Python<'py>, id: &str) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .id_entry(id) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// The id map's entries, in order: `(id, what it holds)`. + fn id_positions<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + let pairs: Vec = registry + .host_indexes() + .id_map + .iter() + .map(|(id, held)| Value::Array(vec![Value::String(id.clone()), held.clone()])) + .collect(); + to_py(py, &Value::Array(pairs)).map(|value| value.into_bound(py)) + } + + /// What the name index holds for `name`, as written, or None when it does not index it. + fn named_positions<'py>( + &self, + py: Python<'py>, + name: &str, + ) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .named(name) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// What the team index holds for `(team, public)`, as written, or None. + fn team_positions<'py>( + &self, + py: Python<'py>, + team: &str, + public: &str, + ) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .team_named(team, public) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// `team_public_model_names_of`: the public names the team index holds for `team`. + fn team_public_names_of(&self, team: &str) -> Vec { + self.router + .registry() + .host_indexes() + .team_public_names_of(team) + } + + /// The positions of the deployments whose `litellm_params.model` is `model`. + fn litellm_model_positions(&self, model: &str) -> Vec { + self.router + .registry() + .deployments(&DeploymentFilter::default()) + .iter() + .enumerate() + .filter(|(_, deployment)| { + deployment.raw_params.get("model").and_then(Value::as_str) == Some(model) + }) + .map(|(position, _)| position) + .collect() + } + + /// `PatternMatchRouter.route` over the wildcard router (or `team`'s): the first pattern's + /// deployments, each with the model the wildcards matched, or None. + #[pyo3(signature = (team, request, filtered=None))] + fn pattern_route<'py>( + &self, + py: Python<'py>, + team: Option<&str>, + request: &str, + filtered: Option>, + ) -> PyResult>> { + let registry = self.router.registry(); + let indexes = registry.host_indexes(); + let router = match team { + None => Some(&indexes.patterns), + Some(team) => indexes + .team_patterns + .iter() + .find(|(held, _)| held == team) + .map(|(_, router)| router), + }; + router + .and_then(|router| router.route(request, filtered.as_deref())) + .map(|routed| to_py(py, &Value::Array(routed)).map(|value| value.into_bound(py))) + .transpose() + } + + /// `get_models_by_pattern`: the models the first matching (request, pattern) pair routes to. + fn models_by_pattern(&self, team: Option<&str>, requests: Vec) -> Option> { + let registry = self.router.registry(); + let indexes = registry.host_indexes(); + let router = match team { + None => Some(&indexes.patterns), + Some(team) => indexes + .team_patterns + .iter() + .find(|(held, _)| held == team) + .map(|(_, router)| router), + }; + router.and_then(|router| router.models_by_pattern(&requests)) + } + + /// The indexes the host hands out (`HostIndexes`): id map, name and team indexes, group and + /// deployment names, team public names, provider default ids and the pattern routers. + fn host_indexes<'py>(&self, py: Python<'py>) -> PyResult> { + let indexes = serde_json::to_value(self.router.registry().host_indexes()) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + to_py(py, &indexes).map(|value| value.into_bound(py)) + } + + /// `provider_default_deployment_ids` as the indexes hold it now: the ids of the wildcard + /// deployments the cooldown check skips. + fn provider_default_ids(&self) -> Vec { + self.router + .registry() + .host_indexes() + .provider_default_ids + .clone() + } + + /// The provider the host resolved when it built the deployment with this id (its + /// `provider` build step), or None when there is no such deployment or it never went + /// through a host build. + fn provider(&self, id: &str) -> Option { + self.router + .registry() + .deployment(id) + .and_then(|deployment| deployment.resolved_provider.clone()) + } + + /// What `name` resolves to: `{"model_group", "kind", "ids"}`, or None for an unknown name. + #[pyo3(signature = (name, specific_deployment=false))] + fn resolve<'py>( + &self, + py: Python<'py>, + name: &str, + specific_deployment: bool, + ) -> PyResult>> { + let registry = self.router.registry(); + let Some(resolution) = registry.resolve( + name, + ResolveScope { + specific_deployment, + }, + ) else { + return Ok(None); + }; + let dict = PyDict::new(py); + dict.set_item("model_group", resolution.group)?; + dict.set_item("kind", format!("{:?}", resolution.kind))?; + dict.set_item( + "ids", + PyList::new(py, resolution.deployments.iter().map(|d| d.id.clone()))?, + )?; + Ok(Some(dict)) + } + + /// The deployments a pick for `model` would resolve to for `request`, before any team, + /// cooldown or strategy filter (`_core_get_routable_deployments`); empty when it refuses. + fn routable<'py>( + &self, + py: Python<'py>, + model: &str, + request: &Bound<'py, PyAny>, + ) -> PyResult> { + let request = config::parse_request(document(request)?, Some(model)) + .map_err(config_error_to_pyerr)?; + let records = litellm_router_core::resolve::routable( + &self.router.registry(), + &self.router.settings().dynamic, + &request, + ) + .iter() + .map(|d| to_py(py, &host::deployment_record(d))) + .collect::>>()?; + PyList::new(py, records) + } + + /// `_get_routing_context` for a pick of `model`: `(strategy, scope, group)`, spelled as a + /// strategy select step spells them. A name that resolves to nothing reads its routing group + /// membership as is. + #[pyo3(signature = (model, routing_strategy=None))] + fn routing_context( + &self, + model: &str, + routing_strategy: Option<&str>, + ) -> (String, &'static str, Option) { + let registry = self.router.registry(); + let settings = self.router.settings(); + let (kind, group) = registry + .resolve(model, ResolveScope::default()) + .map_or((ResolveKind::Group, model.to_owned()), |resolution| { + (resolution.kind, resolution.group) + }); + let (strategy, scope) = routing_context( + ®istry, + &settings.routing_strategy, + routing_strategy, + &kind, + &group, + ); + let (scope, group) = scope.wire(); + (strategy.to_string(), scope, group.map(str::to_owned)) + } + + fn retire(&self) { + self.router.retire(); + } + + /// What may still take over a stream from `origin` that fails before content, and the + /// retry budget for a failure of `kind` (see `stream::policy`). + #[pyo3(signature = (request, origin, kind=None))] + fn stream_policy<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + origin: &Bound<'py, PyAny>, + kind: Option<&str>, + ) -> PyResult> { + stream::policy(py, &self.router, request, origin, kind) + } + + /// Retry a stream from `origin` that failed before content within its group, driven like + /// `start`: `done` when a retry returned a response, `failed` with the newest failure. + fn retry_stream( + &self, + op: &str, + request: &Bound<'_, PyAny>, + origin: &Bound<'_, PyAny>, + failure: &Bound<'_, PyDict>, + ) -> PyResult { + stream::recover( + &self.router, + stream::Recovery::Retry, + op, + request, + origin, + failure, + ) + } + + /// The fallback walk for a stream from `origin` that failed before content, driven like + /// `start`. + fn fallback_stream( + &self, + op: &str, + request: &Bound<'_, PyAny>, + origin: &Bound<'_, PyAny>, + failure: &Bound<'_, PyDict>, + ) -> PyResult { + stream::recover( + &self.router, + stream::Recovery::Fallback, + op, + request, + origin, + failure, + ) + } + + /// `get_num_retries_from_retry_policy` for a failure shaped like an attempt outcome. + #[pyo3(signature = (failure, model_group=None))] + fn retry_policy_retries( + &self, + failure: &Bound<'_, PyDict>, + model_group: Option<&str>, + ) -> PyResult> { + retries::policy_retries(&self.router, failure, model_group) + } + + /// `_core_get_allowed_fails_from_policy` for a failure shaped like an attempt outcome. + fn allowed_fails_for(&self, failure: &Bound<'_, PyDict>) -> PyResult> { + retries::allowed_fails(&self.router, failure) + } + + /// `should_retry_this_error` as a predicate, for a failure shaped like an attempt outcome. + fn should_retry( + &self, + failure: &Bound<'_, PyDict>, + counts: (usize, usize), + fallbacks: (bool, bool, bool), + ) -> PyResult { + retries::should_retry(failure, counts, fallbacks) + } +} + +/// One router call, driven a step at a time by Python. +#[pyclass(frozen)] +pub(crate) struct RouterCall { + driver: Mutex, +} + +impl RouterCall { + fn new( + future: impl std::future::Future> + Send + 'static, + ) -> Self { + Self { + driver: Mutex::new(Driver::new(future)), + } + } +} + +#[pymethods] +impl RouterCall { + /// Hand back the outcome of the last step (`None` on the first call) and return the next + /// step, or the call's end. + #[pyo3(signature = (outcome))] + fn next<'py>( + &self, + py: Python<'py>, + outcome: Option>, + ) -> PyResult> { + let mut driver = self.driver.lock().unwrap_or_else(PoisonError::into_inner); + let response = match driver.pending_call() { + None => match outcome { + None => None, + Some(_) => return Err(drive_error_to_pyerr(DriveError::UnexpectedResponse)), + }, + Some(call) => Some(host_call_response(call, outcome)?), + }; + let step = driver.resume(response).map_err(drive_error_to_pyerr)?; + let dict = step_dict(py, step)?; + dict.set_item("events", events_list(py, driver.take_events())?)?; + Ok(dict) + } +} + +/// `(group, names)` pairs, in the order the groups appear. +type AccessGroups = Vec<(String, Vec)>; + +/// Lookup rows as the host reads them. +fn rows_to_py<'py>(py: Python<'py>, rows: &[Row]) -> PyResult> { + let rows = Value::Array(rows.iter().map(Row::wire).collect()); + to_py(py, &rows).map(|value| value.into_bound(py)) +} + +fn row_to_py<'py>(py: Python<'py>, row: &Row) -> PyResult> { + to_py(py, &row.wire()).map(|value| value.into_bound(py)) +} + +/// `{"aliases": [{alias, target}], "groups": [{name, members}], "wildcard": [requests]}`. +fn list_scope(scope: &Bound<'_, PyAny>) -> PyResult { + serde_json::from_value(document(scope)?) + .map_err(|error| PyValueError::new_err(format!("bad lookup scope: {error}"))) +} + +/// Read Python's outcome for `call`. +fn host_call_response( + call: &HostCall, + outcome: Option>, +) -> PyResult { + let missing = || { + drive_error_to_pyerr(DriveError::MissingResponse { + expected: call.kind(), + }) + }; + // `{"raised": True, ...}`: the step's work raised, reported as an attempt's failure. + if let HostCall::Select { .. } | HostCall::PreRoute { .. } | HostCall::AllowFallback { .. } = + call + && let Some(failure) = outcome.as_ref().and_then(|o| o.cast::().ok()) + && failure.contains("raised")? + { + return match attempt_outcome(failure)? { + Err(error) => Ok(HostCallResponse::Raised(error)), + Ok(_) => Err(PyValueError::new_err("a `raised` outcome is a failure")), + }; + } + match call { + HostCall::Sleep { .. } => Ok(HostCallResponse::Slept), + HostCall::AllowFallback { .. } => Ok(HostCallResponse::AllowFallback( + outcome.ok_or_else(missing)?.extract()?, + )), + HostCall::RefuseFallbacks { .. } => Ok(HostCallResponse::RefuseFallbacks( + outcome.ok_or_else(missing)?.extract()?, + )), + HostCall::Select { + stage: Stage::PreCall, + .. + } => { + let facts = from_py_argument(&outcome.ok_or_else(missing)?)?; + serde_json::from_value(facts) + .map(HostCallResponse::PreCall) + .map_err(|error| PyValueError::new_err(format!("bad pre-call facts: {error}"))) + } + HostCall::Select { + stage: Stage::RoutingPlugins { .. } | Stage::ProviderRequest, + .. + } => { + // `{"models": [...] | None}`: a bare None would read as no outcome. + let outcome = outcome.ok_or_else(missing)?; + let models = outcome.cast::()?.get_item("models")?; + Ok(HostCallResponse::Models(match models { + Some(models) if !models.is_none() => Some(models.extract()?), + _ => None, + })) + } + HostCall::Select { + stage: Stage::CallbackFilters { .. }, + .. + } => { + // The remaining ids, alone or paired with the id a filter prefers, or a failure + // mapping as an attempt reports one. + let outcome = outcome.ok_or_else(missing)?; + if let Ok(failure) = outcome.cast::() { + return match attempt_outcome(failure)? { + Err(error) => Ok(HostCallResponse::Filtered(Err(error))), + Ok(_) => Err(PyValueError::new_err( + "a `callback_filters` outcome is the remaining ids or a failure", + )), + }; + } + if let Ok(answer) = outcome.cast::() { + return Ok(HostCallResponse::Filtered(Ok(answer.extract()?))); + } + Ok(HostCallResponse::Filtered(Ok((outcome.extract()?, None)))) + } + HostCall::Select { + stage: Stage::Strategy { records, .. }, + .. + } if !records.is_empty() => Ok(HostCallResponse::Positions( + outcome.ok_or_else(missing)?.extract()?, + )), + HostCall::Select { .. } => Ok(HostCallResponse::Select( + outcome.ok_or_else(missing)?.extract()?, + )), + HostCall::PreRoute { .. } => pre_route_outcome(outcome).map(HostCallResponse::PreRoute), + HostCall::Attempt { .. } => { + let outcome = outcome.ok_or_else(missing)?; + let outcome = outcome.cast::()?; + let blocked = match outcome.get_item("blocked")? { + Some(value) if !value.is_none() => { + let name: String = value.extract()?; + Some(name.parse::().map_err(|_| { + PyValueError::new_err(format!( + "`{name}` is not a way a response is blocked" + )) + })?) + } + Some(_) | None => None, + }; + let deferred = outcome + .get_item("deferred")? + .is_some_and(|flag| flag.is_truthy().unwrap_or(false)); + let tracks_usage = outcome + .get_item("tracks_usage")? + .map(|flag| flag.extract::>()) + .transpose()? + .flatten(); + attempt_outcome(outcome).map(|result| match result { + Ok(total_tokens) if blocked.is_some() || deferred || tracks_usage.is_some() => { + HostCallResponse::Served { + total_tokens, + blocked, + deferred, + tracks_usage, + } + } + result => HostCallResponse::Attempt(result), + }) + } + HostCall::Cache { op, .. } => cache_outcome(*op, outcome).map(HostCallResponse::Cache), + HostCall::Secrets { .. } => Ok(HostCallResponse::Secrets(from_py_argument( + &outcome.ok_or_else(missing)?, + )?)), + // `None` is a response here (no row, no entry), so a missing outcome reads as one. + HostCall::ModelInfo { .. } => Ok(HostCallResponse::ModelInfo(match outcome { + Some(outcome) if !outcome.is_none() => from_py_argument(&outcome)?, + _ => Value::Null, + })), + HostCall::Build { kind, .. } => { + let outcome = outcome.ok_or_else(missing)?; + let outcome = outcome.cast::()?; + let invalid = outcome + .get_item("invalid")? + .is_some_and(|flag| flag.is_truthy().unwrap_or(false)); + match kind { + StepKind::Provider => Ok(HostCallResponse::Provided(if invalid { + Provided::Invalid + } else { + Provided::Provider(match outcome.get_item("provider")? { + Some(provider) if !provider.is_none() => Some(provider.extract()?), + _ => None, + }) + })), + _ => { + let built = if invalid { + Built::Invalid + } else { + match outcome.get_item("entry")? { + Some(entry) if !entry.is_none() => { + Built::Entry(from_py_argument(&entry)?) + } + _ => Built::Skip, + } + }; + Ok(if *kind == StepKind::Validate { + HostCallResponse::Validated(built) + } else { + HostCallResponse::Added(built) + }) + } + } + } + } +} + +/// A pre-route response: `None`, a model group name, or a custom strategy's deployment document. +fn pre_route_outcome(outcome: Option>) -> PyResult> { + let Some(outcome) = outcome.filter(|outcome| !outcome.is_none()) else { + return Ok(None); + }; + if let Ok(model) = outcome.cast::() { + return Ok(Some(PreRouted::Model(model.to_str()?.to_owned()))); + } + // `(model, group)`: the model to resolve, served as the group a session binding stamped. + if let Ok(pair) = outcome.cast::() + && pair.len() == 2 + { + return Ok(Some(PreRouted::Bound { + model: pair.get_item(0)?.extract()?, + group: pair.get_item(1)?.extract()?, + })); + } + let spec: DeploymentSpec = serde_json::from_value(document(&outcome)?) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + let deployment = Deployment::from_spec(spec) + .map_err(|error| PyValueError::new_err(Report::new(&error).to_string()))?; + Ok(Some(PreRouted::Deployment(Arc::new(deployment)))) +} + +/// `get`: the JSON value or None; `batch_get`: a list of JSON values or None, one per key; +/// `increment`: the new int or float; `set` and `delete`: None. +fn cache_outcome(op: CacheOp, outcome: Option>) -> PyResult { + let outcome = outcome.filter(|outcome| !outcome.is_none()); + let json = |value: Option| value.filter(|value| !value.is_null()); + match op { + CacheOp::Get => outcome + .map(|outcome| from_py_argument::(&outcome)) + .transpose() + .map(|value| CacheOutcome::Value(json(value))), + CacheOp::BatchGet => { + let outcome = outcome.ok_or_else(|| { + PyValueError::new_err("a `batch_get` cache outcome must be a list, not None") + })?; + if !outcome.is_instance_of::() { + return Err(PyValueError::new_err( + "a `batch_get` cache outcome must be a list with one entry per key", + )); + } + let values: Vec> = from_py_argument(&outcome)?; + Ok(CacheOutcome::Values(values.into_iter().map(json).collect())) + } + CacheOp::Increment => outcome + .filter(|outcome| !outcome.is_instance_of::()) + .and_then(|outcome| outcome.extract::().ok()) + .map(CacheOutcome::Number) + .ok_or_else(|| { + PyValueError::new_err("an `increment` cache outcome must be the new int or float") + }), + CacheOp::Set | CacheOp::Delete | CacheOp::IncrementUsage => match outcome { + None => Ok(CacheOutcome::Stored), + Some(_) => Err(PyValueError::new_err(format!( + "a `{op}` cache step has no outcome, so pass None" + ))), + }, + } +} + +/// `{"ok": True, "total_tokens": int}` (an attempt's may add `"blocked": str`, `"deferred": bool` +/// and `"tracks_usage": bool`) or `{"ok": False, "kind": str, "status": int | None, +/// "retry_after": float | None, "num_retries": int | None, "detail": str | None, +/// "outside_deployment": bool, "caller_scoped": bool, "status_absent": bool, +/// "uncounted_request": bool}`. +pub(crate) fn attempt_outcome(outcome: &Bound<'_, PyDict>) -> PyResult> { + fn item<'py, T: FromPyObjectOwned<'py>>( + outcome: &Bound<'py, PyDict>, + key: &str, + ) -> PyResult> { + match outcome.get_item(key)? { + Some(value) if !value.is_none() => Ok(Some(value.extract().map_err(Into::into)?)), + Some(_) | None => Ok(None), + } + } + let ok: bool = item(outcome, "ok")? + .ok_or_else(|| PyValueError::new_err("an attempt outcome needs `ok`"))?; + if ok { + return Ok(Ok(item(outcome, "total_tokens")?.unwrap_or(0))); + } + let kind: String = item(outcome, "kind")? + .ok_or_else(|| PyValueError::new_err("a failed attempt outcome needs `kind`"))?; + let kind: FailureKind = kind + .parse() + .map_err(|_| PyValueError::new_err(format!("`{kind}` is not a router failure kind")))?; + let retry_after = item::(outcome, "retry_after")? + .map(|seconds| { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("`retry_after` must be a non-negative number")) + }) + .transpose()?; + let mut failure = AttemptError::new(kind, item(outcome, "status")?); + failure.retry_after = retry_after; + failure.num_retries = item(outcome, "num_retries")?; + failure.failed_deployment_id = item(outcome, "failed_deployment_id")?; + failure.detail = item::(outcome, "detail")?.map(Into::into); + failure.recorded = item::(outcome, "recorded")?.unwrap_or(false); + failure.names_provider = item::(outcome, "llm_provider")?.unwrap_or(false); + failure.agentic_followup = item::(outcome, "agentic_followup")?.unwrap_or(false); + failure.outside_deployment = item::(outcome, "outside_deployment")?.unwrap_or(false); + failure.status_absent = item::(outcome, "status_absent")?.unwrap_or(false); + failure.uncounted_request = item::(outcome, "uncounted_request")?.unwrap_or(false); + failure.caller_scoped = item::(outcome, "caller_scoped")?.unwrap_or(false); + failure.classes = item::>(outcome, "classes")? + .unwrap_or_default() + .into_iter() + .map(Into::into) + .collect(); + Ok(Err(failure)) +} + +/// `[{"kind": "retries_exhausted", "model_group": str, "attempted": int, "max": int} +/// | {"kind": "fallback", "to": str, "order": int | None, "level": int}]`. +/// The `cause` of a `retry_logged` or `retry_policy_raised` event, with what names it. +fn cause_items(dict: &Bound<'_, PyDict>, cause: RetryCause) -> PyResult<()> { + match cause { + RetryCause::Attempt => dict.set_item("cause", "attempt"), + RetryCause::Level => dict.set_item("cause", "level"), + RetryCause::Other => dict.set_item("cause", "other"), + RetryCause::Pick(picked) => { + dict.set_item("cause", "pick")?; + pick_items(dict, &picked.0) + } + } +} + +/// A pick's error, as the `failed` step names it (`error_kind` and what names it), with the +/// model its message names (`failed_model`) when the error carries one. +fn pick_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> { + error_items(dict, error)?; + match error { + Error::NoHealthyDeployments { model, .. } | Error::UnknownModel { model } => { + dict.set_item("failed_model", model) + } + _ => Ok(()), + } +} + +/// `error_kind` and the fields that name the error. +fn error_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> { + dict.set_item("error", Report::new(error).to_string())?; + dict.set_item("error_kind", <&'static str>::from(ErrorKind::from(error)))?; + match error { + Error::EndpointNotSupported { model, message } => { + dict.set_item("error_model", model)?; + dict.set_item("error_detail", message) + } + Error::OnlyMarkers { model } + | Error::UnknownModel { model } + | Error::NoPassThroughDeployments { model } + | Error::PtuReserved { model } + | Error::MultipleTeams { model } + | Error::DefaultDeploymentRaises { model } + | Error::IndexOutOfRange { model } => dict.set_item("error_model", model), + Error::ContextWindowExceeded { model, detail } => { + dict.set_item("error_model", model)?; + dict.set_item("error_detail", detail) + } + Error::UnsupportedParams { model, params } => { + dict.set_item("error_model", model)?; + dict.set_item("error_params", PyList::new(dict.py(), params)?) + } + Error::RoutingPluginsLeftNothing { models } => { + dict.set_item("error_params", PyList::new(dict.py(), models)?) + } + Error::Reserved { + model, + until, + timezone, + } => { + dict.set_item("error_model", model)?; + dict.set_item("error_detail", format!("{until} {timezone}")) + } + Error::NoHealthyDeployments { + request_ids: Some(ids), + .. + } => dict.set_item("error_request_ids", PyList::new(dict.py(), ids)?), + Error::NonNumericWeight { metric } => dict.set_item("error_detail", metric), + Error::MissingRoutingGroup { group } => dict.set_item("error_model", group), + Error::ModelInfo { exception, detail } => { + dict.set_item("error_exception", *exception)?; + dict.set_item("error_detail", detail) + } + _ => Ok(()), + } +} + +fn events_list(py: Python<'_>, events: Vec) -> PyResult> { + let list = PyList::empty(py); + for event in events { + let dict = PyDict::new(py); + match event { + RouterEvent::RetriesExhausted { + model_group, + attempted, + max, + level, + picked, + } => { + dict.set_item("kind", "retries_exhausted")?; + dict.set_item("model_group", model_group)?; + dict.set_item("attempted", attempted)?; + dict.set_item("max", max)?; + dict.set_item("level", level)?; + dict.set_item("picked", picked)?; + } + RouterEvent::Retry { retry, .. } => { + dict.set_item("kind", "retry")?; + dict.set_item("retry", retry)?; + } + RouterEvent::FallingBack { level, model } => { + dict.set_item("kind", "falling_back")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + } + RouterEvent::OrderChainFailed { level, last_hop } => { + dict.set_item("kind", "order_chain_failed")?; + dict.set_item("level", level)?; + dict.set_item("last_hop", last_hop)?; + } + RouterEvent::FallbacksFailed { + level, + chain, + last_hop, + } => { + dict.set_item("kind", "fallbacks_failed")?; + dict.set_item("level", level)?; + dict.set_item("chain", chain)?; + dict.set_item("last_hop", last_hop)?; + } + RouterEvent::Fallback { + from, + to, + order, + depth, + level, + scope, + parent, + max_fallbacks, + unroutable, + } => { + dict.set_item("kind", "fallback")?; + dict.set_item("from", from)?; + dict.set_item("to", to)?; + dict.set_item("order", order)?; + dict.set_item("depth", depth)?; + dict.set_item("level", level)?; + dict.set_item("scope", scope)?; + dict.set_item("parent", parent)?; + dict.set_item("max_fallbacks", max_fallbacks)?; + if let Some(model) = unroutable { + dict.set_item("unroutable", to_py(py, &model)?)?; + } + } + RouterEvent::Blocked { + level, + model, + blocked, + } => { + dict.set_item("kind", "blocked")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("blocked", <&'static str>::from(blocked))?; + } + RouterEvent::FallbackDone { + from, + to, + level, + scope, + parent, + succeeded, + } => { + dict.set_item("kind", "fallback_done")?; + dict.set_item("from", from)?; + dict.set_item("to", to)?; + dict.set_item("level", level)?; + dict.set_item("scope", scope)?; + dict.set_item("parent", parent)?; + dict.set_item("succeeded", succeeded)?; + } + RouterEvent::RetryLogged { + scope, + level, + model, + attempted_retries, + cause, + } => { + dict.set_item("kind", "retry_logged")?; + dict.set_item("scope", scope)?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("attempted_retries", attempted_retries)?; + cause_items(&dict, cause)?; + } + RouterEvent::RetryPolicyRaised { + level, + model, + cause, + } => { + dict.set_item("kind", "retry_policy_raised")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + cause_items(&dict, cause)?; + } + RouterEvent::RetryCountRaised { + level, + model, + cause, + } => { + dict.set_item("kind", "retry_count_raised")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + cause_items(&dict, cause)?; + } + RouterEvent::RetriesOpened { + scope, + max_retries, + model_group_size, + } => { + dict.set_item("kind", "retries_opened")?; + dict.set_item("scope", scope)?; + dict.set_item("max_retries", max_retries)?; + dict.set_item("model_group_size", model_group_size)?; + } + RouterEvent::RetryBudget { scope, max_retries } => { + dict.set_item("kind", "retry_budget")?; + dict.set_item("scope", scope)?; + dict.set_item("max_retries", max_retries)?; + } + RouterEvent::PickFailed { + level, + model, + error, + } => { + dict.set_item("kind", "pick_failed")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + pick_items(&dict, &error.0)?; + } + RouterEvent::LevelFailed { + level, + model, + error, + handling, + } => { + dict.set_item("kind", "level_failed")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("handling", handling)?; + pick_items(&dict, &error.0)?; + } + RouterEvent::MockFailed { + mock, + level, + model, + num_retries, + } => { + dict.set_item("kind", "mock_failed")?; + dict.set_item("mock", <&'static str>::from(mock))?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("num_retries", num_retries)?; + } + RouterEvent::CooldownSet { + deployment_id, + cooldown, + } => { + dict.set_item("kind", "cooldown_set")?; + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("cooldown_time", cooldown.as_secs_f64())?; + } + RouterEvent::PickStarted { model } => { + dict.set_item("kind", "pick_started")?; + dict.set_item("model", model)?; + } + RouterEvent::Selected { + deployment_id, + pick, + .. + } => { + dict.set_item("kind", "selected")?; + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("pick_seconds", pick.as_secs_f64())?; + } + _ => continue, + } + list.append(dict)?; + } + Ok(list) +} + +fn step_dict(py: Python<'_>, step: Step) -> PyResult> { + let dict = PyDict::new(py); + let kind = match &step { + Step::HostCall(call) => call.kind(), + Step::Done(_) => StepKind::Done, + Step::Failed(_) => StepKind::Failed, + }; + dict.set_item("kind", <&'static str>::from(kind))?; + match step { + Step::HostCall(call) => host_call_items(py, &dict, call)?, + Step::Done(finish) => done_items(&dict, finish)?, + Step::Failed(error) => { + let attempt = error.attempt(); + error_items(&dict, &error)?; + dict.set_item("attempt_failed", attempt.is_some())?; + dict.set_item("status", attempt.and_then(|failure| failure.status))?; + dict.set_item( + "failure_kind", + attempt.map(|failure| <&'static str>::from(failure.kind)), + )?; + } + } + Ok(dict) +} + +fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) -> PyResult<()> { + match call { + HostCall::Attempt { + deployment_id, + deployment, + model, + model_group, + routed, + target_order, + candidate_ids, + retry, + fallback_depth, + fallback_level, + timeout, + params, + scope, + max_retries, + model_group_size, + } => { + dict.set_item("routed", routed)?; + dict.set_item("scope", scope)?; + dict.set_item("max_retries", max_retries)?; + dict.set_item("model_group_size", model_group_size)?; + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("deployment", to_py(py, &deployment)?)?; + dict.set_item("model", model)?; + dict.set_item("model_group", model_group)?; + dict.set_item("target_order", target_order)?; + dict.set_item("candidate_ids", PyList::new(py, candidate_ids)?)?; + dict.set_item("retry", retry)?; + dict.set_item("fallback_depth", fallback_depth)?; + dict.set_item("fallback_level", fallback_level)?; + dict.set_item("timeout", timeout)?; + dict.set_item("params", params.map(|p| to_py(py, &p)).transpose()?) + } + HostCall::Sleep { seconds } => dict.set_item("seconds", seconds), + HostCall::AllowFallback { target } => dict.set_item("target", target), + HostCall::RefuseFallbacks { level, targets } => { + dict.set_item("level", level)?; + let targets = targets.0.iter().cloned().map(Value::Object).collect(); + dict.set_item("targets", to_py(py, &Value::Array(targets))?) + } + HostCall::Select { + stage, + model_group, + name, + scope, + group, + pool, + candidates, + } => { + let (stage, strategy) = match stage { + Stage::Plugins => ("plugins", None), + Stage::Strategy { strategy, records } => { + if !records.is_empty() { + dict.set_item("records", to_py(py, &Value::Array(records))?)?; + } + ("strategy", Some(strategy)) + } + Stage::Tags => ("tags", None), + Stage::PreCall => ("pre_call", None), + Stage::TextOnly => ("text_only", None), + Stage::RouteIncapable { handler } => { + dict.set_item("handler", handler)?; + ("route_incapable", None) + } + Stage::MediaRefusing => ("media_refusing", None), + Stage::EmbeddingBatch => ("embedding_batch", None), + Stage::OverrideCheck(strategy) => ("override_check", Some(strategy)), + Stage::ProviderRequest => ("provider_request", None), + Stage::CallbackFilters { + excluded_ids, + skipped_ids, + target_order, + } => { + dict.set_item("excluded_ids", excluded_ids)?; + dict.set_item("skipped_ids", skipped_ids)?; + dict.set_item("target_order", target_order)?; + ("callback_filters", None) + } + Stage::RoutingPlugins { run } => { + dict.set_item("run", run)?; + ("routing_plugins", None) + } + }; + dict.set_item("pool", PyList::new(py, pool)?)?; + dict.set_item("stage", stage)?; + dict.set_item("strategy", strategy)?; + dict.set_item("model_group", model_group)?; + dict.set_item("model", name)?; + dict.set_item("scope", scope)?; + dict.set_item("group", group)?; + dict.set_item("candidates", PyList::new(py, candidates)?) + } + HostCall::PreRoute { + model, + fallback_depth, + retry, + } => { + dict.set_item("model", model)?; + dict.set_item("fallback_depth", fallback_depth)?; + dict.set_item("retry", retry) + } + HostCall::Cache { + op, + key, + keys, + value, + ttl, + local, + } => { + dict.set_item("op", <&'static str>::from(op))?; + dict.set_item("key", key)?; + dict.set_item("keys", keys.map(|keys| PyList::new(py, keys)).transpose()?)?; + dict.set_item("value", value.map(|value| to_py(py, &value)).transpose()?)?; + dict.set_item("ttl", ttl)?; + dict.set_item("local", local) + } + HostCall::Secrets { index, refs } => { + dict.set_item("index", index)?; + let refs = refs + .into_iter() + .map(|secret| { + let item = PyDict::new(py); + item.set_item("name", secret.name)?; + match secret.at { + SecretAt::Param(key) => { + item.set_item("key", key)?; + item.set_item("env_only", false)?; + } + SecretAt::DataSource { source, key } => { + item.set_item("source", source)?; + item.set_item("key", key)?; + item.set_item("env_only", true)?; + } + } + Ok(item) + }) + .collect::>>()?; + dict.set_item("refs", PyList::new(py, refs)?) + } + HostCall::Build { + kind: _, + path, + index, + entry, + declared_id, + duplicate_id, + provider, + } => { + dict.set_item("path", <&'static str>::from(path))?; + dict.set_item("index", index)?; + dict.set_item("entry", to_py(py, &entry)?)?; + dict.set_item("declared_id", declared_id)?; + dict.set_item("duplicate_id", duplicate_id)?; + dict.set_item("provider", provider) + } + HostCall::ModelInfo { lookup } => { + if let Value::Object(lookup) = lookup { + for (key, value) in lookup { + dict.set_item(key, to_py(py, &value)?)?; + } + } + Ok(()) + } + } +} + +fn done_items(dict: &Bound<'_, PyDict>, finish: Finish) -> PyResult<()> { + match finish { + Finish::Call(response) => { + dict.set_item("deployment_id", response.deployment_id)?; + dict.set_item("model_group", response.model_group)?; + dict.set_item("attempted_retries", response.attempted_retries)?; + dict.set_item("max_retries", response.max_retries)?; + dict.set_item("attempted_fallbacks", response.attempted_fallbacks)?; + dict.set_item("total_tokens", response.output.total_tokens) + } + Finish::Selected { + deployment_id, + model_group, + deployment, + } => { + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("model_group", model_group)?; + dict.set_item("deployment", to_py(dict.py(), &deployment)?) + } + Finish::Reported { cooldown } => { + dict.set_item("cooled", cooldown.is_some())?; + dict.set_item("cooldown_time", cooldown.map(|c| c.as_secs_f64())) + } + Finish::Built(built) => { + let py = dict.py(); + let records = built + .into_iter() + .map(|deployment| { + let item = PyDict::new(py); + item.set_item("index", deployment.index)?; + item.set_item( + "id", + deployment.spec.model_info.get("id").and_then(Value::as_str), + )?; + item.set_item("provider", deployment.spec.provider)?; + Ok(item) + }) + .collect::>>()?; + dict.set_item("built", PyList::new(py, records)?) + } + Finish::Counted { key } => dict.set_item("key", key), + Finish::RouterModelInfo(merged) => { + let py = dict.py(); + match merged { + Some(merged) => { + dict.set_item("info", to_py(py, &Value::Object(merged.info))?)?; + dict.set_item("user_keys", merged.user_keys) + } + None => dict.set_item("info", py.None()), + } + } + Finish::DeploymentModelInfo(info) => { + let py = dict.py(); + match info { + Some(info) => dict.set_item("info", to_py(py, &Value::Object(info))?), + None => dict.set_item("info", py.None()), + } + } + Finish::Value(value) => dict.set_item("value", to_py(dict.py(), &value)?), + Finish::Healthy { + single, + deployment_ids, + deployments, + } => { + dict.set_item("single", single)?; + dict.set_item("deployment_ids", deployment_ids)?; + let records = deployments + .iter() + .map(|d| to_py(dict.py(), d)) + .collect::>>()?; + dict.set_item("deployments", PyList::new(dict.py(), records)?) + } + } +} + +#[cfg(test)] +mod tests { + use pyo3::{ffi::c_str, prelude::*, types::PyDict}; + use rstest::rstest; + + const CONFIG: &str = r#"{"model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}} + ]}"#; + + const PAIR: &str = r#"{"model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}} + ]}"#; + + #[rstest] + fn python_drives_calls_and_selections_through_the_handle() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("native", crate::native_module(py)).unwrap(); + locals.set_item("config", CONFIG).unwrap(); + locals.set_item("pair", PAIR).unwrap(); + py.run( + c_str!( + r#" +import time +store = {} +asked = [] + +def answer_cache(step): + asked.append(step) + op, key = step["op"], step["key"] + if op == "get": + return store.get(key) + if op == "batch_get": + return [store.get(k) for k in step["keys"]] + if op == "set": + store[key] = step["value"] + return None + if op == "increment": + store[key] = store.get(key, 0) + step["value"] + return store[key] + if op == "delete": + store.pop(key, None) + return None + raise AssertionError(op) + +def drive(call, outcome): + step = call.next(outcome) + while step["kind"] == "cache": + step = call.next(answer_cache(step)) + return step + +handle = native.RouterCoreHandle(config, {"num_retries": 1}) +picked = drive(handle.select("primary", {}), None) +assert (picked["kind"], picked["deployment_id"], picked["model_group"]) == ("done", "a1", "primary"), picked +assert [step["op"] for step in asked] == ["batch_get"] +assert list(asked[0]) == ["kind", "op", "key", "keys", "value", "ttl", "local", "events"], list(asked[0]) +call = handle.start("acompletion", '{"model": "primary"}') +version = handle.registry_version() +assert handle.apply({"delete": ["gone"]}) == version + 1 +steps = [drive(call, None)] +steps.append(drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})) +steps.append(drive(call, None)) +steps.append(drive(call, {"ok": True, "total_tokens": 5})) +kinds = [step["kind"] for step in steps] + +cooled = native.RouterCoreHandle(pair, {"num_retries": 0}) +call = cooled.start("acompletion", {"model": "primary"}) +attempt = drive(call, None) +tried = attempt["deployment_id"] +assert attempt["deployment"]["model_info"]["id"] == tried, attempt +assert attempt["deployment"]["litellm_params"]["model"].startswith("openai/"), attempt +assert [d["model_info"]["id"] for d in cooled.deployments()] == ["a1", "a2"] +assert cooled.deployment("a2")["model_name"] == "primary" +assert cooled.resolve("primary")["ids"] == ["a1", "a2"] +assert cooled.resolve("nope") is None +assert drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})["kind"] == "failed" +# The core records the failure itself: the 429 cools the deployment down in the host's store. +assert f"deployment:{tried}:cooldown" in store, store +other = drive(cooled.select("primary", {}), None) +assert other["kind"] == "done" and other["deployment_id"] != tried, other + +call = handle.select("primary", {}) +step = call.next(None) +assert step["op"] == "batch_get" +for bad in ([], [None, None], "nope", None): + try: + call.next(bad) + except ValueError: + pass + else: + raise AssertionError(f"{bad!r} should not answer a batch_get of one key") +assert drive(call, [None])["kind"] == "done" +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let kinds: Vec = locals + .get_item("kinds") + .unwrap() + .unwrap() + .extract() + .unwrap(); + assert_eq!(kinds, ["attempt", "sleep", "attempt", "done"]); + }); + } + + #[rstest] + #[case::unknown_key(r#"{"routing_strategy_args": {}}"#, true)] + #[case::bad_json("{", false)] + #[case::not_an_object("[]", false)] + fn construction_failures_raise_the_class_python_reads( + #[case] config: &str, + #[case] declined: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let error = crate::native_module(py) + .getattr("RouterCoreHandle") + .and_then(|handle| handle.call1((config, "{}"))) + .unwrap_err(); + assert_eq!( + error.is_instance_of::(py), + declined, + "{error}" + ); + assert_eq!( + error.is_instance_of::(py), + !declined, + "{error}" + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/retries.rs b/litellm-rust/crates/python-bridge/src/router/retries.rs new file mode 100644 index 00000000000..949aae96b34 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/retries.rs @@ -0,0 +1,55 @@ +//! The retry decisions `RouterCoreHandle` exposes as members: `get_num_retries_from_retry_policy` +//! and `should_retry_this_error`, decided by router-core for a failure Python describes the way it +//! describes a failed attempt. + +use litellm_router_core::{ + AttemptError, Router, + retry::{RetryContext, resolve_retry_policy, should_retry as core_should_retry}, +}; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; + +fn failure(outcome: &Bound<'_, PyDict>) -> PyResult { + match super::attempt_outcome(outcome)? { + Err(failure) => Ok(failure), + Ok(_) => Err(PyValueError::new_err("a failure needs `ok: False`")), + } +} + +pub(crate) fn policy_retries( + router: &Router, + outcome: &Bound<'_, PyDict>, + model_group: Option<&str>, +) -> PyResult> { + let failure = failure(outcome)?; + let settings = router.settings(); + let policy = match model_group { + Some(group) => resolve_retry_policy( + settings.retry_policy.as_ref(), + group, + Some(&settings.model_group_retry_policy), + ), + None => settings.retry_policy.as_ref(), + }; + Ok(policy.and_then(|policy| policy.retries_for(&failure))) +} + +pub(crate) fn allowed_fails(router: &Router, outcome: &Bound<'_, PyDict>) -> PyResult> { + Ok(router.settings().allowed_fails_for(&failure(outcome)?)) +} + +pub(crate) fn should_retry( + outcome: &Bound<'_, PyDict>, + (healthy, all): (usize, usize), + (context_window, content_policy, regular): (bool, bool, bool), +) -> PyResult { + Ok(core_should_retry( + &failure(outcome)?, + RetryContext { + healthy, + all, + has_context_window_fallbacks: context_window, + has_content_policy_fallbacks: content_policy, + has_regular_fallbacks: regular, + }, + )) +} diff --git a/litellm-rust/crates/python-bridge/src/router/stream.rs b/litellm-rust/crates/python-bridge/src/router/stream.rs new file mode 100644 index 00000000000..f85d73e0778 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/stream.rs @@ -0,0 +1,94 @@ +//! The stream-recovery entry points of `RouterCoreHandle`: what may still take over a stream that +//! failed before content, the same-group stream retry, and the fallback walk re-entered from a +//! stream. The host's stream wrappers call them; the calls are driven like `start`. + +use litellm_router_core::{ + FailureKind, Op, Router, + call::{StreamOrigin, StreamPolicy}, +}; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; + +use super::{RouterCall, attempt_outcome, config, config_error_to_pyerr, document, host::Finish}; + +/// A stream recovery a host wrapper asks for. +#[derive(Clone, Copy)] +pub(super) enum Recovery { + /// `_astream_retry_same_group`. + Retry, + /// The fallback walk re-entered from the stream wrapper. + Fallback, +} + +fn origin(value: &Bound<'_, PyAny>) -> PyResult { + serde_json::from_value(document(value)?) + .map_err(|error| PyValueError::new_err(format!("a stream origin: {error}"))) +} + +fn kind(kind: Option<&str>) -> PyResult> { + kind.map(|kind| { + kind.parse() + .map_err(|_| PyValueError::new_err(format!("`{kind}` is not a router failure kind"))) + }) + .transpose() +} + +/// `{"can_retry", "can_fall_back", "refusal_fallback", "retry_budget": int | None, +/// "policy_applies": bool}`. +pub(super) fn policy<'py>( + py: Python<'py>, + router: &Router, + request: &Bound<'py, PyAny>, + stream_origin: &Bound<'py, PyAny>, + failure_kind: Option<&str>, +) -> PyResult> { + let request = config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?; + let StreamPolicy { + can_retry, + can_fall_back, + refusal_fallback, + retry_budget, + } = router.stream_policy(&request, &origin(stream_origin)?, kind(failure_kind)?); + let dict = PyDict::new(py); + dict.set_item("can_retry", can_retry)?; + dict.set_item("can_fall_back", can_fall_back)?; + dict.set_item("refusal_fallback", refusal_fallback)?; + dict.set_item("retry_budget", retry_budget.map(|(budget, _)| budget))?; + dict.set_item( + "policy_applies", + retry_budget.is_some_and(|(_, applies)| applies), + )?; + Ok(dict) +} + +/// A recovery call for a stream from `stream_origin` that failed with `failure`. +pub(super) fn recover( + router: &Router, + recovery: Recovery, + op: &str, + request: &Bound<'_, PyAny>, + stream_origin: &Bound<'_, PyAny>, + failure: &Bound<'_, PyDict>, +) -> PyResult { + let op = Op::new(op); + let request = config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?; + let stream_origin = origin(stream_origin)?; + let Err(failure) = attempt_outcome(failure)? else { + return Err(PyValueError::new_err("a stream failure needs `ok: False`")); + }; + let router = router.clone(); + Ok(RouterCall::new(async move { + let response = match recovery { + Recovery::Retry => { + router + .retry_stream(op, request, stream_origin, failure) + .await + } + Recovery::Fallback => { + router + .fallback_stream(op, request, stream_origin, failure) + .await + } + }; + response.map(Finish::Call) + })) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs index d78d141f0f2..1bcdc18a765 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use host::ChatCompletionsPythonHost; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, ScopedCache}; +use litellm_cache_response::CachePolicy; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_host::{call::HostedMachine, protocol::Protocol}; use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions}; @@ -57,10 +57,6 @@ impl InferenceRoute for ChatCompletionsRoute { Self::new(http, auth, secrets) } - fn with_cache(self, cache: ScopedCache) -> Self { - self.with_cache(cache) - } - fn machine( self, call: ::Request, diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index ca30c94b02f..1245c126c4e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache}; +use litellm_cache_response::CachePolicy; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{call::HostedMachine, protocol::Protocol}; @@ -142,7 +142,6 @@ pub(super) trait InferenceRoute: Sized + 'static { auth: Arc, secrets: Arc, ) -> Self; - fn with_cache(self, cache: ScopedCache) -> Self; fn machine( self, call: ::Request, @@ -176,12 +175,8 @@ where crate::http::resources().auth.clone(), crate::secrets::source(py)?, ); - let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?; - let route = match cache { - Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)), - None => route, - }; - Ok(route.machine(request, cache_options.policy)) + crate::cache::ensure_cache_disabled(py, arguments, call_type)?; + Ok(route.machine(request, CachePolicy::default())) }, host, hooks, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 27ce8c943f0..c975ac339c2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -24,34 +24,27 @@ use crate::{ ArgumentError, optional_timeout, project_optional_fields, public_response, python_timeout_seconds, }, - python_settings::missing_module, + python_settings::PythonSettings, routes::RequestView, }; const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; +const MODEL_CAPABILITIES_MODULE: &str = "litellm.rust_bridge.host.model_capabilities"; fn merge_headers( forwarded: Option>, extra_headers: Option>, ) -> Option> { - let merged: Map = forwarded - .into_iter() - .flatten() - .chain(extra_headers.into_iter().flatten()) - .collect(); + let merged = litellm_inference_messages::merge_messages_header_sources( + forwarded.into_iter().chain(extra_headers), + ); (!merged.is_empty()).then_some(merged) } -/// Reads `litellm.`, the module global a param spec names, the way -/// `VertexBase.get_vertex_ai_project` reads `litellm.vertex_project`. fn module_global<'py>(py: Python<'py>, name: &str) -> PyResult>> { - let module = match py.import("litellm") { - Ok(module) => module, - Err(error) if missing_module(py, &error, "litellm")? => return Ok(None), - Err(error) => return Err(error), - }; - let value = module.getattr(name)?; - Ok((!value.is_none()).then_some(value)) + PythonSettings::ProviderDefaults + .read_or_unset(py)? + .map_or(Ok(None), |defaults| defaults.attribute(name)) } /// The caller's litellm params, read from the kwargs by the names the typed params declare, @@ -107,10 +100,19 @@ impl MessagesPythonHost { .map(|value| python_timeout_seconds(py, value.unbind())) .transpose()? .flatten(); - let custom_llm_provider = string("custom_llm_provider")?; - let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; let api_key = string("api_key")?; let api_base = string("api_base")?; + // Python's handler resolves a bare model name (`claude-sonnet-4-5`) through + // `get_llm_provider`'s model lists, which the native prefix lookup does not carry. + let custom_llm_provider = match string("custom_llm_provider")? { + Some(provider) => Some(provider), + None => py + .import(ROUTE_HOST_MODULE)? + .getattr("inferred_provider")? + .call1((&model, api_base.as_deref()))? + .extract()?, + }; + let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; let extra_headers = self.merged_headers(py, arguments)?; let provider_specific_header = self.provider_specific_header(py, arguments)?; let litellm_params = project_litellm_params(argument, |name| module_global(py, name))?; @@ -166,7 +168,7 @@ impl MessagesPythonHost { ) -> PyResult { let module = py.import(ROUTE_HOST_MODULE)?; let capabilities: MessagesModelCapabilities = from_py( - &py.import("litellm.rust_bridge.host.model_capabilities")? + &py.import(MODEL_CAPABILITIES_MODULE)? .getattr("anthropic_model_capabilities")? .call1((model, custom_llm_provider))?, )?; @@ -232,9 +234,13 @@ impl PythonBinding for MessagesPythonHost { fn encode_response( &mut self, py: Python<'_>, - response: Box, + response: litellm_inference_messages::MessagesBody, ) -> PyResult> { - public_response(py, ROUTE_HOST_MODULE, response.as_ref()) + public_response(py, ROUTE_HOST_MODULE, &response) + } + + fn deliver_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> { + self.cache.store_response(py, response) } fn encode_stream_head( diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index c788760b268..a9ae01bd7ee 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -50,9 +50,18 @@ fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyR ) } +/// Python's sync Messages handler raises for every native provider before it sends anything +/// or opens a logging object (`BaseLLMHTTPHandler.anthropic_messages_handler`), and Rust copies +/// that until both sides serve sync calls. +const SYNC_NOT_IMPLEMENTED: &str = "anthropic_messages_handler is not implemented for sync calls"; + #[pyfunction] -pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { - run_messages(py, call, false) +#[pyo3(signature = (call))] +pub(crate) fn messages(call: NativeCall<'_>) -> PyResult> { + drop(call); + Err(pyo3::exceptions::PyValueError::new_err( + SYNC_NOT_IMPLEMENTED, + )) } #[pyfunction] diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index c2c3fe51600..73a2922804a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,4 +1,3 @@ -use litellm_auth::SecretValue; use litellm_auth::VertexParams; use litellm_host_python::{from_py, present}; use litellm_inference_ocr::{ @@ -6,6 +5,7 @@ use litellm_inference_ocr::{ wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, }; use litellm_llms::base_llm::ocr::{error::Error, transformation::DROP_PARAMS_PARAM}; +use litellm_secret_value::SecretValue; use pyo3::{prelude::*, types::PyDict}; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index b797e334c6a..019f8452737 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -44,12 +44,9 @@ pub(super) fn project( let value = match (name.as_str(), value) { ("previous_response_id", serde_json::Value::String(id)) => { let decoded: String = py - .import("litellm.responses.utils")? - .getattr("ResponsesAPIRequestUtils")? - .call_method1( - "decode_previous_response_id_to_original_previous_response_id", - (id,), - )? + .import(super::ROUTE_HOST_MODULE)? + .getattr("original_previous_response_id")? + .call1((id,))? .extract()?; serde_json::Value::String(decoded) } diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index 9c045d65c13..cc4d85b4d26 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use host::ResponsesPythonHost; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, ScopedCache}; +use litellm_cache_response::CachePolicy; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_host::{call::HostedMachine, protocol::Protocol}; use litellm_inference_responses::{ResponsesRoute, route::Responses}; @@ -49,10 +49,6 @@ impl InferenceRoute for ResponsesRoute { Self::new(http, auth, secrets) } - fn with_cache(self, cache: ScopedCache) -> Self { - self.with_cache(cache) - } - fn machine( self, call: ::Request, diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 669c10e602a..a2afb42b540 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -2,9 +2,11 @@ use std::{future::Future, pin::Pin, sync::Arc}; use litellm_core_utils::settings::Lookup; use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_secret_value::SecretValue; use litellm_secrets::{ - Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, + Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, }; + use pyo3::{ exceptions::PyException, prelude::*, @@ -13,7 +15,6 @@ use pyo3::{ use super::error::{external_error, read_error}; -const HANDLER_MODULE: &str = "litellm.secret_managers.secret_manager_handler"; const DIAGNOSTICS_MODULE: &str = "litellm.rust_bridge.host.diagnostics"; /// A secret manager whose reads execute in Python: a custom manager, a legacy compatible @@ -59,8 +60,8 @@ impl PythonClient { None => kwargs.set_item("key_management_settings", py.None())?, } let result = py - .import(HANDLER_MODULE)? - .getattr("get_secret_from_manager")? + .import(super::SECRET_MANAGER_MODULE)? + .getattr("read_from_manager")? .call((), Some(&kwargs))?; if result.is_instance_of::() { result.extract().map(Some) @@ -131,7 +132,8 @@ mod tests { use litellm_host_python::PythonContext; - use super::{DIAGNOSTICS_MODULE, HANDLER_MODULE, PythonSecretManager, python_name}; + use super::{DIAGNOSTICS_MODULE, PythonSecretManager, python_name}; + use crate::secrets::SECRET_MANAGER_MODULE; use crate::secrets::python_error; /// `sys.modules` is interpreter-global, so tests that install or rely on the handler module @@ -167,12 +169,12 @@ class Manager: raise failure manager = Manager() import sys, types -for name in ('litellm', 'litellm.secret_managers'): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'): sys.modules.setdefault(name, types.ModuleType(name)) -handler = sys.modules.setdefault('litellm.secret_managers.secret_manager_handler', types.ModuleType('litellm.secret_managers.secret_manager_handler')) -def get_secret_from_manager(**kwargs): +handler = sys.modules.setdefault('litellm.rust_bridge.host.secret_manager', types.ModuleType('litellm.rust_bridge.host.secret_manager')) +def read_from_manager(**kwargs): return kwargs['client'].sync_read_secret(kwargs['secret_name']) -handler.get_secret_from_manager = get_secret_from_manager +handler.read_from_manager = read_from_manager ", Some(&locals), Some(&locals), @@ -315,7 +317,7 @@ sys.modules.setdefault('litellm._logging', logging) }); } - /// Installs a fake `get_secret_from_manager` that records its kwargs, runs `body`, and + /// Installs a fake `read_from_manager` that records its kwargs, runs `body`, and /// removes the fake handler again; parent package stubs persist for concurrent tests. /// Callers hold `handler_guard` before attaching so the GIL is never held while waiting on it. fn with_fake_handler<'py>(py: Python<'py>, body: impl FnOnce(&Bound<'py, PyDict>)) { @@ -323,16 +325,16 @@ sys.modules.setdefault('litellm._logging', logging) py.run( c" import sys, types -previous_handler = sys.modules.get('litellm.secret_managers.secret_manager_handler') +previous_handler = sys.modules.get('litellm.rust_bridge.host.secret_manager') calls = [] -def get_secret_from_manager(**kwargs): +def read_from_manager(**kwargs): calls.append(kwargs) return 'handled-' + kwargs['secret_name'] -handler = types.ModuleType('litellm.secret_managers.secret_manager_handler') -handler.get_secret_from_manager = get_secret_from_manager -for name in ('litellm', 'litellm.secret_managers'): +handler = types.ModuleType('litellm.rust_bridge.host.secret_manager') +handler.read_from_manager = read_from_manager +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host'): sys.modules.setdefault(name, types.ModuleType(name)) -sys.modules['litellm.secret_managers.secret_manager_handler'] = handler +sys.modules['litellm.rust_bridge.host.secret_manager'] = handler ", Some(&locals), Some(&locals), @@ -342,9 +344,9 @@ sys.modules['litellm.secret_managers.secret_manager_handler'] = handler py.run( c" if previous_handler is None: - sys.modules.pop('litellm.secret_managers.secret_manager_handler', None) + sys.modules.pop('litellm.rust_bridge.host.secret_manager', None) else: - sys.modules['litellm.secret_managers.secret_manager_handler'] = previous_handler + sys.modules['litellm.rust_bridge.host.secret_manager'] = previous_handler ", Some(&locals), Some(&locals), @@ -364,7 +366,7 @@ else: with_fake_handler(py, |locals| { locals.set_item("expression", expression).unwrap(); py.run( - c"handler.get_secret_from_manager = lambda **kwargs: eval(expression)", + c"handler.read_from_manager = lambda **kwargs: eval(expression)", Some(locals), Some(locals), ) @@ -415,7 +417,7 @@ else: reader.client.read(py, "API_KEY").unwrap().as_deref(), Some("handled-API_KEY") ); - assert!(py.import(HANDLER_MODULE).is_ok()); + assert!(py.import(SECRET_MANAGER_MODULE).is_ok()); let calls = locals.get_item("calls").unwrap().unwrap(); let call = calls.get_item(0).unwrap().cast_into::().unwrap(); assert!(call.get_item("client").unwrap().unwrap().is(&client)); diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs index 7afd06e2373..c1203dcbd89 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -1,7 +1,9 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets::{SecretManager, SecretManagerState}; -use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, SecretValue}; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem}; + use pyo3::prelude::*; use serde_json::Value; diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs index 57042e6a32e..0840ecdfc6c 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mod.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -21,6 +21,7 @@ use resolved::ResolvedSecrets; use crate::{coercion::FieldSpec, python_settings::PythonSettings}; +const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager"; const NATIVE: FieldSpec = FieldSpec::new("native", |field| field.schema_bool()); /// Where a Rust route reads provider secrets from. Python's `get_secret_str` when the settings diff --git a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs index b04060651ac..47cd36c4015 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs @@ -2,7 +2,7 @@ use super::operations::{PythonMutationError, PythonMutationResponse}; use litellm_error::Report; use litellm_host_python::{json_loads, to_py}; use litellm_secrets::cyberark; -use pyo3::{prelude::*, types::PyDict}; +use pyo3::prelude::*; pub(super) fn mutation_value( result: Result, @@ -62,14 +62,11 @@ fn cyberark_message( http_message(py, "POST", url, status)? } cyberark::Error::Http(source) if authentication => match os_error_code(&source) { - Some(code) => { - let reason = py.import("os")?.getattr("strerror")?.call1((code,))?; - py.import("builtins")? - .getattr("OSError")? - .call1((code, reason))? - .str()? - .extract()? - } + Some(code) => py + .import(super::SECRET_MANAGER_MODULE)? + .getattr("os_error_message")? + .call1((code,))? + .extract()?, None => Report::new(&cyberark::Error::Http(source)).to_string(), }, cyberark::Error::Http(source) if source.is_connect() => { @@ -103,13 +100,8 @@ pub(super) fn http_message( url: &str, status: u16, ) -> PyResult { - let httpx = py.import("httpx")?; - let request = httpx.getattr("Request")?.call1((method, url))?; - let kwargs = PyDict::new(py); - kwargs.set_item("request", request)?; - let response = httpx.getattr("Response")?.call((status,), Some(&kwargs))?; - match response.call_method0("raise_for_status") { - Err(error) => error.value(py).str()?.extract(), - Ok(_) => Ok(format!("HTTP {status}")), - } + py.import(super::SECRET_MANAGER_MODULE)? + .getattr("http_status_message")? + .call1((method, url, status))? + .extract() } diff --git a/litellm-rust/crates/python-bridge/src/secrets/operations.rs b/litellm-rust/crates/python-bridge/src/secrets/operations.rs index b458e3f34e8..027da1e4718 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/operations.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/operations.rs @@ -81,7 +81,7 @@ impl From for PythonMutationError { pub(super) async fn write_python_provider( manager: &SecretManager, name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, ) -> Result { match manager { SecretManager::Cyberark(client) => { @@ -128,7 +128,7 @@ pub(super) async fn rotate_python_provider( manager: &SecretManager, current_name: &str, new_name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, ) -> Result { match manager { SecretManager::Cyberark(client) => { @@ -186,7 +186,7 @@ pub(super) enum PythonMutationResponse { pub(super) async fn write_python_provider_with_context( manager: &SecretManager, name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, context: &litellm_secrets_types::SecretWriteContext, ) -> Result { if let (SecretManager::HashicorpVault(client), SecretOperationContext::Hashicorp(operation)) = @@ -235,7 +235,7 @@ pub(super) async fn rotate_python_provider_with_context( manager: &SecretManager, current_name: &str, new_name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, context: &SecretOperationContext, ) -> Result { if let (SecretManager::HashicorpVault(client), SecretOperationContext::Hashicorp(context)) = diff --git a/litellm-rust/crates/python-bridge/src/secrets/provider.rs b/litellm-rust/crates/python-bridge/src/secrets/provider.rs index ff10f596311..55343fb63bb 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/provider.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/provider.rs @@ -1,7 +1,9 @@ use std::time::Duration; use super::operations::PythonReadRequest; -use litellm_secrets::{KeyManagementSystem, SecretValue}; +use litellm_secret_value::SecretValue; +use litellm_secrets::KeyManagementSystem; + use litellm_secrets_types::{ AwsOperationContext, CyberarkOperationContext, GoogleOperationContext, HashicorpOperationContext, SecretOperationContext, diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index 2f040bff773..a92cf9a497a 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -2,13 +2,13 @@ use std::sync::Arc; use futures_util::future::BoxFuture; use litellm_host_python::{PythonContext, attach_blocking}; -use litellm_secrets::{Error, SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::{Error, source::SecretSource}; + use pyo3::prelude::*; use super::error::external_error; -const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager"; - /// Reads each secret through the host module's `read_secret_str`, so the configured manager, the key /// management settings and the environment fallback behave exactly as they do in Python. pub(super) struct PythonSecrets { @@ -19,7 +19,7 @@ pub(super) struct PythonSecrets { impl PythonSecrets { pub(super) fn new(py: Python<'_>) -> PyResult { Ok(Self::reading_with( - py.import(SECRET_MANAGER_MODULE)? + py.import(super::SECRET_MANAGER_MODULE)? .getattr("read_secret_str")? .unbind(), PythonContext::capture(py)?, diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index deb07175e7a..752b533886e 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -4,10 +4,9 @@ use futures_util::future::BoxFuture; use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::PythonContext; use litellm_http::Client; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; -use litellm_secrets::{ - Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver, SecretValue, -}; +use litellm_secrets::{Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver}; use super::config::SecretManagerSnapshot; diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index b956ec297b6..8061d1682ed 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -12,8 +12,6 @@ use pyo3::{exceptions::PyAttributeError, prelude::*, types::PyDict}; use super::ManagerError; use crate::execution::{run_async_value, run_sync_value}; -const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager"; - #[derive(Clone, PartialEq)] struct Configuration { system: KeyManagementSystem, @@ -98,7 +96,7 @@ impl NativeSecretManager { return Ok(Some(native)); } let config = py - .import(SECRET_MANAGER_MODULE)? + .import(super::SECRET_MANAGER_MODULE)? .getattr("native_secret_manager_config")? .call1((client,))?; if config.is_none() { @@ -307,7 +305,7 @@ impl NativeSecretManager { super::operations::write_python_provider_with_context( &backend, &secret_name, - &litellm_secrets::SecretValue::new(secret_value), + &litellm_secret_value::SecretValue::new(secret_value), &context, ) .await, @@ -365,7 +363,7 @@ impl NativeSecretManager { &backend, ¤t_secret_name, &new_secret_name, - &litellm_secrets::SecretValue::new(new_secret_value), + &litellm_secret_value::SecretValue::new(new_secret_value), &context, ) .await, diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault.rs b/litellm-rust/crates/python-bridge/src/secrets/vault.rs index 30e3bd2c413..5f14727abaa 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault.rs @@ -134,31 +134,21 @@ fn timeout_message( String::new() }); } - let kwargs = pyo3::types::PyDict::new(py); - kwargs.set_item( - "message", - transport_timeout_message( - TimeoutPhase::Read, - context.timeout, - Some(elapsed.as_secs_f64()), - ), - )?; - kwargs.set_item("model", "default-model-name")?; - kwargs.set_item("llm_provider", "litellm-httpx-handler")?; - py.import("litellm.exceptions")? - .getattr("Timeout")? - .call((), Some(&kwargs))? - .str()? + let message = transport_timeout_message( + TimeoutPhase::Read, + context.timeout, + Some(elapsed.as_secs_f64()), + ); + py.import(super::SECRET_MANAGER_MODULE)? + .getattr("timeout_message")? + .call1((message,))? .extract() } fn response_text(py: Python<'_>, body: &[u8]) -> PyResult { - let kwargs = pyo3::types::PyDict::new(py); - kwargs.set_item("content", pyo3::types::PyBytes::new(py, body))?; - py.import("httpx")? - .getattr("Response")? - .call((200,), Some(&kwargs))? - .getattr("text")? + py.import(super::SECRET_MANAGER_MODULE)? + .getattr("response_text")? + .call1((pyo3::types::PyBytes::new(py, body),))? .extract() } diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs index 41658f0baf8..b99dff7aeac 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs @@ -1,9 +1,8 @@ use std::collections::HashMap; -use litellm_secrets::{ - SecretValue, - hashicorp::{HashicorpVault, RawOperationError}, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets::hashicorp::{HashicorpVault, RawOperationError}; + use litellm_secrets_types::{HashicorpOperationContext, SecretWriteContext}; use serde_json::value::RawValue; diff --git a/litellm-rust/crates/python-bridge/src/tokenizer.rs b/litellm-rust/crates/python-bridge/src/tokenizer.rs index 2026b2fdc42..11b21bd5673 100644 --- a/litellm-rust/crates/python-bridge/src/tokenizer.rs +++ b/litellm-rust/crates/python-bridge/src/tokenizer.rs @@ -38,16 +38,19 @@ use litellm_token_counter::huggingface::{ #[cfg(feature = "tiktoken")] use litellm_token_counter::tiktoken::{TiktokenTokenizer, Vocabulary}; +#[cfg(any(feature = "tiktoken", feature = "huggingface"))] +const HOST_MODULE: &str = "litellm.rust_bridge.host.tokenizers"; + #[cfg(feature = "tiktoken")] pub(crate) fn load_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { enter_native()?; - let resource: std::path::PathBuf = - PyModule::import(py, "litellm.litellm_core_utils.tokenizers")? - .getattr("__file__")? - .extract()?; + let resources: std::path::PathBuf = PyModule::import(py, HOST_MODULE)? + .getattr("tiktoken_resource_dir")? + .call0()? + .extract()?; release_gil(py, || { TiktokenTokenizer::from_cached_ranks(encoding, |file| { - std::fs::read_to_string(resource.with_file_name(file)) + std::fs::read_to_string(resources.join(file)) }) }) .map_err(|error| token_count_error_to_pyerr(error.into())) @@ -197,14 +200,9 @@ impl Tokenizer { #[cfg(feature = "huggingface")] { enter_native()?; - let kwargs = PyDict::new(py); - kwargs.set_item("repo_id", identifier)?; - kwargs.set_item("filename", "tokenizer.json")?; - kwargs.set_item("revision", revision)?; - kwargs.set_item("token", token)?; - let path: String = PyModule::import(py, "huggingface_hub")? - .getattr("hf_hub_download")? - .call((), Some(&kwargs))? + let path: String = PyModule::import(py, HOST_MODULE)? + .getattr("download_tokenizer")? + .call1((identifier, revision, token))? .extract()?; let json = release_gil(py, || std::fs::read_to_string(path)) .map_err(TokenizerError::Download)?; diff --git a/litellm-rust/crates/python-compat/benches/formats.rs b/litellm-rust/crates/python-compat/benches/formats.rs index f0b83cb1658..2f23b882e19 100644 --- a/litellm-rust/crates/python-compat/benches/formats.rs +++ b/litellm-rust/crates/python-compat/benches/formats.rs @@ -10,7 +10,7 @@ use std::{hint::black_box, time::Duration}; use criterion::{ - BatchSize, BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main, + BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main, measurement::WallTime, }; use litellm_python_compat::{Value, json, literal::literal_eval, pickle, repr::repr}; @@ -72,14 +72,9 @@ fn formats(c: &mut Criterion) { let mut group = c.benchmark_group("cached_completion"); text_formats(&mut group, &text, &value); binary_formats(&mut group, &value, &pickled); - // `from_json` consumes its input, so each iteration gets a freshly parsed one. group.throughput(Throughput::Bytes(dumped.len() as u64)); - group.bench_function("from_json", |bencher| { - bencher.iter_batched( - || serde_json::from_str::(&dumped).expect("dumps output parses"), - json::from_json, - BatchSize::SmallInput, - ) + group.bench_function("json_loads", |bencher| { + bencher.iter(|| json::loads(black_box(&dumped))) }); group.finish(); } diff --git a/litellm-rust/crates/python-compat/src/error.rs b/litellm-rust/crates/python-compat/src/error.rs index 6547907056e..db09925dde9 100644 --- a/litellm-rust/crates/python-compat/src/error.rs +++ b/litellm-rust/crates/python-compat/src/error.rs @@ -6,6 +6,8 @@ use crate::MAX_DEPTH; pub enum Error { #[error("malformed Python literal at byte {0}")] InvalidLiteral(usize), + #[error("malformed JSON at byte {0}")] + InvalidJson(usize), #[error("unhashable type: '{0}'")] Unhashable(&'static str), #[error("value nests deeper than {MAX_DEPTH} levels")] diff --git a/litellm-rust/crates/python-compat/src/json.rs b/litellm-rust/crates/python-compat/src/json.rs index d3436519ac4..37cc88fec78 100644 --- a/litellm-rust/crates/python-compat/src/json.rs +++ b/litellm-rust/crates/python-compat/src/json.rs @@ -1,14 +1,19 @@ -//! `json.dumps` with CPython's default options, and the JSON value `json.loads` returns. +//! `json.dumps` with CPython's default options, and `json.loads` with its default decoder. //! //! Defaults are `ensure_ascii=True`, `allow_nan=True`, separators `(", ", ": ")`, and no key -//! sorting. Python's `json.loads` is mapped by [`from_json`]; a `serde_json` number that does -//! not fit `i64` or `u64` arrives as a float, where Python would keep an `int`. +//! sorting on the way out, and `strict=True` with `NaN`, `Infinity` and `-Infinity` read as floats +//! on the way in. -use std::fmt::Write; +use std::{ + collections::{HashMap, hash_map::Entry}, + fmt::Write, + str::FromStr, +}; +use num_bigint::BigInt; use serde_json::{Map, Number}; -use crate::{Error, Value, repr::float_repr}; +use crate::{Error, MAX_DEPTH, Value, repr::float_repr}; /// `json.dumps(value)`. pub fn dumps(value: &Value) -> Result { @@ -44,6 +49,260 @@ pub fn from_json(value: serde_json::Value) -> Value { } } +/// `json.loads(text)`. Integers keep full precision, and a repeated key keeps its first position +/// with its last value, as `dict(pairs)` does. An escape that leaves a lone surrogate fails, since +/// a Rust `String` cannot hold one. +pub fn loads(text: &str) -> Result { + let mut decoder = Decoder { + text, + bytes: text.as_bytes(), + pos: 0, + depth: 0, + }; + decoder.skip_whitespace(); + let value = decoder.value()?; + decoder.skip_whitespace(); + if decoder.pos != decoder.bytes.len() { + return Err(decoder.error()); + } + Ok(value) +} + +/// CPython's `json.scanner` and `scanstring`, one pass over the bytes. +struct Decoder<'a> { + text: &'a str, + bytes: &'a [u8], + pos: usize, + depth: usize, +} + +impl Decoder<'_> { + fn error(&self) -> Error { + Error::InvalidJson(self.pos) + } + + fn peek(&self) -> Option { + self.bytes.get(self.pos).copied() + } + + fn eat(&mut self, token: &str) -> bool { + let found = self.bytes[self.pos..].starts_with(token.as_bytes()); + if found { + self.pos += token.len(); + } + found + } + + fn expect(&mut self, token: &str) -> Result<(), Error> { + if self.eat(token) { + Ok(()) + } else { + Err(self.error()) + } + } + + fn skip_whitespace(&mut self) { + while matches!(self.peek(), Some(b' ' | b'\t' | b'\n' | b'\r')) { + self.pos += 1; + } + } + + fn value(&mut self) -> Result { + match self.peek() { + Some(b'"') => self.string().map(Value::Str), + Some(b'[') => self.nested(Self::array), + Some(b'{') => self.nested(Self::object), + _ if self.eat("null") => Ok(Value::None), + _ if self.eat("true") => Ok(Value::Bool(true)), + _ if self.eat("false") => Ok(Value::Bool(false)), + _ if self.eat("NaN") => Ok(Value::Float(f64::NAN)), + _ if self.eat("Infinity") => Ok(Value::Float(f64::INFINITY)), + _ if self.eat("-Infinity") => Ok(Value::Float(f64::NEG_INFINITY)), + _ => self.number(), + } + } + + fn nested(&mut self, parse: fn(&mut Self) -> Result) -> Result { + self.depth += 1; + if self.depth > MAX_DEPTH { + return Err(Error::TooDeep); + } + let value = parse(self)?; + self.depth -= 1; + Ok(value) + } + + fn array(&mut self) -> Result { + self.expect("[")?; + self.skip_whitespace(); + let mut values = Vec::new(); + if self.eat("]") { + return Ok(Value::List(values)); + } + loop { + values.push(self.value()?); + self.skip_whitespace(); + if self.eat("]") { + return Ok(Value::List(values)); + } + self.expect(",")?; + self.skip_whitespace(); + } + } + + fn object(&mut self) -> Result { + self.expect("{")?; + self.skip_whitespace(); + let mut entries = Vec::new(); + let mut positions = HashMap::<_, usize>::new(); + if self.eat("}") { + return Ok(Value::Dict(entries)); + } + loop { + if self.peek() != Some(b'"') { + return Err(self.error()); + } + let key = self.string()?; + self.skip_whitespace(); + self.expect(":")?; + self.skip_whitespace(); + let value = self.value()?; + match positions.entry(key) { + Entry::Occupied(position) => entries[*position.get()].1 = value, + Entry::Vacant(position) => { + entries.push((Value::Str(position.key().clone()), value)); + position.insert(entries.len() - 1); + } + } + self.skip_whitespace(); + if self.eat("}") { + return Ok(Value::Dict(entries)); + } + self.expect(",")?; + self.skip_whitespace(); + } + } + + /// `scanstring` with `strict=True`: control characters must be escaped. + fn string(&mut self) -> Result { + self.expect("\"")?; + let mut out = String::new(); + loop { + let start = self.pos; + while self + .peek() + .is_some_and(|byte| byte != b'"' && byte != b'\\' && byte >= 0x20) + { + self.pos += 1; + } + out.push_str(&self.text[start..self.pos]); + match self.peek() { + Some(b'"') => { + self.pos += 1; + return Ok(out); + } + Some(b'\\') => { + self.pos += 1; + out.push(self.escape()?); + } + _ => return Err(self.error()), + } + } + } + + fn escape(&mut self) -> Result { + let escaped = match self.peek() { + Some(b'"') => '"', + Some(b'\\') => '\\', + Some(b'/') => '/', + Some(b'b') => '\u{8}', + Some(b'f') => '\u{c}', + Some(b'n') => '\n', + Some(b'r') => '\r', + Some(b't') => '\t', + Some(b'u') => { + self.pos += 1; + return self.unicode_escape(); + } + _ => return Err(self.error()), + }; + self.pos += 1; + Ok(escaped) + } + + /// `\uXXXX`, joining a high surrogate with the low one escaped right after it. + fn unicode_escape(&mut self) -> Result { + let unit = self.hex4()?; + let code = if (0xd800..0xdc00).contains(&unit) { + self.expect("\\u")?; + let low = self.hex4()?; + if !(0xdc00..0xe000).contains(&low) { + return Err(self.error()); + } + 0x10000 + ((unit - 0xd800) << 10) + (low - 0xdc00) + } else { + unit + }; + char::from_u32(code).ok_or_else(|| self.error()) + } + + fn hex4(&mut self) -> Result { + let digits = self + .text + .get(self.pos..self.pos + 4) + .filter(|digits| digits.bytes().all(|byte| byte.is_ascii_hexdigit())) + .ok_or_else(|| self.error())?; + self.pos += 4; + u32::from_str_radix(digits, 16).map_err(|_| self.error()) + } + + /// `NUMBER_RE`: `(-?(?:0|[1-9][0-9]*))(\.[0-9]+)?([eE][-+]?[0-9]+)?`, an `int` unless a + /// fraction or exponent matched. + fn number(&mut self) -> Result { + let start = self.pos; + self.eat("-"); + match self.peek() { + Some(b'0') => self.pos += 1, + Some(b'1'..=b'9') => self.digits(), + _ => return Err(Error::InvalidJson(start)), + } + let integer_end = self.pos; + if self.peek() == Some(b'.') && self.bytes.get(self.pos + 1).is_some_and(u8::is_ascii_digit) + { + self.pos += 1; + self.digits(); + } + if matches!(self.peek(), Some(b'e' | b'E')) { + let exponent = self.pos; + self.pos += 1; + if matches!(self.peek(), Some(b'+' | b'-')) { + self.pos += 1; + } + if self.peek().is_some_and(|byte| byte.is_ascii_digit()) { + self.digits(); + } else { + self.pos = exponent; + } + } + let text = &self.text[start..self.pos]; + if self.pos == integer_end { + BigInt::from_str(text) + .map(Value::Int) + .map_err(|_| Error::InvalidJson(start)) + } else { + f64::from_str(text) + .map(Value::Float) + .map_err(|_| Error::InvalidJson(start)) + } + } + + fn digits(&mut self) { + while self.peek().is_some_and(|byte| byte.is_ascii_digit()) { + self.pos += 1; + } + } +} + /// `json.loads(json.dumps(value))` as a `serde_json` value: tuples become arrays and dict /// keys are coerced to strings as `json.dumps` does. Non-finite floats, which Python writes /// as `NaN` and `Infinity`, have no `serde_json` form and fail with [`Error::NonFiniteFloat`]. @@ -166,3 +425,86 @@ fn write_string(out: &mut String, text: &str) { } out.push('"'); } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::{dumps, loads}; + use crate::{Error, Value, repr::repr}; + + // Expected values are the repr of CPython 3.14 json.loads(text) + #[rstest] + #[case::nan("NaN", "nan")] + #[case::infinity("Infinity", "inf")] + #[case::negative_infinity("-Infinity", "-inf")] + #[case::non_finite_in_a_list("[NaN, Infinity, -Infinity]", "[nan, inf, -inf]")] + #[case::beyond_u64("18446744073709551616", "18446744073709551616")] + #[case::exponent_overflow("1e400", "inf")] + #[case::upper_exponent("1E+2", "100.0")] + #[case::negative_zero_int("-0", "0")] + #[case::negative_zero_float("-0.0", "-0.0")] + #[case::repeated_key(r#"{"a": 1, "b": 2, "a": 3}"#, "{'a': 3, 'b': 2}")] + #[case::surrogate_pair(r#""😀""#, "'😀'")] + #[case::upper_hex_surrogate_pair(r#""😀""#, "'😀'")] + #[case::escapes( + r#"{"k": [true, false, null, 1.5, "é\/"]}"#, + "{'k': [True, False, None, 1.5, 'é/']}" + )] + #[case::delete_is_not_a_control_character("\"\x7f\"", "'\\x7f'")] + #[case::surrounding_whitespace(" \n[1]\t", "[1]")] + #[case::empty_containers("[{}, []]", "[{}, []]")] + fn loads_reads_like_python(#[case] text: &str, #[case] expected: &str) { + assert_eq!(repr(&loads(text).unwrap()), expected); + } + + // CPython 3.14 json.loads raises JSONDecodeError on each + #[rstest] + #[case::lowercase_nan("nan")] + #[case::signed_nan("-NaN")] + #[case::plus_sign("+1")] + #[case::bare_minus("-")] + #[case::leading_zero("01")] + #[case::bare_point("1.")] + #[case::bare_exponent("1e")] + #[case::trailing_data("Infinity1")] + #[case::trailing_comma("[1,]")] + #[case::missing_comma("[1 2]")] + #[case::missing_colon(r#"{"a" 1}"#)] + #[case::unquoted_key("{1: 2}")] + #[case::raw_newline_in_string("\"a\nb\"")] + #[case::unknown_escape(r#""\x41""#)] + #[case::empty("")] + #[case::byte_order_mark("\u{feff}1")] + fn loads_rejects_what_python_rejects(#[case] text: &str) { + assert!(matches!(loads(text), Err(Error::InvalidJson(_)))); + } + + #[rstest] + fn a_lone_surrogate_is_a_known_limit() { + assert!(matches!(loads(r#""\ud800""#), Err(Error::InvalidJson(_)))); + } + + #[rstest] + fn loads_reads_back_the_non_finite_floats_dumps_writes() { + let written = dumps(&Value::List(vec![ + Value::Float(f64::NAN), + Value::Float(f64::INFINITY), + Value::Float(f64::NEG_INFINITY), + ])) + .unwrap(); + + assert_eq!(repr(&loads(&written).unwrap()), "[nan, inf, -inf]"); + } + + #[rstest] + fn loads_rejects_nesting_past_max_depth() { + let text = format!( + "{}{}", + "[".repeat(crate::MAX_DEPTH + 1), + "]".repeat(crate::MAX_DEPTH + 1) + ); + + assert!(matches!(loads(&text), Err(Error::TooDeep))); + } +} diff --git a/litellm-rust/crates/python-compat/src/lib.rs b/litellm-rust/crates/python-compat/src/lib.rs index b23aff409f2..2cc0eeac397 100644 --- a/litellm-rust/crates/python-compat/src/lib.rs +++ b/litellm-rust/crates/python-compat/src/lib.rs @@ -7,9 +7,10 @@ //! |---|---| //! | [`literal`] | `ast.literal_eval(text)` | //! | [`repr`] | `repr(value)` and `str(value)` | -//! | [`json`] | `json.dumps(value)`, and `json.loads(json.dumps(value))` as a JSON value | +//! | [`json`] | `json.dumps(value)` and `json.loads(text)` | //! | [`pickle`] | `pickle.loads(data)` and `pickle.dumps(value)` for plain data | //! | [`truthy`] | `bool(value)` | +//! | [`number`] | `int(text)`, `float(text)`, and `int(seconds)` as a `Duration` | //! //! [`Value`] is the closed data model these formats share. Live Python objects //! (descriptors, `__bool__`, `__str__`, callbacks) are out of scope: those belong to the @@ -20,10 +21,12 @@ //! (string hashes are randomized per process). Sets keep their literal order. //! - `str` values are Rust `String`s, so lone surrogates cannot be represented. //! - [`pickle::loads`] decodes `tuple`, `set`, and `frozenset` as lists. +//! - [`number`] reads ASCII digits only, while Python's `int` and `float` accept any Unicode decimal digit. mod error; pub mod json; pub mod literal; +pub mod number; pub mod pickle; pub mod repr; pub mod truthy; diff --git a/litellm-rust/crates/python-compat/src/number.rs b/litellm-rust/crates/python-compat/src/number.rs new file mode 100644 index 00000000000..6e2f067ae25 --- /dev/null +++ b/litellm-rust/crates/python-compat/src/number.rs @@ -0,0 +1,121 @@ +use std::{str::FromStr, time::Duration}; + +use num_bigint::BigInt; + +/// `int(text)` in base 10, or `None` where Python raises `ValueError`. +pub fn int(text: &str) -> Option { + parse(text) +} + +/// `float(text)`, or `None` where Python raises `ValueError`. +pub fn float(text: &str) -> Option { + parse(text) +} + +/// Why [`seconds`] refused a value, as `int(value)` would. Messages are CPython's. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum SecondsError { + /// `OverflowError`. Finite values past `Duration::MAX` land here too, where Python builds a big int. + #[error("cannot convert float infinity to integer")] + Overflow, + /// `ValueError`. + #[error("cannot convert float NaN to integer")] + NotANumber, +} + +/// A float of seconds as a `Duration`, refusing what `int(value)` refuses. Finite negative values clamp to zero. +pub fn seconds(value: f64) -> Result { + match value { + value if value.is_nan() => Err(SecondsError::NotANumber), + value if value.is_infinite() => Err(SecondsError::Overflow), + value => Duration::try_from_secs_f64(value.max(0.0)).map_err(|_| SecondsError::Overflow), + } +} + +fn parse(text: &str) -> Option { + let text = text.trim(); + let digit_at = |index: Option| { + index + .and_then(|index| text.as_bytes().get(index)) + .is_some_and(u8::is_ascii_digit) + }; + let underscores_between_digits = text + .match_indices('_') + .all(|(index, _)| digit_at(index.checked_sub(1)) && digit_at(Some(index + 1))); + if !underscores_between_digits { + return None; + } + text.replace('_', "").parse().ok() +} + +#[cfg(test)] +mod tests { + use num_bigint::BigInt; + use rstest::rstest; + + use std::time::Duration; + + use super::{SecondsError, float, int, seconds}; + + // Expected values are CPython 3.14 int(text) + #[rstest] + #[case::trailing_newline("2\n", Some(2))] + #[case::surrounding_spaces_and_underscore(" 6_380 ", Some(6380))] + #[case::plus_sign("+5", Some(5))] + #[case::negative_zero("-0", Some(0))] + #[case::leading_zeros("007", Some(7))] + #[case::unicode_spaces("\u{a0}7\u{2003}", Some(7))] + #[case::beyond_i64("99999999999999999999", Some(99_999_999_999_999_999_999_i128))] + #[case::double_underscore("1__0", None)] + #[case::leading_underscore("_1", None)] + #[case::trailing_underscore("1_", None)] + #[case::decimal_point("5.0", None)] + #[case::information_separator("\x1c7", None)] + #[case::empty("", None)] + #[case::sign_only("+", None)] + #[case::non_ascii_digit_is_a_known_limit("\u{667}", None)] + fn int_parses_like_python(#[case] text: &str, #[case] expected: Option) { + assert_eq!(int(text), expected.map(BigInt::from)); + } + + // Expected values are CPython 3.14 float(text) + #[rstest] + #[case::underscore_and_newline(" 1_0.5\n", Some(10.5))] + #[case::exponent("1e3", Some(1000.0))] + #[case::leading_point(".5", Some(0.5))] + #[case::trailing_point("5.", Some(5.0))] + #[case::plus_sign("+1.5", Some(1.5))] + #[case::infinity_word("Infinity", Some(f64::INFINITY))] + #[case::infinity_with_space("infinity ", Some(f64::INFINITY))] + #[case::underscore_before_exponent("1_e3", None)] + #[case::underscore_after_exponent("1e_3", None)] + #[case::underscore_after_point("1._5", None)] + #[case::underscore_after_word("inf_", None)] + #[case::hexadecimal("0x10", None)] + #[case::bare_exponent("1.5e", None)] + fn float_parses_like_python(#[case] text: &str, #[case] expected: Option) { + assert_eq!(float(text), expected); + } + + #[rstest] + fn float_reads_a_signed_nan_like_python() { + assert!(float("-NaN").is_some_and(f64::is_nan)); + } + + // CPython 3.14: int(inf) and int(-inf) raise OverflowError, int(nan) raises ValueError + #[rstest] + #[case::fraction(1.5, Ok(Duration::from_millis(1500)))] + #[case::zero(0.0, Ok(Duration::ZERO))] + #[case::negative_clamps_to_zero(-3.0, Ok(Duration::ZERO))] + #[case::infinity(f64::INFINITY, Err(SecondsError::Overflow))] + #[case::negative_infinity(f64::NEG_INFINITY, Err(SecondsError::Overflow))] + #[case::past_duration_max(1e20, Err(SecondsError::Overflow))] + #[case::nan(f64::NAN, Err(SecondsError::NotANumber))] + #[case::negative_nan(-f64::NAN, Err(SecondsError::NotANumber))] + fn seconds_refuses_what_int_refuses( + #[case] value: f64, + #[case] expected: Result, + ) { + assert_eq!(seconds(value), expected); + } +} diff --git a/litellm-rust/crates/python-compat/tests/fixtures.rs b/litellm-rust/crates/python-compat/tests/fixtures.rs index 9c202fdc5f1..3f9c6fd7010 100644 --- a/litellm-rust/crates/python-compat/tests/fixtures.rs +++ b/litellm-rust/crates/python-compat/tests/fixtures.rs @@ -76,6 +76,11 @@ const KNOWN: &[(&str, &str, &str)] = &[ "literal_eval", "deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it", ), + ( + "nested_150", + "json.loads", + "deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it", + ), ( "nested_150", "pickle.loads", @@ -341,3 +346,19 @@ fn to_json_matches_python_json_round_trip(fixtures: &Fixtures) { } mismatches.finish(); } + +/// CPython gives back the same text from `json.dumps(json.loads(text))` for every text its +/// `json.dumps` wrote here, non-finite floats included, so `loads` must too. +#[rstest] +fn loads_reads_back_what_python_json_dumps_wrote(fixtures: &Fixtures) { + let mut mismatches = Mismatches::default(); + for row in &fixtures.rows { + let Some(text) = &row.json else { continue }; + let actual = match json::loads(text).and_then(|value| json::dumps(&value)) { + Ok(text) => text, + Err(error) => format!("error: {error}"), + }; + mismatches.check(row, "json.loads", text, &actual); + } + mismatches.finish(); +} diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index c73a4ed8c4d..6135919c539 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -1,12 +1,12 @@ - The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace - It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default - Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route -- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`) +- Every connection sets TCP_NODELAY, like redis-py's `socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)`. Set every redis-rs timeout and retry explicitly, citing the redis-py 5.3.1 or `litellm/` source next to each constant. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's. `Timeouts` takes `socket_timeout` and `socket_connect_timeout`, and the connect timeout falls back to the socket timeout like redis-py. Standalone reconnects on the next command with no retry of the failed one, matching redis-py's `Retry(NoBackoff(), 0)`. redis-rs keeps a socket whose command timed out, so `Renewing` in `connection.rs` swaps in a fresh `ConnectionManager`, or a fresh cluster connection like redis-py's cluster `aclose`, after a timeout. Without it a socket that went silent times out every command until TCP gives up on it. Cluster keeps redis-py's 3 `cluster_error_retry_attempts` and waits a flat 250ms between them, which is redis-py's CLUSTERDOWN sleep. redis-py's jittered per-node reconnect backoff has no redis-rs equivalent. The retry set also differs, redis-rs retries LOADING and READONLY internally, while redis-py 5.3.1 raises them. Do not claim exact retry parity - Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values ## Errors - `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it -- Classify in `From for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error +- Classify with `Error::classify(error, deployment)`, never a bare `From`, because redis-py (5.3.1) decides by topology. `Transport` is what Python's `_redis_health_error_types` counts: LOADING everywhere, CLUSTERDOWN and MASTERDOWN only on cluster (only `ClusterParser` maps them), READONLY only under Sentinel (where `SentinelManagedConnection` turns it into `ConnectionError` and rediscovers), plus NOAUTH, WRONGPASS and the `ERR` messages redis-py maps to `ConnectionError`, matched exactly like `parse_error`. TRYAGAIN is `Server`, since the cluster client already retried it. Each code gets its own rstest case built from a parsed RESP error - A `Server` error never drops the connection - `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome - Never swallow an error or turn it into a miss. The caller picks the fallback @@ -17,27 +17,61 @@ - Every public call (each command, `eval`, one `run` of a pipeline) takes one admission through `guarded` / `admit` in `client.rs`. Never call the connection around it - `Transport` counts as a failure (`Timeout` when `is_timeout()`, else `Hard`). `Server` and `Decode` count as nothing, since Redis answered and the request was bad. A probe that gets one still reopens, through the dropped-probe rule in `circuit-breaker`. A pipeline counts one failure if any op came back `Transport`, else one success - An open breaker returns `Error::BreakerOpen` without touching the socket -- Python doesn't count MASTERDOWN, TRYAGAIN or READONLY as health failures. Here they are `Transport` and do count ## Namespace - Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace` - The rule never changes, because Python and Rust pods share keys in the same Redis ## Commands -- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- A command is a value, one constructor in `commands.rs` returning `Command`: its alone and pipelined forms and its reply decode. `RedisClient::send`, `Pipeline::with` and `RedisBatch::add` are the only entry points, and each applies its namespace to the keys once. A new command is one constructor, never a method on an executor. `ping` and `subscribe` stay methods. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- A store attaches its value codec with `Command::map`, never a keyspace codec inside this crate - On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test -- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`) +- Every command is the one `RedisCache` sends, argument for argument, so `MONITOR` on a mixed fleet shows one shape. A difference from main's Python is a defect even when the result is the same +- TTLs are `Duration` and are sent in whole seconds (`SET ... EX`, `EXPIRE`, `TTL`), like redis-py's `int(ttl)`. `set_nx` sends `EX` before `NX`, the order redis-py builds +- Numbers go out as `Amount`, written the way redis-py encodes them: an `int` with `str()` and a `float` with `repr()` from `python-compat` +- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` alone is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. Pipelined it is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when an expiry is given, `KeepExisting` lowered to `Refresh`. `Plan` keeps only the cases whose alone and pipelined forms differ ## Scripts and pipelines - A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production -- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async` +- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline a script is sent as `EVALSHA`, like `RedisBatch`. A `NOSCRIPT` reply means that op ran nothing, so it alone runs again after the pipeline through `invoke_async` (`EVALSHA`, `SCRIPT LOAD`, `EVALSHA`), like `RedisBatch`'s `_settle_alone`. It then lands after the ops around it, as it does in Python +- `Plan::Script` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice +- `RedisBatch::pipelined` is one `Pipeline` whose ops answer as typed `Declared`s. A store that sends its own pipeline uses it rather than decoding `run`'s raw `Value`s - A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip -- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order -- Pipeline ops take one key each. A multi-key op would span slots inside one group -- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once +- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key and runs the slots concurrently, returning results in input order. Within a slot, runs of idempotent ops (GET, SET, DEL, EXPIRE) share a request and every script op gets its own, sent in order. redis-rs retries a cluster request whole, so a script op batched with others would be applied twice when only a neighbour failed +- Pipeline ops take one key each. A multi-key op would span slots inside one group. An op can send more than one command for that key (`INCRBYFLOAT` then `EXPIRE`), and settles from the first reply, failing if any of its replies failed +- Command constructors live in `commands.rs` and every executor sends them, so a command is written once -## Tests +## Pub/sub +- Classic `PUBLISH` / `SUBSCRIBE`, never sharded, since redis-py publishes and subscribes through one plain client and cluster broadcasts classic messages to every node. Payloads are raw bytes, and the codec belongs to the store that owns the channel +- Channel names never get the namespace. Python builds `{namespace}:{channel}` itself, so the store does too +- `publish` is a command like any other, guarded by the breaker and reported to events +- `subscribe` returns a `Stream` that opens its own RESP2 `PubSub` connection when first polled, like redis-py's `PubSub`, to the endpoint, the first reachable cluster node, or the master Sentinel names at that moment. It never touches the breaker or events, like Python's subscriber, which reads off its own connection +- After a lost connection or a failed subscribe it waits and resubscribes, 5s doubling to 60s and reset by a successful subscribe, like `AuthCacheInvalidationSubscriber._run`. Messages published while it is down are lost. The connection has no idle ping, so a socket that dies silently is only noticed by TCP +- redis-rs's `PubSub` takes no credentials provider, so a provider is asked for its current credentials once when the subscription connects, which is when redis-py's pubsub connection asks it. The connect is bounded by the connect timeout + +## Batches +- `RedisBatch` is litellm's `RedisBatch` (`litellm/caching/redis_batch.py`). It keeps each declared op as data and, on a standalone or Sentinel Redis, sends them later as one `Pipeline`, so it inherits every pipeline rule above. On a cluster it never pipelines and runs each op on its own, concurrently, the way `_settle_alone` calls the `RedisCache` method behind it. An increment there keeps a TTL the key already has, as `async_increment` does, where the pipelined one refreshes it like `_Increment` +- `mget` is one `MGET` of the keys in order with repeats read once, like `_MGet`, and answers one value per key it was given, in that order. Use it instead of one `get` per key whenever Python reads several keys in one call +- It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local +- `RedisBatch::alone` runs every op on its own like the cluster path, so a fresh alone batch is "send it now" and `Declared` is the one reply handle on every route +- Each op returns a `Declared` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once +- An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis +- `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client - Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes +- Tests that only need the public API live in `tests/.rs`, with client builders in `tests/support`. Inline `#[cfg(test)]` modules are only for private items (`error.rs`, `namespace.rs`) +- Observe connections from the server side (`RedisDeployment::clients`, `cli_on_every_node`), never by reaching into `RedisClient`'s connection - Every command test runs on each topology through `#[values(Topology::Standalone, Topology::Cluster)]`. A test is single-topology only when the behavior exists on one (a `db` number has no cluster meaning) -- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables +- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables. A socket that never answers again comes from `TcpProxy::blackhole_open_connections`, since `thaw` revives the same socket - Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline + +## TLS, credentials and Sentinel +- `TlsOptions` follows redis-py 5.3.1. It verifies the certificate by default, skips the hostname check by default (`ssl_check_hostname=False`), and trusts the system roots plus `ca_certs` +- rustls only reads unencrypted keys, so a PKCS#8 `ENCRYPTED PRIVATE KEY` is decrypted with `key_password` before it reaches redis-rs. A legacy `Proc-Type: 4,ENCRYPTED` PEM key is refused with `KeyError::LegacyEncryption` rather than decrypted by hand-rolled OpenSSL key derivation +- redis-rs loads the system trust store for every TLS connection and fails if that fails, even with `ca_certs` set. Where the macOS keychain is blocked (agent sandboxes), run the TLS tests with `SSL_CERT_FILE=/etc/ssl/cert.pem` +- `Credentials::Static` are sent at connect, and URL userinfo wins over them, like redis-py's `from_url` +- A url's `?protocol=` is ignored and every connection speaks RESP2, which every reply decode assumes +- `Credentials::Provider` wraps a redis-rs `StreamingCredentialsProvider`. redis-rs takes the first item at connect and re-`AUTH`s live connections on every later item, on each cluster node. Providers push, redis-rs never polls. With a provider, redis-rs overwrites URL userinfo at connect, which matches Python stripping `user:pass@` from the URL +- Sentinel discovers the master on each new connection and drops it on a transport failure or a READONLY reply, to one command or inside a pipeline, like redis-py's `SentinelConnectionPool` +- redis-rs sends `ROLE` to the master a Sentinel names before using it, which redis-py never does. That probe gets the data credentials (a provider's current pair) and TLS, so the Sentinel client is built on each discovery. redis-rs gives the probe no way to skip the hostname check, and drops the client certificate when the certificate is not verified, so a TLS master whose certificate does not name the announced address, or that requires a client certificate while verification is off, fails discovery where Python connects +- Sentinel monitors get `sentinel_password`, the data username when that password is set, and the TLS mode and certificates. Python's `_sentinel_auth_kwargs` forwards every other connection kwarg too, and it sends the username even without a sentinel password +- A `Debug` impl never shows a secret. Follow `.agents/skills/rust-redaction` diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index eab57b36423..d3d3226abc6 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -6,12 +6,28 @@ license.workspace = true repository.workspace = true [dependencies] +tracing.workspace = true litellm-circuit-breaker.workspace = true +litellm-secret-value.workspace = true +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true futures-util.workspace = true -redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } +redis = { workspace = true, features = [ + "cluster-async", + "connection-manager", + "sentinel", + "tls-rustls-insecure", + "token-based-authentication", + "tokio-rustls-comp", +] } +pem = "4.0.0" +pkcs8 = { version = "0.10.2", features = ["3des", "encryption", "std"] } +rustls-native-certs.workspace = true +socket2 = "0.6.5" thiserror.workspace = true +tokio = { workspace = true, features = ["rt", "sync", "time"] } [dev-dependencies] litellm-redis-testing.workspace = true +rustls.workspace = true rstest.workspace = true -tokio.workspace = true diff --git a/litellm-rust/crates/redis-client/src/batch.rs b/litellm-rust/crates/redis-client/src/batch.rs new file mode 100644 index 00000000000..84fc1643424 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/batch.rs @@ -0,0 +1,147 @@ +use std::{ + io, + sync::{Arc, Mutex, MutexGuard}, +}; + +use futures_util::future::join_all; +use redis::{RedisError, Value}; +use tokio::sync::oneshot::{self, error::TryRecvError}; + +use crate::{Command, Error, Pipeline, RedisClient, commands::Op}; + +/// Ops declared by many callers and sent together, like litellm's `RedisBatch`: one pipeline on a +/// standalone or Sentinel Redis, and each op on its own on a cluster or in an alone batch. +#[derive(Clone)] +pub struct RedisBatch(Arc); + +#[derive(Clone, Copy)] +enum Mode { + /// litellm's `RedisBatch`: one pipeline, or each op on its own on a cluster. + Batch, + /// Each op on its own, the way the `RedisCache` method behind it sends it. + Alone, + /// One redis-py pipeline on every topology, split by slot on a cluster. + Pipeline, +} + +struct Shared { + client: RedisClient, + mode: Mode, + pending: Mutex, + flushing: tokio::sync::Mutex<()>, +} + +#[derive(Default)] +struct Pending { + ops: Vec, + replies: Vec>>, +} + +type Decode = Box Result + Send>; + +#[must_use] +pub struct Declared { + batch: RedisBatch, + reply: oneshot::Receiver>, + decode: Decode, +} + +impl RedisBatch { + pub(crate) fn new(client: RedisClient) -> Self { + Self::build(client, Mode::Batch) + } + + /// A batch whose every op goes out on its own, the way the `RedisCache` method behind it sends + /// it, so a fresh one is "send it now": the first `resolve` sends it. + pub fn alone(client: RedisClient) -> Self { + Self::build(client, Mode::Alone) + } + + /// A batch sent as one redis-py pipeline, like `async_increment_pipeline`, on every topology. + pub fn pipelined(client: RedisClient) -> Self { + Self::build(client, Mode::Pipeline) + } + + fn build(client: RedisClient, mode: Mode) -> Self { + Self(Arc::new(Shared { + client, + mode, + pending: Mutex::new(Pending::default()), + flushing: tokio::sync::Mutex::new(()), + })) + } + + pub fn add(&self, command: Command) -> Declared { + let (op, decode) = command.prepare(self.0.client.namespace()); + let (sender, reply) = oneshot::channel(); + let mut pending = self.pending(); + pending.ops.push(op); + pending.replies.push(sender); + Declared { + batch: self.clone(), + reply, + decode, + } + } + + pub async fn flush(&self) { + let _flushing = self.0.flushing.lock().await; + let Pending { ops, replies } = std::mem::take(&mut *self.pending()); + if replies.is_empty() { + return; + } + let client = &self.0.client; + let alone = match self.0.mode { + Mode::Alone => true, + Mode::Batch => client.is_cluster(), + Mode::Pipeline => false, + }; + let results = match alone { + true => join_all(ops.into_iter().map(|op| client.run_alone(op))).await, + false => { + client + .run(&Pipeline::of(client.namespace().clone(), ops)) + .await + } + }; + for (reply, result) in replies.into_iter().zip(results) { + let _ = reply.send(result); + } + } + + fn pending(&self) -> MutexGuard<'_, Pending> { + self.0 + .pending + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +impl Declared { + pub async fn resolve(mut self) -> Result { + match self.reply.try_recv() { + Ok(reply) => return reply.and_then(self.decode), + Err(TryRecvError::Closed) => return Err(abandoned()), + Err(TryRecvError::Empty) => {} + } + self.batch.flush().await; + self.settled().await + } + + pub async fn settled(self) -> Result { + let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned())); + reply.and_then(self.decode) + } + + /// Every reply in order, flushing their batch once if nobody has yet. + pub async fn all(declared: impl IntoIterator) -> Vec> { + join_all(declared.into_iter().map(Self::resolve)).await + } +} + +fn abandoned() -> Error { + Error::Transport(RedisError::from(io::Error::new( + io::ErrorKind::Interrupted, + "the batch flush carrying this op was cancelled", + ))) +} diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index a02b0847663..d9a8b3af7e8 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -1,21 +1,43 @@ use std::{ future::Future, + sync::{Arc, Mutex, PoisonError, Weak}, time::{Duration, Instant}, }; +use futures_util::stream::BoxStream; use litellm_circuit_breaker::{Admission, Breaker, Failure}; -use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value}; +use redis::{Cmd, ErrorKind, RedisError, ServerErrorKind, Value}; use crate::{ - ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands, + CallEvent, Command, ConnectError, Error, FailureKind, Namespace, Operation, Outcome, Pipeline, + RedisBatch, RedisConfig, RedisEvents, + commands::{self, Op, Plan, Trips}, connection::Connection, + pubsub::Subscriber, }; +/// One Redis client. Clones share the connection, the breaker and the subscriber. #[derive(Clone)] -pub struct RedisClient { - pub(crate) connection: Connection, +pub struct RedisClient(Arc); + +/// A handle that does not keep the client alive, for a registry that must not outlive its users. +#[derive(Clone)] +pub struct WeakRedisClient(Weak); + +impl WeakRedisClient { + pub fn upgrade(&self) -> Option { + self.0.upgrade().map(RedisClient) + } +} + +struct Shared { + connection: Connection, + subscriber: Subscriber, namespace: Namespace, breaker: Option, + events: Option>, + health_check: Option, + last_used: Mutex, } #[derive(Debug, PartialEq, Eq)] @@ -27,122 +49,152 @@ pub enum Ttl { impl RedisClient { pub fn new(config: RedisConfig) -> Result { - let connection = match &config.endpoint { - Endpoint::Url(url) => { - Connection::standalone(url.as_str().into_connection_info()?, config.timeout)? - } - Endpoint::Host(host) => { - Connection::standalone(host.connection_info()?, config.timeout)? - } - Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout)?, - }; - Ok(Self { + let connection = Connection::new( + &config.target, + &config.credentials, + config.timeouts, + config.keepalive, + )?; + Ok(Self(Arc::new(Shared { + subscriber: Subscriber::new( + &config.target, + &config.credentials, + config.timeouts, + config.keepalive, + &connection, + )?, connection, namespace: config.namespace, breaker: config.breaker, - }) + events: config.events, + health_check: config.health_check, + last_used: Mutex::new(Instant::now()), + }))) + } + + pub fn downgrade(&self) -> WeakRedisClient { + WeakRedisClient(Arc::downgrade(&self.0)) } pub async fn ping(&self) -> Result<(), Error> { - self.query(&commands::ping()).await - } - - pub async fn get(&self, key: &str) -> Result, Error> { - self.query(&commands::get(self.key(key))).await - } - - pub async fn mget(&self, keys: &[&str]) -> Result>, Error> { - if keys.is_empty() { - return Ok(Vec::new()); - } - self.query(&commands::mget(self.keys(keys))).await - } - - pub async fn set( - &self, - key: &str, - value: impl ToSingleRedisArg, - ttl: Duration, - ) -> Result<(), Error> { - self.query(&commands::set(self.key(key), value, ttl)).await - } - - pub async fn set_nx( - &self, - key: &str, - value: impl ToSingleRedisArg, - ttl: Duration, - ) -> Result { - let reply: Option = self - .query(&commands::set_nx(self.key(key), value, ttl)) - .await?; - Ok(reply.is_some()) - } - - pub async fn del(&self, keys: &[&str]) -> Result { - if keys.is_empty() { - return Ok(0); - } - self.query(&commands::del(self.keys(keys))).await - } - - pub async fn expire(&self, key: &str, ttl: Duration) -> Result { - self.query(&commands::expire(self.key(key), ttl)).await - } - - pub async fn ttl(&self, key: &str) -> Result { - let reply: i64 = self.query(&commands::ttl(self.key(key))).await?; - Ok(match reply { - -2 => Ttl::Missing, - -1 => Ttl::Persistent, - millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())), - }) - } - - pub async fn incr_by(&self, key: &str, delta: i64) -> Result { - self.query(&commands::incr_by(self.key(key), delta)).await - } - - pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result { - self.query(&commands::incr_by_float(self.key(key), delta)) - .await - } - - pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result { - self.query(&commands::rpush(self.key(key), values)).await - } - - pub async fn eval( - &self, - script: &Script, - keys: &[&str], - args: impl ToRedisArgs, - ) -> Result { - let keys = self.keys(keys); - let args = args.to_redis_args(); - self.guarded(async { - Ok(script - .invocation(&keys, &args) - .invoke_async(&mut self.connection.clone()) - .await?) + self.guarded(Operation::Command("PING"), async { + commands::ping() + .query_async(&mut self.0.connection.clone()) + .await }) .await } + /// Sends one command on its own, the way the `RedisCache` method behind it sends it. + pub async fn send(&self, command: Command) -> Result { + let (op, decode) = command.prepare(&self.0.namespace); + self.run_alone(op).await.and_then(decode) + } + + pub(crate) async fn run_alone(&self, op: Op) -> Result { + let Op { + operation, plan, .. + } = op; + self.guarded(operation, async { + let mut connection = self.0.connection.clone(); + match plan { + Plan::Reply(value) => Ok(value), + Plan::Commands(commands) => match <[Cmd; 1]>::try_from(commands) { + Ok([command]) => command.query_async(&mut connection).await, + Err(commands) => { + let mut pipeline = redis::pipe(); + for command in commands { + pipeline.add_command(command); + } + let replies: Vec = pipeline.query_async(&mut connection).await?; + Ok(replies.into_iter().next().unwrap_or(Value::Nil)) + } + }, + Plan::Script(call) => { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut connection) + .await + } + Plan::KeepingTtl { + increment, + key, + ttl, + trips, + } => { + let (value, current): (Value, i64) = match trips { + Trips::Shared => { + redis::pipe() + .add_command(increment) + .add_command(commands::ttl_cmd(&key)) + .query_async(&mut connection) + .await? + } + Trips::Separate => ( + increment.query_async(&mut connection).await?, + commands::ttl_cmd(&key).query_async(&mut connection).await?, + ), + }; + if current == -1 { + redis::cmd("EXPIRE") + .arg(&key) + .arg(ttl.as_secs()) + .exec_async(&mut connection) + .await?; + } + Ok(value) + } + } + }) + .await + } + + /// Messages published on `channel`, on a connection of its own. After the connection is lost + /// it reconnects and resubscribes, and messages published meanwhile are not delivered. + pub fn subscribe(&self, channel: &str) -> BoxStream<'static, Vec> { + self.0.subscriber.clone().messages(channel.to_owned()) + } + pub fn pipeline(&self) -> Pipeline { - Pipeline::new(self.namespace.clone()) + Pipeline::new(self.0.namespace.clone()) + } + + pub(crate) fn namespace(&self) -> &Namespace { + &self.0.namespace + } + + pub fn is_cluster(&self) -> bool { + matches!(self.0.connection, Connection::Cluster(_)) + } + + pub fn same_client(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } + + pub fn batch(&self) -> RedisBatch { + RedisBatch::new(self.clone()) } pub async fn run(&self, pipeline: &Pipeline) -> Vec> { + let operation = Operation::Pipeline { + ops: pipeline.len(), + }; + let started = Instant::now(); let admission = match self.admit() { Ok(admission) => admission, Err(_) => { + self.observe( + operation, + started, + Outcome::Failed(FailureKind::BreakerOpen), + ); return (0..pipeline.len()) .map(|_| Err(Error::BreakerOpen)) .collect(); } }; - let results = pipeline.run(&self.connection).await; + self.check_health().await; + let results = pipeline.run(&self.0.connection).await; let transport = results.iter().find_map(|result| { result .as_ref() @@ -150,31 +202,76 @@ impl RedisClient { .filter(|error| matches!(error, Error::Transport(_))) }); settle(admission, transport); + let failure = transport.or_else(|| results.iter().find_map(|result| result.as_ref().err())); + let outcome = failure.map_or(Outcome::Succeeded, |error| { + Outcome::Failed(FailureKind::of(error)) + }); + self.observe(operation, started, outcome); results } - fn key(&self, key: &str) -> String { - self.namespace.apply(key) - } - - fn keys(&self, keys: &[&str]) -> Vec { - keys.iter().map(|key| self.key(key)).collect() - } - - async fn query(&self, command: &Cmd) -> Result { - self.guarded(async { Ok(command.query_async(&mut self.connection.clone()).await?) }) - .await - } - - async fn guarded(&self, call: impl Future>) -> Result { - let admission = self.admit()?; - let result = call.await; - settle(admission, result.as_ref().err()); + async fn guarded( + &self, + operation: Operation, + call: impl Future>, + ) -> Result { + let started = Instant::now(); + let result = match self.admit() { + Ok(admission) => { + self.check_health().await; + let result = call + .await + .map_err(|error| Error::classify(error, self.0.connection.deployment())); + settle(admission, result.as_ref().err()); + result + } + Err(error) => Err(error), + }; + self.observe(operation, started, Outcome::of(&result)); result } + fn observe(&self, operation: Operation, started: Instant, outcome: Outcome) { + *self + .0 + .last_used + .lock() + .unwrap_or_else(PoisonError::into_inner) = Instant::now(); + if let Some(events) = &self.0.events { + events.call(&CallEvent { + operation, + duration: started.elapsed(), + outcome, + }); + } + } + + /// redis-py pings a connection that sat idle longer than `health_check_interval` before + /// sending the next command on it. A ping that times out or finds the socket dropped leaves + /// that connection behind, as a command would, so the command then goes out on a fresh one, + /// where redis-py 5.3.1 disconnects and fails the command with the ping's error. The ping is + /// not a call of the caller's and never reaches the breaker or the events. + async fn check_health(&self) { + let Some(interval) = self.0.health_check else { + return; + }; + let idle = self + .0 + .last_used + .lock() + .unwrap_or_else(PoisonError::into_inner) + .elapsed(); + if idle < interval { + return; + } + if let Err(error) = self.0.connection.ping_every_node().await { + tracing::debug!(%error, "idle Redis connection failed its health check, replacing it"); + } + } + fn admit(&self) -> Result, Error> { - self.breaker + self.0 + .breaker .as_ref() .map(|breaker| { breaker @@ -195,453 +292,15 @@ fn settle(admission: Option, error: Option<&Error>) { admission.failed(Failure::Timeout, Instant::now()) } Some(Error::Transport(_)) => admission.failed(Failure::Hard, Instant::now()), + Some(Error::Server(error)) if proves_redis_answers(error) => admission.succeeded(), Some(Error::Server(_) | Error::Decode(_) | Error::BreakerOpen) => drop(admission), } } -#[cfg(test)] -mod tests { - use std::time::{Duration, Instant}; - - use litellm_redis_testing::RedisServer; - use rstest::rstest; - - use litellm_circuit_breaker::{Breaker, Config, State}; - - use super::{RedisClient, Ttl}; - use crate::{ - DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig, - test_support::{Server, Topology, client_id}, - }; - - #[rstest] - #[tokio::test] - async fn namespaces_isolate_keys_on_the_server( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let a = server.client("a", DEFAULT_TIMEOUT).await; - let b = server.client("b", DEFAULT_TIMEOUT).await; - let bare = server.client("", DEFAULT_TIMEOUT).await; - - a.set("x", "1", Duration::from_secs(60)).await.unwrap(); - - assert_eq!(a.get::("x").await.unwrap().as_deref(), Some("1")); - assert_eq!(a.get::("a:x").await.unwrap().as_deref(), Some("1")); - assert_eq!( - bare.get::("a:x").await.unwrap().as_deref(), - Some("1") - ); - assert_eq!(b.get::("x").await.unwrap(), None); - } - - #[rstest] - #[tokio::test] - async fn a_rejected_command_keeps_the_connection( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.rpush("list", "item").await.unwrap(); - let before = client_id(&client, "list").await; - - let rejected = client.incr_by("list", 1).await; - - assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); - assert_eq!(client.rpush("list", "next").await.unwrap(), 2); - assert_eq!(client_id(&client, "list").await, before); - } - - #[rstest] - #[tokio::test] - async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client - .set("float", "15.0", Duration::from_secs(60)) - .await - .unwrap(); - - let decoded = client.get::("float").await; - - assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); - assert_eq!(client.get::("missing").await.unwrap(), None); - } - - #[rstest] - #[tokio::test] - async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - client - .set("first", "1", Duration::from_secs(60)) - .await - .unwrap(); - client - .set("second", "2", Duration::from_secs(60)) - .await - .unwrap(); - - server.freeze().await; - let timed_out = client.get::("first").await; - server.thaw().await; - - assert!( - matches!(timed_out, Err(Error::Transport(_))), - "{timed_out:?}" - ); - assert_eq!( - client.get::("second").await.unwrap().as_deref(), - Some("2") - ); - } - - #[rstest] - #[tokio::test] - async fn the_same_client_recovers_after_redis_restarts( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - - server.restart().await; - - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.incr_by("after_restart", 1).await { - Ok(value) => break Some(value), - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(200)).await - } - Err(_) => break None, - } - }; - assert_eq!(recovered, Some(1)); - } - - #[rstest] - #[tokio::test] - async fn a_client_created_while_redis_is_down_works_once_redis_is_up( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - server.stop().await; - - let client = server.client("", Duration::from_millis(500)).await; - let while_down = client.ping().await; - server.resume().await; - - assert!( - matches!(while_down, Err(Error::Transport(_))), - "{while_down:?}" - ); - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.ping().await { - Ok(()) => break true, - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(200)).await - } - Err(_) => break false, - } - }; - assert!(recovered); - } - - #[rstest] - #[tokio::test] - async fn empty_batches_succeed_without_asking_redis( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - assert_eq!( - client.mget::(&[]).await.unwrap(), - Vec::>::new() - ); - assert_eq!(client.del(&[]).await.unwrap(), 0); - } - - #[rstest] - #[tokio::test] - async fn mget_and_del_span_slots_in_input_order( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.set("a", "1", Duration::from_secs(60)).await.unwrap(); - client.set("c", "3", Duration::from_secs(60)).await.unwrap(); - - let values = client.mget::(&["c", "b", "a"]).await.unwrap(); - let deleted = client.del(&["a", "b", "c"]).await.unwrap(); - - assert_eq!( - values, - vec![Some("3".to_owned()), None, Some("1".to_owned())] - ); - assert_eq!(deleted, 2); - assert_eq!( - client.mget::(&["a", "c"]).await.unwrap(), - vec![None, None] - ); - } - - #[rstest] - #[tokio::test] - async fn set_nx_only_writes_a_missing_key( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - let first = client - .set_nx("lock", "owner-1", Duration::from_secs(60)) - .await - .unwrap(); - let second = client - .set_nx("lock", "owner-2", Duration::from_secs(60)) - .await - .unwrap(); - - assert!(first); - assert!(!second); - assert_eq!( - client.get::("lock").await.unwrap().as_deref(), - Some("owner-1") - ); - } - - #[rstest] - #[tokio::test] - async fn ttl_reports_missing_persistent_and_remaining_time( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.incr_by("persistent", 1).await.unwrap(); - client - .set("expiring", "1", Duration::from_secs(60)) - .await - .unwrap(); - client.incr_by_float("expired_later", 0.5).await.unwrap(); - assert!( - client - .expire("expired_later", Duration::from_secs(30)) - .await - .unwrap() - ); - - assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing); - assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent); - assert!(matches!( - client.ttl("expiring").await.unwrap(), - Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) - )); - assert!(matches!( - client.ttl("expired_later").await.unwrap(), - Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) - )); - assert!( - !client - .expire("missing", Duration::from_secs(30)) - .await - .unwrap() - ); - } - - #[rstest] - #[tokio::test] - async fn incr_by_float_accumulates_fractions( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - client.incr_by_float("spend", 0.25).await.unwrap(); - - assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75); - } - - #[rstest] - #[tokio::test] - async fn concurrent_commands_on_one_cloned_client_all_apply( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - let tasks: Vec<_> = (0..100) - .map(|_| { - let client = client.clone(); - tokio::spawn(async move { - for _ in 0..10 { - client.incr_by("shared", 1).await.unwrap(); - } - }) - }) - .collect(); - for task in tasks { - task.await.unwrap(); - } - - assert_eq!(client.get::("shared").await.unwrap(), Some(1000)); - } - - fn breaker(failure_threshold: u32) -> Breaker { - Breaker::new( - Config { - failure_threshold, - recovery_timeout: Duration::from_millis(200), - timeout_min_duration: Duration::ZERO, - }, - Box::new(|_| {}), - ) - } - - #[rstest] - #[tokio::test] - async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let breaker = breaker(2); - let client = server - .guarded_client("", Duration::from_millis(300), Some(breaker.clone())) - .await; - - server.freeze().await; - let first = client.get::("k").await; - let second = client.get::("k").await; - let started = Instant::now(); - let third = client.get::("k").await; - let rejected_in = started.elapsed(); - server.thaw().await; - - assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); - assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); - assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); - assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); - - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.get::("k").await { - Ok(value) => break Some(value), - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(50)).await - } - Err(_) => break None, - } - }; - assert_eq!(recovered, Some(None)); - assert!( - matches!(breaker.state(), State::Closed { .. }), - "{:?}", - breaker.state() - ); - } - - #[rstest] - #[tokio::test] - async fn rejected_commands_never_open_the_breaker( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let breaker = breaker(2); - let client = server - .guarded_client("", DEFAULT_TIMEOUT, Some(breaker.clone())) - .await; - client.rpush("list", "item").await.unwrap(); - - let rejected = [ - client.incr_by("list", 1).await, - client.incr_by("list", 1).await, - client.incr_by("list", 1).await, - ]; - - assert!( - matches!(rejected[2], Err(Error::Server(_))), - "{:?}", - rejected[2] - ); - assert!( - matches!(breaker.state(), State::Closed { hard: 0, .. }), - "{:?}", - breaker.state() - ); - } - - #[rstest] - #[tokio::test] - async fn an_open_breaker_fails_every_pipeline_op_without_sending( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server - .guarded_client("", Duration::from_millis(300), Some(breaker(1))) - .await; - let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); - - server.freeze().await; - let failed = client.run(&pipeline).await; - let rejected = client.run(&pipeline).await; - server.thaw().await; - - assert!( - matches!(failed[0], Err(Error::Transport(_))), - "{:?}", - failed[0] - ); - assert!( - matches!(rejected[0], Err(Error::BreakerOpen)), - "{:?}", - rejected[0] - ); - assert!( - matches!(rejected[1], Err(Error::BreakerOpen)), - "{:?}", - rejected[1] - ); - } - - #[rstest] - #[tokio::test] - async fn a_host_endpoint_connects_to_its_database() { - let server = RedisServer::start().await.unwrap(); - let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); - let connect = |endpoint| { - RedisClient::new(RedisConfig { - endpoint, - namespace: Namespace::default(), - timeout: DEFAULT_TIMEOUT, - breaker: None, - }) - }; - let in_db_3 = connect(Endpoint::Host(HostEndpoint { - host: "127.0.0.1".to_owned(), - port, - username: None, - password: None, - tls: false, - db: 3, - })) - .unwrap(); - let in_db_0 = connect(Endpoint::Url(server.url().to_owned())).unwrap(); - let via_url = connect(Endpoint::Url(format!("{}/3", server.url()))).unwrap(); - - in_db_3 - .set("k", "3", Duration::from_secs(60)) - .await - .unwrap(); - - assert_eq!(in_db_0.get::("k").await.unwrap(), None); - assert_eq!( - via_url.get::("k").await.unwrap().as_deref(), - Some("3") - ); - } +// A replica refusing a write is the old primary mid-failover, which must not reset a run of failures +fn proves_redis_answers(error: &RedisError) -> bool { + !matches!( + error.kind(), + ErrorKind::InvalidClientConfig | ErrorKind::Server(ServerErrorKind::ReadOnly) + ) } diff --git a/litellm-rust/crates/redis-client/src/commands.rs b/litellm-rust/crates/redis-client/src/commands.rs index fb8a4a28c94..5f8c1eff3f7 100644 --- a/litellm-rust/crates/redis-client/src/commands.rs +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -1,62 +1,407 @@ use std::time::Duration; -use redis::{Cmd, ToRedisArgs, ToSingleRedisArg}; +use litellm_python_compat::repr::float_repr; +use redis::{ + Cmd, ErrorKind, FromRedisValue, RedisError, RedisWrite, ToRedisArgs, ToSingleRedisArg, Value, +}; + +use crate::{Error, Namespace, Operation, Script, Ttl, script::ScriptCall}; + +/// How an increment treats the key's TTL, the two ways `RedisCache` increments with one. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Expiry { + /// `async_increment(refresh_ttl=True)` and every pipelined increment: `EXPIRE` after each write. + Refresh(Duration), + /// `async_increment(refresh_ttl=False)`: read the `TTL` and `EXPIRE` only a key that has none. + /// Pipelined it is lowered to [`Expiry::Refresh`], like `RedisBatch.increment`. + KeepExisting(Duration), +} + +/// A number written the way redis-py encodes it: an `int` with `str()`, a `float` with `repr()`. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Amount { + Int(i64), + Float(f64), +} + +impl ToRedisArgs for Amount { + fn write_redis_args(&self, out: &mut W) { + match self { + Self::Int(value) => out.write_arg(value.to_string().as_bytes()), + Self::Float(value) => out.write_arg(float_repr(*value).as_bytes()), + } + } +} + +impl ToSingleRedisArg for Amount {} + +type Build = Box) -> Plan + Send>; +type Decode = Box Result + Send>; + +/// One Redis command as a value: what it sends alone and pipelined, and how its reply decodes. Every +/// executor (`RedisClient::send`, `Pipeline::with`, `RedisBatch::add`) applies its namespace to the +/// keys once. +#[must_use] +pub struct Command { + operation: Operation, + keys: Vec, + idempotent: bool, + build: Build, + decode: Decode, +} + +/// What a command sends once its keys carry the namespace. Only the cases whose alone and pipelined +/// forms differ get a variant of their own. +pub(crate) enum Plan { + /// Answered without a round trip, like an empty `MGET` or `DEL`. + Reply(Value), + /// Sent together; the reply is the first command's, failing if any failed. + Commands(Vec), + /// Alone through `invoke_async`, pipelined as `EVALSHA` resent once on `NOSCRIPT`. + Script(ScriptCall), + /// Alone the increment reads the `TTL` and sends `EXPIRE` only to a key without one, pipelined it + /// is the increment then `EXPIRE`. + KeepingTtl { + increment: Cmd, + key: String, + ttl: Duration, + trips: Trips, + }, +} + +#[derive(Clone, Copy)] +pub(crate) enum Trips { + /// `async_increment`: the increment and its `TTL` in one round trip. + Shared, + /// `increment_cache`: the increment and its `TTL` each on its own round trip. + Separate, +} + +/// A command with the namespace applied, ready for an executor to send. +pub(crate) struct Op { + pub(crate) operation: Operation, + pub(crate) key: String, + pub(crate) idempotent: bool, + pub(crate) plan: Plan, +} + +impl Plan { + /// What a pipeline sends for this op. + pub(crate) fn pipelined(&self) -> Vec { + match self { + Self::Reply(_) => Vec::new(), + Self::Commands(commands) => commands.clone(), + Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)], + Self::KeepingTtl { + increment, + key, + ttl, + .. + } => vec![increment.clone(), expire_cmd(key, *ttl)], + } + } +} + +impl Command { + fn new( + operation: Operation, + keys: Vec, + idempotent: bool, + build: impl FnOnce(Vec) -> Plan + Send + 'static, + decode: impl FnOnce(Value) -> Result + Send + 'static, + ) -> Self { + Self { + operation, + keys, + idempotent, + build: Box::new(build), + decode: Box::new(decode), + } + } + + /// Attaches a codec over the decoded reply, such as a store's value format. + pub fn map(self, decode: fn(T) -> Result) -> Command + where + T: 'static, + { + let inner = self.decode; + Command { + operation: self.operation, + keys: self.keys, + idempotent: self.idempotent, + build: self.build, + decode: Box::new(move |value| inner(value).and_then(decode)), + } + } + + pub(crate) fn prepare(self, namespace: &Namespace) -> (Op, Decode) { + let keys: Vec = self.keys.iter().map(|key| namespace.apply(key)).collect(); + let key = keys.first().cloned().unwrap_or_default(); + let op = Op { + operation: self.operation, + key, + idempotent: self.idempotent, + plan: (self.build)(keys), + }; + (op, self.decode) + } +} + +fn decode(value: Value) -> Result { + T::from_redis_value(value).map_err(|error| Error::Decode(error.into())) +} + +fn owned(keys: &[&str]) -> Vec { + keys.iter().map(|key| (*key).to_owned()).collect() +} + +fn single(keys: Vec) -> String { + keys.into_iter().next().unwrap_or_default() +} + +fn expire_cmd(key: &str, ttl: Duration) -> Cmd { + redis::cmd("EXPIRE").arg(key).arg(ttl.as_secs()).to_owned() +} + +pub(crate) fn ttl_cmd(key: &str) -> Cmd { + redis::cmd("TTL").arg(key).to_owned() +} + +/// A reply whose shape a codec over it refuses, as `what` describes. +pub fn unexpected_reply(what: &'static str) -> Error { + Error::Decode(RedisError::from((ErrorKind::UnexpectedReturnType, what))) +} + +pub fn get(key: &str) -> Command> { + Command::new( + Operation::Command("GET"), + vec![key.to_owned()], + true, + |keys| Plan::Commands(vec![redis::cmd("GET").arg(single(keys)).to_owned()]), + decode, + ) +} + +/// One `MGET` of `keys` in order, repeated keys read once, like `RedisBatch.mget`, answering one +/// value per key of `keys` in their order. No keys is no round trip, since Redis rejects an empty +/// `MGET`. Pipelined, every key must share a hash slot on a cluster, and a batch never pipelines +/// there. +pub fn mget(keys: &[&str]) -> Command>> { + let unique = keys.iter().fold(Vec::::new(), |mut unique, key| { + if !unique.iter().any(|seen| seen == key) { + unique.push((*key).to_owned()); + } + unique + }); + let positions: Vec = keys + .iter() + .filter_map(|key| unique.iter().position(|seen| seen == key)) + .collect(); + Command::new( + Operation::Command("MGET"), + unique, + true, + |keys| match keys.is_empty() { + true => Plan::Reply(Value::Array(Vec::new())), + false => Plan::Commands(vec![redis::cmd("MGET").arg(keys).to_owned()]), + }, + move |reply| { + let values: Vec = decode(reply)?; + let in_key_order = positions + .iter() + .map(|&position| values.get(position).cloned()) + .collect::>>() + .ok_or_else(|| unexpected_reply("MGET answered fewer values than keys"))?; + decode(Value::Array(in_key_order)) + }, + ) +} + +// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())` +pub fn set(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<()> { + let value = value.to_redis_args().concat(); + Command::new( + Operation::Command("SET"), + vec![key.to_owned()], + true, + move |keys| { + Plan::Commands(vec![ + redis::cmd("SET") + .arg(single(keys)) + .arg(value) + .arg("EX") + .arg(ttl.as_secs()) + .to_owned(), + ]) + }, + decode, + ) +} + +/// `SET ... EX ... NX`, `EX` before `NX` as redis-py builds it. `true` when this call wrote the key. +pub fn set_nx(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command { + let value = value.to_redis_args().concat(); + Command::new( + Operation::Command("SET"), + vec![key.to_owned()], + false, + move |keys| { + Plan::Commands(vec![ + redis::cmd("SET") + .arg(single(keys)) + .arg(value) + .arg("EX") + .arg(ttl.as_secs()) + .arg("NX") + .to_owned(), + ]) + }, + |value| decode::>(value).map(|reply| reply.is_some()), + ) +} + +/// One multi-key `DEL`, no round trip for no keys. Pipelined, every key must share a hash slot. +pub fn del(keys: &[&str]) -> Command { + Command::new( + Operation::Command("DEL"), + owned(keys), + true, + |keys| match keys.is_empty() { + true => Plan::Reply(Value::Int(0)), + false => Plan::Commands(vec![redis::cmd("DEL").arg(keys).to_owned()]), + }, + decode, + ) +} + +pub fn expire(key: &str, ttl: Duration) -> Command { + Command::new( + Operation::Command("EXPIRE"), + vec![key.to_owned()], + true, + move |keys| Plan::Commands(vec![expire_cmd(&single(keys), ttl)]), + decode, + ) +} + +pub fn ttl(key: &str) -> Command { + Command::new( + Operation::Command("TTL"), + vec![key.to_owned()], + true, + |keys| Plan::Commands(vec![ttl_cmd(&single(keys))]), + |value| { + Ok(match decode::(value)? { + -2 => Ttl::Missing, + -1 => Ttl::Persistent, + seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())), + }) + }, + ) +} + +/// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no TTL, +/// each its own round trip. +pub fn incr_by(key: &str, delta: i64, ttl: Duration) -> Command { + Command::new( + Operation::Command("INCRBY"), + vec![key.to_owned()], + false, + move |keys| { + let key = single(keys); + Plan::KeepingTtl { + increment: redis::cmd("INCRBY").arg(&key).arg(delta).to_owned(), + key, + ttl, + trips: Trips::Separate, + } + }, + decode, + ) +} + +/// `INCRBYFLOAT`, with the TTL handled as `expiry` says, or left alone without one. Alone it is +/// `RedisCache.async_increment`, pipelined `RedisBatch.increment`. +pub fn incr_by_float(key: &str, amount: Amount, expiry: Option) -> Command { + Command::new( + Operation::Command("INCRBYFLOAT"), + vec![key.to_owned()], + false, + move |keys| { + let key = single(keys); + let increment = redis::cmd("INCRBYFLOAT").arg(&key).arg(amount).to_owned(); + match expiry { + None => Plan::Commands(vec![increment]), + Some(Expiry::Refresh(ttl)) => { + Plan::Commands(vec![increment, expire_cmd(&key, ttl)]) + } + Some(Expiry::KeepExisting(ttl)) => Plan::KeepingTtl { + increment, + key, + ttl, + trips: Trips::Shared, + }, + } + }, + decode, + ) +} + +/// A registered script, like redis-py's `Script`. Every key it touches must be in `keys`. +pub fn eval( + script: &Script, + keys: &[&str], + args: impl ToRedisArgs, +) -> Command { + let script = script.clone(); + let args = args.to_redis_args(); + Command::new( + Operation::Script, + owned(keys), + false, + move |keys| Plan::Script(ScriptCall { script, keys, args }), + decode, + ) +} + +/// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the scripts +/// `RedisCache` runs that way rather than through `register_script`. +pub fn eval_source( + source: &'static str, + keys: &[&str], + args: impl ToRedisArgs, +) -> Command { + let args = args.to_redis_args(); + Command::new( + Operation::Script, + owned(keys), + false, + move |keys| { + Plan::Commands(vec![ + redis::cmd("EVAL") + .arg(source) + .arg(keys.len()) + .arg(keys) + .arg(args) + .to_owned(), + ]) + }, + decode, + ) +} + +/// The channel is sent as given, without the namespace, like redis-py's `publish`. +pub fn publish(channel: &str, payload: &[u8]) -> Command<()> { + let command = redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned(); + Command::new( + Operation::Command("PUBLISH"), + Vec::new(), + false, + move |_| Plan::Commands(vec![command]), + |_| Ok(()), + ) +} pub(crate) fn ping() -> Cmd { redis::cmd("PING") } - -pub(crate) fn get(key: String) -> Cmd { - redis::cmd("GET").arg(key).to_owned() -} - -pub(crate) fn mget(keys: Vec) -> Cmd { - redis::cmd("MGET").arg(keys).to_owned() -} - -pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { - redis::cmd("SET") - .arg(key) - .arg(value) - .arg("PX") - .arg(millis(ttl)) - .to_owned() -} - -pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { - redis::cmd("SET") - .arg(key) - .arg(value) - .arg("NX") - .arg("PX") - .arg(millis(ttl)) - .to_owned() -} - -pub(crate) fn del(keys: Vec) -> Cmd { - redis::cmd("DEL").arg(keys).to_owned() -} - -pub(crate) fn expire(key: String, ttl: Duration) -> Cmd { - redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned() -} - -pub(crate) fn ttl(key: String) -> Cmd { - redis::cmd("PTTL").arg(key).to_owned() -} - -pub(crate) fn incr_by(key: String, delta: i64) -> Cmd { - redis::cmd("INCRBY").arg(key).arg(delta).to_owned() -} - -pub(crate) fn incr_by_float(key: String, delta: f64) -> Cmd { - redis::cmd("INCRBYFLOAT").arg(key).arg(delta).to_owned() -} - -pub(crate) fn rpush(key: String, values: impl ToRedisArgs) -> Cmd { - redis::cmd("RPUSH").arg(key).arg(values).to_owned() -} - -fn millis(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index 64934c68891..f43ab4feb85 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -1,152 +1,458 @@ -use std::{fmt, time::Duration}; +use std::{fmt, path::PathBuf, pin::Pin, sync::Arc, time::Duration}; -use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; +use futures_util::{Stream, StreamExt}; + +use redis::{ + BasicAuth, ClientTlsConfig, ConnectionAddr, ConnectionInfo, ErrorKind, IntoConnectionInfo, + ProtocolVersion, RedisConnectionInfo, RedisError, RedisResult, StreamingCredentialsProvider, + TlsCertificates, TlsMode, io::tcp::TcpSettings, +}; use litellm_circuit_breaker::Breaker; +use litellm_secret_value::SecretValue; -use crate::Namespace; +use pkcs8::EncryptedPrivateKeyInfo; +use crate::{ConnectError, KeyError, Namespace, RedisEvents}; +use litellm_redis_types::{Cluster, RedisUrl, Target, TlsOptions}; + +// litellm/constants.py REDIS_SOCKET_TIMEOUT, which _coerce_redis_kwargs_types sets when socket_timeout is unset pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); +// below the idle timeouts of common cloud load balancers (AWS NLB 350s, Azure Load Balancer 240s), which the +// OS default of two hours never probes within +const KEEPALIVE_IDLE: Duration = Duration::from_secs(60); pub struct RedisConfig { - pub endpoint: Endpoint, + pub target: Target, + pub credentials: Credentials, pub namespace: Namespace, - pub timeout: Duration, + pub timeouts: Timeouts, pub breaker: Option, + pub events: Option>, + /// Ping every node before the next command once the client sat idle this long, like + /// redis-py's `health_check_interval`, so a connection the server dropped silently is + /// replaced before a real command pays for it. + pub health_check: Option, + /// Probe a connection with TCP keepalive once it sat idle, like redis-py's `socket_keepalive`. + pub keepalive: bool, } -#[derive(Clone, PartialEq)] -pub enum Endpoint { - Url(String), - Host(HostEndpoint), - Cluster(ClusterEndpoint), +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Timeouts { + pub response: Duration, + pub connect: Duration, } -#[derive(Clone, PartialEq)] -pub struct HostEndpoint { - pub host: String, - pub port: u16, - pub username: Option, - pub password: Option, - pub tls: bool, - pub db: i64, -} - -#[derive(Clone, PartialEq)] -pub struct ClusterEndpoint { - pub nodes: Vec, - pub username: Option, - pub password: Option, - pub tls: bool, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct ClusterNode { - pub host: String, - pub port: u16, -} - -impl ClusterEndpoint { - pub(crate) fn addrs(&self) -> impl Iterator + '_ { - self.nodes - .iter() - .map(|node| addr(&node.host, node.port, self.tls)) - } -} - -impl HostEndpoint { - pub(crate) fn connection_info(&self) -> RedisResult { - let settings = RedisConnectionInfo::default().set_db(self.db); - let settings = match &self.username { - Some(username) => settings.set_username(username), - None => settings, - }; - let settings = match &self.password { - Some(password) => settings.set_password(password), - None => settings, - }; - Ok(addr(&self.host, self.port, self.tls) - .into_connection_info()? - .set_redis_settings(settings)) - } -} - -fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr { - match tls { - true => ConnectionAddr::TcpTls { - host: host.to_owned(), - port, - insecure: false, - tls_params: None, - }, - false => ConnectionAddr::Tcp(host.to_owned(), port), - } -} - -impl fmt::Debug for Endpoint { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Url(url) => formatter - .debug_tuple("Url") - .field(&redacted_url(url)) - .finish(), - Self::Host(host) => formatter - .debug_struct("Host") - .field("host", &host.host) - .field("port", &host.port) - .field("username", &host.username) - .field("password", &host.password.as_ref().map(|_| "")) - .field("tls", &host.tls) - .field("db", &host.db) - .finish(), - Self::Cluster(cluster) => formatter - .debug_struct("Cluster") - .field("nodes", &cluster.nodes) - .field("username", &cluster.username) - .field("password", &cluster.password.as_ref().map(|_| "")) - .field("tls", &cluster.tls) - .finish(), +impl Timeouts { + pub fn new(socket_timeout: Duration, connect_timeout: Option) -> Self { + Self { + response: socket_timeout, + connect: connect_timeout.unwrap_or(socket_timeout), } } } -fn redacted_url(url: &str) -> String { - match url - .split_once("://") - .and_then(|(scheme, rest)| Some((scheme, rest.rsplit_once('@')?.1))) - { - Some((scheme, host)) => format!("{scheme}://@{host}"), - None => url.to_owned(), +impl Default for Timeouts { + fn default() -> Self { + Self::new(DEFAULT_TIMEOUT, None) } } +#[derive(Clone, Debug)] +pub enum Credentials { + Static { + username: Option, + password: Option, + }, + Provider(CredentialsProvider), +} + +#[derive(Clone)] +pub struct CredentialsProvider(pub Arc); + +impl fmt::Debug for CredentialsProvider { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CredentialsProvider") + .finish_non_exhaustive() + } +} + +impl StreamingCredentialsProvider for CredentialsProvider { + fn subscribe(&self) -> Pin> + Send + 'static>> { + self.0.subscribe() + } +} + +impl Credentials { + pub const NONE: Self = Self::Static { + username: None, + password: None, + }; +} + +pub(crate) fn tls_mode(tls: &TlsOptions) -> TlsMode { + match tls.verify_certificate { + true => TlsMode::Secure, + false => TlsMode::Insecure, + } +} + +pub(crate) fn tls_certificates(tls: &TlsOptions) -> Result { + let client_tls = match (&tls.certfile, &tls.keyfile) { + (Some(certfile), keyfile) => Some(ClientTlsConfig { + client_cert: read(certfile)?, + client_key: client_key( + keyfile.as_ref().unwrap_or(certfile), + tls.key_password.as_ref(), + )?, + }), + (None, _) => None, + }; + let from_file: Option> = tls.ca_certs.as_ref().map(read).transpose()?; + let from_directory: Vec> = tls + .ca_path + .as_ref() + .map(read_directory) + .transpose()? + .unwrap_or_default(); + let configured: Vec> = from_file + .into_iter() + .chain(tls.ca_data.iter().map(|pem| pem.clone().into_bytes())) + .chain(from_directory) + .collect(); + Ok(TlsCertificates { + client_tls, + root_cert: (!configured.is_empty()).then(|| with_system_roots(configured.join(&b'\n'))), + }) +} + +/// Every regular file in `path`, in name order, as redis-py hands `ssl_ca_path` to +/// `load_verify_locations(capath=...)` +fn read_directory(path: &PathBuf) -> Result>, ConnectError> { + let tls_file = |source| ConnectError::TlsFile { + path: path.display().to_string(), + source, + }; + let mut files: Vec = std::fs::read_dir(path) + .map_err(tls_file)? + .map(|entry| entry.map(|entry| entry.path())) + .collect::>() + .map_err(tls_file)?; + files.sort(); + files + .iter() + .filter(|file| file.is_file()) + .map(read) + .collect() +} + +pub(crate) fn apply_tls( + tls: &TlsOptions, + info: ConnectionInfo, +) -> Result { + let tls_info = redis::Client::build_with_tls(info, tls_certificates(tls)?)? + .get_connection_info() + .clone(); + let mut addr = match tls_info.addr().clone() { + ConnectionAddr::TcpTls { + host, + port, + tls_params, + .. + } => ConnectionAddr::TcpTls { + host, + port, + insecure: !tls.verify_certificate, + tls_params, + }, + addr => addr, + }; + addr.set_danger_accept_invalid_hostnames(!tls.check_hostname); + Ok(tls_info.set_addr(addr)) +} + +fn with_system_roots(configured: Vec) -> Vec { + let system = rustls_native_certs::load_native_certs(); + let system_pem = system + .certs + .iter() + .map(|cert| pem::encode(&pem::Pem::new("CERTIFICATE", cert.to_vec()))); + configured + .into_iter() + .chain(b"\n".iter().copied()) + .chain(system_pem.flat_map(String::into_bytes)) + .collect() +} + +/// The key as unencrypted PEM, decrypting a PKCS#8 `ENCRYPTED PRIVATE KEY` with `ssl_password` the way +/// `load_cert_chain(password=...)` does, since rustls only reads unencrypted keys. Like OpenSSL, a password next to +/// an unencrypted key is unused +fn client_key(path: &PathBuf, password: Option<&SecretValue>) -> Result, ConnectError> { + let contents = read(path)?; + let key_error = |reason| ConnectError::TlsKey { + path: path.display().to_string(), + reason, + }; + let Ok(sections) = pem::parse_many(&contents) else { + return Ok(contents); + }; + let encrypted = sections + .iter() + .find(|section| section.tag() == "ENCRYPTED PRIVATE KEY"); + match (encrypted, password) { + (Some(_), None) => Err(key_error(KeyError::MissingPassword)), + (Some(section), Some(password)) => { + let decrypted = EncryptedPrivateKeyInfo::try_from(section.contents()) + .and_then(|info| info.decrypt(password.expose())) + .map_err(|source| key_error(KeyError::Decrypt(source)))?; + Ok(pem::encode(&pem::Pem::new("PRIVATE KEY", decrypted.as_bytes())).into_bytes()) + } + (None, _) if sections.iter().any(is_legacy_encrypted) => { + Err(key_error(KeyError::LegacyEncryption)) + } + (None, _) => Ok(contents), + } +} + +fn is_legacy_encrypted(section: &pem::Pem) -> bool { + section + .headers() + .get("Proc-Type") + .is_some_and(|proc_type| proc_type.contains("ENCRYPTED")) +} + +fn read(path: &PathBuf) -> Result, ConnectError> { + std::fs::read(path).map_err(|source| ConnectError::TlsFile { + path: path.display().to_string(), + source, + }) +} + +pub(crate) fn addr(host: &str, port: u16, tls: Option<&TlsOptions>) -> ConnectionAddr { + match tls { + Some(tls) => ConnectionAddr::TcpTls { + host: host.to_owned(), + port, + insecure: !tls.verify_certificate, + tls_params: None, + }, + None => ConnectionAddr::Tcp(host.to_owned(), port), + } +} + +pub(crate) fn cluster_addrs(cluster: &Cluster) -> impl Iterator + '_ { + cluster + .nodes + .iter() + .map(|node| addr(&node.host, node.port, cluster.tls.as_ref())) +} + +pub(crate) fn tcp_settings(keepalive: bool) -> TcpSettings { + let settings = TcpSettings::default().set_nodelay(true); + match keepalive { + true => settings.set_keepalive(socket2::TcpKeepalive::new().with_time(KEEPALIVE_IDLE)), + false => settings, + } +} + +pub(crate) fn connection_info( + target: &Target, + credentials: &Credentials, + keepalive: bool, +) -> Result { + let (info, tls) = match target { + Target::Url(url) => (url_info(url, credentials)?, url.tls.as_ref()), + Target::Standalone(host) => { + let info = addr(&host.host, host.port, host.tls.as_ref()).into_connection_info()?; + let settings = info.redis_settings().clone().set_db(host.db); + (info.set_redis_settings(settings), host.tls.as_ref()) + } + Target::Unix(socket) => { + let info = ConnectionAddr::Unix(socket.path.clone()).into_connection_info()?; + let settings = info.redis_settings().clone().set_db(socket.db); + (info.set_redis_settings(settings), None) + } + Target::Cluster(_) | Target::Sentinel(_) => { + unreachable!("cluster and sentinel endpoints build their own connection info") + } + }; + let info = with_static_credentials(info, credentials).set_tcp_settings(tcp_settings(keepalive)); + match (tls, info.addr()) { + (Some(tls), ConnectionAddr::TcpTls { .. }) => apply_tls(tls, info), + _ => Ok(info), + } +} + +// parse_url reads ?db= over the path, and the url's userinfo over ?username= and ?password=. A credential +// provider supersedes both, as it does in redis-py. ?protocol=3 would switch redis-rs to RESP3, which no +// reply decode here expects, so the protocol is ignored as the resolver warns +fn url_info(url: &RedisUrl, credentials: &Credentials) -> Result { + let info = url.url.expose().into_connection_info()?; + let settings = info + .redis_settings() + .clone() + .set_protocol(ProtocolVersion::RESP2); + let settings = match url.db { + Some(db) => settings.set_db(db), + None => settings, + }; + let info = info.set_redis_settings(settings); + Ok(match credentials { + Credentials::Static { .. } => with_static_credentials( + info, + &Credentials::Static { + username: url.username.clone(), + password: url.password.clone(), + }, + ), + Credentials::Provider(_) => info, + }) +} + +fn with_static_credentials(info: ConnectionInfo, credentials: &Credentials) -> ConnectionInfo { + let settings = static_credentials(info.redis_settings().clone(), credentials); + info.set_redis_settings(settings) +} + +fn static_credentials( + settings: RedisConnectionInfo, + credentials: &Credentials, +) -> RedisConnectionInfo { + let Credentials::Static { username, password } = credentials else { + return settings; + }; + let settings = match (username, settings.username()) { + (Some(username), None) => settings.set_username(username.expose()), + _ => settings, + }; + match (password, settings.password()) { + (Some(password), None) => settings.set_password(password.expose()), + _ => settings, + } +} + +/// `settings` with the credentials a connection opened now sends. A provider is asked once for its +/// current pair, which is when redis-py asks its `credential_provider`. +pub(crate) async fn authenticated( + settings: RedisConnectionInfo, + credentials: &Credentials, +) -> RedisResult { + let Credentials::Provider(provider) = credentials else { + return Ok(static_credentials(settings, credentials)); + }; + let current = provider.subscribe().next().await.ok_or_else(|| { + RedisError::from(( + ErrorKind::AuthenticationFailed, + "the Redis credential provider ended before giving credentials", + )) + })??; + Ok(settings + .set_username(current.username()) + .set_password(current.password())) +} + #[cfg(test)] mod tests { + use litellm_redis_types::{RedisUrl, UnixSocket}; use rstest::rstest; - use super::{Endpoint, HostEndpoint}; + use super::*; #[rstest] - #[case::url_with_password(Endpoint::Url("redis://user:secret@cache:6379/0".to_owned()), "redis://@cache:6379/0")] - #[case::url_with_at_in_password(Endpoint::Url("rediss://:p@ss@cache:6380".to_owned()), "rediss://@cache:6380")] - #[case::host_password( - Endpoint::Host(HostEndpoint { - host: "cache".to_owned(), - port: 6379, - username: None, - password: Some("secret".to_owned()), - tls: false, - db: 0, - }), - "" - )] - fn debug_output_never_shows_the_password(#[case] endpoint: Endpoint, #[case] shown: &str) { - let debug = format!("{endpoint:?}"); + #[case::host(Target::Standalone(crate::Standalone { host: "redis".into(), port: 6379, tls: None, db: 0 }))] + #[case::url(Target::Url(RedisUrl { + url: SecretValue::new("redis://cache:6379".to_owned()), + db: None, + username: None, + password: None, + tls: None, + }))] + fn keepalive_probes_an_idle_socket_only_when_enabled(#[case] target: Target) { + let probed = |keepalive| { + connection_info(&target, &Credentials::NONE, keepalive) + .unwrap() + .tcp_settings() + .keepalive() + .is_some() + }; - assert!( - !debug.contains("secret") && !debug.contains("p@ss"), - "{debug}" + assert!(probed(true)); + assert!(!probed(false)); + } + + #[test] + fn a_unix_socket_replaces_host_port_and_tls_like_redis_py() { + let target = Target::Unix(UnixSocket { + path: "/tmp/redis.sock".into(), + db: 3, + }); + + let info = connection_info(&target, &Credentials::NONE, true).unwrap(); + + assert_eq!(*info.addr(), ConnectionAddr::Unix("/tmp/redis.sock".into())); + assert_eq!(info.redis_settings().db(), 3); + } + + #[rstest] + #[case::the_query_db_wins_over_the_path("redis://cache:6379/1", Some(3), None, 3, None)] + #[case::the_path_db_without_a_query_db("redis://cache:6379/1", None, None, 1, None)] + #[case::the_query_password_fills_a_missing_userinfo( + "redis://cache:6379", + None, + Some("query"), + 0, + Some("query") + )] + #[case::the_userinfo_wins_over_the_query_password( + "redis://:userinfo@cache:6379", + None, + Some("query"), + 0, + Some("userinfo") + )] + fn a_url_reads_its_query_like_parse_url( + #[case] url: &str, + #[case] db: Option, + #[case] password: Option<&str>, + #[case] expected_db: i64, + #[case] expected_password: Option<&str>, + ) { + let target = Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db, + username: None, + password: password.map(SecretValue::new), + }); + let static_password = Credentials::Static { + username: None, + password: Some(SecretValue::new("static")), + }; + + let info = connection_info(&target, &static_password, true).unwrap(); + + assert_eq!(info.redis_settings().db(), expected_db); + assert_eq!( + info.redis_settings().password(), + expected_password.or(Some("static")) ); - assert!(debug.contains(shown), "{debug}"); + } + + #[rstest] + #[case::by_number("redis://cache:6379/?protocol=3")] + #[case::by_name("redis://cache:6379/?protocol=resp3")] + fn a_url_never_switches_the_protocol_from_resp2(#[case] url: &str) { + let target = Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db: None, + username: None, + password: None, + }); + + let info = connection_info(&target, &Credentials::NONE, true).unwrap(); + + assert_eq!(info.redis_settings().protocol(), ProtocolVersion::RESP2); } } diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index aaf26f105da..c1f82490b3b 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -1,53 +1,161 @@ -use std::time::Duration; - -use redis::{ - Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value, - aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig}, - cluster::{ClusterClientBuilder, ClusterConfig}, - cluster_async::ClusterConnection, +use std::{ + sync::{ + Arc, Mutex, PoisonError, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, }; -use crate::ClusterEndpoint; +use redis::{ + AsyncConnectionConfig, Cmd, ConnectionAddr, ConnectionInfo, ErrorKind, Pipeline, + RedisConnectionInfo, RedisError, RedisFuture, RedisResult, ServerErrorKind, TlsCertificates, + TlsMode, Value, + aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig, MultiplexedConnection}, + cluster::{ClusterClientBuilder, ClusterConfig}, + cluster_async::ClusterConnection, + cluster_routing::{MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo}, + sentinel::{SentinelClient, SentinelClientBuilder, SentinelServerType}, +}; -const RETRIES: u32 = 3; +use crate::{ + Cluster, ConnectError, Credentials, Sentinel, Target, Timeouts, + config::{self, addr}, + error::Deployment, +}; + +// redis-py 5.3.1 redis/asyncio/connection.py: without retry options a connection gets Retry(NoBackoff(), 0) +const RECONNECT_RETRIES: usize = 0; +// redis-py 5.3.1 redis/asyncio/cluster.py: cluster_error_retry_attempts=3 +const CLUSTER_RETRIES: u32 = 3; +// redis-py 5.3.1 redis/asyncio/cluster.py sleeps a flat 0.25s before retrying after CLUSTERDOWN. +// redis-rs draws the wait from [min, max), so a 1ms window pins it to 250ms +const CLUSTER_RETRY_WAIT_MS: u64 = 250; #[derive(Clone)] pub(crate) enum Connection { - Standalone(ConnectionManager), - Cluster(ClusterConnection), + Standalone(Renewing), + Cluster(Renewing), + Sentinel(SentinelConnection), } impl Connection { - pub(crate) fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult { - let config = ConnectionManagerConfig::new() - .set_connection_timeout(Some(timeout)) - .set_response_timeout(Some(timeout)) - .set_number_of_retries(RETRIES as usize); - let manager = ConnectionManager::new_lazy_with_config(redis::Client::open(info)?, config)?; - Ok(Self::Standalone(manager)) + pub(crate) fn new( + target: &Target, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + match target { + Target::Url(_) | Target::Standalone(_) | Target::Unix(_) => Self::standalone( + config::connection_info(target, credentials, keepalive)?, + credentials, + timeouts, + ), + Target::Cluster(cluster) => Self::cluster(cluster, credentials, timeouts, keepalive), + Target::Sentinel(sentinel) => Ok(Self::Sentinel(SentinelConnection::new( + sentinel, + credentials, + timeouts, + keepalive, + )?)), + } } - pub(crate) fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult { - let builder = ClusterClientBuilder::new(endpoint.addrs()) - .connection_timeout(timeout) - .response_timeout(timeout) - .retries(RETRIES); - let builder = match &endpoint.username { - Some(username) => builder.username(username), + pub(crate) fn deployment(&self) -> Deployment { + match self { + Self::Standalone(_) => Deployment::Standalone, + Self::Cluster(_) => Deployment::Cluster, + Self::Sentinel(_) => Deployment::Sentinel, + } + } + + fn standalone( + info: ConnectionInfo, + credentials: &Credentials, + timeouts: Timeouts, + ) -> Result { + let config = ConnectionManagerConfig::new() + .set_connection_timeout(Some(timeouts.connect)) + .set_response_timeout(Some(timeouts.response)) + .set_number_of_retries(RECONNECT_RETRIES); + let config = match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + }; + let client = redis::Client::open(info)?; + Ok(Self::Standalone(Renewing::new(move || { + ConnectionManager::new_lazy_with_config(client.clone(), config.clone()) + })?)) + } + + fn cluster( + endpoint: &Cluster, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + let builder = ClusterClientBuilder::new(config::cluster_addrs(endpoint)) + .tcp_settings(config::tcp_settings(keepalive)) + .connection_timeout(timeouts.connect) + .response_timeout(timeouts.response) + .retries(CLUSTER_RETRIES) + .retry_wait_formula(0, 2) + .min_retry_wait(CLUSTER_RETRY_WAIT_MS) + .max_retry_wait(CLUSTER_RETRY_WAIT_MS + 1); + let builder = match &endpoint.tls { + Some(tls) => builder + .tls(config::tls_mode(tls)) + .certs(config::tls_certificates(tls)?) + .danger_accept_invalid_hostnames(!tls.check_hostname), None => builder, }; - let builder = match &endpoint.password { - Some(password) => builder.password(password), - None => builder, + let builder = match credentials { + Credentials::Static { username, password } => { + let builder = match username { + Some(username) => builder.username(username.expose()), + None => builder, + }; + match password { + Some(password) => builder.password(password.expose()), + None => builder, + } + } + Credentials::Provider(provider) => builder.set_credentials_provider(provider.clone()), }; - let config = ClusterConfig::new() - .set_connection_timeout(timeout) - .set_response_timeout(timeout); - Ok(Self::Cluster( - builder - .build()? - .get_pending_async_connection_with_config(config), - )) + // The builder's response timeout bounds the whole request, but with_config overwrites each node's + // response timeout with this one even when unset, so it has to be repeated here + let config = ClusterConfig::new().set_response_timeout(timeouts.response); + let client = builder.build()?; + Ok(Self::Cluster(Renewing::new(move || { + Ok(client.get_pending_async_connection_with_config(config.clone())) + })?)) + } +} + +impl Connection { + /// `PING` on the one multiplexed connection, or on every cluster node, so a dead socket is + /// replaced before a real command uses it. + pub(crate) async fn ping_every_node(&self) -> RedisResult<()> { + let ping = redis::cmd("PING"); + match self { + Self::Cluster(renewing) => { + let (generation, mut connection) = renewing.current(); + let reply = connection + .route_command( + ping, + RoutingInfo::MultiNode(( + MultipleNodeRoutingInfo::AllNodes, + Some(ResponsePolicy::AllSucceeded), + )), + ) + .await; + renewing.settle(generation, &reply); + reply.map(drop) + } + Self::Standalone(_) | Self::Sentinel(_) => { + ping.query_async::(&mut self.clone()).await.map(drop) + } + } } } @@ -56,6 +164,7 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.req_packed_command(cmd), Self::Cluster(connection) => connection.req_packed_command(cmd), + Self::Sentinel(connection) => connection.req_packed_command(cmd), } } @@ -68,6 +177,7 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.req_packed_commands(pipeline, offset, count), Self::Cluster(connection) => connection.req_packed_commands(pipeline, offset, count), + Self::Sentinel(connection) => connection.req_packed_commands(pipeline, offset, count), } } @@ -75,6 +185,441 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.get_db(), Self::Cluster(connection) => connection.get_db(), + Self::Sentinel(connection) => connection.get_db(), } } } + +/// A connection swapped for a fresh one once a command on it times out, like redis-py disconnecting +/// a connection whose read timed out. redis-rs keeps the socket after a timeout and the cluster +/// client retries on it, so a socket that went silent would time out every command until TCP gives +/// up on it. +#[derive(Clone)] +pub(crate) struct Renewing(Arc>); + +type Open = Box RedisResult + Send + Sync>; + +struct RenewingState { + open: Open, + current: Mutex<(u64, C)>, +} + +impl Renewing { + fn new(open: impl Fn() -> RedisResult + Send + Sync + 'static) -> RedisResult { + let first = open()?; + Ok(Self(Arc::new(RenewingState { + open: Box::new(open), + current: Mutex::new((0, first)), + }))) + } + + fn current(&self) -> (u64, C) { + self.0 + .current + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + /// Replaces the connection of `generation` when its reply timed out, unless another timeout + /// already replaced it. + fn settle(&self, generation: u64, reply: &RedisResult) { + if !reply.as_ref().is_err_and(RedisError::is_timeout) { + return; + } + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(PoisonError::into_inner); + if current.0 != generation { + return; + } + match (self.0.open)() { + Ok(fresh) => *current = (generation + 1, fresh), + Err(error) => { + tracing::warn!(%error, "replacing a timed-out Redis connection failed") + } + } + } +} + +impl ConnectionLike for Renewing { + fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { + Box::pin(async move { + let (generation, mut connection) = self.current(); + let reply = connection.req_packed_command(cmd).await; + self.settle(generation, &reply); + reply + }) + } + + fn req_packed_commands<'a>( + &'a mut self, + pipeline: &'a Pipeline, + offset: usize, + count: usize, + ) -> RedisFuture<'a, Vec> { + Box::pin(async move { + let (generation, mut connection) = self.current(); + let replies = connection + .req_packed_commands(pipeline, offset, count) + .await; + self.settle(generation, &replies); + replies + }) + } + + fn get_db(&self) -> i64 { + self.current().1.get_db() + } +} + +#[derive(Clone)] +pub(crate) struct SentinelConnection(Arc); + +struct SentinelState { + discovery: tokio::sync::Mutex<()>, + current: Mutex>, + generations: AtomicU64, + endpoint: Sentinel, + credentials: Credentials, + tls: Option<(TlsMode, TlsCertificates)>, + config: AsyncConnectionConfig, + connect_timeout: Duration, + keepalive: bool, +} + +impl SentinelConnection { + fn new( + endpoint: &Sentinel, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + let tls = match &endpoint.tls { + Some(tls) => Some((config::tls_mode(tls), config::tls_certificates(tls)?)), + None => None, + }; + let state = SentinelState { + discovery: tokio::sync::Mutex::new(()), + current: Mutex::new(None), + generations: AtomicU64::new(0), + endpoint: endpoint.clone(), + credentials: credentials.clone(), + tls, + config: async_config(credentials, timeouts), + connect_timeout: timeouts.connect, + keepalive, + }; + state.sentinel_client(&RedisConnectionInfo::default())?; + Ok(Self(Arc::new(state))) + } + + fn cached(&self) -> Option<(u64, MultiplexedConnection)> { + self.0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + } + + async fn connection(&self) -> RedisResult<(u64, MultiplexedConnection)> { + if let Some(current) = self.cached() { + return Ok(current); + } + let _discovering = self.0.discovery.lock().await; + if let Some(current) = self.cached() { + return Ok(current); + } + let connection = self + .discover() + .await? + .get_multiplexed_async_connection_with_config(&self.0.config) + .await?; + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let generation = self.0.generations.fetch_add(1, Ordering::Relaxed); + *current = Some((generation, connection.clone())); + Ok((generation, connection)) + } + + pub(crate) async fn master(&self) -> RedisResult { + self.discover().await + } + + async fn discover(&self) -> RedisResult { + let master = tokio::time::timeout(self.0.connect_timeout, self.ask_sentinels()) + .await + .map_err(|_| timed_out("asking the Sentinels for the master timed out"))??; + self.master_client(master.get_connection_info().addr()) + } + + /// redis-rs sends `ROLE` to the master a Sentinel names before trusting it, so that probe gets + /// the data credentials and TLS a connection opened now would use. The Sentinel client is built + /// per discovery since a provider's credentials change after `new`. + async fn ask_sentinels(&self) -> RedisResult { + let data_auth = + config::authenticated(RedisConnectionInfo::default(), &self.0.credentials).await?; + self.0.sentinel_client(&data_auth)?.async_get_client().await + } + + fn master_client(&self, master: &ConnectionAddr) -> Result { + let (host, port) = match master { + ConnectionAddr::Tcp(host, port) | ConnectionAddr::TcpTls { host, port, .. } => { + (host.as_str(), *port) + } + _ => { + return Err(RedisError::from(( + ErrorKind::InvalidClientConfig, + "Sentinel named a master that is not a TCP address", + ))); + } + }; + let endpoint = Target::Standalone(crate::Standalone { + host: host.to_owned(), + port, + tls: self.0.endpoint.tls.clone(), + db: self.0.endpoint.db, + }); + let info = config::connection_info(&endpoint, &self.0.credentials, self.0.keepalive) + .map_err(|error| match error { + ConnectError::Redis(error) => error, + ConnectError::TlsFile { path, .. } => RedisError::from(( + ErrorKind::InvalidClientConfig, + "reading a Redis TLS file", + path, + )), + ConnectError::TlsKey { path, reason } => RedisError::from(( + ErrorKind::InvalidClientConfig, + "reading the Redis TLS key", + format!("{path}: {reason}"), + )), + })?; + redis::Client::open(info) + } + + fn forget(&self, generation: u64) { + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if current + .as_ref() + .is_some_and(|(cached, _)| *cached == generation) + { + *current = None; + } + } +} + +impl SentinelState { + fn sentinel_client(&self, data_auth: &RedisConnectionInfo) -> RedisResult { + let endpoint = &self.endpoint; + let sentinels = endpoint + .monitors + .iter() + .map(|node| addr(&node.host, node.port, endpoint.tls.as_ref())); + let builder = SentinelClientBuilder::new( + sentinels, + &endpoint.service_name, + SentinelServerType::Master, + )?; + let (username, password) = match &endpoint.monitor_auth { + Some(auth) => (auth.username.as_ref(), auth.password.as_ref()), + None => (None, None), + }; + let builder = match username { + Some(username) => builder.set_client_to_sentinel_username(username.expose()), + None => builder, + }; + let builder = match password { + Some(password) => builder.set_client_to_sentinel_password(password.expose()), + None => builder, + }; + let builder = match data_auth.username() { + Some(username) => builder.set_client_to_redis_username(username), + None => builder, + }; + let builder = match data_auth.password() { + Some(password) => builder.set_client_to_redis_password(password), + None => builder, + }; + let builder = match &self.tls { + Some((mode, certificates)) => builder + .set_client_to_sentinel_tls_mode(*mode) + .set_client_to_sentinel_certificates(certificates.clone()) + .set_client_to_redis_tls_mode(*mode) + .set_client_to_redis_certificates(certificates.clone()), + None => builder, + }; + builder.build() + } +} + +pub(crate) fn async_config(credentials: &Credentials, timeouts: Timeouts) -> AsyncConnectionConfig { + let config = AsyncConnectionConfig::new() + .set_connection_timeout(Some(timeouts.connect)) + .set_response_timeout(Some(timeouts.response)); + match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + } +} + +fn timed_out(message: &'static str) -> RedisError { + RedisError::from(std::io::Error::new(std::io::ErrorKind::TimedOut, message)) +} + +fn needs_rediscovery(error: &RedisError) -> bool { + matches!( + crate::Error::classify(error.clone(), Deployment::Sentinel), + crate::Error::Transport(_) + ) +} + +fn replied_read_only(values: &[Value]) -> bool { + values.iter().any(|value| { + matches!(value, Value::ServerError(error) if error.kind() == Some(ServerErrorKind::ReadOnly)) + }) +} + +impl ConnectionLike for SentinelConnection { + fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { + Box::pin(async move { + let (generation, mut connection) = self.connection().await?; + let reply = connection.req_packed_command(cmd).await; + let rediscover = match &reply { + Ok(value) => replied_read_only(std::slice::from_ref(value)), + Err(error) => needs_rediscovery(error), + }; + if rediscover { + self.forget(generation); + } + reply + }) + } + + fn req_packed_commands<'a>( + &'a mut self, + pipeline: &'a Pipeline, + offset: usize, + count: usize, + ) -> RedisFuture<'a, Vec> { + Box::pin(async move { + let (generation, mut connection) = self.connection().await?; + let replies = connection + .req_packed_commands(pipeline, offset, count) + .await; + let rediscover = match &replies { + Ok(values) => replied_read_only(values), + Err(error) => needs_rediscovery(error), + }; + if rediscover { + self.forget(generation); + } + replies + }) + } + + fn get_db(&self) -> i64 { + self.0.endpoint.db + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Node; + use litellm_redis_testing::{RedisSentinel, SERVICE_NAME}; + + fn endpoint(server: &RedisSentinel) -> Sentinel { + let (host, port) = server.sentinel(); + Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + } + } + + #[rstest::rstest] + #[tokio::test] + async fn a_late_failure_does_not_discard_the_reconnected_sentinel_master() { + let server = RedisSentinel::start().await.unwrap(); + let mut client = SentinelConnection::new( + &endpoint(&server), + &Credentials::NONE, + Timeouts::default(), + true, + ) + .unwrap(); + let (old_generation, _) = client.connection().await.unwrap(); + let original: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + client.forget(old_generation); + let replacement: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + assert_ne!(original, replacement); + + client.forget(old_generation); + let after_late_failure: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + + assert_eq!(replacement, after_late_failure); + } + + #[rstest::rstest] + #[tokio::test] + async fn a_read_only_reply_to_one_command_rediscovers_the_master() { + let server = RedisSentinel::start().await.unwrap(); + let mut client = SentinelConnection::new( + &endpoint(&server), + &Credentials::NONE, + Timeouts::default(), + true, + ) + .unwrap(); + let replica = redis::Client::open(format!("redis://127.0.0.1:{}", server.replica_port())) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let stale = client.0.generations.fetch_add(1, Ordering::Relaxed); + *client.0.current.lock().unwrap() = Some((stale, replica)); + + let on_replica = redis::cmd("SET") + .arg("k") + .arg("1") + .exec_async(&mut client) + .await; + let on_master = redis::cmd("SET") + .arg("k") + .arg("1") + .exec_async(&mut client) + .await; + + assert_eq!( + on_replica.unwrap_err().kind(), + ErrorKind::Server(ServerErrorKind::ReadOnly) + ); + on_master.unwrap(); + } +} diff --git a/litellm-rust/crates/redis-client/src/error.rs b/litellm-rust/crates/redis-client/src/error.rs index ed3b2c74004..54e55615e60 100644 --- a/litellm-rust/crates/redis-client/src/error.rs +++ b/litellm-rust/crates/redis-client/src/error.rs @@ -13,32 +13,93 @@ pub enum Error { } #[derive(Debug, thiserror::Error)] -#[error("connecting to Redis")] -pub struct ConnectError(#[from] RedisError); +pub enum ConnectError { + #[error("connecting to Redis")] + Redis(#[from] RedisError), + #[error("reading the Redis TLS file {path}")] + TlsFile { + path: String, + #[source] + source: std::io::Error, + }, + #[error("reading the Redis TLS key {path}")] + TlsKey { + path: String, + #[source] + reason: KeyError, + }, +} -impl From for Error { - fn from(error: RedisError) -> Self { - match error.kind() { - ErrorKind::Server( - ServerErrorKind::BusyLoading - | ServerErrorKind::ClusterDown - | ServerErrorKind::MasterDown - | ServerErrorKind::TryAgain - | ServerErrorKind::ReadOnly, - ) => Self::Transport(error), - ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error), - ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error), +#[derive(Debug, thiserror::Error)] +pub enum KeyError { + #[error("the key is encrypted, so ssl_password is required")] + MissingPassword, + #[error("ssl_password does not decrypt the key, or the key uses an unsupported cipher")] + Decrypt(#[source] pkcs8::Error), + #[error( + "the key uses legacy OpenSSL PEM encryption, which Rust cannot read. Convert it to PKCS#8 with `openssl pkcs8 -topk8`" + )] + LegacyEncryption, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Deployment { + Standalone, + Cluster, + Sentinel, +} + +// redis-py 5.3.1 redis/_parsers/base.py EXCEPTION_CLASSES["ERR"] entries that subclass ConnectionError +const CONNECTION_ERR_MESSAGES: &[&str] = &[ + "max number of clients reached", + "invalid password", + "Client sent AUTH, but no password is set", + "AUTH called without any password configured for the default user. Are you sure your configuration is correct?", +]; + +impl Error { + pub(crate) fn classify(error: RedisError, deployment: Deployment) -> Self { + match (error.kind(), deployment) { + (ErrorKind::Parse | ErrorKind::UnexpectedReturnType, _) => Self::Decode(error), + (ErrorKind::Server(kind), _) if is_unavailable(kind, deployment) => { + Self::Transport(error) + } + (ErrorKind::Server(ServerErrorKind::ResponseError), _) + if error + .detail() + .is_some_and(|detail| CONNECTION_ERR_MESSAGES.contains(&detail)) => + { + Self::Transport(error) + } + (ErrorKind::Extension, _) if matches!(error.code(), Some("NOAUTH" | "WRONGPASS")) => { + Self::Transport(error) + } + (ErrorKind::Server(_) | ErrorKind::Extension | ErrorKind::InvalidClientConfig, _) => { + Self::Server(error) + } _ => Self::Transport(error), } } } +fn is_unavailable(kind: ServerErrorKind, deployment: Deployment) -> bool { + matches!( + (kind, deployment), + (ServerErrorKind::BusyLoading, _) + | ( + ServerErrorKind::ClusterDown | ServerErrorKind::MasterDown, + Deployment::Cluster + ) + | (ServerErrorKind::ReadOnly, Deployment::Sentinel) + ) +} + #[cfg(test)] mod tests { use redis::{RedisError, Value}; use rstest::rstest; - use super::Error; + use super::{Deployment, Error}; fn reply(resp: &str) -> RedisError { match redis::parse_redis_value(resp.as_bytes()).unwrap() { @@ -47,23 +108,101 @@ mod tests { } } + fn is_transport(resp: &str, deployment: Deployment) -> bool { + match Error::classify(reply(resp), deployment) { + Error::Transport(_) => true, + Error::Server(_) => false, + other => panic!("{resp:?} classified as {other:?}"), + } + } + #[rstest] #[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")] #[case::no_script("-NOSCRIPT No matching script\r\n")] #[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")] #[case::generic("-ERR value is not an integer or out of range\r\n")] #[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")] - fn rejections_are_server_errors(#[case] resp: &str) { - assert!(matches!(Error::from(reply(resp)), Error::Server(_))); + #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] + #[case::out_of_memory("-OOM command not allowed when used memory > 'maxmemory'.\r\n")] + #[case::auth_arity("-ERR wrong number of arguments for 'auth' command\r\n")] + #[case::max_clients_with_suffix("-ERR max number of clients reached, try later\r\n")] + fn rejections_are_server_errors_on_every_deployment( + #[case] resp: &str, + #[values(Deployment::Standalone, Deployment::Cluster, Deployment::Sentinel)] + deployment: Deployment, + ) { + assert!(!is_transport(resp, deployment)); } #[rstest] #[case::loading("-LOADING Redis is loading the dataset in memory\r\n")] - #[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")] - #[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")] - #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] - #[case::read_only("-READONLY You can't write against a read only replica.\r\n")] - fn unavailable_nodes_are_transport_errors(#[case] resp: &str) { - assert!(matches!(Error::from(reply(resp)), Error::Transport(_))); + #[case::no_auth("-NOAUTH Authentication required.\r\n")] + #[case::wrong_pass("-WRONGPASS invalid username-password pair or user is disabled.\r\n")] + #[case::max_clients("-ERR max number of clients reached\r\n")] + #[case::invalid_password("-ERR invalid password\r\n")] + #[case::no_password_set("-ERR Client sent AUTH, but no password is set\r\n")] + #[case::no_default_password( + "-ERR AUTH called without any password configured for the default user. Are you sure your configuration is correct?\r\n" + )] + fn unavailable_nodes_are_transport_errors_on_every_deployment( + #[case] resp: &str, + #[values(Deployment::Standalone, Deployment::Cluster, Deployment::Sentinel)] + deployment: Deployment, + ) { + assert!(is_transport(resp, deployment)); + } + + #[rstest] + #[case::master_down_on_standalone( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Standalone, + false + )] + #[case::master_down_on_cluster( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Cluster, + true + )] + #[case::cluster_down_on_standalone( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Standalone, + false + )] + #[case::cluster_down_on_cluster( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Cluster, + true + )] + #[case::read_only_on_standalone( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Standalone, + false + )] + #[case::read_only_on_cluster( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Cluster, + false + )] + #[case::master_down_on_sentinel( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Sentinel, + false + )] + #[case::cluster_down_on_sentinel( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Sentinel, + false + )] + #[case::read_only_on_sentinel( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Sentinel, + true + )] + fn topology_decides_like_redis_py( + #[case] resp: &str, + #[case] deployment: Deployment, + #[case] transport: bool, + ) { + assert_eq!(is_transport(resp, deployment), transport); } } diff --git a/litellm-rust/crates/redis-client/src/events.rs b/litellm-rust/crates/redis-client/src/events.rs new file mode 100644 index 00000000000..9ea9ff4defa --- /dev/null +++ b/litellm-rust/crates/redis-client/src/events.rs @@ -0,0 +1,78 @@ +use std::{any::Any, future::Future, sync::Arc, time::Duration}; + +use litellm_circuit_breaker::{State, Transition}; + +use crate::Error; + +/// What the host attached to the calls made inside `tagged`, opaque to the client. +pub type CallTag = Arc; + +tokio::task_local! { + static TAG: CallTag; +} + +pub async fn tagged(tag: CallTag, operation: F) -> F::Output { + TAG.scope(tag, operation).await +} + +/// The tag of the `tagged` scope this task runs in. `RedisEvents::call` runs on the calling task, +/// so a sink reads its call's tag here. +pub fn current_tag() -> Option { + TAG.try_with(Arc::clone).ok() +} + +pub trait RedisEvents: Send + Sync { + fn call(&self, event: &CallEvent); + + fn breaker(&self, transition: &Transition); +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Operation { + Command(&'static str), + Script, + Pipeline { ops: usize }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Outcome { + Succeeded, + Failed(FailureKind), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FailureKind { + Timeout, + Transport, + Server, + Decode, + BreakerOpen, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallEvent { + pub operation: Operation, + pub duration: Duration, + pub outcome: Outcome, +} + +impl FailureKind { + pub fn of(error: &Error) -> Self { + match error { + Error::Transport(error) if error.is_timeout() => Self::Timeout, + Error::Transport(_) => Self::Transport, + Error::Server(_) => Self::Server, + Error::Decode(_) => Self::Decode, + Error::BreakerOpen => Self::BreakerOpen, + } + } +} + +impl Outcome { + pub(crate) fn of(result: &Result) -> Self { + match result { + Ok(_) => Self::Succeeded, + Err(error) => Self::Failed(FailureKind::of(error)), + } + } +} diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 47ff8677e06..2754666c329 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -1,20 +1,30 @@ +mod batch; mod client; -mod commands; +pub mod commands; mod config; mod connection; mod error; +mod events; mod namespace; mod pipeline; +mod pubsub; mod script; -#[cfg(test)] -mod test_support; -pub use client::{RedisClient, Ttl}; -pub use config::{ - ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig, +pub use batch::{Declared, RedisBatch}; +pub use client::{RedisClient, Ttl, WeakRedisClient}; +pub use commands::{Amount, Command, Expiry}; +pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts}; +pub use error::{ConnectError, Error, KeyError}; +pub use events::{ + CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged, +}; +pub use litellm_redis_types::{ + Cluster, Node, RedisUrl, Sentinel, Standalone, StaticAuth, Target, TlsOptions, UnixSocket, }; -pub use error::{ConnectError, Error}; pub use namespace::Namespace; pub use pipeline::Pipeline; -pub use redis::Value; +pub use redis::{ + BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value, + cluster_routing::Slot, +}; pub use script::Script; diff --git a/litellm-rust/crates/redis-client/src/pipeline.rs b/litellm-rust/crates/redis-client/src/pipeline.rs index e54c65e9605..0b6bb6a6673 100644 --- a/litellm-rust/crates/redis-client/src/pipeline.rs +++ b/litellm-rust/crates/redis-client/src/pipeline.rs @@ -1,30 +1,20 @@ -use std::{collections::HashMap, time::Duration}; +use std::collections::HashMap; use futures_util::future::join_all; -use redis::{ - Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value, - cluster_routing::Slot, -}; +use redis::{ErrorKind, RedisError, RedisResult, ServerErrorKind, Value, cluster_routing::Slot}; -use crate::{Error, Namespace, Script, commands, connection::Connection}; +use crate::{ + Command, Error, Namespace, + commands::{Op, Plan}, + connection::Connection, + script::ScriptCall, +}; pub struct Pipeline { namespace: Namespace, ops: Vec, } -struct Op { - command: Cmd, - key: String, - script: Option, -} - -struct ScriptCall { - script: Script, - keys: Vec, - args: Vec>, -} - type Settled = (usize, Result); impl Pipeline { @@ -35,60 +25,13 @@ impl Pipeline { } } - pub fn get(self, key: &str) -> Self { - let key = self.namespace.apply(key); - self.command(commands::get(key.clone()), key) + pub(crate) fn of(namespace: Namespace, ops: Vec) -> Self { + Self { namespace, ops } } - pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self { - let key = self.namespace.apply(key); - self.command(commands::set(key.clone(), value, ttl), key) - } - - pub fn del(self, key: &str) -> Self { - let key = self.namespace.apply(key); - self.command(commands::del(vec![key.clone()]), key) - } - - pub fn expire(self, key: &str, ttl: Duration) -> Self { - let key = self.namespace.apply(key); - self.command(commands::expire(key.clone(), ttl), key) - } - - pub fn incr_by(self, key: &str, delta: i64) -> Self { - let key = self.namespace.apply(key); - self.command(commands::incr_by(key.clone(), delta), key) - } - - pub fn incr_by_float(self, key: &str, delta: f64) -> Self { - let key = self.namespace.apply(key); - self.command(commands::incr_by_float(key.clone(), delta), key) - } - - pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self { - let keys: Vec = keys.iter().map(|key| self.namespace.apply(key)).collect(); - let args = args.to_redis_args(); - let op = Op { - command: script.evalsha(&keys, &args), - key: keys.first().cloned().unwrap_or_default(), - script: Some(ScriptCall { - script: script.clone(), - keys, - args, - }), - }; - self.push(op) - } - - fn command(self, command: Cmd, key: String) -> Self { - self.push(Op { - command, - key, - script: None, - }) - } - - fn push(mut self, op: Op) -> Self { + /// Appends one command in its pipelined form; `RedisClient::run` answers it in order, undecoded. + pub fn with(mut self, command: Command) -> Self { + let (op, _) = command.prepare(&self.namespace); self.ops.push(op); self } @@ -102,7 +45,7 @@ impl Pipeline { let settled = join_all( groups .iter() - .map(|group| self.run_group(connection.clone(), group)), + .map(|batches| self.run_in_order(connection, batches)), ) .await; let mut settled: Vec = settled.into_iter().flatten().collect(); @@ -110,183 +53,131 @@ impl Pipeline { settled.into_iter().map(|(_, result)| result).collect() } - fn groups(&self, connection: &Connection) -> Vec> { + async fn run_in_order(&self, connection: &Connection, batches: &[Vec]) -> Vec { + let mut settled = Vec::with_capacity(batches.iter().map(Vec::len).sum()); + for batch in batches { + settled.extend(self.run_group(connection.clone(), batch).await); + } + settled + } + + fn groups(&self, connection: &Connection) -> Vec>> { match connection { _ if self.ops.is_empty() => Vec::new(), - Connection::Standalone(_) => vec![(0..self.ops.len()).collect()], + Connection::Standalone(_) | Connection::Sentinel(_) => { + vec![vec![(0..self.ops.len()).collect()]] + } Connection::Cluster(_) => { let mut by_slot: HashMap> = HashMap::new(); - for (index, op) in self.ops.iter().enumerate() { + for index in 0..self.ops.len() { by_slot - .entry(Slot::for_key(&op.key)) + .entry(Slot::for_key(&self.ops[index].key)) .or_default() .push(index); } - by_slot.into_values().collect() + by_slot + .into_values() + .map(|indices| self.retry_safe_batches(indices)) + .collect() } } } + fn retry_safe_batches(&self, indices: Vec) -> Vec> { + indices + .chunk_by(|&left, &right| self.ops[left].idempotent && self.ops[right].idempotent) + .map(<[usize]>::to_vec) + .collect() + } + + /// A script the server had not cached ran nothing, so it alone runs again, the way litellm's + /// `RedisBatch` reruns it through redis-py's `Script`: `EVALSHA`, then `SCRIPT LOAD` and + /// `EVALSHA` on `NOSCRIPT`. + async fn settle_script( + &self, + connection: Connection, + index: usize, + reply: RedisResult, + ) -> (usize, RedisResult) { + match (&self.ops[index].plan, reply) { + (Plan::Script(call), Err(error)) if is_missing_script(&error) => { + (index, rerun_missing_script(connection, call).await) + } + (_, reply) => (index, reply), + } + } + async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec { + let deployment = connection.deployment(); let mut pipeline = redis::pipe(); pipeline.ignore_errors(); - for &index in indices { - pipeline.add_command(self.ops[index].command.clone()); - } - let replies: RedisResult>> = - pipeline.query_async(&mut connection).await; + let widths: Vec = indices + .iter() + .map(|&index| { + let commands = self.ops[index].plan.pipelined(); + let width = commands.len(); + for command in commands { + pipeline.add_command(command); + } + width + }) + .collect(); + let replies: RedisResult>> = match pipeline.is_empty() { + true => Ok(Vec::new()), + false => pipeline.query_async(&mut connection).await, + }; let replies = match replies { Ok(replies) => replies, Err(error) => { return indices .iter() - .map(|&index| (index, Err(Error::from(error.clone())))) + .map(|&index| (index, Err(Error::classify(error.clone(), deployment)))) .collect(); } }; - let mut settled = Vec::with_capacity(indices.len()); - for (&index, reply) in indices.iter().zip(replies) { - settled.push((index, self.settle(&mut connection, index, reply).await)); - } - settled - } - - async fn settle( - &self, - connection: &mut Connection, - index: usize, - reply: RedisResult, - ) -> Result { - match (reply, &self.ops[index].script) { - (Err(error), Some(call)) - if error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) => - { - Ok(call - .script - .invocation(&call.keys, &call.args) - .invoke_async(connection) - .await?) - } - (reply, _) => Ok(reply?), - } + let mut replies = replies.into_iter(); + let per_op: Vec<(usize, RedisResult)> = indices + .iter() + .zip(widths) + .map(|(&index, width)| match &self.ops[index].plan { + Plan::Reply(value) => (index, Ok(value.clone())), + _ => (index, first_reply(replies.by_ref().take(width))), + }) + .collect(); + join_all( + per_op + .into_iter() + .map(|(index, reply)| self.settle_script(connection.clone(), index, reply)), + ) + .await + .into_iter() + .map(|(index, reply)| { + ( + index, + reply.map_err(|error| Error::classify(error, deployment)), + ) + }) + .collect() } } -#[cfg(test)] -mod tests { - use std::time::{Duration, Instant}; - - use redis::Value; - use rstest::rstest; - - use crate::{ - DEFAULT_TIMEOUT, Error, Script, - test_support::{Server, Topology, flush_scripts}, - }; - - #[rstest] - #[tokio::test] - async fn each_op_settles_on_its_own_in_input_order( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client - .set("text", "not a number", Duration::from_secs(60)) - .await - .unwrap(); - let pipeline = client - .pipeline() - .incr_by("a", 1) - .incr_by("text", 1) - .incr_by("c", 5) - .set("b", "2", Duration::from_secs(60)) - .get("b") - .incr_by("a", 1) - .get("missing"); - - let results = client.run(&pipeline).await; - - assert_eq!(results.len(), 7); - assert_eq!(results[0].as_ref().unwrap(), &Value::Int(1)); - assert!( - matches!(results[1], Err(Error::Server(_))), - "{:?}", - results[1] - ); - assert_eq!(results[2].as_ref().unwrap(), &Value::Int(5)); - assert_eq!(results[3].as_ref().unwrap(), &Value::Okay); - assert_eq!( - results[4].as_ref().unwrap(), - &Value::BulkString(b"2".to_vec()) - ); - assert_eq!(results[5].as_ref().unwrap(), &Value::Int(2)); - assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); - } - - #[rstest] - #[tokio::test] - async fn a_flushed_script_in_a_pipeline_reruns_alone( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("ns", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.eval::(&incr, &["counter"], 1).await.unwrap(); - flush_scripts(&client).await; - let pipeline = client - .pipeline() - .eval(&incr, &["counter"], 2) - .incr_by("other", 7); - - let results = client.run(&pipeline).await; - - assert_eq!(results[0].as_ref().unwrap(), &Value::Int(3)); - assert_eq!(results[1].as_ref().unwrap(), &Value::Int(7)); - assert_eq!(client.get::("counter").await.unwrap(), Some(3)); - } - - #[rstest] - #[tokio::test] - async fn an_unreachable_server_fails_every_op_as_transport( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); - - server.freeze().await; - let results = client.run(&pipeline).await; - server.thaw().await; - - assert_eq!(results.len(), 2); - assert!( - matches!(results[0], Err(Error::Transport(_))), - "{:?}", - results[0] - ); - assert!( - matches!(results[1], Err(Error::Transport(_))), - "{:?}", - results[1] - ); - } - - #[rstest] - #[tokio::test] - async fn an_empty_pipeline_returns_without_asking_redis( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_secs(2)).await; - - server.freeze().await; - let started = Instant::now(); - let results = client.run(&client.pipeline()).await; - let elapsed = started.elapsed(); - server.thaw().await; - - assert!(results.is_empty()); - assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); - } +async fn rerun_missing_script(mut connection: Connection, call: &ScriptCall) -> RedisResult { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut connection) + .await +} + +fn is_missing_script(error: &RedisError) -> bool { + error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) +} + +// RedisBatch settles an op from its first reply and fails it on any reply that failed +fn first_reply(replies: impl Iterator>) -> RedisResult { + replies + .reduce(|first, next| match (first, next) { + (Err(error), _) | (Ok(_), Err(error)) => Err(error), + (first, Ok(_)) => first, + }) + .unwrap_or(Ok(Value::Nil)) } diff --git a/litellm-rust/crates/redis-client/src/pubsub.rs b/litellm-rust/crates/redis-client/src/pubsub.rs new file mode 100644 index 00000000000..b6b18d3e481 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/pubsub.rs @@ -0,0 +1,225 @@ +use std::{sync::Arc, time::Duration}; + +use futures_util::{StreamExt, stream::BoxStream}; +use redis::{ConnectionInfo, RedisResult, aio::PubSubStream}; + +use crate::{ + ConnectError, Credentials, Standalone, Target, Timeouts, config, + connection::{Connection, SentinelConnection}, +}; + +// litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py: _BACKOFF_INITIAL_SECONDS, _BACKOFF_MAX_SECONDS +const BACKOFF_INITIAL: Duration = Duration::from_secs(5); +const BACKOFF_MAX: Duration = Duration::from_secs(60); + +/// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends +/// `SUBSCRIBE`, authenticated once when it connects. +#[derive(Clone)] +pub(crate) struct Subscriber { + source: Source, + credentials: Credentials, + connect_timeout: Duration, +} + +#[derive(Clone)] +enum Source { + Nodes(Arc<[ConnectionInfo]>), + Sentinel(SentinelConnection), +} + +impl Subscriber { + pub(crate) fn new( + endpoint: &Target, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + connection: &Connection, + ) -> Result { + let source = match (endpoint, connection) { + (_, Connection::Sentinel(sentinel)) => Source::Sentinel(sentinel.clone()), + (Target::Cluster(cluster), _) => Source::Nodes( + cluster + .nodes + .iter() + .map(|node| { + let node = Target::Standalone(Standalone { + host: node.host.clone(), + port: node.port, + tls: cluster.tls.clone(), + db: 0, + }); + config::connection_info(&node, credentials, keepalive) + }) + .collect::>()?, + ), + _ => Source::Nodes(Arc::new([config::connection_info( + endpoint, + credentials, + keepalive, + )?])), + }; + Ok(Self { + source, + credentials: credentials.clone(), + connect_timeout: timeouts.connect, + }) + } + + pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec> { + futures_util::stream::unfold( + (self, channel, None::, BACKOFF_INITIAL), + |(subscriber, channel, live, backoff)| async move { + let mut live = live; + let mut backoff = backoff; + loop { + let Some(current) = live.as_mut() else { + match subscriber.subscribe(&channel).await { + Ok(subscribed) => { + live = Some(subscribed); + backoff = BACKOFF_INITIAL; + } + Err(error) => { + tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(BACKOFF_MAX); + } + } + continue; + }; + match current.next().await { + Some(message) => { + let payload = message.get_payload_bytes().to_vec(); + return Some((payload, (subscriber, channel, live, backoff))); + } + None => live = None, + } + tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(BACKOFF_MAX); + } + }, + ) + .boxed() + } + + async fn subscribe(&self, channel: &str) -> RedisResult { + match &self.source { + Source::Sentinel(sentinel) => { + let master = sentinel.master().await?; + self.subscribe_on(master.get_connection_info().clone(), channel) + .await + } + Source::Nodes(nodes) => { + let mut last = None; + for node in nodes.iter() { + match self.subscribe_on(node.clone(), channel).await { + Ok(live) => return Ok(live), + Err(error) => last = Some(error), + } + } + Err(last.expect("an endpoint always has at least one node")) + } + } + } + + async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult { + let settings = + config::authenticated(info.redis_settings().clone(), &self.credentials).await?; + let client = redis::Client::open(info.set_redis_settings(settings))?; + let mut pubsub = tokio::time::timeout(self.connect_timeout, client.get_async_pubsub()) + .await + .map_err(|_| { + redis::RedisError::from(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "connecting the Redis subscription timed out", + )) + })??; + subscribe_connected(&mut pubsub, channel, self.connect_timeout).await?; + Ok(pubsub.into_on_message()) + } +} + +async fn subscribe_connected( + pubsub: &mut redis::aio::PubSub, + channel: &str, + connect_timeout: Duration, +) -> RedisResult<()> { + tokio::time::timeout(connect_timeout, pubsub.subscribe(channel)) + .await + .map_err(|_| { + redis::RedisError::from(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "subscribing to the Redis channel timed out", + )) + })? +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use futures_util::StreamExt; + use litellm_redis_testing::{RedisServer, TcpProxy}; + use rstest::rstest; + + use super::subscribe_connected; + + #[rstest] + #[tokio::test] + async fn a_stalled_pubsub_subscribe_times_out_and_a_fresh_pubsub_recovers() { + let server = RedisServer::start().await.unwrap(); + let proxy = TcpProxy::start(server.port()).await.unwrap(); + let client = redis::Client::open(format!("redis://127.0.0.1:{}", proxy.port())).unwrap(); + let mut pubsub = tokio::time::timeout(Duration::from_secs(2), client.get_async_pubsub()) + .await + .expect("connecting the initial pubsub should finish") + .unwrap(); + tokio::time::timeout( + Duration::from_secs(2), + subscribe_connected(&mut pubsub, "existing", Duration::from_secs(1)), + ) + .await + .expect("the initial subscription acknowledgement should arrive") + .unwrap(); + + proxy.blackhole_open_connections(); + let result = tokio::time::timeout( + Duration::from_secs(1), + subscribe_connected(&mut pubsub, "stalled", Duration::from_millis(100)), + ) + .await + .expect("the stalled subscription should return before the safety deadline"); + assert!( + matches!(&result, Err(error) if error.is_timeout()), + "{result:?}" + ); + drop(pubsub); + + let mut recovered = tokio::time::timeout(Duration::from_secs(2), client.get_async_pubsub()) + .await + .expect("a fresh pubsub connection should finish") + .unwrap(); + tokio::time::timeout( + Duration::from_secs(2), + subscribe_connected(&mut recovered, "recovered", Duration::from_secs(1)), + ) + .await + .expect("the fresh subscription acknowledgement should arrive") + .unwrap(); + + let publisher = redis::Client::open(server.url()).unwrap(); + let mut connection = publisher.get_multiplexed_async_connection().await.unwrap(); + let subscriber_count: i64 = redis::cmd("PUBLISH") + .arg("recovered") + .arg("payload") + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(subscriber_count, 1); + let message = tokio::time::timeout(Duration::from_secs(2), recovered.on_message().next()) + .await + .expect("the fresh subscription should receive the published payload") + .unwrap(); + assert_eq!(message.get_payload::().unwrap(), "payload"); + } +} diff --git a/litellm-rust/crates/redis-client/src/script.rs b/litellm-rust/crates/redis-client/src/script.rs index c5d73133ecc..6dbdbafe9c9 100644 --- a/litellm-rust/crates/redis-client/src/script.rs +++ b/litellm-rust/crates/redis-client/src/script.rs @@ -1,13 +1,18 @@ +use std::sync::Arc; + use redis::{Cmd, ScriptInvocation}; +/// A Lua script, cheap to clone into every call that runs it. #[derive(Clone)] -pub struct Script(redis::Script); +pub struct Script(Arc); impl Script { pub fn new(source: &str) -> Self { - Self(redis::Script::new(source)) + Self(Arc::new(redis::Script::new(source))) } + /// The `EVALSHA` a pipeline sends, like redis-py. A `NOSCRIPT` reply is settled by + /// [`Script::invocation`] alone, which loads the script and sends it again. pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec]) -> Cmd { redis::cmd("EVALSHA") .arg(self.0.get_hash()) @@ -24,65 +29,10 @@ impl Script { } } -#[cfg(test)] -mod tests { - use std::time::Duration; - - use rstest::rstest; - - use super::Script; - use crate::{ - DEFAULT_TIMEOUT, Error, - test_support::{Server, Topology, flush_scripts}, - }; - - #[rstest] - #[tokio::test] - async fn script_keys_are_namespaced( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("a", DEFAULT_TIMEOUT).await; - let bare = server.client("", DEFAULT_TIMEOUT).await; - let read = Script::new("return redis.call('GET', KEYS[1])"); - bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap(); - - let value: Option = client - .eval(&read, &["x"], Vec::::new()) - .await - .unwrap(); - - assert_eq!(value.as_deref(), Some("1")); - } - - #[rstest] - #[tokio::test] - async fn a_failing_script_is_a_server_error_and_keeps_the_connection( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.rpush("list", "item").await.unwrap(); - - let failed = client.eval::(&incr, &["list"], 1).await; - - assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); - assert_eq!(client.eval::(&incr, &["counter"], 3).await.unwrap(), 3); - } - - #[rstest] - #[tokio::test] - async fn a_flushed_script_is_loaded_again( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.eval::(&incr, &["counter"], 1).await.unwrap(); - - flush_scripts(&client).await; - - assert_eq!(client.eval::(&incr, &["counter"], 1).await.unwrap(), 2); - } +/// A script call in a pipeline, kept apart from plain commands because it is the one entry a +/// pipeline may send again. +pub(crate) struct ScriptCall { + pub(crate) script: Script, + pub(crate) keys: Vec, + pub(crate) args: Vec>, } diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs deleted file mode 100644 index 64a174b1543..00000000000 --- a/litellm-rust/crates/redis-client/src/test_support.rs +++ /dev/null @@ -1,125 +0,0 @@ -use std::time::Duration; - -use litellm_circuit_breaker::Breaker; -use litellm_redis_testing::{RedisCluster, RedisServer}; -use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; - -use crate::{ - ClusterEndpoint, ClusterNode, Endpoint, Namespace, RedisClient, RedisConfig, - connection::Connection, -}; - -#[derive(Clone, Copy, Debug)] -pub(crate) enum Topology { - Standalone, - Cluster, -} - -pub(crate) enum Server { - Standalone(RedisServer), - Cluster(RedisCluster), -} - -impl Server { - pub(crate) async fn start(topology: Topology) -> Self { - match topology { - Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()), - Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()), - } - } - - pub(crate) fn endpoint(&self) -> Endpoint { - match self { - Self::Standalone(server) => Endpoint::Url(server.url().to_owned()), - Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { - nodes: cluster - .nodes() - .map(|(host, port)| ClusterNode { - host: host.to_owned(), - port, - }) - .collect(), - username: None, - password: None, - tls: false, - }), - } - } - - pub(crate) async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient { - self.guarded_client(namespace, timeout, None).await - } - - pub(crate) async fn guarded_client( - &self, - namespace: &str, - timeout: Duration, - breaker: Option, - ) -> RedisClient { - RedisClient::new(RedisConfig { - endpoint: self.endpoint(), - namespace: Namespace::new(namespace), - timeout, - breaker, - }) - .unwrap() - } - - pub(crate) async fn restart(&self) { - match self { - Self::Standalone(server) => server.restart().await.unwrap(), - Self::Cluster(cluster) => cluster.restart().await.unwrap(), - } - } - - pub(crate) async fn stop(&self) { - match self { - Self::Standalone(server) => server.stop().await.unwrap(), - Self::Cluster(cluster) => cluster.stop().await.unwrap(), - } - } - - pub(crate) async fn resume(&self) { - match self { - Self::Standalone(server) => server.resume().await.unwrap(), - Self::Cluster(cluster) => cluster.resume().await.unwrap(), - } - } - - pub(crate) async fn freeze(&self) { - match self { - Self::Standalone(server) => server.freeze().await.unwrap(), - Self::Cluster(cluster) => cluster.freeze().await.unwrap(), - } - } - - pub(crate) async fn thaw(&self) { - match self { - Self::Standalone(server) => server.thaw().await.unwrap(), - Self::Cluster(cluster) => cluster.thaw().await.unwrap(), - } - } -} - -pub(crate) async fn client_id(client: &RedisClient, key: &str) -> i64 { - let command = redis::cmd("CLIENT").arg("ID").to_owned(); - match client.connection.clone() { - Connection::Standalone(mut connection) => { - command.query_async(&mut connection).await.unwrap() - } - Connection::Cluster(mut connection) => { - let node = Route::with_key(key, SlotAddr::Master); - let routing = RoutingInfo::SingleNode(SingleNodeRoutingInfo::SpecificNode(node)); - redis::from_redis_value(connection.route_command(command, routing).await.unwrap()) - .unwrap() - } - } -} - -pub(crate) async fn flush_scripts(client: &RedisClient) { - redis::cmd("SCRIPT") - .arg("FLUSH") - .query_async::<()>(&mut client.connection.clone()) - .await - .unwrap(); -} diff --git a/litellm-rust/crates/redis-client/tests/batch.rs b/litellm-rust/crates/redis-client/tests/batch.rs new file mode 100644 index 00000000000..45ca419951e --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/batch.rs @@ -0,0 +1,330 @@ +mod support; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_redis_client::commands; +use litellm_redis_client::{ + Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Declared, Error, Expiry, Namespace, Operation, + RedisBatch, RedisClient, RedisConfig, RedisEvents, Script, Timeouts, Ttl, +}; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, endpoint}; + +const MINUTE: Option = Some(Expiry::KeepExisting(Duration::from_secs(60))); + +// RedisBatch pipelines a standalone flush and runs each op on its own on a cluster +fn expected_trips(topology: Topology, alone: &[Operation]) -> Vec { + match topology { + Topology::Standalone => vec![Operation::Pipeline { ops: alone.len() }], + Topology::Cluster => alone.to_vec(), + } +} + +#[derive(Default)] +struct RoundTrips(Mutex>); + +impl RoundTrips { + fn seen(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +impl RedisEvents for RoundTrips { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(event.operation); + } + + fn breaker(&self, _transition: &Transition) {} +} + +fn observed(server: &RedisDeployment) -> (RedisClient, Arc) { + let trips = Arc::new(RoundTrips::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(server), + credentials: Credentials::NONE, + namespace: Namespace::new("ns"), + timeouts: Timeouts::default(), + breaker: None, + events: Some(trips.clone()), + health_check: None, + keepalive: true, + }) + .unwrap(); + (client, trips) +} + +#[rstest] +#[tokio::test] +async fn declared_ops_wait_for_the_flush_and_share_one_round_trip( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(3), MINUTE)); + let total = batch.add(commands::incr_by_float("total", Amount::Float(1.5), MINUTE)); + let written = batch.add(commands::set("other", "v", Duration::from_secs(60))); + let read = batch.add(commands::get::("value")); + + let outside = connect(&server, "ns", DEFAULT_TIMEOUT); + assert_eq!( + outside.send(commands::get::("counter")).await.unwrap(), + None + ); + assert!(trips.seen().is_empty()); + + batch.flush().await; + + let expected = expected_trips( + topology, + &[ + Operation::Command("INCRBYFLOAT"), + Operation::Command("INCRBYFLOAT"), + Operation::Command("SET"), + Operation::Command("GET"), + ], + ); + let mut seen = trips.seen(); + seen.sort_by_key(|operation| format!("{operation:?}")); + let mut sorted = expected.clone(); + sorted.sort_by_key(|operation| format!("{operation:?}")); + assert_eq!(seen, sorted); + assert_eq!(counter.resolve().await.unwrap(), 3.0); + assert_eq!(total.resolve().await.unwrap(), 1.5); + written.resolve().await.unwrap(); + assert_eq!(read.resolve().await.unwrap(), None); + assert_eq!(trips.seen().len(), expected.len()); +} + +#[rstest] +#[tokio::test] +async fn resolving_before_a_flush_sends_everything_declared_so_far( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let first = batch.add(commands::incr_by_float("a", Amount::Int(1), MINUTE)); + let second = batch.add(commands::incr_by_float("b", Amount::Int(2), MINUTE)); + let increment = Operation::Command("INCRBYFLOAT"); + + assert_eq!(second.resolve().await.unwrap(), 2.0); + let late = batch.add(commands::incr_by_float("a", Amount::Int(10), MINUTE)); + assert_eq!(first.resolve().await.unwrap(), 1.0); + assert_eq!( + trips.seen(), + expected_trips(topology, &[increment, increment]) + ); + + assert_eq!(late.resolve().await.unwrap(), 11.0); + let both: Vec = expected_trips(topology, &[increment, increment]) + .into_iter() + .chain(expected_trips(topology, &[increment])) + .collect(); + assert_eq!(trips.seen(), both); + batch.flush().await; + assert_eq!(trips.seen().len(), both.len()); +} + +#[rstest] +#[tokio::test] +async fn a_pipelined_batch_is_one_pipeline_on_every_topology( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = RedisBatch::pipelined(client); + let declared = + ["a", "b", "c"].map(|key| batch.add(commands::incr_by_float(key, Amount::Int(2), MINUTE))); + + let totals = Declared::all(declared).await; + + assert_eq!( + totals.into_iter().map(Result::unwrap).collect::>(), + [2.0; 3] + ); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 3 }]); +} + +#[rstest] +#[tokio::test] +async fn each_declared_op_gets_its_own_typed_result( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set( + "text", + "not a number", + Duration::from_secs(60), + )) + .await + .unwrap(); + let batch = client.batch(); + let rejected = batch.add(commands::incr_by_float("text", Amount::Int(1), MINUTE)); + let undecodable = batch.add(commands::get::("text")); + let missing = batch.add(commands::get::("missing")); + let script = batch.add(commands::eval::>( + &Script::new("return {KEYS[1], ARGV[1]}"), + &["k"], + "arg", + )); + let removed = batch.add(commands::del(&["text"])); + + batch.flush().await; + + assert!(matches!(rejected.resolve().await, Err(Error::Server(_)))); + assert!(matches!(undecodable.resolve().await, Err(Error::Decode(_)))); + assert_eq!(missing.resolve().await.unwrap(), None); + assert_eq!(script.resolve().await.unwrap(), ["k", "arg"]); + assert_eq!(removed.resolve().await.unwrap(), 1); +} + +#[rstest] +#[tokio::test] +async fn an_op_whose_flush_was_cancelled_resolves_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(5)); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE)); + + server.freeze().await.unwrap(); + let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; + server.thaw().await.unwrap(); + + assert!(cancelled.is_err()); + assert!(matches!(counter.resolve().await, Err(Error::Transport(_)))); +} + +#[rstest] +#[tokio::test] +async fn a_settled_op_waits_for_someone_else_to_flush_and_sends_nothing_itself( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let counter = tokio::spawn( + batch + .add(commands::incr_by_float("counter", Amount::Int(4), MINUTE)) + .settled(), + ); + + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(!counter.is_finished()); + assert!(trips.seen().is_empty()); + + batch.flush().await; + + assert_eq!(counter.await.unwrap().unwrap(), 4.0); + assert_eq!( + trips.seen(), + expected_trips(topology, &[Operation::Command("INCRBYFLOAT")]) + ); +} + +#[rstest] +#[tokio::test] +async fn a_settled_op_whose_flush_was_cancelled_resolves_as_transport() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(5)); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE)); + + server.freeze().await.unwrap(); + let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; + server.thaw().await.unwrap(); + + assert!(cancelled.is_err()); + assert!(matches!(counter.settled().await, Err(Error::Transport(_)))); +} + +#[rstest] +#[tokio::test] +async fn clones_are_the_same_client_and_a_second_connection_is_not() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let twin = connect(&server, "", DEFAULT_TIMEOUT); + + assert!(client.same_client(&client.clone())); + assert!(!client.same_client(&twin)); +} + +#[rstest] +#[tokio::test] +async fn mget_reads_each_key_once_and_answers_every_key_in_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("a", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("c", "3", Duration::from_secs(60))) + .await + .unwrap(); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + let batch = client.batch(); + + let values = batch.add(commands::mget::(&["c", "missing", "a", "c"])); + batch.flush().await; + + assert_eq!( + values.resolve().await.unwrap(), + [ + Some("3".to_owned()), + None, + Some("1".to_owned()), + Some("3".to_owned()) + ] + ); + let stats = support::command_stats(&server).await; + assert!(support::calls(&stats, "mget") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "get"), 0, "{stats}"); + assert_eq!(support::keys_read(&server).await, 3); +} + +// RedisBatch's _Increment refreshes the TTL in a pipeline, while its cluster path calls +// async_increment, which keeps a TTL the key already has +#[rstest] +#[case::pipelined(Topology::Standalone, false)] +#[case::alone_on_a_cluster(Topology::Cluster, true)] +#[tokio::test] +async fn a_batched_increment_treats_the_ttl_like_redis_batch( + #[case] topology: Topology, + #[case] keeps_the_ttl: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("counter", "1", Duration::from_secs(30))) + .await + .unwrap(); + let batch = client.batch(); + + let counter = batch.add(commands::incr_by_float( + "counter", + Amount::Int(1), + Some(Expiry::KeepExisting(Duration::from_secs(600))), + )); + batch.flush().await; + + assert_eq!(counter.resolve().await.unwrap(), 2.0); + let Ttl::Expires(left) = client.send(commands::ttl("counter")).await.unwrap() else { + panic!("expected an expiring counter") + }; + assert_eq!(left <= Duration::from_secs(30), keeps_the_ttl, "{left:?}"); +} diff --git a/litellm-rust/crates/redis-client/tests/client.rs b/litellm-rust/crates/redis-client/tests/client.rs new file mode 100644 index 00000000000..e9c6cd7e712 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/client.rs @@ -0,0 +1,1114 @@ +mod support; + +use std::{ + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use litellm_redis_testing::{RedisCluster, RedisDeployment, RedisServer, Topology, eventually}; +use rstest::rstest; + +use litellm_circuit_breaker::{Breaker, Config, State}; + +use litellm_redis_client::commands; +use litellm_redis_client::{ + Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Error, Expiry, FailureKind, Namespace, + Operation, Outcome, RedisClient, RedisConfig, RedisEvents, RedisUrl, Script, Standalone, + Target, Timeouts, Ttl, Value, +}; +use litellm_secret_value::SecretValue; +use support::{KEEP, connect, connect_guarded, connections, endpoint}; + +async fn persist(client: &RedisClient, key: &str) { + let set = Script::new("return redis.call('SET', KEYS[1], '0')"); + client + .send(commands::eval::<()>(&set, &[key], Vec::::new())) + .await + .unwrap(); +} + +#[derive(Clone, Copy, Debug)] +enum Counter { + IncrBy, + IncrByFloat, +} + +const MINUTE: Duration = Duration::from_secs(60); + +async fn write(client: &RedisClient, counter: Counter, key: &str, expiry: Expiry) { + match (counter, expiry) { + (Counter::IncrBy, Expiry::KeepExisting(ttl)) => { + client.send(commands::incr_by(key, 1, ttl)).await.unwrap(); + } + (Counter::IncrBy, Expiry::Refresh(_)) => unreachable!("increment_cache never refreshes"), + (Counter::IncrByFloat, expiry) => { + client + .send(commands::incr_by_float( + key, + Amount::Float(0.5), + Some(expiry), + )) + .await + .unwrap(); + } + } +} + +async fn persist_as(client: &RedisClient, key: &str) { + let source = "return redis.call('SET', KEYS[1], '1')"; + client + .send(commands::eval::( + &Script::new(source), + &[key], + Vec::::new(), + )) + .await + .unwrap(); +} + +fn remaining(ttl: Ttl) -> Duration { + match ttl { + Ttl::Expires(left) => left, + other => panic!("expected an expiring key, got {other:?}"), + } +} + +#[rstest] +#[case::increment_cache(Counter::IncrBy, Expiry::KeepExisting(MINUTE))] +#[case::async_increment(Counter::IncrByFloat, Expiry::KeepExisting(MINUTE))] +#[case::async_increment_refreshing(Counter::IncrByFloat, Expiry::Refresh(MINUTE))] +#[tokio::test] +async fn a_counter_write_expires_a_new_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[case] counter: Counter, + #[case] expiry: Expiry, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + write(&client, counter, "fresh", expiry).await; + + let left = remaining(client.send(commands::ttl("fresh")).await.unwrap()); + assert!(left > Duration::from_secs(50) && left <= MINUTE, "{left:?}"); +} + +#[rstest] +#[tokio::test] +async fn keep_existing_leaves_a_running_ttl_and_refresh_resets_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let short = Expiry::Refresh(Duration::from_secs(30)); + let long = Duration::from_secs(600); + write(&client, Counter::IncrByFloat, "kept", short).await; + write(&client, Counter::IncrByFloat, "refreshed", short).await; + write( + &client, + Counter::IncrBy, + "kept_int", + Expiry::KeepExisting(Duration::from_secs(30)), + ) + .await; + + write( + &client, + Counter::IncrByFloat, + "kept", + Expiry::KeepExisting(long), + ) + .await; + write( + &client, + Counter::IncrByFloat, + "refreshed", + Expiry::Refresh(long), + ) + .await; + write( + &client, + Counter::IncrBy, + "kept_int", + Expiry::KeepExisting(long), + ) + .await; + + assert!( + remaining(client.send(commands::ttl("kept")).await.unwrap()) <= Duration::from_secs(30) + ); + assert!( + remaining(client.send(commands::ttl("kept_int")).await.unwrap()) <= Duration::from_secs(30) + ); + assert!( + remaining(client.send(commands::ttl("refreshed")).await.unwrap()) + > Duration::from_secs(500) + ); +} + +#[rstest] +#[tokio::test] +async fn keep_existing_expires_a_key_that_had_no_ttl( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(Counter::IncrBy, Counter::IncrByFloat)] counter: Counter, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + persist_as(&client, "persistent").await; + + write(&client, counter, "persistent", KEEP).await; + + let left = remaining(client.send(commands::ttl("persistent")).await.unwrap()); + assert!(left > Duration::from_secs(50), "{left:?}"); +} + +#[rstest] +#[tokio::test] +async fn counters_send_the_commands_redis_cache_sends( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + client + .send(commands::incr_by("int", 1, MINUTE)) + .await + .unwrap(); + client + .send(commands::incr_by_float("kept", Amount::Int(1), Some(KEEP))) + .await + .unwrap(); + client + .send(commands::incr_by_float( + "refreshed", + Amount::Float(0.5), + Some(Expiry::Refresh(MINUTE)), + )) + .await + .unwrap(); + client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + "piped", + Amount::Float(0.25), + Some(Expiry::Refresh(MINUTE)), + )) + .with(commands::incr_by_float("bare", Amount::Float(0.25), None)), + ) + .await; + + let stats = support::command_stats(&server).await; + let calls = |command: &str| support::calls(&stats, command); + assert_eq!(calls("incrby"), 1); + assert_eq!(calls("incrbyfloat"), 4); + assert_eq!(calls("ttl"), 2); + assert_eq!(calls("expire"), 4); + assert_eq!( + calls("eval") + calls("evalsha") + calls("pttl") + calls("pexpire"), + 0 + ); +} + +#[derive(Default)] +struct Recorded(Mutex>); + +#[rstest] +#[case::beyond_double_precision(9_007_199_254_740_992_i64, 1_i64)] +#[case::negative_beyond_double_precision(-9_007_199_254_740_992_i64, -1_i64)] +#[case::maximum_integer(i64::MAX - 1, 1_i64)] +#[case::minimum_integer(i64::MIN + 1, -1_i64)] +#[tokio::test] +async fn integer_increments_preserve_every_bit( + #[case] initial: i64, + #[case] delta: i64, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("integer", initial, Duration::from_secs(60))) + .await + .unwrap(); + let expected = initial + delta; + let result = client + .send(commands::incr_by("integer", delta, MINUTE)) + .await + .unwrap(); + assert_eq!(result, expected); + assert_eq!( + client.send(commands::get::("integer")).await.unwrap(), + Some(expected) + ); +} + +#[rstest] +#[tokio::test] +async fn outage_events_report_the_failure_kind( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(false, true)] timeout: bool, + #[values(false, true)] pipeline: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(300), None), + breaker: None, + events: Some(recorded.clone()), + credentials: Credentials::NONE, + health_check: None, + keepalive: true, + }) + .unwrap(); + client.ping().await.unwrap(); + if timeout { + server.freeze().await.unwrap(); + } else { + server.stop().await.unwrap(); + } + let result = if pipeline { + client + .run(&client.pipeline().with(commands::get::("missing"))) + .await + .remove(0) + .map(drop) + } else { + client + .send(commands::get::("missing")) + .await + .map(drop) + }; + if timeout { + server.thaw().await.unwrap(); + } else { + server.resume().await.unwrap(); + } + assert!(matches!(result, Err(Error::Transport(_)))); + let expected = if timeout { + FailureKind::Timeout + } else { + FailureKind::Transport + }; + assert_eq!( + recorded.0.lock().unwrap().last().unwrap().outcome, + Outcome::Failed(expected) + ); +} + +impl RedisEvents for Recorded { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(*event); + } + + fn breaker(&self, _transition: &litellm_circuit_breaker::Transition) {} +} + +#[rstest] +#[tokio::test] +async fn every_call_reports_its_operation_and_outcome( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: Some(recorded.clone()), + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let _ = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + let _ = client.send(commands::get::("missing")).await; + let _ = client + .run( + &client + .pipeline() + .with(commands::get::("a")) + .with(commands::get::("b")), + ) + .await; + let failures = client + .run(&client.pipeline().with(commands::incr_by_float( + "list", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + ))) + .await; + assert!(matches!(failures[0], Err(Error::Server(_)))); + + let seen: Vec<(Operation, Outcome)> = recorded + .0 + .lock() + .unwrap() + .iter() + .map(|event| (event.operation, event.outcome)) + .collect(); + assert_eq!( + seen, + vec![ + (Operation::Command("SET"), Outcome::Succeeded), + ( + Operation::Command("INCRBY"), + Outcome::Failed(FailureKind::Server) + ), + (Operation::Command("GET"), Outcome::Succeeded), + (Operation::Pipeline { ops: 2 }, Outcome::Succeeded), + ( + Operation::Pipeline { ops: 1 }, + Outcome::Failed(FailureKind::Server) + ), + ] + ); +} + +#[rstest] +#[tokio::test] +async fn namespaces_isolate_keys_on_the_server( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let a = connect(&server, "a", DEFAULT_TIMEOUT); + let b = connect(&server, "b", DEFAULT_TIMEOUT); + let bare = connect(&server, "", DEFAULT_TIMEOUT); + + a.send(commands::set("x", "1", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + a.send(commands::get::("x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!( + a.send(commands::get::("a:x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!( + bare.send(commands::get::("a:x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!(b.send(commands::get::("x")).await.unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_command_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + let before = connections(&server).await; + + let rejected = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + + assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); + assert_eq!( + client + .send(commands::get::("list")) + .await + .unwrap() + .as_deref(), + Some("item") + ); + assert_eq!(connections(&server).await, before); +} + +#[rstest] +#[tokio::test] +async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("float", "15.0", Duration::from_secs(60))) + .await + .unwrap(); + + let decoded = client.send(commands::get::("float")).await; + + assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); + assert_eq!( + client.send(commands::get::("missing")).await.unwrap(), + None + ); +} + +#[rstest] +#[tokio::test] +async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + client + .send(commands::set("first", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("second", "2", Duration::from_secs(60))) + .await + .unwrap(); + + server.freeze().await.unwrap(); + let timed_out = client.send(commands::get::("first")).await; + server.thaw().await.unwrap(); + + assert!( + matches!(timed_out, Err(Error::Transport(_))), + "{timed_out:?}" + ); + assert_eq!( + client + .send(commands::get::("second")) + .await + .unwrap() + .as_deref(), + Some("2") + ); +} + +#[rstest] +#[tokio::test] +async fn the_same_client_recovers_after_redis_restarts( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + + server.restart().await.unwrap(); + + let recovered = eventually(async || { + client + .send(commands::incr_by( + "after_restart", + 1, + Duration::from_secs(60), + )) + .await + .ok() + }) + .await; + assert_eq!(recovered, 1); +} + +#[rstest] +#[tokio::test] +async fn a_client_created_while_redis_is_down_works_once_redis_is_up( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server.stop().await.unwrap(); + + let client = connect(&server, "", Duration::from_millis(500)); + let started = Instant::now(); + let while_down = client.ping().await; + let failed_in = started.elapsed(); + server.resume().await.unwrap(); + + assert!(failed_in < Duration::from_secs(2), "{failed_in:?}"); + + assert!( + matches!(while_down, Err(Error::Transport(_))), + "{while_down:?}" + ); + let recovered = eventually(async || client.ping().await.ok()).await; + assert_eq!(recovered, ()); +} + +#[rstest] +#[tokio::test] +async fn empty_batches_succeed_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + assert_eq!( + client.send(commands::mget::(&[])).await.unwrap(), + Vec::>::new() + ); + assert_eq!(client.send(commands::del(&[])).await.unwrap(), 0); +} + +#[rstest] +#[tokio::test] +async fn mget_and_del_span_slots_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("a", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("c", "3", Duration::from_secs(60))) + .await + .unwrap(); + + let values = client + .send(commands::mget::(&["c", "b", "a"])) + .await + .unwrap(); + let deleted = client.send(commands::del(&["a", "b", "c"])).await.unwrap(); + + assert_eq!( + values, + vec![Some("3".to_owned()), None, Some("1".to_owned())] + ); + assert_eq!(deleted, 2); + assert_eq!( + client + .send(commands::mget::(&["a", "c"])) + .await + .unwrap(), + vec![None, None] + ); +} + +#[rstest] +#[tokio::test] +async fn set_nx_only_writes_a_missing_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + let first = client + .send(commands::set_nx("lock", "owner-1", Duration::from_secs(60))) + .await + .unwrap(); + let second = client + .send(commands::set_nx("lock", "owner-2", Duration::from_secs(60))) + .await + .unwrap(); + + assert!(first); + assert!(!second); + assert_eq!( + client + .send(commands::get::("lock")) + .await + .unwrap() + .as_deref(), + Some("owner-1") + ); +} + +#[rstest] +#[tokio::test] +async fn ttl_reports_missing_persistent_and_remaining_time( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + persist(&client, "persistent").await; + client + .send(commands::set("expiring", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::incr_by_float( + "expired_later", + Amount::Float(0.5), + Some(KEEP), + )) + .await + .unwrap(); + assert!( + client + .send(commands::expire("expired_later", Duration::from_secs(30))) + .await + .unwrap() + ); + + assert_eq!( + client.send(commands::ttl("missing")).await.unwrap(), + Ttl::Missing + ); + assert_eq!( + client.send(commands::ttl("persistent")).await.unwrap(), + Ttl::Persistent + ); + assert!(matches!( + client.send(commands::ttl("expiring")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); + assert!(matches!( + client.send(commands::ttl("expired_later")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) + )); + assert!( + !client + .send(commands::expire("missing", Duration::from_secs(30))) + .await + .unwrap() + ); +} + +#[rstest] +#[tokio::test] +async fn incr_by_float_accumulates_fractions( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + client + .send(commands::incr_by_float( + "spend", + Amount::Float(0.25), + Some(KEEP), + )) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::incr_by_float( + "spend", + Amount::Float(0.5), + Some(KEEP) + )) + .await + .unwrap(), + 0.75 + ); +} + +#[rstest] +#[tokio::test] +async fn concurrent_commands_on_one_cloned_client_all_apply( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + let tasks: Vec<_> = (0..100) + .map(|_| { + let client = client.clone(); + tokio::spawn(async move { + for _ in 0..10 { + client + .send(commands::incr_by("shared", 1, Duration::from_secs(60))) + .await + .unwrap(); + } + }) + }) + .collect(); + for task in tasks { + task.await.unwrap(); + } + + assert_eq!( + client.send(commands::get::("shared")).await.unwrap(), + Some(1000) + ); +} + +fn breaker(failure_threshold: u32) -> Breaker { + Breaker::new( + Config { + failure_threshold, + recovery_timeout: Duration::from_millis(200), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ) +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + + server.freeze().await.unwrap(); + let first = client.send(commands::get::("k")).await; + let second = client.send(commands::get::("k")).await; + let started = Instant::now(); + let third = client.send(commands::get::("k")).await; + let rejected_in = started.elapsed(); + server.thaw().await.unwrap(); + + assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); + assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); + assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + + let recovered = eventually(async || client.send(commands::get::("k")).await.ok()).await; + assert_eq!(recovered, None); + assert!( + matches!(breaker.state(), State::Closed { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn rejected_commands_never_open_the_breaker( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded(&server, "", DEFAULT_TIMEOUT, Some(breaker.clone())); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let rejected = [ + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + ]; + + assert!( + matches!(rejected[2], Err(Error::Server(_))), + "{:?}", + rejected[2] + ); + assert!( + matches!(breaker.state(), State::Closed { hard: 0, .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn a_probe_redis_rejects_still_proves_redis_answers_and_closes_the_breaker( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + assert!( + matches!(breaker.state(), State::Open { .. }), + "{:?}", + breaker.state() + ); + tokio::time::sleep(Duration::from_millis(250)).await; + + let probe = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + + assert!(matches!(probe, Err(Error::Server(_))), "{probe:?}"); + assert!( + matches!(breaker.state(), State::Closed { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn a_replica_refusing_a_write_does_not_reset_a_run_of_failures() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + client.send(commands::get::("k")).await.unwrap(); + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + server + .cli_on_every_node(&["REPLICAOF", "127.0.0.1", "1"]) + .await + .unwrap(); + + let refused = client + .send(commands::set("k", "v", Duration::from_secs(60))) + .await; + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + + assert!(matches!(refused, Err(Error::Server(_))), "{refused:?}"); + assert!( + matches!(breaker.state(), State::Open { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_fails_every_pipeline_op_without_sending( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker(1))); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )); + + server.freeze().await.unwrap(); + let failed = client.run(&pipeline).await; + let rejected = client.run(&pipeline).await; + server.thaw().await.unwrap(); + + assert!( + matches!(failed[0], Err(Error::Transport(_))), + "{:?}", + failed[0] + ); + assert!( + matches!(rejected[0], Err(Error::BreakerOpen)), + "{:?}", + rejected[0] + ); + assert!( + matches!(rejected[1], Err(Error::BreakerOpen)), + "{:?}", + rejected[1] + ); +} + +#[rstest] +#[tokio::test] +async fn a_host_endpoint_connects_to_its_database() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let connect = |target| { + RedisClient::new(RedisConfig { + target, + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + }; + let url = |url: String| { + Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db: None, + username: None, + password: None, + }) + }; + let in_db_3 = connect(Target::Standalone(Standalone { + host: "127.0.0.1".to_owned(), + port, + tls: None, + db: 3, + })) + .unwrap(); + let in_db_0 = connect(url(server.url().to_owned())).unwrap(); + let via_url = connect(url(format!("{}/3", server.url()))).unwrap(); + + in_db_3 + .send(commands::set("k", "3", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + in_db_0.send(commands::get::("k")).await.unwrap(), + None + ); + assert_eq!( + via_url + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("3") + ); +} + +async fn ping_calls(cluster: &RedisCluster) -> Vec { + cluster + .cli_outputs(&["INFO", "commandstats"]) + .await + .unwrap() + .iter() + .map(|info| { + info.lines() + .find_map(|line| line.strip_prefix("cmdstat_ping:calls=")) + .and_then(|rest| rest.split(',').next()) + .and_then(|calls| calls.trim().parse().ok()) + .unwrap_or(0) + }) + .collect() +} + +fn health_checked( + server: &RedisDeployment, + interval: Duration, + timeout: Duration, + breaker: Option, + events: Option>, +) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(server), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(timeout, None), + breaker, + events: events.map(|events| events as Arc), + health_check: Some(interval), + keepalive: true, + }) + .unwrap() +} + +// redis-py 5.3.1 redis/asyncio/connection.py check_health: a connection idle longer than +// health_check_interval gets a PING before the next command, a busy one does not +#[rstest] +#[tokio::test] +async fn an_idle_client_pings_every_node_before_the_next_command_like_redis_py() { + let server = RedisDeployment::start(Topology::Cluster).await.unwrap(); + let RedisDeployment::Cluster(cluster) = &server else { + unreachable!("started a cluster") + }; + let client = health_checked(&server, Duration::from_secs(3), DEFAULT_TIMEOUT, None, None); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = ping_calls(cluster).await; + + client + .send(commands::set("k", "2", Duration::from_secs(60))) + .await + .unwrap(); + let busy = ping_calls(cluster).await; + tokio::time::sleep(Duration::from_millis(3500)).await; + client + .send(commands::set("k", "3", Duration::from_secs(60))) + .await + .unwrap(); + let idle = ping_calls(cluster).await; + + assert_eq!(busy, before); + assert_eq!( + idle, + before.iter().map(|calls| calls + 1).collect::>(), + "before {before:?}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_failed_health_check_is_neither_a_call_nor_a_breaker_failure() { + let server = RedisDeployment::start(Topology::Cluster).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let breaker = breaker(2); + let client = health_checked( + &server, + Duration::from_millis(200), + Duration::from_millis(300), + Some(breaker.clone()), + Some(recorded.clone()), + ); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + + server.freeze().await.unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + let failed = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + + assert!(matches!(failed, Err(Error::Transport(_))), "{failed:?}"); + assert!( + matches!(breaker.state(), State::Closed { timeouts: 1, .. }), + "{:?}", + breaker.state() + ); + assert_eq!(recorded.0.lock().unwrap().len(), 2); +} + +#[rstest] +#[case::int(Amount::Int(1), "1")] +#[case::float_keeps_its_point(Amount::Float(1.0), "1.0")] +#[case::small_float(Amount::Float(1e-5), "1e-05")] +#[case::large_float(Amount::Float(1e16), "1e+16")] +#[case::fraction(Amount::Float(0.1 + 0.2), "0.30000000000000004")] +fn amounts_are_written_like_redis_py_encodes_them(#[case] amount: Amount, #[case] expected: &str) { + assert_eq!( + redis::ToRedisArgs::to_redis_args(&amount), + vec![expected.as_bytes().to_vec()] + ); +} diff --git a/litellm-rust/crates/redis-client/tests/config.rs b/litellm-rust/crates/redis-client/tests/config.rs new file mode 100644 index 00000000000..101c6529e94 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/config.rs @@ -0,0 +1,290 @@ +use rstest::rstest; + +use litellm_redis_client::{Credentials, RedisUrl, Sentinel, StaticAuth, Target, TlsOptions}; +use litellm_secret_value::SecretValue; + +#[rstest] +#[case::url(format!("{:?}", Target::Url(RedisUrl { url: SecretValue::new("redis://user:credential-value@cache:6379/0"), tls: None, db: None, username: None, password: None })))] +#[case::static_password(format!("{:?}", Credentials::Static { username: None, password: Some(SecretValue::new("credential-value")) }))] +#[case::sentinel_password(format!("{:?}", Target::Sentinel(Sentinel { + monitors: Vec::new(), + service_name: "primary".to_owned(), + monitor_auth: Some(StaticAuth { username: None, password: Some(SecretValue::new("credential-value")) }), + tls: None, + db: 0, +})))] +fn debug_output_never_shows_the_secret(#[case] debug: String) { + assert!(!debug.contains("credential-value"), "{debug}"); + assert!(debug.contains("[REDACTED]"), "{debug}"); +} + +#[derive(Clone, Copy, Debug)] +enum Ca { + Trusted, + TrustedData, + TrustedDirectory, + SystemRoots, +} + +fn directory_holding(ca: &std::path::Path) -> std::path::PathBuf { + let directory = std::env::temp_dir().join(format!("litellm-ca-dir-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + std::fs::copy(ca, directory.join("ca.pem")).unwrap(); + directory +} + +#[derive(Clone, Copy, Debug)] +enum ClientKey { + Absent, + Plain, + Encrypted { password: Option<&'static str> }, +} + +async fn tls_server(require_client_certificate: bool) -> litellm_redis_testing::RedisTlsServer { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + litellm_redis_testing::RedisTlsServer::start(require_client_certificate) + .await + .unwrap() +} + +fn tls_client( + server: &litellm_redis_testing::RedisTlsServer, + host: &str, + tls: TlsOptions, +) -> Result { + litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: host.to_owned(), + port: server.port(), + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) +} + +fn tls_options( + files: &litellm_redis_testing::TlsFiles, + ca: Ca, + verify_certificate: bool, + check_hostname: bool, + client_key: ClientKey, +) -> TlsOptions { + let (keyfile, key_password) = match client_key { + ClientKey::Absent => (None, None), + ClientKey::Plain => (Some(files.client_key.clone()), None), + ClientKey::Encrypted { password } => ( + Some(files.encrypted_client_key.clone()), + password.map(SecretValue::new), + ), + }; + TlsOptions { + ca_certs: matches!(ca, Ca::Trusted).then(|| files.ca.clone()), + ca_data: matches!(ca, Ca::TrustedData).then(|| std::fs::read_to_string(&files.ca).unwrap()), + ca_path: matches!(ca, Ca::TrustedDirectory).then(|| directory_holding(&files.ca)), + certfile: keyfile.is_some().then(|| files.client_cert.clone()), + keyfile, + key_password, + verify_certificate, + check_hostname, + } +} + +async fn ping_over_tls( + require_client_certificate: bool, + host: &str, + ca: Ca, + verify_certificate: bool, + check_hostname: bool, + client_key: ClientKey, +) -> Result<(), litellm_redis_client::Error> { + let server = tls_server(require_client_certificate).await; + let tls = tls_options( + server.files(), + ca, + verify_certificate, + check_hostname, + client_key, + ); + tls_client(&server, host, tls).unwrap().ping().await +} + +#[rstest] +#[case::trusted_ca_without_hostname_check_like_redis_py( + "127.0.0.1", + Ca::Trusted, + true, + false, + true +)] +#[case::trusted_ca_as_pem_text_like_ssl_ca_data("127.0.0.1", Ca::TrustedData, true, false, true)] +#[case::trusted_ca_in_a_directory_like_ssl_ca_path( + "127.0.0.1", + Ca::TrustedDirectory, + true, + false, + true +)] +#[case::hostname_check_with_a_matching_name("localhost", Ca::Trusted, true, true, true)] +#[case::hostname_check_against_an_unlisted_address("127.0.0.1", Ca::Trusted, true, true, false)] +#[case::an_untrusted_ca_is_refused("127.0.0.1", Ca::SystemRoots, true, false, false)] +#[case::cert_reqs_none_skips_verification("127.0.0.1", Ca::SystemRoots, false, false, true)] +#[tokio::test] +async fn tls_options_decide_which_servers_are_trusted( + #[case] host: &str, + #[case] ca: Ca, + #[case] verify_certificate: bool, + #[case] check_hostname: bool, + #[case] connects: bool, +) { + let result = ping_over_tls( + false, + host, + ca, + verify_certificate, + check_hostname, + ClientKey::Absent, + ) + .await; + + assert_eq!(result.is_ok(), connects, "{result:?}"); + if let Err(error) = result { + assert!( + matches!(error, litellm_redis_client::Error::Transport(_)), + "{error:?}" + ); + } +} + +#[rstest] +#[case::with_the_client_certificate(ClientKey::Plain, true)] +#[case::with_an_encrypted_key_and_its_ssl_password( + ClientKey::Encrypted { password: Some(litellm_redis_testing::CLIENT_KEY_PASSWORD) }, + true +)] +#[case::without_a_client_certificate(ClientKey::Absent, false)] +#[tokio::test] +async fn mutual_tls_sends_the_configured_client_certificate( + #[case] client_key: ClientKey, + #[case] connects: bool, +) { + let result = ping_over_tls(true, "127.0.0.1", Ca::Trusted, true, false, client_key).await; + + assert_eq!(result.is_ok(), connects, "{result:?}"); +} + +#[rstest] +#[case::without_ssl_password(None, "ssl_password is required")] +#[case::with_the_wrong_ssl_password(Some("wrong-password"), "does not decrypt the key")] +#[tokio::test] +async fn an_encrypted_key_that_cannot_be_decrypted_fails_before_connecting( + #[case] password: Option<&'static str>, + #[case] reason: &str, +) { + let server = tls_server(true).await; + let tls = tls_options( + server.files(), + Ca::Trusted, + true, + false, + ClientKey::Encrypted { password }, + ); + + let error = tls_client(&server, "127.0.0.1", tls).err().unwrap(); + + let litellm_redis_client::ConnectError::TlsKey { + path, + reason: cause, + } = &error + else { + panic!("expected a TLS key error, got {error:?}"); + }; + assert_eq!( + *path, + server.files().encrypted_client_key.display().to_string() + ); + assert!(cause.to_string().contains(reason), "{cause}"); +} + +#[rstest] +fn a_missing_tls_file_names_the_path() { + let tls = TlsOptions { + ca_certs: Some("/nonexistent/ca.pem".into()), + ..TlsOptions::verified() + }; + + let error = litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: "127.0.0.1".to_owned(), + port: 6379, + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .err() + .unwrap(); + + assert!(error.to_string().contains("/nonexistent/ca.pem"), "{error}"); +} + +#[rstest] +fn a_legacy_encrypted_pem_key_is_refused_with_a_conversion_hint() { + let directory = std::env::temp_dir().join(format!("litellm-legacy-key-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + let keyfile = directory.join("client.key"); + std::fs::write( + &keyfile, + "-----BEGIN RSA PRIVATE KEY-----\nProc-Type: 4,ENCRYPTED\nDEK-Info: AES-256-CBC,00112233445566778899AABBCCDDEEFF\n\nAAAA\n-----END RSA PRIVATE KEY-----\n", + ) + .unwrap(); + let tls = TlsOptions { + certfile: Some(keyfile.clone()), + keyfile: Some(keyfile), + key_password: Some(SecretValue::new("legacy-password")), + ..TlsOptions::verified() + }; + + let error = litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: "127.0.0.1".to_owned(), + port: 6379, + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .err() + .unwrap(); + + assert!( + matches!( + &error, + litellm_redis_client::ConnectError::TlsKey { + reason: litellm_redis_client::KeyError::LegacyEncryption, + .. + } + ), + "{error:?}" + ); +} diff --git a/litellm-rust/crates/redis-client/tests/connection.rs b/litellm-rust/crates/redis-client/tests/connection.rs new file mode 100644 index 00000000000..4efa6038cf0 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/connection.rs @@ -0,0 +1,329 @@ +mod support; + +use std::{collections::BTreeSet, sync::Arc, time::Duration}; + +use litellm_redis_testing::{ + ClientConnection, RedisDeployment, RedisSentinel, RedisServer, SERVICE_NAME, Topology, + eventually, +}; +use rstest::rstest; + +use litellm_redis_client::commands; +use litellm_redis_client::{ + Credentials, CredentialsProvider, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, + Standalone, Target, Timeouts, +}; +use litellm_secret_value::SecretValue; +use support::{Rotating, connect_as, connections, proxied}; + +#[rstest] +#[tokio::test] +async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect_uses_the_latest( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = Arc::new(Rotating::new("alice", "first")); + let client = connect_as( + &server, + Credentials::Provider(CredentialsProvider(provider.clone())), + ); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = connections(&server).await; + + provider.rotate("bob", "second"); + let rotated = eventually(async || { + let now = connections(&server).await; + (users(&now) == BTreeSet::from(["bob"])).then_some(now) + }) + .await; + server + .cli_on_every_node(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + server + .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let after_reconnect = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + + let reconnected = connections(&server).await; + + assert_eq!(users(&before), BTreeSet::from(["alice"])); + assert_eq!(ids(&rotated), ids(&before)); + assert_eq!(after_reconnect, "1"); + assert!(!reconnected.is_empty()); + assert_eq!(users(&reconnected), BTreeSet::from(["bob"])); +} + +fn users(connections: &[ClientConnection]) -> BTreeSet<&str> { + connections + .iter() + .map(|connection| connection.user.as_str()) + .collect() +} + +fn ids(connections: &[ClientConnection]) -> Vec<(u16, u64)> { + connections + .iter() + .map(|connection| (connection.node, connection.id)) + .collect() +} + +#[rstest] +#[tokio::test] +async fn an_out_of_range_database_is_a_server_rejection() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let client = RedisClient::new(RedisConfig { + target: Target::Standalone(Standalone { + host: "127.0.0.1".into(), + port, + tls: None, + db: 999, + }), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(2), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + let result = client.ping().await; + assert!(matches!(result, Err(Error::Server(_))), "{result:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_credential_is_a_transport_failure() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + server + .cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let with_password = |password: &str| { + connect_as( + &server, + Credentials::Static { + username: None, + password: Some(SecretValue::new(password)), + }, + ) + }; + with_password("right").ping().await.unwrap(); + + let rejected = with_password("wrong").ping().await; + + assert!(matches!(rejected, Err(Error::Transport(_))), "{rejected:?}"); +} + +fn sentinel_client(sentinel: &RedisSentinel, credentials: Credentials) -> RedisClient { + let (host, port) = sentinel.sentinel(); + RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(1), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +#[rstest] +#[tokio::test] +async fn sentinel_writes_follow_the_master_through_a_failover() { + let sentinel = RedisSentinel::start().await.unwrap(); + let client = sentinel_client(&sentinel, Credentials::NONE); + client + .send(commands::set("before", "1", Duration::from_secs(600))) + .await + .unwrap(); + eventually(async || { + let replica = + redis::Client::open(format!("redis://127.0.0.1:{}", sentinel.replica_port())).ok()?; + let mut replica = replica.get_multiplexed_async_connection().await.ok()?; + let value: Option = redis::cmd("GET") + .arg("before") + .query_async(&mut replica) + .await + .ok()?; + value + }) + .await; + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(600))) + .await + .ok() + }) + .await; + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!( + client + .send(commands::get::("before")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + let promoted = redis::Client::open(format!("redis://127.0.0.1:{promoted}")).unwrap(); + let mut promoted = promoted.get_multiplexed_async_connection().await.unwrap(); + let after: Option = redis::cmd("GET") + .arg("after") + .query_async(&mut promoted) + .await + .unwrap(); + assert_eq!(after.as_deref(), Some("2")); +} + +const DATA_PASSWORD: &str = "data-secret"; + +#[rstest] +#[tokio::test] +async fn sentinel_reaches_a_master_that_requires_a_password() { + let sentinel = RedisSentinel::start_with_password(DATA_PASSWORD) + .await + .unwrap(); + let client = sentinel_client( + &sentinel, + Credentials::Static { + username: None, + password: Some(SecretValue::new(DATA_PASSWORD)), + }, + ); + + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("1") + ); +} + +#[rstest] +#[tokio::test] +async fn sentinel_rediscovers_with_the_credentials_the_provider_holds_now() { + let sentinel = RedisSentinel::start_with_password(DATA_PASSWORD) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = Arc::new(Rotating::new("alice", "first")); + let client = sentinel_client( + &sentinel, + Credentials::Provider(CredentialsProvider(provider.clone())), + ); + client + .send(commands::set("k", "1", Duration::from_secs(600))) + .await + .unwrap(); + + provider.rotate("bob", "second"); + sentinel + .cli_on_data_nodes(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let after_rediscovery = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + + assert_eq!(after_rediscovery, "1"); +} + +#[rstest] +#[tokio::test] +async fn a_timed_out_command_replaces_a_connection_that_stopped_answering( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let proxies = proxied(&server).await; + let client = RedisClient::new(RedisConfig { + target: proxies.target.clone(), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(500), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("k", "1", Duration::from_secs(600))) + .await + .unwrap(); + + proxies.blackhole_open_connections(); + let stalled = client.send(commands::get::("k")).await; + let next = client.send(commands::get::("k")).await; + + assert!( + matches!(&stalled, Err(Error::Transport(error)) if error.is_timeout()), + "{stalled:?}" + ); + assert_eq!(next.unwrap().as_deref(), Some("1")); +} diff --git a/litellm-rust/crates/redis-client/tests/pipeline.rs b/litellm-rust/crates/redis-client/tests/pipeline.rs new file mode 100644 index 00000000000..5aed70b42dd --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pipeline.rs @@ -0,0 +1,362 @@ +mod support; + +use std::time::{Duration, Instant}; + +use litellm_redis_client::Value; +use litellm_redis_client::commands; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; + +use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Expiry, Script}; +use support::connect; + +async fn begin_migration(cluster: &litellm_redis_testing::RedisCluster, key: &str) { + let (_, first) = cluster.nodes().next().unwrap(); + let slot = cluster + .cli_on_node(first, &["CLUSTER", "KEYSLOT", key]) + .await + .unwrap(); + let slot = slot.trim(); + let number: u16 = slot.parse().unwrap(); + let nodes = cluster + .cli_on_node(first, &["CLUSTER", "NODES"]) + .await + .unwrap(); + let owner = nodes.lines().find(|line| line.split_whitespace().skip(8).any(|range| { + let (start, end) = range.split_once('-').unwrap_or((range, range)); + matches!((start.parse::(), end.parse::()), (Ok(start), Ok(end)) if (start..=end).contains(&number)) + })).unwrap(); + let target = nodes.lines().find(|line| *line != owner).unwrap(); + let address = |line: &str| { + let value = line.split_whitespace().nth(1).unwrap(); + value + .split_once('@') + .unwrap() + .0 + .rsplit_once(':') + .unwrap() + .1 + .parse::() + .unwrap() + }; + let source_port = address(owner); + let target_port = address(target); + let source_id = owner.split_whitespace().next().unwrap(); + let target_id = target.split_whitespace().next().unwrap(); + assert_eq!( + cluster + .cli_on_node(source_port, &["GET", key]) + .await + .unwrap() + .trim(), + "0" + ); + assert_eq!( + cluster + .cli_on_node( + target_port, + &["CLUSTER", "SETSLOT", slot, "IMPORTING", source_id] + ) + .await + .unwrap() + .trim(), + "OK" + ); + assert_eq!( + cluster + .cli_on_node( + source_port, + &["CLUSTER", "SETSLOT", slot, "MIGRATING", target_id] + ) + .await + .unwrap() + .trim(), + "OK" + ); + let state = cluster + .cli_on_node(source_port, &["CLUSTER", "NODES"]) + .await + .unwrap(); + assert!( + state.contains(&format!("[{slot}->-{target_id}]")), + "{state}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() { + let cluster = litellm_redis_testing::RedisCluster::start().await.unwrap(); + let server = RedisDeployment::Cluster(cluster); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let counter = "{migration}:counter"; + client + .send(commands::set(counter, 0, Duration::from_secs(60))) + .await + .unwrap(); + let RedisDeployment::Cluster(cluster) = &server else { + unreachable!() + }; + begin_migration(cluster, counter).await; + let both_keys = Script::new("return redis.call('MGET', KEYS[1], KEYS[2])"); + + let results = client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + counter, + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::eval::( + &both_keys, + &[counter, "{migration}:missing"], + Vec::::new(), + )), + ) + .await; + + assert_eq!( + client.send(commands::get::(counter)).await.unwrap(), + Some(1) + ); + assert_eq!( + results[0].as_ref().unwrap(), + &Value::BulkString(b"1".to_vec()) + ); + assert!( + matches!(&results[1], Err(Error::Server(error)) if error.code() == Some("TRYAGAIN")), + "{:?}", + results[1] + ); +} + +#[rstest] +#[tokio::test] +async fn each_op_settles_on_its_own_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set( + "text", + "not a number", + Duration::from_secs(60), + )) + .await + .unwrap(); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "text", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(5), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::set("b", "2", Duration::from_secs(60))) + .with(commands::get::("b")) + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::get::("missing")); + + let results = client.run(&pipeline).await; + + assert_eq!(results.len(), 7); + assert_eq!( + results[0].as_ref().unwrap(), + &Value::BulkString(b"1".to_vec()) + ); + assert!( + matches!(results[1], Err(Error::Server(_))), + "{:?}", + results[1] + ); + assert_eq!( + results[2].as_ref().unwrap(), + &Value::BulkString(b"5".to_vec()) + ); + assert_eq!(results[3].as_ref().unwrap(), &Value::Okay); + assert_eq!( + results[4].as_ref().unwrap(), + &Value::BulkString(b"2".to_vec()) + ); + assert_eq!( + results[5].as_ref().unwrap(), + &Value::BulkString(b"2".to_vec()) + ); + assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); +} + +#[rstest] +#[case::cached(true, Value::BulkString(b"2".to_vec()))] +#[case::lost_by_the_server(false, Value::Nil)] +#[tokio::test] +async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py( + #[case] cached: bool, + #[case] read_in_the_same_pipeline: Value, +) { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::eval::(&incr, &["warmup"], 1)) + .await + .unwrap(); + if !cached { + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + } + let pipeline = client + .pipeline() + .with(commands::eval::(&incr, &["counter"], 2)) + .with(commands::get::("counter")); + + let results = client.run(&pipeline).await; + + assert_eq!(results[0].as_ref().unwrap(), &Value::Int(2)); + assert_eq!(results[1].as_ref().unwrap(), &read_in_the_same_pipeline); + assert_eq!( + client.send(commands::get::("counter")).await.unwrap(), + Some(2) + ); +} + +#[rstest] +#[tokio::test] +async fn an_unreachable_server_fails_every_op_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )); + + server.freeze().await.unwrap(); + let results = client.run(&pipeline).await; + server.thaw().await.unwrap(); + + assert_eq!(results.len(), 2); + assert!( + matches!(results[0], Err(Error::Transport(_))), + "{:?}", + results[0] + ); + assert!( + matches!(results[1], Err(Error::Transport(_))), + "{:?}", + results[1] + ); +} + +#[rstest] +#[tokio::test] +async fn an_empty_pipeline_returns_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(2)); + + server.freeze().await.unwrap(); + let started = Instant::now(); + let results = client.run(&client.pipeline()).await; + let elapsed = started.elapsed(); + server.thaw().await.unwrap(); + + assert!(results.is_empty()); + assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_lost( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&deployment, "scripts", DEFAULT_TIMEOUT); + let count = Script::new("return redis.call('INCR', KEYS[1])"); + deployment + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + deployment + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + let first = client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + "{a}spend", + Amount::Float(0.5), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::eval::( + &count, + &["{a}"], + Vec::::new(), + )) + .with(commands::eval::( + &count, + &["b"], + Vec::::new(), + )), + ) + .await; + let second = client + .run(&client.pipeline().with(commands::eval::( + &count, + &["{a}"], + Vec::::new(), + ))) + .await; + + let replies: Vec = first + .into_iter() + .chain(second) + .map(Result::unwrap) + .collect(); + assert_eq!( + replies, + [ + Value::BulkString(b"0.5".to_vec()), + Value::Int(1), + Value::Int(1), + Value::Int(2) + ] + ); + assert_eq!( + client.send(commands::get::("{a}spend")).await.unwrap(), + Some(0.5) + ); + let stats = support::command_stats(&deployment).await; + assert!(support::calls(&stats, "evalsha") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "eval"), 0, "{stats}"); +} diff --git a/litellm-rust/crates/redis-client/tests/pubsub.rs b/litellm-rust/crates/redis-client/tests/pubsub.rs new file mode 100644 index 00000000000..2ad7cb38a5a --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pubsub.rs @@ -0,0 +1,189 @@ +mod support; + +use std::time::{Duration, Instant}; + +use futures_util::StreamExt; +use litellm_circuit_breaker::{Breaker, Config}; +use litellm_redis_client::commands; +use litellm_redis_client::{ + Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, + Target, Timeouts, +}; +use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually}; +use rstest::rstest; +use support::{connect, connect_guarded}; +use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + +const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation"; + +fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver> { + let (sender, receiver) = unbounded_channel(); + let mut messages = client.subscribe(channel); + tokio::spawn(async move { + while let Some(message) = messages.next().await { + if sender.send(message).is_err() { + return; + } + } + }); + receiver +} + +async fn delivered( + publisher: &RedisClient, + channel: &str, + received: &mut UnboundedReceiver>, + payload: &[u8], +) { + eventually(async || { + publisher + .send(commands::publish(channel, payload)) + .await + .ok()?; + let deadline = tokio::time::Instant::now() + Duration::from_millis(200); + while let Ok(Some(message)) = tokio::time::timeout_at(deadline, received.recv()).await { + if message == payload { + return Some(()); + } + } + None + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn a_subscriber_receives_raw_bytes_on_the_channel_name_without_the_namespace( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let publisher = connect(&server, "ns", DEFAULT_TIMEOUT); + let plain = connect(&server, "", DEFAULT_TIMEOUT); + let mut received = listen(&plain, CHANNEL); + + delivered( + &publisher, + CHANNEL, + &mut received, + b"{\"cache_key\":\"k\"}\xff", + ) + .await; +} + +#[rstest] +#[tokio::test] +async fn a_subscriber_resubscribes_after_redis_restarts( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let mut received = listen(&client, CHANNEL); + delivered(&client, CHANNEL, &mut received, b"before").await; + + server.restart().await.unwrap(); + + delivered(&client, CHANNEL, &mut received, b"after").await; +} + +#[rstest] +#[tokio::test] +async fn a_sentinel_subscriber_receives_from_the_master() { + let server = RedisSentinel::start().await.unwrap(); + let (host, port) = server.sentinel(); + let client = RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + let mut received = listen(&client, CHANNEL); + + delivered(&client, CHANNEL, &mut received, b"payload").await; +} + +async fn publish_calls(server: &RedisDeployment) -> Vec { + let replies = match server { + RedisDeployment::Standalone(server) => { + vec![server.cli(&["INFO", "commandstats"]).await.unwrap()] + } + RedisDeployment::Cluster(cluster) => cluster + .cli_outputs(&["INFO", "commandstats"]) + .await + .unwrap(), + }; + replies + .iter() + .flat_map(|reply| reply.lines()) + .filter(|line| line.starts_with("cmdstat_publish:")) + .map(str::to_owned) + .collect() +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_rejects_a_publish_without_sending_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = Breaker::new( + Config { + failure_threshold: 2, + recovery_timeout: Duration::from_secs(60), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ); + let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker)); + + server.freeze().await.unwrap(); + let opened = [client.ping().await, client.ping().await]; + let started = Instant::now(); + let rejected = client.send(commands::publish(CHANNEL, b"payload")).await; + let rejected_in = started.elapsed(); + server.thaw().await.unwrap(); + + assert!( + opened + .iter() + .all(|ping| matches!(ping, Err(Error::Transport(_)))), + "{opened:?}" + ); + assert!(matches!(rejected, Err(Error::BreakerOpen)), "{rejected:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + assert_eq!(publish_calls(&server).await, Vec::::new()); +} + +// redis-py's PubSub subscribes on a RESP2 connection, so it never sends HELLO +#[rstest] +#[tokio::test] +async fn a_subscription_speaks_resp2_like_redis_py( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let mut received = listen(&client, CHANNEL); + + delivered(&client, CHANNEL, &mut received, b"evict").await; + + let stats = support::command_stats(&server).await; + assert!(support::calls(&stats, "subscribe") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "hello"), 0, "{stats}"); +} diff --git a/litellm-rust/crates/redis-client/tests/script.rs b/litellm-rust/crates/redis-client/tests/script.rs new file mode 100644 index 00000000000..ea6d1c5ff78 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/script.rs @@ -0,0 +1,133 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; + +use litellm_redis_client::Script; +use litellm_redis_client::commands; +use litellm_redis_client::{DEFAULT_TIMEOUT, Error, Ttl}; +use support::connect; + +#[rstest] +#[tokio::test] +async fn script_keys_are_namespaced( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "a", DEFAULT_TIMEOUT); + let bare = connect(&server, "", DEFAULT_TIMEOUT); + let read = Script::new("return redis.call('GET', KEYS[1])"); + bare.send(commands::set("a:x", "1", Duration::from_secs(60))) + .await + .unwrap(); + + let value: Option = client + .send(commands::eval(&read, &["x"], Vec::::new())) + .await + .unwrap(); + + assert_eq!(value.as_deref(), Some("1")); +} + +#[rstest] +#[tokio::test] +async fn a_failing_script_is_a_server_error_and_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let failed = client + .send(commands::eval::(&incr, &["list"], 1)) + .await; + + assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); + assert_eq!( + client + .send(commands::eval::(&incr, &["counter"], 3)) + .await + .unwrap(), + 3 + ); +} + +#[rstest] +#[tokio::test] +async fn a_flushed_script_is_loaded_again( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::eval::(&incr, &["counter"], 1)) + .await + .unwrap(); + + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::eval::(&incr, &["counter"], 1)) + .await + .unwrap(), + 2 + ); +} + +const INCREMENT_WITH_FLOOR: &str = "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) \ +if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end \ +if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end \ +return count"; + +#[rstest] +#[tokio::test] +async fn pythons_increment_with_floor_script_floors_at_zero_and_expires_in_one_call( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(false, true)] pipelined: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + let floor = Script::new(INCREMENT_WITH_FLOOR); + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + + let count = match pipelined { + false => client + .send(commands::eval::(&floor, &["slots"], (-5, 60))) + .await + .unwrap(), + true => { + let results = client + .run(&client.pipeline().with(commands::eval::( + &floor, + &["slots"], + (-5, 60), + ))) + .await; + redis::from_redis_value(results[0].as_ref().unwrap().clone()).unwrap() + } + }; + + assert_eq!(count, 0); + assert_eq!( + client.send(commands::get::("slots")).await.unwrap(), + Some(0) + ); + assert!(matches!( + client.send(commands::ttl("slots")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); +} diff --git a/litellm-rust/crates/redis-client/tests/support/mod.rs b/litellm-rust/crates/redis-client/tests/support/mod.rs new file mode 100644 index 00000000000..9cf209dce4a --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/support/mod.rs @@ -0,0 +1,206 @@ +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::{pin::Pin, time::Duration}; + +use futures_util::Stream; +use litellm_circuit_breaker::Breaker; +use litellm_redis_client::{ + BasicAuth, Cluster, Credentials, Expiry, Namespace, Node, RedisClient, RedisConfig, RedisError, + RedisUrl, Standalone, StreamingCredentialsProvider, Target, Timeouts, +}; +use litellm_redis_testing::{ClientConnection, RedisDeployment, TcpProxy}; +use litellm_secret_value::SecretValue; +use tokio::sync::watch; + +pub const KEEP: Expiry = Expiry::KeepExisting(Duration::from_secs(60)); + +pub fn endpoint(deployment: &RedisDeployment) -> Target { + match deployment { + RedisDeployment::Standalone(server) => Target::Url(RedisUrl { + url: SecretValue::new(server.url()), + tls: None, + db: None, + username: None, + password: None, + }), + RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster { + nodes: cluster + .nodes() + .map(|(host, port)| Node { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + } +} + +pub fn connect(deployment: &RedisDeployment, namespace: &str, timeout: Duration) -> RedisClient { + connect_guarded(deployment, namespace, timeout, None) +} + +pub fn connect_guarded( + deployment: &RedisDeployment, + namespace: &str, + timeout: Duration, + breaker: Option, +) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(deployment), + credentials: Credentials::NONE, + namespace: Namespace::new(namespace), + timeouts: Timeouts::new(timeout, None), + breaker, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +pub fn connect_as(deployment: &RedisDeployment, credentials: Credentials) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(deployment), + credentials, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(2), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +pub async fn connections(deployment: &RedisDeployment) -> Vec { + deployment.clients().await.unwrap() +} + +/// A `TcpProxy` in front of every node, and a target that reaches the nodes only through them. A +/// cluster announces the proxy ports, so its slot map and redirects lead there too. +pub struct Proxied { + proxies: Vec, + pub target: Target, +} + +impl Proxied { + pub fn blackhole_open_connections(&self) { + self.proxies + .iter() + .for_each(TcpProxy::blackhole_open_connections); + } +} + +pub async fn proxied(deployment: &RedisDeployment) -> Proxied { + match deployment { + RedisDeployment::Standalone(server) => { + let proxy = TcpProxy::start(server.port()).await.unwrap(); + let target = Target::Standalone(Standalone { + host: "127.0.0.1".to_owned(), + port: proxy.port(), + tls: None, + db: 0, + }); + Proxied { + proxies: vec![proxy], + target, + } + } + RedisDeployment::Cluster(cluster) => { + let mut proxies = Vec::new(); + for (_, port) in cluster.nodes() { + proxies.push(TcpProxy::start(port).await.unwrap()); + } + let ports: [u16; 3] = std::array::from_fn(|node| proxies[node].port()); + cluster.announce_ports(ports).await.unwrap(); + let target = Target::Cluster(Cluster { + nodes: ports + .iter() + .map(|&port| Node { + host: "127.0.0.1".to_owned(), + port, + }) + .collect(), + tls: None, + }); + Proxied { proxies, target } + } + } +} + +/// `INFO

` from every node, concatenated. +async fn info(deployment: &RedisDeployment, section: &str) -> String { + match deployment { + RedisDeployment::Standalone(server) => server.cli(&["INFO", section]).await.unwrap(), + RedisDeployment::Cluster(cluster) => { + let mut stats = String::new(); + for (_, port) in cluster.nodes() { + stats.push_str(&cluster.cli_on_node(port, &["INFO", section]).await.unwrap()); + } + stats + } + } +} + +/// `INFO commandstats` from every node, so a test can count what was actually sent. +pub async fn command_stats(deployment: &RedisDeployment) -> String { + info(deployment, "commandstats").await +} + +/// Keys looked up on every node since the last `CONFIG RESETSTAT`, found or not. +pub async fn keys_read(deployment: &RedisDeployment) -> u64 { + let stats = info(deployment, "stats").await; + ["keyspace_hits:", "keyspace_misses:"] + .iter() + .flat_map(|field| { + stats + .lines() + .filter_map(move |line| line.trim().strip_prefix(field)?.parse::().ok()) + }) + .sum() +} + +/// How many times `command` reached a server, summed across the nodes in `stats`. +pub fn calls(stats: &str, command: &str) -> u64 { + let prefix = format!("cmdstat_{command}:calls="); + stats + .lines() + .filter_map(|line| line.trim().strip_prefix(&prefix)) + .filter_map(|rest| rest.split(',').next()?.parse::().ok()) + .sum() +} + +pub struct Rotating(watch::Sender); + +impl Rotating { + pub fn new(username: &str, password: &str) -> Self { + Self(watch::Sender::new(BasicAuth::new( + username.to_owned(), + password.to_owned(), + ))) + } + + pub fn rotate(&self, username: &str, password: &str) { + self.0 + .send_replace(BasicAuth::new(username.to_owned(), password.to_owned())); + } +} + +impl StreamingCredentialsProvider for Rotating { + fn subscribe( + &self, + ) -> Pin> + Send + 'static>> { + let receiver = self.0.subscribe(); + Box::pin(futures_util::stream::unfold( + (receiver, true), + |(mut receiver, first)| async move { + if !first { + receiver.changed().await.ok()?; + } + let credentials = receiver.borrow_and_update().clone(); + Some((Ok(credentials), (receiver, false))) + }, + )) + } +} diff --git a/litellm-rust/crates/redis-connect/AGENTS.md b/litellm-rust/crates/redis-connect/AGENTS.md new file mode 100644 index 00000000000..8e9d1302eaf --- /dev/null +++ b/litellm-rust/crates/redis-connect/AGENTS.md @@ -0,0 +1,11 @@ +- Turns a `litellm_redis_types::Resolution` into a `RedisClient`, the same way for every host. It depends on `litellm-redis-types`, never on `litellm-config`, so a host that resolves Redis some other way connects the same way. `redis-client` stays cloud-agnostic, so credential providers and cloud token sources are built here +- The host passes only what it alone can build in `HostServices`: the event sink, the HTTP client and the environment lookup. The breaker config comes from `Resolution::breaker`, which `config` reads from `REDIS_CIRCUIT_BREAKER_*` with Python's fallback, so a bad value keeps the default instead of failing connect. Each cloud source asks its auth crate for one explicit identity, the way the Python SDK builds it: `AwsAuthService::default_chain_credentials` (botocore's default chain), `CredentialSource::application_default` (`google.auth.default`) and `AzureIdentity` (the class `_build_azure_credential` picks). LiteLLM's provider settings such as `AWS_ROLE_NAME`, `AZURE_AD_TOKEN` or `VERTEXAI_CREDENTIALS` belong to Bedrock, Azure OpenAI and Vertex and never reach those constructors. `Auth::Provider` (a handle for a Python `credential_provider` object) fails at connect with an error naming the key, since no host can supply one yet +- `RedisHosts` is the one place a host shares Redis state. Breakers are keyed by `Endpoint`, so every caller on a server sees one health signal whatever its db or credentials. A client takes only breaker settings and admits through its server's breaker, so no caller can pair a client with a breaker of its own, and clients are keyed by `Resolution` alone. Both are held weakly, so a breaker or client is freed with its last user, and the lock is held while a client connects so each is built once. A caller holding the Python GIL releases it first. Hosts own one instance (python-bridge per process, the gateway per `connect`) and never keep a registry of their own +- Every expiring credential goes through the one refresh loop in `refresh.rs`. A cloud source only fetches. Never add a second loop or let a source sleep. The expiry deadline lives only in the published `Latest::Ready`, so expiring it leaves `Ready` and cannot fire again. Never keep a copy of the deadline in the loop +- AWS SDK identity settings must be in the process environment. `aws-config` cannot take the host lookup, so config-only overrides for profiles, role/web-identity settings, credentials, files and metadata sources fail before connecting rather than silently using another identity. Equal process/config values are allowed. Bedrock-only settings remain unrelated to Redis +- redis-rs subscribes once for the connection handshake and again for re-authentication, so a new connection sends a redundant initial AUTH. Unchanged refreshes do not trigger another AUTH. Do not skip a provider's first or alternating subscriptions, because concurrent connections and retries each need the current token. Removing the initial duplicate requires an upstream connection-level fix +- The loop starts on the first `subscribe()`, which redis-rs calls inside its runtime, so `connect` is safe off a runtime +- A token-only credential (Azure without `REDIS_USERNAME`, GCP) is sent as `AUTH default `, because redis-rs always sends a username with a credentials provider. On Redis 6+ that is the same as Python's `AUTH ` +- Tests for the refresh loop use a fake source and real Redis ACL users. Cloud sources are proven against the live service, unit tests only pin request shapes we own +- Rotation tests must retain the same server-side connection IDs across re-AUTH, then force a reconnect after removing the old user and verify the new identity still works +- Live cloud tests live in `tests/live/`. Set `LITELLM_REDIS_AWS_LIVE_CONFIG`, `LITELLM_REDIS_AZURE_LIVE_CONFIG` or `LITELLM_REDIS_GCP_LIVE_CONFIG` to a JSON file containing `redis` settings and `hold_seconds`. Set the waiting period from the deployment's token lifetime when checking expiry. Optional fixtures return early when unset, and configured tests run under plain `cargo test`. An unconfigured run never counts as live proof diff --git a/litellm-rust/crates/redis-connect/Cargo.toml b/litellm-rust/crates/redis-connect/Cargo.toml new file mode 100644 index 00000000000..29b3293b856 --- /dev/null +++ b/litellm-rust/crates/redis-connect/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "litellm-redis-connect" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-aws.workspace = true +litellm-auth-azure.workspace = true +litellm-auth-gcp.workspace = true +litellm-auth-types.workspace = true +litellm-error.workspace = true +litellm-secret-value.workspace = true +strum.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-types.workspace = true +litellm-core-utils.workspace = true +litellm-http.workspace = true +litellm-redis-client.workspace = true +futures-util.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +veil.workspace = true +time.workspace = true +url.workspace = true +tokio = { workspace = true, features = ["sync", "time", "macros"] } +tracing.workspace = true + +[dev-dependencies] +litellm-redis-resolve.workspace = true +litellm-http = { workspace = true, features = ["test-support"] } +litellm-redis-testing.workspace = true +rustls.workspace = true +rstest.workspace = true +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] } diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs new file mode 100644 index 00000000000..d11f6ba47ff --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -0,0 +1,152 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_circuit_breaker::{Breaker, Config as BreakerConfig}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{ + Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Namespace, RedisClient, RedisConfig, + RedisEvents, Target, Timeouts, +}; +use litellm_redis_types::{Auth, HealthCheck, Resolution, StaticAuth}; + +use crate::{ + CredentialSource, Error, RefreshingCredentials, + sources::{AzureAd, ElastiCacheIam, GcpIam}, +}; + +pub type SharedEnvironment = Arc; + +pub struct HostServices { + pub events: Option>, + pub http: litellm_http::Client, + pub environment: SharedEnvironment, + pub credential_provider: Option>, +} + +// litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only +const CLUSTER_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(25); + +/// A breaker whose transitions reach `events`, for every client that talks to one Redis. +pub(crate) fn breaker(config: BreakerConfig, events: Option>) -> Breaker { + Breaker::new( + config, + Box::new(move |transition| { + if let Some(events) = &events { + events.breaker(transition); + } + }), + ) +} + +pub fn connect( + resolved: Resolution, + host: HostServices, + breaker: Option, +) -> Result { + tracing::info!(target = ?resolved.target, namespace = ?resolved.namespace, "redis client configured"); + let credentials = credentials(resolved.auth, &host)?; + let events = host.events; + let health_check = health_check(&resolved.target, resolved.health_check); + Ok(RedisClient::new(RedisConfig { + target: resolved.target, + credentials, + namespace: resolved.namespace.map(Namespace::new).unwrap_or_default(), + timeouts: Timeouts::new( + resolved.timeout.unwrap_or(DEFAULT_TIMEOUT), + resolved.connect_timeout, + ), + breaker, + events, + health_check, + keepalive: resolved.keepalive, + })?) +} + +fn health_check(target: &Target, configured: HealthCheck) -> Option { + match (target, configured) { + (_, HealthCheck::Every(interval)) => Some(interval), + (_, HealthCheck::Disabled) => None, + (Target::Cluster(_), HealthCheck::Unset) => Some(CLUSTER_HEALTH_CHECK_INTERVAL), + (_, HealthCheck::Unset) => None, + } +} + +fn credentials(auth: Auth, host: &HostServices) -> Result { + let refreshing = |provider: RefreshingCredentials| { + Credentials::Provider(CredentialsProvider(Arc::new(provider))) + }; + Ok(match auth { + Auth::Static(StaticAuth { username, password }) => { + Credentials::Static { username, password } + } + Auth::Provider(_) => match &host.credential_provider { + Some(source) => refreshing(RefreshingCredentials::new(Arc::clone(source))), + None => return Err(Error::NoHostCredentialProvider), + }, + Auth::AzureAd(auth) => refreshing(RefreshingCredentials::new(AzureAd::new( + auth, + host.environment.clone(), + ))), + Auth::AwsIam(auth) => refreshing(RefreshingCredentials::new(ElastiCacheIam::new( + auth, + &*host.environment, + )?)), + Auth::GcpIam { service_account } => refreshing(RefreshingCredentials::new(GcpIam::new( + host.http.clone(), + service_account, + host.environment.clone(), + ))), + }) +} + +#[cfg(test)] +mod tests { + use litellm_redis_client::{Cluster, Node, Standalone}; + use rstest::rstest; + + use super::*; + + fn cluster() -> Target { + Target::Cluster(Cluster { + nodes: vec![Node { + host: "node".to_owned(), + port: 7000, + }], + tls: None, + }) + } + + fn standalone() -> Target { + Target::Standalone(Standalone { + host: "redis".to_owned(), + port: 6379, + tls: None, + db: 0, + }) + } + + #[rstest] + #[case::an_unset_cluster_takes_the_python_default( + cluster(), + HealthCheck::Unset, + Some(CLUSTER_HEALTH_CHECK_INTERVAL) + )] + #[case::zero_keeps_a_cluster_off(cluster(), HealthCheck::Disabled, None)] + #[case::an_unset_standalone_stays_off(standalone(), HealthCheck::Unset, None)] + #[case::a_configured_interval_wins( + cluster(), + HealthCheck::Every(Duration::from_secs(5)), + Some(Duration::from_secs(5)) + )] + #[case::a_negative_interval_checks_every_command( + standalone(), + HealthCheck::Every(Duration::ZERO), + Some(Duration::ZERO) + )] + fn health_checks_follow_redis_py( + #[case] target: Target, + #[case] configured: HealthCheck, + #[case] expected: Option, + ) { + assert_eq!(health_check(&target, configured), expected); + } +} diff --git a/litellm-rust/crates/redis-connect/src/error.rs b/litellm-rust/crates/redis-connect/src/error.rs new file mode 100644 index 00000000000..98343a03fb8 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/error.rs @@ -0,0 +1,102 @@ +#[derive(veil::Redact, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Connect(#[from] litellm_redis_client::ConnectError), + #[error( + "Redis setting credential_provider names a Python object, which this host cannot supply" + )] + NoHostCredentialProvider, + #[error( + "Redis credential setting {setting} must be set in the process environment, not only environment_variables" + )] + ProcessEnvironmentRequired { setting: &'static str }, + #[error("{stage} failed{}", http_status_suffix(*http_status))] + Credential { + stage: CredentialStage, + http_status: Option, + #[source] + #[redact(with = "[REDACTED]")] + cause: litellm_error::Source, + }, + #[error("credential acquisition exceeded {timeout:?}")] + CredentialTimeout { + timeout: std::time::Duration, + #[source] + cause: tokio::time::error::Elapsed, + }, + #[error("Redis credential has expired")] + ExpiredCredential, + #[error("no {provider} credential is available for Redis")] + MissingCredential { provider: &'static str }, +} + +#[derive(Clone, Copy, Debug, strum::Display, strum::VariantArray)] +pub enum CredentialStage { + #[strum(serialize = "caller identity")] + Identity, + #[strum(serialize = "Redis token signing")] + Signing, + #[strum(serialize = "Redis token exchange")] + Exchange, + #[strum(serialize = "Redis token response")] + Response, + #[strum(serialize = "Redis credential_provider")] + Provider, +} + +impl Error { + pub fn credential( + stage: CredentialStage, + ) -> impl FnOnce(E) -> Self { + move |cause| Self::Credential { + stage, + http_status: None, + cause: litellm_error::Source::new(cause), + } + } + + pub(crate) fn credential_response( + stage: CredentialStage, + status: u16, + ) -> impl FnOnce(litellm_http::Error) -> Self { + move |cause| Self::Credential { + stage, + http_status: Some(status), + cause: litellm_error::Source::new(cause), + } + } + + pub(crate) fn upstream_status(&self) -> Option { + match self { + Self::Credential { http_status, .. } => *http_status, + Self::Connect(_) + | Self::NoHostCredentialProvider + | Self::ProcessEnvironmentRequired { .. } + | Self::CredentialTimeout { .. } + | Self::ExpiredCredential + | Self::MissingCredential { .. } => None, + } + } + + pub(crate) fn kind(&self) -> &'static str { + match self { + Self::Connect(_) => "connection", + Self::NoHostCredentialProvider => "host-credential-provider", + Self::ProcessEnvironmentRequired { .. } => "process-environment-required", + Self::Credential { stage, .. } => match stage { + CredentialStage::Identity => "caller-identity", + CredentialStage::Signing => "token-signing", + CredentialStage::Exchange => "token-exchange", + CredentialStage::Response => "token-response", + CredentialStage::Provider => "credential-provider", + }, + Self::CredentialTimeout { .. } => "acquisition-timeout", + Self::ExpiredCredential => "expired-credential", + Self::MissingCredential { .. } => "missing-credential", + } + } +} + +fn http_status_suffix(status: Option) -> String { + status.map_or_else(String::new, |status| format!(" (HTTP {status})")) +} diff --git a/litellm-rust/crates/redis-connect/src/events.rs b/litellm-rust/crates/redis-connect/src/events.rs new file mode 100644 index 00000000000..a556b7d370b --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/events.rs @@ -0,0 +1,129 @@ +use std::{ + sync::Mutex, + time::{Duration, Instant}, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{CallEvent, FailureKind, Outcome, RedisEvents}; + +const DEFAULT_TIMEOUT_LOG_INTERVAL: Duration = Duration::from_secs(5); + +pub struct TracingEvents { + throttle: TimeoutLogThrottle, +} + +impl TracingEvents { + pub fn new(lookup: &dyn Lookup) -> Self { + let interval = lookup + .get("REDIS_TIMEOUT_LOG_INTERVAL") + .and_then(|seconds| seconds.parse::().ok()) + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .unwrap_or(DEFAULT_TIMEOUT_LOG_INTERVAL); + Self { + throttle: TimeoutLogThrottle::new(interval), + } + } +} + +impl RedisEvents for TracingEvents { + fn call(&self, event: &CallEvent) { + let Outcome::Failed(kind) = event.outcome else { + return; + }; + match kind { + FailureKind::BreakerOpen => { + tracing::debug!(operation = ?event.operation, "redis call rejected by the open circuit breaker") + } + FailureKind::Timeout => match self.throttle.admit(Instant::now()) { + None => { + tracing::debug!(operation = ?event.operation, duration = ?event.duration, "redis call timed out") + } + Some(0) => { + tracing::error!(operation = ?event.operation, duration = ?event.duration, "redis call timed out") + } + Some(suppressed) => tracing::error!( + operation = ?event.operation, + duration = ?event.duration, + suppressed, + "redis call timed out, more redis timeouts since the previous line were logged at debug" + ), + }, + FailureKind::Transport | FailureKind::Server | FailureKind::Decode => { + tracing::error!(operation = ?event.operation, failure = ?kind, "redis call failed") + } + } + } + + fn breaker(&self, transition: &Transition) { + match transition.to { + State::Open { .. } => tracing::warn!( + generation = transition.generation, + "redis circuit breaker opened, fast-failing redis calls" + ), + State::HalfOpen => tracing::info!( + generation = transition.generation, + "redis circuit breaker half open, probing redis" + ), + State::Closed { .. } => tracing::info!( + generation = transition.generation, + "redis circuit breaker closed, redis recovered" + ), + } + } +} + +struct TimeoutLogThrottle { + interval: Duration, + state: Mutex<(Option, u64)>, +} + +impl TimeoutLogThrottle { + fn new(interval: Duration) -> Self { + Self { + interval, + state: Mutex::new((None, 0)), + } + } + + fn admit(&self, now: Instant) -> Option { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let (last_logged_at, suppressed) = &mut *state; + if last_logged_at.is_some_and(|last| now.duration_since(last) < self.interval) { + *suppressed += 1; + return None; + } + *last_logged_at = Some(now); + Some(std::mem::take(suppressed)) + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use rstest::rstest; + + use super::TimeoutLogThrottle; + + #[rstest] + fn admits_one_line_per_interval_and_counts_the_rest() { + let throttle = TimeoutLogThrottle::new(Duration::from_secs(5)); + let start = Instant::now(); + let at = |seconds: u64| start + Duration::from_secs(seconds); + + let admitted = [ + throttle.admit(at(0)), + throttle.admit(at(1)), + throttle.admit(at(4)), + throttle.admit(at(5)), + throttle.admit(at(6)), + throttle.admit(at(20)), + ]; + + assert_eq!(admitted, [Some(0), None, None, Some(2), None, Some(1)]); + } +} diff --git a/litellm-rust/crates/redis-connect/src/hosts.rs b/litellm-rust/crates/redis-connect/src/hosts.rs new file mode 100644 index 00000000000..11e69d45aff --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/hosts.rs @@ -0,0 +1,241 @@ +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, WeakBreaker}; +use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient}; +use litellm_redis_types::{Endpoint, Resolution}; + +use crate::{Error, HostServices, breaker, connect}; + +/// The Redis servers one host talks to. Every caller on a server shares its breaker, and callers with the same +/// settings share one client, which admits through its server's breaker. Both are held weakly, so each is freed +/// with its last user and the next caller builds a fresh one. +#[derive(Default)] +pub struct RedisHosts { + breakers: Mutex, WeakBreaker)>>, + clients: Mutex>, +} + +#[derive(PartialEq)] +struct ClientKey { + resolved: Resolution, + guarded: bool, +} + +impl RedisHosts { + /// The breaker of the server `endpoint` names, built from `config` by its first caller. Without an endpoint + /// the breaker is this caller's own, which `breakers` still lists. + pub fn breaker( + &self, + endpoint: Option, + config: BreakerConfig, + events: Option>, + ) -> Breaker { + let mut breakers = live(&self.breakers, WeakBreaker::upgrade); + if let Some(shared) = breakers + .iter() + .filter(|(key, _)| key.is_some() && *key == endpoint) + .find_map(|(_, weak)| weak.upgrade()) + { + return shared; + } + let built = breaker(config, events); + breakers.push((endpoint, built.downgrade())); + built + } + + /// The client for `resolved`, connected by its first caller and guarded by its server's breaker unless `breaker` + /// is none. The lock is held while it connects, so concurrent callers with the same settings connect once. + pub fn client( + &self, + resolved: Resolution, + host: HostServices, + breaker: Option, + ) -> Result { + let breaker = breaker + .map(|config| self.breaker(resolved.target.endpoint(), config, host.events.clone())); + let key = ClientKey { + resolved, + guarded: breaker.is_some(), + }; + let mut clients = live(&self.clients, WeakRedisClient::upgrade); + if let Some(shared) = clients + .iter() + .filter(|(entry, _)| *entry == key) + .find_map(|(_, weak)| weak.upgrade()) + { + return Ok(shared); + } + let client = connect(key.resolved.clone(), host, breaker)?; + clients.push((key, client.downgrade())); + Ok(client) + } + + /// Every breaker a caller still holds, for a state gauge. + pub fn breakers(&self) -> Vec { + live(&self.breakers, WeakBreaker::upgrade) + .iter() + .filter_map(|(_, weak)| weak.upgrade()) + .collect() + } +} + +fn live( + entries: &Mutex>, + upgrade: fn(&W) -> Option, +) -> MutexGuard<'_, Vec<(K, W)>> { + let mut entries = entries.lock().unwrap_or_else(PoisonError::into_inner); + entries.retain(|(_, weak)| upgrade(weak).is_some()); + entries +} + +#[cfg(test)] +mod tests { + use std::{ + sync::Arc, + time::{Duration, Instant}, + }; + + use futures_util::FutureExt; + use litellm_circuit_breaker::{Config, Failure, State}; + use litellm_redis_types::{ + Auth, Endpoint, HealthCheck, Resolution, Standalone, StaticAuth, Target, + }; + use rstest::rstest; + + use super::RedisHosts; + use crate::HostServices; + + const CONFIG: Config = Config { + failure_threshold: 1, + recovery_timeout: Duration::from_secs(60), + timeout_min_duration: Duration::ZERO, + }; + + fn target(host: &str, db: i64) -> Target { + Target::Standalone(Standalone { + host: host.to_owned(), + port: 6379, + tls: None, + db, + }) + } + + fn endpoint(host: &str, db: i64) -> Option { + target(host, db).endpoint() + } + + fn resolved(host: &str, namespace: Option<&str>) -> Resolution { + Resolution { + target: target(host, 0), + auth: Auth::Static(StaticAuth::default()), + namespace: namespace.map(str::to_owned), + timeout: None, + connect_timeout: None, + health_check: HealthCheck::Unset, + keepalive: true, + ignored: Vec::new(), + } + } + + fn services() -> HostServices { + HostServices { + events: None, + http: litellm_http::Client::for_test(), + environment: Arc::new(|_: &str| None), + credential_provider: None, + } + } + + fn open(breaker: &litellm_circuit_breaker::Breaker) { + breaker.failed(Failure::Hard, Instant::now()); + } + + #[rstest] + fn callers_on_one_server_share_its_breaker_whatever_their_db() { + let hosts = RedisHosts::default(); + let first = hosts.breaker(endpoint("r1", 0), CONFIG, None); + let second = hosts.breaker(endpoint("r1", 1), CONFIG, None); + let other = hosts.breaker(endpoint("r2", 0), CONFIG, None); + + open(&first); + + assert!(matches!(second.state(), State::Open { .. })); + assert!(matches!(other.state(), State::Closed { .. })); + } + + #[rstest] + fn a_breaker_without_an_endpoint_is_listed_but_never_shared() { + let hosts = RedisHosts::default(); + let own = hosts.breaker(None, CONFIG, None); + let shared = hosts.breaker(endpoint("r1", 0), CONFIG, None); + + open(&own); + + assert!(matches!(shared.state(), State::Closed { .. })); + assert_eq!(hosts.breakers().len(), 2); + } + + #[rstest] + fn a_server_whose_callers_are_all_gone_gets_a_fresh_breaker() { + let hosts = RedisHosts::default(); + let first = hosts.breaker(endpoint("r1", 0), CONFIG, None); + open(&first); + drop(first); + + let again = hosts.breaker(endpoint("r1", 0), CONFIG, None); + + assert!(matches!(again.state(), State::Closed { .. })); + assert_eq!(hosts.breakers().len(), 1); + } + + #[rstest] + #[tokio::test] + async fn callers_with_the_same_settings_share_one_client() { + let hosts = RedisHosts::default(); + let first = hosts + .client(resolved("r1", Some("a")), services(), Some(CONFIG)) + .unwrap(); + + let same = hosts + .client(resolved("r1", Some("a")), services(), Some(CONFIG)) + .unwrap(); + let other_namespace = hosts + .client(resolved("r1", Some("b")), services(), Some(CONFIG)) + .unwrap(); + + assert!(first.same_client(&same)); + assert!(!first.same_client(&other_namespace)); + } + + #[rstest] + #[tokio::test] + async fn a_client_admits_through_its_servers_breaker() { + let hosts = RedisHosts::default(); + let breaker = hosts.breaker(endpoint("r1", 1), CONFIG, None); + let client = hosts + .client(resolved("r1", None), services(), Some(CONFIG)) + .unwrap(); + + open(&breaker); + + assert!(matches!( + client.ping().now_or_never(), + Some(Err(litellm_redis_client::Error::BreakerOpen)) + )); + } + + #[rstest] + #[tokio::test] + async fn a_client_without_a_breaker_is_never_shared_with_a_guarded_one() { + let hosts = RedisHosts::default(); + let guarded = hosts + .client(resolved("r1", None), services(), Some(CONFIG)) + .unwrap(); + + let unguarded = hosts + .client(resolved("r1", None), services(), None) + .unwrap(); + + assert!(!guarded.same_client(&unguarded)); + } +} diff --git a/litellm-rust/crates/redis-connect/src/lib.rs b/litellm-rust/crates/redis-connect/src/lib.rs new file mode 100644 index 00000000000..7e37f721679 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/lib.rs @@ -0,0 +1,14 @@ +mod connect; +mod error; +mod events; +mod hosts; +mod refresh; +mod sources; + +pub(crate) use connect::breaker; +pub use connect::{HostServices, SharedEnvironment, connect}; +pub use error::{CredentialStage, Error}; +pub use events::TracingEvents; +pub use hosts::RedisHosts; +pub(crate) use refresh::RefreshingCredentials; +pub use refresh::{CredentialSource, Expiring, SourceFuture}; diff --git a/litellm-rust/crates/redis-connect/src/refresh.rs b/litellm-rust/crates/redis-connect/src/refresh.rs new file mode 100644 index 00000000000..1e05d3d6289 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/refresh.rs @@ -0,0 +1,783 @@ +use std::{ + future::{Future, pending}, + pin::Pin, + sync::{Arc, Mutex}, + time::{Duration, SystemTime}, +}; + +use futures_util::{Stream, stream}; +use litellm_redis_client::{BasicAuth, RedisError, StreamingCredentialsProvider}; +use litellm_secret_value::SecretValue; +use tokio::{sync::watch, time::Instant}; + +use crate::Error; + +const TOKEN_ONLY_USERNAME: &str = "default"; +const UNKNOWN_LIFETIME: Duration = Duration::from_secs(300); +const MIN_REFRESH: Duration = Duration::from_secs(1); +const ACQUISITION_TIMEOUT: Duration = Duration::from_secs(30); +const FIRST_RETRY: Duration = Duration::from_secs(1); +const MAX_RETRY: Duration = Duration::from_secs(60); + +#[derive(Clone, PartialEq, Eq)] +pub struct Expiring { + pub username: Option, + pub password: SecretValue, + pub expires_at: Option, +} + +pub type SourceFuture<'a> = Pin> + Send + 'a>>; + +pub trait CredentialSource: Send + Sync + 'static { + fn fetch(&self) -> SourceFuture<'_>; +} + +impl CredentialSource for Arc { + fn fetch(&self) -> SourceFuture<'_> { + (**self).fetch() + } +} + +#[derive(Clone)] +enum Latest { + Pending, + Ready { + username: String, + password: SecretValue, + expires_at: Option, + }, + Unavailable(Arc), +} + +impl Latest { + fn is_same_token(&self, other: &Self) -> bool { + matches!( + (self, other), + ( + Self::Ready { username, password, .. }, + Self::Ready { username: other_username, password: other_password, .. }, + ) if username == other_username && password == other_password + ) + } + + fn expires_at(&self) -> Option { + match self { + Self::Ready { expires_at, .. } => *expires_at, + Self::Pending | Self::Unavailable(_) => None, + } + } +} + +type Start = Box; + +#[derive(Clone)] +pub struct RefreshingCredentials { + latest: watch::Receiver, + start: Arc>>, +} + +impl RefreshingCredentials { + pub fn new(source: impl CredentialSource) -> Self { + let (sender, latest) = watch::channel(Latest::Pending); + let start: Start = Box::new(move || drop(tokio::spawn(refresh(source, sender)))); + Self { + latest, + start: Arc::new(Mutex::new(Some(start))), + } + } + + fn start_once(&self) { + let start = self + .start + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(start) = start { + start(); + } + } +} + +impl StreamingCredentialsProvider for RefreshingCredentials { + fn subscribe( + &self, + ) -> Pin> + Send + 'static>> { + self.start_once(); + Box::pin(stream::unfold( + (self.latest.clone(), true), + |(mut receiver, first)| async move { + if !first { + receiver.changed().await.ok()?; + } + let latest = receiver + .wait_for(|latest| !matches!(latest, Latest::Pending)) + .await + .ok()? + .clone(); + let item = match latest { + Latest::Ready { + username, password, .. + } => Ok(BasicAuth::new(username, password.expose().to_owned())), + Latest::Unavailable(error) => Err(public_error(&error)), + Latest::Pending => unreachable!("wait_for skips Pending"), + }; + Some((item, (receiver, false))) + }, + )) + } +} + +fn public_error(error: &Error) -> RedisError { + RedisError::from(( + litellm_redis_client::RedisErrorKind::AuthenticationFailed, + "fetching Redis credentials failed", + error.to_string(), + )) +} + +struct Fresh { + latest: Latest, + refresh_in: Duration, +} + +fn fresh(credentials: Expiring, now: SystemTime) -> Result { + let remaining = credentials + .expires_at + .map(|expires_at| expires_at.duration_since(now).unwrap_or_default()); + if remaining == Some(Duration::ZERO) { + return Err(Error::ExpiredCredential); + } + Ok(Fresh { + latest: Latest::Ready { + username: credentials + .username + .unwrap_or_else(|| TOKEN_ONLY_USERNAME.to_owned()), + password: credentials.password, + expires_at: remaining.map(|remaining| Instant::now() + remaining), + }, + refresh_in: refresh_after(credentials.expires_at, now), + }) +} + +async fn until_expiry( + future: impl Future, + sender: &watch::Sender, +) -> Option { + tokio::pin!(future); + loop { + let expires_at = sender.borrow().expires_at(); + let expiration = async { + match expires_at { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => pending().await, + } + }; + tokio::select! { + biased; + () = sender.closed() => return None, + () = expiration => { + sender.send_replace(Latest::Unavailable(Arc::new(Error::ExpiredCredential))); + } + result = &mut future => return Some(result), + } + } +} + +async fn refresh(source: impl CredentialSource, sender: watch::Sender) { + let mut retry = FIRST_RETRY; + loop { + let acquisition = async { + tokio::time::timeout(ACQUISITION_TIMEOUT, source.fetch()) + .await + .map_err(|cause| Error::CredentialTimeout { + timeout: ACQUISITION_TIMEOUT, + cause, + })? + .and_then(|credentials| fresh(credentials, SystemTime::now())) + }; + let Some(result) = until_expiry(acquisition, &sender).await else { + return; + }; + let wait = match result { + Ok(Fresh { latest, refresh_in }) => { + retry = FIRST_RETRY; + sender.send_if_modified(|current| { + let changed = !current.is_same_token(&latest); + *current = latest; + changed + }); + refresh_in + } + Err(error) => { + tracing::warn!(failure = error.kind(), upstream_status = error.upstream_status(), retry_in = ?retry, "fetching Redis credentials failed"); + if !matches!(*sender.borrow(), Latest::Ready { .. }) { + sender.send_replace(Latest::Unavailable(Arc::new(error))); + } + let wait = retry; + retry = (retry * 2).min(MAX_RETRY); + wait + } + }; + if until_expiry(tokio::time::sleep(wait), &sender) + .await + .is_none() + { + return; + } + } +} + +fn refresh_after(expires_at: Option, now: SystemTime) -> Duration { + let remaining = expires_at + .map(|expires_at| expires_at.duration_since(now).unwrap_or_default()) + .unwrap_or(UNKNOWN_LIFETIME); + (remaining * 3 / 4).max(MIN_REFRESH) +} + +#[cfg(test)] +mod tests { + use litellm_redis_client::commands; + use std::{ + collections::BTreeSet, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::{Duration, SystemTime}, + }; + + use futures_util::StreamExt; + use litellm_redis_client::StreamingCredentialsProvider; + use litellm_secret_value::SecretValue; + use rstest::rstest; + + use super::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture, refresh_after}; + use litellm_redis_client::{ + Cluster, Credentials, CredentialsProvider, Namespace, Node, RedisClient, RedisConfig, + RedisUrl, Sentinel, Target, Timeouts, + }; + use litellm_redis_testing::{ + ClientConnection, RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually, + }; + + #[rstest] + #[case::three_quarters_of_the_lifetime( + Some(Duration::from_secs(900)), + Duration::from_secs(675) + )] + #[case::already_expired(Some(Duration::ZERO), Duration::from_secs(1))] + #[case::unknown_expiry(None, Duration::from_secs(225))] + fn refreshes_well_before_expiry( + #[case] lifetime: Option, + #[case] expected: Duration, + ) { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000); + + assert_eq!( + refresh_after(lifetime.map(|lifetime| now + lifetime), now), + expected + ); + } + + struct Scripted { + replies: Mutex>>, + calls: Arc, + } + + impl CredentialSource for Scripted { + fn fetch(&self) -> SourceFuture<'_> { + self.calls.fetch_add(1, Ordering::SeqCst); + let reply = self.replies.lock().unwrap().remove(0); + Box::pin(async move { + reply.map_err(|reason| { + crate::Error::credential(crate::error::CredentialStage::Identity)( + std::io::Error::other(reason), + ) + }) + }) + } + } + + fn token(password: &str, lifetime: Duration) -> Result { + Ok(Expiring { + username: None, + password: SecretValue::new(password), + expires_at: Some(SystemTime::now() + lifetime), + }) + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn pushes_each_refresh_and_keeps_the_last_token_through_failures() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + token("first", Duration::from_secs(8)), + Err("token endpoint down"), + token("second", Duration::from_secs(3600)), + ]), + calls: calls.clone(), + }); + let mut stream = provider.subscribe(); + + let first = stream.next().await.unwrap().unwrap(); + let second = stream.next().await.unwrap().unwrap(); + let late = provider.subscribe().next().await.unwrap().unwrap(); + + assert_eq!((first.username(), first.password()), ("default", "first")); + assert_eq!(second.password(), "second"); + assert_eq!(late.password(), "second"); + assert_eq!(calls.load(Ordering::SeqCst), 3); + } + + #[rstest] + fn building_outside_a_runtime_fetches_nothing_until_a_connection_subscribes() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![token("first", Duration::from_secs(3600))]), + calls: calls.clone(), + }); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let before = calls.load(Ordering::SeqCst); + let first = runtime.block_on(async { provider.subscribe().next().await.unwrap().unwrap() }); + + assert_eq!(before, 0); + assert_eq!(first.password(), "first"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn a_failure_before_any_token_fails_the_connection_instead_of_hanging() { + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + Err("no identity"), + token("late", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + + assert!(stream.next().await.unwrap().is_err()); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "late"); + } + + struct Blocking { + first: Mutex>, + started: Arc, + cancelled: Arc, + } + + struct Cancellation(Arc); + + impl Drop for Cancellation { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + impl CredentialSource for Blocking { + fn fetch(&self) -> SourceFuture<'_> { + let first = self.first.lock().unwrap().take(); + Box::pin(async move { + if let Some(first) = first { + return Ok(first); + } + let _guard = Cancellation(self.cancelled.clone()); + self.started.notify_one(); + std::future::pending().await + }) + } + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn a_refresh_that_returns_the_same_token_is_not_pushed_again() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + token("same", Duration::from_secs(8)), + token("same", Duration::from_secs(3600)), + ]), + calls: calls.clone(), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "same"); + + let pushed = tokio::time::timeout(Duration::from_secs(60), stream.next()).await; + + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert!(pushed.is_err(), "an unchanged token was pushed again"); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn expiration_is_published_while_refresh_is_stalled() { + let started = Arc::new(tokio::sync::Notify::new()); + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(Some(token("first", Duration::from_secs(4)).unwrap())), + started: started.clone(), + cancelled: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "first"); + let start = tokio::time::Instant::now(); + started.notified().await; + let expired = stream.next().await.unwrap().unwrap_err(); + assert!(start.elapsed() <= Duration::from_secs(5)); + assert!(expired.to_string().contains("expired")); + let fresh = provider.subscribe().next().await.unwrap().unwrap_err(); + assert!(fresh.to_string().contains("expired")); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn expiration_is_published_once_and_later_failures_keep_their_cause() { + let replies = std::iter::once(token("first", Duration::from_secs(8))) + .chain(std::iter::repeat_n(Err("token endpoint down"), 16)) + .collect(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(replies), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "first"); + + let failures: Vec = stream + .take(5) + .map(|item| item.unwrap_err().to_string()) + .collect() + .await; + + assert!(failures[0].contains("expired"), "{failures:?}"); + assert!( + failures[1..] + .iter() + .all(|failure| failure.contains("caller identity") && !failure.contains("expired")), + "{failures:?}" + ); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn initial_acquisition_has_a_bounded_wait() { + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(None), + started: Arc::new(tokio::sync::Notify::new()), + cancelled: Arc::new(AtomicUsize::new(0)), + }); + let error = provider.subscribe().next().await.unwrap().unwrap_err(); + assert!(error.to_string().contains("exceeded")); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn dropping_all_subscribers_cancels_inflight_acquisition() { + let started = Arc::new(tokio::sync::Notify::new()); + let cancelled = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(None), + started: started.clone(), + cancelled: cancelled.clone(), + }); + let stream = provider.subscribe(); + started.notified().await; + drop(stream); + drop(provider); + tokio::task::yield_now().await; + assert_eq!(cancelled.load(Ordering::SeqCst), 1); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn an_expired_fetch_is_rejected_and_a_later_valid_token_recovers() { + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + Ok(Expiring { + username: None, + password: SecretValue::new("expired"), + expires_at: Some(SystemTime::UNIX_EPOCH), + }), + token("recovered", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert!( + stream + .next() + .await + .unwrap() + .unwrap_err() + .to_string() + .contains("expired") + ); + assert_eq!( + stream.next().await.unwrap().unwrap().password(), + "recovered" + ); + } + + #[rstest] + fn credential_failure_keeps_its_typed_cause_but_hides_it_at_the_redis_boundary() { + let error = crate::Error::credential(crate::error::CredentialStage::Identity)( + std::io::Error::new(std::io::ErrorKind::PermissionDenied, "credential-value"), + ); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .kind(), + std::io::ErrorKind::PermissionDenied + ); + let debug = format!("{error:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("[REDACTED]")); + let public = super::public_error(&error).to_string(); + assert!(public.contains("caller identity")); + assert!(!public.contains("credential-value")); + } + + #[rstest] + fn credential_http_failures_expose_status_without_exposing_the_response_body() { + let error = crate::Error::credential_response(crate::CredentialStage::Exchange, 503)( + litellm_http::Error::upstream(503, "credential-value".into()), + ); + let public = super::public_error(&error).to_string(); + assert!(public.contains("HTTP 503")); + assert!(!public.contains("credential-value")); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .upstream_response() + .unwrap() + .0, + 503 + ); + } + + fn endpoint(deployment: &RedisDeployment) -> Target { + match deployment { + RedisDeployment::Standalone(server) => Target::Url(RedisUrl { + url: SecretValue::new(server.url()), + tls: None, + db: None, + username: None, + password: None, + }), + RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster { + nodes: cluster + .nodes() + .map(|(host, port)| Node { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + } + } + + fn user(username: &str, password: &str, lifetime: Duration) -> Result { + Ok(Expiring { + username: Some(username.to_owned()), + password: SecretValue::new(password), + expires_at: Some(SystemTime::now() + lifetime), + }) + } + + fn connection_users(connections: &[ClientConnection]) -> BTreeSet<&str> { + connections + .iter() + .map(|connection| connection.user.as_str()) + .collect() + } + + fn connection_ids(connections: &[ClientConnection]) -> BTreeSet<(u16, u64)> { + connections + .iter() + .map(|connection| (connection.node, connection.id)) + .collect() + } + + #[rstest] + #[tokio::test] + async fn a_refreshed_token_reauthenticates_live_connections_before_the_old_one_expires( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + user("alice", "first", Duration::from_secs(2)), + user("bob", "second", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))), + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = server.clients().await.unwrap(); + + let after = eventually(async || { + let connections = server.clients().await.ok()?; + (!connections.is_empty() && connection_users(&connections) == BTreeSet::from(["bob"])) + .then_some(connections) + }) + .await; + server + .cli_on_every_node(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + + assert!(!before.is_empty()); + assert_eq!(connection_users(&before), BTreeSet::from(["alice"])); + assert_eq!(connection_ids(&after), connection_ids(&before)); + assert_eq!( + client + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + server + .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let value = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + let reconnected = server.clients().await.unwrap(); + + assert_eq!(value, "1"); + assert!(!reconnected.is_empty()); + assert_eq!(connection_users(&reconnected), BTreeSet::from(["bob"])); + assert!(connection_ids(&reconnected).is_disjoint(&connection_ids(&after))); + } + + /// Cloud credentials rotate while Sentinel fails over: the refreshed user must be what the + /// client presents to the master Sentinel promotes, not the one it connected with first. + #[rstest] + #[tokio::test] + async fn a_refreshed_user_reaches_the_master_sentinel_promotes() { + let sentinel = RedisSentinel::start().await.unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + user("alice", "first", Duration::from_secs(2)), + user("bob", "second", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let (host, port) = sentinel.sentinel(); + let client = RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))), + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(1), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("before", "1", Duration::from_secs(60))) + .await + .unwrap(); + let master = sentinel.master_port().await.unwrap(); + let app_users = |connections: &[ClientConnection]| -> BTreeSet { + connections + .iter() + .filter(|connection| connection.user != "default") + .map(|connection| connection.user.clone()) + .collect() + }; + eventually(async || { + let connections = sentinel.clients_on(master).await.ok()?; + (app_users(&connections) == BTreeSet::from(["bob".to_owned()])).then_some(()) + }) + .await; + sentinel + .cli_on_data_nodes(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(60))) + .await + .ok() + }) + .await; + let on_promoted = sentinel.clients_on(promoted).await.unwrap(); + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!( + app_users(&on_promoted), + BTreeSet::from(["bob".to_owned()]), + "{on_promoted:?}" + ); + assert_eq!( + client + .send(commands::get::("after")) + .await + .unwrap() + .as_deref(), + Some("2") + ); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/aws.rs b/litellm-rust/crates/redis-connect/src/sources/aws.rs new file mode 100644 index 00000000000..1e87bdd1b12 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/aws.rs @@ -0,0 +1,188 @@ +use std::time::{Duration, SystemTime}; + +use litellm_auth_aws::{AwsAuthService, SigningScope, presign_url}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_types::AwsIamAuth; + +use crate::Error; +use crate::error::CredentialStage; +use litellm_secret_value::SecretValue; + +use crate::{CredentialSource, Expiring, SourceFuture}; + +const SERVICE: &str = "elasticache"; +const TOKEN_LIFETIME: Duration = Duration::from_secs(900); + +const SDK_IDENTITY_ENVIRONMENT: &[&str] = &[ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_SHARED_CREDENTIALS_FILE", + "AWS_CONFIG_FILE", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_ARN", + "AWS_ROLE_SESSION_NAME", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_AUTHORIZATION_TOKEN", + "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", + "AWS_EC2_METADATA_DISABLED", + "AWS_EC2_METADATA_SERVICE_ENDPOINT", + "AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE", + "AWS_ENDPOINT_URL", + "AWS_ENDPOINT_URL_STS", +]; + +pub(crate) struct ElastiCacheIam { + service: AwsAuthService, + auth: AwsIamAuth, +} + +impl ElastiCacheIam { + pub(crate) fn new(auth: AwsIamAuth, environment: &dyn Lookup) -> Result { + validate_environment(environment, &|name: &str| std::env::var(name).ok())?; + Ok(Self { + service: AwsAuthService::default(), + auth, + }) + } +} + +fn validate_environment(environment: &dyn Lookup, process: &dyn Lookup) -> Result<(), Error> { + match SDK_IDENTITY_ENVIRONMENT.iter().find(|name| { + environment.get(name).is_some_and(|value| { + Some(SecretValue::new(value)) != process.get(name).map(SecretValue::new) + }) + }) { + Some(setting) => Err(Error::ProcessEnvironmentRequired { setting }), + None => Ok(()), + } +} + +pub(crate) fn connect_url(auth: &AwsIamAuth) -> String { + let query = [ + Some(("Action", "connect")), + Some(("User", auth.user_name.as_str())), + auth.serverless + .then_some(("ResourceType", "ServerlessCache")), + ]; + let mut url = String::from("https://"); + url.push_str(&auth.cache_name.to_lowercase()); + url.push_str("/?"); + url.push_str( + &url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(query.into_iter().flatten()) + .finish(), + ); + url +} + +impl CredentialSource for ElastiCacheIam { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let credentials = self + .service + .default_chain_credentials(&self.auth.region) + .await + .map_err(Error::credential(CredentialStage::Identity))?; + let now = SystemTime::now(); + let signed = presign_url( + &connect_url(&self.auth), + TOKEN_LIFETIME, + SigningScope { + region: &self.auth.region, + service: SERVICE, + credentials: &credentials, + time: now, + }, + ) + .map_err(Error::credential(CredentialStage::Signing))?; + let token_expiry = now + TOKEN_LIFETIME; + Ok(Expiring { + username: Some(self.auth.user_name.clone()), + password: SecretValue::new(signed.trim_start_matches("https://")), + expires_at: Some( + credentials + .expiry() + .map_or(token_expiry, |expiry| expiry.min(token_expiry)), + ), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use litellm_redis_types::AwsIamAuth; + use rstest::rstest; + + use super::{connect_url, validate_environment}; + + #[rstest] + #[case::profile("AWS_PROFILE")] + #[case::role("AWS_ROLE_ARN")] + #[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")] + #[case::source_key("AWS_ACCESS_KEY_ID")] + #[case::source_secret("AWS_SECRET_ACCESS_KEY")] + #[case::session("AWS_SESSION_TOKEN")] + #[case::profile_path("AWS_SHARED_CREDENTIALS_FILE")] + #[case::container("AWS_CONTAINER_CREDENTIALS_FULL_URI")] + fn configured_identity_overrides_cannot_silently_use_an_ambient_identity( + #[case] setting: &'static str, + ) { + let configured = |name: &str| (name == setting).then(|| "credential-value".to_owned()); + let process = |name: &str| (name == setting).then(|| "ambient".to_owned()); + let error = validate_environment(&configured, &process).unwrap_err(); + + assert!(matches!( + error, + crate::Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting + )); + assert!(!error.to_string().contains("credential-value")); + assert!(!format!("{error:?}").contains("credential-value")); + } + + #[rstest] + #[case::process_environment("AWS_PROFILE", Some("profile"), Some("profile"))] + #[case::ambient_only("AWS_PROFILE", None, Some("profile"))] + #[case::unconfigured("AWS_PROFILE", None, None)] + #[case::unrelated_bedrock_setting("AWS_ROLE_NAME", Some("bedrock"), None)] + fn unchanged_or_unrelated_settings_preserve_the_sdk_chain( + #[case] setting: &str, + #[case] configured: Option<&str>, + #[case] ambient: Option<&str>, + ) { + let configured = |name: &str| { + (name == setting) + .then_some(configured) + .flatten() + .map(str::to_owned) + }; + let ambient = |name: &str| { + (name == setting) + .then_some(ambient) + .flatten() + .map(str::to_owned) + }; + assert!(validate_environment(&configured, &ambient).is_ok()); + } + + #[rstest] + #[case::provisioned(false, "https://cache-1/?Action=connect&User=iam+user%2F1")] + #[case::serverless( + true, + "https://cache-1/?Action=connect&User=iam+user%2F1&ResourceType=ServerlessCache" + )] + fn the_connect_url_matches_python_urlencode(#[case] serverless: bool, #[case] expected: &str) { + let auth = AwsIamAuth { + user_name: "iam user/1".to_owned(), + cache_name: "Cache-1".to_owned(), + region: "us-east-1".to_owned(), + serverless, + }; + + assert_eq!(connect_url(&auth), expected); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/azure.rs b/litellm-rust/crates/redis-connect/src/sources/azure.rs new file mode 100644 index 00000000000..6dc72e6ce72 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/azure.rs @@ -0,0 +1,64 @@ +use litellm_auth_azure::{AzureAuthService, AzureIdentity}; +use litellm_auth_types::ResolvedCredential; +use litellm_redis_types::AzureAdAuth; + +use crate::Error; +use crate::error::CredentialStage; + +use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture}; + +const AZURE_REDIS_SCOPE: &str = "https://redis.azure.com/.default"; + +pub(crate) struct AzureAd { + service: AzureAuthService, + identity: AzureIdentity, + username: Option, + lookup: SharedEnvironment, +} + +impl AzureAd { + pub(crate) fn new(auth: AzureAdAuth, lookup: SharedEnvironment) -> Self { + // _build_azure_credential in litellm/_redis.py picks the credential class the same way + let identity = match (auth.client_id, auth.tenant_id, auth.client_secret) { + (Some(client_id), Some(tenant_id), Some(client_secret)) => { + AzureIdentity::ClientSecret { + tenant_id, + client_id, + client_secret, + } + } + (Some(client_id), _, _) => AzureIdentity::ManagedIdentity { client_id }, + (None, _, _) => AzureIdentity::Default, + }; + Self { + service: AzureAuthService::default(), + identity, + username: auth.username, + lookup, + } + } +} + +impl CredentialSource for AzureAd { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let credential = self + .service + .identity_token(&self.identity, AZURE_REDIS_SCOPE, &|name| { + self.lookup.get(name) + }) + .await + .map_err(Error::credential(CredentialStage::Identity))? + .ok_or(Error::MissingCredential { provider: "Azure" })?; + let (token, expires_at) = match credential { + ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on), + ResolvedCredential::Static(token) => (token, None), + }; + Ok(Expiring { + username: self.username.clone(), + password: token, + expires_at, + }) + }) + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/gcp.rs b/litellm-rust/crates/redis-connect/src/sources/gcp.rs new file mode 100644 index 00000000000..9fb238709a1 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/gcp.rs @@ -0,0 +1,152 @@ +use std::time::SystemTime; + +use litellm_auth_gcp::{CredentialSource as GcpCredentialSource, GcpAuth}; +use serde::Deserialize; + +use crate::Error; +use crate::error::CredentialStage; +use litellm_secret_value::SecretValue; + +use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture}; + +const IAM_CREDENTIALS: &str = "https://iamcredentials.googleapis.com/v1"; +const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +pub(crate) struct GcpIam { + http: litellm_http::Client, + identity: GcpAuth, + service_account: String, + lookup: SharedEnvironment, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct AccessToken { + access_token: String, + expire_time: String, +} + +impl GcpIam { + pub(crate) fn new( + http: litellm_http::Client, + service_account: String, + lookup: SharedEnvironment, + ) -> Self { + Self { + http, + identity: GcpAuth::default(), + service_account, + lookup, + } + } +} + +fn token_url(service_account: &str) -> String { + match service_account.starts_with("projects/") { + true => format!("{IAM_CREDENTIALS}/{service_account}:generateAccessToken"), + false => format!( + "{IAM_CREDENTIALS}/projects/-/serviceAccounts/{service_account}:generateAccessToken" + ), + } +} + +impl CredentialSource for GcpIam { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let lookup = |name: &str| self.lookup.get(name); + let caller = self + .identity + .access_token(GcpCredentialSource::application_default(&lookup)) + .await + .map_err(Error::credential(CredentialStage::Identity))?; + let response = self + .http + .post(token_url(&self.service_account)) + .bearer_auth(caller.expose()) + .json(&serde_json::json!({ "scope": [CLOUD_PLATFORM_SCOPE] })) + .send() + .await + .map_err(Error::credential(CredentialStage::Exchange))?; + let status = response.status().as_u16(); + let reply: AccessToken = response + .error_for_status() + .map_err(litellm_http::Error::from) + .map_err(Error::credential_response( + CredentialStage::Exchange, + status, + ))? + .json() + .await + .map_err(litellm_http::Error::from) + .map_err(Error::credential(CredentialStage::Response))?; + let expires_at = time::OffsetDateTime::parse( + &reply.expire_time, + &time::format_description::well_known::Rfc3339, + ) + .map(SystemTime::from) + .map_err(Error::credential(CredentialStage::Response))?; + Ok(Expiring { + username: None, + password: SecretValue::new(reply.access_token), + expires_at: Some(expires_at), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::{GcpIam, token_url}; + use crate::CredentialSource; + + #[rstest] + #[case::bare_email( + "redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + #[case::resource_name( + "projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + #[case::project_scoped_resource_name( + "projects/project/serviceAccounts/redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/project/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + fn a_bare_service_account_email_is_addressed_as_a_resource_name( + #[case] service_account: &str, + #[case] expected: &str, + ) { + assert_eq!(token_url(service_account), expected); + } + + #[rstest] + #[tokio::test] + async fn the_caller_identity_comes_from_application_default_credentials_not_vertex_settings() { + let environment = |name: &str| { + match name { + "VERTEXAI_CREDENTIALS" => { + Some(r#"{"type":"service_account","project_id":"vertex"}"#) + } + "GOOGLE_APPLICATION_CREDENTIALS" => Some("/nonexistent/redis-adc.json"), + _ => None, + } + .map(str::to_owned) + }; + let source = GcpIam::new( + litellm_http::Client::for_test(), + "projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com".to_owned(), + std::sync::Arc::new(environment), + ); + + let Err(error) = source.fetch().await else { + panic!("a missing credentials file must fail"); + }; + + assert_eq!( + litellm_error::find::(&error).map(std::io::Error::kind), + Some(std::io::ErrorKind::NotFound) + ); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/mod.rs b/litellm-rust/crates/redis-connect/src/sources/mod.rs new file mode 100644 index 00000000000..988ad3e1d23 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/mod.rs @@ -0,0 +1,7 @@ +mod aws; +mod azure; +mod gcp; + +pub(crate) use aws::ElastiCacheIam; +pub(crate) use azure::AzureAd; +pub(crate) use gcp::GcpIam; diff --git a/litellm-rust/crates/redis-connect/tests/connect.rs b/litellm-rust/crates/redis-connect/tests/connect.rs new file mode 100644 index 00000000000..b1d541f4a5e --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/connect.rs @@ -0,0 +1,230 @@ +use litellm_redis_client::commands; +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use litellm_redis_connect::{ + CredentialSource, Error, Expiring, HostServices, SourceFuture, connect, +}; +use litellm_redis_testing::{ + RedisDeployment, RedisSentinel, RedisServer, RedisTlsServer, SERVICE_NAME, Topology, eventually, +}; +use litellm_redis_types::{ProviderHandle, Resolution}; +use litellm_secret_value::SecretValue; +use rstest::rstest; + +// a host hands its credential_provider object over as a handle, so these tests do the same +fn resolved(kwargs: serde_json::Value) -> Resolution { + let provider = kwargs + .get("credential_provider") + .map(|_| ProviderHandle::new(1)); + let kwargs = serde_json::from_value::>(kwargs).unwrap(); + litellm_redis_resolve::interpret(&litellm_redis_resolve::kwargs(kwargs), provider, None) + .unwrap() +} + +fn services() -> HostServices { + HostServices { + events: None, + http: litellm_http::Client::for_test(), + environment: Arc::new(|_: &str| None), + credential_provider: None, + } +} + +struct HostPassword(&'static str); + +impl CredentialSource for HostPassword { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + Ok(Expiring { + username: None, + password: SecretValue::new(self.0), + expires_at: None, + }) + }) + } +} + +fn port(url: &str) -> u16 { + url.rsplit(':').next().unwrap().parse().unwrap() +} + +#[rstest] +#[case::trusted_ca(true, true)] +#[case::system_roots_only(false, false)] +#[tokio::test] +async fn python_tls_kwargs_reach_the_server(#[case] with_ca: bool, #[case] connects: bool) { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + let server = RedisTlsServer::start(false).await.unwrap(); + let kwargs = serde_json::json!({ + "host": server.host(), + "port": server.port(), + "ssl": true, + "ssl_ca_certs": with_ca.then(|| server.files().ca.display().to_string()), + "socket_timeout": 2, + }); + + let pinged = connect(resolved(kwargs), services(), None) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} + +#[rstest] +#[case::right_password("right", true)] +#[case::wrong_password("wrong", false)] +#[tokio::test] +async fn static_credentials_authenticate( + #[case] password: &str, + #[case] connects: bool, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let kwargs = match &server { + RedisDeployment::Standalone(server) => serde_json::json!({ + "host": "127.0.0.1", + "port": port(server.url()), + "password": password, + "socket_timeout": 2, + }), + RedisDeployment::Cluster(cluster) => serde_json::json!({ + "startup_nodes": cluster.nodes().map(|(host, port)| serde_json::json!({ + "host": host, + "port": port, + })).collect::>(), + "password": password, + "socket_timeout": 2, + }), + }; + + let pinged = connect(resolved(kwargs), services(), None) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} + +#[rstest] +#[tokio::test] +async fn python_sentinel_kwargs_follow_the_master_through_a_failover() { + let sentinel = RedisSentinel::start().await.unwrap(); + let (host, sentinel_port) = sentinel.sentinel(); + let kwargs = serde_json::json!({ + "sentinel_nodes": [[host, sentinel_port]], + "service_name": SERVICE_NAME, + "socket_timeout": 1, + }); + let client = connect(resolved(kwargs), services(), None).unwrap(); + client + .send(commands::set("before", "1", Duration::from_secs(600))) + .await + .unwrap(); + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(600))) + .await + .ok() + }) + .await; + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!(sentinel.master_port().await.unwrap(), promoted); + assert_eq!( + client + .send(commands::get::("after")) + .await + .unwrap() + .as_deref(), + Some("2") + ); +} + +#[rstest] +fn a_python_credential_provider_without_a_host_source_fails_at_boot() { + let error = connect( + resolved(serde_json::json!({"host": "redis", "credential_provider": true})), + services(), + None, + ) + .err() + .unwrap(); + + assert!( + matches!(error, Error::NoHostCredentialProvider), + "{error:?}" + ); +} + +#[rstest] +#[case::profile("AWS_PROFILE")] +#[case::role("AWS_ROLE_ARN")] +#[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")] +fn config_only_aws_identities_fail_before_connecting(#[case] setting: &'static str) { + let configured = format!("{}-override", std::env::var(setting).unwrap_or_default()); + let host = HostServices { + environment: Arc::new(move |name: &str| (name == setting).then(|| configured.clone())), + ..services() + }; + let error = connect( + resolved(serde_json::json!({ + "host": "redis", "ssl": true, "aws_iam_auth": true, + "aws_iam_user_name": "user", "aws_iam_cache_name": "cache", + "aws_iam_region": "us-east-1", + })), + host, + None, + ) + .err() + .expect("must not select ambient credentials"); + + assert!(matches!( + error, + Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting + )); +} + +#[rstest] +#[case::right_password("right", true)] +#[case::wrong_password("wrong", false)] +#[tokio::test] +async fn a_python_credential_provider_authenticates_through_the_host_source( + #[case] password: &'static str, + #[case] connects: bool, +) { + let server = RedisServer::start().await.unwrap(); + server + .cli(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let kwargs = serde_json::json!({ + "host": "127.0.0.1", + "port": port(server.url()), + "password": "ignored", + "credential_provider": true, + "socket_timeout": 2, + }); + + let pinged = connect( + resolved(kwargs), + HostServices { + credential_provider: Some(Arc::new(HostPassword(password))), + ..services() + }, + None, + ) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} diff --git a/litellm-rust/crates/redis-connect/tests/live/aws.rs b/litellm-rust/crates/redis-connect/tests/live/aws.rs new file mode 100644 index 00000000000..da5ec323392 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/aws.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_AWS_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/azure.rs b/litellm-rust/crates/redis-connect/tests/live/azure.rs new file mode 100644 index 00000000000..a7f678edef4 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/azure.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_AZURE_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/gcp.rs b/litellm-rust/crates/redis-connect/tests/live/gcp.rs new file mode 100644 index 00000000000..0965a7a9e9a --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/gcp.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_GCP_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/main.rs b/litellm-rust/crates/redis-connect/tests/live/main.rs new file mode 100644 index 00000000000..90e7637530e --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/main.rs @@ -0,0 +1,111 @@ +mod aws; +mod azure; +mod gcp; + +use std::{collections::BTreeMap, path::PathBuf, sync::Arc, time::Duration}; + +use litellm_http::{ + ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, + media::PublicDnsResolver, +}; +use litellm_redis_client::commands; +use litellm_redis_connect::{HostServices, connect}; +use serde::Deserialize; + +#[derive(Deserialize)] +struct LiveConfig { + redis: BTreeMap, + hold_seconds: u64, +} + +fn config(variable: &str) -> Option { + let path = PathBuf::from(std::env::var_os(variable)?); + Some( + serde_json::from_slice(&std::fs::read(path).expect("read live Redis config")) + .expect("parse live Redis config"), + ) +} + +async fn roundtrip(config: LiveConfig) { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + let lookup = |name: &str| std::env::var(name).ok(); + let resolved = litellm_redis_resolve::interpret( + &litellm_redis_resolve::kwargs(config.redis), + None, + lookup("REDIS_USERNAME").as_deref(), + ) + .expect("resolve live Redis config"); + let http_settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&lookup)]); + let pool = HttpClientPool::new(Arc::new(PublicDnsResolver)); + let http = pool + .client( + &Resolution::from(&http_settings).config, + ClientVariant::Provider, + ) + .expect("build production HTTP client"); + let services = || HostServices { + events: None, + http: http.clone(), + environment: Arc::new(|name: &str| std::env::var(name).ok()), + credential_provider: None, + }; + let client = connect(resolved.clone(), services(), None).expect("build live Redis client"); + let marker = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let key = format!("litellm:redis-connect:live:{}:{marker}", std::process::id()); + let ttl = Duration::from_secs(config.hold_seconds + 120); + + client.ping().await.expect("authenticate and ping"); + client + .send(commands::set(&key, "before-refresh", ttl)) + .await + .expect("set before refresh"); + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + + tokio::time::sleep(Duration::from_secs(config.hold_seconds)).await; + + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + let fresh = connect(resolved, services(), None).expect("build a fresh live Redis client"); + assert_eq!( + fresh + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + fresh + .send(commands::set(&key, "after-refresh", ttl)) + .await + .expect("set after refresh"); + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("after-refresh") + ); + client + .send(commands::del(&[&key])) + .await + .expect("delete live test key"); +} diff --git a/litellm-rust/crates/redis-resolve/AGENTS.md b/litellm-rust/crates/redis-resolve/AGENTS.md new file mode 100644 index 00000000000..b7ce8278975 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/AGENTS.md @@ -0,0 +1,31 @@ +## Interpreting merged kwargs + +- The boundary is merged redis-py kwargs in, `litellm_redis_types::Resolution` out. `interpret` takes the dict a host already merged the way `_merge_redis_kwargs` in `litellm/_redis.py` does, with its `REDIS_*` values and `get_secret` fallbacks filled, and stands in for redis-py's constructors and `from_url`. It never reads the environment or a secret, so both hosts share it. Every `Entry` keeps the source a warning or error names, which `kwargs` sets to the kwarg name and the gateway merge sets to `REDIS_` or the fallback name +- The `Resolution` holds the target, the auth and the timeouts Rust models, and every setting the Rust client ignores as an `Ignored` with its source and consequence, which the host logs. A host passes its `credential_provider` object as a `ProviderHandle`, and a `credential_provider` key without one is `Error::ProviderWithoutHandle`. Python's Azure AD auth reads `REDIS_USERNAME` from `os.environ` instead of the merged kwargs, so it is a required argument of `interpret` that each host reads from its own process environment, the gateway through its `Lookup` and python-bridge through Python's `os.environ`. An empty one is unset, like `or None` +- `options.rs` reads the kwargs into the typed `Options` and collects the ignored ones. `endpoint.rs`, `auth.rs` and `tls.rs` pick the target, the credentials and TLS from `Options` alone +- Rust never introspects redis-py. Every Redis kwarg redis-py or `litellm/_redis.py` accepts is one row in `SETTINGS` in `src/keys.rs`, and that row is the whole answer for the kwarg: its `Handling` (`Read`, `Warn` with the consequence the warning names, or `Reject`), whether a url query carries it (`FromUrl`), whether `_get_redis_env_kwarg_mapping` reads it as `REDIS_`, and the names `_merge_redis_kwargs` falls back to. The last two columns are only read by the gateway merge, and they stay on the row so there is one table. A `Read` row is a field of `Options`, which an inline test pins. A kwarg with no row is not a Redis setting, so it is dropped without a warning +- `redis-py-kwargs.json` holds the Python facts each row answers: which redis-py versions name the kwarg in which signature, and whether Python reads its `REDIS_*` name under the locked redis-py. Never edit it by hand. `redis_py_kwargs.py` next to it regenerates it, and its docstring says how to run it. A kwarg litellm reads that no signature or `REDIS_*` name covers goes in its `LITELLM_ONLY`, and a new redis-py version in its `VERSIONS` and the `redis-compat` CI matrix together. `test_redis_py_kwargs_catalog_matches_the_installed_redis_py` in `tests/unit/test_redis.py` checks it against each version the `redis-compat` CI matrix installs, and `tests/catalog.rs` fails until `SETTINGS` has exactly one row per kwarg in it with the same `environment`. A redis-py bump that adds a kwarg therefore cannot land without a row deciding how Rust handles it +- The contract is `resolves_like_python_redis` in `tests/interpret.rs` and its mirror `test_redis_settings_resolve_like_the_rust_resolver_expects` in `tests/unit/test_redis.py`. Each case holds the merged dict, the sources the merge gave it, the `REDIS_USERNAME` the host passes, and the expected outcome written next to it. `Resolves` is the view both sides build. `Refuses` is a config Rust refuses and Python connects with, a disagreement on purpose. `Fails` is a config redis-py 5.3.1 already raises on, building the client or its first connection, so Rust failing changes no deployment that works today. Its Rust side names the class `Error::raised_building_the_python_client` answers, `None` when Python only raises once a command connects, and its Python mirror `_Fails` asserts the class and message redis-py raises. The Python side asserts that `_merge_redis_kwargs` returns that merged dict and reads the effective redis-py 5.3.1 connection objects without connecting, and `merges_like_python_redis` in `gateway/src/redis/surface_tests.rs` pins the same merged dicts from the Rust merge. A new rule is a new case in every file, never a generated fixture. Python and Rust pods share one Redis during the switch, so a change here is a change to a wire contract +- A url carries its own TLS. redis-py's url allowlist is the plain `Connection` signature, so every `ssl_*` setting outside the url is dropped, and a credential provider strips the url's userinfo. A `rediss://` url's `ssl_cert_reqs`, `ssl_ca_certs`, `ssl_ca_data`, `ssl_ca_path`, `ssl_certfile`, `ssl_keyfile`, `ssl_password` and `ssl_check_hostname` query parameters are read like `parse_url` reads them +- TLS options and `db` are only read for the endpoint that uses them, so a bad `ssl_cert_reqs` never fails a url target. A cluster with a non-zero `db`, and Sentinel with a `host` or `port`, are errors because redis-py raises building those clients +- An empty value (`""`, `[]`, `{}`) is the same as an unset one, like Python's `x or ...`. A null is unset too, except `ssl_cert_reqs`, where redis-py reads `None` as `CERT_NONE`, so a null skips certificate verification +- `ssl_cert_reqs` takes `none`, `optional` and `required` case-sensitively like redis-py's lookup, so `None` (what `REDIS_SSL_CERT_REQS=None` gives) is `Error::Invalid`, and an `ssl.VerifyMode` int, where 0 skips verification and 1 or 2 verifies. `optional` verifies, because rustls has no optional mode +- Precedence matches Python. The target is cluster, then URL, then Sentinel (`sentinel_nodes` plus `service_name`), then host. Auth is a `credential_provider` handle, then GCP, then Azure, then AWS, and the first one wins with a warning +- Numbers are read with `litellm_python_compat::number`, so whitespace and digit underscores parse like Python's `int()` and `float()`. A value they cannot parse (`port: abc`, `socket_timeout: slow`) is dropped with a warning and the default is used, like `_coerce_redis_kwargs_types`. A number that parses but is out of range (negative seconds, a port above 65535) fails with `Error::Invalid` +- Flags (`ssl`, `ssl_check_hostname`, `aws_iam_auth`, `aws_iam_serverless`) are read one way on every surface. `true`, `1` and `yes` are true and `false`, `0`, `no`, `off`, `f` and `n` are false, the text redis-py's `_str_to_bool` and the proxy's pydantic models agree on. Any other text is `Error::Invalid`, so a TLS setting is never guessed +- The resolver fails on a setting it cannot honor only when that setting is a security check. A `Warn` row is dropped and returned as an `Ignored` naming its source (`ssl_ciphers` or `REDIS_SSL_CIPHERS`) and the consequence on its row, because a Python deployment that worked must keep working when its stores move to Rust. Only a value that is malformed (`Error::Invalid` or `Error::UrlQuery`), a target that is missing or a `Reject` row that asks for its check fails. Promoting a `Warn` row to `Read` is then a pure improvement with no behavior cliff, while moving one to `Reject` can fail a config that connects today + - A `Warn` row with a `met` check is dropped silently when its value asks for nothing more than Rust does, so `ssl_min_version` of TLS 1.2 or lower never warns + - A `Reject` row is a security check rustls cannot perform that Python's locked async client already refuses with a `TypeError`, so failing closed breaks no proxy that runs today. `tests/catalog.rs` fails if the locked redis-py's async, cluster or connection signature accepts a `Reject` row, so a redis-py bump forces the decision again. A row's `Asked` says which values ask for the check, from the kwargs, `REDIS_*` or a url query, and those are `Error::Unenforceable` + - `ssl_validate_ocsp` and `ssl_validate_ocsp_stapled` ask when true, because rustls has no client-side OCSP. False asks for nothing and is dropped without a warning + - `ssl_include_verify_flags` and `ssl_exclude_verify_flags` ask when set, because rustls takes no X.509 verify flags such as CRL checking. An empty list asks for nothing + - `ssl_ca_data` and `ssl_ca_path` are read into `TlsOptions` and trusted next to `ssl_ca_certs`. `ssl_password` is read into `TlsOptions::key_password`, which `redis-client` uses to decrypt a PKCS#8 key like the sync client's `load_cert_chain(password=...)`. redis-py 5.3.1's async client refuses it, so this is a superset of the path the proxy uses, like `ssl_ca_path` + - `unix_socket_path` resolves to `Target::Unix`, which like redis-py's `UnixDomainSocketConnection` replaces host, port and TLS. Python still requires `host` next to it, so Rust does too + - A url's `?db=`, `?username=`, `?password=`, `?socket_timeout=`, `?socket_connect_timeout=` and `?health_check_interval=` are read like `parse_url` reads them, so `?db=` wins over the path, the userinfo wins over `?username=` and `?password=`, and the timeouts and health check win over the kwargs like `from_url`. A repeated parameter takes its first value like `parse_qs`. A number `int()` or `float()` cannot read is `Error::UrlQuery`, as `parse_url` raises `ValueError`. Any other query parameter, including one with no row, is ignored naming `url ?`, because redis-py turns each one into a client kwarg. Which ones a url carries is the `FromUrl` column + - `socket_keepalive` is a flag that defaults to true on every target, like `_get_redis_cluster_kwargs` defaults it for a cluster. Python leaves it off for a standalone or Sentinel client, so an unset value there is a difference on purpose, because without keepalive a load balancer drops an idle connection silently. False turns it off. `socket_keepalive_options` is warned, and `redis-client` probes after its own idle time instead + - `health_check_interval` is read as an `int` of seconds like `_coerce_redis_kwargs_types` into `HealthCheck`. Zero is `Disabled` and stays off on every target, a negative interval checks before every command like redis-py, and only `Unset` takes a default. `redis-connect` gives an unset cluster Python's 25s default and `redis-client` pings every node before the next command once the client sat idle that long, outside the breaker and the call events +- `Error::raised_building_the_python_client` is how python-bridge treats a config Rust cannot use. Where Python raises building its client, the bridge raises the same class from `NativeRedisClient.connect`. Where Python builds the client and only fails once a command connects, which `RedisCache` catches and logs on every call, the stores stay on memory and the bridge logs an error +- Cases where Rust and Python disagree on purpose are cases in both contract tests with a different expected outcome on each side + - `ssl_ca_path`: redis-py 5.3.1's async `Redis` has no such parameter and raises `TypeError` building the client, while Rust trusts the directory like `load_verify_locations(capath=...)` + - `on`, `t` and `y` for a flag are `Error::Invalid` on every surface. pydantic reads them as true for `coordination_redis` and `_str_to_bool` reads them as false everywhere else, so Rust refuses to pick one. Other unrecognized text is false in Python and an error here + - `ssl_cert_reqs` as an `ssl.VerifyMode` int: redis-py 5.3.1's sync client takes it, while its async `RedisSSLContext` leaves `cert_reqs` unset and raises `AttributeError` once a command connects. Rust reads it like the sync client + - Empty `sentinel_nodes` with no `REDIS_SENTINEL_NODES` is no target at all, so it fails as "either host or url must be specified". redis-py keeps the empty list and fails building the Sentinel client diff --git a/litellm-rust/crates/redis-resolve/Cargo.toml b/litellm-rust/crates/redis-resolve/Cargo.toml new file mode 100644 index 00000000000..b3d3fc01e44 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-redis-resolve" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true +litellm-secret-value.workspace = true +serde_json.workspace = true +strum.workspace = true +thiserror.workspace = true +tracing.workspace = true +url.workspace = true + +[dev-dependencies] +rstest.workspace = true +serde.workspace = true diff --git a/litellm-rust/crates/redis-resolve/redis-py-kwargs.json b/litellm-rust/crates/redis-resolve/redis-py-kwargs.json new file mode 100644 index 00000000000..22c4495818b --- /dev/null +++ b/litellm-rust/crates/redis-resolve/redis-py-kwargs.json @@ -0,0 +1,118 @@ +{ + "meta": { + "description": "What redis-py and litellm/_redis.py accept for every Redis kwarg the Rust resolver can be handed. SETTINGS in src/keys.rs has exactly one row per arg here", + "redis_py_versions": [ + "5.3.1", + "6.4.0", + "7.4.1", + "8.0.1" + ], + "python": "3.12", + "method": "python redis_py_kwargs.py next to this file, see its docstring. It reads inspect.getfullargspec(inspect.unwrap(cls.__init__)) args and kwonlyargs, minus self", + "signatures": { + "sync": "redis.Redis.__init__. _get_redis_kwargs uses it as the config allowlist and the REDIS_ environment mapping", + "async": "redis.asyncio.Redis.__init__, which the proxy constructs. A kwarg in sync but not async raises TypeError when the client is built", + "connection": "redis.Connection or redis.asyncio.Connection MRO, the url allowlist from _get_redis_url_kwargs", + "cluster": "redis.asyncio.RedisCluster.__init__, the allowlist from _get_redis_cluster_kwargs" + }, + "environment": "true when _get_redis_env_kwarg_mapping reads the arg as REDIS_ under the locked redis-py", + "redis_py": "each redis-py version whose signatures name the arg, with which ones. An empty object means only litellm reads it", + "checked_by": "tests/unit/test_redis.py checks the installed redis-py's column and the REDIS_* names against this file, and tests/catalog.rs checks SETTINGS against it", + "redis_py_locked": "5.3.1" + }, + "args": { + "address_remap": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "auto_close_connection_pool": {"environment": false, "redis_py": {"5.3.1": ["async"], "6.4.0": ["async"], "7.4.1": ["async"], "8.0.1": ["async"]}}, + "aws_iam_auth": {"environment": true, "redis_py": {}}, + "aws_iam_cache_name": {"environment": true, "redis_py": {}}, + "aws_iam_region": {"environment": true, "redis_py": {}}, + "aws_iam_serverless": {"environment": true, "redis_py": {}}, + "aws_iam_user_name": {"environment": true, "redis_py": {}}, + "azure_client_id": {"environment": true, "redis_py": {}}, + "azure_client_secret": {"environment": true, "redis_py": {}}, + "azure_redis_ad_token": {"environment": true, "redis_py": {}}, + "azure_tenant_id": {"environment": true, "redis_py": {}}, + "cache": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "cache_config": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "charset": {"environment": true, "redis_py": {"5.3.1": ["sync"]}}, + "client_name": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "cluster_error_retry_attempts": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "command_packer": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "connection_class": {"environment": false, "redis_py": {}}, + "connection_error_retry_attempts": {"environment": false, "redis_py": {"5.3.1": ["cluster"]}}, + "connection_pool": {"environment": false, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "credential_provider": {"environment": false, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "db": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "decode_responses": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "driver_info": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "dynamic_startup_nodes": {"environment": false, "redis_py": {"6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "encoder_class": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "encoding": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "encoding_errors": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "errors": {"environment": true, "redis_py": {"5.3.1": ["sync"]}}, + "event_dispatcher": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "gcp_service_account": {"environment": true, "redis_py": {}}, + "gcp_ssl_ca_certs": {"environment": true, "redis_py": {}}, + "health_check_interval": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "host": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "legacy_responses": {"environment": false, "redis_py": {"8.0.1": ["sync", "async", "connection", "cluster"]}}, + "lib_name": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "lib_version": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "load_balancing_strategy": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "maint_notifications_config": {"environment": false, "redis_py": {"7.4.1": ["sync", "connection"], "8.0.1": ["sync", "connection"]}}, + "maint_notifications_pool_handler": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "maintenance_notification_hash": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "maintenance_state": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "max_connections": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "namespace": {"environment": false, "redis_py": {}}, + "orig_host_address": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "orig_socket_connect_timeout": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "orig_socket_timeout": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "oss_cluster_maint_notifications_handler": {"environment": false, "redis_py": {"7.4.1": ["sync", "connection"], "8.0.1": ["sync", "connection"]}}, + "parser": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "parser_class": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "password": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "path": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "policy_resolver": {"environment": false, "redis_py": {"7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "port": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "protocol": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "read_from_replicas": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "redis_connect_func": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection"], "6.4.0": ["sync", "async", "connection"], "7.4.1": ["sync", "async", "connection"], "8.0.1": ["sync", "async", "connection"]}}, + "reinitialize_steps": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "require_full_coverage": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "retry": {"environment": false, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "retry_on_error": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "retry_on_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection"], "6.4.0": ["sync", "async", "connection"], "7.4.1": ["sync", "async", "connection"], "8.0.1": ["sync", "async", "connection"]}}, + "sentinel_nodes": {"environment": false, "redis_py": {}}, + "sentinel_password": {"environment": false, "redis_py": {}}, + "service_name": {"environment": false, "redis_py": {}}, + "single_connection_client": {"environment": true, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "socket_connect_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_keepalive": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_keepalive_options": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_read_size": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_type": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "ssl": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_certs": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_data": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_path": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "ssl_cert_reqs": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_certfile": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_check_hostname": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ciphers": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_exclude_verify_flags": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_include_verify_flags": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_keyfile": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_min_version": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ocsp_context": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_ocsp_expected_cert": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_password": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "ssl_validate_ocsp": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_validate_ocsp_stapled": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "startup_nodes": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "unix_socket_path": {"environment": true, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "url": {"environment": true, "redis_py": {}}, + "username": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}} + } +} diff --git a/litellm-rust/crates/redis-resolve/redis_py_kwargs.py b/litellm-rust/crates/redis-resolve/redis_py_kwargs.py new file mode 100644 index 00000000000..9459cfd3761 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/redis_py_kwargs.py @@ -0,0 +1,110 @@ +"""Regenerates redis-py-kwargs.json, the Python facts SETTINGS in src/keys.rs answers. + +Run it from the repository root in the project environment, which has the locked redis-py and litellm: + + uv run python litellm-rust/crates/redis-resolve/redis_py_kwargs.py + +It reads every other redis-py version's signatures in a throwaway environment through +`uv run --no-project --with redis==`, so nothing else is installed. Then update SETTINGS until +tests/catalog.rs passes. +""" + +import inspect +import json +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +VERSIONS: Final = ("5.3.1", "6.4.0", "7.4.1", "8.0.1") +PYTHON: Final = "3.12" +CATALOG: Final = Path(__file__).with_name("redis-py-kwargs.json") +# Keys litellm/_redis.py or its callers read that neither a redis-py signature nor the REDIS_* mapping names +LITELLM_ONLY: Final = frozenset( + {"connection_class", "namespace", "sentinel_nodes", "sentinel_password", "service_name"} +) + +META: Final = { + "description": "What redis-py and litellm/_redis.py accept for every Redis kwarg the Rust resolver can be handed. SETTINGS in src/keys.rs has exactly one row per arg here", + "redis_py_versions": list(VERSIONS), + "python": PYTHON, + "method": "python redis_py_kwargs.py next to this file, see its docstring. It reads inspect.getfullargspec(inspect.unwrap(cls.__init__)) args and kwonlyargs, minus self", + "signatures": { + "sync": "redis.Redis.__init__. _get_redis_kwargs uses it as the config allowlist and the REDIS_ environment mapping", + "async": "redis.asyncio.Redis.__init__, which the proxy constructs. A kwarg in sync but not async raises TypeError when the client is built", + "connection": "redis.Connection or redis.asyncio.Connection MRO, the url allowlist from _get_redis_url_kwargs", + "cluster": "redis.asyncio.RedisCluster.__init__, the allowlist from _get_redis_cluster_kwargs", + }, + "environment": "true when _get_redis_env_kwarg_mapping reads the arg as REDIS_ under the locked redis-py", + "redis_py": "each redis-py version whose signatures name the arg, with which ones. An empty object means only litellm reads it", + "checked_by": "tests/unit/test_redis.py checks the installed redis-py's column and the REDIS_* names against this file, and tests/catalog.rs checks SETTINGS against it", +} + + +def _named_args(*classes: type) -> frozenset[str]: + return frozenset( + name + for cls in classes + for spec in (inspect.getfullargspec(inspect.unwrap(cls.__init__)),) + for name in spec.args + spec.kwonlyargs + if name != "self" + ) + + +def arg_signatures() -> dict[str, list[str]]: + """Each kwarg the installed redis-py names, with the signatures that name it.""" + import redis + import redis.asyncio as async_redis + + connection_classes: Final = (*redis.Connection.__mro__, *async_redis.Connection.__mro__) + signatures: Final = { + "sync": _named_args(redis.Redis), + "async": _named_args(async_redis.Redis), + "connection": _named_args(*(cls for cls in connection_classes if cls is not object)), + "cluster": _named_args(async_redis.RedisCluster), + } + names: Final = sorted(frozenset().union(*signatures.values())) + return {name: [kind for kind, named in signatures.items() if name in named] for name in names} + + +def _arg_signatures_of(version: str) -> Mapping[str, list[str]]: + command: Final = ("uv", "run", "--no-project", "--python", PYTHON, "--with", f"redis=={version}") + script: Final = ("python", "-I", __file__, "--arg-signatures") + return json.loads(subprocess.run((*command, *script), check=True, capture_output=True, text=True).stdout) + + +def catalog() -> dict[str, object]: + import redis + + from litellm._redis import _get_redis_env_kwarg_mapping + + environment: Final = frozenset(_get_redis_env_kwarg_mapping().values()) + by_version: Final = {version: _arg_signatures_of(version) for version in VERSIONS} + names: Final = sorted(LITELLM_ONLY | environment | frozenset().union(*by_version.values())) + return { + "meta": {**META, "redis_py_locked": redis.__version__}, + "args": { + name: { + "environment": name in environment, + "redis_py": {version: args[name] for version, args in by_version.items() if name in args}, + } + for name in names + }, + } + + +# one line per arg, so a version bump reads as a diff of the args it touched +def render(catalog: Mapping[str, object]) -> str: + meta: Final = json.dumps(catalog["meta"], indent=2).replace("\n", "\n ") + args: Final = catalog["args"] + assert isinstance(args, dict) + rows: Final = ",\n".join(f" {json.dumps(name)}: {json.dumps(arg)}" for name, arg in args.items()) + return f'{{\n "meta": {meta},\n "args": {{\n{rows}\n }}\n}}\n' + + +if __name__ == "__main__": + if sys.argv[1:] == ["--arg-signatures"]: + sys.stdout.write(json.dumps(arg_signatures())) + else: + CATALOG.write_text(render(catalog())) diff --git a/litellm-rust/crates/redis-resolve/src/auth.rs b/litellm-rust/crates/redis-resolve/src/auth.rs new file mode 100644 index 00000000000..620f1a2e11b --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/auth.rs @@ -0,0 +1,95 @@ +use litellm_redis_types::{Auth, AwsIamAuth, AzureAdAuth, ProviderHandle, StaticAuth}; + +use crate::{ + Error, + options::{self, AwsOptions, Options, flag, secret, text}, +}; + +pub(super) fn redis_auth( + options: &Options, + provider: Option, + redis_username: Option<&str>, +) -> Result { + match (provider, &options.credential_provider) { + (Some(handle), _) => return Ok(Auth::Provider(handle)), + (None, Some(_)) => return Err(Error::ProviderWithoutHandle), + (None, None) => {} + } + let gcp = text(options.gcp.service_account.as_ref()); + let azure = text(options.azure.ad_token.as_ref()) + .is_some_and(|enabled| enabled.eq_ignore_ascii_case("true")); + let aws = flag(options.aws.auth.as_ref())?; + match (gcp, azure, aws) { + (Some(service_account), azure, aws) => { + if azure { + tracing::warn!( + "Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. Using GCP IAM. Remove one to avoid misconfiguration." + ); + } + if aws { + tracing::warn!( + "Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured for Redis. Using GCP IAM. Remove one to avoid misconfiguration." + ); + } + Ok(Auth::GcpIam { service_account }) + } + (None, true, aws) => { + if aws { + tracing::warn!( + "Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured for Redis. Using Azure AD. Remove one to avoid misconfiguration." + ); + } + Ok(Auth::AzureAd(AzureAdAuth { + client_id: text(options.azure.client_id.as_ref()), + tenant_id: text(options.azure.tenant_id.as_ref()), + client_secret: secret(options.azure.client_secret.as_ref()), + username: redis_username + .filter(|username| !username.is_empty()) + .map(str::to_owned), + })) + } + (None, false, true) => aws_iam(&options.aws, uses_tls(options)?).map(Auth::AwsIam), + (None, false, false) => Ok(Auth::Static(StaticAuth { + username: secret(options.username.as_ref()), + password: secret(options.password.as_ref()), + })), + } +} + +// _uses_tls: a url's scheme decides unless cluster nodes are set, otherwise the ssl kwarg does +fn uses_tls(options: &Options) -> Result { + match (&options.startup_nodes, text(options.url.as_ref())) { + (None, Some(url)) => Ok(url + .split_once("://") + .is_some_and(|(scheme, _)| scheme.eq_ignore_ascii_case("rediss"))), + _ => flag(options.ssl.as_ref()), + } +} + +fn aws_iam(aws: &AwsOptions, uses_tls: bool) -> Result { + if !uses_tls { + return Err(Error::AwsIamWithoutTls); + } + match ( + text(aws.user_name.as_ref()), + text(aws.cache_name.as_ref()), + text(aws.region.as_ref()), + ) { + (Some(user_name), Some(cache_name), Some(region)) => Ok(AwsIamAuth { + user_name, + cache_name, + region, + serverless: options::flag(aws.serverless.as_ref())?, + }), + (user_name, cache_name, region) => Err(Error::AwsIamIncomplete { + missing: [ + ("aws_iam_user_name", user_name.is_none()), + ("aws_iam_cache_name", cache_name.is_none()), + ("aws_iam_region", region.is_none()), + ] + .into_iter() + .filter_map(|(name, missing)| missing.then_some(name)) + .collect(), + }), + } +} diff --git a/litellm-rust/crates/redis-resolve/src/endpoint.rs b/litellm-rust/crates/redis-resolve/src/endpoint.rs new file mode 100644 index 00000000000..e3722b1b7cc --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/endpoint.rs @@ -0,0 +1,180 @@ +use litellm_redis_types::{ + Auth, Cluster, Node, RedisUrl, Sentinel, Standalone, StaticAuth, Target, TlsOptions, UnixSocket, +}; +use litellm_secret_value::SecretValue; + +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, + options::{Options, as_text, flag, integer, invalid, secret, text}, + tls::{self, QuerySettings}, +}; + +/// The target, first match wins like `get_redis_client`: cluster, url, Sentinel, then host. A +/// url also returns the settings its query carries. +pub(super) fn target(options: &Options, auth: &Auth) -> Result<(Target, QuerySettings), Error> { + if let Some(nodes) = cluster_nodes(options.startup_nodes.as_ref())? { + reject_cluster_db(options.db.as_ref())?; + let cluster = Cluster { + nodes, + tls: tls_for(options, auth)?, + }; + return Ok((Target::Cluster(cluster), QuerySettings::default())); + } + if let Some(url) = text(options.url.as_ref()) { + let query = tls::url_query(&url)?; + let url = RedisUrl { + url: SecretValue::new(match auth { + Auth::Static(_) => url, + _ => without_userinfo(&url), + }), + tls: query.tls, + db: query.db, + username: query.username, + password: query.password, + }; + return Ok((Target::Url(url), query.settings)); + } + if let Some(sentinels) = options.sentinel.nodes.as_ref() { + let sentinel = sentinel(options, auth, sentinels)?; + return Ok((Target::Sentinel(sentinel), QuerySettings::default())); + } + let host = text(options.host.as_ref()).ok_or(Error::NoTarget)?; + let target = match text(options.unix_socket_path.as_ref()) { + Some(path) => Target::Unix(UnixSocket { + path: path.into(), + db: db(options)?, + }), + None => Target::Standalone(Standalone { + host, + port: integer(options.port.as_ref())?.unwrap_or(6379), + tls: tls_for(options, auth)?, + db: db(options)?, + }), + }; + Ok((target, QuerySettings::default())) +} + +fn tls_for(options: &Options, auth: &Auth) -> Result, Error> { + if !flag(options.ssl.as_ref())? { + return Ok(None); + } + let gcp_ca_certs = match auth { + Auth::GcpIam { .. } => options.gcp.ssl_ca_certs.as_ref(), + _ => None, + }; + tls::tls_options(&options.tls, gcp_ca_certs).map(Some) +} + +fn db(options: &Options) -> Result { + integer(options.db.as_ref()).map(|db| db.unwrap_or(0)) +} + +fn sentinel(options: &Options, auth: &Auth, sentinels: &Entry) -> Result { + let monitors = sentinel_nodes(&sentinels.value) + .ok_or_else(|| invalid(&sentinels.source, Expected::SentinelNodes))?; + let service_name = + text(options.sentinel.service_name.as_ref()).ok_or(Error::NoSentinelService)?; + // redis.asyncio.Sentinel passes the connection kwargs to each monitor's Redis(host=..., port=...), so a + // host or port setting raises TypeError when litellm builds the client + if let Some(entry) = options.host.as_ref().or(options.port.as_ref()) { + return Err(Error::SentinelWithHost { + setting: entry.source.clone(), + }); + } + Ok(Sentinel { + monitors, + service_name, + monitor_auth: monitor_auth(options, auth), + tls: tls_for(options, auth)?, + db: db(options)?, + }) +} + +// _sentinel_auth_kwargs gives the monitors the data node's connection kwargs with sentinel_password as the +// password, so a static username reaches them too +fn monitor_auth(options: &Options, auth: &Auth) -> Option { + let password = secret(options.sentinel.password.as_ref())?; + let username = match auth { + Auth::Static(StaticAuth { username, .. }) => username.clone(), + _ => None, + }; + Some(StaticAuth { + username, + password: Some(password), + }) +} + +fn cluster_nodes(entry: Option<&Entry>) -> Result>, Error> { + let Some(entry) = entry else { + return Ok(None); + }; + let nodes = nodes(&entry.value, |node| { + node_address( + node.get("host").and_then(as_text), + node.get("port").and_then(as_text), + ) + }) + .ok_or_else(|| invalid(&entry.source, Expected::StartupNodes))?; + Ok((!nodes.is_empty()).then_some(nodes)) +} + +fn sentinel_nodes(value: &Value) -> Option> { + let nodes = nodes(value, |pair| match pair { + Value::Array(pair) if pair.len() == 2 => node_address(as_text(&pair[0]), as_text(&pair[1])), + _ => None, + })?; + (!nodes.is_empty()).then_some(nodes) +} + +fn nodes(value: &Value, parse: impl Fn(&Value) -> Option) -> Option> { + match value { + Value::String(json) => serde_json::from_str::>(json) + .ok()? + .iter() + .map(parse) + .collect(), + Value::Array(nodes) => nodes.iter().map(parse).collect(), + _ => None, + } +} + +fn node_address(host: Option, port: Option) -> Option { + Some(Node { + host: host?, + port: port?.parse().ok()?, + }) +} + +// redis-py's RedisCluster raises unless db is 0 or None +fn reject_cluster_db(db: Option<&Entry>) -> Result<(), Error> { + let Some(entry) = db else { + return Ok(()); + }; + let is_zero = match &entry.value { + Value::String(_) => matches!(integer::(Some(entry)), Ok(None | Some(0))), + Value::Bool(flag) => !flag, + Value::Number(number) => number.as_f64() == Some(0.0), + _ => false, + }; + match is_zero { + true => Ok(()), + false => Err(Error::DbInCluster { + setting: entry.source.clone(), + }), + } +} + +// _get_redis_url_kwargs keeps only plain Connection kwargs, so redis-py drops every ssl_* setting for a url +// and a credential provider replaces the url's userinfo, which _url_without_userinfo strips +fn without_userinfo(url: &str) -> String { + let Some((scheme, rest)) = url.split_once("://") else { + return url.to_owned(); + }; + let (authority, tail) = rest.split_at(rest.find(['/', '?', '#']).unwrap_or(rest.len())); + let host = authority + .rsplit_once('@') + .map_or(authority, |(_, host)| host); + format!("{scheme}://{host}{tail}") +} diff --git a/litellm-rust/crates/redis-resolve/src/error.rs b/litellm-rust/crates/redis-resolve/src/error.rs new file mode 100644 index 00000000000..b13b65edfb5 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/error.rs @@ -0,0 +1,80 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid Redis setting {setting}: expected {expected}")] + Invalid { setting: String, expected: Expected }, + #[error("invalid Redis setting {setting}: redis-py's parse_url cannot read it as a number")] + UrlQuery { setting: String }, + #[error("either host or url must be specified for redis")] + NoTarget, + #[error("service_name is required with sentinel_nodes")] + NoSentinelService, + #[error("{setting} cannot be combined with sentinel_nodes, which redis-py rejects at startup")] + SentinelWithHost { setting: String }, + #[error("{setting} must be 0 in cluster mode")] + DbInCluster { setting: String }, + #[error("credential_provider is a host object, and the host passed no handle for it")] + ProviderWithoutHandle, + #[error( + "{setting} asks for {check}, which rustls cannot perform, so Redis is not connected without them" + )] + Unenforceable { + setting: String, + check: &'static str, + }, + #[error("AWS ElastiCache IAM Redis authentication requires TLS")] + AwsIamWithoutTls, + #[error("AWS ElastiCache IAM Redis authentication requires: {}", missing.join(", "))] + AwsIamIncomplete { missing: Vec<&'static str> }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)] +pub enum Expected { + #[strum(serialize = "an integer")] + Integer, + #[strum(serialize = "a boolean")] + Boolean, + #[strum(serialize = "a non-negative number of seconds")] + Seconds, + #[strum(serialize = "none, optional, required, null or an ssl.VerifyMode value")] + CertRequirement, + #[strum(serialize = "a list of objects with host and port")] + StartupNodes, + #[strum(serialize = "a non-empty list of [host, port] pairs")] + SentinelNodes, +} + +/// The exception class redis-py 5.3.1 raises building its client. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PythonException { + ValueError, + TypeError, + RedisClusterException, +} + +impl Error { + /// What `_get_redis_client_logic` and redis-py's async client constructors raise on the same kwargs, or `None` + /// when Python builds the client and fails only once a command opens a connection, or never fails. + pub fn raised_building_the_python_client(&self) -> Option { + match self { + Self::Invalid { expected, .. } => match expected { + Expected::StartupNodes | Expected::SentinelNodes => { + Some(PythonException::ValueError) + } + Expected::Integer + | Expected::Boolean + | Expected::Seconds + | Expected::CertRequirement => None, + }, + Self::UrlQuery { .. } + | Self::NoTarget + | Self::NoSentinelService + | Self::AwsIamWithoutTls + | Self::AwsIamIncomplete { .. } => Some(PythonException::ValueError), + Self::SentinelWithHost { .. } | Self::Unenforceable { .. } => { + Some(PythonException::TypeError) + } + Self::DbInCluster { .. } => Some(PythonException::RedisClusterException), + Self::ProviderWithoutHandle => None, + } + } +} diff --git a/litellm-rust/crates/redis-resolve/src/keys.rs b/litellm-rust/crates/redis-resolve/src/keys.rs new file mode 100644 index 00000000000..c4897701d6f --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/keys.rs @@ -0,0 +1,304 @@ +use litellm_redis_types::Ignored; +use serde_json::Value; + +/// How the Rust client honors a redis-py kwarg. +#[derive(Clone, Copy, Debug)] +pub enum Handling { + /// Modeled by the resolver into a field of `Options`. + Read, + /// Dropped and returned as an `Ignored` naming `consequence`, unless `met` says the value asks for nothing the + /// Rust client does not already do. + Warn { + consequence: &'static str, + met: Option bool>, + }, + /// A security check the Rust client cannot perform and Python's async client already refuses, so asking for it + /// fails instead of connecting without it. + Reject { check: &'static str, asked: Asked }, +} + +/// Which values of a `Reject` row ask for its check. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Asked { + /// A flag read as true. False asks for nothing. + WhenTrue, + /// Any present value, like a list of flags. An empty one asks for nothing. + WhenSet, +} + +/// Whether redis-py's `parse_url` reads the setting from a url's query and the Rust resolver does too. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FromUrl { + No, + Any, + Rediss, +} + +#[derive(Clone, Copy, Debug)] +pub struct Setting { + pub kwarg: &'static str, + pub handling: Handling, + pub from_url: FromUrl, + /// `_get_redis_env_kwarg_mapping` reads it as `REDIS_`. + pub environment: bool, + /// The names `_merge_redis_kwargs` reads with `get_secret` when the merged kwargs lack the setting. + pub fallback: &'static [&'static str], +} + +impl Setting { + const fn new(kwarg: &'static str, handling: Handling) -> Self { + Self { + kwarg, + handling, + from_url: FromUrl::No, + environment: false, + fallback: &[], + } + } + + const fn env(self) -> Self { + Self { + environment: true, + ..self + } + } + + const fn in_url(self, from_url: FromUrl) -> Self { + Self { from_url, ..self } + } + + const fn or_secrets(self, fallback: &'static [&'static str]) -> Self { + Self { fallback, ..self } + } + + pub fn environment_name(&self) -> Option { + self.environment + .then(|| format!("REDIS_{}", self.kwarg.to_ascii_uppercase())) + } + + pub fn is_read(&self) -> bool { + matches!(self.handling, Handling::Read) + } + + /// The warning a set value earns, or `None` when the setting is read, refused or already met. + pub(crate) fn warning(&self, value: &Value) -> Option<&'static str> { + match self.handling { + Handling::Warn { consequence, met } if !met.is_some_and(|met| met(value)) => { + Some(consequence) + } + _ => None, + } + } +} + +const fn read(kwarg: &'static str) -> Setting { + Setting::new(kwarg, Handling::Read) +} + +const fn warn(kwarg: &'static str, consequence: &'static str) -> Setting { + Setting::new( + kwarg, + Handling::Warn { + consequence, + met: None, + }, + ) +} + +const fn warn_unless( + kwarg: &'static str, + consequence: &'static str, + met: fn(&Value) -> bool, +) -> Setting { + Setting::new( + kwarg, + Handling::Warn { + consequence, + met: Some(met), + }, + ) +} + +const fn reject(kwarg: &'static str, check: &'static str, asked: Asked) -> Setting { + Setting::new(kwarg, Handling::Reject { check, asked }) +} + +const NO_EQUIVALENT: &str = "the Rust client has no equivalent"; +const OWN_STACK: &str = "the Rust client uses its own pool, retries and events"; +const OWN_CONNECTION: &str = "it tunes redis-py's own connection, and the Rust client has its own"; +const OWN_CLUSTER: &str = "the Rust cluster client discovers, routes and retries on its own"; +const MAINTENANCE: &str = "the Rust client does not follow server maintenance notifications"; +const RUSTLS: &str = "rustls negotiates TLS 1.2+ with its own ECDHE AEAD suites"; +const REVOCATION: &str = "OCSP revocation checks"; +const OCSP: &str = "it only configures ssl_validate_ocsp, which the Rust client refuses when set"; +const VERIFY_FLAGS: &str = "X.509 verify flags such as CRL checking"; + +/// Every Redis kwarg redis-py or `litellm/_redis.py` accepts, one row each. `redis-py-kwargs.json` records which +/// redis-py versions and `REDIS_*` names accept them, and `tests/catalog.rs` keeps the two in step. +pub const SETTINGS: &[Setting] = &[ + warn("address_remap", OWN_CLUSTER), + warn("auto_close_connection_pool", OWN_STACK), + read("aws_iam_auth").env(), + read("aws_iam_cache_name").env(), + read("aws_iam_region") + .env() + .or_secrets(&["AWS_REGION", "AWS_DEFAULT_REGION"]), + read("aws_iam_serverless").env(), + read("aws_iam_user_name").env(), + read("azure_client_id") + .env() + .or_secrets(&["AZURE_CLIENT_ID"]), + read("azure_client_secret") + .env() + .or_secrets(&["AZURE_CLIENT_SECRET"]), + read("azure_redis_ad_token") + .env() + .or_secrets(&["REDIS_AZURE_AD_TOKEN"]), + read("azure_tenant_id") + .env() + .or_secrets(&["AZURE_TENANT_ID"]), + warn("cache", NO_EQUIVALENT).env(), + warn("cache_config", NO_EQUIVALENT).env(), + warn("charset", NO_EQUIVALENT).env(), + warn("client_name", NO_EQUIVALENT).env(), + warn("cluster_error_retry_attempts", OWN_CLUSTER), + warn("command_packer", OWN_CONNECTION), + warn("connection_class", OWN_STACK), + warn("connection_error_retry_attempts", OWN_CLUSTER), + warn("connection_pool", OWN_STACK), + read("credential_provider"), + read("db").env().in_url(FromUrl::Any), + warn("decode_responses", NO_EQUIVALENT).env(), + warn("driver_info", NO_EQUIVALENT), + warn("dynamic_startup_nodes", OWN_CLUSTER), + warn("encoder_class", OWN_CONNECTION), + warn("encoding", NO_EQUIVALENT).env(), + warn("encoding_errors", NO_EQUIVALENT).env(), + warn("errors", NO_EQUIVALENT).env(), + warn("event_dispatcher", OWN_STACK).env(), + read("gcp_service_account") + .env() + .or_secrets(&["REDIS_GCP_SERVICE_ACCOUNT"]), + read("gcp_ssl_ca_certs") + .env() + .or_secrets(&["REDIS_GCP_SSL_CA_CERTS"]), + read("health_check_interval").env().in_url(FromUrl::Any), + read("host").env(), + warn("legacy_responses", NO_EQUIVALENT), + warn("lib_name", NO_EQUIVALENT).env(), + warn("lib_version", NO_EQUIVALENT).env(), + warn("load_balancing_strategy", OWN_CLUSTER), + warn("maint_notifications_config", MAINTENANCE), + warn("maint_notifications_pool_handler", MAINTENANCE), + warn("maintenance_notification_hash", MAINTENANCE), + warn("maintenance_state", MAINTENANCE), + warn("max_connections", "the Rust client has no connection cap").env(), + read("namespace"), + warn("orig_host_address", MAINTENANCE), + warn("orig_socket_connect_timeout", MAINTENANCE), + warn("orig_socket_timeout", MAINTENANCE), + warn("oss_cluster_maint_notifications_handler", MAINTENANCE), + warn("parser", OWN_CONNECTION), + warn("parser_class", OWN_CONNECTION), + read("password").env().in_url(FromUrl::Any), + warn("path", OWN_CLUSTER), + warn("policy_resolver", OWN_CLUSTER), + read("port").env(), + warn("protocol", NO_EQUIVALENT).env(), + warn("read_from_replicas", OWN_CLUSTER), + warn( + "redis_connect_func", + "the Python object is not callable from Rust, so static username and password are used", + ) + .env(), + warn("reinitialize_steps", OWN_CLUSTER), + warn("require_full_coverage", OWN_CLUSTER), + warn("retry", OWN_STACK), + warn("retry_on_error", NO_EQUIVALENT).env(), + warn("retry_on_timeout", NO_EQUIVALENT).env(), + read("sentinel_nodes").or_secrets(&["REDIS_SENTINEL_NODES"]), + read("sentinel_password").or_secrets(&["REDIS_SENTINEL_PASSWORD"]), + read("service_name").or_secrets(&["REDIS_SERVICE_NAME"]), + warn("single_connection_client", NO_EQUIVALENT).env(), + read("socket_connect_timeout").env().in_url(FromUrl::Any), + read("socket_keepalive").env(), + warn("socket_keepalive_options", NO_EQUIVALENT).env(), + warn("socket_read_size", OWN_CONNECTION), + read("socket_timeout").env().in_url(FromUrl::Any), + warn("socket_type", OWN_CONNECTION), + read("ssl").env(), + read("ssl_ca_certs").env().in_url(FromUrl::Rediss), + read("ssl_ca_data").env().in_url(FromUrl::Rediss), + read("ssl_ca_path").env().in_url(FromUrl::Rediss), + read("ssl_cert_reqs").env().in_url(FromUrl::Rediss), + read("ssl_certfile").env().in_url(FromUrl::Rediss), + read("ssl_check_hostname").env().in_url(FromUrl::Rediss), + warn("ssl_ciphers", RUSTLS).env(), + reject("ssl_exclude_verify_flags", VERIFY_FLAGS, Asked::WhenSet), + reject("ssl_include_verify_flags", VERIFY_FLAGS, Asked::WhenSet), + read("ssl_keyfile").env().in_url(FromUrl::Rediss), + warn_unless("ssl_min_version", RUSTLS, is_met_by_rustls).env(), + warn("ssl_ocsp_context", OCSP).env(), + warn("ssl_ocsp_expected_cert", OCSP).env(), + read("ssl_password").env().in_url(FromUrl::Rediss), + reject("ssl_validate_ocsp", REVOCATION, Asked::WhenTrue).env(), + reject("ssl_validate_ocsp_stapled", REVOCATION, Asked::WhenTrue).env(), + read("startup_nodes").or_secrets(&["REDIS_CLUSTER_NODES"]), + read("unix_socket_path").env(), + read("url").env(), + read("username").env().in_url(FromUrl::Any), +]; + +// rustls only speaks TLS 1.2 and 1.3, so a minimum of 1.2 or lower is already met. redis-py takes an +// ssl.TLSVersion, which reaches Rust as its name or its value +fn is_met_by_rustls(value: &Value) -> bool { + match value { + Value::Number(number) => matches!(number.as_i64(), Some(-2 | 769 | 770 | 771)), + Value::String(text) => matches!( + text.strip_prefix("TLSVersion.") + .unwrap_or(text) + .to_ascii_uppercase() + .as_str(), + "MINIMUM_SUPPORTED" | "TLSV1" | "TLSV1_1" | "TLSV1_2" + ), + _ => false, + } +} + +pub fn setting(kwarg: &str) -> Option<&'static Setting> { + SETTINGS.iter().find(|setting| setting.kwarg == kwarg) +} + +/// The check a `Reject` row stands for and which values ask for it. +pub(crate) fn refusal(kwarg: &str) -> Option<(&'static str, Asked)> { + match setting(kwarg)?.handling { + Handling::Reject { check, asked } => Some((check, asked)), + Handling::Read | Handling::Warn { .. } => None, + } +} + +/// A warned setting the kwargs set, with what dropping it means. A kwarg with no row is not a Redis setting, so it +/// is dropped without a warning. +pub(crate) fn ignored(source: &str, kwarg: &str, value: &Value) -> Option { + let consequence = setting(kwarg)?.warning(value)?; + Some(Ignored { + setting: source.to_owned(), + consequence, + }) +} + +/// A url query parameter the resolver does not read. redis-py turns every one into a client kwarg, so even one with +/// no row is reported. +pub(crate) fn ignored_query(source: &str, name: &str, value: &Value) -> Option { + let consequence = match setting(name) { + None => NO_EQUIVALENT, + Some(setting) if setting.is_read() => { + "the Rust client does not read it from this url's query" + } + Some(setting) => setting.warning(value)?, + }; + Some(Ignored { + setting: source.to_owned(), + consequence, + }) +} diff --git a/litellm-rust/crates/redis-resolve/src/lib.rs b/litellm-rust/crates/redis-resolve/src/lib.rs new file mode 100644 index 00000000000..b7aa5d18f64 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/lib.rs @@ -0,0 +1,95 @@ +use std::{collections::BTreeMap, time::Duration}; + +use litellm_redis_types::{Auth, HealthCheck, ProviderHandle, Resolution, StaticAuth, Target}; +use serde_json::Value; + +mod auth; +mod endpoint; +mod error; +mod keys; +mod options; +mod tls; + +pub use error::{Error, Expected, PythonException}; +pub use keys::{Asked, FromUrl, Handling, SETTINGS, Setting, setting}; +pub use options::is_present; + +use options::Options; + +/// One merged redis-py kwarg and the name a warning or error reports it under. +#[derive(Clone, Debug, PartialEq)] +pub struct Entry { + pub value: Value, + pub source: String, +} + +pub type Kwargs = BTreeMap; + +/// Kwargs whose every entry is reported under its own kwarg name. +pub fn kwargs(fields: impl IntoIterator) -> Kwargs { + fields + .into_iter() + .map(|(kwarg, value)| { + let entry = Entry { + value, + source: kwarg.clone(), + }; + (kwarg, entry) + }) + .collect() +} + +/// What redis-py's constructors make of kwargs a host already merged the way `_get_redis_client_logic` does. A +/// host passes its `credential_provider` object as `provider`, and `REDIS_USERNAME` as its process environment +/// holds it, which Python's Azure AD auth reads from `os.environ` instead of the merged kwargs. +pub fn interpret( + kwargs: &Kwargs, + provider: Option, + redis_username: Option<&str>, +) -> Result { + let (options, ignored) = Options::read(kwargs)?; + let auth = auth::redis_auth(&options, provider, redis_username)?; + let (target, query) = endpoint::target(&options, &auth)?; + let auth = match (&target, auth) { + (Target::Url(_), Auth::Static(StaticAuth { username, .. })) => Auth::Static(StaticAuth { + username, + password: None, + }), + (_, auth) => auth, + }; + let timeout = query.timeout.map_or_else( + || options::seconds(options.socket_timeout.as_ref()), + |timeout| Ok(Some(timeout)), + )?; + let connect_timeout = query.connect_timeout.map_or_else( + || options::seconds(options.socket_connect_timeout.as_ref()), + |timeout| Ok(Some(timeout)), + )?; + let health_check_interval = query.health_check_interval.map_or_else( + || options::integer::(options.health_check_interval.as_ref()), + |interval| Ok(Some(interval)), + )?; + Ok(Resolution { + target, + auth, + namespace: options::text(options.namespace.as_ref()), + timeout, + connect_timeout, + health_check: health_check(health_check_interval), + keepalive: options + .socket_keepalive + .as_ref() + .map_or(Ok(true), |entry| options::flag(Some(entry)))?, + ignored: ignored.into_iter().chain(query.ignored).collect(), + }) +} + +// redis-py checks health only when the interval is truthy, so zero turns it off and a negative interval checks +// before every command +fn health_check(seconds: Option) -> HealthCheck { + match seconds { + None => HealthCheck::Unset, + Some(0) => HealthCheck::Disabled, + Some(seconds) => HealthCheck::Every(Duration::from_secs(seconds.try_into().unwrap_or(0))), + } +} diff --git a/litellm-rust/crates/redis-resolve/src/options.rs b/litellm-rust/crates/redis-resolve/src/options.rs new file mode 100644 index 00000000000..f2539e46ca5 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/options.rs @@ -0,0 +1,281 @@ +use std::time::Duration; + +use litellm_python_compat::{BigInt, number}; +use litellm_redis_types::Ignored; +use litellm_secret_value::SecretValue; +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, Kwargs, + keys::{self, Asked}, +}; + +/// Every setting the resolver reads. A setting is `None` when it is unset, null or empty, except that a null +/// `tls.cert_reqs` is kept because redis-py reads `ssl_cert_reqs=None` as `CERT_NONE`. +pub(super) struct Options { + pub(super) url: Option, + pub(super) host: Option, + pub(super) port: Option, + pub(super) db: Option, + pub(super) unix_socket_path: Option, + pub(super) startup_nodes: Option, + pub(super) sentinel: SentinelOptions, + pub(super) username: Option, + pub(super) password: Option, + pub(super) credential_provider: Option, + pub(super) gcp: GcpOptions, + pub(super) azure: AzureOptions, + pub(super) aws: AwsOptions, + pub(super) ssl: Option, + pub(super) tls: TlsInput, + pub(super) namespace: Option, + pub(super) socket_timeout: Option, + pub(super) socket_connect_timeout: Option, + pub(super) health_check_interval: Option, + pub(super) socket_keepalive: Option, +} + +pub(super) struct SentinelOptions { + pub(super) nodes: Option, + pub(super) service_name: Option, + pub(super) password: Option, +} + +pub(super) struct GcpOptions { + pub(super) service_account: Option, + pub(super) ssl_ca_certs: Option, +} + +pub(super) struct AzureOptions { + pub(super) ad_token: Option, + pub(super) client_id: Option, + pub(super) tenant_id: Option, + pub(super) client_secret: Option, +} + +pub(super) struct AwsOptions { + pub(super) auth: Option, + pub(super) user_name: Option, + pub(super) cache_name: Option, + pub(super) region: Option, + pub(super) serverless: Option, +} + +/// The TLS settings, read from the kwargs or from a `rediss://` url's query. +pub(super) struct TlsInput { + pub(super) cert_reqs: Option, + pub(super) ca_certs: Option, + pub(super) ca_data: Option, + pub(super) ca_path: Option, + pub(super) certfile: Option, + pub(super) keyfile: Option, + pub(super) key_password: Option, + pub(super) check_hostname: Option, +} + +/// The settings the kwargs set that the Rust client drops, once none of them asks for a check it cannot perform. +fn ignored(kwargs: &Kwargs) -> Result, Error> { + refuse_unenforceable(kwargs.iter())?; + Ok(kwargs + .iter() + .filter_map(|(key, entry)| keys::ignored(&entry.source, key, &entry.value)) + .collect()) +} + +impl Options { + /// The typed settings and the ones the Rust client drops. + pub(super) fn read(kwargs: &Kwargs) -> Result<(Self, Vec), Error> { + let ignored = ignored(kwargs)?; + let options = Self::from_lookup(|key| kwargs.get(key).cloned()); + Ok((options, ignored)) + } + + fn from_lookup(lookup: impl Fn(&'static str) -> Option) -> Self { + let get = |key| lookup(key).filter(|entry| is_present(&entry.value)); + Self { + url: get("url"), + host: get("host"), + port: get("port"), + db: get("db"), + unix_socket_path: get("unix_socket_path"), + startup_nodes: get("startup_nodes"), + sentinel: SentinelOptions { + nodes: get("sentinel_nodes"), + service_name: get("service_name"), + password: get("sentinel_password"), + }, + username: get("username"), + password: get("password"), + credential_provider: get("credential_provider"), + gcp: GcpOptions { + service_account: get("gcp_service_account"), + ssl_ca_certs: get("gcp_ssl_ca_certs"), + }, + azure: AzureOptions { + ad_token: get("azure_redis_ad_token"), + client_id: get("azure_client_id"), + tenant_id: get("azure_tenant_id"), + client_secret: get("azure_client_secret"), + }, + aws: AwsOptions { + auth: get("aws_iam_auth"), + user_name: get("aws_iam_user_name"), + cache_name: get("aws_iam_cache_name"), + region: get("aws_iam_region"), + serverless: get("aws_iam_serverless"), + }, + ssl: get("ssl"), + tls: TlsInput { + cert_reqs: lookup("ssl_cert_reqs") + .filter(|entry| entry.value.is_null() || is_present(&entry.value)), + ca_certs: get("ssl_ca_certs"), + ca_data: get("ssl_ca_data"), + ca_path: get("ssl_ca_path"), + certfile: get("ssl_certfile"), + keyfile: get("ssl_keyfile"), + key_password: get("ssl_password"), + check_hostname: get("ssl_check_hostname"), + }, + namespace: get("namespace"), + socket_timeout: get("socket_timeout"), + socket_connect_timeout: get("socket_connect_timeout"), + health_check_interval: get("health_check_interval"), + socket_keepalive: get("socket_keepalive"), + } + } +} + +/// A `Reject` row whose value asks for its check fails, because Rust cannot perform it. +pub(super) fn refuse_unenforceable<'a>( + entries: impl IntoIterator, +) -> Result<(), Error> { + entries + .into_iter() + .filter_map(|(key, entry)| Some((keys::refusal(key)?, entry))) + .try_for_each(|((check, asked), entry)| { + let asks = match asked { + Asked::WhenTrue => flag(Some(entry))?, + Asked::WhenSet => is_present(&entry.value), + }; + match asks { + true => Err(Error::Unenforceable { + setting: entry.source.clone(), + check, + }), + false => Ok(()), + } + }) +} + +// Python reads most of these with `x or fallback`, so an empty value counts as unset everywhere +pub fn is_present(value: &Value) -> bool { + match value { + Value::Null => false, + Value::String(text) => !text.is_empty(), + Value::Array(items) => !items.is_empty(), + Value::Object(entries) => !entries.is_empty(), + Value::Bool(_) | Value::Number(_) => true, + } +} + +pub(super) fn invalid(setting: &str, expected: Expected) -> Error { + Error::Invalid { + setting: setting.to_owned(), + expected, + } +} + +pub(super) fn as_text(value: &Value) -> Option { + match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} + +pub(super) fn text(entry: Option<&Entry>) -> Option { + entry.and_then(|entry| as_text(&entry.value)) +} + +pub(super) fn secret(entry: Option<&Entry>) -> Option { + text(entry).map(SecretValue::new) +} + +/// A flag every Python surface reads the same way. `_str_to_bool` takes only `true`, `1` and `yes` +/// as true while the proxy's pydantic models also take `on`, `t` and `y`, so text they disagree on, +/// or text neither accepts, is an error instead of a guess. +pub(super) fn flag(entry: Option<&Entry>) -> Result { + let Some(entry) = entry else { + return Ok(false); + }; + match &entry.value { + Value::Bool(flag) => Ok(*flag), + Value::Number(number) => Ok(number.as_f64().is_some_and(|number| number != 0.0)), + Value::String(text) => match text.to_ascii_lowercase().as_str() { + "true" | "1" | "yes" => Ok(true), + "false" | "0" | "no" | "off" | "f" | "n" => Ok(false), + _ => Err(invalid(&entry.source, Expected::Boolean)), + }, + _ => Err(invalid(&entry.source, Expected::Boolean)), + } +} + +pub(super) fn integer>(entry: Option<&Entry>) -> Result, Error> { + coerced(entry, Expected::Integer, number::int, |number| { + T::try_from(number).ok() + }) +} + +pub(super) fn seconds(entry: Option<&Entry>) -> Result, Error> { + coerced(entry, Expected::Seconds, number::float, |seconds| { + Duration::try_from_secs_f64(seconds).ok() + }) +} + +fn coerced( + entry: Option<&Entry>, + expected: Expected, + parse: impl Fn(&str) -> Option, + convert: impl Fn(N) -> Option, +) -> Result, Error> { + let Some((entry, text)) = entry.and_then(|entry| Some((entry, as_text(&entry.value)?))) else { + return Ok(None); + }; + match parse(&text) { + Some(number) => convert(number) + .map(Some) + .ok_or_else(|| invalid(&entry.source, expected)), + None => { + tracing::warn!(setting = %entry.source, "Redis setting is not a number, so the default is used like redis-py"); + Ok(None) + } + } +} + +#[cfg(test)] +mod tests { + use std::{cell::RefCell, collections::BTreeSet}; + + use rstest::rstest; + + use super::Options; + use crate::SETTINGS; + + // a Read row that Options never asks for would be neither modeled nor passed through nor warned about + #[rstest] + fn options_ask_for_exactly_the_read_rows() { + let asked = RefCell::new(BTreeSet::new()); + let _ = Options::from_lookup(|kwarg| { + asked.borrow_mut().insert(kwarg); + None + }); + + let read: BTreeSet<&str> = SETTINGS + .iter() + .filter(|setting| setting.is_read()) + .map(|setting| setting.kwarg) + .collect(); + assert_eq!(asked.into_inner(), read); + } +} diff --git a/litellm-rust/crates/redis-resolve/src/tls.rs b/litellm-rust/crates/redis-resolve/src/tls.rs new file mode 100644 index 00000000000..a46a6ccc650 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/tls.rs @@ -0,0 +1,189 @@ +use std::time::Duration; + +use litellm_python_compat::number; +use litellm_redis_types::{Ignored, TlsOptions}; +use litellm_secret_value::SecretValue; + +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, + keys::{self, FromUrl}, + options::{self, TlsInput, flag, invalid, secret, text}, +}; + +/// `gcp_ca_certs` replaces `ssl_ca_certs` when GCP IAM is on, as `_get_redis_client_logic` does. +pub(super) fn tls_options( + input: &TlsInput, + gcp_ca_certs: Option<&Entry>, +) -> Result { + Ok(TlsOptions { + ca_certs: text(gcp_ca_certs.or(input.ca_certs.as_ref())).map(Into::into), + ca_data: text(input.ca_data.as_ref()), + ca_path: text(input.ca_path.as_ref()).map(Into::into), + certfile: text(input.certfile.as_ref()).map(Into::into), + keyfile: text(input.keyfile.as_ref()).map(Into::into), + key_password: secret(input.key_password.as_ref()), + verify_certificate: verify_certificate(input.cert_reqs.as_ref())?, + check_hostname: flag(input.check_hostname.as_ref())?, + }) +} + +// RedisSSLContext looks text up case-sensitively and takes an ssl.VerifyMode as is +fn verify_certificate(cert_reqs: Option<&Entry>) -> Result { + let Some(entry) = cert_reqs else { + return Ok(true); + }; + match &entry.value { + // Explicit None intentionally means CERT_NONE, matching redis-py rather than the unset default. + Value::Null => Ok(false), + Value::String(text) if text == "none" => Ok(false), + Value::String(text) if text == "optional" || text == "required" => Ok(true), + Value::Number(mode) => match mode.as_i64() { + Some(0) => Ok(false), + Some(1 | 2) => Ok(true), + _ => Err(invalid(&entry.source, Expected::CertRequirement)), + }, + _ => Err(invalid(&entry.source, Expected::CertRequirement)), + } +} + +/// What redis-py's `parse_url` reads from a url's query on top of the url itself. +pub(super) struct UrlQuery { + pub(super) tls: Option, + pub(super) db: Option, + pub(super) username: Option, + pub(super) password: Option, + pub(super) settings: QuerySettings, +} + +/// The connection settings a url's query sets, which win over the kwargs because `from_url` applies them last. +#[derive(Default)] +pub(super) struct QuerySettings { + pub(super) timeout: Option, + pub(super) connect_timeout: Option, + pub(super) health_check_interval: Option, + pub(super) ignored: Vec, +} + +pub(super) fn url_query(url: &str) -> Result { + let Ok(url) = url::Url::parse(url) else { + return Ok(UrlQuery { + tls: None, + db: None, + username: None, + password: None, + settings: QuerySettings::default(), + }); + }; + let tls_scheme = url.scheme().eq_ignore_ascii_case("rediss"); + let pairs: Vec<(String, Entry)> = url + .query_pairs() + .filter(|(_, value)| !value.is_empty()) + .map(|(name, value)| { + ( + name.to_string(), + Entry { + value: Value::String(value.into_owned()), + source: format!("url ?{name}"), + }, + ) + }) + .collect(); + // parse_qs keeps every value of a repeated parameter and parse_url reads the first + let query: Vec<(String, Entry)> = pairs + .iter() + .enumerate() + .filter(|(index, (name, _))| !pairs[..*index].iter().any(|(earlier, _)| earlier == name)) + .map(|(_, pair)| pair.clone()) + .collect(); + options::refuse_unenforceable(query.iter().map(|(name, entry)| (name, entry)))?; + let (query, ignored): (Vec<_>, Vec<_>) = query.into_iter().partition(|(name, _)| { + keys::setting(name).is_some_and(|setting| match setting.from_url { + FromUrl::Any => true, + FromUrl::Rediss => tls_scheme, + FromUrl::No => false, + }) + }); + let get = |key: &str| { + query + .iter() + .find(|(name, _)| name == key) + .map(|(_, entry)| entry.clone()) + }; + let seconds = |key: &str| { + query_number( + get(key).as_ref(), + number::float, + Expected::Seconds, + |seconds| Duration::try_from_secs_f64(seconds).ok(), + ) + }; + let integer = |key: &str| { + query_number( + get(key).as_ref(), + number::int, + Expected::Integer, + |number| i64::try_from(number).ok(), + ) + }; + let db = integer("db")?; + let settings = QuerySettings { + timeout: seconds("socket_timeout")?, + connect_timeout: seconds("socket_connect_timeout")?, + health_check_interval: integer("health_check_interval")?, + ignored: ignored + .iter() + .filter_map(|(name, entry)| keys::ignored_query(&entry.source, name, &entry.value)) + .collect(), + }; + let tls = match tls_scheme { + true => Some(TlsOptions { + ca_certs: text(get("ssl_ca_certs").as_ref()).map(Into::into), + ca_data: text(get("ssl_ca_data").as_ref()), + ca_path: text(get("ssl_ca_path").as_ref()).map(Into::into), + certfile: text(get("ssl_certfile").as_ref()).map(Into::into), + keyfile: text(get("ssl_keyfile").as_ref()).map(Into::into), + key_password: secret(get("ssl_password").as_ref()), + verify_certificate: verify_certificate(get("ssl_cert_reqs").as_ref())?, + check_hostname: text(get("ssl_check_hostname").as_ref()) + .is_some_and(|text| url_bool(&text)), + }), + false => None, + }; + Ok(UrlQuery { + tls, + db, + username: secret(get("username").as_ref()), + password: secret(get("password").as_ref()), + settings, + }) +} + +// parse_url reads these with int() or float() and raises ValueError building the client on anything else +fn query_number( + entry: Option<&Entry>, + parse: impl Fn(&str) -> Option, + expected: Expected, + convert: impl Fn(N) -> Option, +) -> Result, Error> { + let Some(entry) = entry else { + return Ok(None); + }; + let number = text(Some(entry)) + .and_then(|text| parse(&text)) + .ok_or_else(|| Error::UrlQuery { + setting: entry.source.clone(), + })?; + convert(number) + .map(Some) + .ok_or_else(|| invalid(&entry.source, expected)) +} + +// redis-py's to_bool: only 0, F, FALSE, N and NO (any case) are false +fn url_bool(text: &str) -> bool { + !matches!( + text.to_ascii_uppercase().as_str(), + "" | "0" | "F" | "FALSE" | "N" | "NO" + ) +} diff --git a/litellm-rust/crates/redis-resolve/tests/catalog.rs b/litellm-rust/crates/redis-resolve/tests/catalog.rs new file mode 100644 index 00000000000..b26ad0eebfb --- /dev/null +++ b/litellm-rust/crates/redis-resolve/tests/catalog.rs @@ -0,0 +1,129 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_redis_resolve::{Asked, Error, Handling, SETTINGS, Setting, interpret, kwargs}; +use rstest::rstest; +use serde::Deserialize; +use serde_json::json; + +#[derive(Deserialize)] +struct Catalog { + meta: Meta, + args: BTreeMap, +} + +#[derive(Deserialize)] +struct Meta { + redis_py_locked: String, +} + +#[derive(Deserialize)] +struct Arg { + environment: bool, + redis_py: BTreeMap>, +} + +// the signatures of the clients the proxy builds, where an unknown kwarg raises TypeError +const PROXY_SIGNATURES: &[&str] = &["async", "cluster", "connection"]; + +fn rejected() -> impl Iterator { + SETTINGS + .iter() + .filter_map(|setting| match setting.handling { + Handling::Reject { asked, .. } => Some((setting, asked)), + Handling::Read | Handling::Warn { .. } => None, + }) +} + +fn catalog() -> Catalog { + serde_json::from_str(include_str!("../redis-py-kwargs.json")).unwrap() +} + +fn host_with(kwarg: &str, value: serde_json::Value) -> litellm_redis_resolve::Kwargs { + kwargs([ + ("host".to_owned(), json!("redis")), + (kwarg.to_owned(), value), + ]) +} + +#[rstest] +fn every_catalogued_kwarg_has_exactly_one_settings_row() { + let rows: Vec<&str> = SETTINGS.iter().map(|setting| setting.kwarg).collect(); + let unique: BTreeSet<&str> = rows.iter().copied().collect(); + assert_eq!(rows.len(), unique.len(), "a kwarg has two SETTINGS rows"); + + let catalog = catalog(); + let catalogued: BTreeSet<&str> = catalog.args.keys().map(String::as_str).collect(); + assert_eq!(unique, catalogued); +} + +#[rstest] +fn settings_read_redis_environment_names_exactly_where_python_does() { + let catalog = catalog(); + + let differing: Vec<&str> = SETTINGS + .iter() + .filter(|setting| { + catalog + .args + .get(setting.kwarg) + .is_some_and(|arg| arg.environment != setting.environment) + }) + .map(|setting| setting.kwarg) + .collect(); + assert!(differing.is_empty(), "{differing:?}"); +} + +#[rstest] +fn every_warned_setting_is_reported() { + let warned = SETTINGS + .iter() + .filter(|setting| matches!(setting.handling, Handling::Warn { met: None, .. })); + + for setting in warned { + let resolved = interpret(&host_with(setting.kwarg, json!("set")), None, None).unwrap(); + + let reported: Vec<&str> = resolved + .ignored + .iter() + .map(|ignored| ignored.setting.as_str()) + .collect(); + assert_eq!(reported, vec![setting.kwarg]); + } +} + +#[rstest] +fn every_rejected_setting_fails_when_asked_for_and_not_otherwise() { + for (setting, asked) in rejected() { + let (asking, not_asking) = match asked { + Asked::WhenTrue => (json!(true), json!(false)), + Asked::WhenSet => (json!(["VERIFY_CRL_CHECK_LEAF"]), json!([])), + }; + + let error = interpret(&host_with(setting.kwarg, asking), None, None).unwrap_err(); + assert!( + matches!(&error, Error::Unenforceable { setting: source, .. } if source == setting.kwarg), + "{error:?}" + ); + assert!(interpret(&host_with(setting.kwarg, not_asking), None, None).is_ok()); + } +} + +#[rstest] +fn python_already_refuses_every_rejected_setting() { + let catalog = catalog(); + + let accepted: Vec<&str> = rejected() + .map(|(setting, _)| setting.kwarg) + .filter(|kwarg| { + catalog.args[*kwarg] + .redis_py + .get(&catalog.meta.redis_py_locked) + .is_some_and(|signatures| { + signatures + .iter() + .any(|signature| PROXY_SIGNATURES.contains(&signature.as_str())) + }) + }) + .collect(); + assert!(accepted.is_empty(), "{accepted:?}"); +} diff --git a/litellm-rust/crates/redis-resolve/tests/interpret.rs b/litellm-rust/crates/redis-resolve/tests/interpret.rs new file mode 100644 index 00000000000..fd59ec2acb1 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/tests/interpret.rs @@ -0,0 +1,790 @@ +use std::collections::BTreeMap; + +use litellm_redis_resolve::{Entry, Error, Expected, Kwargs, PythonException, interpret, kwargs}; +use litellm_redis_types::{ + Auth, HealthCheck, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions, +}; +use rstest::rstest; +use serde_json::json; + +use Outcome::{Fails, Refuses, Resolves}; +use PythonException::{RedisClusterException, TypeError, ValueError}; + +const PROVIDER: ProviderHandle = ProviderHandle::new(7); + +fn merged(fields: serde_json::Value, sources: &[(&str, &str)]) -> Kwargs { + let fields = serde_json::from_value::>(fields).unwrap(); + kwargs(fields) + .into_iter() + .map(|(kwarg, entry)| { + let source = sources + .iter() + .find(|(named, _)| *named == kwarg) + .map_or(entry.source, |(_, source)| (*source).to_owned()); + (kwarg, Entry { source, ..entry }) + }) + .collect() +} + +// a host hands its credential_provider object over as a handle, so the cases do the same +fn interpret_merged( + fields: serde_json::Value, + sources: &[(&str, &str)], + redis_username: Option<&str>, +) -> Result { + let has_provider = fields + .get("credential_provider") + .is_some_and(|provider| !provider.is_null()); + interpret( + &merged(fields, sources), + has_provider.then_some(PROVIDER), + redis_username, + ) +} + +fn interpret_kwargs(fields: serde_json::Value) -> Result { + interpret_merged(fields, &[], None) +} + +/// What a contract case expects of the resolver, next to what its Python mirror expects of redis-py. +enum Outcome { + Resolves(serde_json::Value), + /// Rust refuses a config Python connects with, a disagreement on purpose. + Refuses(&'static str), + /// redis-py 5.3.1 already raises on the config, building the client or its first connection, so failing + /// here breaks no deployment that works today. The class is what it raises building the client, and `None` + /// when it only raises once a command connects. + Fails(&'static str, Option), +} + +#[rstest] +fn a_provider_handle_beats_every_cloud_auth_and_static_credentials() { + let resolved = interpret_kwargs(json!({ + "host": "redis", "password": "pass", "credential_provider": "host-object", + "gcp_service_account": "account", "aws_iam_auth": true, "azure_redis_ad_token": "true", + })) + .unwrap(); + + assert_eq!(resolved.auth, Auth::Provider(PROVIDER)); +} + +#[rstest] +fn a_rediss_url_reads_ssl_password_from_its_query() { + let resolved = + interpret_kwargs(json!({"url": "rediss://cache:6380/0?ssl_password=query-key-secret"})) + .unwrap(); + + let Target::Url(parsed) = resolved.target else { + panic!("expected a url target, got {:?}", resolved.target); + }; + let key_password = parsed.tls.and_then(|tls| tls.key_password); + assert_eq!(expose(key_password.as_ref()), Some("query-key-secret")); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +#[case::read_on_a_plain_url( + "redis://cache:6379/1?db=3&username=query-user&password=query-secret", + Some(3), + Some("query-user"), + Some("query-secret") +)] +#[case::absent("rediss://cache:6380/1", None, None, None)] +fn a_url_query_carries_db_and_credentials_like_parse_url( + #[case] url: &str, + #[case] db: Option, + #[case] username: Option<&str>, + #[case] password: Option<&str>, +) { + let resolved = interpret_kwargs(json!({"url": url})).unwrap(); + + let Target::Url(parsed) = resolved.target else { + panic!("expected a url target, got {:?}", resolved.target); + }; + assert_eq!(parsed.db, db); + assert_eq!(expose(parsed.username.as_ref()), username); + assert_eq!(expose(parsed.password.as_ref()), password); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +fn a_url_query_db_that_is_not_an_integer_fails_like_parse_url() { + let error = interpret_kwargs(json!({"url": "redis://cache:6379?db=first"})).unwrap_err(); + + assert!(matches!(error, Error::UrlQuery { setting } if setting == "url ?db")); +} + +#[rstest] +#[case::tls_1_2_is_already_met(json!("TLSv1_2"), false)] +#[case::the_enum_name(json!("TLSVersion.TLSv1_2"), false)] +#[case::the_enum_value(json!(771), false)] +#[case::tls_1_3_is_not_enforced(json!("TLSv1_3"), true)] +fn ssl_min_version_warns_only_when_rustls_cannot_meet_it( + #[case] version: serde_json::Value, + #[case] warns: bool, +) { + let resolved = + interpret_kwargs(json!({"host": "redis", "ssl": true, "ssl_min_version": version})) + .unwrap(); + + assert_eq!( + !resolved.ignored.is_empty(), + warns, + "{:?}", + resolved.ignored + ); +} + +#[rstest] +#[case::unset(json!({"host": "redis"}), true)] +#[case::empty_like_unset(json!({"host": "redis", "socket_keepalive": ""}), true)] +#[case::enabled(json!({"host": "redis", "socket_keepalive": "yes"}), true)] +#[case::disabled(json!({"host": "redis", "socket_keepalive": false}), false)] +#[case::disabled_as_text(json!({"host": "redis", "socket_keepalive": "0"}), false)] +#[case::disabled_on_a_cluster(json!({"startup_nodes": [{"host": "node-a", "port": 7000}], "socket_keepalive": "false"}), false)] +fn keepalive_stays_on_unless_socket_keepalive_turns_it_off( + #[case] kwargs: serde_json::Value, + #[case] keepalive: bool, +) { + let resolved = interpret_kwargs(kwargs).unwrap(); + + assert_eq!(resolved.keepalive, keepalive); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +fn a_socket_keepalive_no_python_surface_agrees_on_is_invalid() { + let error = interpret_kwargs(json!({"host": "redis", "socket_keepalive": "on"})).unwrap_err(); + + assert!( + matches!(&error, Error::Invalid { setting, expected: Expected::Boolean } if setting == "socket_keepalive"), + "{error:?}" + ); +} + +fn tls_view(tls: Option<&TlsOptions>) -> serde_json::Value { + tls.map_or(serde_json::Value::Null, |tls| { + serde_json::json!({ + "ca_certs": tls.ca_certs, + "ca_data": tls.ca_data, + "ca_path": tls.ca_path, + "certfile": tls.certfile, + "keyfile": tls.keyfile, + "verify_certificate": tls.verify_certificate, + "check_hostname": tls.check_hostname, + }) + }) +} + +fn nodes_view(nodes: &[Node]) -> Vec { + nodes + .iter() + .map(|node| format!("{}:{}", node.host, node.port)) + .collect() +} + +fn expose(value: Option<&litellm_secret_value::SecretValue>) -> Option<&str> { + value.map(litellm_secret_value::SecretValue::expose) +} + +fn view(resolved: &Resolution) -> serde_json::Value { + let target = match &resolved.target { + Target::Unix(socket) => serde_json::json!({ + "mode": "unix", "path": socket.path, "tls": null, + }), + Target::Standalone(host) => serde_json::json!({ + "mode": "host", + "host": host.host, + "port": host.port, + "db": host.db, + "tls": tls_view(host.tls.as_ref()), + }), + Target::Url(url) => serde_json::json!({ + "mode": "url", + "url": url.url.expose(), + "tls": tls_view(url.tls.as_ref()), + }), + Target::Cluster(cluster) => serde_json::json!({ + "mode": "cluster", + "nodes": nodes_view(&cluster.nodes), + "tls": tls_view(cluster.tls.as_ref()), + }), + Target::Sentinel(sentinel) => serde_json::json!({ + "mode": "sentinel", + "nodes": nodes_view(&sentinel.monitors), + "service_name": sentinel.service_name, + "sentinel_password": expose(sentinel.monitor_auth.as_ref().and_then(|auth| auth.password.as_ref())), + "db": sentinel.db, + "tls": tls_view(sentinel.tls.as_ref()), + }), + }; + let auth = match &resolved.auth { + Auth::Static(StaticAuth { username, password }) => serde_json::json!({ + "kind": "static", + "username": expose(username.as_ref()), + "password": expose(password.as_ref()), + }), + Auth::Provider(_) => serde_json::json!({"kind": "host"}), + Auth::GcpIam { service_account } => { + serde_json::json!({"kind": "gcp_iam", "service_account": service_account}) + } + Auth::AzureAd(azure) => { + serde_json::json!({"kind": "azure_ad", "username": azure.username}) + } + Auth::AwsIam(aws) => serde_json::json!({ + "kind": "aws_iam", + "user_name": aws.user_name, + "cache_name": aws.cache_name, + "region": aws.region, + "serverless": aws.serverless, + }), + }; + let serde_json::Value::Object(target) = target else { + unreachable!("every target view is an object") + }; + let rest = serde_json::json!({ + "auth": auth, + "socket_timeout": resolved.timeout.map_or(5.0, |timeout| timeout.as_secs_f64()), + "socket_connect_timeout": resolved.connect_timeout.or(resolved.timeout).map_or(5.0, |timeout| timeout.as_secs_f64()), + "health_check_interval": match (&resolved.target, resolved.health_check) { + (_, HealthCheck::Every(interval)) => Some(interval.as_secs_f64()), + (Target::Cluster(_), HealthCheck::Unset) => Some(25.0), + (_, HealthCheck::Unset | HealthCheck::Disabled) => None, + }, + }); + let serde_json::Value::Object(rest) = rest else { + unreachable!("the view is an object") + }; + serde_json::Value::Object(target.into_iter().chain(rest).collect()) +} + +fn expected_view(fields: serde_json::Value) -> serde_json::Value { + let serde_json::Value::Object(fields) = fields else { + unreachable!("expected views are objects") + }; + let defaults = json!({ + "auth": {"kind": "static", "username": null, "password": null}, + "tls": null, + "socket_timeout": 5.0, + "socket_connect_timeout": 5.0, + "health_check_interval": null, + }); + let serde_json::Value::Object(defaults) = defaults else { + unreachable!("the defaults are an object") + }; + serde_json::Value::Object(defaults.into_iter().chain(fields).collect()) +} + +#[rstest] +#[case::host_defaults( + json!({"host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::empty_cluster_nodes_use_the_environment( + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::empty_sentinel_nodes_use_the_environment( + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))), +)] +#[case::empty_service_name_uses_the_environment( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[("service_name", "REDIS_SERVICE_NAME")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))), +)] +#[case::plaintext_ignores_tls_requirements( + json!({"host": "redis", "ssl_cert_reqs": "strict"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::an_ambiguous_flag_is_refused( + json!({"host": "redis", "ssl": "on"}), + &[], + None, + Refuses("invalid Redis setting ssl: expected a boolean"), +)] +#[case::empty_sentinel_nodes_without_the_environment_name_no_target( + json!({"sentinel_nodes": [], "service_name": "primary"}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +#[case::ssl_ca_path_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_ca_path": "/ca"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/ca", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::ssl_validate_ocsp_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), + &[], + None, + Fails("ssl_validate_ocsp asks for OCSP revocation checks, which rustls cannot perform, so Redis is not connected without them", Some(TypeError)), +)] +#[case::ssl_include_verify_flags_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_include_verify_flags": ["VERIFY_CRL_CHECK_LEAF"]}), + &[], + None, + Fails("ssl_include_verify_flags asks for X.509 verify flags such as CRL checking, which rustls cannot perform, so Redis is not connected without them", Some(TypeError)), +)] +#[case::ssl_password_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_password": "key-secret"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::host_with_every_basic_setting( + json!({"db": 2, "host": "redis", "password": "pass", "port": 6380, "socket_timeout": 1.5, "ssl": true, "username": "user"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 2, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": "pass"}, "socket_timeout": 1.5, "socket_connect_timeout": 1.5}))), +)] +#[case::environment_host_port_and_password( + json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}), + &[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "environment", "port": 6390, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": null, "password": "environment-secret"}}))), +)] +#[case::string_numbers_are_coerced( + json!({"db": "3", "host": "redis", "port": "6380", "socket_timeout": "2.5"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 3, "socket_timeout": 2.5, "socket_connect_timeout": 2.5}))), +)] +#[case::padded_and_underscored_numbers_parse_like_python_int_and_float( + json!({"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}), + &[], + None, + Resolves(expected_view(json!({"db": 2, "host": "redis", "mode": "host", "port": 6380, "socket_timeout": 10.5, "socket_connect_timeout": 10.5}))), +)] +#[case::padded_environment_numbers_parse_like_python_int( + json!({"port": "6_380", "host": "redis", "db": "3\n"}), + &[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")], + None, + Resolves(expected_view(json!({"db": 3, "host": "redis", "mode": "host", "port": 6380}))), +)] +#[case::socket_timeout_from_the_environment( + json!({"socket_timeout": "0.75", "host": "redis"}), + &[("socket_timeout", "REDIS_SOCKET_TIMEOUT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 0.75, "socket_connect_timeout": 0.75}))), +)] +#[case::environment_url( + json!({"url": "redis://environment:6379/1"}), + &[("url", "REDIS_URL")], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://environment:6379/1"}))), +)] +#[case::environment_url_beats_the_environment_host( + json!({"host": "ignored", "url": "redis://from-url:6379"}), + &[("host", "REDIS_HOST"), ("url", "REDIS_URL")], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://from-url:6379"}))), +)] +#[case::an_explicit_host_drops_the_environment_url( + json!({"host": "configured"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "configured", "port": 6379, "db": 0}))), +)] +#[case::a_url_drops_host_db_and_password_and_keeps_the_username( + json!({"db": 4, "host": "dropped", "password": "dropped", "url": "rediss://cache:6380", "username": "user"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": null}}))), +)] +#[case::a_url_ignores_tls_settings_outside_the_url( + json!({"ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true, "url": "rediss://cache:6380"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::startup_nodes_beat_a_url( + json!({"password": "cluster-secret", "startup_nodes": [{"host": "node-a", "port": 7000}], "url": "redis://single:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "auth": {"kind": "static", "username": null, "password": "cluster-secret"}, "health_check_interval": 25.0}))), +)] +#[case::environment_cluster_nodes( + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"], "health_check_interval": 25.0}))), +)] +#[case::cluster_tls_options( + json!({"ssl": "true", "ssl_ca_certs": "/ca.pem", "ssl_check_hostname": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": true}, "health_check_interval": 25.0}))), +)] +#[case::sentinel_from_settings( + json!({"password": "data-secret", "sentinel_nodes": [["sentinel-a", 26379], ["sentinel-b", 26380]], "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379", "sentinel-b:26380"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0, "auth": {"kind": "static", "username": null, "password": "data-secret"}}))), +)] +#[case::sentinel_from_the_environment( + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0}))), +)] +#[case::a_host_breaks_sentinel( + json!({"host": "ignored", "sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[], + None, + Fails("host cannot be combined with sentinel_nodes, which redis-py rejects at startup", Some(TypeError)), +)] +#[case::a_url_beats_sentinels( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary", "url": "redis://single:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://single:6379"}))), +)] +#[case::an_explicit_credential_provider_beats_cloud_auth( + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "host"}}))), +)] +#[case::gcp_beats_azure_and_aws( + json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "gcp_iam", "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}}))), +)] +#[case::gcp_ca_certs_replace_the_configured_ones_under_ssl( + json!({"gcp_service_account": "account", "gcp_ssl_ca_certs": "/gcp-ca.pem", "host": "redis", "ssl": true, "ssl_ca_certs": "/configured-ca.pem"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/gcp-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "gcp_iam", "service_account": "account"}}))), +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[], + Some("environment-user"), + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": "environment-user"}}))), +)] +#[case::an_empty_environment_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}), + &[("username", "REDIS_USERNAME")], + Some(""), + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))), +)] +#[case::azure_without_a_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))), +)] +#[case::azure_needs_the_literal_true( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": "1", "azure_tenant_id": "tenant", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::aws_iam_over_tls_with_the_default_region( + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}), + &[("aws_iam_region", "AWS_DEFAULT_REGION")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "us-east-1", "serverless": false}}))), +)] +#[case::aws_region_beats_the_default_region( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}), + &[("aws_iam_region", "AWS_REGION")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-central-1", "serverless": false}}))), +)] +#[case::aws_iam_serverless_over_a_rediss_url_drops_its_userinfo( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "eu-west-1", "aws_iam_serverless": true, "aws_iam_user_name": "iam-user", "url": "rediss://user:pass@cache-1:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache-1:6379", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-west-1", "serverless": true}}))), +)] +#[case::aws_iam_without_tls_fails( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "us-east-1", "aws_iam_user_name": "iam-user", "host": "cache-1"}), + &[], + None, + Fails("AWS ElastiCache IAM Redis authentication requires TLS", Some(ValueError)), +)] +#[case::aws_iam_without_a_region_fails( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}), + &[], + None, + Fails("AWS ElastiCache IAM Redis authentication requires: aws_iam_region", Some(ValueError)), +)] +#[case::tls_options_from_settings( + json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))), +)] +#[case::a_null_cert_requirement_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": false}}))), +)] +#[case::a_verify_mode_of_none_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": 0}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": false}}))), +)] +#[case::a_verify_mode_of_optional_verifies( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": 1}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::a_cert_requirement_spelled_like_python_none_fails( + json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], + None, + Fails("invalid Redis setting REDIS_SSL_CERT_REQS: expected none, optional, required, null or an ssl.VerifyMode value", None), +)] +#[case::tls_options_from_the_environment( + json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/env-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::connect_timeout_is_separate_from_the_socket_timeout( + json!({"host": "redis", "socket_connect_timeout": 0.25, "socket_timeout": 2.5}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 2.5, "socket_connect_timeout": 0.25}))), +)] +#[case::environment_references_resolve( + json!({"host": "referenced", "password": "referenced-secret"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "referenced", "port": 6379, "db": 0, "auth": {"kind": "static", "username": null, "password": "referenced-secret"}}))), +)] +#[case::no_target_fails( + json!({"port": 6379}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +#[case::client_certificates( + json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": "/client.key", "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::a_rediss_url_takes_tls_options_from_its_query( + json!({"url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes", "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))), +)] +#[case::a_url_query_carries_socket_timeout( + json!({"url": "redis://cache:6379/0?socket_timeout=2"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2", "socket_timeout": 2.0, "socket_connect_timeout": 2.0}))), +)] +#[case::url_query_settings_win_over_the_kwargs( + json!({"health_check_interval": 4, "socket_connect_timeout": 7, "socket_timeout": 1, "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3", "socket_timeout": 3.0, "socket_connect_timeout": 0.5, "health_check_interval": 9.0}))), +)] +#[case::a_repeated_url_query_parameter_takes_its_first_value( + json!({"url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9", "socket_timeout": 2.0, "socket_connect_timeout": 2.0}))), +)] +#[case::a_url_query_number_parse_url_cannot_read( + json!({"url": "redis://cache:6379/0?socket_timeout=slow"}), + &[], + None, + Fails("invalid Redis setting url ?socket_timeout: redis-py's parse_url cannot read it as a number", Some(ValueError)), +)] +#[case::tls_settings_outside_a_url_never_fail_it( + json!({"db": "99999999999999999999", "ssl_cert_reqs": "strict", "url": "rediss://cache:6380"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::db_is_not_read_for_a_cluster( + json!({"db": 99_999_999_999_i64, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_true( + json!({"db": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_false( + json!({"db": false, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::cluster_db_fraction( + json!({"db": 0.5, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_float_zero( + json!({"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::cluster_db_nonempty_list( + json!({"db": [1], "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_empty_list( + json!({"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::ssl_ciphers( + json!({"host": "redis", "ssl": true, "ssl_ciphers": "HIGH"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::retry_on_timeout( + json!({"host": "redis", "retry_on_timeout": true}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_port( + json!({"host": "redis", "port": "abc"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_db( + json!({"db": "first", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::socket_timeout_not_a_number( + json!({"host": "redis", "socket_timeout": "slow"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_port_from_the_environment( + json!({"port": "abc", "host": "redis"}), + &[("host", "REDIS_HOST"), ("port", "REDIS_PORT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::protocol_3( + json!({"host": "redis", "protocol": 3}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::retry_on_error_from_the_environment( + json!({"retry_on_error": true, "host": "redis"}), + &[("retry_on_error", "REDIS_RETRY_ON_ERROR")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::unix_socket_path( + json!({"host": "redis", "unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Resolves(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))), +)] +#[case::unix_socket_path_keeps_tls_settings_out_of_the_target( + json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Resolves(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))), +)] +#[case::ssl_ca_data( + json!({"host": "redis", "ssl": true, "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::ssl_ca_path_from_the_environment( + json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/etc/ssl/private-cas", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::health_check_interval_is_whole_seconds( + json!({"health_check_interval": "7", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "health_check_interval": 7.0}))), +)] +#[case::health_check_interval_zero_is_off( + json!({"health_check_interval": 0, "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::health_check_interval_zero_keeps_a_cluster_off( + json!({"health_check_interval": 0, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))), +)] +#[case::a_cluster_rejects_a_nonzero_db( + json!({"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::unix_socket_path_without_a_host_fails_like_python( + json!({"unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +fn resolves_like_python_redis( + #[case] merged: serde_json::Value, + #[case] sources: &[(&str, &str)], + #[case] redis_username: Option<&str>, + #[case] expected: Outcome, +) { + let outcome = interpret_merged(merged, sources, redis_username); + let rendered = outcome + .as_ref() + .map(view) + .map_err(|error| error.to_string()); + + match expected { + Resolves(expected) => assert_eq!(rendered, Ok(expected)), + Refuses(message) => assert_eq!(rendered, Err(message.to_owned())), + Fails(message, raised) => { + assert_eq!(rendered, Err(message.to_owned())); + assert_eq!( + outcome + .err() + .and_then(|error| error.raised_building_the_python_client()), + raised + ); + } + } +} diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md new file mode 100644 index 00000000000..963f7ff168f --- /dev/null +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -0,0 +1,13 @@ +- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`) +- An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on, or `None` when it returns a settlement the caller waits on, like `SpendCounterStore::increment_many_post_call` +- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s +- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec +- A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` +- Every family works with no Redis, answering from memory. `attach` adds the Redis tier at most once and keeps memory, as `DualCache.attach_redis_cache` does, so a host builds the store at startup and never rebuilds it to connect +- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not +- Key suffixes and wire spellings derive from strum on the family's enums. A value the family stores in Redis implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`. An argument or reply is a typed record, never a positional tuple: a frozen Python dataclass read in `python-bridge` by a `#[derive(FromPyObject)]` struct with `#[pyo3(attribute)]` fields and converted into the family's type with `From`, and a reply returned as a frozen `#[pyclass]`. The bridge reads an enum through its strum spelling, never a hand-written match +- Wire contracts are hand-written cases with expected values, inlined as the same literals in this crate's test and in `tests/unit/redis/stores/`, checked here against real Redis and there against Python's own codec. Change both sides together +- `rate_limits.rs` keeps the v3 limiter's Lua byte for byte (`rate_limits/lua.rs`) so pods on either stack share one script SHA; the admission decisions stay in the Python limiter +- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. Each op exists once and takes a `Send` (`Now`, `PreCall`, `PostCall`) picked by the caller, and the store calls `send.batch(&client)`. A store never encodes a phase in a method name. A post-call method exists only where a follow-up runs after the write settles (`token_increment_post_call`), as a thin wrapper over the shared op. `Send::Now` is an alone batch, so a script call outside a request goes out alone, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here +- Every store sends the commands main's `RedisCache` sends for that operation: `MGET` where Python reads several keys at once, `SET ... EX`, and `INCRBYFLOAT` with `EXPIRE` or `TTL`. A `MONITOR` diff against main is a parity check, and any difference is a defect +- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml new file mode 100644 index 00000000000..1dd61d76d3c --- /dev/null +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "litellm-redis-store" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-python-compat.workspace = true +litellm-redis-client.workspace = true +strum.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tracing.workspace = true + +[dev-dependencies] +litellm-circuit-breaker.workspace = true +litellm-redis-testing.workspace = true +litellm-secret-value.workspace = true +litellm-tracing.workspace = true +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/redis-store/src/auth_cache/keys.rs b/litellm-rust/crates/redis-store/src/auth_cache/keys.rs new file mode 100644 index 00000000000..feb87a40b00 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/keys.rs @@ -0,0 +1,8 @@ +/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every +/// other object under a prefixed key. +pub fn is_key_object(key: &str) -> bool { + key.len() == 64 + && key + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/mod.rs b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs new file mode 100644 index 00000000000..b5a83b5aec6 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs @@ -0,0 +1,258 @@ +mod keys; +mod types; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_python_compat::{Value, json}; +use litellm_redis_client::{Declared, RedisBatch, commands}; + +pub use self::keys::is_key_object; +pub use self::types::Settings; +use crate::{ + Error, Number, Pending, RequestRedis, Send, + dual_cache::DualCache, + memory::{Clock, lock, system_clock}, + remote::{Remote, RemoteSlot}, + request::ready, +}; + +/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it. +const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"; + +/// The proxy's auth object cache with `enable_redis_auth_cache` on: `UserApiKeyCache`'s two +/// `DualCache` partitions, so churn in management objects cannot evict user-key objects, in front of +/// one Redis shared by every pod. Keys are the logical cache keys `UserApiKeyCache` callers already +/// use; values are what `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`. +pub struct AuthCache { + key_objects: DualCache, + objects: DualCache, + redis: Arc, + settings: Mutex, +} + +impl AuthCache { + pub fn new(remote: Option, settings: Settings) -> Self { + Self::with_clock(remote, settings, system_clock()) + } + + pub fn with_clock(remote: Option, settings: Settings, now: Clock) -> Self { + let redis = Arc::new(RemoteSlot::new(remote)); + let partition = || DualCache::new(redis.clone(), settings.memory(), now.clone()); + Self { + key_objects: partition(), + objects: partition(), + redis, + settings: Mutex::new(settings), + } + } + + pub fn has_redis(&self) -> bool { + self.redis.get().is_some() + } + + pub fn settings(&self) -> Settings { + *lock(&self.settings) + } + + /// `UserApiKeyCache.update_cache_ttl` and `update_in_memory_max_size` on both partitions, which + /// keep what they hold. + pub fn reconfigure(&self, settings: Settings) { + *lock(&self.settings) = settings; + self.key_objects.reconfigure(settings.memory()); + self.objects.reconfigure(settings.memory()); + } + + pub fn get_local(&self, key: &str) -> Result, Error> { + Ok(self.partition(key).memory().get(key)?) + } + + /// A memory write with `ttl`, or `default_in_memory_ttl` when it names none. + pub fn set_local(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { + let partition = self.partition(key); + Ok(partition + .memory() + .set(key, value, ttl.or(partition.default_in_memory_ttl()))?) + } + + pub fn delete_local(&self, key: &str) -> Result<(), Error> { + Ok(self.partition(key).memory().delete(key)?) + } + + pub fn flush_local(&self) -> Result<(), Error> { + self.key_objects.memory().flush()?; + Ok(self.objects.memory().flush()?) + } + + /// Entries held by the management-object and the user-key partitions. + pub fn len_local(&self) -> Result<(usize, usize), Error> { + Ok(( + self.objects.memory().len()?, + self.key_objects.memory().len()?, + )) + } + + /// `DualCache.async_get_cache`. A key `request` already read as missing from Redis is not read + /// again. + pub async fn get( + &self, + request: Option<&RequestRedis>, + key: &str, + ) -> Result, Error> { + self.partition(key).get(request, key, None).await + } + + /// `DualCache.async_set_cache`, writing `json.dumps(value)`. + pub async fn set(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { + self.partition(key).set(key, value, ttl).await + } + + /// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`. + pub async fn set_str( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), Error> { + self.partition(key).set_str(key, value, ttl).await + } + + /// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks. + pub fn declare_set( + &self, + send: Send<'_>, + key: &str, + value: Value, + ttl: Option, + ) -> Result, Error> { + self.partition(key).declare_set(send, key, value, ttl) + } + + /// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks. + pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { + self.partition(key).declare_delete(send, key) + } + + /// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then + /// the rest in another, each with one TTL. + pub async fn set_many( + &self, + entries: Vec<(String, Value)>, + ttl: Option, + ) -> Result<(), Error> { + let (key_objects, others): (Vec<_>, Vec<_>) = + entries.into_iter().partition(|(key, _)| is_key_object(key)); + self.key_objects.set_pipeline(key_objects, ttl).await?; + self.objects.set_pipeline(others, ttl).await + } + + /// The auth prefetch write-back: memory first, an entry without a TTL taking + /// `default_in_memory_ttl`, then `RedisCache.async_set_cache_pipeline_with_ttls`, which encodes + /// every value before it sends anything and gives an entry without a TTL the connection's + /// default. On a request batch the `SET`s are declared and nothing is left to wait for. + /// [`Send::Now`] is one pipeline of its own rather than one round trip per entry. + pub fn write_back( + &self, + send: Send<'_>, + entries: Vec<(String, Value, Option)>, + ) -> Result, Error> { + let encoded = self + .redis + .get() + .map(|_| { + entries + .iter() + .map(|(_, value, _)| json::dumps(value)) + .collect::, _>>() + }) + .transpose(); + let writes = entries + .into_iter() + .map(|(key, value, ttl)| { + self.set_local(&key, value, ttl)?; + Ok((key, ttl)) + }) + .collect::, Error>>()?; + let (Some(remote), Some(encoded)) = (self.redis.get(), encoded?) else { + return Ok(ready(())); + }; + let sets = writes.iter().zip(encoded).map(|((key, ttl), value)| { + commands::set(key, value, ttl.unwrap_or(remote.default_ttl)) + }); + let Send::Now = send else { + let batch = send.batch(&remote.client); + sets.for_each(|set| drop(batch.add(set))); + return Ok(ready(())); + }; + let pipeline = RedisBatch::pipelined(remote.client.clone()); + let written: Vec<_> = sets.map(|set| pipeline.add(set)).collect(); + Ok(Box::pin(async move { + Declared::all(written) + .await + .into_iter() + .try_for_each(|reply| Ok(reply?)) + })) + } + + /// `UserApiKeyCache.async_delete_cache_keys`: the key-object partition and the rest as two + /// concurrent `DualCache.async_delete_cache_keys`. Both run even when one fails. + pub async fn delete(&self, keys: &[String]) -> Result<(), Error> { + let (key_objects, others): (Vec, Vec) = + keys.iter().cloned().partition(|key| is_key_object(key)); + let (key_objects, others) = tokio::join!( + self.key_objects.delete_keys(&key_objects), + self.objects.delete_keys(&others) + ); + key_objects.and(others) + } + + /// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on + /// the batch `send` picks, noting the keys Redis lacks as the request's misses. `None` without + /// Redis. Any failed read or undecodable value fails the whole read, as a failed `MGET` does. + pub fn get_many(&self, send: Send<'_>, keys: &[String]) -> Option>>> { + let client = self.redis.client()?; + let misses = send.misses(client); + let read = send.batch(client).add(commands::mget::>( + &keys.iter().map(String::as_str).collect::>(), + )); + let keys = keys.to_vec(); + Some(Box::pin(async move { + let raw = read.resolve().await?; + if let Some(misses) = misses { + misses.note( + keys.into_iter() + .zip(&raw) + .filter(|(_, raw)| raw.is_none()) + .map(|(key, _)| key), + ); + } + raw.into_iter().map(|raw| Ok(types::decode(raw)?)).collect() + })) + } + + /// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds. + /// `None` without Redis or for an absent entry. + pub async fn increment_if_present( + &self, + key: &str, + amount: Number, + ) -> Result, Error> { + let Some(client) = self.redis.client() else { + return Ok(None); + }; + let amount = litellm_redis_client::Amount::from(amount); + Ok(client + .send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount)) + .await?) + } + + fn partition(&self, key: &str) -> &DualCache { + if is_key_object(key) { + &self.key_objects + } else { + &self.objects + } + } +} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/types.rs b/litellm-rust/crates/redis-store/src/auth_cache/types.rs new file mode 100644 index 00000000000..e933ac19af9 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/types.rs @@ -0,0 +1,35 @@ +use std::time::Duration; + +use litellm_python_compat::Value; + +use crate::{DecodeError, codec::python_value, dual_cache::MemorySettings}; + +/// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Settings { + pub memory_entries: usize, + pub memory_ttl: Option, +} + +impl Default for Settings { + /// `InMemoryCache`'s 200 entries. + fn default() -> Self { + Self { + memory_entries: 200, + memory_ttl: None, + } + } +} + +impl Settings { + pub(super) fn memory(self) -> MemorySettings { + MemorySettings { + entries: Some(self.memory_entries), + ttl: self.memory_ttl, + } + } +} + +pub(super) fn decode(raw: Option>) -> Result, DecodeError> { + raw.as_deref().map(python_value).transpose() +} diff --git a/litellm-rust/crates/redis-store/src/codec.rs b/litellm-rust/crates/redis-store/src/codec.rs new file mode 100644 index 00000000000..21cde7d0b81 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/codec.rs @@ -0,0 +1,74 @@ +use litellm_python_compat::{Value, json, literal, repr::to_str}; + +use crate::{Error, FromValue, IntoValue, Number}; + +#[derive(Debug, thiserror::Error)] +pub enum DecodeError { + #[error("Redis value is not UTF-8")] + NotUtf8, + #[error("Redis value is neither JSON nor a Python literal: {0}")] + Unparsed(#[source] litellm_python_compat::Error), + #[error("Redis value has the wrong shape: {0}")] + Shape(&'static str), +} + +/// `RedisCache._get_cache_logic`: `json.loads`, falling back to `ast.literal_eval`. +pub(crate) fn python_value(raw: &[u8]) -> Result { + let text = std::str::from_utf8(raw).map_err(|_| DecodeError::NotUtf8)?; + json::loads(text).or_else(|_| literal::literal_eval(text).map_err(DecodeError::Unparsed)) +} + +/// A cached value as `RedisCache._get_cache_logic` reads it. +pub(crate) fn decode(raw: &[u8]) -> Result { + V::from_value(&python_value(raw)?) +} + +/// `json.dumps(value)`, what `RedisCache.async_set_cache` writes. +pub(crate) fn dumps(value: V) -> Result { + Ok(json::dumps(&value.into_value())?) +} + +/// `str(value)`, what the sync `RedisCache.set_cache` writes. +pub(crate) fn str(value: V) -> String { + to_str(&value.into_value()) +} + +/// A number written by Python, `shape` naming what it was meant to be when it is not one. +pub(crate) fn number(raw: &[u8], shape: &'static str) -> Result { + Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape(shape)) +} + +#[cfg(test)] +mod tests { + use litellm_python_compat::{Value, json, repr::repr}; + use rstest::rstest; + + use super::{number, python_value}; + use crate::Number; + + // Expected values are the repr of main's RedisCache._get_cache_logic on these bytes + #[rstest] + #[case::nan(b"NaN", "nan")] + #[case::infinity(b"Infinity", "inf")] + #[case::negative_infinity(b"-Infinity", "-inf")] + #[case::nan_inside_a_document(br#"{"spend": NaN}"#, "{'spend': nan}")] + #[case::python_literal_fallback(b"{'a': (1, 2)}", "{'a': (1, 2)}")] + fn reads_what_python_reads(#[case] raw: &[u8], #[case] expected: &str) { + assert_eq!(repr(&python_value(raw).unwrap()), expected); + } + + #[rstest] + #[case::nan(f64::NAN)] + #[case::infinity(f64::INFINITY)] + #[case::negative_infinity(f64::NEG_INFINITY)] + fn a_non_finite_number_reads_back_as_written(#[case] value: f64) { + let written = json::dumps(&Value::Float(value)).unwrap(); + + let read = number(written.as_bytes(), "a number").unwrap(); + + assert_eq!( + repr(&read.to_value()), + repr(&Number::Float(value).to_value()) + ); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs b/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs new file mode 100644 index 00000000000..71dfda49233 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs @@ -0,0 +1,242 @@ +use std::{ + collections::{HashMap, VecDeque}, + sync::Mutex, + time::Duration, +}; + +use litellm_redis_client::{Declared, Error as RedisError, commands}; + +use super::{DualCache, MemorySettings}; +use crate::{Error, FromValue, Send, codec, memory::lock, request::Misses}; + +/// A batch-reading family's `DualCache` settings, as the host reads them from `litellm` globals. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TierPolicy { + /// Python's `redis_batch_cache_expiry`: a key memory misses is read from Redis at most this often. + pub refresh_after: Duration, + /// Python's `default_max_redis_batch_cache_size`: how many keys the refresh throttle tracks. + pub tracked_reads: usize, + /// `litellm.default_in_memory_ttl`, which a value read from Redis keeps in memory. + pub default_in_memory_ttl: Option, +} + +impl TierPolicy { + pub(crate) fn memory(&self) -> MemorySettings { + MemorySettings { + entries: None, + ttl: self.default_in_memory_ttl, + } + } +} + +/// `DualCache.last_redis_batch_access_time` with its expiry, which only `async_batch_get_cache` reads. +pub(crate) struct BatchReads { + refresh_after: Duration, + recent: Mutex, +} + +impl BatchReads { + pub(super) fn new(policy: TierPolicy) -> Self { + Self { + refresh_after: policy.refresh_after, + recent: Mutex::new(RecentReads::new(policy.tracked_reads)), + } + } +} + +impl DualCache { + /// `DualCache.async_batch_get_cache`: memory first, then one `MGET` for the misses whose refresh + /// is due, written back to memory. A failed Redis read serves memory, as `RedisCache` swallows + /// it, and only a breaker refusal releases the refresh it reserved. + pub(crate) async fn batch_get(&self, keys: &[String]) -> Result>, Error> { + let declared = self.declare_batch_get(Send::Now, keys)?; + self.resolve_batch_get(declared).await + } + + /// The first half of `batch_get`: reserves the due refreshes and declares their `MGET` on the + /// batch `send` picks, without sending anything. + pub(crate) fn declare_batch_get( + &self, + send: Send<'_>, + keys: &[String], + ) -> Result { + let fetch = match self.redis.client() { + None => None, + Some(client) => { + let reservation = self.reserve(keys, &self.read_memory(keys)?); + (!reservation.keys.is_empty()).then(|| { + let keys: Vec<&str> = reservation.keys.iter().map(String::as_str).collect(); + let replies = send.batch(client).add(commands::mget::>(&keys)); + Fetch { + reservation, + replies, + misses: send.misses(client), + } + }) + } + }; + Ok(DeclaredRead { + keys: keys.to_vec(), + fetch, + }) + } + + /// The second half of `batch_get`. Memory is read again here, so a value written to memory + /// since the declaration wins over what Redis answered. The keys Redis lacks are noted as the + /// request's misses. An undecodable value serves memory for the whole read, where `main`'s + /// `RedisCache` returns the raw bytes and `DualCache` backfills memory with them. + pub(crate) async fn resolve_batch_get( + &self, + declared: DeclaredRead, + ) -> Result>, Error> { + let DeclaredRead { keys, fetch } = declared; + let Some(Fetch { + reservation, + replies, + misses, + }) = fetch + else { + return self.read_memory(&keys); + }; + let fetched = replies.resolve().await; + let local = self.read_memory(&keys)?; + let fetched = match fetched { + Ok(fetched) => fetched, + Err(RedisError::BreakerOpen) => { + self.release(reservation); + return Ok(local); + } + Err(error) => { + crate::failure_log::redis_failure!( + &error, + "Redis read failed, serving the memory tier" + ); + return Ok(local); + } + }; + if let Some(misses) = misses { + misses.note( + reservation + .keys + .iter() + .zip(&fetched) + .filter(|(_, raw)| raw.is_none()) + .map(|(key, _)| key.clone()), + ); + } + let decoded = match fetched + .iter() + .map(|raw| raw.as_deref().map(codec::decode).transpose()) + .collect::, _>>() + { + Ok(decoded) => decoded, + Err(error) => { + tracing::warn!(%error, "a Redis value did not decode, serving the memory tier"); + return Ok(local); + } + }; + let read = reservation + .keys + .into_iter() + .zip(decoded) + .collect::>(); + keys.iter() + .zip(local) + .map(|(key, local)| match local { + Some(value) => Ok(Some(value)), + None => self.backfill(key, read.get(key).cloned().flatten(), None), + }) + .collect() + } + + fn read_memory(&self, keys: &[String]) -> Result>, Error> { + keys.iter().map(|key| Ok(self.memory.get(key)?)).collect() + } + + fn reserve(&self, keys: &[String], local: &[Option]) -> Reservation { + let now = self.now(); + let refresh_after = self.reads.refresh_after; + let mut recent = lock(&self.reads.recent); + let due = keys + .iter() + .zip(local) + .filter(|(_, value)| value.is_none()) + .filter_map(|(key, _)| { + let previous = recent.get(key); + let due = previous.is_none_or(|read| now.saturating_sub(read) >= refresh_after); + due.then(|| { + recent.insert(key.clone(), now); + (key.clone(), previous) + }) + }) + .collect::>(); + let (keys, previous) = due.into_iter().unzip(); + Reservation { keys, previous } + } + + fn release(&self, reservation: Reservation) { + let mut recent = lock(&self.reads.recent); + for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) { + match previous { + Some(read) => recent.insert(key, read), + None => recent.remove(&key), + } + } + } +} + +/// The `MGET` `declare_batch_get` put on a batch, with the refreshes it reserved. +pub(crate) struct DeclaredRead { + keys: Vec, + fetch: Option, +} + +struct Fetch { + reservation: Reservation, + replies: Declared>>>, + misses: Option, +} + +struct Reservation { + keys: Vec, + previous: Vec>, +} + +/// Python's `LimitedSizeOrderedDict`: insertion-ordered, and any write at capacity drops the +/// oldest entry first, even when the written key is already present. +struct RecentReads { + capacity: usize, + order: VecDeque, + times: HashMap, +} + +impl RecentReads { + fn new(capacity: usize) -> Self { + Self { + capacity, + order: VecDeque::new(), + times: HashMap::new(), + } + } + + fn get(&self, key: &str) -> Option { + self.times.get(key).copied() + } + + fn insert(&mut self, key: String, time: Duration) { + if self.times.len() >= self.capacity + && let Some(oldest) = self.order.pop_front() + { + self.times.remove(&oldest); + } + if self.times.insert(key.clone(), time).is_none() { + self.order.push_back(key); + } + } + + fn remove(&mut self, key: &str) { + if self.times.remove(key).is_some() { + self.order.retain(|tracked| tracked != key); + } + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/increment.rs b/litellm-rust/crates/redis-store/src/dual_cache/increment.rs new file mode 100644 index 00000000000..705bc71b2c1 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/increment.rs @@ -0,0 +1,129 @@ +use std::time::Duration; + +use futures_util::future::join_all; +use litellm_redis_client::{Declared, Expiry, RedisBatch, commands}; + +use super::DualCache; +use crate::{Error, Number, RequestRedis, Send, failure_log::redis_failure, memory::Counter}; + +const INCREMENT_FAILED: &str = + "Redis async_increment_cache failed, falling back to in-memory result"; +const PIPELINE_FAILED: &str = + "Redis async_increment_cache_pipeline failed, falling back to in-memory result"; + +/// A `RedisPipelineIncrementOperation`. +pub(crate) struct PipelineIncrement { + pub(crate) key: String, + pub(crate) amount: Number, + pub(crate) ttl: Option, +} + +impl DualCache { + /// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL, else + /// setting `ttl` or Redis's default. `None` when memory refuses the increment, which then never + /// reaches Redis, and the memory count when Redis fails. Both only log, as there. + pub(crate) async fn increment( + &self, + key: &str, + amount: Number, + ttl: Option, + ) -> Option { + let local = self + .memory + .increment(key, amount, ttl) + .map_err(|error| redis_failure!(&Error::from(error), INCREMENT_FAILED)) + .ok()?; + let Some(remote) = self.redis.get() else { + return Some(local); + }; + let expiry = Expiry::KeepExisting(ttl.unwrap_or(remote.default_ttl)); + let increment = commands::incr_by_float(key, amount.into(), Some(expiry)); + match remote.client.send(increment).await { + Ok(total) => Some(Number::Float(total)), + Err(error) => { + redis_failure!(&error, INCREMENT_FAILED); + Some(local) + } + } + } + + /// `DualCache.async_increment_cache_pipeline`: memory, then one pipeline of `INCRBYFLOAT` with + /// `EXPIRE` per increment. `None` when memory refuses an increment, which leaves the ones before + /// it applied and sends nothing, and the memory totals when Redis fails. + pub(crate) async fn increment_pipeline( + &self, + increments: &[PipelineIncrement], + ) -> Option> { + let local = increments + .iter() + .map(|increment| { + self.memory + .increment(&increment.key, increment.amount, increment.ttl) + }) + .collect::, _>>() + .map_err(|error| redis_failure!(&Error::from(error), PIPELINE_FAILED)) + .ok()?; + let Some(client) = self.redis.client() else { + return Some(local); + }; + let pipeline = RedisBatch::pipelined(client.clone()); + let totals = increments + .iter() + .map(|increment| { + pipeline.add(commands::incr_by_float( + &increment.key, + increment.amount.into(), + increment.ttl.map(Expiry::Refresh), + )) + }) + .collect::>(); + match Declared::all(totals) + .await + .into_iter() + .collect::, _>>() + { + Ok(totals) => Some(totals.into_iter().map(Number::Float).collect()), + Err(error) => { + redis_failure!(&error, PIPELINE_FAILED); + Some(local) + } + } + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, stopping at the first + /// increment memory refuses and only logging it, then each increment on the request's post-call + /// batch, whose failure is logged once it settles. + pub(crate) fn increment_pipeline_post_call( + &self, + request: &RequestRedis, + increments: &[PipelineIncrement], + ) { + if let Err(error) = increments.iter().try_for_each(|increment| { + self.memory + .increment(&increment.key, increment.amount, increment.ttl) + .map(drop) + }) { + tracing::warn!(%error, "in-memory increment failed"); + } + let Some(batch) = self.redis.batch(Send::PostCall(request)) else { + return; + }; + let declared = increments + .iter() + .map(|increment| { + batch.add(commands::incr_by_float( + &increment.key, + increment.amount.into(), + increment.ttl.map(Expiry::KeepExisting), + )) + }) + .collect::>(); + request.spawn(async move { + join_all(declared.into_iter().map(Declared::settled)) + .await + .into_iter() + .filter_map(Result::err) + .for_each(|error| redis_failure!(&error, "deferred Redis increment failed")); + }); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/mod.rs b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs new file mode 100644 index 00000000000..b2c277a118e --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs @@ -0,0 +1,279 @@ +mod batch_read; +mod increment; +mod read; +mod write; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +pub use self::batch_read::TierPolicy; +pub(crate) use self::{batch_read::BatchReads, increment::PipelineIncrement}; +use crate::{ + memory::{Clock, Memory, lock}, + remote::RemoteSlot, +}; + +/// The memory tier's `InMemoryCache` size and the `default_in_memory_ttl` a write or a Redis +/// backfill takes when it names no TTL. `None` keeps `InMemoryCache`'s own default. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct MemorySettings { + pub(crate) entries: Option, + pub(crate) ttl: Option, +} + +/// Python's `DualCache`: a per-worker `InMemoryCache` in front of the Redis every pod shares. Each +/// method is the `DualCache` method it names. A family whose Python reached past `DualCache` to its +/// `in_memory_cache` or `redis_cache` uses `memory` and `redis` the same way. +pub(crate) struct DualCache { + memory: Memory, + redis: Arc, + default_in_memory_ttl: Mutex>, + now: Clock, + reads: R, +} + +impl DualCache { + /// `redis` may be shared with the family's other caches, as one `RedisCache` is attached to each. + pub(crate) fn new(redis: Arc, memory: MemorySettings, now: Clock) -> Self { + Self { + memory: Memory::new(memory.entries, &now), + redis, + default_in_memory_ttl: Mutex::new(memory.ttl), + now, + reads: (), + } + } + + /// A cache that batch-reads, with the `async_batch_get_cache` throttle `policy` configures. + pub(crate) fn batch_reading( + redis: Arc, + policy: TierPolicy, + now: Clock, + ) -> DualCache { + let Self { + memory, + redis, + default_in_memory_ttl, + now, + reads: (), + } = Self::new(redis, policy.memory(), now); + DualCache { + memory, + redis, + default_in_memory_ttl, + now, + reads: BatchReads::new(policy), + } + } +} + +impl DualCache { + pub(crate) fn memory(&self) -> &Memory { + &self.memory + } + + pub(crate) fn redis(&self) -> &RemoteSlot { + &self.redis + } + + pub(crate) fn now(&self) -> Duration { + (self.now)() + } + + pub(crate) fn default_in_memory_ttl(&self) -> Option { + *lock(&self.default_in_memory_ttl) + } + + /// `DualCache.update_cache_ttl` and `update_in_memory_max_size` together: both apply in place, + /// and the entries memory already holds stay. + pub(crate) fn reconfigure(&self, memory: MemorySettings) { + *lock(&self.default_in_memory_ttl) = memory.ttl; + self.memory.update_max_size(memory.entries); + } +} + +#[cfg(test)] +mod tests { + use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, + }; + + use rstest::rstest; + use tokio::runtime::Handle; + + use super::{DualCache, MemorySettings, PipelineIncrement}; + use crate::{Clock, Number, RequestRedis, memory::Counter, remote::RemoteSlot}; + + #[derive(Clone, Debug, PartialEq)] + enum Entry { + Number(Number), + Text, + } + + impl From for Entry { + fn from(number: Number) -> Self { + Self::Number(number) + } + } + + impl Counter for Entry { + fn as_number(&self) -> Option { + match self { + Self::Number(number) => Some(*number), + Self::Text => None, + } + } + } + + fn clock() -> (Clock, Arc) { + let seconds = Arc::new(AtomicU64::new(1_000)); + let read = seconds.clone(); + let clock: Clock = Arc::new(move || Duration::from_secs(read.load(Ordering::SeqCst))); + (clock, seconds) + } + + fn without_redis( + default_in_memory_ttl: Option, + now: Clock, + ) -> DualCache { + DualCache::new( + Arc::new(RemoteSlot::new(None)), + MemorySettings { + entries: None, + ttl: default_in_memory_ttl, + }, + now, + ) + } + + fn increment(key: &str, amount: Number) -> PipelineIncrement { + PipelineIncrement { + key: key.to_owned(), + amount, + ttl: None, + } + } + + // InMemoryCache's own default TTL is 600 s, which an entry outlives here only when the cache + // configured none + #[rstest] + #[case::default_in_memory_ttl(Some(Duration::from_secs(5)), None)] + #[case::in_memory_cache_default(None, Some(Number::Int(1)))] + #[tokio::test] + async fn a_set_without_a_ttl_takes_default_in_memory_ttl( + #[case] default_in_memory_ttl: Option, + #[case] after_six_seconds: Option, + ) { + let (now, seconds) = clock(); + let cache = without_redis(default_in_memory_ttl, now); + + cache.set("k", Number::Int(1), None).await.unwrap(); + seconds.fetch_add(6, Ordering::SeqCst); + + assert_eq!(cache.memory().get("k").unwrap(), after_six_seconds); + } + + #[rstest] + #[case::absent_starts_at_int_zero(None, Number::Int(2), Number::Int(2))] + #[case::int_plus_int_stays_int(Some(Number::Int(1)), Number::Int(2), Number::Int(3))] + #[case::a_float_makes_a_float(Some(Number::Int(1)), Number::Float(0.5), Number::Float(1.5))] + fn a_local_increment_adds_as_python_adds( + #[case] stored: Option, + #[case] amount: Number, + #[case] expected: Number, + ) { + let (now, _) = clock(); + let cache = without_redis::(None, now); + if let Some(stored) = stored { + cache.memory().set("k", stored, None).unwrap(); + } + + let total = cache.memory().increment("k", amount, None).unwrap(); + + assert_eq!(total, expected); + assert_eq!(cache.memory().get("k").unwrap(), Some(expected)); + } + + #[test] + fn a_local_increment_of_a_non_number_fails_and_keeps_the_entry() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + cache.memory().set("k", Entry::Text, None).unwrap(); + + let total = cache.memory().increment("k", Number::Int(1), None); + + assert!(total.is_err()); + assert_eq!(cache.memory().get("k").unwrap(), Some(Entry::Text)); + } + + // async_increment_cache and async_increment_cache_pipeline catch the in-memory failure with the + // Redis one, so the caller gets None and nothing is sent + #[tokio::test] + async fn an_increment_memory_refuses_answers_none() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + cache.memory().set("text", Entry::Text, None).unwrap(); + + let single = cache.increment("text", Number::Int(1), None).await; + let pipelined = cache + .increment_pipeline(&[ + increment("before", Number::Int(1)), + increment("text", Number::Int(1)), + ]) + .await; + + assert_eq!(single, None); + assert_eq!(pipelined, None); + assert_eq!(cache.memory().get("text").unwrap(), Some(Entry::Text)); + } + + // update_cache_ttl and update_in_memory_max_size change the settings in place, so what memory + // already holds survives and the next write takes the new TTL + #[tokio::test] + async fn reconfigure_keeps_entries_and_applies_the_new_ttl() { + let (now, seconds) = clock(); + let cache = without_redis(None, now); + cache.set("kept", Number::Int(1), None).await.unwrap(); + + cache.reconfigure(MemorySettings { + entries: Some(10), + ttl: Some(Duration::from_secs(5)), + }); + cache.set("short", Number::Int(2), None).await.unwrap(); + seconds.fetch_add(6, Ordering::SeqCst); + + assert_eq!(cache.memory().get("kept").unwrap(), Some(Number::Int(1))); + assert_eq!(cache.memory().get("short").unwrap(), None); + } + + // async_increment_cache_pipeline_post_call: InMemoryCache.async_increment_pipeline raises at the + // first entry it cannot add to, so later increments never reach memory, and DualCache only logs + #[tokio::test] + async fn a_post_call_increment_stops_at_the_first_entry_memory_refuses() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + cache.memory().set("text", Entry::Text, None).unwrap(); + + cache.increment_pipeline_post_call( + &request, + &[ + increment("before", Number::Int(1)), + increment("text", Number::Int(1)), + increment("after", Number::Int(1)), + ], + ); + + assert_eq!( + cache.memory().get("before").unwrap(), + Some(Entry::Number(Number::Int(1))) + ); + assert_eq!(cache.memory().get("after").unwrap(), None); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/read.rs b/litellm-rust/crates/redis-store/src/dual_cache/read.rs new file mode 100644 index 00000000000..14156483c6f --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/read.rs @@ -0,0 +1,55 @@ +use std::time::Duration; + +use litellm_redis_client::commands; + +use super::DualCache; +use crate::{Error, FromValue, RequestRedis, codec, failure_log::redis_failure}; + +impl DualCache { + /// `DualCache.async_get_cache`: memory, else `None` for a key `request` already read as missing + /// from Redis, else a `GET` written back to memory for `ttl` or `default_in_memory_ttl`. A failed + /// or undecodable Redis read is a miss, as there. + pub(crate) async fn get( + &self, + request: Option<&RequestRedis>, + key: &str, + ttl: Option, + ) -> Result, Error> { + if let Some(value) = self.memory.get(key)? { + return Ok(Some(value)); + } + let Some(client) = self.redis.client() else { + return Ok(None); + }; + if request.is_some_and(|request| request.misses(client).contains(key)) { + return Ok(None); + } + let value = client + .send(commands::get::>(key)) + .await + .map_err(Error::from) + .and_then(|raw| Ok(raw.as_deref().map(codec::decode).transpose()?)) + .unwrap_or_else(|error| { + redis_failure!(&error, "LiteLLM Cache: exception in async_get_cache"); + None + }); + self.backfill(key, value, ttl) + } +} + +impl DualCache { + /// `_backfill_kwargs`: a value read from Redis stays in memory for `ttl` or + /// `default_in_memory_ttl`. + pub(super) fn backfill( + &self, + key: &str, + value: Option, + ttl: Option, + ) -> Result, Error> { + if let Some(value) = &value { + self.memory + .set(key, value.clone(), ttl.or(self.default_in_memory_ttl()))?; + } + Ok(value) + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/write.rs b/litellm-rust/crates/redis-store/src/dual_cache/write.rs new file mode 100644 index 00000000000..6fcc6f32741 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/write.rs @@ -0,0 +1,178 @@ +use std::time::Duration; + +use litellm_redis_client::{Declared, RedisBatch, commands}; + +use super::DualCache; +use crate::{ + Error, IntoValue, Pending, Send, codec, + failure_log::redis_failure, + request::{pending, ready}, +}; + +/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `async_delete_cache_keys` deletes at once. +const DELETE_CHUNK: usize = 1000; + +impl DualCache { + /// `DualCache.async_set_cache`: memory, then a `SET` of `json.dumps(value)`. A Redis failure, or + /// a value `json.dumps` refuses, only logs. + pub(crate) async fn set( + &self, + key: &str, + value: V, + ttl: Option, + ) -> Result<(), Error> { + self.set_now(key, value, ttl, codec::dumps).await + } + + /// `DualCache.set_cache`, whose sync `RedisCache.set_cache` writes `str(value)`. + pub(crate) async fn set_str( + &self, + key: &str, + value: V, + ttl: Option, + ) -> Result<(), Error> { + self.set_now(key, value, ttl, |value| Ok(codec::str(value))) + .await + } + + /// `DualCache._set_on_batch`: memory now, the `SET` on the batch `send` picks. Memory and Redis + /// both take `ttl`, else `default_in_memory_ttl`, and Redis falls back to its own default. + pub(crate) fn declare_set( + &self, + send: Send<'_>, + key: &str, + value: V, + ttl: Option, + ) -> Result, Error> { + self.declare_encoded(send, key, value, ttl, codec::dumps) + } + + /// `DualCache.async_set_cache_pipeline`: memory, then one pipeline of `SET`s. A value + /// `json.dumps` refuses aborts the whole pipeline, and any failure only logs, as `RedisCache` + /// does. + pub(crate) async fn set_pipeline( + &self, + entries: Vec<(String, V)>, + ttl: Option, + ) -> Result<(), Error> { + let ttl = ttl.or(self.default_in_memory_ttl()); + let encoded = self.redis.get().map(|_| { + entries + .iter() + .map(|(_, value)| codec::dumps(value.clone())) + .collect::, Error>>() + }); + let keys = entries + .into_iter() + .map(|(key, value)| { + self.memory.set(&key, value, ttl)?; + Ok(key) + }) + .collect::, Error>>()?; + let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else { + return Ok(()); + }; + let failure = match encoded { + Err(error) => Some(error), + Ok(values) => { + let pipeline = RedisBatch::pipelined(remote.client.clone()); + let ttl = ttl.unwrap_or(remote.default_ttl); + let written = keys + .iter() + .zip(values) + .map(|(key, value)| pipeline.add(commands::set(key, value, ttl))) + .collect::>(); + Declared::all(written) + .await + .into_iter() + .find_map(Result::err) + .map(Error::from) + } + }; + if let Some(error) = failure { + redis_failure!( + &error, + "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS", + ); + } + Ok(()) + } + + async fn set_now( + &self, + key: &str, + value: V, + ttl: Option, + encode: impl FnOnce(V) -> Result, + ) -> Result<(), Error> { + if let Err(error) = self + .declare_encoded(Send::Now, key, value, ttl, encode)? + .await + { + redis_failure!(&error, "LiteLLM Cache: exception in async add_cache"); + } + Ok(()) + } + + fn declare_encoded( + &self, + send: Send<'_>, + key: &str, + value: V, + ttl: Option, + encode: impl FnOnce(V) -> Result, + ) -> Result, Error> { + let ttl = ttl.or(self.default_in_memory_ttl()); + let encoded = self.redis.get().map(|_| encode(value.clone())); + self.memory.set(key, value, ttl)?; + let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else { + return Ok(ready(())); + }; + let encoded = match encoded { + Ok(encoded) => encoded, + Err(error) => { + redis_failure!(&error, "a cache value did not encode, kept in memory only"); + return Ok(ready(())); + } + }; + if let Some(misses) = send.misses(&remote.client) { + misses.forget(key); + } + Ok(pending(send.batch(&remote.client).add(commands::set( + key, + encoded, + ttl.unwrap_or(remote.default_ttl), + )))) + } +} + +impl DualCache { + /// `DualCache.async_delete_cache_pre_call`: memory now, the `DEL` on the batch `send` picks, + /// after which the request reads the key as missing. + pub(crate) fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { + self.memory.delete(key)?; + let Some(client) = self.redis.client() else { + return Ok(ready(())); + }; + if let Some(misses) = send.misses(client) { + misses.note([key.to_owned()]); + } + let deleted = send.batch(client).add(commands::del(&[key])); + Ok(Box::pin( + async move { Ok(deleted.resolve().await.map(drop)?) }, + )) + } + + /// `DualCache.async_delete_cache_keys`: memory, then one multi-key `DEL` per chunk, in order. + pub(crate) async fn delete_keys(&self, keys: &[String]) -> Result<(), Error> { + keys.iter().try_for_each(|key| self.memory.delete(key))?; + let Some(client) = self.redis.client() else { + return Ok(()); + }; + for chunk in keys.chunks(DELETE_CHUNK) { + let chunk: Vec<&str> = chunk.iter().map(String::as_str).collect(); + client.send(commands::del(&chunk)).await?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/redis-store/src/error.rs b/litellm-rust/crates/redis-store/src/error.rs new file mode 100644 index 00000000000..fdc722855bb --- /dev/null +++ b/litellm-rust/crates/redis-store/src/error.rs @@ -0,0 +1,63 @@ +use crate::DecodeError; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Redis(#[from] litellm_redis_client::Error), + #[error(transparent)] + Memory(#[from] litellm_cache::Error), + #[error(transparent)] + Decode(#[from] DecodeError), + #[error("the value is not JSON serializable: {0}")] + Encode(#[from] litellm_python_compat::Error), +} + +impl Error { + pub fn kind(&self) -> ErrorKind { + match self { + Self::Redis(error) => error.into(), + Self::Memory(_) => ErrorKind::Memory, + Self::Decode(_) => ErrorKind::Decode, + Self::Encode(_) => ErrorKind::Encode, + } + } +} + +/// The failure class a host reports with a store error, `timeout` and `breaker_open` being the ones a +/// caller may drop. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr)] +pub enum ErrorKind { + #[strum(serialize = "breaker_open")] + BreakerOpen, + #[strum(serialize = "timeout")] + Timeout, + #[strum(serialize = "transport")] + Transport, + #[strum(serialize = "rejected")] + Rejected, + #[strum(serialize = "decode")] + Decode, + #[strum(serialize = "encode")] + Encode, + #[strum(serialize = "memory")] + Memory, +} + +impl From<&Error> for ErrorKind { + fn from(error: &Error) -> Self { + error.kind() + } +} + +impl From<&litellm_redis_client::Error> for ErrorKind { + fn from(error: &litellm_redis_client::Error) -> Self { + use litellm_redis_client::Error; + match error { + Error::BreakerOpen => Self::BreakerOpen, + Error::Transport(error) if error.is_timeout() => Self::Timeout, + Error::Transport(_) => Self::Transport, + Error::Server(_) => Self::Rejected, + Error::Decode(_) => Self::Decode, + } + } +} diff --git a/litellm-rust/crates/redis-store/src/failure_log.rs b/litellm-rust/crates/redis-store/src/failure_log.rs new file mode 100644 index 00000000000..e6bf5187cfd --- /dev/null +++ b/litellm-rust/crates/redis-store/src/failure_log.rs @@ -0,0 +1,18 @@ +/// `log_redis_failure`: an open breaker is DEBUG, anything else WARNING. The event's `redis.failure` +/// field names the class, so the host can rate-limit timeout lines with the throttle redis-py shares. +/// A macro so the event carries its caller's target, file and line, as `stacklevel=2` does in Python. +macro_rules! redis_failure { + ($error:expr, $message:expr $(,)?) => {{ + let error = $error; + let kind = $crate::ErrorKind::from(error); + let failure: &'static str = kind.into(); + match kind { + $crate::ErrorKind::BreakerOpen => { + tracing::debug!(redis.failure = failure, "{}: {}", $message, error) + } + _ => tracing::warn!(redis.failure = failure, "{}: {}", $message, error), + } + }}; +} + +pub(crate) use redis_failure; diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs new file mode 100644 index 00000000000..69600aac2a6 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -0,0 +1,24 @@ +pub mod auth_cache; +mod codec; +mod dual_cache; +mod error; +mod failure_log; +mod memory; +mod number; +pub mod rate_limits; +mod remote; +mod request; +pub mod router_cooldowns; +pub mod router_usage; +pub mod sensitive_route_pins; +pub mod spend_counters; +mod value; + +pub use codec::DecodeError; +pub use dual_cache::TierPolicy; +pub use error::{Error, ErrorKind}; +pub use memory::{Clock, system_clock}; +pub use number::Number; +pub use remote::Remote; +pub use request::{Pending, RequestRedis, Send}; +pub use value::{FromValue, IntoValue}; diff --git a/litellm-rust/crates/redis-store/src/memory.rs b/litellm-rust/crates/redis-store/src/memory.rs new file mode 100644 index 00000000000..f89d5926760 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/memory.rs @@ -0,0 +1,122 @@ +use std::{ + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::Error; +use litellm_cache_memory::InMemoryCache; + +use crate::Number; + +pub type Clock = Arc Duration + Send + Sync>; + +pub fn system_clock() -> Clock { + Arc::new(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + }) +} + +/// A store's memory tier: Python's `InMemoryCache`, with its 200 entries and 600 s TTL unless the +/// store names others, plus one lock for every read-modify-write on it. +pub struct Memory { + cache: InMemoryCache, + updates: Mutex<()>, +} + +impl Memory { + pub fn new(entries: Option, now: &Clock) -> Self { + let now = now.clone(); + Self { + cache: InMemoryCache::with_clock(entries, None, move || now()), + updates: Mutex::new(()), + } + } + + pub fn get(&self, key: &str) -> Result, Error> { + self.cache.get_cache(key) + } + + pub fn set(&self, key: &str, value: V, ttl: Option) -> Result<(), Error> { + self.cache.set_cache(key, value, ttl).map(drop) + } + + pub fn delete(&self, key: &str) -> Result<(), Error> { + self.cache.delete_cache(key) + } + + pub fn flush(&self) -> Result<(), Error> { + self.cache.flush_cache() + } + + pub fn len(&self) -> Result { + self.cache.len() + } + + pub fn expires_at(&self, key: &str) -> Result, Error> { + self.cache.expires_at(key) + } + + pub fn entries(&self) -> Result, Error> { + self.cache.entries() + } + + /// `DualCache.update_in_memory_max_size`: `None` restores the 200-entry default. + pub fn update_max_size(&self, entries: Option) { + self.cache.update_max_size_in_memory(entries); + } + + /// Read-modify-write of one entry, atomic against every other update of this tier. + pub fn update( + &self, + key: &str, + ttl: Option, + update: impl FnOnce(Option) -> Result, + ) -> Result { + let _guard = self.guard(); + let value = update(self.get(key)?)?; + self.set(key, value.clone(), ttl)?; + Ok(value) + } + + /// Held across a read and a conditional write that `update` cannot express. + pub fn guard(&self) -> MutexGuard<'_, ()> { + lock(&self.updates) + } +} + +/// A value `InMemoryCache.increment_cache` can add to. +pub(crate) trait Counter: Clone + From { + fn as_number(&self) -> Option; +} + +impl Counter for Number { + fn as_number(&self) -> Option { + Some(*self) + } +} + +impl Memory { + /// `InMemoryCache.increment_cache`: `(current or 0) + amount`, refusing an entry that is not a + /// number and leaving it in place. + pub(crate) fn increment( + &self, + key: &str, + amount: Number, + ttl: Option, + ) -> Result { + let _guard = self.guard(); + let current = match self.get(key)? { + None => Number::Int(0), + Some(value) => value.as_number().ok_or(Error::InvalidEntry(None))?, + }; + let total = current + amount; + self.set(key, V::from(total), ttl)?; + Ok(total) + } +} + +pub(crate) fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} diff --git a/litellm-rust/crates/redis-store/src/number.rs b/litellm-rust/crates/redis-store/src/number.rs new file mode 100644 index 00000000000..7f3e9f4d10b --- /dev/null +++ b/litellm-rust/crates/redis-store/src/number.rs @@ -0,0 +1,83 @@ +use std::{cmp::Ordering, ops::Add}; + +use litellm_python_compat::{BigInt, Value, json, repr}; + +/// A Python `int` or `float`, kept apart because `repr` and `json.dumps` write them differently. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Number { + Int(i64), + Float(f64), +} + +impl From for litellm_redis_client::Amount { + fn from(number: Number) -> Self { + match number { + Number::Int(value) => Self::Int(value), + Number::Float(value) => Self::Float(value), + } + } +} + +impl Number { + pub fn as_f64(self) -> f64 { + match self { + Self::Int(value) => value as f64, + Self::Float(value) => value, + } + } + + /// Python's `int(value)`, which truncates toward zero. + pub fn truncate(self) -> i64 { + match self { + Self::Int(value) => value, + Self::Float(value) => value.trunc() as i64, + } + } + + pub fn to_value(self) -> Value { + match self { + Self::Int(value) => Value::Int(BigInt::from(value)), + Self::Float(value) => Value::Float(value), + } + } + + /// `json.dumps`, which `RedisCache.async_set_cache` writes. + pub fn to_json(self) -> String { + let value = self.to_value(); + json::dumps(&value).unwrap_or_else(|_| repr::to_str(&value)) + } + + pub fn from_value(value: &Value) -> Option { + match value { + Value::Int(value) => i64::try_from(value).ok().map(Self::Int), + Value::Float(value) => Some(Self::Float(*value)), + _ => None, + } + } + + /// Python's `min` over numbers: the first of the smallest, keeping its type. + pub fn min(values: impl IntoIterator) -> Option { + values + .into_iter() + .fold(None, |smallest, value| match smallest { + Some(smallest) + if value.as_f64().partial_cmp(&smallest.as_f64()) != Some(Ordering::Less) => + { + Some(smallest) + } + _ => Some(value), + }) + } +} + +/// Python's `+`: two ints stay an int, anything with a float is a float. +impl Add for Number { + type Output = Self; + + fn add(self, other: Self) -> Self { + match (self, other) { + (Self::Int(left), Self::Int(right)) => Self::Int(left + right), + (left, right) => Self::Float(left.as_f64() + right.as_f64()), + } + } +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/keys.rs b/litellm-rust/crates/redis-store/src/rate_limits/keys.rs new file mode 100644 index 00000000000..7f5714a19ad --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/keys.rs @@ -0,0 +1,74 @@ +/// A rate-limit descriptor, `{descriptor:value}`, the hash tag every key of its window shares. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Scope { + pub descriptor: String, + pub value: String, +} + +impl Scope { + pub(super) fn key(&self, suffix: &str) -> String { + format!("{{{}:{}}}:{suffix}", self.descriptor, self.value) + } + + pub(super) fn window_key(&self) -> String { + self.key(WINDOW) + } + + pub(super) fn gauge_key(&self) -> String { + self.key(GAUGE) + } +} + +const WINDOW: &str = "window"; +const GAUGE: &str = "max_parallel_requests"; + +/// A counter's key suffix, which is also how the Python limiter names it. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)] +pub enum Kind { + #[strum(serialize = "requests")] + Requests, + #[strum(serialize = "tokens")] + Tokens, +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Counter { + pub scope: Scope, + pub kind: Kind, +} + +impl Counter { + pub(super) fn key(&self) -> String { + self.scope.key(self.kind.into()) + } +} + +/// One value of the limiter's keyspace, as both tiers address it. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum Entry { + Window(Scope), + Counter(Counter), + Gauge(Scope), +} + +impl Entry { + /// The entry the Python limiter names by its key suffix: `window`, `max_parallel_requests`, or a + /// counter's `Kind`. + pub fn new(suffix: &str, scope: Scope) -> Result { + match suffix { + WINDOW => Ok(Self::Window(scope)), + GAUGE => Ok(Self::Gauge(scope)), + kind => kind + .parse() + .map(|kind| Self::Counter(Counter { scope, kind })), + } + } + + pub(super) fn key(&self) -> String { + match self { + Self::Window(scope) => scope.window_key(), + Self::Counter(counter) => counter.key(), + Self::Gauge(scope) => scope.gauge_key(), + } + } +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/lua.rs b/litellm-rust/crates/redis-store/src/rate_limits/lua.rs new file mode 100644 index 00000000000..96451212d35 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/lua.rs @@ -0,0 +1,260 @@ +//! The limiter's Lua, byte for byte as `main`'s Python registers it, so both share one script SHA. + +pub(super) const BATCH_RATE_LIMITER: &str = r#" +local results = {} +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) + +-- Process each window/counter pair +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local increment_value = 1 + + -- Check if window exists and is valid + local window_start = redis.call('GET', window_key) + if not window_start or (now - tonumber(window_start)) >= window_size then + -- Reset window and counter + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment_value) + redis.call('EXPIRE', window_key, window_size) + redis.call('EXPIRE', counter_key, window_size) + table.insert(results, tostring(now)) -- window_start + table.insert(results, increment_value) -- counter + else + local counter = redis.call('INCR', counter_key) + -- This happens when window_key exists but counter_key doesn't (e.g., tokens key + -- created after requests key when both share the same window_key) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 then + redis.call('EXPIRE', counter_key, window_size) + end + table.insert(results, window_start) -- window_start + table.insert(results, counter) -- counter + end +end + +return results +"#; + +pub(super) const BATCH_COUNTER_READ: &str = r#"return redis.call('MGET', unpack(KEYS))"#; + +pub(super) const CHECK_AND_INCREMENT_BY_N: &str = r#" +-- Atomic check-and-increment-by-N across one or more descriptors. +-- All-or-nothing: if any descriptor would exceed its limit, no counter is +-- modified. +-- +-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied +-- timestamp so that window resets are deterministic across replicas with +-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the +-- TOCTOU window across multi-replica deployments. +-- +-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. +-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: +-- ARGV[(i-1)*4 + 1] = limit +-- ARGV[(i-1)*4 + 2] = increment +-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) +-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) +-- +-- Return on success: +-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... } +-- Return on over-limit: { 1, descriptor_index, current_counter, limit } +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local descriptor_count = #KEYS / 2 +local reset_windows = {} + +-- Pass 1: read state, validate. Abort without writing if any over limit. +local descriptor_state = {} +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[arg_base]) + local increment = tonumber(ARGV[arg_base + 1]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_start = redis.call('GET', window_key) + local window_expired = (not window_start) or + ((now - tonumber(window_start)) >= window_size) + + local current_counter + if window_expired then + current_counter = 0 + else + current_counter = tonumber(redis.call('GET', counter_key) or 0) + end + + local blocked + if increment > 0 then + blocked = current_counter + increment > limit + else + blocked = current_counter >= limit + end + if blocked then + return { 1, i, current_counter, limit } + end + + descriptor_state[i] = { window_expired, current_counter, window_start } +end + +-- Pass 2: all checks passed. Apply increments. +local results = { 0 } +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_expired = descriptor_state[i][1] + local active_window_start + + if window_expired then + active_window_start = now + if not reset_windows[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset_windows[window_key] = true + end + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + else + active_window_start = tonumber(descriptor_state[i][3]) + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, new_counter) + end + table.insert(results, active_window_start) +end + +return results +"#; + +pub(super) const WINDOW_GUARDED_TOKEN_INCREMENT: &str = r#" +local results = {} +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local arg_base = ((i - 1) / 2) * 3 + 1 + local expected_window_start = ARGV[arg_base] + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local active_window_start = redis.call('GET', window_key) + + if active_window_start and active_window_start == expected_window_start then + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, 1) + table.insert(results, new_counter) + else + table.insert(results, 0) + table.insert(results, tonumber(redis.call('GET', counter_key) or 0)) + end +end +return results +"#; + +pub(super) const PARALLEL_ACQUIRE: &str = r#" +-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge. +-- Each gauge key is a sorted set of per-request slot ids scored by acquire +-- time (Redis server clock). In-flight requests are counted by ZCARD after +-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM +-- counters the gauge is never reset while requests are in flight, a +-- rejected request never occupies a slot, and a slot leaked by a crashed +-- worker self-heals after the slot TTL even under continuous traffic. +-- +-- KEYS: one gauge zset key per descriptor. +-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id). +-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +for i = 1, #KEYS do + local limit = tonumber(ARGV[(i - 1) * 3 + 1]) + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl) + local in_flight = redis.call('ZCARD', KEYS[i]) + if in_flight + 1 > limit then + return { 1, i, in_flight, limit } + end +end +local results = { 0 } +for i = 1, #KEYS do + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + local slot_id = ARGV[(i - 1) * 3 + 3] + redis.call('ZADD', KEYS[i], now, slot_id) + redis.call('EXPIRE', KEYS[i], slot_ttl) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const PARALLEL_RELEASE: &str = r#" +-- Release one slot per gauge key by removing this request's slot id. +-- ZREM of an absent member (or key) is a no-op, so a release without a +-- matching acquire (proxy-side rejection, double-fired callback, slot +-- already expired) can never free a slot owned by another request. +-- KEYS: gauge zset keys. ARGV: per-key slot_id. +-- Returns the remaining in-flight count per key. +local results = {} +for i = 1, #KEYS do + redis.call('ZREM', KEYS[i], ARGV[i]) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const PARALLEL_COUNT: &str = r#" +-- Read the current in-flight count per gauge key (prunes expired slots +-- first so leaked slots do not inflate the reading). +-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local results = {} +for i = 1, #KEYS do + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i])) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const TOKEN_INCREMENT: &str = r#" +local results = {} + +-- Process each key/increment_value/ttl triplet +for i = 1, #KEYS do + local key = KEYS[i] + local increment_value = tonumber(ARGV[i * 2 - 1]) + local ttl_seconds = tonumber(ARGV[i * 2]) + + -- Increment the value + local new_value = redis.call('INCRBYFLOAT', key, increment_value) + + -- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL + -- ttl_seconds can be 0 (no TTL) or positive (set TTL) + if ttl_seconds and ttl_seconds > 0 then + local current_ttl = redis.call('TTL', key) + if current_ttl == -1 then + redis.call('EXPIRE', key, ttl_seconds) + end + end + + table.insert(results, new_value) +end + +return results +"#; diff --git a/litellm-rust/crates/redis-store/src/rate_limits/mod.rs b/litellm-rust/crates/redis-store/src/rate_limits/mod.rs new file mode 100644 index 00000000000..71b67d2ecc0 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/mod.rs @@ -0,0 +1,304 @@ +mod keys; +mod lua; +mod scripts; +mod types; + +use std::{sync::Arc, time::Duration}; + +use litellm_redis_client::{Declared, RedisBatch, RedisClient, Slot}; + +pub use self::keys::{Counter, Entry, Kind, Scope}; +pub use self::types::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, + Increment, Local, WindowCount, +}; +use crate::{ + Error, Number, Pending, RequestRedis, Send, + dual_cache::{DualCache, MemorySettings, PipelineIncrement}, + memory::{Clock, system_clock}, + remote::{Remote, RemoteSlot}, + request::pending, +}; + +/// The v3 rate limiter's keyspace: its memory tier, its Lua scripts and the plain counter writes +/// `DualCache` made for it on `main`. The admission logic stays with the caller. A script needs +/// Redis, so without one it declares nothing and answers `None`. +pub struct RateLimits { + cache: DualCache, +} + +impl RateLimits { + /// `backfill_ttl` is the `default_in_memory_ttl` of the `DualCache` this tier replaces. An + /// increment without a TTL takes the remote's default. + pub fn new(remote: Option, backfill_ttl: Option) -> Self { + Self::with_clock(remote, backfill_ttl, system_clock()) + } + + pub fn with_clock(remote: Option, backfill_ttl: Option, now: Clock) -> Self { + Self { + cache: DualCache::new( + Arc::new(RemoteSlot::new(remote)), + MemorySettings { + entries: None, + ttl: backfill_ttl, + }, + now, + ), + } + } + + pub fn has_redis(&self) -> bool { + self.cache.redis().get().is_some() + } + + /// `_group_keys_by_hash_tag`: indices of `scopes` sharing a cluster slot, in first-seen order, + /// or one group when the client is not a cluster and grouping is not forced. + pub fn slot_groups(&self, scopes: &[Scope], force: bool) -> Vec> { + let by_slot = force + || self + .cache + .redis() + .client() + .is_some_and(RedisClient::is_cluster); + if !by_slot { + return vec![(0..scopes.len()).collect()]; + } + scopes + .iter() + .enumerate() + .fold( + Vec::<(Slot, Vec)>::new(), + |mut groups, (index, scope)| { + let slot = Slot::for_key(scope.window_key()); + match groups.iter_mut().find(|(seen, _)| *seen == slot) { + Some((_, members)) => members.push(index), + None => groups.push((slot, vec![index])), + } + groups + }, + ) + .into_iter() + .map(|(_, members)| members) + .collect() + } + + pub fn get_local(&self, entry: &Entry) -> Result, Error> { + Ok(self.cache.memory().get(&entry.key())?) + } + + pub fn set_local( + &self, + entry: &Entry, + value: Local, + ttl: Option, + ) -> Result<(), Error> { + Ok(self.cache.memory().set(&entry.key(), value, ttl)?) + } + + /// `DualCache.async_get_cache`, backfilling memory for `backfill_ttl` or the store's default. + /// A key `request` already read as missing from Redis is not read again. + pub async fn get( + &self, + request: Option<&RequestRedis>, + entry: &Entry, + backfill_ttl: Option, + ) -> Result, Error> { + self.cache.get(request, &entry.key(), backfill_ttl).await + } + + pub fn sliding_window( + &self, + send: Send<'_>, + counters: &[Counter], + now: i64, + window_size: i64, + ) -> Option>> { + Some(pending(scripts::sliding_window( + &self.batch(send)?, + counters, + now, + window_size, + ))) + } + + pub fn read_counters( + &self, + send: Send<'_>, + counters: &[Counter], + ) -> Option>>> { + Some(pending(scripts::read_counters( + &self.batch(send)?, + counters, + ))) + } + + pub fn check_and_increment( + &self, + send: Send<'_>, + counters: &[AtomicCounter], + ) -> Option> { + Some(pending(scripts::check_and_increment( + &self.batch(send)?, + counters, + ))) + } + + pub fn guarded_increment( + &self, + send: Send<'_>, + increment: &GuardedIncrement, + ) -> Option> { + Some(pending(scripts::guarded_increment( + &self.batch(send)?, + increment, + ))) + } + + pub fn parallel_acquire( + &self, + send: Send<'_>, + gauges: &[Gauge], + slot_ttl: Duration, + slot_id: &str, + ) -> Option> { + Some(pending(scripts::parallel_acquire( + &self.batch(send)?, + gauges, + slot_ttl, + slot_id, + ))) + } + + pub fn parallel_release( + &self, + send: Send<'_>, + gauges: &[Scope], + slot_id: &str, + ) -> Option>> { + Some(pending(self.declare_release(send, gauges, slot_id)?)) + } + + pub fn parallel_count( + &self, + send: Send<'_>, + gauges: &[Scope], + slot_ttl: Duration, + ) -> Option>> { + Some(pending(scripts::parallel_count( + &self.batch(send)?, + gauges, + slot_ttl, + ))) + } + + pub fn token_increment( + &self, + send: Send<'_>, + increments: &[Increment], + ) -> Option>> { + Some(pending(self.declare_token_increment(send, increments)?)) + } + + /// `TOKEN_INCREMENT_SCRIPT` on the request's post-call batch. A failed script falls back to + /// `increment_many` for these increments once the batch has gone out. `false` without Redis. + pub fn token_increment_post_call( + self: &Arc, + request: &RequestRedis, + increments: Vec, + ) -> bool { + let Some(declared) = self.declare_token_increment(Send::PostCall(request), &increments) + else { + return false; + }; + let store = Arc::clone(self); + request.spawn(async move { + let Err(error) = declared.settled().await else { + return; + }; + crate::failure_log::redis_failure!( + &error, + "TTL preservation failed, falling back to regular pipeline" + ); + store.increment_many(&increments).await; + }); + true + } + + /// `PARALLEL_RELEASE_SCRIPT` on the request's post-call batch; the caller has already freed + /// the slot in memory. `false` without Redis. + pub fn parallel_release_post_call( + &self, + request: &RequestRedis, + gauges: &[Scope], + slot_id: &str, + ) -> bool { + let Some(declared) = self.declare_release(Send::PostCall(request), gauges, slot_id) else { + return false; + }; + request.spawn(async move { + if let Err(error) = declared.settled().await { + crate::failure_log::redis_failure!( + &error, + "parallel_release_script failed, the slot stays released in memory only" + ); + } + }); + true + } + + /// `DualCache.async_increment_cache`. A Redis failure answers with the memory count, and an + /// entry memory cannot add to with `None`. + pub async fn increment(&self, increment: &Increment) -> Option { + self.cache + .increment(&increment.counter.key(), increment.amount, increment.ttl) + .await + } + + /// `DualCache.async_increment_cache_pipeline`, whose failures only log. + pub async fn increment_many(&self, increments: &[Increment]) { + self.cache.increment_pipeline(&pipeline(increments)).await; + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, the Redis increments on + /// the request's post-call batch. + pub fn increment_many_post_call(&self, request: &RequestRedis, increments: &[Increment]) { + self.cache + .increment_pipeline_post_call(request, &pipeline(increments)); + } + + fn declare_release( + &self, + send: Send<'_>, + gauges: &[Scope], + slot_id: &str, + ) -> Option>> { + Some(scripts::parallel_release( + &self.batch(send)?, + gauges, + slot_id, + )) + } + + fn declare_token_increment( + &self, + send: Send<'_>, + increments: &[Increment], + ) -> Option>> { + Some(scripts::token_increment(&self.batch(send)?, increments)) + } + + fn batch(&self, send: Send<'_>) -> Option { + self.cache.redis().batch(send) + } +} + +fn pipeline(increments: &[Increment]) -> Vec { + increments + .iter() + .map(|increment| PipelineIncrement { + key: increment.counter.key(), + amount: increment.amount, + ttl: increment.ttl, + }) + .collect() +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs new file mode 100644 index 00000000000..04bdad38268 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs @@ -0,0 +1,303 @@ +use std::{sync::LazyLock, time::Duration}; + +use litellm_python_compat::repr::float_repr; +use litellm_redis_client::{ + Command, Declared, Error, RedisBatch, Script, Value, + commands::{self, unexpected_reply}, +}; + +use super::{ + keys::{Counter, Scope}, + lua, + types::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, + Increment, WindowCount, + }, +}; +use crate::Number; + +static BATCH_RATE_LIMITER: LazyLock', + "deprecation_source": "javascript:alert(1)", + } + ) + email = ModelDeprecationEmail( + model=model, + usage_model=model.litellm_model, + recipient_email="owner@example.com", + milestone=DeprecationEmailMilestone.THIRTY_DAYS, + key_count=3, + requests=12, + team_names=("Research",), + start_date=date(2026, 10, 3), + end_date=NOW.date(), + ) + rendered = render_model_deprecation_email(email, base_url="https://gateway.example") + assert "/public/assets/logos/litellm_logo.png" in rendered.html + assert "litellm-listing.s3.amazonaws.com" not in rendered.html + assert "