scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge remote-tracking branch 'origin/main' into litellm_bedrock_embedding_region

Browse source
This commit is contained in:
Krishnatejaswi S 2026-11-04 15:12:34 +05:30
commit 00f5408b8c
No known key found for this signature in database
GPG key ID: 5F7D3E5324DAC64E
1986 changed files with 179367 additions and 55982 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

77
.agents/skills/database-access/SKILL.md
View file

@ -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/<feature>.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

101
.agents/skills/database-access/references/python-fastapi.md Normal file
View file

@ -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/)

119
.agents/skills/database-access/references/python-repository.md Normal file
View file

@ -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/<feature>.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/<feature>.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/<feature>_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 `<Feature>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.<table>` 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/<feature>.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

4
.agents/skills/database-access/references/rust-axum.md Normal file
View file

@ -0,0 +1,4 @@
# Rust and Axum
- This should use same `crates/db` what is exposed to Python world.
- Details are TBD.

7
.circleci/AGENTS.md Normal file
View file

@ -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

58
.circleci/config.yml
View file

@ -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
@ -559,7 +571,7 @@ jobs:
resource_class: medium
working_directory: ~/project
steps:
- checkout
- checkout_over_https
- run:
name: Generate LiteLLM master key
command: |
@ -659,7 +671,7 @@ jobs:
resource_class: large
working_directory: ~/project
steps:
- checkout
- checkout_over_https
- run:
name: Generate LiteLLM master key
command: |
@ -740,7 +752,7 @@ jobs:
resource_class: large
working_directory: ~/project
steps:
- checkout
- checkout_over_https
- skip_if_unrelated_changes:
category: client
- setup_google_dns
@ -794,7 +806,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 +952,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 +1078,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:
@ -1122,7 +1135,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 +1211,7 @@ jobs:
resource_class: medium
working_directory: ~/project
steps:
- checkout
- checkout_over_https
- skip_if_unrelated_changes
- attach_workspace:
at: ~/project
@ -1371,7 +1385,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 +1433,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 +1508,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

2
.circleci/scripts/run_integration.sh
View file

@ -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

16
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -4,9 +4,21 @@
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 -->
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
## Handwritten 1-3 sentence description
<!-- The author writes a plain summary in their own words before requesting review.
Wrap newly AI-written summary text in the make-me-human-readable tags below, including backticks.
The author reviews or rewrites the draft and removes the tags.
When updating an existing PR, preserve the author's wording and removed tags.
Changes elsewhere in the body do not make an unchanged summary an AI draft again.
Add tags again only around newly AI-written replacement summary text, never merely to match this template.
This preserves author review across later updates (2026-11-02, #49779).
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
No bullets, tables, or code names here. Everything else goes inside the Details block below -->
`<make-me-human-readable>`
<blah>
`</make-me-human-readable>`
<details>
<summary>Details</summary>

4
.github/workflows/compat-matrix-image.yml vendored
View file

@ -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:
@ -36,7 +36,7 @@ jobs:
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)"
version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)"
/opt/litellm/tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" /tmp/claude-cli
/tmp/claude-cli/claude --version
'

2
.github/workflows/create-rc-branch.yml vendored
View file

@ -2,7 +2,7 @@ name: Create RC Branch
on:
schedule:
- cron: "0 3 * * 5"
- cron: "0 3 * * 1"
timezone: "America/Los_Angeles"
workflow_dispatch:

30
.github/workflows/deploy-dev-on-main.yml vendored
View file

@ -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

1
.github/workflows/osv-scan.yml vendored
View file

@ -5,6 +5,7 @@ on:
branches:
- main
- "litellm_**"
- "rc/**"
schedule:
- cron: "23 6 * * *"
workflow_dispatch:

2
.github/workflows/test-linting.yml vendored
View file

@ -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 \

10
.github/workflows/test-rust-db.yml vendored
View file

@ -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

20
.github/workflows/test-rust.yml vendored
View file

@ -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

14
.github/workflows/test-unit.yml vendored
View file

@ -73,9 +73,21 @@ jobs:
timeout-minutes: 5
uses: ./.github/actions/cache-cargo-build
- name: Set up sccache
if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true'
uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a
with:
version: "v0.18.0"
- name: Build editable Rust bridge
if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true'
timeout-minutes: 15
env:
RUSTC_WRAPPER: sccache
SCCACHE_GHA_ENABLED: "true"
SCCACHE_GHA_VERSION: litellm-ci-v1
SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }}
CARGO_INCREMENTAL: "0"
run: |
mkdir -p rust-bridge-dist
uv run --no-project --with maturin==1.15.0 python -c 'import maturin; maturin.build_editable("rust-bridge-dist")'
@ -429,6 +441,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
@ -565,6 +578,7 @@ jobs:
tests/unit/proxy/rerank_endpoints
tests/unit/proxy/realtime_endpoints
tests/unit/proxy/ui_crud_endpoints
tests/unit/proxy/deps
tests/unit/proxy/config_resolvers
tests/unit/proxy/utils
python-version: "3.12"

85
.github/workflows/weekly_load_anomaly.yml vendored
View file

@ -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

4
AGENTS.md
View file

@ -41,6 +41,8 @@ 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 <tip SHA>`, 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
@ -90,6 +92,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: <what bounds it>`
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):

2
CONTRIBUTING.md
View file

@ -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

3
README.md
View file

@ -444,6 +444,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 +697,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**

245
cookbook/gpt_realtime_translate.py Normal file
View file

@ -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())

31
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -10,6 +10,7 @@ from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
from pydantic import TypeAdapter, ValidationError
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
@ -50,6 +51,7 @@ def _error_names_resource(error: Exception, resource_id: str) -> bool:
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
_FILE_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
@ -550,6 +552,10 @@ class CheckBatchCost:
llm_provider: str,
bare_model_name: str,
) -> Optional[Tuple[str, str]]:
stored_model_id: Final = self._get_stored_model_id(job)
if stored_model_id is not None and self._batch_deployment_exists(stored_model_id):
return stored_model_id, job.unified_object_id
deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider)
if deployment_id is None:
verbose_proxy_logger.info(
@ -662,23 +668,36 @@ class CheckBatchCost:
)
@staticmethod
def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
from litellm.types.utils import LiteLLMBatch
def _parse_file_object(job: "ManagedObjectRow") -> Optional[dict[str, object]]:
file_object: Final = job.file_object
try:
parsed_file_object: Final[object] = (
return _FILE_OBJECT_ADAPTER.validate_python(
json.loads(file_object) if isinstance(file_object, str) else file_object
)
except (json.JSONDecodeError, ValueError):
except (json.JSONDecodeError, ValueError, ValidationError):
return None
if not isinstance(parsed_file_object, dict):
@classmethod
def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]:
from litellm.types.utils import LiteLLMBatch
parsed_file_object: Final = cls._parse_file_object(job)
if parsed_file_object is None:
return None
try:
return LiteLLMBatch.model_validate(parsed_file_object).input_file_id
except Exception:
return None
@classmethod
def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]:
"""The deployment that created the batch, written into file_object by the managed
files hook at create time. Older rows have no such key and fall back to matching
the bare model name against the router."""
parsed_file_object: Final = cls._parse_file_object(job)
model_id: Final = parsed_file_object.get("model_id") if parsed_file_object else None
return model_id if isinstance(model_id, str) and model_id else None
async def _track_completed_batch_cost(
self,
job: "ManagedObjectRow",

1
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -200,6 +200,7 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
unpriced: Final = _without_read_price(billed.response)
cost: Final = logging_obj.response_cost_calculator(result=unpriced)
end_time: Final = datetime.now()
logging_obj.capture_usage_before_logging(unpriced, logging_obj.start_time, end_time)
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=logging_obj.async_success_handler(
result=unpriced, start_time=logging_obj.start_time, end_time=end_time

234
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -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 (
@ -94,6 +97,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,
@ -374,6 +378,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: ...
@ -388,6 +398,134 @@ class _RouterFileCallKwargs(TypedDict, total=False):
custom_llm_provider: ReadOnly[str | None]
_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", "_litellm_internal_model_credentials"})
@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:
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
delete_data: Final = {
**{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
await llm_router.afile_delete(
model=model_id,
file_id=provider_file_id,
**cast(_RouterFileCallKwargs, delete_data),
)
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:
return prisma_client.db
@ -479,6 +617,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 +655,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 +859,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 +890,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 +929,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},
@ -1989,16 +2165,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 +2183,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()
@ -2170,6 +2355,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 +2832,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(

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20261104000000_skill_zip_upload/migration.sql Normal file
View file

@ -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
);

9
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1615,9 +1615,18 @@ model LiteLLM_ClaudeCodePluginTable {
updated_at DateTime? @default(now()) @updatedAt
created_by String?
uploaded_archive LiteLLM_ClaudeCodePluginUpload?
@@map("LiteLLM_ClaudeCodePluginTable")
}
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.

43
litellm-rust/.agents/skills/rust-redaction/SKILL.md Normal file
View file

@ -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 `"<redacted>"`, `"<configured>"` 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<SecretValue>,
pub password: Option<SecretValue>,
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<dyn Trait>` 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

2
litellm-rust/AGENTS.md
View file

@ -1,5 +1,7 @@
# Rust workspace rules
For keeping secrets out of Debug output, follow [.agents/skills/rust-redaction/SKILL.md](.agents/skills/rust-redaction/SKILL.md)
For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md)
For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md)

137
litellm-rust/Cargo.lock generated
View file

@ -3380,6 +3380,7 @@ dependencies = [
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-auth-types",
"litellm-secret-value",
]
[[package]]
@ -3416,6 +3417,7 @@ dependencies = [
"azure_identity",
"litellm-auth-types",
"litellm-error",
"litellm-secret-value",
"moka",
"rstest",
"serde_json",
@ -3435,6 +3437,7 @@ dependencies = [
"http 1.4.2",
"litellm-auth-types",
"litellm-error",
"litellm-secret-value",
"moka",
"rstest",
"serde",
@ -3450,10 +3453,10 @@ name = "litellm-auth-types"
version = "0.1.0"
dependencies = [
"litellm-error",
"litellm-secret-value",
"rstest",
"serde",
"serde_json",
"subtle",
"thiserror 2.0.19",
"tokio",
"veil",
@ -3523,6 +3526,7 @@ dependencies = [
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"litellm-secret-value",
"percent-encoding",
"reqwest",
"rstest",
@ -3565,37 +3569,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 +3576,8 @@ dependencies = [
"litellm-cache",
"litellm-cache-gcs",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-http",
"py_literal",
"redis",
"redis-test",
"rstest",
"serde",
"serde_json",
@ -3646,32 +3616,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 +3644,9 @@ dependencies = [
name = "litellm-config"
version = "0.1.0"
dependencies = [
"litellm-auth-types",
"litellm-error",
"litellm-router-types",
"litellm-secret-value",
"rstest",
"serde",
"serde_json",
@ -3804,7 +3756,6 @@ dependencies = [
"envy",
"futures-util",
"http-body-util",
"litellm-auth-types",
"litellm-circuit-breaker",
"litellm-config",
"litellm-error",
@ -3817,6 +3768,7 @@ dependencies = [
"litellm-llms",
"litellm-redis-client",
"litellm-redis-testing",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"rstest",
@ -3839,9 +3791,9 @@ dependencies = [
"axum",
"axum-login",
"futures-util",
"litellm-auth-types",
"litellm-config",
"litellm-error",
"litellm-secret-value",
"litellm-secrets",
"rstest",
"serde",
@ -3879,6 +3831,7 @@ dependencies = [
"litellm-llms",
"litellm-llms-types",
"litellm-router",
"litellm-secret-value",
"litellm-secrets",
"rstest",
"serde",
@ -3895,9 +3848,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 +3865,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 +3895,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 +3917,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 +4047,7 @@ dependencies = [
"litellm-inference-testing",
"litellm-llms",
"litellm-llms-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"rstest",
@ -4124,6 +4078,7 @@ dependencies = [
"litellm-llms",
"litellm-llms-types",
"litellm-router-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"reqwest",
@ -4154,6 +4109,7 @@ dependencies = [
"litellm-inference-testing",
"litellm-llms",
"litellm-llms-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"mime_guess",
@ -4188,6 +4144,7 @@ dependencies = [
"litellm-inference-testing",
"litellm-llms",
"litellm-llms-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"reqwest",
@ -4207,6 +4164,7 @@ dependencies = [
"litellm-auth",
"litellm-http",
"litellm-inference",
"litellm-secret-value",
"litellm-secrets",
"tokio",
]
@ -4254,6 +4212,7 @@ dependencies = [
"litellm-llms-types",
"litellm-python-compat",
"litellm-router-types",
"litellm-secret-value",
"litellm-secrets",
"reqwest",
"rstest",
@ -4328,8 +4287,6 @@ dependencies = [
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core-utils",
@ -4347,6 +4304,7 @@ dependencies = [
"litellm-llms-types",
"litellm-model-catalog",
"litellm-router-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-secrets-aws",
"litellm-secrets-types",
@ -4357,7 +4315,6 @@ dependencies = [
"prost",
"pyo3",
"pyo3-async-runtimes",
"redis",
"reqwest",
"rstest",
"serde",
@ -4396,6 +4353,7 @@ dependencies = [
"futures-util",
"litellm-circuit-breaker",
"litellm-redis-testing",
"litellm-secret-value",
"redis",
"rstest",
"thiserror 2.0.19",
@ -4432,12 +4390,23 @@ name = "litellm-router-types"
version = "0.1.0"
dependencies = [
"litellm-auth-types",
"litellm-secret-value",
"rstest",
"serde",
"serde_json",
"serde_with",
]
[[package]]
name = "litellm-secret-value"
version = "0.1.0"
dependencies = [
"rstest",
"serde",
"subtle",
"veil",
]
[[package]]
name = "litellm-secrets"
version = "0.1.0"
@ -4451,6 +4420,7 @@ dependencies = [
"litellm-error",
"litellm-http",
"litellm-python-compat",
"litellm-secret-value",
"litellm-secrets-aws",
"litellm-secrets-azure",
"litellm-secrets-cyberark",
@ -4479,6 +4449,7 @@ dependencies = [
"base64 0.22.1",
"litellm-auth-aws",
"litellm-core-utils",
"litellm-secret-value",
"litellm-secrets-types",
"litellm-tracing",
"rstest",
@ -4498,6 +4469,7 @@ dependencies = [
"litellm-auth-types",
"litellm-core-utils",
"litellm-http",
"litellm-secret-value",
"litellm-secrets-types",
"percent-encoding",
"reqwest",
@ -4519,6 +4491,7 @@ dependencies = [
"base64 0.22.1",
"litellm-core-utils",
"litellm-http",
"litellm-secret-value",
"litellm-secrets-types",
"litellm-tracing",
"moka",
@ -4549,6 +4522,7 @@ dependencies = [
"litellm-auth-types",
"litellm-core-utils",
"litellm-http",
"litellm-secret-value",
"litellm-secrets-types",
"moka",
"percent-encoding",
@ -4570,6 +4544,7 @@ dependencies = [
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-secret-value",
"litellm-secrets-types",
"reqwest",
"rstest",
@ -4589,7 +4564,7 @@ dependencies = [
name = "litellm-secrets-types"
version = "0.1.0"
dependencies = [
"litellm-auth-types",
"litellm-secret-value",
"moka",
"rstest",
"serde",
@ -5752,17 +5727,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 +5900,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"
@ -6424,15 +6376,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"

4
litellm-rust/Cargo.toml
View file

@ -37,6 +37,7 @@ litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
litellm-framer = { path = "crates/framer" }
litellm-auth = { path = "crates/auth" }
litellm-auth-types = { path = "crates/auth-types" }
litellm-secret-value = { path = "crates/secret-value" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
litellm-auth-gcp = { path = "crates/auth-gcp" }
@ -56,11 +57,9 @@ 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-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" }
@ -96,7 +95,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"] }

1
litellm-rust/crates/auth-azure/Cargo.toml
View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-secret-value.workspace = true
litellm-auth-types.workspace = true
moka.workspace = true

6
litellm-rust/crates/auth-azure/src/native.rs
View file

@ -13,7 +13,8 @@ use azure_identity::{
use sha2::{Digest, Sha256};
use litellm_auth_types::{Error, ErrorDetail};
use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced};
use litellm_auth_types::{InputSource, ResolvedCredential, Sourced};
use litellm_secret_value::SecretValue;
use super::credential_provider_cache::{
AzureCredentialProviderCache, AzureCredentialProviderCacheKey,
@ -498,7 +499,8 @@ mod tests {
use azure_core::{Bytes, Result};
use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
use litellm_auth_types::{InputSource, SecretValue, Sourced};
use litellm_auth_types::{InputSource, Sourced};
use litellm_secret_value::SecretValue;
fn deployment<T>(value: T) -> Sourced<T> {
Sourced::new(value, InputSource::Deployment)

9
litellm-rust/crates/auth-azure/src/resolve.rs
View file

@ -1,7 +1,9 @@
use litellm_auth_types::{
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
SecretValue, Sourced, TokenProviderHandle,
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced,
TokenProviderHandle,
};
use litellm_secret_value::SecretValue;
use litellm_auth_types::{Error, ErrorDetail};
use std::future::Future;
use std::pin::Pin;
@ -500,8 +502,9 @@ mod tests {
use litellm_auth_types::ResolvedCredential;
use litellm_auth_types::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef,
CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced,
CredentialResolver, CredentialResolverHandle, InputSource, Sourced,
};
use litellm_secret_value::SecretValue;
#[derive(Debug)]
struct FileResolver;

4
litellm-rust/crates/auth-azure/src/types.rs
View file

@ -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;

1
litellm-rust/crates/auth-gcp/Cargo.toml
View file

@ -9,6 +9,7 @@ repository.workspace = true
google-sdk = ["dep:google-cloud-auth", "dep:http"]
[dependencies]
litellm-secret-value.workspace = true
litellm-auth-types.workspace = true
base64.workspace = true

2
litellm-rust/crates/auth-gcp/src/auth.rs
View file

@ -258,7 +258,7 @@ fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
#[cfg(test)]
mod tests {
use litellm_auth_types::SecretValue;
use litellm_secret_value::SecretValue;
use super::*;

4
litellm-rust/crates/auth-gcp/src/config.rs
View file

@ -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;

2
litellm-rust/crates/auth-types/Cargo.toml
View file

@ -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

5
litellm-rust/crates/auth-types/src/credential.rs
View file

@ -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;

2
litellm-rust/crates/auth-types/src/lib.rs
View file

@ -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};

3
litellm-rust/crates/auth-types/src/policy.rs
View file

@ -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,

3
litellm-rust/crates/auth-types/src/token.rs
View file

@ -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 {

6
litellm-rust/crates/auth-types/tests/token.rs
View file

@ -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 {

3
litellm-rust/crates/auth/Cargo.toml
View file

@ -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

3
litellm-rust/crates/auth/tests/facade.rs
View file

@ -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,

1
litellm-rust/crates/cache-gcs/Cargo.toml
View file

@ -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

4
litellm-rust/crates/cache-gcs/src/token.rs
View file

@ -1,7 +1,9 @@
use std::{future::Future, pin::Pin};
use litellm_auth_gcp::{VertexAuth, VertexConfig};
use litellm_auth_types::{InputSource, SecretValue, Sourced};
use litellm_auth_types::{InputSource, Sourced};
use litellm_secret_value::SecretValue;
use litellm_cache::Error;
pub trait TokenSource: Send + Sync + 'static {

19
litellm-rust/crates/cache-redis-semantic/Cargo.toml
View file

@ -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

404
litellm-rust/crates/cache-redis-semantic/src/cache.rs
View file

@ -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<u8>,
prompt: &str,
vector: &[f32],
ttl: Option<Duration>,
) -> 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<SemanticLookup<Vec<u8>>, 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::<redis::Value>(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<E, S, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
codec: S,
inner: Arc<Inner>,
}
impl<E: Embedder, S: CacheCodec> RedisSemanticCache<E, S> {
pub fn new(
url: &str,
embedder: E,
codec: S,
config: RedisSemanticConfig,
) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
codec,
inner: Arc::new(Inner::new(config, timestamp)),
})
}
}
impl<E, S, C> RedisSemanticCache<E, S, C>
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<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
Ok(SemanticLookup {
value: lookup
.value
.map(|bytes| self.codec.decode(&bytes))
.transpose()?,
similarity: lookup.similarity,
})
}
}
impl<E, S, C> BaseCache for RedisSemanticCache<E, S, C>
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<Duration> {
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<Option<Self::Value>, 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<Option<Self::Value>, 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<E, S, C> SemanticCache for RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, 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<SemanticLookup<Self::Value>, 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<u8> {
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)
}

8
litellm-rust/crates/cache-redis-semantic/src/config.rs
View file

@ -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,
}

209
litellm-rust/crates/cache-redis-semantic/src/index.rs
View file

@ -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 `<name>_isolated`, recreated when that one is stale too.
pub(crate) struct Index {
name: String,
resolved: OnceLock<String>,
}
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<String, Error> {
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<String, Error> {
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<Option<bool>, Error> {
let info = match redis::cmd("FT.INFO")
.arg(name)
.query::<redis::Value>(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<String>,
field_type: Option<String>,
dim: Option<f64>,
data_type: Option<String>,
distance_metric: Option<String>,
}
fn attribute(value: &redis::Value) -> Option<Attribute> {
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::<Vec<_>>();
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"))
})
}

7
litellm-rust/crates/cache-redis-semantic/src/lib.rs
View file

@ -1,7 +0,0 @@
mod cache;
mod config;
mod index;
mod reply;
pub use cache::RedisSemanticCache;
pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig};

57
litellm-rust/crates/cache-redis-semantic/src/reply.rs
View file

@ -1,57 +0,0 @@
pub(crate) fn string_value(value: &redis::Value) -> Option<String> {
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<f64> {
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<String> {
field_value(fields, name).and_then(string_value)
}
pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
field_value(fields, name).and_then(number_value)
}
pub(crate) fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
match field_value(fields, name)? {
redis::Value::BulkString(bytes) => Some(bytes.clone()),
redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
_ => None,
}
}

778
litellm-rust/crates/cache-redis-semantic/tests/cache.rs
View file

@ -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<FakeEmbedder, JsonCodec<Value>, 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<Value>) -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(Value::Array(messages)),
..Default::default()
}
}
fn cache(commands: Vec<MockCmd>, 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<u8> {
serde_json::to_vec(value).unwrap()
}
fn vector_bytes(vector: &[f32]) -> Vec<u8> {
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::<redis::Value, _>(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>) -> 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 {
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<redis::Value>,
) -> 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<u8>) -> 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<Duration>,
#[case] expire: Option<u64>,
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<Option<Value>, 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<Option<Value>, 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<redis::Value>,
#[case] value: Option<Value>,
#[case] similarity: Option<f64>,
#[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<Duration>) {
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<String> {
std::env::var("LITELLM_REDIS_STACK_URL").ok()
}
fn live_cache(
url: &str,
index_name: &str,
prompt: &str,
vector: Vec<f32>,
) -> RedisSemanticCache<FakeEmbedder, JsonCodec<Value>> {
RedisSemanticCache::new(
url,
FakeEmbedder::new(&[(prompt, vector.as_slice())]),
JsonCodec::<Value>::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<String>, 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<String>, 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}"
);
}

63
litellm-rust/crates/cache-redis-semantic/tests/contract.rs
View file

@ -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<PreparedEmbedding, JsonCodec<Value>, 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;
}

302
litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs
View file

@ -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<Mutex<Vec<(String, Option<Value>)>>>;
/// 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<String, Vec<f32>>,
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<Vec<f32>, 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<Vec<f32>, Error> {
self.embed(prompt, metadata)
}
}
struct FakeIndex {
prefix: Vec<u8>,
dims: usize,
vector_field: String,
}
#[derive(Default)]
struct SearchState {
indexes: HashMap<String, FakeIndex>,
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
}
/// 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<Mutex<SearchState>>,
}
impl FakeSearch {
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
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::<Vec<_>>();
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<redis::Value> {
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<Vec<redis::Value>> {
let replies = parse_commands(commands)
.into_iter()
.map(|args| self.run(args))
.collect::<redis::RedisResult<Vec<_>>>()?;
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<f32> {
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::<f64>();
let norm = |vector: &[f32]| {
vector
.iter()
.map(|value| f64::from(*value).powi(2))
.sum::<f64>()
.sqrt()
};
dot / (norm(left) * norm(right))
}
/// Splits a packed RESP request into each command's arguments.
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
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::<usize>().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
}

20
litellm-rust/crates/cache-redis/Cargo.toml
View file

@ -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

192
litellm-rust/crates/cache-redis/src/cache.rs
View file

@ -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<S, C = redis::Connection> {
pub(crate) connections: Arc<Connections<C>>,
pub(crate) default_ttl: Duration,
pub(crate) codec: S,
pub(crate) namespace: Option<String>,
pub(crate) topology: RedisTopology,
/// The server's major version, read from `INFO` once, like Python's `redis_version`.
pub(crate) major_version: Arc<OnceLock<u32>>,
}
impl<S: CacheCodec> RedisCache<S> {
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
Self::connect(url, &RedisTopology::Standalone, default_ttl, codec)
}
pub fn connect(
url: &str,
topology: &RedisTopology,
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
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<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(connection: C, default_ttl: Option<Duration>, 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<String>) -> 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<String> {
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<Duration>) -> u64 {
ttl_seconds(ttl.unwrap_or(self.default_ttl))
}
pub(crate) fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
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<T, F>(&self, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
Connections::run_blocking(Arc::clone(&self.connections), operation).await
}
pub(crate) fn namespaced_pattern(&self) -> Result<String, Error> {
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<Option<S::Value>, 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<BatchEntry<S::Value>, 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<u32> {
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::<f64>().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)
}

109
litellm-rust/crates/cache-redis/src/claim.rs
View file

@ -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<Option<Vec<u8>>, 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<S: CacheCodec>(
connection: &mut ConnectionRef<'_>,
codec: &S,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
ttl: u64,
) -> Result<S::Value, Error>
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::<bool>(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<S, C> ClaimCache for RedisCache<S, C>
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<S::Value, Error> {
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<S::Value>,
context: ExactCacheContext,
) -> Result<S::Value, Error> {
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
}
}

513
litellm-rust/crates/cache-redis/src/connection.rs
View file

@ -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<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
}
impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
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<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + 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<Self, Error> {
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<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, 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<M, T>(pool: &r2d2::Pool<M>)
where
M: r2d2::ManageConnection<Connection = PooledConnection<T>>,
{
let mut idle = Vec::new();
while let Some(mut connection) = pool.try_get() {
connection.failed = true;
idle.push(connection);
}
}
pub(crate) struct PooledConnection<C> {
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<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(Error::unavailable)
}
}
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection<redis::Connection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
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::<String>(&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<Self, Error> {
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::<Result<Vec<_>, _>>()?;
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<ConnectionInfo, Error> {
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<ClusterConnection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
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::<String>(&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<redis::Value> {
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<Vec<redis::Value>> {
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<T: redis::FromRedisValue>(
&mut self,
pipeline: &redis::Pipeline,
) -> Result<T, Error> {
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<String>) -> Result<bool, Error>,
) -> Result<(), Error> {
let pages = match self {
Self::Node(connection) => {
let page = scan_command(0, pattern, count)
.query::<ScanPage>(*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<bool, redis::RedisError> {
let command = redis::cmd("PING");
match self {
Self::Node(connection) => command
.query::<String>(*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<String, Error> {
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::<String>(node).map_err(Error::unavailable)?,
redis::from_redis_value::<String>(reply).map_err(Error::unavailable)?,
))
})
.collect::<Result<Vec<(String, String)>, Error>>()?;
replies.sort();
Ok(replies
.into_iter()
.map(|(_, reply)| reply)
.collect::<Vec<_>>()
.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<ScanPage, Error> {
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<String>);
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, 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::<String>(node).map_err(Error::unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?;
let page = redis::from_redis_value::<ScanPage>(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<Vec<redis::Value>>);
impl redis::ConnectionLike for Replies {
fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult<redis::Value> {
Err((redis::ErrorKind::Client, "replies hold a pipeline only").into())
}
fn req_packed_commands(
&mut self,
_: &[u8],
_: usize,
_: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
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
}
}

205
litellm-rust/crates/cache-redis/src/counter.rs
View file

@ -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<S, C> CounterCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
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<f64, Error> {
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<IncrementOperation>,
) -> Result<Vec<f64>, 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<f64, Error> {
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<S, C> CountReadCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, 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<String>) -> Result<Vec<Option<i64>>, 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<Option<i64>, Error> {
redis::from_redis_value(value).map_err(Error::invalid_entry)
}
impl<S, C> BoundedCounterCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result<i64, Error> {
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<i64, Error> {
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<Duration>,
) -> Result<f64, Error> {
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<i64, Error> {
redis::cmd("EVAL")
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable)
}

63
litellm-rust/crates/cache-redis/src/keys.rs
View file

@ -1,63 +0,0 @@
use std::time::Duration;
use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache};
use crate::cache::RedisCache;
impl<S, C> TtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
let key = self.namespaced_key(key);
let ttl = self
.run(move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
.map_err(Error::unavailable)
})
.await?;
Ok(u64::try_from(ttl).ok().map(Duration::from_secs))
}
}
impl<S, C> RefreshTtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_refresh_ttl(&self, key: &str, ttl: Option<Duration>) -> Result<bool, Error> {
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<S, C> ScanCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, 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
}
}

15
litellm-rust/crates/cache-redis/src/lib.rs
View file

@ -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};

86
litellm-rust/crates/cache-redis/src/lifecycle.rs
View file

@ -1,86 +0,0 @@
use litellm_cache::{
CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache,
DisconnectCache, Error, PingCache,
};
use crate::{cache::RedisCache, topology::RedisTopology};
impl<S, C> PingCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn sync_ping(&self) -> Result<bool, Error> {
self.execute(|connection| connection.ping().map_err(Error::unavailable))
}
async fn ping(&self) -> Result<bool, Error> {
self.run(|connection| connection.ping().map_err(Error::unavailable))
.await
}
}
impl<S, C> ConnectionCache for RedisCache<S, C>
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<CacheConnectionResult, Error> {
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<S, C> DisconnectCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn disconnect(&self) -> Result<(), Error> {
self.connections.disconnect();
Ok(())
}
}
impl<S, C> ClientInfoCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type ClientList = String;
type Info = String;
fn client_list(&self) -> Result<Self::ClientList, Error> {
self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
}
fn info(&self) -> Result<Self::Info, Error> {
self.execute(|connection| connection.node_text(&redis::cmd("INFO")))
}
}

226
litellm-rust/crates/cache-redis/src/queue.rs
View file

@ -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<RedisArg>;
pub type RedisLpopOperation = PopOperation;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RedisLpopResult {
Missing,
Value(Vec<u8>),
Values(Vec<Vec<u8>>),
}
impl<S, C> SetCache for RedisCache<S, C>
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<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
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<S, C> QueueCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type QueueValue = RedisArg;
type PopResult = RedisLpopResult;
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
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<Self::QueueValue>,
max_len: usize,
) -> Result<usize, Error> {
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<RedisRpushOperation>,
) -> Result<Vec<usize>, 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<usize>) -> Result<Self::PopResult, Error> {
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::<redis::Value>(connection)
.map_err(Error::unavailable)
})
.await?;
lpop_result(value, count.is_some())
}
async fn async_lpop_pipeline(
&self,
operations: Vec<RedisLpopOperation>,
) -> Result<Vec<Self::PopResult>, 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::<Vec<_>>();
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::<Vec<redis::Value>>(&pipeline))
.await?
.into_iter()
.zip(multiple)
.map(|(value, multiple)| lpop_result(value, multiple))
.collect()
}
}
fn lpop(key: String, count: Option<usize>) -> 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<Vec<u8>, Error> {
match value {
redis::Value::BulkString(bytes) => Ok(bytes),
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
_ => Err(Error::InvalidEntry(None)),
}
}
impl<S, C> RedisCache<S, C>
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<RedisLpopResult, Error> {
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::<Vec<redis::Value>>(&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<RedisLpopResult, Error> {
match value {
redis::Value::Nil => Ok(RedisLpopResult::Missing),
redis::Value::Array(values) if multiple => values
.into_iter()
.map(redis_bytes)
.collect::<Result<Vec<_>, _>>()
.map(RedisLpopResult::Values),
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
_ => Err(Error::InvalidEntry(None)),
}
}

136
litellm-rust/crates/cache-redis/src/script.rs
View file

@ -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<u8>),
Integer(i64),
Float(f64),
}
impl From<&str> for RedisArg {
fn from(value: &str) -> Self {
Self::Bytes(value.as_bytes().to_vec())
}
}
impl From<String> for RedisArg {
fn from(value: String) -> Self {
Self::Bytes(value.into_bytes())
}
}
impl From<Vec<u8>> for RedisArg {
fn from(value: Vec<u8>) -> Self {
Self::Bytes(value)
}
}
impl From<i64> for RedisArg {
fn from(value: i64) -> Self {
Self::Integer(value)
}
}
impl From<f64> for RedisArg {
fn from(value: f64) -> Self {
Self::Float(value)
}
}
impl redis::ToRedisArgs for RedisArg {
fn write_redis_args<W>(&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<C> {
connections: Arc<Connections<C>>,
namespace: Option<String>,
source: String,
}
impl<C> CacheScript for RedisScript<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Argument = RedisArg;
type Output = redis::Value;
async fn invoke(
&self,
keys: Vec<String>,
arguments: Vec<Self::Argument>,
) -> Result<Self::Output, Error> {
let keys = keys
.into_iter()
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
.collect::<Vec<_>>();
let source = self.source.clone();
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
eval(connection, &source, keys, arguments)
})
.await
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub async fn async_eval(
&self,
script: String,
keys: Vec<String>,
arguments: Vec<RedisArg>,
) -> Result<redis::Value, Error> {
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<String>,
arguments: Vec<RedisArg>,
) -> Result<redis::Value, Error> {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(Error::unavailable)
}
impl<S, C> ScriptCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Script = RedisScript<C>;
fn async_register_script(&self, source: String) -> Self::Script {
RedisScript {
connections: Arc::clone(&self.connections),
namespace: self.namespace.clone(),
source,
}
}
}

232
litellm-rust/crates/cache-redis/src/store.rs
View file

@ -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<S, C> BaseCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
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<Option<Self::Value>, 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<Option<Self::Value>, 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<S, C> TtlPipelineCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_set_cache_pipeline_with_ttls(
&self,
entries: Vec<(String, Self::Value, Option<Duration>)>,
) -> 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<S, C> BatchCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_cache(
&self,
keys: &[String],
_: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, 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<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, 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<String>,
) -> Result<Vec<redis::Value>, Error> {
redis::cmd("MGET")
.arg(keys)
.query(connection)
.map_err(Error::unavailable)
}
impl<S, C> DeleteCache for RedisCache<S, C>
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<S, C> BulkDeleteCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
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<S, C> FlushCache for RedisCache<S, C>
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<S, C> FlushAllCache for RedisCache<S, C>
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)
})
}

14
litellm-rust/crates/cache-redis/src/topology.rs
View file

@ -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<RedisNode>,
},
}

1164
litellm-rust/crates/cache-redis/tests/cache.rs
View file

File diff suppressed because it is too large Load diff

563
litellm-rust/crates/cache-redis/tests/cluster.rs
View file

@ -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<JsonCodec<f64>>;
#[fixture]
fn cache(#[default("cache")] label: &str) -> Option<JsonCache> {
cluster_cache(label, Duration::from_secs(120), JsonCodec::new())
}
#[fixture]
fn counter(#[default("counter")] label: &str) -> Option<Counter> {
cluster_cache(label, Duration::from_secs(60), JsonCodec::new())
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
fn multi_slot_keys(count: usize) -> Vec<String> {
let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
let slots: HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
assert!(slots.len() > 1, "keys must span multiple slots");
keys
}
fn seconds(seconds: u64) -> Option<Duration> {
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<RedisNode>,
) {
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<JsonCache>,
) {
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<JsonCache>,
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<JsonCache>,
counter: Option<Counter>,
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<String> = 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::<Vec<_>>());
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<Vec<u8>> = (0..=index)
.take(2)
.map(|value| value.to_string().into_bytes())
.collect();
assert_eq!(values, expected);
}
other => panic!("queue {index}: {other:?}"),
}
}
let counters: Vec<String> = 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<f64> = (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<JsonCache>,
) {
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<JsonCache>,
#[from(cache)]
#[with("other")]
other: Option<JsonCache>,
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<String> = 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<JsonCache>) {
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<JsonCache>,
) {
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<Counter>,
#[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<Counter>,
#[with("claim")] cache: Option<JsonCache>,
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);
}

95
litellm-rust/crates/cache-redis/tests/contract.rs
View file

@ -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<C>(contract: Contract, cache: &JsonCache<C>)
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;
}
}
}
}

231
litellm-rust/crates/cache-redis/tests/support/mod.rs
View file

@ -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<Vec<u8>, Error> {
if *value > 127 {
return Err(Error::InvalidEntry(None));
}
Ok(vec![self.0, *value])
}
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
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<u8>, Vec<u8>>,
}
impl FakeRedis {
fn run(&mut self, command: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
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<Vec<Vec<u8>>> {
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<redis::Value> {
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<Vec<redis::Value>> {
let replies = commands(cmd)
.into_iter()
.map(|command| self.run(command))
.collect::<redis::RedisResult<Vec<_>>>()?;
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<C = redis::Connection> = RedisCache<JsonCodec<serde_json::Value>, C>;
pub fn fake_cache(namespace: &str) -> JsonCache<FakeRedis> {
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<RedisTopology> {
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<S: CacheCodec>(
label: &str,
default_ttl: Duration,
codec: S,
) -> Option<RedisCache<S>> {
let topology = cluster_topology()?;
Some(
RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec)
.expect("cluster connection")
.with_namespace(Some(unique_namespace(label))),
)
}

3
litellm-rust/crates/cache-response/Cargo.toml
View file

@ -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

72
litellm-rust/crates/cache-response/tests/connection.rs
View file

@ -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<dyn ConnectionProbe> =
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<dyn ConnectionProbe> = service.clone();
let exact: Arc<dyn ExactResponseCache> = 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.

97
litellm-rust/crates/cache-response/tests/response.rs
View file

@ -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<ResponseCache<InMemoryCache<CacheEntry>>>;
@ -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(

17
litellm-rust/crates/cache-response/tests/support/mod.rs
View file

@ -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<ResponseCacheCodec, MockRedisConnection>;
#[fixture]
pub fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
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<MockCmd>, namespace: Option<&str>) -> ResponseCache<MockedRedis> {
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)),
))
}

22
litellm-rust/crates/cache-valkey-semantic/Cargo.toml
View file

@ -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

245
litellm-rust/crates/cache-valkey-semantic/src/cache.rs
View file

@ -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<E, S, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
codec: S,
config: ValkeySemanticConfig,
index_dimension: Arc<Mutex<Option<usize>>>,
}
impl<E: Embedder, S: CacheCodec> ValkeySemanticCache<E, S> {
pub fn new(
url: &str,
embedder: E,
codec: S,
config: ValkeySemanticConfig,
) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
codec,
config,
index_dimension: Arc::new(Mutex::new(None)),
})
}
}
impl<E, S, C> ValkeySemanticCache<E, S, C>
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<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
Ok(SemanticLookup {
value: lookup
.value
.map(|bytes| self.codec.decode(&bytes))
.transpose()?,
similarity: lookup.similarity,
})
}
}
impl<E, S, C> ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec + Clone,
C: redis::ConnectionLike + Send + 'static,
{
/// The same index and connections behind a different embedder.
pub fn with_embedder<E2: Embedder>(&self, embedder: E2) -> ValkeySemanticCache<E2, S, C> {
ValkeySemanticCache {
connections: Arc::clone(&self.connections),
embedder,
codec: self.codec.clone(),
config: self.config.clone(),
index_dimension: Arc::clone(&self.index_dimension),
}
}
}
impl<E, S, C> BaseCache for ValkeySemanticCache<E, S, C>
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<Duration> {
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<Option<Self::Value>, 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<Option<Self::Value>, 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<E, S, C> SemanticCache for ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, 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<SemanticLookup<Self::Value>, 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)
}
}

8
litellm-rust/crates/cache-valkey-semantic/src/config.rs
View file

@ -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,
}

117
litellm-rust/crates/cache-valkey-semantic/src/index.rs
View file

@ -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<Mutex<Option<usize>>>,
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::<String>(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::<redis::Value>(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<usize> {
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::<Vec<_>>();
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()
})
})
}

7
litellm-rust/crates/cache-valkey-semantic/src/lib.rs
View file

@ -1,7 +0,0 @@
mod cache;
mod config;
mod index;
mod search;
pub use cache::ValkeySemanticCache;
pub use config::{DEFAULT_INDEX_NAME, ValkeySemanticConfig};

161
litellm-rust/crates/cache-valkey-semantic/src/search.rs
View file

@ -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<u8> {
embedding
.iter()
.flat_map(|value| value.to_le_bytes())
.collect()
}
/// `HSET` a fresh `<prefix><scope>:<uuid4>` 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<u8>,
vector: Vec<u8>,
ttl: Option<Duration>,
) -> Result<(), Error> {
ensure_index(connection, index, vector.len() / size_of::<f32>())?;
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<u8>,
) -> Result<SemanticLookup<Vec<u8>>, Error> {
ensure_index(connection, index, vector.len() / size_of::<f32>())?;
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::<redis::Value>(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<u8>)>;
fn search_fields(value: redis::Value) -> Result<Option<SearchFields>, 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::<Result<Vec<_>, Error>>()
.map(Some)
}
fn parse_i64(value: &redis::Value) -> Result<i64, Error> {
value_text(value)
.ok_or(Error::InvalidEntry(None))?
.parse()
.map_err(Error::invalid_entry)
}
fn parse_f64(value: &[u8]) -> Result<f64, Error> {
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<String> {
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<Vec<u8>, 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)),
}
}

430
litellm-rust/crates/cache-valkey-semantic/tests/cache.rs
View file

@ -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<Mutex<Vec<Vec<u8>>>>;
type RecordingCache = ValkeySemanticCache<FakeEmbedder, JsonCodec<Value>, 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::<Vec<_>>()
.join("\n")
}
}
fn config() -> ValkeySemanticConfig {
ValkeySemanticConfig {
similarity_threshold: 0.8,
index_name: "test".into(),
}
}
fn recording(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> 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<redis::Value> {
Ok(redis::Value::SimpleString("OK".into()))
}
fn already_exists() -> redis::RedisResult<redis::Value> {
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<u8> {
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<Value>,
#[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::<Vec<_>>();
assert_eq!(
prompts,
expected.into_iter().map(str::to_owned).collect::<Vec<_>>()
);
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<Duration>,
#[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<redis::Value>,
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<Option<Value>, Error>,
context: SemanticCacheContext,
) {
let recording = recording([ok(), Ok(reply)]);
let kind = |result: Result<Option<Value>, 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<redis::Value>,
#[case] value: Option<Value>,
#[case] similarity: Option<f64>,
#[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]
);
}

62
litellm-rust/crates/cache-valkey-semantic/tests/contract.rs
View file

@ -1,62 +0,0 @@
//! `overwrite_replaces` does not apply: like Python, every write is a new `<scope>:<uuid4>`
//! 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<PreparedEmbedding, JsonCodec<Value>, 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;
}

352
litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs
View file

@ -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<Mutex<Vec<(String, Option<Value>)>>>;
/// 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<String, Vec<f32>>,
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<Vec<f32>, 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<Vec<f32>, Error> {
self.embed(prompt, metadata)
}
}
struct FakeIndex {
prefix: Vec<u8>,
dims: usize,
vector_field: String,
}
#[derive(Default)]
struct SearchState {
indexes: HashMap<String, FakeIndex>,
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
}
/// 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<Mutex<SearchState>>,
}
impl FakeSearch {
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
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::<Vec<_>>();
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<redis::Value> {
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<Vec<redis::Value>> {
let replies = parse_commands(commands)
.into_iter()
.map(|args| self.run(args))
.collect::<redis::RedisResult<Vec<_>>>()?;
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<f32> {
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::<f64>();
let norm = |vector: &[f32]| {
vector
.iter()
.map(|value| f64::from(*value).powi(2))
.sum::<f64>()
.sqrt()
};
dot / (norm(left) * norm(right))
}
/// Splits a packed RESP request into each command's arguments.
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
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::<usize>().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<Mutex<Vec<Vec<u8>>>>,
replies: Mutex<VecDeque<redis::RedisResult<redis::Value>>>,
}
impl RecordingConnection {
pub fn new(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Self {
Self {
requests: Arc::default(),
replies: Mutex::new(replies.into_iter().collect()),
}
}
pub fn requests(&self) -> Arc<Mutex<Vec<Vec<u8>>>> {
Arc::clone(&self.requests)
}
fn reply(&self) -> redis::RedisResult<redis::Value> {
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<redis::Value> {
self.requests.lock().unwrap().push(command.to_vec());
self.reply()
}
fn req_packed_commands(
&mut self,
command: &[u8],
_offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
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
}
}

2
litellm-rust/crates/callbacks-legacy-python/Cargo.toml
View file

@ -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

18
litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

@ -33,6 +33,17 @@
"api_key",
"additional_args"
],
"record_response": [
"logger",
"status",
"headers",
"body"
],
"record_api_window": [
"logger",
"start",
"end"
],
"defers_async_logging": [
"logger"
],
@ -91,6 +102,13 @@
"logger",
"head"
],
"stamp_stream_head": [
"head",
"logger",
"kwargs",
"start_time",
"end_time"
],
"stream_success": [
"logger",
"url_route",

119
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -57,6 +57,8 @@ pub struct LegacyLogging {
asynchronous: bool,
internal: bool,
cache_key: Option<String>,
/// When the provider request went out, until its window is recorded (Messages only).
api_started: Option<f64>,
}
fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
@ -87,6 +89,60 @@ 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<Bound<'py, PyDict>> {
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<serde_json::Map<String, Value>> {
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)
}
/// 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(()),
}
}
@ -111,6 +167,16 @@ impl LegacyLogging {
}
}
/// 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())
}
}
}
fn stream_billing(&self) -> Option<PassThroughStream> {
match self.operation {
LoggingOperation::Messages => Some(PassThroughStream {
@ -377,7 +443,13 @@ impl LegacyLogging {
context: &RequestContext,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
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::<PyDict>()?;
@ -400,11 +472,14 @@ impl LegacyLogging {
self.logger()?.pre_call(
py,
self.input_description(),
context.api_key.as_ref().map(|api_key| api_key.expose()),
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::<String>()?, value.extract::<String>()?)))
@ -455,7 +530,7 @@ impl LegacyLogging {
py,
(
&logger,
self.call.kwargs(),
self.logging_kwargs(py, true)?,
&facts.provider.model,
logger.getattr("optional_params")?,
params,
@ -483,6 +558,16 @@ impl LegacyLogging {
py: Python<'_>,
raw: &RawResponse,
) -> PyResult<HookStep<Self, ()>> {
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()
@ -522,6 +607,9 @@ impl LegacyLogging {
) -> PyResult<HookStep<Self, ()>> {
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 +631,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(),
@ -925,11 +1027,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::*;
@ -1117,9 +1219,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(
@ -1869,9 +1969,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

51
litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs
View file

@ -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,9 +16,10 @@ pub trait LegacyCallbacks {
fn update_from_kwargs(
&self,
py: Python<'_>,
kwargs: &Py<PyDict>,
kwargs: &Bound<'_, PyDict>,
wire: &WireRequest,
context: &RequestContext,
request_params: Option<&serde_json::Map<String, Value>>,
) -> PyResult<()>;
/// `Logging.pre_call`.
@ -31,6 +33,20 @@ pub trait LegacyCallbacks {
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,
@ -83,12 +99,13 @@ impl LegacyCallbacks for PythonLogger {
fn update_from_kwargs(
&self,
py: Python<'_>,
kwargs: &Py<PyDict>,
kwargs: &Bound<'_, PyDict>,
wire: &WireRequest,
context: &RequestContext,
request_params: Option<&serde_json::Map<String, Value>>,
) -> 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 +114,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,
(
@ -145,6 +162,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<'_>,

6
litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -44,6 +44,10 @@ pub(crate) enum Logging {
PreCall,
#[strum(serialize = "post_call")]
PostCall,
#[strum(serialize = "record_response")]
RecordResponse,
#[strum(serialize = "record_api_window")]
RecordApiWindow,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "defer_success")]
@ -77,6 +81,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")]

7
litellm-rust/crates/callbacks-legacy-python/src/test_support.rs
View file

@ -65,6 +65,12 @@ FAKES = {
'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),
'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls(
@ -85,6 +91,7 @@ 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),
'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: logger.record(
'stream_success', list(chunks)
),

2
litellm-rust/crates/config/Cargo.toml
View file

@ -6,7 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth-types.workspace = true
litellm-secret-value.workspace = true
litellm-router-types.workspace = true
serde.workspace = true
serde_json.workspace = true

2
litellm-rust/crates/config/src/mcp.rs
View file

@ -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;

2
litellm-rust/crates/config/src/redis.rs
View file

@ -1,7 +1,7 @@
use std::{collections::BTreeMap, fmt, time::Duration};
use litellm_auth_types::SecretValue;
use litellm_router_types::Spelled;
use litellm_secret_value::SecretValue;
use serde::{Deserialize, Deserializer, de::Error as _};
use crate::{AdditionalFields, Error, Object, Value};

2
litellm-rust/crates/config/src/settings.rs
View file

@ -1,6 +1,6 @@
use std::fmt;
use litellm_auth_types::SecretValue;
use litellm_secret_value::SecretValue;
use serde::Deserialize;
use crate::{

2
litellm-rust/crates/core-utils/src/core_helpers.rs
View file

@ -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")]

1
litellm-rust/crates/cost/AGENTS.md Normal file
View file

@ -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

10
litellm-rust/crates/cost/src/price.rs
View file

@ -19,6 +19,7 @@ pub enum ServiceTier {
Balanced,
}
/// A prompt-size range. `OverPromptTokens(n)` is spelled `_above_<n>_tokens` or `_above_<n / 1000>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<Price, UnitPrice>,
@ -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<UnitPrice>,
}
/// Multipliers applied on top of a priced call, never prices themselves
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Adjustments(BTreeMap<Uplift, Multiplier>);
@ -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.<geo>`, lowercased as on main

3
litellm-rust/crates/cost/src/price/meter.rs
View file

@ -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,

2
litellm-rust/crates/gateway-auth/Cargo.toml
View file

@ -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

Some files were not shown because too many files have changed in this diff Show more