Merge origin/main into litellm_customer_key_rate_limits
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled

This commit is contained in:
devin-ai-integration[bot] 2026-11-10 01:40:59 +00:00
commit 0e30fd3df9
1477 changed files with 108215 additions and 25203 deletions

View file

@ -182,6 +182,7 @@ jobs:
tests/unit/google_genai
tests/unit/router_strategy
tests/unit/router_utils
tests/unit/redis
tests/unit/proxy/common_utils/test_cache_aware_routing.py
tests/unit/enterprise/enterprise_callbacks/send_emails
tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py
@ -1034,9 +1035,9 @@ jobs:
run: |
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
else
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
@ -1293,7 +1294,7 @@ jobs:
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
@ -1496,9 +1497,9 @@ jobs:
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
run: |
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
else
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
fi

View file

@ -52,7 +52,6 @@ If you ever write any human-facing text (pull requests, issues, commit messages,
- don't use "—". Don't use ";" or ":" unless it would be nonsensical not to. Instead, reach for ",", ".", conjunction words, etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them also feels AI-y. A word cap does not penalize you for adding more sentences
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a ".". Exception is that docs (i.e., litellm-docs) should have trailing "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure

View file

@ -46,7 +46,43 @@ For a managed workspace rollout, use OpenAI's [gateway configuration guide](http
## Native compatibility and Pi
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name.
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials. A ChatGPT model keeps its upstream name and needs no deployment of its own.
### Mix ChatGPT and LiteLLM models
On the native route, Codex's model picker shows OpenAI's catalog first and then the other chat and responses models the virtual key can use. Pick one of those and Codex's requests go through LiteLLM's `/v1/responses` with the key's usual routing, budgets, guardrails, and logging. The proxy removes the ChatGPT bearer and account header from that request, so the other provider and your callbacks never receive them. A name is routed this way when another provider serves it as an exact `model_name`, through a router `model_group_alias`, or as a team's public model name. A ChatGPT model, a name that only a wildcard deployment matches, or a name OpenAI's catalog lists for the caller's ChatGPT account keeps the passthrough, so a model OpenAI adds later still uses the subscription even when a deployment of another provider has the same name. The proxy learns the account's list from the catalog Codex fetches, or asks OpenAI with the caller's login. When it cannot get the list, for example because OpenAI is down, the request times out, or the client is not Codex, every request keeps the passthrough for the next minute, as it did before LiteLLM models were added to the picker. When the key can use many models, the picker lists the ones that fit in Codex's catalog size limit, and the proxy logs a warning with how many it left out.
Codex's picker sends a `none` reasoning effort for a LiteLLM model whose catalog entry lists no effort levels. With `litellm_settings: drop_params: true`, the proxy drops it for a model that does not support `none`, and the provider applies its own default. A model whose cost map entry or `model_info` sets `supports_none_reasoning_effort: true` keeps it. Without `drop_params`, the effort reaches the provider unchanged.
Codex also sends its ChatGPT app tools with every turn, which on a LiteLLM model adds a few hundred thousand prompt tokens to each turn. Set `apps = false` under `[features]` in Codex's `config.toml` to turn them off.
Only a key that may call `/v1/responses` gets these models. A key whose `allowed_routes` covers only the native routes sees OpenAI's catalog unchanged, and every request it sends uses the subscription.
Codex sends earlier reasoning items back with each turn. The proxy tags the reasoning that a LiteLLM model returns on this route, OpenAI and Azure deployments included. When a session switches to a ChatGPT model, the proxy drops those items before forwarding the request to OpenAI, which would reject them. Messages, tool calls, and ChatGPT's own reasoning are kept, and a LiteLLM model still gets its own reasoning back.
When the catalog adds LiteLLM models for a key, its `ETag` is a digest of OpenAI's ETag and those models, and that key's responses carry the same value in `X-Models-Etag`. Codex keeps its cached catalog while the two match and fetches it again when OpenAI's catalog or the key's models change. A key whose catalog adds no LiteLLM models keeps OpenAI's ETags.
### Restrict native models with an access group
A key with a model allowlist needs a rule for ChatGPT names that OpenAI adds over time. Put the subscription wildcards in an access group under `chatgpt/` names:
```yaml
model_list:
- model_name: chatgpt/gpt-*
litellm_params:
model: chatgpt/gpt-*
model_info:
access_groups: ["chatgpt-native"]
- model_name: chatgpt/codex-*
litellm_params:
model: chatgpt/codex-*
model_info:
access_groups: ["chatgpt-native"]
```
Then add `chatgpt-native` to the key's models, next to any other model names it should reach. On the native route, a key that holds only the `chatgpt/` form of a name is checked as that form, so a request for `gpt-5.6-sol` is allowed when the key may call `chatgpt/gpt-5.6-sol`. The key can call every ChatGPT model that matches the wildcards, including ones the proxy configuration does not name.
On `/v1` these wildcards match only a request that asks for a `chatgpt/` name, so a request for a plain `gpt-` or `codex-` name still goes where it went before, such as a `*` catch-all deployment. A `/v1` request for a `chatgpt/` name without a ChatGPT login falls back to the proxy's own saved login if it has one, so keep that login unset as described above.
### Configure Pi
@ -71,7 +107,7 @@ Keep Pi's subscription login as the provider credential. The header value names
Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key needs the `chatgpt-native` access group above or that exact name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate.

View file

@ -182,7 +182,7 @@ class ManagedObjectCreatorAttribution:
team_id: Final = getattr(job, "team_id", None)
rows: Final = _CreatorRows(
key=await self._live_key_row(api_key, job.created_by),
team=await self._live_team_row(team_id),
team=await self._live_team_row(team_id, job.created_at),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
return ResponseCreator(
@ -269,10 +269,15 @@ class ManagedObjectCreatorAttribution:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
return None
async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
async def _live_team_row(
self, team_id: str | None, created_before: datetime | None = None
) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
return None
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
team: Final = await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
if team is not None and created_before is not None and team.created_at > created_before:
return None
return team
def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.81"
version = "0.1.83"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.81"
version = "0.1.83"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.112"
version = "0.4.114"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.112"
version = "0.4.114"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -0,0 +1,53 @@
---
name: rust-python-boundary
description: Decide where Python-shaped input becomes typed Rust in litellm-rust. Use when reading config.yaml values, os.environ/ references, REDIS_* or other environment layering, get_secret fallbacks, or kwargs and objects handed over from Python, when coercing text to a bool, int or float the way Python does, or when adding a dependency on litellm-python-compat
metadata:
internal: true
---
# Python boundary
Python-shaped input becomes typed Rust in exactly two crates. Every crate past them handles typed values, so no crate re-guesses how Python would have read a value
## The two boundaries
- `crates/config` reads config.yaml, which is Python-shaped: `os.environ/` references, scalars that Helm or the environment turned into text, raw kwargs maps. It also owns the environment layering and `get_secret` fallbacks Python applies while it loads that config. It has no Python runtime, so it reproduces Python value semantics with `litellm-python-compat`
- `crates/python-bridge` receives live Python objects, such as `RedisCache.redis_kwargs`. It reads them through PyO3 with the real protocol (`is_truthy()`, `str()`, `extract`), so it never imitates Python
- Both produce the same typed input, defined in the data crate the consumer already depends on (`redis-types` for Redis), so a downstream crate cannot tell which host it runs under
## Python value semantics
Reading a value the way Python itself reads it. This lives only at the two boundaries
- Truthiness and `x or fallback`, `str(x)`, `int(text)` and `float(text)`
- Each surface's true and false text: `_str_to_bool`, pydantic's lax bool, `get_secret`'s `true` and `false`
- `os.environ/` references and `REDIS_<KWARG>` style environment layering
A boundary knows which surface a value came from. `general_settings.coordination_redis` is a pydantic model and `cache_params` goes through `_str_to_bool`, so the boundary applies that surface's rule exactly instead of refusing the text two surfaces disagree on
Each typed field keeps where its value came from (the key, `REDIS_<KWARG>` or the fallback name), so an error or warning downstream still names it
## What stays downstream
A crate that stands in for a Python library keeps that library's semantics, because they are its compatibility contract and not value coercion. `redis-resolve` stands in for redis-py, so it keeps
- Which target wins, and the defaults redis-py applies
- How `parse_url` reads a url, including `int()` on `?db=` and redis-py's `to_bool` on the TLS query parameters
- Which kwargs redis-py accepts, and whether the Rust client reads, warns about or refuses each one
Its input is typed. A rule it applies reads typed fields, never a `serde_json::Value` it has to coerce first. A value it only carries back to Python untouched, like `passthrough`, may stay untyped
## litellm-python-compat
The rule is about where config and kwargs get coerced, not about who may depend on the crate
- Its data formats (`pickle`, `json.dumps`, `repr`) are fine anywhere Rust reads or writes state Python also reads or writes, such as cache values in a shared Redis
- Its coercions (`truthy`, `number`) belong at the two boundaries, and in a crate that stands in for a Python library when that library calls them itself
## Existing code
Some code still coerces Python-shaped values past the boundary, for example the `REDIS_*` merge in `gateway/src/redis/layers.rs` and the flag and number readers in `redis-resolve/src/options.rs`. Do not add more. When you change one, move the coercion to its boundary instead of extending it in place
## Tests
Pin a boundary's coercion with named cases next to the Python function it mirrors, in the Rust test and in its Python twin, like the Redis contract cases in `redis-resolve/tests/interpret.rs` and `tests/unit/test_redis.py`. A downstream crate's tests build its typed input directly and never feed it raw text to coerce

View file

@ -8,6 +8,10 @@ For string-valued enums and their Serde conversions, follow [.agents/skills/rust
For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md)
For where config.yaml values, environment layering and kwargs from Python become typed Rust, follow [.agents/skills/rust-python-boundary/SKILL.md](.agents/skills/rust-python-boundary/SKILL.md)
The process-wide rustls `CryptoProvider` (ring) is installed once by the host. rust-proxy does it at the top of `main` through `litellm_gateway::install_crypto_provider`, and python-bridge must do the same in its module init before it builds any rustls config. Library crates never install it. The AWS SDK links aws-lc-rs next to ring, so without an installed default the first rustls config a library builds, such as redis-rs TLS, panics
For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches
Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all`

212
litellm-rust/Cargo.lock generated
View file

@ -8,6 +8,17 @@ version = "2.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
[[package]]
name = "aes"
version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
dependencies = [
"cfg-if",
"cipher",
"cpufeatures 0.2.17",
]
[[package]]
name = "ahash"
version = "0.8.12"
@ -982,6 +993,15 @@ dependencies = [
"hybrid-array",
]
[[package]]
name = "block-padding"
version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93"
dependencies = [
"generic-array",
]
[[package]]
name = "bollard"
version = "0.20.2"
@ -1134,6 +1154,15 @@ dependencies = [
"rustversion",
]
[[package]]
name = "cbc"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6"
dependencies = [
"cipher",
]
[[package]]
name = "cc"
version = "1.3.0"
@ -1210,6 +1239,16 @@ dependencies = [
"half",
]
[[package]]
name = "cipher"
version = "0.4.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
dependencies = [
"crypto-common 0.1.7",
"inout",
]
[[package]]
name = "clap"
version = "4.6.7"
@ -1774,6 +1813,15 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "des"
version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e"
dependencies = [
"cipher",
]
[[package]]
name = "digest"
version = "0.10.7"
@ -3105,6 +3153,16 @@ dependencies = [
"serde_core",
]
[[package]]
name = "inout"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
dependencies = [
"block-padding",
"generic-array",
]
[[package]]
name = "ipnet"
version = "2.12.0"
@ -3663,8 +3721,7 @@ dependencies = [
"strum",
"tempfile",
"thiserror 2.0.19",
"tracing",
"tracing-subscriber",
"veil",
]
[[package]]
@ -3767,6 +3824,7 @@ dependencies = [
"http-body-util",
"litellm-circuit-breaker",
"litellm-config",
"litellm-core-utils",
"litellm-error",
"litellm-gateway-auth",
"litellm-gateway-forward",
@ -3777,11 +3835,15 @@ dependencies = [
"litellm-inference",
"litellm-llms",
"litellm-redis-client",
"litellm-redis-connect",
"litellm-redis-resolve",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"litellm-secrets",
"litellm-tracing",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"tempfile",
@ -3792,6 +3854,7 @@ dependencies = [
"tower-sessions-moka-store",
"tracing",
"uuid",
"veil",
]
[[package]]
@ -4318,11 +4381,13 @@ dependencies = [
"criterion",
"fancy-regex 0.19.2",
"futures-util",
"indexmap 2.14.0",
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-error",
"litellm-host",
@ -4337,6 +4402,12 @@ dependencies = [
"litellm-llms",
"litellm-llms-types",
"litellm-model-catalog",
"litellm-python-compat",
"litellm-redis-client",
"litellm-redis-connect",
"litellm-redis-resolve",
"litellm-redis-store",
"litellm-redis-types",
"litellm-router-core",
"litellm-router-types",
"litellm-secret-value",
@ -4352,6 +4423,7 @@ dependencies = [
"pyo3-async-runtimes",
"reqwest",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"serde_with",
@ -4387,23 +4459,110 @@ version = "0.1.0"
dependencies = [
"futures-util",
"litellm-circuit-breaker",
"litellm-python-compat",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"pem",
"pkcs8",
"redis",
"rstest",
"rustls 0.23.42",
"rustls-native-certs",
"socket2 0.6.5",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-redis-connect"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-auth-types",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-redis-client",
"litellm-redis-resolve",
"litellm-redis-testing",
"litellm-redis-types",
"litellm-secret-value",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"time",
"tokio",
"tracing",
"url",
"veil",
]
[[package]]
name = "litellm-redis-resolve"
version = "0.1.0"
dependencies = [
"litellm-python-compat",
"litellm-redis-types",
"litellm-secret-value",
"rstest",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"tracing",
"url",
]
[[package]]
name = "litellm-redis-store"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-cache",
"litellm-cache-memory",
"litellm-circuit-breaker",
"litellm-python-compat",
"litellm-redis-client",
"litellm-redis-testing",
"litellm-secret-value",
"litellm-tracing",
"rstest",
"strum",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-redis-testing"
version = "0.1.0"
dependencies = [
"pkcs8",
"rcgen",
"tempfile",
"testcontainers-modules",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-redis-types"
version = "0.1.0"
dependencies = [
"litellm-secret-value",
"rstest",
"url",
]
[[package]]
name = "litellm-router"
version = "0.1.0"
@ -5286,6 +5445,16 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4"
[[package]]
name = "pbkdf2"
version = "0.12.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2"
dependencies = [
"digest 0.10.7",
"hmac 0.12.1",
]
[[package]]
name = "pem"
version = "4.0.0"
@ -5396,6 +5565,22 @@ dependencies = [
"spki",
]
[[package]]
name = "pkcs5"
version = "0.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6"
dependencies = [
"aes",
"cbc",
"der",
"des",
"pbkdf2",
"scrypt",
"sha2 0.10.9",
"spki",
]
[[package]]
name = "pkcs8"
version = "0.10.2"
@ -5403,6 +5588,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
dependencies = [
"der",
"pkcs5",
"rand_core 0.6.4",
"spki",
]
@ -5590,6 +5777,7 @@ version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
dependencies = [
"indexmap 2.14.0",
"libc",
"once_cell",
"portable-atomic",
@ -6411,6 +6599,15 @@ version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
[[package]]
name = "salsa20"
version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213"
dependencies = [
"cipher",
]
[[package]]
name = "same-file"
version = "1.0.6"
@ -6474,6 +6671,17 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
[[package]]
name = "scrypt"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0516a385866c09368f0b5bcd1caff3366aace790fcd46e2bb032697bb172fd1f"
dependencies = [
"pbkdf2",
"salsa20",
"sha2 0.10.9",
]
[[package]]
name = "sct"
version = "0.7.1"

View file

@ -73,7 +73,11 @@ litellm-host-python = { path = "crates/host-python" }
litellm-python-compat = { path = "crates/python-compat" }
litellm-circuit-breaker = { path = "crates/circuit-breaker" }
litellm-redis-client = { path = "crates/redis-client" }
litellm-redis-connect = { path = "crates/redis-connect" }
litellm-redis-store = { path = "crates/redis-store" }
litellm-redis-testing = { path = "crates/redis-testing" }
litellm-redis-resolve = { path = "crates/redis-resolve" }
litellm-redis-types = { path = "crates/redis-types" }
litellm-cost = { path = "crates/cost" }
litellm-model-catalog = { path = "crates/model-catalog" }
litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" }

View file

@ -2,7 +2,10 @@ use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap, HashSet},
hash::Hash,
sync::{Arc, Mutex, MutexGuard},
sync::{
Arc, Mutex, MutexGuard,
atomic::{AtomicUsize, Ordering},
},
time::{Duration, SystemTime, UNIX_EPOCH},
};
@ -31,7 +34,7 @@ struct CacheState<V> {
pub struct InMemoryCache<V: Clone> {
state: Mutex<CacheState<V>>,
max_size_in_memory: usize,
max_size_in_memory: AtomicUsize,
default_ttl: Duration,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
@ -74,7 +77,9 @@ impl<V: Clone> InMemoryCache<V> {
expirations: HashMap::new(),
expiration_heap: BinaryHeap::new(),
}),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
max_size_in_memory: AtomicUsize::new(
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
),
default_ttl: default_ttl
.filter(|ttl| !ttl.is_zero())
.unwrap_or(DEFAULT_TTL),
@ -94,7 +99,7 @@ impl<V: Clone> InMemoryCache<V> {
value: V,
ttl: Option<Duration>,
) -> Result<CacheWrite, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(CacheWrite::Disabled);
}
let now = (self.now)();
@ -129,7 +134,7 @@ impl<V: Clone> InMemoryCache<V> {
pub fn evict_cache(&self) -> Result<(), Error> {
let now = (self.now)();
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, None);
Self::evict(&mut state, self.max_size_in_memory(), now, None);
Ok(())
}
@ -170,7 +175,16 @@ impl<V: Clone> InMemoryCache<V> {
}
pub fn max_size_in_memory(&self) -> usize {
self.max_size_in_memory
self.max_size_in_memory.load(Ordering::Relaxed)
}
/// `DualCache.update_in_memory_max_size`: `None` restores the default. Entries already held stay
/// until the next write evicts down to the new size.
pub fn update_max_size_in_memory(&self, max_size_in_memory: Option<usize>) {
self.max_size_in_memory.store(
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
Ordering::Relaxed,
);
}
pub fn max_entry_bytes(&self) -> Option<usize> {
@ -200,6 +214,23 @@ impl<V: Clone> InMemoryCache<V> {
.collect())
}
/// Every live entry, for callers that inspect the whole tier.
pub fn entries(&self) -> Result<Vec<(String, V)>, Error> {
let now = (self.now)();
let state = self.lock()?;
Ok(state
.values
.iter()
.filter(|(key, _)| {
state
.expirations
.get(*key)
.is_none_or(|expiration| *expiration >= now)
})
.map(|(key, value)| (key.clone(), value.clone()))
.collect())
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.lock()?;
Self::remove(&mut state, key);
@ -276,7 +307,7 @@ impl<V: Clone> InMemoryCache<V> {
ttl: Option<Duration>,
now: Duration,
) -> Result<CacheWrite, Error> {
Self::evict(state, self.max_size_in_memory, now, None);
Self::evict(state, self.max_size_in_memory(), now, None);
if !self.check_value_size(&value)? {
return Ok(CacheWrite::TooLarge);
}
@ -300,12 +331,12 @@ where
eligible: &[V],
context: ExactCacheContext,
) -> Result<V, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(candidate);
}
let now = (self.now)();
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
Self::evict(&mut state, self.max_size_in_memory(), now, Some(key));
let existing = state
.values
.get(key)
@ -335,7 +366,7 @@ impl CounterCache for InMemoryCache<f64> {
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(amount);
}
let now = (self.now)();
@ -408,7 +439,7 @@ where
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
if self.max_size_in_memory == 0 {
if self.max_size_in_memory() == 0 {
return Ok(values);
}
let now = (self.now)();

View file

@ -27,6 +27,10 @@
"api_key",
"additional_args"
],
"messages_input": [
"body",
"kwargs"
],
"post_call": [
"logger",
"original_response",
@ -47,9 +51,18 @@
"defers_async_logging": [
"logger"
],
"claim_async_success_log": [
"logger"
],
"mark_async_success_enqueued": [
"logger"
],
"defer_success": [
"logger",
"pending"
"pending",
"response",
"start",
"end"
],
"sync_success_for_async_call": [
"logger",
@ -117,7 +130,8 @@
"chunks",
"start",
"end",
"first_chunk"
"first_chunk",
"stream_teardown"
],
"stream_failure": [
"logger",

View file

@ -14,7 +14,7 @@ use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDateTime, PyDict, PyList},
types::{PyDateTime, PyDict, PyList, PyString},
};
use serde_json::Value;
@ -22,7 +22,7 @@ use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call,
python::Streaming,
python::{Logging, Streaming},
setup,
};
@ -137,6 +137,40 @@ impl LegacyLogging {
Some(params)
}
// Compared as Python compares them, so a caller's `0` is kept over a typed float's `0.0`.
fn caller_value<'py>(
&self,
py: Python<'py>,
name: &str,
sent: &Value,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match self.call.lookup(py, name)? {
Some(value) if from_py::<Value>(&value).is_ok() && value.eq(to_py(py, sent)?)? => {
Ok(Some(value))
}
_ => Ok(None),
}
}
// Python's Messages handler logs the caller's own optional params, not their typed copies.
fn as_called(&self, py: Python<'_>, context: &RequestContext) -> PyResult<RequestContext> {
let mut context = context.clone();
if self.operation != LoggingOperation::Messages {
return Ok(context);
}
for (name, sent) in context
.optional_params
.as_object_mut()
.into_iter()
.flatten()
{
if let Some(value) = self.caller_value(py, name, sent)? {
*sent = from_py(&value)?;
}
}
Ok(context)
}
/// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync
/// client is untimed.
fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> {
@ -158,13 +192,20 @@ impl LegacyLogging {
}
}
fn input_description(&self) -> &'static str {
match self.operation {
fn logged_input<'py>(
&self,
py: Python<'py>,
body: &Bound<'py, PyDict>,
) -> PyResult<Bound<'py, PyAny>> {
let description = match self.operation {
LoggingOperation::Completion => "Chat completions",
LoggingOperation::Responses => "Responses",
LoggingOperation::Messages => "Messages",
LoggingOperation::Messages => {
return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py)));
}
LoggingOperation::Ocr => "OCR document processing",
}
};
Ok(PyString::new(py, description).into_any())
}
/// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key.
@ -251,30 +292,49 @@ impl LegacyLogging {
if !self.asynchronous {
return pending().sync(py);
}
if !self.internal
&& self
if self.internal {
return Ok(());
}
let defers = logger.defers_async_logging(py);
if logger.claim_async_success_log(py)? {
let enqueues = self
.call
.kwargs()
.bind(py)
.get_item("fallbacks")?
.is_none_or(|value| value.is_none())
{
if logger.defers_async_logging(py) {
.is_none_or(|value| value.is_none());
if defers {
let pending = Py::new(
py,
PendingLogging {
pending: Some(pending()),
enqueues,
},
)?;
logger.defer_success(py, pending.bind(py).as_any())?;
} else {
return logger.defer_success(
py,
pending.bind(py).as_any(),
&self.response,
&self.start,
&self.end,
);
}
if enqueues {
pending().asynchronous(py)?;
}
}
if defers {
return Ok(());
}
logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
fn stream_success(
&self,
py: Python<'_>,
stream: &DeliveredStream,
teardown: bool,
) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.stream_billing().ok_or(HostError::MissingState)?;
let billed = Streaming::Success.call(
@ -291,6 +351,7 @@ impl LegacyLogging {
&self.start,
&self.end,
&stream.first_chunk,
teardown,
),
);
match billed {
@ -442,6 +503,7 @@ impl LegacyLogging {
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
let context = &self.as_called(py, context)?;
let logger = self.logger()?;
logger.update_from_kwargs(
py,
@ -454,12 +516,17 @@ impl LegacyLogging {
.into_bound(py)
.cast_into::<PyDict>()?;
for (name, sent) in wire.body.as_object().into_iter().flatten() {
if let Some(value) = self.call.lookup(py, name)?
&& from_py::<Value>(&value).is_ok_and(|caller| caller == *sent)
{
if let Some(value) = self.caller_value(py, name, sent)? {
body.set_item(name, value)?;
}
}
if self.operation == LoggingOperation::Messages {
// Python's Messages handler merges its transformed request body into the call details.
logger
.object(py)
.getattr("model_call_details")?
.call_method1("update", (&body,))?;
}
let headers = PyDict::new(py);
for (name, value) in &wire.headers {
headers.set_item(name, value)?;
@ -469,9 +536,10 @@ impl LegacyLogging {
headers: headers.clone().unbind(),
context: context.clone(),
});
let input = self.logged_input(py, &body)?;
self.logger()?.pre_call(
py,
self.input_description(),
&input,
self.logged_api_key(context),
&body,
&headers,
@ -525,7 +593,11 @@ impl LegacyLogging {
.getattr("litellm_params")?
.cast_into::<PyDict>()?
.copy()?;
params.set_item("custom_llm_provider", &facts.provider.provider)?;
if self.operation != LoggingOperation::Messages {
// Python's Messages handler hands the provider to the logger beside its params, never
// in them, and the proxy's logger carries over to the next fallback hop.
params.set_item("custom_llm_provider", &facts.provider.provider)?;
}
crate::python::Logging::Update.call(
py,
(
@ -548,8 +620,17 @@ impl LegacyLogging {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
};
details.set_item("cache_hit", self.cache_key.is_some())?;
if self.cache_key.is_some() {
details.set_item("cache_hit", true)?;
} else {
details.call_method1("pop", ("cache_hit", py.None()))?;
}
details.set_item("cache_key", self.cache_key.as_deref())?;
if let (LoggingOperation::Messages, ResultSource::Provider, Some(request)) =
(self.operation, &facts.source, &self.request)
{
details.call_method1("update", (request.body.bind(py),))?;
}
Ok(HookStep::Ready(()))
}
@ -592,7 +673,7 @@ impl LegacyLogging {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
match &self.stream {
Some(stream) => self.stream_success(py, stream)?,
Some(stream) => self.stream_success(py, stream, false)?,
None => self.dispatch_success(py)?,
}
Ok(HookStep::Ready(()))
@ -662,8 +743,28 @@ impl LegacyLogging {
}
}
impl LegacyLogging {
// A caller that goes away mid-stream cancels the call before it ends: Python's Messages
// stream iterator still bills the chunks it read, straight away rather than deferred.
fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> {
if !self.asynchronous || self.end.is_some() {
return Ok(());
}
let Some(stream) = self.stream.take() else {
return Ok(());
};
self.end = Some(datetime(py, epoch_seconds())?);
self.stream_success(py, &stream, true)
}
}
impl PythonOwned for LegacyLogging {
fn close(&mut self, py: Python<'_>) {
if self.logger.is_some()
&& let Err(error) = self.bill_abandoned_stream(py)
{
error.write_unraisable(py, None);
}
if let Some(logger) = self.logger.take()
&& let Err(error) = logger.restore_context(py)
{
@ -1114,10 +1215,10 @@ logger._provider_reported_cost = None
};
let previous = litellm_host::interceptors::ExecutionFacts {
provider: facts.provider.clone(),
cost: litellm_host::interceptors::Cost::Reported {
amount: serde_json::Number::from_f64(9.0).unwrap(),
cost: litellm_host::interceptors::Cost::CacheHit,
source: litellm_host::interceptors::ResultSource::Cache {
key: "earlier-hit".into(),
},
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &previous).unwrap();
local(&locals, "logger")
@ -1135,13 +1236,59 @@ logger._provider_reported_cost = None
c"
assert logger.model_call_details.get('response_cost') is None, logger.model_call_details
assert logger._provider_reported_cost == reported, logger._provider_reported_cost
assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details
assert logger.model_call_details.get('cache_hit') is (True if cache_hit else None), logger.model_call_details
assert logger.update['custom_llm_provider'] == 'edenai', logger.update
",
);
});
}
#[rstest]
#[case::messages(crate::LoggingOperation::Messages, false)]
#[case::ocr(crate::LoggingOperation::Ocr, true)]
fn result_ready_puts_the_provider_in_the_logged_params_except_for_messages(
#[case] operation: crate::LoggingOperation,
#[case] in_params: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
run(
py,
&locals,
c"
logger.litellm_params = {}
logger.optional_params = {}
logger.model_call_details = {}
logger._provider_reported_cost = None
",
);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
operation,
..legacy_call(py, &locals, false)
};
let facts = litellm_host::interceptors::ExecutionFacts {
provider: litellm_host::interceptors::ProviderIdentity {
model: "native-test".into(),
provider: "anthropic".into(),
},
cost: litellm_host::interceptors::Cost::Deferred,
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &facts).unwrap();
locals.set_item("in_params", in_params).unwrap();
run(
py,
&locals,
c"
assert logger.update['custom_llm_provider'] == 'anthropic', logger.update
assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params, logger.update
",
);
});
}
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk";
@ -1276,7 +1423,10 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update
#[rstest]
#[case::completion(crate::LoggingOperation::Completion, "Chat completions")]
#[case::responses(crate::LoggingOperation::Responses, "Responses")]
#[case::messages(crate::LoggingOperation::Messages, "Messages")]
#[case::messages(
crate::LoggingOperation::Messages,
r#"[{'role': 'user', 'content': '{"pages": [1]}'}]"#
)]
#[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")]
fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases(
#[case] operation: crate::LoggingOperation,
@ -1317,7 +1467,7 @@ prepared = {'pages': replacement}
&locals,
c"
assert logger.pre['complete_input_dict']['pages'] is replacement
assert logger.pre_input == description
assert str(logger.pre_input) == description
assert original == [0]
",
);
@ -1873,10 +2023,15 @@ mod terminal_tests {
#[case::async_listened(
true,
c"",
&["async_success_handler", "enqueued", "sync_success_for_async_call"]
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call"]
)]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &[])]
#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
#[case::async_deferred_with_fallbacks(
true,
c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}",
&[]
)]
fn success_reaches_the_logging_handlers(
#[case] asynchronous: bool,
#[case] script: &CStr,
@ -1902,12 +2057,88 @@ mod terminal_tests {
&locals,
c"
assert all(value is response for name, value in logger.calls if name.endswith('_handler'))
assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False)
assert hasattr(logger, 'enqueue_deferred_logging') == getattr(logger, '_defer_async_logging', False)
assert getattr(logger, 'deferred_success_result', (response,))[0] is response
",
);
});
}
#[rstest]
#[case::accepted(c"", &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"])]
#[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
fn a_deferred_success_runs_sync_callbacks_only_once_released(
#[case] script: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"response = object()\nlogger._defer_async_logging = True",
);
run(py, &locals, script);
let mut logging = logged(py, &locals, true);
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
",
);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::deferred(
c"logger._defer_async_logging = True",
c"logger.enqueue_deferred_logging()",
&["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"]
)]
#[case::deferred_and_blocked(c"logger._defer_async_logging = True", c"", &[])]
#[case::immediate(
c"",
c"",
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call", "sync_success_for_async_call"]
)]
fn a_fallback_attempt_on_the_same_logger_leaves_the_first_success_claimed(
#[case] script: &CStr,
#[case] release: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"refused = object()\nresponse = refused");
run(py, &locals, script);
let mut first = logged(py, &locals, true);
succeed(py, &locals, &mut first);
run(py, &locals, c"response = object()");
let mut fallback = logged(py, &locals, true);
succeed(py, &locals, &mut fallback);
run(py, &locals, release);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
run(
py,
&locals,
c"assert all(value is refused for name, value in logger.calls if name == 'async_success_handler')",
);
});
}
#[rstest]
fn dropped_observations_preserve_deferred_success_and_response_identity() {
Python::initialize();
@ -1933,10 +2164,10 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy
));
assert_eq!(sender.dropped_events(), 1);
run(py, &locals, c"
assert logger.names() == ['sync_success_for_async_call'], logger.calls
logger._native_pending_logging.release(True)
logger._native_pending_logging.release(True)
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
assert logger.calls[1][1] is response
");
@ -1981,6 +2212,45 @@ assert chunks[1] is last
});
}
#[rstest]
#[case::cancelled_mid_stream(None, &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
#[case::after_success(Some(true), &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
#[case::after_failure(Some(false), &["stamp_stream_head", "stream_opened", "stream_failure", "restore"])]
fn closing_an_unfinished_stream_bills_its_delivered_chunks_once(
#[case] ended: Option<bool>,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"first = b'first'\nresponse = None\nfailure = ValueError('provider')\nhead = {}",
);
let mut logging = LegacyLogging {
operation: crate::LoggingOperation::Messages,
..logged(py, &locals, true)
};
logging
.on_stream_open(py, &local(&locals, "head").unbind())
.unwrap();
logging
.on_stream_chunk(py, &local(&locals, "first").unbind())
.unwrap();
match ended {
Some(true) => drop(succeed(py, &locals, &mut logging)),
Some(false) => drop(fail(py, &locals, &mut logging)),
None => {}
}
litellm_host_python::PythonOwned::close(&mut logging, py);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::synchronous(false, &["failure_handler"])]
#[case::asynchronous(true, &[])]
@ -2009,7 +2279,7 @@ assert chunks[1] is last
}
#[rstest::rstest]
fn internal_async_calls_skip_the_async_success_fan_out() {
fn internal_async_calls_skip_success_logging() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
@ -2018,11 +2288,7 @@ assert chunks[1] is last
..logged(py, &locals, true)
};
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"assert logger.names() == ['sync_success_for_async_call'], logger.calls",
);
run(py, &locals, c"assert logger.calls == [], logger.calls");
});
}

View file

@ -26,7 +26,7 @@ pub trait LegacyCallbacks {
fn pre_call(
&self,
py: Python<'_>,
input: &str,
input: &Bound<'_, PyAny>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
@ -59,7 +59,16 @@ pub trait LegacyCallbacks {
fn defers_async_logging(&self, py: Python<'_>) -> bool;
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool>;
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()>;
fn sync_success_for_async_call(
&self,
@ -148,7 +157,7 @@ impl LegacyCallbacks for PythonLogger {
fn pre_call(
&self,
py: Python<'_>,
input: &str,
input: &Bound<'_, PyAny>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
@ -203,8 +212,21 @@ impl LegacyCallbacks for PythonLogger {
.unwrap_or(false)
}
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending))?;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool> {
Logging::ClaimAsyncSuccess
.call(py, (self.object(py),))?
.extract()
}
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending, response, start, end))?;
Ok(())
}
@ -252,6 +274,7 @@ impl LegacyCallbacks for PythonLogger {
) -> PyResult<()> {
let coroutine =
Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?;
Logging::MarkAsyncSuccessEnqueued.call(py, (self.object(py),))?;
let enqueue = Logging::Enqueue.call(py, (&coroutine,));
if enqueue.is_err()
&& let Err(error) = coroutine.call_method0("close")

View file

@ -22,21 +22,35 @@ impl PendingSuccess {
self.logger
.enqueue_success(py, &self.response, &self.start, &self.end)
}
pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> {
self.logger
.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
}
#[pyclass]
pub(crate) struct PendingLogging {
pub(crate) pending: Option<PendingSuccess>,
pub(crate) enqueues: bool,
}
#[pymethods]
impl PendingLogging {
fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> {
let pending = slf.borrow_mut().pending.take();
if let Some(pending) = pending
&& success
{
match pending.asynchronous(py) {
fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> {
let (pending, enqueues) = {
let mut slot = slf.borrow_mut();
(slot.pending.take(), slot.enqueues)
};
if let Some(pending) = pending {
let released = pending.sync_for_async_call(py).and_then(|()| {
if enqueues {
pending.asynchronous(py)
} else {
Ok(())
}
});
match released {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(pending.logger.object(py)));
}
@ -87,6 +101,7 @@ mod tests {
start: py.None(),
end: Some(py.None()),
}),
enqueues: true,
},
)
.unwrap();
@ -108,7 +123,7 @@ observed = []
def on_enqueue(coroutine):
observed.append(marker.get())
pending.release(True)
pending()
logger.on_enqueue = on_enqueue
",
@ -118,28 +133,12 @@ logger.on_enqueue = on_enqueue
&locals,
c"
marker.set('release')
pending.release(True)
pending.release(True)
pending()
pending()
assert observed == ['release'], observed
assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
",
);
});
}
#[test]
fn a_blocked_release_drops_the_success_for_good() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"");
run(
py,
&locals,
c"
pending.release(False)
pending.release(True)
assert logger.calls == [], logger.calls
assert logger.calls[1][1] is response
",
);
});
@ -168,7 +167,7 @@ logger.on_enqueue = on_enqueue
locals
.set_item("failure", py.eval(failure, None, Some(&locals)).unwrap())
.unwrap();
let released = local(&locals, "pending").call_method1("release", (true,));
let released = local(&locals, "pending").call0();
match released {
Ok(_) => assert!(!propagates),
Err(error) => {
@ -181,8 +180,8 @@ logger.on_enqueue = on_enqueue
py,
&locals,
c"
pending.release(True)
assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls
pending()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued', 'closed'], logger.calls
assert unraisable_from(logger) == ([] if propagates else [failure])
",
);

View file

@ -42,6 +42,8 @@ pub(crate) enum Logging {
Update,
#[strum(serialize = "pre_call")]
PreCall,
#[strum(serialize = "messages_input")]
MessagesInput,
#[strum(serialize = "post_call")]
PostCall,
#[strum(serialize = "record_response")]
@ -50,6 +52,8 @@ pub(crate) enum Logging {
RecordApiWindow,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "claim_async_success_log")]
ClaimAsyncSuccess,
#[strum(serialize = "defer_success")]
DeferSuccess,
#[strum(serialize = "sync_success_for_async_call")]
@ -58,6 +62,8 @@ pub(crate) enum Logging {
SubmitSuccess,
#[strum(serialize = "async_success_handler")]
AsyncSuccessHandler,
#[strum(serialize = "mark_async_success_enqueued")]
MarkAsyncSuccessEnqueued,
#[strum(serialize = "enqueue_logging")]
Enqueue,
#[strum(serialize = "failure_handler")]

View file

@ -62,6 +62,7 @@ FAKES = {
custom_llm_provider=provider,
),
'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args),
'messages_input': lambda body, kwargs: [{'role': 'user', 'content': json.dumps(body)}],
'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call(
original_response, api_key, additional_args
),
@ -72,7 +73,11 @@ FAKES = {
llm_api_duration_ms=(end - start) * 1000
),
'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)),
'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending),
'claim_async_success_log': lambda logger: logger.claim_async_success_log(),
'defer_success': lambda logger, pending, response, start, end: (
setattr(logger, 'deferred_success_result', (response, start, end)),
setattr(logger, 'enqueue_deferred_logging', pending),
),
'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls(
response, start, end
),
@ -81,6 +86,7 @@ FAKES = {
)(error, ''.join(traceback.format_exception(error)), start, end),
'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)),
'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end),
'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None),
'enqueue_logging': lambda coroutine: coroutine.enqueue(),
'restore_context': lambda logger: logger.record('restore', None),
'custom_pricing_fields': lambda: ('ocr_cost_per_page',),
@ -92,7 +98,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': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record(
'stream_success', list(chunks)
),
'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error),
@ -129,6 +135,7 @@ class StubCoroutine:
class StubLogger:
def __init__(self):
self.calls = []
self.model_call_details = {}
self.hooks = {}
self.on_enqueue = lambda coroutine: None
@ -156,6 +163,10 @@ class StubLogger:
self.record('async_success_handler', response)
return StubCoroutine(self)
def claim_async_success_log(self):
claimed, self.claimed = getattr(self, 'claimed', False), True
return not claimed
def handle_sync_success_callbacks_for_async_calls(self, response, start, end):
self.record('sync_success_for_async_call', response)

View file

@ -1,13 +1,14 @@
- `fsm` is a generic state machine runtime and knows nothing about breakers. `breaker` is a generic circuit breaker and knows nothing about Redis. Deciding which errors count as failures belongs to the consumer
- `Machine::step` stays pure. `&self` is config and time comes in on the event, never read inside the machine. The only clock read is `Admission`'s drop, at the call-site boundary
- Transition rules are one `match (state, event)` per machine, with no macro DSL and no `_` arm that would hide a new state from the compiler
- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes, matching Python bumping `generation` only when `_state` changes. Data inside a variant (failure counters, `since`) changes without a transition
- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes. Data inside a variant (failure counters, `since`) changes without a transition
- The observer runs under the lock so transitions are seen in generation order. Keep it to logging or a metric
- A disabled breaker is `Option<Breaker>::None` at the consumer, not a flag inside the machine. `Config` comes in already resolved, reading `REDIS_CIRCUIT_BREAKER_*` env is the consumer's job
## Parity with Python `RedisCircuitBreaker` (`litellm/caching/redis_cache.py`)
- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against the Python breaker. A behavior change here needs the same check, by driving the Python class with the same steps
- Divergences, each on purpose
## The only Redis breaker
- Python's `RedisCache` admits through this breaker via `NativeRedisBreaker` (`python-bridge/src/redis/breaker.rs`), sharing one per Redis server with the native client through `redis-connect`'s `RedisHosts`, so the Python `RedisCircuitBreaker` it replaced is gone
- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against that Python class and now pins its behavior. A behavior change here changes what every Redis caller in the proxy sees
- Divergences from the old Python class, each on purpose
- monotonic `Instant` instead of `time.time()`, so a wall clock jump can't open or close the breaker
- a probe `Admission` dropped without an outcome records a hard failure and reopens. Python's probe that dies on `CancelledError` (a `BaseException`) records nothing and leaves the breaker HALF_OPEN for good, and a dropped Rust future is common
- `_failure_count` is dropped, it only fed a log line

View file

@ -1,5 +1,5 @@
use std::{
sync::Arc,
sync::{Arc, Weak},
time::{Duration, Instant},
};
@ -128,6 +128,9 @@ impl Machine for Config {
#[derive(Clone)]
pub struct Breaker(Arc<Shared<Config>>);
#[derive(Clone)]
pub struct WeakBreaker(Weak<Shared<Config>>);
#[derive(Debug, PartialEq, Eq)]
pub struct Rejected;
@ -165,6 +168,16 @@ impl Breaker {
pub fn state(&self) -> State {
self.0.state()
}
pub fn downgrade(&self) -> WeakBreaker {
WeakBreaker(Arc::downgrade(&self.0))
}
}
impl WeakBreaker {
pub fn upgrade(&self) -> Option<Breaker> {
self.0.upgrade().map(Breaker)
}
}
impl Admission {

View file

@ -1,5 +1,7 @@
mod breaker;
mod fsm;
pub use breaker::{Admission, Breaker, Config, Event, Failure, Output, Rejected, State};
pub use breaker::{
Admission, Breaker, Config, Event, Failure, Output, Rejected, State, WeakBreaker,
};
pub use fsm::{Machine, Observer, Shared, Transition};

View file

@ -0,0 +1,3 @@
## Redis settings
- This crate only deserializes YAML. `RedisSettings` is the raw `coordination_redis` or `cache_params` map and `CacheSettings` tells a `type: redis` cache apart from the rest. Reading those maps the way `litellm/_redis.py` does belongs to `gateway/src/redis`, and turning merged kwargs into a `Resolution` belongs to `redis-resolve`

View file

@ -6,17 +6,16 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-secret-value.workspace = true
litellm-router-types.workspace = true
litellm-secret-value.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_yaml_ng = "0.10.0"
strum.workspace = true
thiserror.workspace = true
tracing.workspace = true
veil.workspace = true
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
serde_json.workspace = true
tempfile.workspace = true
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }

View file

@ -11,11 +11,4 @@ pub enum Error {
path: PathBuf,
source: serde_yaml_ng::Error,
},
#[error("Redis host does not yet support configured setting {0}")]
RedisSetting(String),
#[error("invalid Redis setting {setting}: {reason}")]
RedisInvalid {
setting: &'static str,
reason: &'static str,
},
}

View file

@ -14,10 +14,7 @@ pub use error::Error;
pub use litellm_router_types::{LitellmParams, Spelled};
pub use mcp::{McpAuth, McpServer, McpTransport};
pub use model::Model;
pub use redis::{
CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings,
ResolvedRedis, environment_has_target,
};
pub use redis::{CacheSettings, RedisSettings};
pub use settings::{
ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings,
TracingStoreSettings,

View file

@ -1,195 +1,21 @@
use std::{collections::BTreeMap, fmt, time::Duration};
use std::collections::BTreeMap;
use litellm_router_types::Spelled;
use litellm_secret_value::SecretValue;
use serde::{Deserialize, Deserializer, de::Error as _};
use veil::Redact;
use crate::{AdditionalFields, Error, Object, Value};
use crate::{Object, Value};
pub type Lookup<'a> = &'a dyn Fn(&str) -> Option<String>;
const SUPPORTED: &[&str] = &[
"url",
"host",
"port",
"username",
"password",
"ssl",
"db",
"socket_timeout",
"startup_nodes",
"namespace",
];
const IGNORED: &[&str] = &[
"max_connections",
"single_connection_client",
"health_check_interval",
"socket_keepalive",
"socket_keepalive_options",
"decode_responses",
"encoding",
"encoding_errors",
"charset",
"errors",
"client_name",
"lib_name",
"lib_version",
"sentinel_password",
"gcp_ssl_ca_certs",
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
];
const WARN_IGNORED: &[&str] = &[
"max_connections",
"health_check_interval",
"socket_keepalive",
"socket_keepalive_options",
];
const REDIS_PY_KWARGS: &[&str] = &[
"aws_iam_auth",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
"aws_iam_user_name",
"azure_client_id",
"azure_client_secret",
"azure_redis_ad_token",
"azure_tenant_id",
"cache",
"cache_config",
"charset",
"client_name",
"db",
"decode_responses",
"encoding",
"encoding_errors",
"errors",
"event_dispatcher",
"gcp_service_account",
"gcp_ssl_ca_certs",
"health_check_interval",
"host",
"lib_name",
"lib_version",
"max_connections",
"password",
"port",
"protocol",
"redis_connect_func",
"retry_on_error",
"retry_on_timeout",
"single_connection_client",
"socket_connect_timeout",
"socket_keepalive",
"socket_keepalive_options",
"socket_timeout",
"ssl",
"ssl_ca_certs",
"ssl_ca_data",
"ssl_ca_path",
"ssl_cert_reqs",
"ssl_certfile",
"ssl_check_hostname",
"ssl_ciphers",
"ssl_keyfile",
"ssl_min_version",
"ssl_ocsp_context",
"ssl_ocsp_expected_cert",
"ssl_password",
"ssl_validate_ocsp",
"ssl_validate_ocsp_stapled",
"unix_socket_path",
"url",
"username",
];
const UNSUPPORTED_EXTRA: &[&str] = &[
"credential_provider",
"sentinel_nodes",
"retry",
"connection_pool",
];
const CACHE_ONLY: &[&str] = &[
"type",
"mode",
"ttl",
"default_in_memory_ttl",
"default_in_redis_ttl",
"supported_call_types",
"similarity_threshold",
"max_messages",
"_backend",
];
#[derive(Clone, Default)]
/// A Redis section as `config.yaml` spells it, before any surface reads it.
#[derive(Clone, Default, Deserialize, Redact)]
#[serde(transparent)]
pub struct RedisSettings {
pub url: Option<SecretValue>,
pub host: Option<Spelled<String>>,
pub port: Option<Spelled<u16>>,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub ssl: Option<Spelled<bool>>,
pub db: Option<Spelled<i64>>,
pub socket_timeout: Option<Spelled<f64>>,
pub startup_nodes: Option<Value>,
pub namespace: Option<Spelled<String>>,
pub additional_fields: AdditionalFields,
unparsed: DeferredInputs,
#[redact(with = "[REDACTED]")]
fields: BTreeMap<String, Value>,
}
impl<'de> Deserialize<'de> for RedisSettings {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let fields = BTreeMap::<String, Value>::deserialize(deserializer)?;
let scalar = |key: &str| fields.get(key).and_then(as_text);
Ok(Self {
url: scalar("url").map(SecretValue::new),
host: scalar("host").map(Spelled::Value),
port: spelled(scalar("port")),
username: scalar("username").map(SecretValue::new),
password: scalar("password").map(SecretValue::new),
ssl: spelled(scalar("ssl")),
db: spelled(scalar("db")),
socket_timeout: spelled(scalar("socket_timeout")),
startup_nodes: fields
.get("startup_nodes")
.filter(|value| !value.is_null())
.cloned(),
namespace: scalar("namespace").map(Spelled::Value),
additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)),
unparsed: deferred_inputs(&fields),
})
}
}
impl fmt::Debug for RedisSettings {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("RedisSettings")
.field("url", &self.url)
.field("host", &self.host)
.field("port", &self.port)
.field("username", &self.username)
.field("password", &self.password)
.field("ssl", &self.ssl)
.field("db", &self.db)
.field("socket_timeout", &self.socket_timeout)
.field(
"startup_nodes",
&self.startup_nodes.as_ref().map(|_| "<configured>"),
)
.field("namespace", &self.namespace)
.field("additional_fields", &self.additional_fields.keys())
.field("unparsed", &self.unparsed.keys())
.finish()
impl RedisSettings {
pub fn fields(&self) -> &BTreeMap<String, Value> {
&self.fields
}
}
@ -216,452 +42,22 @@ impl<'de> Deserialize<'de> for CacheSettings {
}
}
#[derive(Clone, PartialEq)]
pub struct ResolvedRedis {
pub endpoint: RedisEndpoint,
pub namespace: Option<String>,
pub timeout: Option<Duration>,
}
#[cfg(test)]
mod tests {
use super::RedisSettings;
use rstest::rstest;
impl fmt::Debug for ResolvedRedis {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ResolvedRedis")
.field("endpoint", &self.endpoint)
.field("namespace", &self.namespace)
.field("timeout", &self.timeout)
.finish()
#[rstest]
#[case::scalar(serde_json::json!("credential-value"))]
#[case::object(serde_json::json!({"secret": "credential-value"}))]
#[case::array(serde_json::json!(["credential-value"]))]
fn credentials_are_redacted(#[case] password: serde_json::Value) {
let settings: RedisSettings = serde_json::from_value(serde_json::json!({
"host": "redis", "password": password,
}))
.unwrap();
let debug = format!("{settings:?}");
assert!(!debug.contains("credential-value"));
assert!(debug.contains("[REDACTED]"));
}
}
#[derive(Clone, PartialEq)]
pub enum RedisEndpoint {
Url(SecretValue),
Host(RedisHost),
Cluster(RedisCluster),
}
impl fmt::Debug for RedisEndpoint {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Url(_) => formatter.debug_tuple("Url").field(&"<redacted>").finish(),
Self::Host(host) => host.fmt(formatter),
Self::Cluster(cluster) => cluster.fmt(formatter),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisHost {
pub host: String,
pub port: u16,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub tls: bool,
pub db: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisCluster {
pub nodes: Vec<RedisNode>,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub tls: bool,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisNode {
pub host: String,
pub port: u16,
}
#[derive(Clone)]
struct Entry {
value: Value,
source: String,
}
type Kwargs = BTreeMap<String, Entry>;
impl RedisSettings {
pub fn resolve(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
self.resolve_with(lookup, false)
}
pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
self.resolve_with(lookup, true)
}
fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result<ResolvedRedis, Error> {
let block = self.kwargs(cache);
let block: Kwargs = block
.into_iter()
.filter_map(|(key, entry)| {
let value = environment_reference(&entry.value, lookup)?;
Some((key, Entry { value, ..entry }))
})
.collect();
let environment = environment_kwargs(lookup);
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
.iter()
.any(|key| block.contains_key(*key));
let drop_environment_url = names_a_target && !block.contains_key("url");
let merged: Kwargs = environment
.into_iter()
.filter(|(key, _)| !(key == "url" && drop_environment_url))
.chain(block)
.collect();
reject_unsupported(&merged, lookup)?;
warn_ignored(&merged);
let startup_nodes = match merged.get("startup_nodes") {
Some(entry) => Some(cluster_nodes(&entry.value)?),
None => lookup("REDIS_CLUSTER_NODES")
.map(|nodes| cluster_nodes(&Value::String(nodes)))
.transpose()?,
};
let username = secret_value(&merged, "username");
let password = secret_value(&merged, "password");
let tls = bool_value(&merged, "ssl");
let endpoint = match (startup_nodes, text_value(&merged, "url")) {
(Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster {
nodes,
username,
password,
tls,
}),
(_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)),
_ => RedisEndpoint::Host(RedisHost {
host: text_value(&merged, "host").ok_or(Error::RedisInvalid {
setting: "host",
reason: "either host or url must be specified for redis",
})?,
port: integer(&merged, "port")?.unwrap_or(6379),
username,
password,
tls,
db: integer(&merged, "db")?.unwrap_or(0),
}),
};
let timeout = merged
.get("socket_timeout")
.and_then(|entry| as_text(&entry.value))
.map(|seconds| {
seconds
.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
.ok_or(Error::RedisInvalid {
setting: "socket_timeout",
reason: "must be a non-negative number of seconds",
})
})
.transpose()?;
Ok(ResolvedRedis {
endpoint,
namespace: text_value(&merged, "namespace"),
timeout,
})
}
fn kwargs(&self, cache: bool) -> Kwargs {
let known = [
(
"url",
self.url.as_ref().map(|value| value.expose().to_owned()),
),
("host", spelled_text(self.host.as_ref())),
("port", spelled_text(self.port.as_ref())),
(
"username",
self.username
.as_ref()
.map(|value| value.expose().to_owned()),
),
(
"password",
self.password
.as_ref()
.map(|value| value.expose().to_owned()),
),
("ssl", spelled_text(self.ssl.as_ref())),
("db", spelled_text(self.db.as_ref())),
("socket_timeout", spelled_text(self.socket_timeout.as_ref())),
("namespace", spelled_text(self.namespace.as_ref())),
];
let known = known.into_iter().filter_map(|(key, value)| {
value.map(|value| {
(
key.to_owned(),
Entry {
value: Value::String(value),
source: key.to_owned(),
},
)
})
});
let startup_nodes = self.startup_nodes.clone().map(|value| {
(
"startup_nodes".to_owned(),
Entry {
value,
source: "startup_nodes".to_owned(),
},
)
});
let extras = self.additional_fields.iter().filter_map(|(key, value)| {
if cache && CACHE_ONLY.contains(&key.as_str()) {
return None;
}
let key = match (cache, key.as_str()) {
(true, "redis_startup_nodes") => "startup_nodes".to_owned(),
_ => key.clone(),
};
Some((
key.clone(),
Entry {
value: value.clone(),
source: key,
},
))
});
let unparsed = self.unparsed.iter().map(|(key, value)| {
(
key.clone(),
Entry {
value: value.0.clone(),
source: key.clone(),
},
)
});
extras
.chain(unparsed)
.chain(known)
.chain(startup_nodes)
.collect()
}
}
pub fn environment_has_target(lookup: Lookup<'_>) -> bool {
[
"REDIS_HOST",
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
]
.iter()
.any(|name| lookup(name).is_some())
}
fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> {
let active_auth = [
merged
.get("gcp_service_account")
.map(|entry| entry.source.as_str())
.or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")),
text_value(merged, "azure_redis_ad_token")
.or_else(|| lookup("REDIS_AZURE_AD_TOKEN"))
.filter(|value| value.eq_ignore_ascii_case("true"))
.map(|_| {
merged
.get("azure_redis_ad_token")
.map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str())
}),
bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()),
merged
.get("sentinel_nodes")
.map(|entry| entry.source.as_str())
.or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")),
];
if let Some(setting) = active_auth.into_iter().flatten().next() {
return Err(Error::RedisSetting(setting.to_owned()));
}
match merged
.iter()
.find(|(key, _)| is_unsupported(key))
.map(|(_, entry)| entry.source.clone())
{
Some(setting) => Err(Error::RedisSetting(setting)),
None => Ok(()),
}
}
fn warn_ignored(merged: &Kwargs) {
for setting in WARN_IGNORED {
if merged.contains_key(*setting) {
tracing::warn!(setting, "Redis setting is accepted but ignored");
}
}
}
fn is_unsupported(key: &str) -> bool {
let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key);
let handled = SUPPORTED.contains(&key)
|| IGNORED.contains(&key)
|| matches!(
key,
"gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth"
);
known && !handled
}
fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs {
REDIS_PY_KWARGS
.iter()
.filter_map(|key| {
let source = format!("REDIS_{}", key.to_ascii_uppercase());
let value = lookup(&source)?;
Some((
(*key).to_owned(),
Entry {
value: Value::String(value),
source,
},
))
})
.collect()
}
fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option<Value> {
match value {
Value::String(text) => match text.strip_prefix("os.environ/") {
Some(name) => lookup(name).map(Value::String),
None => Some(value.clone()),
},
Value::Null => None,
_ => Some(value.clone()),
}
}
fn as_text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}
fn text_value(kwargs: &Kwargs, key: &str) -> Option<String> {
kwargs.get(key).and_then(|entry| as_text(&entry.value))
}
fn secret_value(kwargs: &Kwargs, key: &str) -> Option<SecretValue> {
text_value(kwargs, key).map(SecretValue::new)
}
fn bool_value(kwargs: &Kwargs, key: &str) -> bool {
match kwargs.get(key).map(|entry| &entry.value) {
Some(Value::Bool(flag)) => *flag,
Some(value) => as_text(value)
.is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")),
None => false,
}
}
fn integer<T: std::str::FromStr>(kwargs: &Kwargs, key: &'static str) -> Result<Option<T>, Error> {
kwargs
.get(key)
.and_then(|entry| as_text(&entry.value))
.map(|text| {
text.parse().map_err(|_| Error::RedisInvalid {
setting: key,
reason: "must be an integer",
})
})
.transpose()
}
const INVALID_NODES: Error = Error::RedisInvalid {
setting: "startup_nodes",
reason: "must be a list of {host, port} objects",
};
fn cluster_nodes(value: &Value) -> Result<Vec<RedisNode>, Error> {
match value {
Value::String(json) => {
let nodes: Vec<serde_json::Map<String, serde_json::Value>> =
serde_json::from_str(json).map_err(|_| INVALID_NODES)?;
nodes
.iter()
.map(|node| {
cluster_node(
node.get("host").and_then(json_text),
node.get("port").and_then(json_text),
)
})
.collect()
}
Value::Sequence(nodes) => nodes
.iter()
.map(|node| {
cluster_node(
node.get("host").and_then(as_text),
node.get("port").and_then(as_text),
)
})
.collect(),
_ => Err(INVALID_NODES),
}
}
fn cluster_node(host: Option<String>, port: Option<String>) -> Result<RedisNode, Error> {
match (host, port.and_then(|port| port.parse().ok())) {
(Some(host), Some(port)) => Ok(RedisNode { host, port }),
_ => Err(INVALID_NODES),
}
}
fn json_text(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::String(text) => Some(text.clone()),
serde_json::Value::Number(number) => Some(number.to_string()),
_ => None,
}
}
fn spelled_text<T: ToString>(value: Option<&Spelled<T>>) -> Option<String> {
value.map(|value| match value {
Spelled::Value(value) => value.to_string(),
Spelled::Text(value) => value.clone(),
})
}
fn spelled<T: std::str::FromStr>(text: Option<String>) -> Option<Spelled<T>> {
text.map(|text| match text.parse() {
Ok(value) => Spelled::Value(value),
Err(_) => Spelled::Text(text),
})
}
fn cloned(
fields: &BTreeMap<String, Value>,
keep: impl Fn(&str, &Value) -> bool,
) -> AdditionalFields {
fields
.iter()
.filter(|(key, value)| keep(key, value))
.map(|(key, value)| (key.clone(), value.clone()))
.collect()
}
#[derive(Clone)]
struct DeferredInput(Value);
type DeferredInputs = BTreeMap<String, DeferredInput>;
fn deferred_inputs(fields: &BTreeMap<String, Value>) -> DeferredInputs {
fields
.iter()
.filter(|(key, value)| is_unparsed(key, value))
.map(|(key, value)| (key.clone(), DeferredInput(value.clone())))
.collect()
}
fn is_unparsed(key: &str, value: &Value) -> bool {
SUPPORTED.contains(&key)
&& match key {
"startup_nodes" => value.is_null(),
_ => as_text(value).is_none(),
}
}

View file

@ -440,3 +440,29 @@ fn mcp_config_redacts_nested_credentials_and_preserves_policy_for_validation() {
fn rejects_invalid_typed_mcp_settings(#[case] setting: &str) {
assert!(Config::from_yaml(&format!("mcp_servers:\n docs:\n {setting}\n")).is_err());
}
#[rstest]
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
let config = Config::from_yaml(
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Redis(_))
));
}
#[rstest]
fn non_redis_cache_settings_remain_opaque() {
let config = Config::from_yaml(
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Other(_))
));
}

View file

@ -1,405 +0,0 @@
use std::{
collections::HashMap,
fmt::Write,
sync::{Arc, Mutex},
};
use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis};
use rstest::rstest;
use tracing::{Event, Subscriber, field::Visit};
use tracing_subscriber::{Layer, layer::Context, prelude::*};
fn resolve(
yaml: &str,
environment: &[(&str, &str)],
) -> Result<litellm_config::ResolvedRedis, Error> {
let environment: HashMap<String, String> = environment
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
let config = Config::from_yaml(yaml).unwrap();
config
.general_settings
.coordination_redis
.as_ref()
.unwrap_or(&RedisSettings::default())
.resolve(&|name| environment.get(name).cloned())
}
fn resolve_cache(yaml: &str) -> Result<ResolvedRedis, Error> {
let config = Config::from_yaml(yaml).unwrap();
match config.litellm_settings.cache_params.unwrap() {
CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None),
CacheSettings::Other(_) => panic!("expected a Redis cache"),
}
}
fn cluster_ports(resolved: &ResolvedRedis) -> Vec<u16> {
match &resolved.endpoint {
RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
endpoint => panic!("expected a cluster, got {endpoint:?}"),
}
}
#[rstest]
#[case::host(
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
&[],
"host:redis:6380:2:true:app:1500"
)]
#[case::environment_url(
"model_list: []",
&[("REDIS_URL", "redis://environment:6379")],
"url:redis://environment:6379"
)]
#[case::explicit_host_drops_environment_url(
"general_settings:\n coordination_redis:\n host: configured",
&[("REDIS_URL", "redis://environment:6379")],
"host:configured:6379:0:false::none"
)]
#[case::yaml_cluster(
"general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"",
&[],
"cluster:node-a:7000,node-b:7001"
)]
#[case::json_environment_cluster(
"model_list: []",
&[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)],
"cluster:node-a:7000,node-b:7001"
)]
#[case::missing_reference_keeps_environment_fallback(
"general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
#[case::configured_values_beat_the_same_environment_values(
"general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3",
&[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")],
"host:configured:6380:3:false::none"
)]
#[case::environment_url_beats_the_environment_host(
"model_list: []",
&[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")],
"url:redis://environment:6379"
)]
#[case::startup_nodes_beat_a_url(
"general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000",
&[],
"cluster:node-a:7000"
)]
#[case::environment_references_resolve_the_host_and_port(
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
"host:referenced:6400:0:false::none"
)]
#[case::a_null_host_falls_back_to_the_environment(
"general_settings:\n coordination_redis:\n host: null",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
fn resolves_redis_through_the_parsed_config(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolved = resolve(yaml, environment).unwrap();
let description = match resolved.endpoint {
RedisEndpoint::Url(url) => format!("url:{}", url.expose()),
RedisEndpoint::Host(host) => format!(
"host:{}:{}:{}:{}:{}:{}",
host.host,
host.port,
host.db,
host.tls,
resolved.namespace.unwrap_or_default(),
resolved
.timeout
.map(|timeout| timeout.as_millis().to_string())
.unwrap_or_else(|| "none".to_owned())
),
RedisEndpoint::Cluster(cluster) => format!(
"cluster:{}",
cluster
.nodes
.iter()
.map(|node| format!("{}:{}", node.host, node.port))
.collect::<Vec<_>>()
.join(",")
),
};
assert_eq!(description, expected);
}
#[rstest]
#[case::custom_tls_file(
"general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem",
&[],
"ssl_ca_certs"
)]
#[case::gcp_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account",
&[],
"gcp_service_account"
)]
#[case::gcp_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")],
"REDIS_GCP_SERVICE_ACCOUNT"
)]
#[case::azure_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true",
&[],
"azure_redis_ad_token"
)]
#[case::azure_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")],
"REDIS_AZURE_AD_TOKEN"
)]
#[case::aws_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true",
&[],
"aws_iam_auth"
)]
#[case::aws_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")],
"REDIS_AWS_IAM_AUTH"
)]
#[case::sentinel_from_config(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]",
&[],
"sentinel_nodes"
)]
#[case::sentinel_from_environment(
"model_list: []",
&[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")],
"REDIS_SENTINEL_NODES"
)]
fn rejected_settings_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(matches!(error, Error::RedisSetting(setting) if setting == expected));
}
#[rstest]
fn an_explicit_host_keeps_environment_credentials() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: configured",
&[
("REDIS_PASSWORD", "environment-secret"),
("REDIS_SSL", "True"),
],
)
.unwrap();
let RedisEndpoint::Host(host) = resolved.endpoint else {
panic!("expected a host endpoint");
};
assert_eq!(
host.password.as_ref().map(|password| password.expose()),
Some("environment-secret")
);
assert!(host.tls);
}
#[rstest]
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
#[case::a_structured_host_with_an_environment_target(
"general_settings:\n coordination_redis:\n host: []",
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
)]
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
let error = resolve(yaml, environment).unwrap_err();
assert!(matches!(
error,
Error::RedisInvalid {
setting: "host",
..
}
));
}
#[rstest]
fn the_cache_alias_names_cluster_nodes() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn canonical_startup_nodes_beat_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
#[rstest]
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
#[rstest]
fn rejected_environment_settings_never_print_the_value() {
let error = resolve(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CERT_REQS", "private-option-value"),
],
)
.unwrap_err();
let message = error.to_string();
assert!(message.contains("REDIS_SSL_CERT_REQS"));
assert!(!message.contains("private-option-value"));
}
#[rstest]
#[case::negative("-1")]
#[case::not_a_number("slow")]
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
let yaml = format!(
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
);
let error = resolve(&yaml, &[]).unwrap_err();
assert!(matches!(
error,
Error::RedisInvalid {
setting: "socket_timeout",
..
}
));
}
#[rstest]
fn defaults_socket_timeout_for_the_gateway() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis",
&[],
)
.unwrap();
assert_eq!(resolved.timeout, None);
}
#[rstest]
fn debug_redacts_redis_secrets_and_url() {
let config = Config::from_yaml(
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password",
)
.unwrap();
let settings = config.general_settings.coordination_redis.unwrap();
let resolved = settings.resolve(&|_| None).unwrap();
let debug = format!("{settings:?} {resolved:?}");
assert!(!debug.contains("redis://user:password"));
assert!(!debug.contains("credential-user"));
assert!(!debug.contains("credential-password"));
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
}
#[rstest]
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
let config = Config::from_yaml(
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Redis(_))
));
}
#[rstest]
fn non_redis_cache_settings_remain_opaque() {
let config = Config::from_yaml(
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Other(_))
));
}
#[rstest]
fn legacy_unknown_settings_do_not_change_resolution() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
&[],
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
struct WarningCapture(Arc<Mutex<Vec<String>>>);
impl<S: Subscriber> Layer<S> for WarningCapture {
fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) {
let mut fields = Fields::default();
event.record(&mut fields);
self.0.lock().unwrap().push(fields.0);
}
}
#[derive(Default)]
struct Fields(String);
impl Visit for Fields {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
write!(self.0, "{}={value:?} ", field.name()).unwrap();
}
fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
write!(self.0, "{}={value} ", field.name()).unwrap();
}
}
#[rstest]
fn ignored_setting_warning_names_the_setting_without_its_value() {
let events = Arc::new(Mutex::new(Vec::new()));
let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone()));
tracing::subscriber::with_default(subscriber, || {
resolve(
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
&[],
)
.unwrap();
});
let warnings = events.lock().unwrap().join(" ");
assert!(warnings.contains("max_connections"));
assert!(!warnings.contains("47"));
}

View file

@ -5,8 +5,11 @@
- Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate
## Redis
- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
- `build_redis` builds the Redis clients once at boot (`src/redis/mod.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
- Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands
- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped
- `src/redis/layers.rs` is the Rust copy of `_merge_redis_kwargs` in `litellm/_redis.py`, which only this host needs because the Python host merges in Python. It puts the configured kwargs over their `REDIS_*` values and fills the `get_secret` fallbacks from the `SETTINGS` rows in `redis-resolve`, under the same conditions Python reads them, then hands the result to `litellm_redis_resolve::interpret`. Values come from a `SecretSnapshot` filled once at boot, read like `get_secret` (secret manager, then environment, with `true` and `false` read as bools), and the only straight environment read is Azure's `REDIS_USERNAME`
- A surface is a `ConfiguredRedis` constructor: `coordination` for `general_settings.coordination_redis`, which rejects a malformed `port` like the proxy's `CoordinationRedisParams` pydantic model, `cache_params` for `litellm_settings.cache_params`, which drops the cache-only keys and reads the `redis_startup_nodes` alias, and `redis_cache_kwargs` for a `RedisCache.redis_kwargs` dict in the contract tests. A new surface is a new constructor, never a flag
- Nulls depend on the surface. `coordination_redis` drops every null like `model_dump(exclude_none=True)`, `cache_params` keeps a null as an explicit `None` that hides the `REDIS_*` value unless `Cache` or `RedisCache` names that parameter, and `RedisCache.redis_kwargs` keeps every null. A missing `os.environ/` reference is a null, as `load_config` resolves it
- A setting the Rust client cannot honor warns and is dropped, and only a security check it cannot enforce fails boot, per `redis-resolve/AGENTS.md`
- Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python
- `/health/readiness` never touches Redis, matching the proxy

View file

@ -20,9 +20,14 @@ futures-util.workspace = true
litellm-gateway-auth.workspace = true
litellm-gateway-ui.workspace = true
litellm-config.workspace = true
litellm-core-utils.workspace = true
litellm-error.workspace = true
litellm-circuit-breaker.workspace = true
litellm-redis-client.workspace = true
litellm-redis-connect.workspace = true
litellm-redis-resolve.workspace = true
litellm-redis-types.workspace = true
rustls.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
litellm-secrets.workspace = true
@ -32,10 +37,12 @@ serde.workspace = true
tracing.workspace = true
tokio = { workspace = true, features = ["signal"] }
uuid.workspace = true
veil.workspace = true
tower-sessions-moka-store = "0.15.0"
[dev-dependencies]
litellm-redis-testing.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
base64.workspace = true
rstest.workspace = true
tempfile.workspace = true

View file

@ -12,13 +12,12 @@ pub enum Error {
Mcp(#[from] litellm_gateway_mcp::ConnectError),
#[error(transparent)]
Config(#[from] litellm_config::Error),
#[error("invalid Redis setting {setting}: {reason}")]
RedisInvalid {
setting: &'static str,
reason: &'static str,
},
#[error(transparent)]
RedisConnect(#[from] litellm_redis_client::ConnectError),
Redis(#[from] litellm_redis_resolve::Error),
#[error("Redis setting {setting} is not plain data")]
RedisNotPlainData { setting: String },
#[error(transparent)]
RedisConnect(#[from] litellm_redis_connect::Error),
#[error("MCP requires working gateway authentication")]
Auth(#[from] litellm_gateway_auth::Error),
}

View file

@ -68,15 +68,29 @@ pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, Error> {
Ok(Arc::new(Gateway::new(resources, http, secrets, router)?))
}
pub async fn build_redis(config: &Config) -> Result<RedisClients, Error> {
let environment = secrets::environment_values(&config.environment_variables)?;
let lookup = |name: &str| {
pub fn install_crypto_provider() {
if rustls::crypto::CryptoProvider::get_default().is_none() {
let _ = rustls::crypto::ring::default_provider().install_default();
}
}
pub async fn build_redis(
config: &Config,
pool: &HttpClientPool,
http: &litellm_http::HttpClientConfig,
) -> Result<RedisClients, Error> {
let environment = Arc::new(secrets::environment_values(&config.environment_variables)?);
let lookup: litellm_redis_connect::SharedEnvironment = Arc::new(move |name: &str| {
environment
.get(name)
.map(|value| value.expose().to_owned())
.or_else(|| std::env::var(name).ok())
});
let resources = redis::Resources {
http: pool.client(http, ClientVariant::Provider)?,
lookup: lookup.clone(),
};
redis::connect(config, &lookup).await
redis::connect(config, &*lookup, &resources).await
}
pub async fn build_mcp(

View file

@ -41,6 +41,7 @@ impl Sink for StderrSink {
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
litellm_gateway::install_crypto_provider();
if let Some(forward) = litellm_gateway_forward::Settings::from_env()? {
let log: LogSettings = settings::from_iter(std::env::vars_os())?;
Logger::new(StderrSink {
@ -55,7 +56,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
Logger::new(StderrSink { level }).install_global()?;
let config = Config::load(settings.litellm_config)?;
let inference = litellm_gateway::build_inference(&config)?;
let _redis = litellm_gateway::build_redis(&config).await?;
let _redis =
litellm_gateway::build_redis(&config, &inference.resources.pool, &inference.http).await?;
let ui = match std::env::var_os("LITELLM_UI_PATH") {
Some(directory) => {
let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?;

View file

@ -0,0 +1,144 @@
use std::collections::BTreeMap;
use litellm_config::{RedisSettings, Value};
use litellm_core_utils::settings::Lookup;
use litellm_redis_resolve::{Expected, Kwargs};
use litellm_redis_types::{ProviderHandle, Resolution};
use veil::Redact;
use super::{layers, secrets::SecretSnapshot};
use crate::Error;
const CACHE_ONLY: &[&str] = &[
"type",
"mode",
"ttl",
"default_in_memory_ttl",
"default_in_redis_ttl",
"supported_call_types",
"similarity_threshold",
"max_messages",
"_backend",
];
/// The redis-py kwargs one surface hands `_get_redis_client_logic`, before `REDIS_*` and secrets.
#[derive(Clone, Redact)]
pub(crate) struct ConfiguredRedis {
#[redact(with = "[REDACTED]")]
pub(super) fields: BTreeMap<String, Value>,
pub(super) nulls: Nulls,
pub(super) credential_provider: Option<ProviderHandle>,
}
/// Whether a null setting reaches redis-py as an explicit `None`, which outranks its `REDIS_*` value.
#[derive(Clone, Copy, Debug)]
pub(super) enum Nulls {
/// `coordination_redis` reaches redis-py through `model_dump(exclude_none=True)`.
Dropped,
/// `Cache` drops a null only for the parameters it names and passes any other key on.
KeptUnlessNamedByCache,
/// `RedisCache.redis_kwargs` already holds only the nulls redis-py receives.
#[cfg(test)]
Kept,
}
impl ConfiguredRedis {
/// `general_settings.coordination_redis`, which the proxy validates as a pydantic model first.
/// An `os.environ/` port is only known after merging, where it is read like any other port.
pub(crate) fn coordination(
settings: &RedisSettings,
) -> Result<Self, litellm_redis_resolve::Error> {
if let Some(port) = settings.fields().get("port").and_then(text)
&& !port.starts_with("os.environ/")
&& port.parse::<u16>().is_err()
{
return Err(litellm_redis_resolve::Error::Invalid {
setting: "port".to_owned(),
expected: Expected::Integer,
});
}
Ok(Self {
fields: settings.fields().clone(),
nulls: Nulls::Dropped,
credential_provider: None,
})
}
/// `litellm_settings.cache_params` with `type: redis`, as `litellm.Cache` hands it to `RedisCache`.
pub(crate) fn cache_params(settings: &RedisSettings) -> Self {
let fields = settings.fields();
let alias = fields
.get("redis_startup_nodes")
.filter(|_| !fields.contains_key("startup_nodes"))
.map(|nodes| ("startup_nodes".to_owned(), nodes.clone()));
let fields = fields
.iter()
.filter(|(key, _)| {
key.as_str() != "redis_startup_nodes" && !CACHE_ONLY.contains(&key.as_str())
})
.map(|(key, value)| (key.clone(), value.clone()))
.chain(alias)
.collect();
Self {
fields,
nulls: Nulls::KeptUnlessNamedByCache,
credential_provider: None,
}
}
/// `RedisCache.redis_kwargs`, the dict every Python surface ends in.
#[cfg(test)]
pub(crate) fn redis_cache_kwargs(fields: BTreeMap<String, Value>) -> Self {
Self {
fields,
nulls: Nulls::Kept,
credential_provider: None,
}
}
/// The handle the host holds for the `credential_provider` object in these settings.
#[cfg(test)]
pub(crate) fn with_credential_provider(self, handle: ProviderHandle) -> Self {
Self {
credential_provider: Some(handle),
..self
}
}
/// No configured settings, so only `REDIS_*` and the other environment fallbacks apply.
pub(crate) fn environment_only() -> Self {
Self {
fields: BTreeMap::new(),
nulls: Nulls::Dropped,
credential_provider: None,
}
}
/// `_merge_redis_kwargs` in `litellm/_redis.py`: the kwargs over their `REDIS_*` values with every
/// `get_secret` fallback filled.
pub(crate) fn merged(&self, secrets: &SecretSnapshot) -> Result<Kwargs, Error> {
layers::merged(self, secrets)
}
/// The environment is only for what Python reads straight from `os.environ`, `REDIS_USERNAME` for Azure.
pub(crate) fn resolve(
&self,
secrets: &SecretSnapshot,
environment: &dyn Lookup,
) -> Result<Resolution, Error> {
Ok(litellm_redis_resolve::interpret(
&self.merged(secrets)?,
self.credential_provider,
environment.get("REDIS_USERNAME").as_deref(),
)?)
}
}
pub(super) fn text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}

View file

@ -0,0 +1,226 @@
use litellm_config::Value;
use litellm_core_utils::settings::Lookup;
use litellm_redis_resolve::{Entry, Kwargs, SETTINGS, is_present, setting};
use super::configured::{ConfiguredRedis, Nulls};
use crate::Error;
const NAMED_BY_CACHE: &[&str] = &[
"host",
"port",
"password",
"socket_timeout",
"startup_nodes",
"namespace",
"gcp_service_account",
"gcp_ssl_ca_certs",
];
/// How `_merge_redis_kwargs` reads a fallback: `get_secret` turns `true` and `false` into a bool, and
/// `get_secret_str` treats that bool as unset.
#[derive(Clone, Copy)]
enum Read {
Secret,
SecretStr,
}
const TARGETS: &[(&str, Read)] = &[
("startup_nodes", Read::Secret),
("sentinel_nodes", Read::Secret),
("sentinel_password", Read::SecretStr),
("service_name", Read::Secret),
];
const IDENTITIES: &[(&str, Read)] = &[
("gcp_service_account", Read::SecretStr),
("gcp_ssl_ca_certs", Read::SecretStr),
("azure_redis_ad_token", Read::Secret),
];
const AZURE: &[(&str, Read)] = &[
("azure_client_id", Read::SecretStr),
("azure_tenant_id", Read::SecretStr),
("azure_client_secret", Read::SecretStr),
];
const AWS: &[(&str, Read)] = &[("aws_iam_region", Read::SecretStr)];
/// The configured settings over their `REDIS_*` values, `{**environment_kwargs, **env_overrides}`
/// in `_merge_redis_kwargs`, with the `get_secret` fallbacks it reads filled in. A null left in the result is
/// an explicit `None` that hid its environment value.
pub(super) fn merged(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result<Kwargs, Error> {
let block: Kwargs = configured
.fields
.iter()
.map(|(key, value)| (key, environment_reference(value, lookup)))
.filter(|(key, value)| !value.is_null() || keeps_null(configured.nulls, key))
.map(|(key, value)| {
let value = serde_json::to_value(&value).map_err(|_| Error::RedisNotPlainData {
setting: key.clone(),
})?;
let entry = Entry {
value,
source: key.clone(),
};
Ok((key.clone(), entry))
})
.collect::<Result<_, Error>>()?;
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
.iter()
.any(|key| is_set(&block, key));
let drop_environment_url = names_a_target && !is_set(&block, "url");
let layered: Kwargs = environment(lookup)
.filter(|(key, _)| !(key == "url" && drop_environment_url))
.chain(block)
.collect();
let fallbacks = fallbacks(&layered, configured.credential_provider.is_some(), lookup);
Ok(layered.into_iter().chain(fallbacks).collect())
}
// The fallbacks are read under the same conditions as _merge_redis_kwargs, so a secret Python never asks
// for is never filled in
fn fallbacks(layered: &Kwargs, has_provider: bool, lookup: &dyn Lookup) -> Kwargs {
let fill = |kwargs: &[(&str, Read)]| -> Kwargs {
kwargs
.iter()
.filter_map(|(kwarg, read)| fallback(layered, kwarg, *read, lookup))
.collect()
};
let targets = fill(TARGETS);
if has_provider || is_set(layered, "credential_provider") {
return targets;
}
let identities = fill(IDENTITIES);
let value = |kwarg: &str| {
layered
.get(kwarg)
.filter(|entry| is_present(&entry.value))
.or_else(|| identities.get(kwarg))
.map(|entry| &entry.value)
};
let gcp = value("gcp_service_account").is_some();
let azure = value("azure_redis_ad_token")
.is_some_and(|token| text(token).is_some_and(|token| token.eq_ignore_ascii_case("true")));
let aws = layered
.get("aws_iam_auth")
.is_some_and(|entry| coerces_to_true(&entry.value));
let azure_settings = match azure && !gcp {
true => fill(AZURE),
false => Kwargs::new(),
};
let aws_settings = match aws && !gcp && !azure {
true => fill(AWS),
false => Kwargs::new(),
};
targets
.into_iter()
.chain(identities)
.chain(azure_settings)
.chain(aws_settings)
.collect()
}
fn fallback(
layered: &Kwargs,
kwarg: &str,
read: Read,
lookup: &dyn Lookup,
) -> Option<(String, Entry)> {
if layered
.get(kwarg)
.is_some_and(|entry| is_present(&entry.value))
{
return None;
}
let names = setting(kwarg).map_or(&[][..], |setting| setting.fallback);
names.iter().find_map(|name| {
let value = lookup
.get(name)
.filter(|value| !value.is_empty())
.map(secret)?;
match (read, &value) {
(Read::SecretStr, serde_json::Value::Bool(_)) => None,
_ => {
let entry = Entry {
value,
source: (*name).to_owned(),
};
Some((kwarg.to_owned(), entry))
}
}
})
}
// get_secret reads an environment value of true or false, in any case and padding, as a bool
fn secret_flag(value: &str) -> Option<bool> {
match value.trim().to_ascii_lowercase().as_str() {
"true" => Some(true),
"false" => Some(false),
_ => None,
}
}
fn secret(value: String) -> serde_json::Value {
secret_flag(&value).map_or(serde_json::Value::String(value), serde_json::Value::Bool)
}
// _coerces_to_true: text is read like _str_to_bool, anything else by its truthiness
fn coerces_to_true(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::String(text) => {
matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")
}
serde_json::Value::Null => false,
serde_json::Value::Bool(flag) => *flag,
serde_json::Value::Number(number) => number.as_f64().is_some_and(|number| number != 0.0),
serde_json::Value::Array(items) => !items.is_empty(),
serde_json::Value::Object(entries) => !entries.is_empty(),
}
}
fn text(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::String(text) => Some(text.clone()),
serde_json::Value::Number(number) => Some(number.to_string()),
serde_json::Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}
fn keeps_null(nulls: Nulls, key: &str) -> bool {
match nulls {
Nulls::Dropped => false,
Nulls::KeptUnlessNamedByCache => !NAMED_BY_CACHE.contains(&key),
#[cfg(test)]
Nulls::Kept => true,
}
}
fn environment(lookup: &dyn Lookup) -> impl Iterator<Item = (String, Entry)> + '_ {
SETTINGS.iter().filter_map(|setting| {
let source = setting.environment_name()?;
let value = lookup.get(&source)?;
let entry = Entry {
value: secret(value),
source,
};
Some((setting.kwarg.to_owned(), entry))
})
}
// load_config resolves os.environ/ to None when unset, so a missing reference is a null like any other
fn environment_reference(value: &Value, lookup: &dyn Lookup) -> Value {
match value {
Value::String(text) => match text.strip_prefix("os.environ/") {
Some(name) => lookup.get(name).map_or(Value::Null, |value| {
secret_flag(&value).map_or(Value::String(value), Value::Bool)
}),
None => value.clone(),
},
_ => value.clone(),
}
}
fn is_set(kwargs: &Kwargs, key: &str) -> bool {
kwargs.get(key).is_some_and(|entry| !entry.value.is_null())
}

View file

@ -1,19 +1,31 @@
use std::time::Duration;
use std::{sync::Arc, time::Duration};
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State};
use litellm_config::{
CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value,
environment_has_target,
};
use litellm_redis_client::{
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient,
RedisConfig,
};
use litellm_circuit_breaker::Config as BreakerConfig;
use litellm_config::{CacheSettings, Config, RedisSettings, Value};
use litellm_core_utils::settings::Lookup;
use litellm_redis_client::{RedisClient, RedisEvents};
use litellm_redis_connect::{HostServices, RedisHosts, SharedEnvironment, TracingEvents};
use litellm_redis_resolve::Expected;
use litellm_redis_types::Resolution;
use crate::Error;
mod configured;
mod layers;
mod secrets;
#[cfg(test)]
mod surface_tests;
use configured::ConfiguredRedis;
use secrets::SecretSnapshot;
pub(crate) struct Resources {
pub(crate) http: litellm_http::Client,
pub(crate) lookup: SharedEnvironment,
}
const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2);
type Params = ResolvedRedis;
type Params = Resolution;
pub struct RedisClients {
pub coordination: Option<RedisClient>,
@ -34,12 +46,15 @@ enum Target {
Environment(Params),
}
pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
pub(crate) fn plan(config: &Config, lookup: &dyn Lookup) -> Result<Plan, Error> {
let response_cache = redis_cache_params(config)
.map(|settings| settings.resolve_cache(lookup))
.map(|settings| resolve(&ConfiguredRedis::cache_params(settings), lookup))
.transpose()?;
let coordination = match &config.general_settings.coordination_redis {
Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)),
Some(settings) => Some(Target::Configured(resolve(
&ConfiguredRedis::coordination(settings)?,
lookup,
)?)),
None => match (&response_cache, config.litellm_settings.cache == Some(true)) {
(Some(cache), _) => Some(Target::Configured(cache.clone())),
(None, true) => environment(lookup)?.map(Target::Configured),
@ -57,12 +72,30 @@ pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
})
}
pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients, Error> {
pub(crate) async fn connect(
config: &Config,
lookup: &dyn Lookup,
resources: &Resources,
) -> Result<RedisClients, Error> {
let plan = plan(config, lookup)?;
let breaker = breaker_config(lookup)?;
let events: Arc<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
let hosts = RedisHosts::default();
let client = |params: Params| -> Result<RedisClient, Error> {
Ok(hosts.client(
params,
HostServices {
events: Some(events.clone()),
http: resources.http.clone(),
environment: resources.lookup.clone(),
credential_provider: None,
},
breaker,
)?)
};
let coordination = match plan.coordination {
Some(Target::Configured(params)) => Some(client(params, breaker)?),
Some(Target::Environment(params)) => match client(params, breaker) {
Some(Target::Configured(params)) => Some(client(params)?),
Some(Target::Environment(params)) => match client(params) {
Ok(client) => reachable(client).await,
Err(_) => {
tracing::warn!(
@ -75,13 +108,10 @@ pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients
};
let transaction_buffer = match (&coordination, plan.transaction_buffer_fallback) {
(Some(coordination), _) => Some(coordination.clone()),
(None, Some(params)) => Some(client(params, breaker)?),
(None, Some(params)) => Some(client(params)?),
(None, None) => None,
};
let response_cache = plan
.response_cache
.map(|params| client(params, breaker))
.transpose()?;
let response_cache = plan.response_cache.map(client).transpose()?;
Ok(RedisClients {
coordination,
transaction_buffer,
@ -98,17 +128,38 @@ fn redis_cache_params(config: &Config) -> Option<&RedisSettings> {
}
}
fn environment(lookup: Lookup<'_>) -> Result<Option<Params>, Error> {
fn resolve(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result<Params, Error> {
let secrets = SecretSnapshot::read(configured.secret_names(), lookup);
let resolution = configured.resolve(&secrets, lookup)?;
for ignored in &resolution.ignored {
tracing::warn!(
"Redis setting {} is not supported by the Rust client and is ignored: {}",
ignored.setting,
ignored.consequence
);
}
Ok(resolution)
}
fn environment(lookup: &dyn Lookup) -> Result<Option<Params>, Error> {
match environment_has_target(lookup) {
true => RedisSettings::default()
.resolve(lookup)
.map(Some)
.map_err(Into::into),
true => Ok(Some(resolve(&ConfiguredRedis::environment_only(), lookup)?)),
false => Ok(None),
}
}
fn optional_environment(lookup: Lookup<'_>) -> Option<Params> {
fn environment_has_target(lookup: &dyn Lookup) -> bool {
[
"REDIS_HOST",
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
]
.iter()
.any(|name| lookup.get(name).is_some())
}
fn optional_environment(lookup: &dyn Lookup) -> Option<Params> {
match environment(lookup) {
Ok(params) => params,
Err(error) => {
@ -149,122 +200,101 @@ async fn reachable(client: RedisClient) -> Option<RedisClient> {
}
}
fn client(params: Params, breaker: Option<BreakerConfig>) -> Result<RedisClient, Error> {
tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured");
let endpoint = match params.endpoint {
RedisEndpoint::Url(url) => Endpoint::Url(url),
RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint {
host: host.host,
port: host.port,
username: host.username,
password: host.password,
tls: host.tls,
db: host.db,
}),
RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint {
nodes: cluster
.nodes
.into_iter()
.map(|node| ClusterNode {
host: node.host,
port: node.port,
})
.collect(),
username: cluster.username,
password: cluster.password,
tls: cluster.tls,
}),
};
Ok(RedisClient::new(RedisConfig {
endpoint,
namespace: params.namespace.map(Namespace::new).unwrap_or_default(),
timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT),
breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))),
})?)
fn parse_seconds(text: &str) -> Option<Duration> {
text.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
}
fn log_transition(transition: &litellm_circuit_breaker::Transition<State>) {
match transition.to {
State::Open { .. } => tracing::warn!(
generation = transition.generation,
"redis circuit breaker opened, fast-failing redis calls"
),
State::HalfOpen => tracing::info!(
generation = transition.generation,
"redis circuit breaker half open, probing redis"
),
State::Closed { .. } => tracing::info!(
generation = transition.generation,
"redis circuit breaker closed, redis recovered"
),
}
}
fn breaker_config(lookup: Lookup<'_>) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED")
fn breaker_config(lookup: &dyn Lookup) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup
.get("REDIS_CIRCUIT_BREAKER_ENABLED")
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
if !enabled {
return Ok(None);
}
let defaults = BreakerConfig::default();
let seconds = |name: &'static str, default: Duration| -> Result<Duration, Error> {
lookup(name)
.map(|value| {
value
.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
.ok_or(Error::RedisInvalid {
setting: name,
reason: "must be a non-negative number of seconds",
})
})
.unwrap_or(Ok(default))
};
let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD")
Ok(Some(BreakerConfig {
failure_threshold: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
Expected::Integer,
|text| text.parse().ok(),
)?
.unwrap_or(defaults.failure_threshold),
recovery_timeout: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.recovery_timeout),
timeout_min_duration: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.timeout_min_duration),
}))
}
fn environment_value<T>(
lookup: &dyn Lookup,
name: &str,
expected: Expected,
parse: impl Fn(&str) -> Option<T>,
) -> Result<Option<T>, litellm_redis_resolve::Error> {
lookup
.get(name)
.map(|value| {
value.parse().map_err(|_| Error::RedisInvalid {
setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
reason: "must be an integer",
parse(&value).ok_or_else(|| litellm_redis_resolve::Error::Invalid {
setting: name.to_owned(),
expected,
})
})
.unwrap_or(Ok(defaults.failure_threshold))?;
Ok(Some(BreakerConfig {
failure_threshold,
recovery_timeout: seconds(
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
defaults.recovery_timeout,
)?,
timeout_min_duration: seconds(
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
defaults.timeout_min_duration,
)?,
}))
.transpose()
}
#[cfg(test)]
mod tests {
use litellm_core_utils::settings::Lookup;
use std::{collections::HashMap, time::Duration};
use litellm_config::{Config, RedisEndpoint, RedisHost};
use litellm_config::Config;
use litellm_redis_client::{Standalone, StaticAuth, commands};
use litellm_redis_testing::RedisServer;
use litellm_redis_types::{Auth, HealthCheck};
use rstest::rstest;
use super::{Params, Plan, Target, connect, plan};
use std::sync::Arc;
use super::{Params, Plan, Resources, Target, plan};
use crate::Error;
async fn connect(config: &Config, lookup: &dyn Lookup) -> Result<super::RedisClients, Error> {
let resources = Resources {
http: litellm_http::Client::for_test(),
lookup: Arc::new(|_: &str| None),
};
super::connect(config, lookup, &resources).await
}
fn host(name: &str, port: u16) -> Params {
Params {
endpoint: RedisEndpoint::Host(RedisHost {
target: litellm_redis_client::Target::Standalone(Standalone {
host: name.to_owned(),
port,
username: None,
password: None,
tls: false,
tls: None,
db: 0,
}),
auth: Auth::Static(StaticAuth::default()),
namespace: None,
timeout: None,
connect_timeout: None,
health_check: HealthCheck::Unset,
keepalive: true,
ignored: Vec::new(),
}
}
@ -273,7 +303,7 @@ mod tests {
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
plan(&Config::from_yaml(yaml).unwrap(), &|name| {
plan(&Config::from_yaml(yaml).unwrap(), &|name: &str| {
environment.get(name).cloned()
})
}
@ -350,36 +380,44 @@ mod tests {
port_of(&server)
))
.unwrap();
let clients = connect(&config, &|_| None).await.unwrap();
let clients = connect(&config, &|_: &str| None).await.unwrap();
let coordination = clients.coordination.unwrap();
let buffer = clients.transaction_buffer.unwrap();
let cache = clients.response_cache.unwrap();
let unnamespaced = connect(&coordination_config(&server), &|_| None)
let unnamespaced = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap()
.coordination
.unwrap();
coordination
.set("rate-limit", "1", Duration::from_secs(60))
.send(commands::set("rate-limit", "1", Duration::from_secs(60)))
.await
.unwrap();
buffer
.set("spend", "2", Duration::from_secs(60))
.send(commands::set("spend", "2", Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(
cache.get::<String>("rate-limit").await.unwrap().as_deref(),
cache
.send(commands::get::<String>("rate-limit"))
.await
.unwrap()
.as_deref(),
Some("1")
);
assert_eq!(
cache.get::<String>("spend").await.unwrap().as_deref(),
cache
.send(commands::get::<String>("spend"))
.await
.unwrap()
.as_deref(),
Some("2")
);
assert_eq!(
unnamespaced
.get::<String>("responses:rate-limit")
.send(commands::get::<String>("responses:rate-limit"))
.await
.unwrap()
.as_deref(),
@ -387,7 +425,7 @@ mod tests {
);
assert_eq!(
unnamespaced
.get::<String>("responses:spend")
.send(commands::get::<String>("responses:spend"))
.await
.unwrap()
.as_deref(),
@ -395,16 +433,21 @@ mod tests {
);
assert!(
unnamespaced
.get::<String>("rate-limit")
.send(commands::get::<String>("rate-limit"))
.await
.unwrap()
.is_none()
);
assert!(
unnamespaced
.send(commands::get::<String>("spend"))
.await
.unwrap()
.is_none()
);
assert!(unnamespaced.get::<String>("spend").await.unwrap().is_none());
}
#[rstest]
#[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
#[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])]
fn redis_migration_regression_malformed_optional_environment_is_absent(
#[case] environment: &[(&str, &str)],
@ -413,54 +456,45 @@ mod tests {
}
#[rstest]
fn redis_migration_regression_explicit_malformed_setting_still_fails() {
#[case::optional_environment("model_list: []", &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
#[case::transaction_buffer_environment(
"general_settings:\n use_redis_transaction_buffer: true",
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")]
)]
#[case::response_cache(
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: env\n port: not-a-port",
&[]
)]
fn a_malformed_port_outside_the_coordination_block_falls_back_to_6379_like_redis_py(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let plan = resolved(yaml, environment).unwrap();
let target = match (plan.coordination, plan.transaction_buffer_fallback) {
(Some(Target::Configured(params) | Target::Environment(params)), _)
| (None, Some(params)) => params,
(None, None) => panic!("expected a Redis target"),
};
assert_eq!(target.target, host("env", 6379).target);
}
#[rstest]
#[case::port("port: not-a-port", "port")]
#[case::ssl("ssl: maybe", "ssl")]
fn a_malformed_coordination_block_fails_like_the_pydantic_model(
#[case] setting: &str,
#[case] expected: &str,
) {
let error = resolved(
"general_settings:\n coordination_redis:\n host: coord\n port: not-a-port",
&format!("general_settings:\n coordination_redis:\n host: coord\n {setting}"),
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() {
let error = resolved(
"general_settings:\n use_redis_transaction_buffer: true",
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
fn redis_migration_regression_invalid_response_cache_still_fails() {
let error = resolved(
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port",
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
assert!(
matches!(error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == expected)
);
}
#[rstest]
@ -468,7 +502,7 @@ mod tests {
async fn redis_migration_regression_invalid_optional_environment_url_is_absent() {
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| {
let clients = connect(&config, &|name: &str| {
(name == "REDIS_URL").then(|| "not a redis URL".to_owned())
})
.await
@ -506,7 +540,7 @@ mod tests {
let server = RedisServer::start().await.unwrap();
server.stop().await.unwrap();
let clients = connect(&coordination_config(&server), &|_| None)
let clients = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap();
let coordination = clients.coordination.unwrap();
@ -521,12 +555,12 @@ mod tests {
#[tokio::test]
async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() {
let server = RedisServer::start().await.unwrap();
let clients = connect(&coordination_config(&server), &|_| None)
let clients = connect(&coordination_config(&server), &|_: &str| None)
.await
.unwrap();
let coordination = clients.coordination.unwrap();
coordination
.set("written", "1", Duration::from_secs(60))
.send(commands::set("written", "1", Duration::from_secs(60)))
.await
.unwrap();
@ -537,6 +571,30 @@ mod tests {
assert!(buffer.ping().await.is_ok());
}
#[rstest]
#[tokio::test]
async fn an_environment_only_tls_redis_is_reached_with_its_ca() {
crate::install_crypto_provider();
let server = litellm_redis_testing::RedisTlsServer::start(false)
.await
.unwrap();
let ca = server.files().ca.display().to_string();
let port = server.port().to_string();
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some(server.host().to_owned()),
"REDIS_PORT" => Some(port.clone()),
"REDIS_SSL" => Some("true".to_owned()),
"REDIS_SSL_CA_CERTS" => Some(ca.clone()),
_ => None,
})
.await
.unwrap();
assert!(clients.coordination.is_some());
}
#[rstest]
#[tokio::test]
async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() {
@ -545,7 +603,7 @@ mod tests {
server.stop().await.unwrap();
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| match name {
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
"REDIS_PORT" => Some(port.clone()),
_ => None,
@ -564,7 +622,7 @@ mod tests {
let port = port_of(&server);
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| match name {
let clients = connect(&config, &|name: &str| match name {
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
"REDIS_PORT" => Some(port.clone()),
_ => None,

View file

@ -0,0 +1,61 @@
use std::collections::{BTreeMap, BTreeSet};
use litellm_core_utils::settings::Lookup;
use litellm_secret_value::SecretValue;
use litellm_config::Value;
use litellm_redis_resolve::SETTINGS;
use super::configured::ConfiguredRedis;
#[derive(Clone, Debug, Default)]
pub(crate) struct SecretSnapshot(BTreeMap<String, Option<SecretValue>>);
impl SecretSnapshot {
pub(crate) fn read(names: BTreeSet<String>, lookup: &dyn Lookup) -> Self {
Self(
names
.into_iter()
.map(|name| {
let value = lookup.get(&name).map(SecretValue::new);
(name, value)
})
.collect(),
)
}
}
impl Lookup for SecretSnapshot {
fn get(&self, name: &str) -> Option<String> {
match self.0.get(name) {
Some(value) => value.as_ref().map(|value| value.expose().to_owned()),
None => {
debug_assert!(
false,
"{name} is missing from ConfiguredRedis::secret_names"
);
tracing::error!(
name,
"a Redis secret was read without being in the snapshot"
);
None
}
}
}
}
impl ConfiguredRedis {
pub(crate) fn secret_names(&self) -> BTreeSet<String> {
let references = self.fields.values().filter_map(|value| match value {
Value::String(text) => text.strip_prefix("os.environ/").map(str::to_owned),
_ => None,
});
let declared = SETTINGS.iter().flat_map(|setting| {
setting
.environment_name()
.into_iter()
.chain(setting.fallback.iter().map(|name| (*name).to_owned()))
});
declared.chain(references).collect()
}
}

View file

@ -0,0 +1,964 @@
use std::collections::{BTreeMap, HashMap};
use litellm_config::{CacheSettings, Config, Value};
use litellm_redis_resolve::Expected;
use litellm_redis_types::{
Auth, AwsIamAuth, AzureAdAuth, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions,
};
use litellm_secret_value::SecretValue;
use rstest::rstest;
use serde_json::json;
use super::{configured::ConfiguredRedis, secrets::SecretSnapshot};
use crate::Error;
fn resolution(
configured: &ConfiguredRedis,
environment: &HashMap<String, String>,
) -> Result<Resolution, Error> {
let lookup = |name: &str| environment.get(name).cloned();
configured.resolve(
&SecretSnapshot::read(configured.secret_names(), &lookup),
&lookup,
)
}
const PROVIDER: ProviderHandle = ProviderHandle::new(7);
fn redis_cache_kwargs(kwargs: serde_json::Value) -> ConfiguredRedis {
let has_provider = kwargs
.get("credential_provider")
.is_some_and(|provider| !provider.is_null());
let configured = ConfiguredRedis::redis_cache_kwargs(
serde_json::from_value::<BTreeMap<String, Value>>(kwargs).unwrap(),
);
match has_provider {
true => configured.with_credential_provider(PROVIDER),
false => configured,
}
}
fn coordination_resolution(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
let environment: HashMap<String, String> = environment
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
let config = Config::from_yaml(yaml).unwrap();
let configured = config
.general_settings
.coordination_redis
.unwrap_or_default();
let configured = ConfiguredRedis::coordination(&configured)?;
resolution(&configured, &environment)
}
fn resolve(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
coordination_resolution(yaml, environment)
}
fn resolve_cache(yaml: &str) -> Result<Resolution, Error> {
resolve_cache_in(yaml, &[])
}
fn resolve_cache_in(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
let lookup = |name: &str| {
environment
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
};
let config = Config::from_yaml(yaml).unwrap();
let configured = match config.litellm_settings.cache_params.unwrap() {
CacheSettings::Redis(settings) => ConfiguredRedis::cache_params(&settings),
CacheSettings::Other(_) => panic!("expected a Redis cache"),
};
configured.resolve(
&SecretSnapshot::read(configured.secret_names(), &lookup),
&lookup,
)
}
fn cluster_ports(resolved: &Resolution) -> Vec<u16> {
match &resolved.target {
Target::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
target => panic!("expected a cluster, got {target:?}"),
}
}
#[rstest]
#[case::host(
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
&[],
"host:redis:6380:2:true:app:1500"
)]
#[case::environment_references_resolve_the_host_and_port(
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
"host:referenced:6400:0:false::none"
)]
#[case::a_null_host_falls_back_to_the_environment(
"general_settings:\n coordination_redis:\n host: null",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
fn resolves_redis_through_the_parsed_config(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolved = resolve(yaml, environment).unwrap();
let description = match resolved.target {
Target::Unix(socket) => format!("unix:{}", socket.path.display()),
Target::Url(url) => format!("url:{}", url.url.expose()),
Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name),
Target::Standalone(host) => format!(
"host:{}:{}:{}:{}:{}:{}",
host.host,
host.port,
host.db,
host.tls.is_some(),
resolved.namespace.unwrap_or_default(),
resolved
.timeout
.map(|timeout| timeout.as_millis().to_string())
.unwrap_or_else(|| "none".to_owned())
),
Target::Cluster(cluster) => format!(
"cluster:{}",
cluster
.nodes
.iter()
.map(|node| format!("{}:{}", node.host, node.port))
.collect::<Vec<_>>()
.join(",")
),
};
assert_eq!(description, expected);
}
#[rstest]
#[case::tls_ciphers(
"general_settings:\n coordination_redis:\n host: redis\n ssl_ciphers: HIGH",
&[],
"ssl_ciphers"
)]
#[case::ssl_ciphers_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_SSL_CIPHERS", "HIGH")],
"REDIS_SSL_CIPHERS"
)]
#[case::connection_pool(
"general_settings:\n coordination_redis:\n host: redis\n connection_pool: shared",
&[],
"connection_pool"
)]
fn unsupported_settings_resolve_and_are_named_by_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolution = coordination_resolution(yaml, environment).unwrap();
assert_eq!(
resolution.target,
Target::Standalone(litellm_redis_types::Standalone {
host: "redis".into(),
port: 6379,
db: 0,
tls: None,
})
);
let [ignored] = resolution.ignored.as_slice() else {
panic!("expected one ignored setting, got {:?}", resolution.ignored);
};
assert_eq!(ignored.setting, expected);
assert!(!ignored.consequence.is_empty());
}
#[rstest]
#[case::kwarg(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_validate_ocsp: true",
&[],
"ssl_validate_ocsp"
)]
#[case::stapled_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_SSL", "true"), ("REDIS_SSL_VALIDATE_OCSP_STAPLED", "yes")],
"REDIS_SSL_VALIDATE_OCSP_STAPLED"
)]
#[case::rediss_url_query(
"general_settings:\n coordination_redis:\n url: rediss://cache:6380/0?ssl_validate_ocsp=true",
&[],
"url ?ssl_validate_ocsp"
)]
fn an_ocsp_check_fails_instead_of_connecting_without_it(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] setting: &str,
) {
let error = coordination_resolution(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Unenforceable { setting: named, .. }) if named == setting),
"{error:?}"
);
}
#[rstest]
#[case::kwarg(" host: redis\n ssl: true\n ssl_validate_ocsp: false")]
#[case::rediss_url_query(" url: rediss://cache:6380/0?ssl_validate_ocsp=false")]
fn an_ocsp_flag_set_to_false_asks_for_nothing(#[case] block: &str) {
let yaml = format!("general_settings:\n coordination_redis:\n{block}");
let resolution = coordination_resolution(&yaml, &[]).unwrap();
assert!(resolution.ignored.is_empty(), "{:?}", resolution.ignored);
}
fn static_auth(username: Option<&str>, password: Option<&str>) -> Auth {
Auth::Static(StaticAuth {
username: username.map(litellm_secret_value::SecretValue::new),
password: password.map(litellm_secret_value::SecretValue::new),
})
}
#[rstest]
fn an_explicit_host_keeps_environment_credentials() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: configured",
&[
("REDIS_PASSWORD", "environment-secret"),
("REDIS_SSL", "True"),
],
)
.unwrap();
assert_eq!(resolved.auth, static_auth(None, Some("environment-secret")));
assert!(matches!(resolved.target, Target::Standalone(host) if host.tls.is_some()));
}
fn azure(username: Option<&str>) -> Auth {
Auth::AzureAd(AzureAdAuth {
client_id: Some("client".to_owned()),
tenant_id: None,
client_secret: None,
username: username.map(str::to_owned),
})
}
fn aws(serverless: bool) -> Auth {
Auth::AwsIam(AwsIamAuth {
user_name: "iam-user".to_owned(),
cache_name: "cache-1".to_owned(),
region: "us-east-1".to_owned(),
serverless,
})
}
#[rstest]
#[case::static_credentials(
"general_settings:\n coordination_redis:\n host: redis\n username: user\n password: pass",
&[],
static_auth(Some("user"), Some("pass"))
)]
#[case::gcp_beats_azure_and_aws(
"general_settings:\n coordination_redis:\n host: redis\n password: pass\n azure_redis_ad_token: \"True\"\n aws_iam_auth: true",
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com")],
Auth::GcpIam { service_account: "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com".to_owned() }
)]
#[case::azure_beats_aws_and_takes_the_environment_username(
"general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client\n aws_iam_auth: true",
&[("REDIS_USERNAME", "environment-user")],
azure(Some("environment-user"))
)]
#[case::azure_without_an_environment_username_sends_only_the_token(
"general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client",
&[],
azure(None)
)]
#[case::azure_needs_the_literal_true(
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: \"1\"",
&[],
static_auth(None, None)
)]
#[case::aws_with_a_region_from_the_environment(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n password: dropped\n aws_iam_auth: \"yes\"\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1",
&[("AWS_DEFAULT_REGION", "us-east-1")],
aws(false)
)]
#[case::aws_serverless_over_a_rediss_url(
"general_settings:\n coordination_redis:\n url: rediss://cache-1:6379\n aws_iam_auth: true\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1\n aws_iam_region: us-east-1\n aws_iam_serverless: true",
&[],
aws(true)
)]
fn auth_follows_python_precedence(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: Auth,
) {
assert_eq!(resolve(yaml, environment).unwrap().auth, expected);
}
#[rstest]
#[case::coordination(
"general_settings:\n coordination_redis:\n host: redis\n credential_provider: host-object",
&[]
)]
#[case::without_a_target(
"general_settings:\n coordination_redis:\n credential_provider: host-object",
&[("REDIS_HOST", "redis")]
)]
fn a_credential_provider_without_a_host_handle_fails(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(
error,
Error::Redis(litellm_redis_resolve::Error::ProviderWithoutHandle)
),
"{error}"
);
}
#[rstest]
#[case::without_tls(
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r",
"requires TLS"
)]
#[case::a_plain_url_ignores_ssl(
"general_settings:\n coordination_redis:\n url: redis://cache:6379\n ssl: true\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r",
"requires TLS"
)]
#[case::missing_settings(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n aws_iam_auth: true\n aws_iam_cache_name: c",
"requires: aws_iam_user_name, aws_iam_region"
)]
fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) {
let error = resolve(yaml, &[]).unwrap_err();
assert!(error.to_string().contains(message), "{error}");
}
#[rstest]
#[case::defaults_verify_the_certificate_but_not_the_hostname(
" ssl: true",
&[],
Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false })
)]
#[case::every_option(
" ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_password: key-secret\n ssl_cert_reqs: none\n ssl_check_hostname: true",
&[],
Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), key_password: Some(SecretValue::new("key-secret")), verify_certificate: false, check_hostname: true })
)]
#[case::options_from_the_environment(
" ssl: true",
&[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional"), ("REDIS_SSL_PASSWORD", "env-key-secret")],
Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: Some(SecretValue::new("env-key-secret")), verify_certificate: true, check_hostname: false })
)]
#[case::gcp_ca_certs_replace_the_configured_ones(
" ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem",
&[],
Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false })
)]
#[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)]
fn tls_options_resolve_like_redis_py(
#[case] block: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: Option<TlsOptions>,
) {
let yaml = format!("general_settings:\n coordination_redis:\n host: redis\n{block}");
let resolved = resolve(&yaml, environment).unwrap();
assert!(matches!(resolved.target, Target::Standalone(host) if host.tls == expected));
}
#[rstest]
fn an_unknown_cert_requirement_fails() {
let error = resolve(
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_cert_reqs: strict",
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::CertRequirement }) if setting == "ssl_cert_reqs"
));
}
#[rstest]
#[case::from_config(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[sentinel-a, 26379], [sentinel-b, \"26380\"]]\n service_name: primary\n sentinel_password: sentinel-secret\n db: 2",
&[]
)]
#[case::from_the_environment(
"model_list: []",
&[
("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379], ["sentinel-b", "26380"]]"#),
("REDIS_SERVICE_NAME", "primary"),
("REDIS_SENTINEL_PASSWORD", "sentinel-secret"),
("REDIS_DB", "2"),
]
)]
fn sentinel_nodes_resolve_to_a_sentinel_endpoint(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
) {
let Target::Sentinel(sentinel) = resolve(yaml, environment).unwrap().target else {
panic!("expected a sentinel endpoint");
};
assert_eq!(
sentinel.monitors,
vec![
Node {
host: "sentinel-a".to_owned(),
port: 26379
},
Node {
host: "sentinel-b".to_owned(),
port: 26380
},
]
);
assert_eq!(sentinel.service_name, "primary");
assert_eq!(
sentinel.monitor_auth,
Some(StaticAuth {
username: None,
password: Some(litellm_secret_value::SecretValue::new("sentinel-secret")),
})
);
assert_eq!(sentinel.db, 2);
}
#[rstest]
#[case::a_url_beats_sentinels(
"general_settings:\n coordination_redis:\n url: redis://single:6379\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
"url:redis://single:6379"
)]
fn sentinel_precedence(#[case] yaml: &str, #[case] expected: &str) {
let resolved = resolve(yaml, &[]).unwrap();
let description = match resolved.target {
Target::Url(url) => format!("url:{}", url.url.expose()),
Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name),
other => panic!("unexpected {other:?}"),
};
assert_eq!(description, expected);
}
#[rstest]
fn sentinels_without_a_service_name_fail() {
let error = resolve(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]",
&[],
)
.unwrap_err();
assert!(
matches!(
error,
Error::Redis(litellm_redis_resolve::Error::NoSentinelService)
),
"{error:?}"
);
}
#[rstest]
#[case::configured_host(
"general_settings:\n coordination_redis:\n host: ignored\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
&[],
"host"
)]
#[case::environment_port(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
&[("REDIS_PORT", "6380")],
"REDIS_PORT"
)]
fn a_host_or_port_cannot_reach_sentinel(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::SentinelWithHost { setting }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::not_pairs(
"general_settings:\n coordination_redis:\n sentinel_nodes: [{host: s, port: 26379}]\n service_name: primary",
&[],
"sentinel_nodes"
)]
#[case::from_the_environment(
"model_list: []",
&[("REDIS_SENTINEL_NODES", "not-json"), ("REDIS_SERVICE_NAME", "primary")],
"REDIS_SENTINEL_NODES"
)]
fn invalid_sentinel_nodes_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::SentinelNodes }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::from_config("general_settings:\n coordination_redis:\n startup_nodes: [{host: a}]", &[], "startup_nodes")]
#[case::from_the_environment("model_list: []", &[("REDIS_CLUSTER_NODES", "not-json")], "REDIS_CLUSTER_NODES")]
#[case::out_of_range_port_from_the_environment("model_list: []", &[("REDIS_HOST", "redis"), ("REDIS_PORT", "70000")], "REDIS_PORT")]
fn invalid_values_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == source),
"{error:?}"
);
}
#[rstest]
#[case::falls_back_to_the_socket_timeout(" socket_timeout: 1.5", Some(1500), None)]
#[case::explicit_connect_timeout(
" socket_timeout: 1.5\n socket_connect_timeout: 0.25",
Some(1500),
Some(250)
)]
fn connect_timeout_is_resolved_separately(
#[case] block: &str,
#[case] timeout_ms: Option<u128>,
#[case] connect_ms: Option<u128>,
) {
let resolved = resolve(
&format!("general_settings:\n coordination_redis:\n host: redis\n{block}"),
&[],
)
.unwrap();
assert_eq!(
resolved.timeout.map(|timeout| timeout.as_millis()),
timeout_ms
);
assert_eq!(
resolved.connect_timeout.map(|timeout| timeout.as_millis()),
connect_ms
);
}
#[rstest]
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
#[case::a_structured_host_with_an_environment_target(
"general_settings:\n coordination_redis:\n host: []",
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
)]
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
let error = resolve(yaml, environment).unwrap_err();
assert!(
matches!(error, Error::Redis(litellm_redis_resolve::Error::NoTarget)),
"{error:?}"
);
}
#[rstest]
fn the_cache_alias_names_cluster_nodes() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn canonical_startup_nodes_beat_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
fn unsupported_environment_settings_never_carry_the_value() {
let resolution = coordination_resolution(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CIPHERS", "private-option-value"),
],
)
.unwrap();
let ignored = format!("{:?}", resolution.ignored);
assert!(ignored.contains("REDIS_SSL_CIPHERS"), "{ignored}");
assert!(!ignored.contains("private-option-value"), "{ignored}");
}
#[rstest]
#[case::negative("-1")]
#[case::infinite("inf")]
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
let yaml = format!(
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
);
let error = resolve(&yaml, &[]).unwrap_err();
assert!(matches!(
error,
Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::Seconds }) if setting == "socket_timeout"
));
}
#[rstest]
fn debug_redacts_redis_secrets_and_url() {
let config = Config::from_yaml(
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password\n sentinel_kwargs: {password: credential-passthrough}",
)
.unwrap();
let settings = config.general_settings.coordination_redis.unwrap();
let configured = ConfiguredRedis::coordination(&settings).unwrap();
let resolved = resolution(&configured, &HashMap::new()).unwrap();
let debug = format!("{settings:?} {configured:?} {resolved:?}");
assert!(!debug.contains("redis://user:password"));
assert!(!debug.contains("credential-user"));
assert!(!debug.contains("credential-password"));
assert!(!debug.contains("credential-passthrough"));
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
}
#[rstest]
fn legacy_unknown_settings_do_not_change_resolution() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
&[],
)
.unwrap();
assert!(matches!(resolved.target, Target::Standalone(_)));
}
#[rstest]
#[case::configured(
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
&[],
"max_connections"
)]
#[case::from_the_environment(
"general_settings:\n coordination_redis:\n host: redis",
&[("REDIS_MAX_CONNECTIONS", "47")],
"REDIS_MAX_CONNECTIONS"
)]
fn ignored_settings_are_named_by_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] source: &str,
) {
let resolution = coordination_resolution(yaml, environment).unwrap();
let settings: Vec<&str> = resolution
.ignored
.iter()
.map(|ignored| ignored.setting.as_str())
.collect();
assert_eq!(settings, vec![source]);
}
fn snapshot(values: &[(&str, &str)]) -> SecretSnapshot {
let values: HashMap<&str, &str> = values.iter().copied().collect();
SecretSnapshot::read(
ConfiguredRedis::environment_only().secret_names(),
&|name: &str| values.get(name).map(|value| (*value).to_owned()),
)
}
#[rstest]
fn secret_manager_values_reach_the_resolver_through_the_snapshot() {
let settings = redis_cache_kwargs(json!({"host": "os.environ/COORD_HOST"}));
let secrets = SecretSnapshot::read(settings.secret_names(), &|name: &str| match name {
"COORD_HOST" => Some("from-the-manager".to_owned()),
"REDIS_PASSWORD" => Some("manager-secret".to_owned()),
_ => None,
});
let resolved = settings.resolve(&secrets, &|_: &str| None).unwrap();
assert!(
matches!(&resolved.target, Target::Standalone(host) if host.host == "from-the-manager")
);
assert_eq!(resolved.auth, static_auth(None, Some("manager-secret")));
}
#[rstest]
#[case::from_the_environment(&[], &[("REDIS_USERNAME", "environment-user")], Some("environment-user"))]
#[case::never_from_the_secret_manager(&[("REDIS_USERNAME", "manager-user")], &[], None)]
fn the_azure_username_comes_from_the_process_environment_like_python(
#[case] secrets: &[(&str, &str)],
#[case] environment: &[(&str, &str)],
#[case] expected: Option<&str>,
) {
let settings = redis_cache_kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"}));
let environment: HashMap<&str, &str> = environment.iter().copied().collect();
let resolved = settings
.resolve(&snapshot(secrets), &|name: &str| {
environment.get(name).map(|value| (*value).to_owned())
})
.unwrap();
assert!(
matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected),
"{:?}",
resolved.auth
);
}
#[rstest]
#[case::coordination_drops_the_null(
"general_settings:\n coordination_redis:\n host: redis\n ssl: null",
false,
true
)]
#[case::cache_passes_an_unnamed_null_on(
"litellm_settings:\n cache_params:\n type: redis\n host: redis\n ssl: null",
true,
false
)]
#[case::cache_drops_a_null_it_names(
"litellm_settings:\n cache_params:\n type: redis\n host: null\n ssl: true",
true,
true
)]
fn a_null_overrides_the_environment_only_where_python_passes_it_on(
#[case] yaml: &str,
#[case] cache: bool,
#[case] tls: bool,
) {
let environment = [("REDIS_HOST", "environment"), ("REDIS_SSL", "true")];
let resolved = match cache {
true => resolve_cache_in(yaml, &environment),
false => resolve(yaml, &environment),
}
.unwrap();
let Target::Standalone(host) = resolved.target else {
panic!("expected a standalone target, got {:?}", resolved.target);
};
assert_eq!(host.tls.is_some(), tls);
}
#[rstest]
#[case::empty_cluster_nodes_use_the_environment(
json!({"startup_nodes": []}),
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)],
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}),
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
)]
#[case::empty_sentinel_nodes_use_the_environment(
json!({"sentinel_nodes": [], "service_name": "primary"}),
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)],
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}),
&[("sentinel_nodes", "REDIS_SENTINEL_NODES")],
)]
#[case::empty_service_name_uses_the_environment(
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}),
&[("REDIS_SERVICE_NAME", "primary")],
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}),
&[("service_name", "REDIS_SERVICE_NAME")],
)]
#[case::environment_host_port_and_password(
json!({}),
&[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")],
json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}),
&[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")],
)]
#[case::padded_environment_numbers_parse_like_python_int(
json!({}),
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "6_380"), ("REDIS_DB", "3\n")],
json!({"port": "6_380", "host": "redis", "db": "3\n"}),
&[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
)]
#[case::socket_timeout_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SOCKET_TIMEOUT", "0.75")],
json!({"socket_timeout": "0.75", "host": "redis"}),
&[("socket_timeout", "REDIS_SOCKET_TIMEOUT")],
)]
#[case::environment_url(
json!({}),
&[("REDIS_URL", "redis://environment:6379/1")],
json!({"url": "redis://environment:6379/1"}),
&[("url", "REDIS_URL")],
)]
#[case::environment_url_beats_the_environment_host(
json!({}),
&[("REDIS_HOST", "ignored"), ("REDIS_URL", "redis://from-url:6379")],
json!({"host": "ignored", "url": "redis://from-url:6379"}),
&[("host", "REDIS_HOST"), ("url", "REDIS_URL")],
)]
#[case::an_explicit_host_drops_the_environment_url(
json!({"host": "configured"}),
&[("REDIS_URL", "redis://environment:6379")],
json!({"host": "configured"}),
&[],
)]
#[case::environment_cluster_nodes(
json!({}),
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)],
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}),
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
)]
#[case::sentinel_from_the_environment(
json!({}),
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")],
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}),
&[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")],
)]
#[case::an_explicit_credential_provider_beats_cloud_auth(
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
&[("REDIS_AZURE_AD_TOKEN", "true")],
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
&[],
)]
#[case::gcp_beats_azure_and_aws(
json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")],
json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
&[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")],
)]
#[case::azure_beats_aws_and_takes_the_environment_username(
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
&[("REDIS_USERNAME", "environment-user")],
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
&[],
)]
#[case::an_empty_environment_username_sends_only_the_token(
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}),
&[("REDIS_USERNAME", "")],
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}),
&[("username", "REDIS_USERNAME")],
)]
#[case::aws_iam_over_tls_with_the_default_region(
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}),
&[("AWS_DEFAULT_REGION", "us-east-1")],
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}),
&[("aws_iam_region", "AWS_DEFAULT_REGION")],
)]
#[case::aws_region_beats_the_default_region(
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}),
&[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")],
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}),
&[("aws_iam_region", "AWS_REGION")],
)]
#[case::a_null_cert_requirement_skips_verification(
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
&[("REDIS_SSL_CERT_REQS", "required")],
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
&[],
)]
#[case::a_cert_requirement_spelled_like_python_none_fails(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CERT_REQS", "None")],
json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
)]
#[case::tls_options_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")],
json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
)]
#[case::environment_references_resolve(
json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}),
&[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")],
json!({"host": "referenced", "password": "referenced-secret"}),
&[],
)]
#[case::malformed_port_from_the_environment(
json!({}),
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")],
json!({"port": "abc", "host": "redis"}),
&[("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
)]
#[case::retry_on_error_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_RETRY_ON_ERROR", "true")],
json!({"retry_on_error": true, "host": "redis"}),
&[("retry_on_error", "REDIS_RETRY_ON_ERROR")],
)]
#[case::ssl_ca_path_from_the_environment(
json!({"host": "redis"}),
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")],
json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}),
&[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")],
)]
fn merges_like_python_redis(
#[case] kwargs: serde_json::Value,
#[case] environment: &[(&str, &str)],
#[case] merged: serde_json::Value,
#[case] sources: &[(&str, &str)],
) {
let configured = redis_cache_kwargs(kwargs);
let lookup = |name: &str| {
environment
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
};
let kwargs = configured
.merged(&SecretSnapshot::read(configured.secret_names(), &lookup))
.unwrap();
let values: serde_json::Map<String, serde_json::Value> = kwargs
.iter()
.map(|(kwarg, entry)| (kwarg.clone(), entry.value.clone()))
.collect();
let renamed: Vec<(&str, &str)> = kwargs
.iter()
.filter(|(kwarg, entry)| *kwarg != &entry.source)
.map(|(kwarg, entry)| (kwarg.as_str(), entry.source.as_str()))
.collect();
assert_eq!(serde_json::Value::Object(values), merged);
assert_eq!(renamed, sources);
}

View file

@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned {
) -> PyResult<Py<PyAny>>;
/// What the stream carries at hand-off, as the caller's stream receives it.
/// The unary response after the deployment transform, before success hooks run.
fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> {
Ok(())
}
fn encode_stream_head(
&mut self,
py: Python<'_>,

View file

@ -399,8 +399,16 @@ where
Ok(step) => step,
Err(error) => return self.failure(error, FailureOrigin::Host).await,
};
match self.settle(step).await {
Ok(response) => self.succeeded(response).await,
let response = match self.settle(step).await {
Ok(response) => response,
Err(error) => return self.failure(error, FailureOrigin::Host).await,
};
let delivered = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.deliver_response(py, response.bind(py)))
};
match delivered {
Ok(()) => self.succeeded(response).await,
Err(error) => self.failure(error, FailureOrigin::Host).await,
}
}

View file

@ -37,7 +37,7 @@ pub use owned::PythonOwned;
pub use runtime::{
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value,
ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync,
run_sync_value, runtime_started,
run_sync_value, runtime_handle, runtime_started,
};
pub use services::PythonHostCalls;

View file

@ -68,6 +68,11 @@ fn runtime() -> PyResult<&'static Runtime> {
Ok(pyo3_async_runtimes::tokio::get_runtime())
}
/// The native runtime, for work that must outlive the call that schedules it.
pub fn runtime_handle() -> PyResult<Handle> {
Ok(runtime()?.handle().clone())
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
where

View file

@ -16,6 +16,15 @@ pub enum TlsSource {
#[error(transparent)]
pub struct Error(Arc<Repr>);
// A streamed body the connection ended early, after the response head arrived; `expected` is
// the declared `content-length`, `None` for a chunked body.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BodyCut {
pub received: u64,
pub expected: Option<u64>,
pub reset: bool,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum Repr {
#[error("could not read {}", path.display())]
@ -56,6 +65,12 @@ pub(crate) enum Repr {
Connect(#[source] reqwest::Error),
#[error("upstream network error")]
Network(#[source] reqwest::Error),
#[error("upstream network error")]
BodyCut {
#[source]
source: reqwest::Error,
cut: BodyCut,
},
#[error("upstream request timed out")]
Timeout(#[source] reqwest::Error),
/// No bytes arrived within the read timeout, which the caller enforces per read; `body` says
@ -201,6 +216,37 @@ impl Error {
matches!(&*self.0, Repr::ReadTimeout { body: true, .. })
}
pub fn body_read(error: reqwest::Error, received: u64, expected: Option<u64>) -> Self {
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(&error);
let kind = std::iter::from_fn(|| {
let current = source?;
source = current.source();
Some(current)
})
.find_map(|cause| cause.downcast_ref::<io::Error>().map(io::Error::kind));
let reset = match kind {
Some(io::ErrorKind::UnexpectedEof) => false,
Some(io::ErrorKind::ConnectionReset) => true,
_ => return error.into(),
};
Repr::BodyCut {
source: error.without_url(),
cut: BodyCut {
received,
expected,
reset,
},
}
.into()
}
pub fn body_cut(&self) -> Option<BodyCut> {
match &*self.0 {
Repr::BodyCut { cut, .. } => Some(*cut),
_ => None,
}
}
pub fn computed_header(name: impl Into<String>) -> Self {
Repr::ComputedHeader(name.into()).into()
}
@ -226,6 +272,7 @@ impl Error {
pub fn dispatched(&self) -> bool {
match &*self.0 {
Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -285,6 +332,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -318,6 +366,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -355,6 +404,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -394,6 +444,7 @@ impl Error {
| Repr::UnsignableHeader(_)
| Repr::Signature(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -437,6 +488,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Status { .. }
| Repr::WebSocketUrl(_)
| Repr::WebSocketScheme(_)
@ -473,6 +525,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::WebSocketUrl(_)
@ -513,6 +566,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::WebSocketUrl(_)
| Repr::WebSocketScheme(_)
| Repr::WebSocketSchemeRewrite(_)
@ -548,6 +602,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -742,4 +797,53 @@ mod tests {
assert_eq!(error.upstream_response(), Some((429, body.as_str())));
assert_eq!(error.status(), Some(429));
}
#[rstest]
#[case::chunked("transfer-encoding: chunked\r\n\r\n5\r\nhello\r\n", None)]
#[case::content_length("content-length: 100\r\n\r\nhello", Some(100))]
#[tokio::test]
async fn a_connection_that_ends_mid_body_is_a_body_cut(
#[case] tail: &'static str,
#[case] expected: Option<u64>,
) {
use std::io::{Read as _, Write as _};
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
let address = listener.local_addr().expect("address");
let server = std::thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("accept");
let mut request = [0_u8; 1024];
let _ = stream.read(&mut request).expect("request");
stream
.write_all(format!("HTTP/1.1 200 OK\r\n{tail}").as_bytes())
.expect("response");
});
let mut response = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}/"))
.send()
.await
.expect("head");
server.join().expect("server");
let mut received = 0_u64;
let error = loop {
match response.chunk().await {
Ok(Some(chunk)) => received += chunk.len() as u64,
Ok(None) => panic!("the body must end with an error"),
Err(error) => break crate::Error::body_read(error, received, expected),
}
};
assert_eq!(
error.body_cut(),
Some(super::BodyCut {
received: 5,
expected,
reset: false,
})
);
assert!(error.dispatched());
assert_eq!(error.to_string(), "upstream network error");
}
}

View file

@ -18,7 +18,7 @@ pub mod websocket;
pub use client::Client;
pub use config::{ClientIdentity, HttpClientConfig, Resolution, Verify};
pub use error::{Error, TlsSource};
pub use error::{BodyCut, Error, TlsSource};
pub use pool::{ClientVariant, HttpClientPool};
pub use proxy::EnvironmentProxies;
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};

View file

@ -1,5 +1,5 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_http::request::string_headers as shared_string_headers;
use litellm_llms::ErrorDetail;
use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG;
use litellm_llms::{
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
@ -19,8 +19,6 @@ use serde_json::{Map, Value};
use super::Error;
const HEADER_CONTEXT: &str = "messages";
#[derive(Clone, Copy)]
pub(crate) struct MessagesProvider {
provider: LlmProviders,
@ -75,7 +73,52 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option<M
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(name, value)| match value {
Value::String(value) => Ok((name, value)),
value => Err(Error::invalid_request(ErrorDetail::invalid(
"extra_headers",
NonStringHeader {
actual: python_type_name(&value),
name,
},
))),
})
.collect()
}
/// A header value the caller set to something other than a string, named by its Python type
/// since the caller wrote it in Python.
#[derive(Debug)]
struct NonStringHeader {
name: String,
actual: &'static str,
}
impl std::fmt::Display for NonStringHeader {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"header {} must be a string, got {}",
self.name, self.actual
)
}
}
impl std::error::Error for NonStringHeader {}
fn python_type_name(value: &Value) -> &'static str {
match value {
Value::Null => "NoneType",
Value::Bool(_) => "bool",
Value::Number(number) if number.is_f64() => "float",
Value::Number(_) => "int",
Value::String(_) => "str",
Value::Array(_) => "list",
Value::Object(_) => "dict",
}
}
#[cfg(test)]
@ -85,7 +128,6 @@ mod tests {
use rstest::rstest;
use super::{messages_provider, string_headers};
use crate::Error;
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
@ -113,17 +155,28 @@ mod tests {
assert!(messages_provider(provider, model).is_none());
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
#[rstest]
#[case::int(json!(12345), "int")]
#[case::float(json!(1.5), "float")]
#[case::bool(json!(true), "bool")]
#[case::null(json!(null), "NoneType")]
#[case::list(json!([1]), "list")]
#[case::dict(json!({"a": 1}), "dict")]
fn non_string_header_is_a_bad_request_naming_the_header_and_its_type(
#[case] value: serde_json::Value,
#[case] python_type: &str,
) {
let headers = json!({"x-kept": "v", "x-count": value})
.as_object()
.unwrap()
.clone();
let error = string_headers(Some(headers)).expect_err("non-string header refused");
assert!(error.is_request(), "{error:?}");
assert_eq!(error.upstream_response(), None);
let rejection = error.rejected_request().expect("a rejected request");
assert_eq!(
err,
Error::Headers(litellm_http::request::HeaderError {
context: "messages",
name: "x-count".to_string(),
actual: "number",
})
litellm_error::Report::new(rejection).to_string(),
format!("invalid extra_headers: header x-count must be a string, got {python_type}")
);
}
}

View file

@ -26,10 +26,10 @@ use litellm_llms_types::formats::messages::{
};
use litellm_secrets::source::Secrets;
use litellm_tracing::ByteChunk;
use serde_json::Value;
use serde_json::{Map, Value};
use super::{
Error, MessagesCall, MessagesCallResponse, MessagesRoute,
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
common_utils::{MessagesProvider, merge_messages_header_sources, string_headers},
constants::MESSAGES_TIMEOUT_SECS,
route::Messages,
@ -85,7 +85,10 @@ pub(super) async fn execute(
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => provider.config().reported_cost(response),
MessagesCallResponse::Complete(response) => match response {
MessagesBody::Message(message) => provider.config().reported_cost(message),
MessagesBody::Unchecked(_) => None,
},
MessagesCallResponse::Stream { .. } => None,
};
context
@ -131,6 +134,7 @@ async fn provider_call(
let transform_context =
MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params);
let filtered = requested_optional_params(sanitized, provider, &transform_context)?;
let requested_params = serde_json::to_value(&filtered.params).map_err(serialize_failure)?;
let trimmed =
without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?;
@ -173,7 +177,7 @@ async fn provider_call(
let request_context = RequestContext {
model: body.model.clone(),
custom_llm_provider: provider.as_str().to_string(),
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
optional_params: requested_params,
secret_fields: Vec::new(),
api_key: api_key.map(SecretValue::new),
};
@ -346,8 +350,7 @@ async fn call_provider(
// body it cannot parse, and the public mapping classifies it by that status.
return Err(litellm_http::Error::upstream(status, text).into());
}
decode_response(config, &identity.model, &text)
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete)
}
fn serialize_failure(err: serde_json::Error) -> Error {
@ -383,13 +386,20 @@ fn timed_chunks(
response: reqwest::Response,
timeout: Duration,
) -> impl futures_util::Stream<Item = Result<Bytes, litellm_http::Error>> + Send + 'static {
futures_util::stream::try_unfold(response, move |mut response| async move {
let chunk = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)?
.map_err(litellm_http::Error::from)?;
Ok(chunk.map(|chunk| (chunk, response)))
})
let expected = response.content_length();
futures_util::stream::try_unfold(
(response, 0_u64),
move |(mut response, received)| async move {
let chunk = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)?
.map_err(|error| litellm_http::Error::body_read(error, received, expected))?;
Ok(chunk.map(|chunk| {
let received = received + chunk.len() as u64;
(chunk, (response, received))
}))
},
)
}
/// The provider's error, its body read under `timeout`: a stream has no overall deadline once
@ -411,15 +421,23 @@ fn decode_response(
config: &dyn BaseMessagesConfig,
model: &str,
text: &str,
) -> Result<MessagesResponse, Error> {
let response = serde_json::from_str(text).map_err(|err| {
Error::invalid_response(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
err,
))
})?;
) -> Result<MessagesBody, Error> {
let response = match serde_json::from_str::<MessagesResponse>(text) {
Ok(response) => response,
Err(err) => {
let Ok(mut body) = serde_json::from_str::<Map<String, Value>>(text) else {
return Err(Error::invalid_response(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
err,
)));
};
body.entry("stop_sequence").or_insert(Value::Null);
return Ok(MessagesBody::Unchecked(body));
}
};
config
.transform_anthropic_messages_response(model, response)
.map(|message| MessagesBody::Message(Box::new(message)))
.map_err(Error::from)
.map_err(Error::after_send)
}

View file

@ -17,8 +17,8 @@ use crate::common_utils::messages_provider;
pub use common_utils::merge_messages_header_sources;
pub use litellm_inference::Error;
pub use types::{
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
messages_body,
MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping,
litellm_params, messages_body,
};
#[derive(Clone)]

View file

@ -5,11 +5,10 @@ use litellm_host::{
call::{HostedCompletion, HostedMachine, hosted_call},
protocol::Protocol,
};
use litellm_llms_types::formats::messages::MessagesResponse;
use super::{Error, MessagesCall};
use super::{Error, MessagesBody, MessagesCall};
pub type MessagesOutput = HostedCompletion<Box<MessagesResponse>>;
pub type MessagesOutput = HostedCompletion<MessagesBody>;
/// The upstream response as the caller sees it at stream hand-off, before any chunk.
pub struct MessagesStreamHead {
@ -19,7 +18,7 @@ pub struct MessagesStreamHead {
pub struct Messages;
impl Protocol for Messages {
type Response = Box<MessagesResponse>;
type Response = MessagesBody;
type Error = Error;
type Request = MessagesCall;
type HostCall = Infallible;

View file

@ -42,7 +42,16 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error {
}
pub type MessagesCallResponse =
CallOutput<Box<MessagesResponse>, super::route::MessagesStreamHead, Bytes, Error>;
CallOutput<MessagesBody, super::route::MessagesStreamHead, Bytes, Error>;
// Python's `transform_anthropic_messages_response` returns any JSON object the provider sends
// with a 200, message or not, so one that is not a message passes through as sent.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessagesBody {
Message(Box<MessagesResponse>),
Unchecked(Map<String, Value>),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct MessagesShaping {
@ -144,6 +153,7 @@ mod tests {
supports_reasoning: true,
supports_adaptive_thinking: true,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: true,
supports_sampling_params: false,

View file

@ -1,7 +1,9 @@
use std::{sync::Arc, time::Duration};
use futures_util::TryStreamExt;
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping};
use litellm_inference_messages::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesShaping,
};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use rstest::{fixture, rstest};
@ -74,7 +76,10 @@ async fn run(call: MessagesCall) -> Result<MessagesCallResponse, Error> {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.unwrap() {
MessagesCallResponse::Complete(message) => *message,
MessagesCallResponse::Complete(message) => match message {
MessagesBody::Message(message) => *message,
MessagesBody::Unchecked(_) => panic!("expected a message"),
},
MessagesCallResponse::Stream { .. } => panic!("expected a completed message"),
}
}

View file

@ -50,42 +50,47 @@ async fn messages_cache_identity_includes_provider_native_parameters(
#[case] original: Value,
#[case] changed: Value,
) {
use litellm_inference_messages::route::Messages;
use litellm_inference_messages::{MessagesBody, route::Messages};
use litellm_llms_types::formats::messages::MessagesResponse;
let calls = AtomicUsize::new(0);
for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] {
let response =
execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
Some(&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
let call = calls.fetch_add(1, Ordering::SeqCst);
Ok(Box::new(serde_json::from_value::<MessagesResponse>(json!({
"id":call.to_string(), "type":"message", "role":"assistant", "model":"test",
"content":[{"type":"text","text":format!("answer {call}")}],
"stop_reason":"end_turn", "stop_sequence":null
})).unwrap()))
let response = execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
)
.await
.unwrap();
Some(&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
let call = calls.fetch_add(1, Ordering::SeqCst);
Ok(MessagesBody::Message(Box::new(
serde_json::from_value::<MessagesResponse>(json!({
"id":call.to_string(), "type":"message", "role":"assistant", "model":"test",
"content":[{"type":"text","text":format!("answer {call}")}],
"stop_reason":"end_turn", "stop_sequence":null
}))
.unwrap(),
)))
},
)
.await
.unwrap();
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(response.id, expected_call.to_string());
assert_eq!(
response.content[0]["text"],
@ -253,7 +258,7 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
#[case] cost: Value,
#[case] expected: Option<f64>,
) {
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
use litellm_inference_messages::{MessagesBody, MessagesCall, MessagesCallResponse};
let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test",
"content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null,
"usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost});
@ -277,6 +282,9 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
panic!("expected Messages responses");
};
assert_eq!(first, second);
let MessagesBody::Message(first) = first else {
panic!("expected a message");
};
assert_eq!(first.extra.get("cost"), Some(&cost));
assert_eq!(support::received(&upstream).await.len(), 1);
let facts = hooks.facts.lock().unwrap();

View file

@ -4,7 +4,9 @@ use futures_util::{FutureExt, TryStreamExt};
use litellm_host::interceptors::{
ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest,
};
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute};
use litellm_inference_messages::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
};
use litellm_inference_testing::live::{LiveResources, model, within_deadline};
use litellm_llms_types::formats::messages::MessagesResponse;
use serde_json::{Value, json};
@ -111,7 +113,10 @@ pub(super) fn call(provider: &str, payload: Value) -> MessagesCall {
pub(super) fn complete(response: MessagesCallResponse) -> Box<MessagesResponse> {
match response {
MessagesCallResponse::Complete(message) => message,
MessagesCallResponse::Complete(message) => match message {
MessagesBody::Message(message) => message,
MessagesBody::Unchecked(_) => panic!("a complete request returned no message"),
},
MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"),
}
}

View file

@ -6,9 +6,11 @@ use litellm_host::{
interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest},
lifecycle::CallEvent,
};
use litellm_inference_messages::{MessagesCallResponse, route::Messages};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities;
use litellm_llms::base_llm::messages::context::{
MessagesModelCapabilities as AnthropicModelCapabilities, SupportedEffortTiers,
};
use rstest::rstest;
use super::*;
@ -338,7 +340,10 @@ async fn response_mode_follows_the_intercepted_request(
}
MessagesCallResponse::Complete(message) => {
assert_eq!(expected_stream, Some(false));
assert_eq!(*message, serde_json::from_value(message_body()).unwrap());
assert_eq!(
message,
MessagesBody::Message(serde_json::from_value(message_body()).unwrap())
);
}
}
Ok::<_, Error>(())
@ -419,7 +424,7 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure(
assert!(host.raw_responses().is_empty());
}
/// Python logs `optional_params` as what it is about to send, so a dropped param must
/// Python logs `optional_params` after its requested-param gate, so a dropped param must
/// not resurface in callbacks.
#[rstest]
#[tokio::test]
@ -450,3 +455,54 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(optional_params, json!({"max_tokens": 16}));
}
#[rstest]
#[case::reasoning_effort_on_an_adaptive_model(
AnthropicModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
supports_output_config: true,
effort_tiers: SupportedEffortTiers { low: true, ..SupportedEffortTiers::default() },
..AnthropicModelCapabilities::default()
},
json!({"reasoning_effort": "low"})
)]
#[case::adaptive_payload_on_a_legacy_model(
AnthropicModelCapabilities {
supports_reasoning: true,
..AnthropicModelCapabilities::default()
},
json!({"max_tokens": 3000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}})
)]
#[tokio::test]
async fn the_request_context_carries_the_requested_reasoning_not_its_translation(
call: MessagesCall,
#[case] capabilities: AnthropicModelCapabilities,
#[case] fields: Value,
) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
capabilities,
..MessagesShaping::default()
},
..with_fields(call, fields.clone())
},
upstream.uri(),
));
run_through(&host).await.expect("messages call succeeds");
let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap())
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(
optional_params,
Value::Object(
object(json!({"max_tokens": 16}))
.into_iter()
.chain(object(fields))
.collect()
)
);
}

View file

@ -5,7 +5,7 @@ use std::{
use litellm_http::{HttpSettings, Resolution};
use litellm_inference_messages::{
Error, MessagesCall, MessagesSettings, MessagesShaping,
Error, MessagesBody, MessagesCall, MessagesSettings, MessagesShaping,
route::{Messages, MessagesMachine, MessagesOutput},
};
use litellm_inference_testing::RecordingSecrets;
@ -119,7 +119,10 @@ async fn run(call: MessagesCall) -> Result<MessagesOutput, Error> {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.expect("messages call succeeds") {
MessagesOutput::Complete(message) => *message,
MessagesOutput::Complete(message) => match message {
MessagesBody::Message(message) => *message,
MessagesBody::Unchecked(_) => panic!("a call returned no message"),
},
MessagesOutput::StreamEnded | MessagesOutput::Detached => {
panic!("a non-streaming call returned a stream")
}

View file

@ -2,7 +2,7 @@ use litellm_host::{
interceptors::{ExecutionFacts, ResultSource},
lifecycle::ExecutionEvent,
};
use litellm_inference_messages::{MessagesCallResponse, messages_body};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body};
use litellm_inference_testing::{
RecordingSecrets, http_config, no_secrets, provider_http, resources,
};
@ -48,6 +48,9 @@ async fn calls_defer_execution_until_polled(
let MessagesCallResponse::Complete(response) = future.await.unwrap() else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
@ -240,10 +243,42 @@ async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: Me
assert_eq!(error.upstream_response(), Some((200, "not json")));
}
#[rstest]
#[case::error_object(
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}),
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"},
"stop_sequence": null}),
)]
#[case::message_missing_id(
json!({"type": "message", "content": [], "stop_sequence": "###"}),
json!({"type": "message", "content": [], "stop_sequence": "###"}),
)]
#[tokio::test]
async fn a_success_body_that_is_not_a_message_passes_through(
call: MessagesCall,
#[case] body: Value,
#[case] expected: Value,
) {
let upstream = upstream([json_response(body)]).await;
let response = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.expect("any JSON object is returned");
let MessagesOutput::Complete(response) = response else {
panic!("expected a completed response");
};
assert_eq!(serde_json::to_value(response).unwrap(), expected);
}
#[rstest]
#[tokio::test]
async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) {
let upstream = upstream([json_response(json!({"unexpected": true}))]).await;
let upstream = upstream([json_response(json!(["unexpected"]))]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
@ -320,6 +355,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes
let MessagesCallResponse::Complete(message) = response else {
panic!("a non-streaming request returns a message");
};
let MessagesBody::Message(message) = message else {
panic!("expected a message");
};
assert_eq!(message.id, "msg_1");
let sent = only_request(&upstream).await;
assert_eq!(sent.header("x-api-key"), Some("sk-ant"));
@ -412,6 +450,9 @@ async fn route_uses_injected_dependencies_and_optional_cache(
else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
@ -484,6 +525,9 @@ async fn cache_overrides_preserve_the_routes_isolated_scope(call: MessagesCall)
else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(response.id, expected["id"].as_str().unwrap());
}
assert_eq!(received(&upstream).await.len(), 2);

View file

@ -164,6 +164,14 @@ impl Error {
}
}
pub fn body_cut(&self) -> Option<litellm_http::BodyCut> {
match self {
Self::Llm(error) => error.body_cut(),
Self::Http(error) => error.body_cut(),
_ => None,
}
}
/// `Some(connect)` when the call timed out, where `connect` says the connection was never
/// established; `None` for every other error.
pub fn timeout(&self) -> Option<bool> {

View file

@ -1913,6 +1913,7 @@ mod tests {
supports_reasoning: false,
supports_adaptive_thinking: false,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: false,
supports_sampling_params: true,

View file

@ -2,8 +2,8 @@ use litellm_llms_types::{
formats::{
chat_completions::ReasoningEffort,
messages::{
AdaptiveThinking, EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig,
ThinkingConfig, ThinkingDisplay,
AdaptiveThinking, DisabledThinking, EffortLevel, MessagesOptionalParams,
MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay,
},
},
recognized::Recognized,
@ -172,7 +172,11 @@ fn translate_reasoning_effort(
) else {
return Ok(MessagesRequest {
params: MessagesOptionalParams {
thinking: None,
thinking: (context.capabilities.thinking_on_by_default
&& !context.capabilities.thinking_always_on)
.then(|| {
Recognized::Known(ThinkingConfig::Disabled(DisabledThinking::default()))
}),
output_config: None,
..request.params
},
@ -748,6 +752,33 @@ mod tests {
);
}
#[rstest]
#[case::default_thinking(true, false, Some(serde_json::json!({"type": "disabled"}))) ]
#[case::legacy(false, false, None)]
#[case::always_on(true, true, None)]
fn none_disables_only_optional_default_thinking(
#[case] thinking_on_by_default: bool,
#[case] thinking_always_on: bool,
#[case] thinking: Option<Value>,
) {
let capabilities = MessagesModelCapabilities {
thinking_on_by_default,
thinking_always_on,
..opus_4_7()
};
let expected = match thinking {
Some(thinking) => serde_json::json!({"max_tokens": 1024, "thinking": thinking}),
None => serde_json::json!({"max_tokens": 1024}),
};
assert_eq!(
translate(
capabilities,
serde_json::json!({"max_tokens": 1024, "reasoning_effort": "none"})
),
Ok(request(expected))
);
}
#[rstest]
#[case::adaptive_model(opus_4_7())]
#[case::effort_capable_model(opus_4_5())]

View file

@ -32,6 +32,8 @@ pub struct MessagesModelCapabilities {
#[serde(default)]
pub thinking_always_on: bool,
#[serde(default)]
pub thinking_on_by_default: bool,
#[serde(default)]
pub supports_legacy_thinking: bool,
#[serde(default)]
pub supports_output_config: bool,
@ -55,6 +57,7 @@ impl Default for MessagesModelCapabilities {
supports_reasoning: false,
supports_adaptive_thinking: false,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: false,
supports_sampling_params: true,

View file

@ -103,6 +103,13 @@ impl Error {
matches!(self, Self::Http(error) if error.is_body_read_timeout())
}
pub fn body_cut(&self) -> Option<litellm_http::BodyCut> {
match self {
Self::Http(error) => error.body_cut(),
_ => None,
}
}
pub fn is_timeout(&self) -> Option<bool> {
match self {
Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()),

View file

@ -142,6 +142,7 @@ impl Catalog {
let mut sections = Vec::new();
let mut entries = IndexMap::with_capacity(document.len());
let mut alias_lists = Vec::new();
let mut pointer_rows = Vec::new();
for (index, (name, value)) in document.into_iter().enumerate() {
let fields = match value {
Value::Object(fields) if !is_section(&name) => fields,
@ -157,16 +158,34 @@ impl Catalog {
continue;
}
};
let (model, aliases, row_problems) = read_row(&fields);
let row = read_row(&fields);
let readable = row.problems.is_empty();
problems.extend(
row_problems
row.problems
.into_iter()
.map(|problem| problem.prepend(&name)),
);
if let Some(aliases) = aliases {
if let Some(target) = row.alias_of {
if readable {
pointer_rows.push((name.clone(), target));
}
sections.push(Section {
index,
key: name,
value: Value::Object(fields),
});
continue;
}
if let Some(aliases) = row.aliases {
alias_lists.push((name.clone(), aliases));
}
entries.insert(name, ModelEntry { fields, model });
entries.insert(
name,
ModelEntry {
fields,
model: row.model,
},
);
}
if entries.is_empty() {
return Err(Error::new("", Repr::Empty));
@ -190,6 +209,21 @@ impl Catalog {
}
}
}
for (alias, target) in pointer_rows {
if !entries.contains_key(&target) {
problems.push(Error::from_key(&alias, Repr::AliasTargetMissing { target }));
} else if let Some(taken_by) = aliases.get(&alias) {
problems.push(Error::from_key(
&alias,
Repr::AliasTaken {
alias: alias.clone(),
taken_by: taken_by.clone(),
},
));
} else {
aliases.insert(alias, target);
}
}
let fallback_generalizations = sections
.iter()
@ -242,8 +276,8 @@ impl Catalog {
}
pub fn row_errors(model: &str, fields: &Fields) -> Vec<Error> {
let (_, _, errors) = read_row(fields);
errors
read_row(fields)
.problems
.into_iter()
.map(|error| error.prepend(model))
.collect()
@ -255,7 +289,13 @@ impl Catalog {
(0..self.rows().entries.len() + self.sections.len()).map_while(move |index| match sections
.next_if(|section| section.index == index)
{
Some(section) => Some((section.key.as_str(), Written::Section(&section.value))),
Some(section) => Some((
section.key.as_str(),
match &section.value {
Value::Object(fields) if !is_section(&section.key) => Written::Entry(fields),
value => Written::Section(value),
},
)),
None => entries
.next()
.map(|(name, entry)| (name.as_str(), Written::Entry(&entry.fields))),
@ -406,16 +446,41 @@ impl Catalog {
}
}
fn read_row(fields: &Fields) -> (Option<Model>, Option<Vec<String>>, Vec<Error>) {
struct RowRead {
model: Option<Model>,
aliases: Option<Vec<String>>,
alias_of: Option<String>,
problems: Vec<Error>,
}
fn read_row(fields: &Fields) -> RowRead {
let wire = EntryWire::read(fields);
if let Some(target) = wire.attributes.alias_of.clone() {
let other_fields = fields
.keys()
.filter(|key| *key != "alias_of")
.map(|key| Error::from_key(key, Repr::AliasRowField));
return RowRead {
model: None,
aliases: None,
alias_of: Some(target),
problems: wire.problems.into_iter().chain(other_fields).collect(),
};
}
let aliases = wire.attributes.aliases.clone();
let timezone_error = unknown_timezone(&wire);
match Model::try_from(wire) {
Ok(model) => (Some(model), aliases, timezone_error.into_iter().collect()),
let (model, problems) = match Model::try_from(wire) {
Ok(model) => (Some(model), timezone_error.into_iter().collect()),
Err(mut errors) => {
errors.extend(timezone_error);
(None, aliases, errors)
(None, errors)
}
};
RowRead {
model,
aliases,
alias_of: None,
problems,
}
}

View file

@ -67,6 +67,12 @@ pub(crate) enum Repr {
#[error("alias {alias:?} already belongs to {taken_by}")]
#[strum_discriminants(strum(serialize = "alias_taken"))]
AliasTaken { alias: String, taken_by: String },
#[error("alias_of {target:?} is not a model name")]
#[strum_discriminants(strum(serialize = "alias_target_missing"))]
AliasTargetMissing { target: String },
#[error("an alias_of row carries no other field")]
#[strum_discriminants(strum(serialize = "alias_row_field"))]
AliasRowField,
#[error("{first} and {second} price the same meter under the same conditions")]
#[strum_discriminants(strum(serialize = "duplicate_price"))]
DuplicatePrice { first: String, second: String },
@ -141,6 +147,8 @@ impl Error {
| Repr::UnknownTimezone { .. }
| Repr::AliasShadowsModel { .. }
| Repr::AliasTaken { .. }
| Repr::AliasTargetMissing { .. }
| Repr::AliasRowField
| Repr::DuplicatePrice { .. }
| Repr::SuffixInsideBlock { .. }
| Repr::TierRange
@ -327,6 +335,8 @@ mod tests {
#[case::unknown_timezone(Repr::UnknownTimezone { zone: "invalid".into() }, Rejects::Row)]
#[case::alias_shadows_model(Repr::AliasShadowsModel { alias: "model".into() }, Rejects::Row)]
#[case::alias_taken(Repr::AliasTaken { alias: "alias".into(), taken_by: "owner".into() }, Rejects::Row)]
#[case::alias_target_missing(Repr::AliasTargetMissing { target: "owner".into() }, Rejects::Row)]
#[case::alias_row_field(Repr::AliasRowField, Rejects::Row)]
#[case::duplicate_price(Repr::DuplicatePrice { first: "a".into(), second: "b".into() }, Rejects::Row)]
#[case::suffix_inside_block(Repr::SuffixInsideBlock { block: "tiered_pricing" }, Rejects::Row)]
#[case::tier_range(Repr::TierRange, Rejects::Row)]

View file

@ -104,9 +104,18 @@ pub fn registry_json_schema() -> Schema {
"title".into(),
"LiteLLM model prices and context window registry".into(),
);
definitions.insert(
"aliasRow".into(),
json!({
"type": "object",
"required": ["alias_of"],
"properties": {"alias_of": {"type": "string"}},
"additionalProperties": false,
}),
);
root.insert(
"additionalProperties".into(),
json!({"$ref": "#/$defs/modelEntry"}),
json!({"anyOf": [{"$ref": "#/$defs/modelEntry"}, {"$ref": "#/$defs/aliasRow"}]}),
);
root.insert("$defs".into(), Value::Object(definitions));
Schema::from(root.clone())
@ -153,6 +162,7 @@ fn decorate_model_entry(schema: &mut Value) {
let rates = json!({FLAT_KEY_PATTERN: {"type": "number", "minimum": 0}});
let object = schema.as_object_mut().unwrap();
object.insert("required".into(), json!(["litellm_provider"]));
object.insert("not".into(), json!({"required": ["alias_of"]}));
object.insert("additionalProperties".into(), Value::Bool(true));
object.insert(
"patternProperties".into(),

View file

@ -47,6 +47,8 @@ macro_rules! attributes {
}
attributes! {
/// Model this row is another name for; such a row carries no other key.
alias_of: String,
/// Other names this entry is served under.
aliases: Vec<String>,
audio_transcription_config: String,

View file

@ -139,6 +139,18 @@ fn a_bad_entry_is_reported_by_model_and_key_and_kept_raw() {
json!({"First": {"litellm_provider": "t", "aliases": ["shared"]}, "Second": {"litellm_provider": "t", "aliases": ["shared"]}}),
r#"Second: alias "shared" already belongs to First"#
)]
#[case::alias_of_target_missing(
json!({"Real": {"litellm_provider": "t"}, "ghost-latest": {"alias_of": "Ghost"}}),
r#"ghost-latest: alias_of "Ghost" is not a model name"#
)]
#[case::alias_of_row_with_other_fields(
json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real", "mode": "chat"}}),
"real-latest.mode: an alias_of row carries no other field"
)]
#[case::alias_of_claimed_by_a_list(
json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real"}, "Other": {"litellm_provider": "t", "aliases": ["real-latest"]}}),
r#"real-latest: alias "real-latest" already belongs to Other"#
)]
#[case::aliases_not_a_list(
json!({"a": {"litellm_provider": "t", "aliases": "bad"}}),
"a.aliases: invalid value: invalid type: string \"bad\", expected a sequence"
@ -399,3 +411,70 @@ fn a_bad_value_keeps_its_json_error_as_the_source_and_names_its_key_and_field()
assert_eq!(error.key().as_deref(), Some("open/router"));
assert_eq!(error.field().as_deref(), Some("input_cost_per_token"));
}
#[rstest]
fn an_alias_of_row_is_served_as_its_target_and_written_as_itself() {
let catalog = catalog(json!({
"Real": {"litellm_provider": "t", "input_cost_per_token": 1e-6},
"real-latest": {"alias_of": "Real"},
"After": {"litellm_provider": "t"}
}));
let found = catalog.lookup("REAL-LATEST").unwrap();
assert_eq!(
(found.matched_key, found.canonical_key),
("real-latest", "Real")
);
assert_eq!(
found.entry.field("input_cost_per_token"),
Some(&json!(1e-6))
);
assert_eq!(catalog.model_count(), 2);
assert_eq!(
catalog.served().map(|(key, _)| key).collect::<Vec<_>>(),
["Real", "After", "real-latest"]
);
assert_eq!(
written(&catalog),
r#"{"Real":{"litellm_provider":"t","input_cost_per_token":1e-6},"real-latest":{"alias_of":"Real"},"After":{"litellm_provider":"t"}}"#
);
}
#[rstest]
fn a_dangling_alias_of_row_is_rejected_and_dropped_from_the_written_document() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t"},
"ghost-latest": {"alias_of": "Ghost"}
}));
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert!(catalog.lookup("ghost-latest").is_none());
assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#);
}
#[rstest]
fn an_alias_of_row_with_other_fields_is_never_served() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t"},
"real-latest": {"alias_of": "Real", "mode": "chat"}
}));
assert!(parsed.value.lookup("real-latest").is_none());
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert!(catalog.lookup("real-latest").is_none());
assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#);
}
#[rstest]
fn a_pointer_row_that_repeats_a_listed_alias_is_dropped_and_the_model_is_kept() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t", "aliases": ["real-latest"]},
"real-latest": {"alias_of": "Real"}
}));
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert_eq!(catalog.lookup("real-latest").unwrap().canonical_key, "Real");
assert_eq!(
written(&catalog),
r#"{"Real":{"litellm_provider":"t","aliases":["real-latest"]}}"#
);
}

View file

@ -35,6 +35,7 @@ fn generated_registry_schema_validates_checked_in_catalog() {
#[case(json!({"example": {"litellm_provider": "test"}}))]
#[case(json!({"example": {"litellm_provider": "test", "future_field": true}}))]
#[case(json!({"sample_spec": {"litellm_provider": "placeholder"}}))]
#[case(json!({"example": {"litellm_provider": "test"}, "example-latest": {"alias_of": "example"}}))]
fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) {
assert!(registry_validator().is_valid(&document));
}
@ -54,6 +55,9 @@ fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value)
#[case::empty_windows(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}))]
#[case::invalid_weekday(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}))]
#[case::invalid_aliases(json!({"litellm_provider": "test", "aliases": "wrong"}))]
#[case::alias_of_on_a_model(json!({"litellm_provider": "test", "alias_of": "other"}))]
#[case::alias_row_with_other_fields(json!({"alias_of": "other", "mode": "chat"}))]
#[case::alias_of_not_a_string(json!({"alias_of": 4}))]
#[case::non_object_model(json!(4))]
fn generated_registry_schema_rejects_invalid_entries(#[case] entry: Value) {
assert!(!registry_validator().is_valid(&json!({"example": entry})));

View file

@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"]
tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
rustls.workspace = true
litellm-secret-value.workspace = true
fancy-regex.workspace = true
litellm-tracing.workspace = true
@ -30,6 +31,14 @@ bytes.workspace = true
futures-util.workspace = true
litellm-cache.workspace = true
litellm-error.workspace = true
litellm-circuit-breaker.workspace = true
litellm-redis-client.workspace = true
litellm-redis-connect.workspace = true
litellm-redis-store.workspace = true
litellm-redis-resolve.workspace = true
litellm-python-compat.workspace = true
litellm-redis-types.workspace = true
tracing.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
@ -52,7 +61,8 @@ litellm-router-core.workspace = true
litellm-host-python.workspace = true
litellm-model-catalog.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
indexmap = "2.14.0"
pyo3 = { workspace = true, features = ["indexmap"] }
prost.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true

View file

@ -6,6 +6,8 @@ pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
use pyo3::{exceptions::PyRuntimeError, prelude::*};
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
#[derive(Debug, thiserror::Error)]
pub(crate) enum HandleError {
#[error("this inference route does not support configured caches")]

View file

@ -4,6 +4,6 @@ This directory lets Rust inference use a selected Python cache. `service.rs` imp
Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details
Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. The one exception is the asynchronous unary messages store: Python's `_async_set_cache` writes the returned response from `create_cache_write_task` after the router has stamped its `_hidden_params`, so `host.rs` defers that store until the deployment transform has returned the response the caller receives and hands that response to `litellm.rust_bridge.host.cache.store_response`, which schedules the same write task. Lookups, sync stores and stream stores stay in the caller's task. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation

View file

@ -9,7 +9,7 @@ use pyo3::{
use serde_json::Value;
use super::service::CacheCall;
use crate::cache::HandleError;
use crate::cache::{CACHE_HOST_MODULE, HandleError};
enum Pending {
Lookup(Reply<Result<Option<Value>, Error>>),
@ -21,6 +21,7 @@ pub(crate) struct PythonCache {
arguments: Option<Py<PyDict>>,
pending: Option<Pending>,
asynchronous: bool,
deferred: bool,
}
impl PythonCache {
@ -30,6 +31,7 @@ impl PythonCache {
arguments: None,
pending: None,
asynchronous,
deferred: false,
}
}
@ -64,7 +66,14 @@ impl PythonCache {
None,
)
}
CacheCall::Store { value, reply } => {
CacheCall::Store {
unary: true, reply, ..
} if self.asynchronous => {
self.deferred = true;
reply.send(Ok(()));
return Ok(None);
}
CacheCall::Store { value, reply, .. } => {
self.pending = Some(Pending::Store(reply));
(
if self.asynchronous {
@ -89,6 +98,19 @@ impl PythonCache {
self.resume(py, result)
}
pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
if !std::mem::take(&mut self.deferred) {
return Ok(());
}
let (Some(cache), Some(arguments)) = (self.cache.as_ref(), self.arguments.as_ref()) else {
return Err(HandleError::Unconfigured.into());
};
py.import(CACHE_HOST_MODULE)?
.getattr("store_response")?
.call((cache, response), Some(&arguments.bind(py).copy()?))?;
Ok(())
}
pub fn resume(
&mut self,
py: Python<'_>,
@ -124,6 +146,7 @@ impl PythonCache {
pub fn close(&mut self) {
self.pending = None;
self.deferred = false;
self.cache = None;
self.arguments = None;
}

View file

@ -45,6 +45,7 @@ pub(crate) enum CacheCall {
},
Store {
value: Value,
unary: bool,
reply: Reply<Result<(), Error>>,
},
}
@ -101,8 +102,13 @@ where
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'a>> {
Box::pin(async move {
let value = to_python(value)?;
let unary = value.get(STREAM_EVENTS_KEY).is_none();
self.services
.call(|reply| CacheCall::Store { value, reply })
.call(|reply| CacheCall::Store {
value,
unary,
reply,
})
.await
.map_err(Error::unavailable)?
})

View file

@ -1,11 +1,10 @@
use super::{HandleError, python};
use super::{CACHE_HOST_MODULE, HandleError, python};
use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache};
use litellm_host::{
machine::{HostServices, MachineFault},
protocol::Protocol,
};
use pyo3::{prelude::*, types::PyDict};
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
pub(crate) struct Cached<P>(std::marker::PhantomData<P>);

View file

@ -29,6 +29,8 @@ const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param";
/// The phase label (`Connection` or `Read`) of a timed-out call, for Python's timeout text, or
/// `Body` for a streamed body read that timed out after the head arrived.
const MESSAGES_TIMEOUT: &str = "messages_timeout";
// `(received, expected, reset)` of a streamed body the connection ended early.
const MESSAGES_BODY_CUT: &str = "messages_body_cut";
/// No credential reached the provider, which Python raises as `AuthenticationError`.
const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential";
/// Python refuses the request only once its logger knows the resolved model and provider (the
@ -134,6 +136,9 @@ pub(crate) fn route_failure(
};
value.setattr(MESSAGES_TIMEOUT, phase)?;
}
if let Some(cut) = error.body_cut() {
value.setattr(MESSAGES_BODY_CUT, (cut.received, cut.expected, cut.reset))?;
}
if error.is_missing_credential() {
value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?;
}

View file

@ -12,6 +12,7 @@ mod marshal;
mod model_registry;
mod preflight;
mod python_settings;
mod redis;
mod router;
mod routes;
mod secrets;
@ -37,6 +38,31 @@ mod _native {
#[pymodule_export]
use crate::model_registry::bundled_model_cost_map;
#[pymodule_export]
use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states};
#[pymodule_export]
use crate::redis::client::NativeRedisClient;
#[pymodule_export]
use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes};
#[pymodule_export]
use crate::redis::store::rate_limits::{
NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired,
PyGuardedReply, PyWindowCount,
};
#[pymodule_export]
use crate::redis::store::request::{NativeRequestRedis, NativeSend};
#[pymodule_export]
use crate::redis::store::router_usage::NativeRouterUsage;
#[pymodule_export]
use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins;
#[pymodule_export]
use crate::redis::store::spend_counters::{
NativeSpendCounters, PyPostCallSettlement, PySpendRead,
};
#[pymodule_export]
use crate::redis::store::{
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
};
#[pymodule_export]
use crate::router::{RouterCall, RouterCoreHandle};
#[pymodule_export]
use crate::routes::audio_transcription::{atranscription, transcription};
@ -66,6 +92,7 @@ mod _native {
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
let _ = rustls::crypto::ring::default_provider().install_default();
let py = module.py();
let dict = module.dict();
dict.set_item(
@ -111,6 +138,29 @@ mod tests {
"aresponses",
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"NativeRedisClient",
"NativeRedisBreaker",
"NativeAdmission",
"NativeRequestRedis",
"NativeSend",
"NativeAuthCache",
"NativeRateLimits",
"WindowCount",
"AtomicApplied",
"AtomicOverLimit",
"GuardedReply",
"GaugesAcquired",
"GaugeOverLimit",
"NativePending",
"NativeRouterCooldowns",
"NativeRouterUsage",
"NativeSensitiveRoutePins",
"NativeRedisStoreError",
"NativeSpendCounters",
"PostCallSettlement",
"SpendRead",
"LocalSizes",
"native_redis_breaker_states",
"NativeClickHouseSpendConfig",
"NativeClickHouseSpendStorage",
"NativeDatabase",
@ -140,4 +190,18 @@ mod tests {
assert_eq!(public_names, expected);
});
}
#[rstest::rstest]
fn module_import_installs_ring_as_the_process_crypto_provider() {
Python::initialize();
Python::attach(|py| {
native_module(py);
});
let installed = rustls::crypto::CryptoProvider::get_default()
.expect("importing _native should install a process crypto provider");
assert_eq!(
format!("{installed:?}"),
format!("{:?}", rustls::crypto::ring::default_provider())
);
}
}

View file

@ -13,6 +13,8 @@ pub(crate) enum PythonSettings {
UrlPolicy,
#[strum(serialize = "provider_defaults")]
ProviderDefaults,
#[strum(serialize = "router_tiers")]
RouterTiers,
#[strum(serialize = "secret_manager")]
SecretManager,
#[strum(serialize = "secret_manager_binding")]

View file

@ -0,0 +1,134 @@
use std::{
sync::{Arc, Mutex, PoisonError},
time::{Duration, Instant},
};
use litellm_circuit_breaker::{Admission, Breaker, Config, Failure};
use litellm_redis_client::RedisEvents;
use litellm_redis_types::Endpoint;
use pyo3::{exceptions::PyValueError, prelude::*};
use crate::redis::{
client::{RedisSpec, resolved_endpoint},
events::{ServiceLoggingEvents, state_label},
hosts,
};
/// `litellm.redis.store_settings.RedisBreakerSettings`, filled from `litellm.constants`.
#[derive(FromPyObject)]
pub(crate) struct BreakerSettings {
failure_threshold: u32,
recovery_timeout: f64,
timeout_min_duration: f64,
}
impl BreakerSettings {
pub(crate) fn config(&self) -> PyResult<Config> {
let seconds = |seconds: f64| {
Duration::try_from_secs_f64(seconds)
.map_err(|error| PyValueError::new_err(error.to_string()))
};
Ok(Config {
failure_threshold: self.failure_threshold,
recovery_timeout: seconds(self.recovery_timeout)?,
timeout_min_duration: seconds(self.timeout_min_duration)?,
})
}
}
/// The one health signal for a Redis, shared by its `RedisCache` and its native client, so a Redis the
/// cache found wedged is fast-failed by every store too.
#[pyclass(frozen, name = "NativeRedisBreaker")]
pub(crate) struct NativeRedisBreaker(pub(crate) Breaker);
impl NativeRedisBreaker {
fn build(
py: Python<'_>,
settings: &BreakerSettings,
endpoint: Option<Endpoint>,
) -> PyResult<Self> {
let config = settings.config()?;
let runtime = litellm_host_python::runtime_handle()?;
let events: Arc<dyn RedisEvents> =
Arc::new(ServiceLoggingEvents::transitions_only(py, runtime)?);
Ok(Self(py.detach(|| {
hosts().breaker(endpoint, config, Some(events))
})))
}
}
#[pymethods]
impl NativeRedisBreaker {
#[new]
fn new(py: Python<'_>, settings: BreakerSettings) -> PyResult<Self> {
Self::build(py, &settings, None)
}
/// The breaker of the Redis server `redis_kwargs` names, which every cache and client on that server shares
/// whatever its db, credentials or namespace. The first caller's settings win.
#[staticmethod]
fn shared(
py: Python<'_>,
redis_kwargs: RedisSpec,
settings: BreakerSettings,
) -> PyResult<Self> {
let endpoint = resolved_endpoint(py, &redis_kwargs)?;
Self::build(py, &settings, endpoint)
}
fn admit(&self) -> Option<NativeAdmission> {
self.0
.admit(Instant::now())
.ok()
.map(|admission| NativeAdmission(Mutex::new(Some(admission))))
}
fn failed(&self, timeout: bool) {
self.0.failed(failure(timeout), Instant::now());
}
#[getter]
fn state(&self) -> &'static str {
state_label(&self.0.state())
}
}
/// One admitted call. Settling it twice records nothing more, and dropping an unsettled probe reopens the breaker.
#[pyclass(frozen, name = "NativeAdmission")]
pub(crate) struct NativeAdmission(Mutex<Option<Admission>>);
impl NativeAdmission {
fn take(&self) -> Option<Admission> {
self.0.lock().unwrap_or_else(PoisonError::into_inner).take()
}
}
#[pymethods]
impl NativeAdmission {
fn succeeded(&self) {
if let Some(admission) = self.take() {
admission.succeeded();
}
}
fn failed(&self, timeout: bool) {
if let Some(admission) = self.take() {
admission.failed(failure(timeout), Instant::now());
}
}
}
fn failure(timeout: bool) -> Failure {
match timeout {
true => Failure::Timeout,
false => Failure::Hard,
}
}
#[pyfunction]
pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> {
py.detach(|| hosts().breakers())
.iter()
.map(|breaker| state_label(&breaker.state()))
.collect()
}

View file

@ -0,0 +1,47 @@
use std::sync::Arc;
use litellm_host_python::{PythonContext, attach_blocking};
use litellm_redis_connect::{CredentialSource, CredentialStage, Error, Expiring, SourceFuture};
use litellm_secret_value::SecretValue;
use pyo3::prelude::*;
const MODULE: &str = "litellm.rust_bridge.host.redis_credentials";
/// A redis-py `CredentialProvider` the SDK caller passed in `cache_kwargs`, asked through the host
/// module the way redis-py's async connection asks it.
pub(crate) struct PythonCredentialSource {
provider: Arc<Py<PyAny>>,
context: PythonContext,
}
impl PythonCredentialSource {
pub(crate) fn new(provider: Py<PyAny>, context: PythonContext) -> Self {
Self {
provider: Arc::new(provider),
context,
}
}
}
impl CredentialSource for PythonCredentialSource {
fn fetch(&self) -> SourceFuture<'_> {
let provider = Arc::clone(&self.provider);
let context = self.context.clone();
Box::pin(async move {
let (username, password) = attach_blocking(context, move |py| {
py.import(MODULE)?
.getattr("credentials")?
.call1((provider.bind(py),))?
.extract::<(Option<String>, String)>()
})
.await
.and_then(|credentials| credentials)
.map_err(Error::credential(CredentialStage::Provider))?;
Ok(Expiring {
username,
password: SecretValue::new(password),
expires_at: None,
})
})
}
}

View file

@ -0,0 +1,337 @@
mod credentials;
use std::{collections::BTreeMap, sync::Arc, time::Duration};
use litellm_circuit_breaker::Config as BreakerConfig;
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_host_python::PythonContext;
use litellm_python_compat::number::{self, SecondsError};
use litellm_redis_client::{RedisClient, RedisEvents};
use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment};
use litellm_redis_resolve::{Kwargs, PythonException};
use litellm_redis_types::{Endpoint, ProviderHandle, Resolution};
use pyo3::{
exceptions::{PyOverflowError, PyTypeError, PyValueError},
prelude::*,
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple},
};
use self::credentials::PythonCredentialSource;
use crate::redis::{breaker::BreakerSettings, events::ServiceLoggingEvents, hosts};
pyo3::import_exception!(redis.exceptions, RedisClusterException);
/// The kwargs `_merge_redis_kwargs` returns for a `RedisCache`, which the resolver reads the way
/// redis-py's constructors do. A key whose value is a Python object, such as `retry`, reaches the
/// resolver by name only, so it can name the setting it cannot honor. A `credential_provider` object
/// is also kept, so the connector can ask it for credentials, and reaches the resolver as a handle.
pub(crate) struct RedisSpec {
kwargs: Kwargs,
credential_provider: Option<Py<PyAny>>,
}
impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
type Error = PyErr;
fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let py = kwargs.py();
let source = PyDict::new(py);
let mapping = kwargs.cast::<PyMapping>()?;
source.update(&mapping)?;
let plain = PyDict::new(py);
let mut host_objects = Vec::new();
for (name, value) in source.iter() {
match value.is_none() || is_plain(&value) {
true => plain.set_item(&name, value)?,
false => host_objects.push(name.extract::<String>()?),
}
}
let credential_provider = source
.get_item("credential_provider")?
.filter(|provider| !provider.is_none() && !is_plain(provider))
.map(Bound::unbind);
let fields: BTreeMap<String, serde_json::Value> =
litellm_host_python::from_py_argument(plain.as_any())?;
let objects = host_objects
.into_iter()
.map(|name| (name, serde_json::Value::Bool(true)));
Ok(Self {
kwargs: litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)),
credential_provider,
})
}
}
/// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built.
#[derive(FromPyObject)]
pub(crate) struct StoreSettings {
redis_kwargs: RedisSpec,
namespace: Option<String>,
default_ttl: f64,
breaker: Option<BreakerSettings>,
}
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
value.is_instance_of::<PyString>()
|| value.is_instance_of::<PyBool>()
|| value.is_instance_of::<PyInt>()
|| value.is_instance_of::<PyFloat>()
|| value.is_instance_of::<PyList>()
|| value.is_instance_of::<PyTuple>()
|| value.is_instance_of::<PyDict>()
}
/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool
/// and its breaker. A config redis-py refuses building its client raises the same class here. One it only fails on
/// once a command connects, where `RedisCache` logs every failed call, leaves the family on memory with an error.
fn shared_client(
py: Python<'_>,
spec: RedisSpec,
namespace: Option<String>,
breaker: Option<BreakerConfig>,
) -> PyResult<Option<RedisClient>> {
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
let logger = crate::logger::capture(py);
let resolved = match logger.scope(|| resolve_spec(py, &spec))? {
Ok(resolution) => {
logger.scope(|| {
for ignored in &resolution.ignored {
litellm_tracing::warn!(
"Redis setting {} is not supported by the Rust client and is ignored: {}",
ignored.setting,
ignored.consequence
);
}
});
Resolution {
namespace,
..resolution
}
}
Err(error) => return unresolved(py, &error),
};
let runtime = litellm_host_python::runtime_handle()?;
let events: Arc<dyn RedisEvents> = Arc::new(ServiceLoggingEvents::new(py, runtime.clone())?);
let host = HostServices {
events: Some(events),
http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
environment,
credential_provider: spec
.credential_provider
.as_ref()
.map(|provider| -> PyResult<Arc<dyn CredentialSource>> {
Ok(Arc::new(PythonCredentialSource::new(
provider.clone_ref(py),
PythonContext::capture(py)?,
)))
})
.transpose()?,
};
let connected = py.detach(|| {
let _runtime = runtime.enter();
hosts().client(resolved, host, breaker)
});
match connected {
Ok(client) => Ok(Some(client)),
Err(error) => Ok(unconnected(py, &error)),
}
}
/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store
/// never sees the redis-py kwargs or the namespace.
#[pyclass(frozen, name = "NativeRedisClient")]
pub(crate) struct NativeRedisClient {
pub(crate) client: RedisClient,
/// `RedisCache.default_ttl`, the TTL a store's write takes when its caller gives none.
pub(crate) default_ttl: Duration,
}
#[pymethods]
impl NativeRedisClient {
/// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory.
#[staticmethod]
fn connect(py: Python<'_>, settings: StoreSettings) -> PyResult<Option<Self>> {
let StoreSettings {
redis_kwargs,
namespace,
default_ttl,
breaker,
} = settings;
let default_ttl = number::seconds(default_ttl).map_err(seconds_error)?;
let breaker = breaker.as_ref().map(BreakerSettings::config).transpose()?;
Ok(
shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self {
client,
default_ttl,
}),
)
}
#[getter]
fn default_ttl_seconds(&self) -> f64 {
self.default_ttl.as_secs_f64()
}
}
/// The Redis server `spec` names, whatever its db, credentials or namespace, or none when the Rust client cannot reach it.
pub(crate) fn resolved_endpoint(py: Python<'_>, spec: &RedisSpec) -> PyResult<Option<Endpoint>> {
Ok(resolve_spec(py, spec)?
.ok()
.and_then(|resolution| resolution.target.endpoint()))
}
// A live registry entry's client keeps its provider alive, so no live entry shares its address
fn resolve_spec(
py: Python<'_>,
spec: &RedisSpec,
) -> PyResult<Result<Resolution, litellm_redis_resolve::Error>> {
let handle = spec
.credential_provider
.as_ref()
.map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64));
resolve(&spec.kwargs, handle, &py.import("os")?.getattr("environ")?)
}
/// What redis-py makes of the kwargs, with `REDIS_USERNAME` read from `environ` the way Python's Azure AD auth
/// reads `os.environ`.
fn resolve(
kwargs: &Kwargs,
provider: Option<ProviderHandle>,
environ: &Bound<'_, PyAny>,
) -> PyResult<Result<Resolution, litellm_redis_resolve::Error>> {
let redis_username: Option<String> = environ
.call_method1("get", ("REDIS_USERNAME",))?
.extract()?;
Ok(litellm_redis_resolve::interpret(
kwargs,
provider,
redis_username.as_deref(),
))
}
fn unresolved(
py: Python<'_>,
error: &litellm_redis_resolve::Error,
) -> PyResult<Option<RedisClient>> {
let message = || litellm_error::Report::new(error).to_string();
match error.raised_building_the_python_client() {
Some(PythonException::ValueError) => Err(PyValueError::new_err(message())),
Some(PythonException::TypeError) => Err(PyTypeError::new_err(message())),
Some(PythonException::RedisClusterException) => {
Err(RedisClusterException::new_err(message()))
}
None => Ok(memory_only(py, error)),
}
}
// redis-py reads TLS files and fetches credentials only once a command connects, so none of these fails building
// its client
fn unconnected(py: Python<'_>, error: &litellm_redis_connect::Error) -> Option<RedisClient> {
match error {
litellm_redis_connect::Error::Connect(_)
| litellm_redis_connect::Error::NoHostCredentialProvider
| litellm_redis_connect::Error::ProcessEnvironmentRequired { .. }
| litellm_redis_connect::Error::Credential { .. }
| litellm_redis_connect::Error::CredentialTimeout { .. }
| litellm_redis_connect::Error::ExpiredCredential
| litellm_redis_connect::Error::MissingCredential { .. } => memory_only(py, error),
}
}
fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option<RedisClient> {
let report = litellm_error::Report::new(error).to_string();
crate::logger::capture(py).scope(|| {
litellm_tracing::error!(
"Redis is not usable from the Rust stores, their state stays in memory: {report}"
)
});
None
}
fn seconds_error(error: SecondsError) -> PyErr {
match error {
SecondsError::Overflow => PyOverflowError::new_err(error.to_string()),
SecondsError::NotANumber => PyValueError::new_err(error.to_string()),
}
}
#[cfg(test)]
mod tests {
use std::ffi::CStr;
use litellm_redis_types::Auth;
use rstest::rstest;
use serde_json::json;
use super::*;
fn kwargs(fields: serde_json::Value) -> Kwargs {
let fields: BTreeMap<String, serde_json::Value> = serde_json::from_value(fields).unwrap();
litellm_redis_resolve::kwargs(fields)
}
#[rstest]
#[case::set(c"environ = {'REDIS_USERNAME': 'azure-user'}", Some("azure-user"))]
#[case::empty_is_unset(c"environ = {'REDIS_USERNAME': ''}", None)]
fn azure_auth_takes_its_username_from_the_environ_like_python(
#[case] source: &CStr,
#[case] expected: Option<&str>,
) {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
py.run(source, Some(&locals), Some(&locals)).unwrap();
let environ = locals.get_item("environ").unwrap().unwrap();
let kwargs = kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"}));
let resolved = resolve(&kwargs, None, &environ).unwrap().unwrap();
assert!(
matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected),
"{:?}",
resolved.auth
);
});
}
#[rstest]
#[case::no_target(json!({"port": 6379}), Some("ValueError"))]
#[case::unenforceable(json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), Some("TypeError"))]
#[case::fails_only_once_a_command_connects(
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": "None"}),
None
)]
fn a_config_raises_where_redis_py_raises_building_its_client(
#[case] fields: serde_json::Value,
#[case] raised: Option<&str>,
) {
Python::initialize();
Python::attach(|py| {
let error = litellm_redis_resolve::interpret(&kwargs(fields), None, None).unwrap_err();
let outcome = unresolved(py, &error)
.map(|client| client.is_some())
.map_err(|error| error.get_type(py).name().unwrap().to_string());
assert_eq!(
outcome,
raised.map_or(Ok(false), |class| Err(class.to_owned()))
);
});
}
#[rstest]
#[case::overflow(f64::INFINITY, "OverflowError")]
#[case::not_a_number(f64::NAN, "ValueError")]
fn a_default_ttl_python_cannot_make_an_int_raises_like_int(
#[case] seconds: f64,
#[case] raised: &str,
) {
Python::initialize();
Python::attach(|py| {
let error = number::seconds(seconds).map_err(seconds_error).unwrap_err();
assert_eq!(error.get_type(py).name().unwrap().to_string(), raised);
});
}
}

View file

@ -0,0 +1,403 @@
use std::{
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_circuit_breaker::{State, Transition};
use litellm_host_python::{PythonContext, attach_blocking};
use litellm_redis_client::{
CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged,
};
use pyo3::{prelude::*, types::PyModule};
use tokio::runtime::Handle;
/// The asyncio loop a Python caller ran on. `RedisCache` scheduled its service events there with
/// `asyncio.create_task`, so the native events go there too instead of a new loop per event.
#[derive(Clone)]
pub(crate) struct EventLoop(Arc<Py<PyAny>>);
impl EventLoop {
pub(crate) fn running(py: Python<'_>) -> Option<Self> {
pyo3_async_runtimes::get_running_loop(py)
.ok()
.map(|event_loop| Self(Arc::new(event_loop.unbind())))
}
}
/// The Python caller of one store operation. Its context carries `with_service_target` and the
/// ambient OTel span, and its `call_type` is the label `RedisCache` reports, so dashboards keep
/// their series.
pub(crate) struct RedisCall {
pub(crate) context: PythonContext,
pub(crate) event_loop: Option<EventLoop>,
pub(crate) call_type: &'static str,
}
impl RedisCall {
pub(crate) fn tag(self) -> CallTag {
Arc::new(self)
}
pub(crate) async fn scope<F: Future>(self, operation: F) -> F::Output {
tagged(self.tag(), operation).await
}
}
/// Reports every native Redis call and breaker transition to `ServiceLogging`.
pub(crate) struct ServiceLoggingEvents {
record_call: Arc<Py<PyAny>>,
record_breaker: Arc<Py<PyAny>>,
context: PythonContext,
/// The loop the client was built on, for calls whose caller had none running.
event_loop: Option<EventLoop>,
runtime: Handle,
}
impl ServiceLoggingEvents {
const MODULE: &str = "litellm.rust_bridge.host.redis_events";
pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult<Self> {
Self::from_module(
py,
&py.import(Self::MODULE)?,
runtime,
EventLoop::running(py),
)
}
/// For a breaker, whose transitions never run on the caller's loop, so building it never looks one up.
pub(crate) fn transitions_only(py: Python<'_>, runtime: Handle) -> PyResult<Self> {
Self::from_module(py, &py.import(Self::MODULE)?, runtime, None)
}
fn from_module(
py: Python<'_>,
module: &Bound<'_, PyModule>,
runtime: Handle,
event_loop: Option<EventLoop>,
) -> PyResult<Self> {
Ok(Self {
record_call: Arc::new(module.getattr("record_redis_call")?.unbind()),
record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()),
context: PythonContext::capture(py)?,
event_loop,
runtime,
})
}
fn report(&self, context: PythonContext, record: Arc<Py<PyAny>>, args: ReportArgs) {
self.runtime.spawn(async move {
let reported = attach_blocking(context, move |py| match args {
ReportArgs::Call(call) => record
.bind(py)
.call1((
call.event_loop
.as_ref()
.map(|event_loop| event_loop.0.bind(py)),
call.call_type,
call.start_time,
call.end_time,
call.failure,
))
.map(drop),
ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop),
})
.await
.and_then(|reported| reported);
if let Err(error) = reported {
tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed");
}
});
}
}
/// One call as `record_redis_call` takes it, its times in `time.time()` seconds.
struct CallReport {
event_loop: Option<EventLoop>,
call_type: &'static str,
start_time: f64,
end_time: f64,
failure: Option<&'static str>,
}
enum ReportArgs {
Call(CallReport),
Breaker(&'static str),
}
/// The wall-clock start and end of a call that just finished after `duration`.
fn wall_clock_span(duration: Duration) -> (f64, f64) {
let end = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default();
(
end.saturating_sub(duration).as_secs_f64(),
end.as_secs_f64(),
)
}
fn operation_call_type(operation: Operation) -> &'static str {
match operation {
Operation::Command(name) => name,
Operation::Script => "eval",
Operation::Pipeline { .. } => "pipeline",
}
}
fn failure_label(outcome: Outcome) -> Option<&'static str> {
match outcome {
Outcome::Succeeded => None,
Outcome::Failed(FailureKind::Timeout) => Some("timeout"),
Outcome::Failed(FailureKind::Transport) => Some("transport"),
Outcome::Failed(FailureKind::Server) => Some("server"),
Outcome::Failed(FailureKind::Decode) => Some("decode"),
Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"),
}
}
pub(crate) fn state_label(state: &State) -> &'static str {
match state {
State::Closed { .. } => "closed",
State::Open { .. } => "open",
State::HalfOpen => "half_open",
}
}
impl RedisEvents for ServiceLoggingEvents {
fn call(&self, event: &CallEvent) {
let scoped = current_tag().and_then(|tag| tag.downcast::<RedisCall>().ok());
let (context, event_loop, call_type) = match scoped {
Some(call) => (
call.context.clone(),
call.event_loop.clone().or_else(|| self.event_loop.clone()),
call.call_type,
),
None => (
self.context.clone(),
self.event_loop.clone(),
operation_call_type(event.operation),
),
};
let (start_time, end_time) = wall_clock_span(event.duration);
self.report(
context,
self.record_call.clone(),
ReportArgs::Call(CallReport {
event_loop,
call_type,
start_time,
end_time,
failure: failure_label(event.outcome),
}),
);
}
fn breaker(&self, transition: &Transition<State>) {
self.report(
self.context.clone(),
self.record_breaker.clone(),
ReportArgs::Breaker(state_label(&transition.to)),
);
}
}
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
use litellm_host_python::PythonContext;
use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents};
use pyo3::{
prelude::*,
types::{PyCFunction, PyModule, PyString},
};
use rstest::rstest;
use tokio::runtime::Handle;
use super::{EventLoop, RedisCall, ServiceLoggingEvents};
const RECORDER: &str = "
import contextvars, time
target = contextvars.ContextVar('target', default=None)
seen = []
def record_redis_call(event_loop, call_type, start_time, end_time, failure):
just_now = abs(time.time() - end_time) < 60
seen.append((event_loop, call_type, round(end_time - start_time, 3), just_now, failure, target.get()))
def record_redis_breaker(state):
seen.append(('breaker', state))
";
fn stand_in_loop(py: Python<'_>, name: &str) -> EventLoop {
EventLoop(Arc::new(PyString::new(py, name).into_any().unbind()))
}
fn recorder(name: &str, client_loop: Option<&str>) -> (ServiceLoggingEvents, Py<PyModule>) {
Python::initialize();
Python::attach(|py| {
let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap();
let module =
PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name)
.unwrap();
let events = ServiceLoggingEvents::from_module(
py,
&module,
Handle::current(),
client_loop.map(|client_loop| stand_in_loop(py, client_loop)),
)
.unwrap();
(events, module.unbind())
})
}
async fn seen(module: &Py<PyModule>, count: usize) -> String {
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
loop {
let seen: (usize, String) = Python::attach(|py| {
let seen = module.bind(py).getattr("seen").unwrap();
(seen.len().unwrap(), seen.repr().unwrap().to_string())
});
if seen.0 >= count || tokio::time::Instant::now() > deadline {
return seen.1;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
fn scoped_call(namespace: &Py<PyModule>, caller_loop: Option<&str>) -> RedisCall {
Python::attach(|py| {
namespace
.bind(py)
.getattr("target")
.unwrap()
.call_method1("set", ("spend_counter",))
.unwrap();
RedisCall {
context: PythonContext::capture(py).unwrap(),
event_loop: caller_loop.map(|caller_loop| stand_in_loop(py, caller_loop)),
call_type: "async_increment",
}
})
}
#[rstest]
#[case::on_the_callers_loop(Some("caller loop"), "caller loop")]
#[case::on_the_clients_loop_when_the_caller_had_none(None, "client loop")]
#[tokio::test(flavor = "multi_thread")]
async fn a_scoped_call_reports_the_store_labels_and_its_times_in_the_callers_context(
#[case] caller_loop: Option<&str>,
#[case] expected_loop: &str,
) {
let (events, namespace) = recorder(
expected_loop.replace(' ', "_").as_str(),
Some("client loop"),
);
let call = scoped_call(&namespace, caller_loop);
let event = CallEvent {
operation: Operation::Command("INCRBYFLOAT"),
duration: Duration::from_millis(250),
outcome: Outcome::Failed(FailureKind::Timeout),
};
call.scope(async { events.call(&event) }).await;
assert_eq!(
seen(&namespace, 1).await,
format!(
"[('{expected_loop}', 'async_increment', 0.25, True, 'timeout', 'spend_counter')]"
)
);
}
#[rstest]
#[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")]
#[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")]
#[tokio::test(flavor = "multi_thread")]
async fn breaker_transitions_report_python_state_names(
#[case] to: litellm_circuit_breaker::State,
#[case] expected: &str,
) {
let (events, namespace) = recorder(expected, None);
events.breaker(&litellm_circuit_breaker::Transition {
from: litellm_circuit_breaker::State::HalfOpen,
to,
generation: 1,
});
assert_eq!(
seen(&namespace, 1).await,
format!("[('breaker', '{expected}')]")
);
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn a_transition_on_a_thread_without_a_runtime_is_still_reported() {
let (events, namespace) = recorder("plain_thread", None);
std::thread::spawn(move || {
events.breaker(&litellm_circuit_breaker::Transition {
from: litellm_circuit_breaker::State::HalfOpen,
to: litellm_circuit_breaker::State::HalfOpen,
generation: 1,
});
})
.join()
.unwrap();
assert_eq!(seen(&namespace, 1).await, "[('breaker', 'half_open')]");
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn an_unscoped_call_falls_back_to_the_command_name_and_the_clients_loop() {
let (events, namespace) = recorder("unscoped", None);
events.call(&CallEvent {
operation: Operation::Pipeline { ops: 3 },
duration: Duration::from_millis(5),
outcome: Outcome::Succeeded,
});
assert_eq!(
seen(&namespace, 1).await,
"[(None, 'pipeline', 0.005, True, None, None)]"
);
}
#[rstest]
fn running_is_the_loop_the_caller_runs_on_and_none_outside_one() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::from_code(
py,
c"import asyncio\nasync def probe_inside(probe):\n return probe() is asyncio.get_running_loop()\n",
c"event_loop_probe",
c"event_loop_probe",
)
.unwrap();
let probe = PyCFunction::new_closure(py, None, None, |_, _| {
Python::attach(|py| -> PyResult<Option<Py<PyAny>>> {
Ok(EventLoop::running(py).map(|event_loop| event_loop.0.clone_ref(py)))
})
})
.unwrap();
let coroutine = module
.getattr("probe_inside")
.unwrap()
.call1((probe,))
.unwrap();
let inside: bool = py
.import("asyncio")
.unwrap()
.call_method1("run", (coroutine,))
.unwrap()
.extract()
.unwrap();
assert!(inside);
assert!(EventLoop::running(py).is_none());
});
}
}

View file

@ -0,0 +1,15 @@
use std::sync::LazyLock;
use litellm_redis_connect::RedisHosts;
pub(crate) mod breaker;
pub(crate) mod client;
mod events;
pub(crate) mod store;
static HOSTS: LazyLock<RedisHosts> = LazyLock::new(RedisHosts::default);
/// The Redis servers this process talks to, shared by every `RedisCache` and native store in it.
pub(crate) fn hosts() -> &'static RedisHosts {
&HOSTS
}

View file

@ -0,0 +1,265 @@
use std::sync::Arc;
use litellm_python_compat::Value;
use litellm_redis_store::auth_cache::{AuthCache, Settings};
use pyo3::prelude::*;
use super::{
NativePending, PyNumber, kind_error, pending, redis_async, redis_sync, remote,
request::{NativeRequestRedis, NativeSend},
ttl,
value::PyValue,
};
use crate::redis::client::NativeRedisClient;
fn values(values: Vec<Option<Value>>) -> Vec<Option<PyValue>> {
values.into_iter().map(|value| value.map(PyValue)).collect()
}
/// `UserApiKeyCache` with `enable_redis_auth_cache` on. The sync methods back `DualCache`'s sync
/// API and block with the GIL released, like redis-py's sync client.
#[pyclass(frozen, name = "NativeAuthCache")]
pub(crate) struct NativeAuthCache(Arc<AuthCache>);
#[pymethods]
impl NativeAuthCache {
#[new]
#[pyo3(signature = (client, *, memory_entries, memory_ttl_seconds = None))]
fn new(
client: Option<PyRef<'_, NativeRedisClient>>,
memory_entries: usize,
memory_ttl_seconds: Option<f64>,
) -> PyResult<Self> {
Ok(Self(Arc::new(AuthCache::new(
client.map(|client| remote(&client)),
Settings {
memory_entries,
memory_ttl: ttl(memory_ttl_seconds)?,
},
))))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
/// `request` answers a key its pre-call reads already found absent without another `GET`.
#[pyo3(signature = (key, request = None))]
fn get<'py>(
&self,
py: Python<'py>,
key: String,
request: Option<&NativeRequestRedis>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let request = request.map(|request| request.0.clone());
redis_async(py, "async_get_cache", kind_error, async move {
store
.get(request.as_ref(), &key)
.await
.map(|value| value.map(PyValue))
})
}
fn get_sync(&self, py: Python<'_>, key: String) -> PyResult<Option<PyValue>> {
let store = self.0.clone();
redis_sync(py, "get_cache", kind_error, async move {
store.get(None, &key).await.map(|value| value.map(PyValue))
})
}
/// `update_cache_ttl` and `update_in_memory_max_size`: both partitions keep what they hold.
#[pyo3(signature = (*, memory_entries, memory_ttl_seconds = None))]
fn reconfigure(&self, memory_entries: usize, memory_ttl_seconds: Option<f64>) -> PyResult<()> {
self.0.reconfigure(Settings {
memory_entries,
memory_ttl: ttl(memory_ttl_seconds)?,
});
Ok(())
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set<'py>(
&self,
py: Python<'py>,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&key, value.0, ttl).await
})
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set_sync(
&self,
py: Python<'_>,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_sync(py, "set_cache", kind_error, async move {
store.set_str(&key, value.0, ttl).await
})
}
/// Memory now and the `SET` on the batch `send` picks.
#[pyo3(signature = (key, value, ttl_seconds, send))]
fn declare_set(
&self,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
send: &NativeSend,
) -> PyResult<NativePending> {
let written = self
.0
.declare_set(send.send(), &key, value.0, ttl(ttl_seconds)?)
.map_err(kind_error)?;
Ok(pending(
"async_set_cache",
kind_error,
written,
std::convert::identity,
))
}
fn declare_delete(&self, key: String, send: &NativeSend) -> PyResult<NativePending> {
let deleted = self
.0
.declare_delete(send.send(), &key)
.map_err(kind_error)?;
Ok(pending(
"async_delete_cache",
kind_error,
deleted,
std::convert::identity,
))
}
#[pyo3(signature = (entries, ttl_seconds = None))]
fn set_many<'py>(
&self,
py: Python<'py>,
entries: Vec<(String, PyValue)>,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let entries = entries
.into_iter()
.map(|(key, value)| (key, value.0))
.collect();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache_pipeline", kind_error, async move {
store.set_many(entries, ttl).await
})
}
/// On a request batch it returns before anything is sent, and `Now` is a pipeline of its own.
fn write_back<'py>(
&self,
py: Python<'py>,
entries: Vec<(String, PyValue, Option<f64>)>,
send: &NativeSend,
) -> PyResult<Bound<'py, PyAny>> {
let entries = entries
.into_iter()
.map(|(key, value, seconds)| Ok((key, value.0, ttl(seconds)?)))
.collect::<PyResult<_>>()?;
let written = self
.0
.write_back(send.send(), entries)
.map_err(kind_error)?;
redis_async(
py,
"async_set_cache_pipeline_with_ttls",
kind_error,
written,
)
}
fn delete<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_delete_cache", kind_error, async move {
store.delete(&keys).await
})
}
fn delete_sync(&self, py: Python<'_>, keys: Vec<String>) -> PyResult<()> {
let store = self.0.clone();
redis_sync(py, "delete_cache", kind_error, async move {
store.delete(&keys).await
})
}
/// One `MGET` of keys the caller found missing in memory, on the batch `send` picks; `None`
/// without Redis.
fn get_many(&self, keys: Vec<String>, send: &NativeSend) -> Option<NativePending> {
let read = self.0.get_many(send.send(), &keys)?;
Some(pending("async_batch_get_cache", kind_error, read, values))
}
fn increment_if_present<'py>(
&self,
py: Python<'py>,
key: String,
amount: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment_if_present", kind_error, async move {
store.increment_if_present(&key, amount.0).await
})
}
fn get_local(&self, key: &str) -> PyResult<Option<PyValue>> {
Ok(self.0.get_local(key).map_err(kind_error)?.map(PyValue))
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set_local(&self, key: &str, value: PyValue, ttl_seconds: Option<f64>) -> PyResult<()> {
self.0
.set_local(key, value.0, ttl(ttl_seconds)?)
.map_err(kind_error)
}
fn delete_local(&self, key: &str) -> PyResult<()> {
self.0.delete_local(key).map_err(kind_error)
}
fn flush_local(&self) -> PyResult<()> {
self.0.flush_local().map_err(kind_error)
}
fn len_local(&self) -> PyResult<PyLocalSizes> {
let (objects, key_objects) = self.0.len_local().map_err(kind_error)?;
Ok(PyLocalSizes {
objects,
key_objects,
})
}
}
/// Entries the management-object and the user-key partitions hold in memory.
#[pyclass(frozen, eq, get_all, name = "LocalSizes")]
#[derive(PartialEq)]
pub(crate) struct PyLocalSizes {
objects: usize,
key_objects: usize,
}
#[pymethods]
impl PyLocalSizes {
#[new]
fn new(objects: usize, key_objects: usize) -> Self {
Self {
objects,
key_objects,
}
}
}

View file

@ -0,0 +1,271 @@
pub(crate) mod auth_cache;
mod policy;
pub(crate) mod rate_limits;
pub(crate) mod request;
pub(crate) mod router_cooldowns;
pub(crate) mod router_usage;
pub(crate) mod sensitive_route_pins;
pub(crate) mod spend_counters;
mod value;
use value::PyValue;
use std::{
pin::Pin,
sync::{Mutex, PoisonError},
time::Duration,
};
use litellm_host_python::PythonContext;
use litellm_python_compat::number::{SecondsError, seconds};
use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote};
use pyo3::{
IntoPyObjectExt,
exceptions::{PyOverflowError, PyRuntimeError, PyTypeError, PyValueError},
prelude::*,
types::{PyBool, PyFloat, PyInt},
};
use crate::execution::{run_async_value, run_sync_value};
use crate::redis::{
client::NativeRedisClient,
events::{EventLoop, RedisCall},
};
/// The labels `RedisCache` reports for this operation, in the caller's context so its
/// `service_target` reaches the service event, dispatched on the loop the caller runs on.
fn call(py: Python<'_>, call_type: &'static str) -> PyResult<RedisCall> {
Ok(RedisCall {
context: PythonContext::capture(py)?,
event_loop: EventLoop::running(py),
call_type,
})
}
/// Runs one Redis operation as a Python awaitable, reported to `ServiceLogging` as `call_type`.
fn redis_async<'py, T, E: 'static>(
py: Python<'py>,
call_type: &'static str,
error: fn(E) -> PyErr,
operation: impl Future<Output = Result<T, E>> + Send + 'static,
) -> PyResult<Bound<'py, PyAny>>
where
T: for<'a> IntoPyObject<'a> + Send + 'static,
{
let scope = call(py, call_type)?;
run_async_value(
py,
async move { scope.scope(operation).await.map_err(error) },
)
}
/// Runs one Redis operation with the GIL released, for the router's sync paths.
fn redis_sync<T, E: 'static>(
py: Python<'_>,
call_type: &'static str,
error: fn(E) -> PyErr,
operation: impl Future<Output = Result<T, E>> + Send + 'static,
) -> PyResult<T>
where
T: Send + 'static,
{
let scope = call(py, call_type)?;
run_sync_value(
py,
async move { scope.scope(operation).await.map_err(error) },
)
}
/// A TTL in seconds, raising what `int(seconds)` raises for infinity and NaN.
fn duration(value: f64) -> PyResult<Duration> {
seconds(value).map_err(|error| match error {
SecondsError::Overflow => PyOverflowError::new_err(error.to_string()),
SecondsError::NotANumber => PyValueError::new_err(error.to_string()),
})
}
fn ttl(value: Option<f64>) -> PyResult<Option<Duration>> {
value.map(duration).transpose()
}
fn remote(client: &NativeRedisClient) -> Remote {
Remote {
client: client.client.clone(),
default_ttl: client.default_ttl,
}
}
type Resolve = Pin<Box<dyn Future<Output = PyResult<Py<PyAny>>> + Send>>;
/// A Redis call declared on a request's batch. `resolve` flushes that batch if it has not gone
/// out and answers the call, once.
#[pyclass(frozen, name = "NativePending")]
pub(crate) struct NativePending {
resolve: Mutex<Option<Resolve>>,
call_type: &'static str,
}
impl NativePending {
pub(crate) fn new<T, E: 'static>(
call_type: &'static str,
error: fn(E) -> PyErr,
reply: impl Future<Output = Result<T, E>> + Send + 'static,
) -> Self
where
T: for<'a> IntoPyObject<'a> + Send + 'static,
{
let resolve = async move {
let reply = reply.await.map_err(error)?;
Python::attach(|py| reply.into_py_any(py))
};
Self {
resolve: Mutex::new(Some(Box::pin(resolve))),
call_type,
}
}
}
#[pymethods]
impl NativePending {
fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let resolve = self
.resolve
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take()
.ok_or_else(|| PyRuntimeError::new_err("this Redis call was already resolved"))?;
redis_async(py, self.call_type, std::convert::identity, resolve)
}
}
pyo3::create_exception!(
_native,
NativeRedisStoreError,
pyo3::exceptions::PyException,
"A native Redis store operation failed. Args are (kind, message), kind spelled as litellm-redis-store's ErrorKind."
);
/// What every native store raises, its kind telling a caller which failures it may drop.
fn kind_error(error: Error) -> PyErr {
let kind: &'static str = error.kind().into();
NativeRedisStoreError::new_err((kind, litellm_error::Report::new(&error).to_string()))
}
/// A store op declared on the batch the caller's `NativeSend` picked, answered as `reply` shapes it.
fn pending<T, R>(
call_type: &'static str,
error: fn(Error) -> PyErr,
declared: Pending<T>,
reply: fn(T) -> R,
) -> NativePending
where
T: Send + 'static,
R: for<'a> IntoPyObject<'a> + Send + 'static,
{
NativePending::new(call_type, error, async move { declared.await.map(reply) })
}
/// A Python `int` or `float` across the bridge, keeping which one it is.
#[derive(Clone, Copy, PartialEq)]
pub(crate) struct PyNumber(Number);
impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
if value.is_instance_of::<PyInt>() && !value.is_instance_of::<PyBool>() {
return Ok(Self(Number::Int(value.extract()?)));
}
if value.is_instance_of::<PyFloat>() {
return Ok(Self(Number::Float(value.extract()?)));
}
Err(PyTypeError::new_err("expected an int or a float"))
}
}
impl<'py> IntoPyObject<'py> for PyNumber {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
Ok(match self.0 {
Number::Int(value) => value.into_pyobject(py)?.into_any(),
Number::Float(value) => value.into_pyobject(py)?.into_any(),
})
}
}
/// A store enum read from its wire spelling, so the bridge names no variant.
pub(crate) struct Spelled<T>(pub(crate) T);
impl<'a, 'py, T: std::str::FromStr> FromPyObject<'a, 'py> for Spelled<T> {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let spelling: String = value.extract()?;
spelling
.parse()
.map(Self)
.map_err(|_| PyValueError::new_err(format!("unknown spelling {spelling:?}")))
}
}
/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue`
/// so the bridge never names a field. A shape the store refuses is `ValueError`.
pub(crate) struct PyData<T>(pub(crate) T);
impl<'a, 'py, T: FromValue> FromPyObject<'a, 'py> for PyData<T> {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let PyValue(value) = value.extract()?;
T::from_value(&value)
.map(Self)
.map_err(|error| PyValueError::new_err(error.to_string()))
}
}
impl<'py, T: IntoValue> IntoPyObject<'py> for PyData<T> {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
PyValue(self.0.into_value()).into_pyobject(py)
}
}
fn unwrap_all<T>(values: Vec<PyData<T>>) -> Vec<T> {
values.into_iter().map(|value| value.0).collect()
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use pyo3::prelude::*;
use rstest::rstest;
use super::ttl;
#[rstest]
#[case::infinity(f64::INFINITY, "OverflowError")]
#[case::nan(f64::NAN, "ValueError")]
fn a_ttl_int_refuses_raises_what_int_raises(#[case] seconds: f64, #[case] raised: &str) {
Python::initialize();
Python::attach(|py| {
let error = ttl(Some(seconds)).unwrap_err();
assert_eq!(error.get_type(py).name().unwrap().to_string(), raised);
});
}
#[rstest]
#[case::fraction(Some(1.5), Some(Duration::from_millis(1500)))]
#[case::negative_clamps_to_zero(Some(-1.0), Some(Duration::ZERO))]
#[case::absent(None, None)]
fn a_finite_ttl_converts(#[case] seconds: Option<f64>, #[case] expected: Option<Duration>) {
assert_eq!(ttl(seconds).unwrap(), expected);
}
}

View file

@ -0,0 +1,151 @@
use std::time::Duration;
use litellm_redis_store::TierPolicy;
use pyo3::prelude::*;
use crate::{
coercion::{Field, FieldSpec, ProjectionError},
python_settings::{PythonSettings, Snapshot},
};
const COOLDOWN_REFRESH: FieldSpec<Duration> =
FieldSpec::new("cooldown_redis_read_interval", seconds);
const USAGE_REFRESH: FieldSpec<Duration> = FieldSpec::new("usage_redis_read_interval", seconds);
const TRACKED_READS: FieldSpec<usize> = FieldSpec::new("max_redis_batch_cache_size", count);
const MEMORY_TTL: FieldSpec<Option<Duration>> =
FieldSpec::new("default_in_memory_ttl", optional_seconds);
/// The tier values this Python process runs with, from its environment and `litellm` globals.
#[derive(Debug)]
pub(crate) struct RouterTierPolicies {
pub(crate) cooldowns: TierPolicy,
pub(crate) usage: TierPolicy,
}
pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult<RouterTierPolicies> {
Ok(project(&PythonSettings::RouterTiers.read(py)?)?)
}
fn project(snapshot: &Snapshot<'_>) -> Result<RouterTierPolicies, ProjectionError> {
let tracked_reads = snapshot.read(&TRACKED_READS)?;
let default_in_memory_ttl = snapshot.read(&MEMORY_TTL)?;
Ok(RouterTierPolicies {
cooldowns: TierPolicy {
refresh_after: snapshot.read(&COOLDOWN_REFRESH)?,
tracked_reads,
default_in_memory_ttl,
},
usage: TierPolicy {
refresh_after: snapshot.read(&USAGE_REFRESH)?,
tracked_reads,
default_in_memory_ttl,
},
})
}
fn seconds(field: &Field<'_>) -> Result<Duration, ProjectionError> {
match field.value().extract::<f64>() {
Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => {
Ok(Duration::from_secs_f64(seconds))
}
_ => Err(field.invalid("a non-negative number of seconds")),
}
}
fn optional_seconds(field: &Field<'_>) -> Result<Option<Duration>, ProjectionError> {
match field.value().is_none() {
true => Ok(None),
false => seconds(field).map(Some),
}
}
fn count(field: &Field<'_>) -> Result<usize, ProjectionError> {
field
.value()
.extract::<usize>()
.map_err(|_| field.invalid("a non-negative integer"))
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use pyo3::{prelude::*, types::PyDict};
use super::project;
use crate::python_settings::PythonSettings;
fn projected(py: Python<'_>, fields: &str) -> Result<super::RouterTierPolicies, String> {
let locals = PyDict::new(py);
let source = std::ffi::CString::new(format!(
"import types\nsettings = types.SimpleNamespace({fields})"
))
.unwrap();
py.run(&source, Some(&locals), Some(&locals)).unwrap();
let snapshot =
PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap());
project(&snapshot).map_err(|error| PyErr::from(error).to_string())
}
#[test]
fn the_host_values_are_each_familys_refresh_and_tracked_reads() {
Python::initialize();
Python::attach(|py| {
let policies = projected(
py,
"cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42, default_in_memory_ttl=None",
)
.unwrap();
assert_eq!(
policies.cooldowns.refresh_after,
Duration::from_millis(2500)
);
assert_eq!(policies.usage.refresh_after, Duration::from_secs(30));
assert_eq!(policies.cooldowns.tracked_reads, 42);
assert_eq!(policies.usage.tracked_reads, 42);
assert_eq!(policies.cooldowns.default_in_memory_ttl, None);
});
}
#[test]
fn a_set_default_in_memory_ttl_reaches_both_familys_backfill() {
Python::initialize();
Python::attach(|py| {
let policies = projected(
py,
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=7.5",
)
.unwrap();
assert_eq!(
policies.cooldowns.default_in_memory_ttl,
Some(Duration::from_millis(7500))
);
assert_eq!(
policies.usage.default_in_memory_ttl,
Some(Duration::from_millis(7500))
);
});
}
#[test]
fn a_value_python_would_not_run_with_is_a_configuration_error() {
Python::initialize();
Python::attach(|py| {
let negative = projected(
py,
"cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=None",
)
.unwrap_err();
let text = projected(
py,
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many', default_in_memory_ttl=None",
)
.unwrap_err();
assert!(
negative.contains("cooldown_redis_read_interval"),
"{negative}"
);
assert!(text.contains("max_redis_batch_cache_size"), "{text}");
});
}
}

View file

@ -0,0 +1,577 @@
use std::{sync::Arc, time::Duration};
use indexmap::IndexMap;
use litellm_redis_store::{
Pending,
rate_limits::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge, GuardedIncrement,
GuardedReply, Increment, Kind, Local, RateLimits, Scope, WindowCount,
},
};
use pyo3::{exceptions::PyValueError, prelude::*};
use super::{
NativePending, PyNumber, Spelled, kind_error, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl,
};
use crate::redis::client::NativeRedisClient;
/// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`;
/// every Redis call is a future.
#[pyclass(frozen, name = "NativeRateLimits")]
pub(crate) struct NativeRateLimits(Arc<RateLimits>);
/// An `int` or `float` counter, a `str` window start, or a `{slot_id: acquired_at}` registry.
#[derive(FromPyObject)]
enum PyLocal {
Number(PyNumber),
Text(String),
Slots(IndexMap<String, f64>),
}
impl From<PyLocal> for Local {
fn from(local: PyLocal) -> Self {
match local {
PyLocal::Number(number) => Self::Number(number.0),
PyLocal::Text(text) => Self::Text(text),
PyLocal::Slots(slots) => Self::Slots(slots.into_iter().collect()),
}
}
}
/// What the memory tier holds, back as the Python value it was written as.
struct LocalValue(Local);
impl<'py> IntoPyObject<'py> for LocalValue {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
Ok(match self.0 {
Local::Number(number) => PyNumber(number).into_pyobject(py)?,
Local::Text(text) => text.into_pyobject(py)?.into_any(),
Local::Slots(slots) => slots
.into_iter()
.collect::<IndexMap<_, _>>()
.into_pyobject(py)?
.into_any(),
})
}
}
/// `litellm.redis.stores.rate_limits.RateLimitScope`.
#[derive(FromPyObject)]
struct PyScope {
descriptor_key: String,
descriptor_value: String,
}
impl From<PyScope> for Scope {
fn from(scope: PyScope) -> Self {
Self {
descriptor: scope.descriptor_key,
value: scope.descriptor_value,
}
}
}
/// `RateLimitCounter`.
#[derive(FromPyObject)]
struct PyCounter {
scope: PyScope,
kind: Spelled<Kind>,
}
impl From<PyCounter> for Counter {
fn from(counter: PyCounter) -> Self {
Self {
scope: counter.scope.into(),
kind: counter.kind.0,
}
}
}
/// `RateLimitWindow`, `RateLimitCounter` or `RateLimitGauge`, told apart by the key suffix each names.
#[derive(FromPyObject)]
struct PyEntry {
suffix: String,
scope: PyScope,
}
impl TryFrom<PyEntry> for Entry {
type Error = PyErr;
fn try_from(entry: PyEntry) -> PyResult<Self> {
Self::new(&entry.suffix, entry.scope.into())
.map_err(|_| PyValueError::new_err(format!("unknown key suffix {:?}", entry.suffix)))
}
}
/// `CounterIncrement`.
#[derive(FromPyObject)]
struct PyIncrement {
counter: PyCounter,
amount: PyNumber,
ttl: Option<u64>,
}
impl From<PyIncrement> for Increment {
fn from(increment: PyIncrement) -> Self {
Self {
counter: increment.counter.into(),
amount: increment.amount.0,
ttl: increment.ttl.map(Duration::from_secs),
}
}
}
/// `AtomicCounter`.
#[derive(FromPyObject)]
struct PyAtomicCounter {
counter: PyCounter,
limit: i64,
increment: i64,
ttl: i64,
window_size: i64,
}
impl From<PyAtomicCounter> for AtomicCounter {
fn from(counter: PyAtomicCounter) -> Self {
Self {
counter: counter.counter.into(),
limit: counter.limit,
increment: counter.increment,
ttl_seconds: counter.ttl,
window_size_seconds: counter.window_size,
}
}
}
fn all<P, T: From<P>>(records: Vec<P>) -> Vec<T> {
records.into_iter().map(T::from).collect()
}
#[pyclass(frozen, eq, get_all, name = "WindowCount")]
#[derive(PartialEq)]
pub(crate) struct PyWindowCount {
window_start: String,
count: i64,
}
#[pymethods]
impl PyWindowCount {
#[new]
fn new(window_start: String, count: i64) -> Self {
Self {
window_start,
count,
}
}
}
/// Per counter, the count after the increment and the start of the window it landed in.
#[pyclass(frozen, eq, get_all, name = "AtomicApplied")]
#[derive(PartialEq)]
pub(crate) struct PyAtomicApplied {
counters: Vec<(i64, i64)>,
}
#[pymethods]
impl PyAtomicApplied {
#[new]
fn new(counters: Vec<(i64, i64)>) -> Self {
Self { counters }
}
}
#[pyclass(frozen, eq, get_all, name = "AtomicOverLimit")]
#[derive(PartialEq)]
pub(crate) struct PyAtomicOverLimit {
index: usize,
current: i64,
limit: i64,
}
#[pymethods]
impl PyAtomicOverLimit {
#[new]
fn new(index: usize, current: i64, limit: i64) -> Self {
Self {
index,
current,
limit,
}
}
}
#[derive(IntoPyObject)]
enum PyAtomicReply {
Applied(PyAtomicApplied),
OverLimit(PyAtomicOverLimit),
}
#[pyclass(frozen, eq, get_all, name = "GuardedReply")]
#[derive(PartialEq)]
pub(crate) struct PyGuardedReply {
applied: bool,
count: i64,
}
#[pymethods]
impl PyGuardedReply {
#[new]
fn new(applied: bool, count: i64) -> Self {
Self { applied, count }
}
}
#[pyclass(frozen, eq, get_all, name = "GaugesAcquired")]
#[derive(PartialEq)]
pub(crate) struct PyGaugesAcquired {
in_flight: Vec<i64>,
}
#[pymethods]
impl PyGaugesAcquired {
#[new]
fn new(in_flight: Vec<i64>) -> Self {
Self { in_flight }
}
}
#[pyclass(frozen, eq, get_all, name = "GaugeOverLimit")]
#[derive(PartialEq)]
pub(crate) struct PyGaugeOverLimit {
index: usize,
in_flight: i64,
limit: i64,
}
#[pymethods]
impl PyGaugeOverLimit {
#[new]
fn new(index: usize, in_flight: i64, limit: i64) -> Self {
Self {
index,
in_flight,
limit,
}
}
}
#[derive(IntoPyObject)]
enum PyAcquireReply {
Acquired(PyGaugesAcquired),
OverLimit(PyGaugeOverLimit),
}
fn window_counts(counts: Vec<WindowCount>) -> Vec<PyWindowCount> {
counts
.into_iter()
.map(
|WindowCount {
window_start,
count,
}| PyWindowCount {
window_start,
count,
},
)
.collect()
}
fn atomic(reply: AtomicReply) -> PyAtomicReply {
match reply {
AtomicReply::Applied(applied) => PyAtomicReply::Applied(PyAtomicApplied {
counters: applied
.into_iter()
.map(
|Applied {
count,
window_start,
}| (count, window_start),
)
.collect(),
}),
AtomicReply::OverLimit {
index,
current,
limit,
} => PyAtomicReply::OverLimit(PyAtomicOverLimit {
index,
current,
limit,
}),
}
}
fn guarded(GuardedReply { applied, count }: GuardedReply) -> PyGuardedReply {
PyGuardedReply { applied, count }
}
fn acquire(reply: AcquireReply) -> PyAcquireReply {
match reply {
AcquireReply::Acquired(in_flight) => {
PyAcquireReply::Acquired(PyGaugesAcquired { in_flight })
}
AcquireReply::OverLimit {
index,
in_flight,
limit,
} => PyAcquireReply::OverLimit(PyGaugeOverLimit {
index,
in_flight,
limit,
}),
}
}
/// One script call declared on the batch the caller's `NativeSend` picked; `None` without Redis.
fn script<T, R>(
call_type: &'static str,
declared: Option<Pending<T>>,
reply: fn(T) -> R,
) -> Option<NativePending>
where
T: Send + 'static,
R: for<'a> IntoPyObject<'a> + Send + 'static,
{
Some(super::pending(call_type, kind_error, declared?, reply))
}
#[pymethods]
impl NativeRateLimits {
#[new]
#[pyo3(signature = (client, backfill_ttl_seconds = None))]
fn new(
client: Option<PyRef<'_, NativeRedisClient>>,
backfill_ttl_seconds: Option<f64>,
) -> PyResult<Self> {
Ok(Self(Arc::new(RateLimits::new(
client.map(|client| remote(&client)),
ttl(backfill_ttl_seconds)?,
))))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
fn slot_groups(&self, scopes: Vec<PyScope>, force: bool) -> Vec<Vec<usize>> {
self.0.slot_groups(&all(scopes), force)
}
fn get_local(&self, entry: PyEntry) -> PyResult<Option<LocalValue>> {
Ok(self
.0
.get_local(&entry.try_into()?)
.map_err(kind_error)?
.map(LocalValue))
}
#[pyo3(signature = (entry, value, ttl_seconds))]
fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option<f64>) -> PyResult<()> {
self.0
.set_local(&entry.try_into()?, value.into(), ttl(ttl_seconds)?)
.map_err(kind_error)
}
/// `request` answers a key its pre-call reads already found absent without another `GET`.
#[pyo3(signature = (entry, *, backfill_ttl_seconds = None, request = None))]
fn get<'py>(
&self,
py: Python<'py>,
entry: PyEntry,
backfill_ttl_seconds: Option<f64>,
request: Option<&NativeRequestRedis>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let entry = Entry::try_from(entry)?;
let backfill_ttl = ttl(backfill_ttl_seconds)?;
let request = request.map(|request| request.0.clone());
redis_async(py, "async_get_cache", kind_error, async move {
store
.get(request.as_ref(), &entry, backfill_ttl)
.await
.map(|local| local.map(LocalValue))
})
}
fn sliding_window(
&self,
counters: Vec<PyCounter>,
now: i64,
window_size: i64,
send: &NativeSend,
) -> Option<NativePending> {
script(
"batch_rate_limiter_script",
self.0
.sliding_window(send.send(), &all(counters), now, window_size),
window_counts,
)
}
fn read_counters(&self, counters: Vec<PyCounter>, send: &NativeSend) -> Option<NativePending> {
script(
"batch_counter_read_script",
self.0.read_counters(send.send(), &all(counters)),
std::convert::identity,
)
}
fn check_and_increment(
&self,
counters: Vec<PyAtomicCounter>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"check_and_increment_by_n_script",
self.0.check_and_increment(send.send(), &all(counters)),
atomic,
)
}
fn guarded_increment(
&self,
counter: PyCounter,
expected_window_start: String,
amount: PyNumber,
ttl_seconds: Option<u64>,
send: &NativeSend,
) -> Option<NativePending> {
let increment = GuardedIncrement {
counter: counter.into(),
expected_window_start,
amount: amount.0,
ttl: ttl_seconds.map(Duration::from_secs),
};
script(
"window_guarded_token_increment_script",
self.0.guarded_increment(send.send(), &increment),
guarded,
)
}
fn parallel_acquire(
&self,
gauges: Vec<(PyScope, i64)>,
slot_ttl_seconds: u64,
slot_id: &str,
send: &NativeSend,
) -> Option<NativePending> {
let gauges: Vec<Gauge> = gauges
.into_iter()
.map(|(scope, limit)| Gauge {
scope: scope.into(),
limit,
})
.collect();
script(
"parallel_acquire_script",
self.0.parallel_acquire(
send.send(),
&gauges,
Duration::from_secs(slot_ttl_seconds),
slot_id,
),
acquire,
)
}
fn parallel_release(
&self,
gauges: Vec<PyScope>,
slot_id: &str,
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_release_script",
self.0.parallel_release(send.send(), &all(gauges), slot_id),
std::convert::identity,
)
}
fn parallel_count(
&self,
gauges: Vec<PyScope>,
slot_ttl_seconds: u64,
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_count_script",
self.0.parallel_count(
send.send(),
&all(gauges),
Duration::from_secs(slot_ttl_seconds),
),
std::convert::identity,
)
}
fn token_increment(
&self,
increments: Vec<PyIncrement>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"token_increment_script",
self.0.token_increment(send.send(), &all(increments)),
std::convert::identity,
)
}
/// `False` without Redis.
fn token_increment_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<PyIncrement>,
) -> bool {
self.0
.token_increment_post_call(&request.0, all(increments))
}
/// `False` without Redis.
fn parallel_release_post_call(
&self,
request: &NativeRequestRedis,
gauges: Vec<PyScope>,
slot_id: &str,
) -> bool {
self.0
.parallel_release_post_call(&request.0, &all(gauges), slot_id)
}
fn increment<'py>(
&self,
py: Python<'py>,
increment: PyIncrement,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increment = Increment::from(increment);
redis_async(py, "async_increment", kind_error, async move {
Ok(store.increment(&increment).await.map(PyNumber))
})
}
fn increment_many<'py>(
&self,
py: Python<'py>,
increments: Vec<PyIncrement>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increments: Vec<Increment> = all(increments);
redis_async(py, "async_increment_pipeline", kind_error, async move {
store.increment_many(&increments).await;
Ok(())
})
}
fn increment_many_post_call(&self, request: &NativeRequestRedis, increments: Vec<PyIncrement>) {
self.0
.increment_many_post_call(&request.0, &all(increments));
}
}

View file

@ -0,0 +1,60 @@
use litellm_redis_store::{RequestRedis, Send};
use pyo3::prelude::*;
use super::{call, duration};
use crate::execution::run_async_value;
/// One request's native Redis batches. Python builds one per request scope and passes it to each
/// request-aware store call.
#[pyclass(frozen, name = "NativeRequestRedis")]
pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis);
#[pymethods]
impl NativeRequestRedis {
#[new]
fn new(py: Python<'_>, post_call_deadline_seconds: f64) -> PyResult<Self> {
let runtime = litellm_host_python::runtime_handle()?;
let deadline = duration(post_call_deadline_seconds)?;
let deadline_tag = call(py, "post_call_redis_batch")?.tag();
Ok(Self(crate::logger::capture(py).scope(|| {
RequestRedis::with_deadline_tag(runtime, deadline, Some(deadline_tag))
})))
}
fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let request = self.0.clone();
let scope = call(py, "request_redis_batch")?;
run_async_value(py, async move {
scope.scope(request.flush_pre()).await;
Ok(())
})
}
fn flush_post<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let request = self.0.clone();
let scope = call(py, "post_call_redis_batch")?;
run_async_value(py, async move {
scope.scope(request.flush_post()).await;
Ok(())
})
}
}
/// The Python face of `Send`: which batch a store op rides. Each variant carries exactly what it
/// needs, so there is no phase string and no phase without a request.
#[pyclass(frozen, name = "NativeSend")]
pub(crate) enum NativeSend {
Now(),
PreCall(Py<NativeRequestRedis>),
PostCall(Py<NativeRequestRedis>),
}
impl NativeSend {
pub(crate) fn send(&self) -> Send<'_> {
match self {
Self::Now() => Send::Now,
Self::PreCall(request) => Send::PreCall(&request.get().0),
Self::PostCall(request) => Send::PostCall(&request.get().0),
}
}
}

View file

@ -0,0 +1,182 @@
use std::{collections::HashMap, sync::Arc};
use litellm_redis_store::{
Send,
router_cooldowns::{Cooldown, RouterCooldowns},
};
use pyo3::{prelude::*, types::PyTuple};
use super::{
NativePending, PyData, PyNumber, duration, kind_error, pending, policy::router_tier_policies,
redis_async, redis_sync, remote, request::NativeSend,
};
use crate::redis::client::NativeRedisClient;
/// One router's cooldown store. The router's sync completion path calls `put`, `active_sync`,
/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py; the
/// `async_` twins return awaitables for async callers.
#[pyclass(frozen, name = "NativeRouterCooldowns")]
pub(crate) struct NativeRouterCooldowns(Arc<RouterCooldowns>);
/// `[(model_id, cooldown), ...]` the way `RedisCache` hands the router its batch read.
struct PyActive(Vec<(String, Cooldown)>);
impl<'py> IntoPyObject<'py> for PyActive {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
active_to_py(py, self.0)?
.into_pyobject(py)
.map(Bound::into_any)
}
}
fn active_to_py(
py: Python<'_>,
active: Vec<(String, Cooldown)>,
) -> PyResult<Vec<Bound<'_, PyTuple>>> {
active
.into_iter()
.map(|(model_id, cooldown)| {
PyTuple::new(
py,
[
model_id.into_pyobject(py)?.into_any(),
PyData(cooldown).into_pyobject(py)?,
],
)
})
.collect()
}
#[pymethods]
impl NativeRouterCooldowns {
#[new]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
let remote = client.map(|client| remote(&client));
let policy = router_tier_policies(py)?.cooldowns;
Ok(Self(Arc::new(RouterCooldowns::new(remote, policy))))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
fn put(
&self,
py: Python<'_>,
model_id: String,
cooldown: PyData<Cooldown>,
ttl_seconds: f64,
) -> PyResult<()> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_sync(py, "set_cache", kind_error, async move {
store.put_str(&model_id, &cooldown.0, ttl).await
})
}
fn increment_allowed_fails(
&self,
py: Python<'_>,
model_id: String,
suffix: Option<String>,
ttl_seconds: f64,
) -> PyResult<i64> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_sync(py, "increment_cache", kind_error, async move {
store
.increment_allowed_fails_int(&model_id, suffix.as_deref(), ttl)
.await
})
}
fn async_put<'py>(
&self,
py: Python<'py>,
model_id: String,
cooldown: PyData<Cooldown>,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.put(&model_id, &cooldown.0, ttl).await
})
}
fn async_increment_allowed_fails<'py>(
&self,
py: Python<'py>,
model_id: String,
suffix: Option<String>,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_increment", kind_error, async move {
store
.increment_allowed_fails(&model_id, suffix.as_deref(), ttl)
.await
})
}
/// The deployments still cooling down, read on the batch `send` picks.
fn active(&self, model_ids: Vec<String>, send: &NativeSend) -> PyResult<NativePending> {
let active = self.0.active(send.send(), &model_ids).map_err(kind_error)?;
Ok(pending(
"async_batch_get_cache",
kind_error,
active,
PyActive,
))
}
/// `active` on its own with the GIL released, for the router's sync completion path.
fn active_sync<'py>(
&self,
py: Python<'py>,
model_ids: Vec<String>,
) -> PyResult<Vec<Bound<'py, PyTuple>>> {
let active = self.0.active(Send::Now, &model_ids).map_err(kind_error)?;
let active = redis_sync(py, "batch_get_cache", kind_error, active)?;
active_to_py(py, active)
}
fn minimum(&self, py: Python<'_>, model_ids: Vec<String>) -> PyResult<Option<PyNumber>> {
let store = self.0.clone();
redis_sync(py, "batch_get_cache", kind_error, async move {
store
.minimum(&model_ids)
.await
.map(|minimum| minimum.map(PyNumber))
})
}
fn delete_local(&self, model_id: &str) -> PyResult<()> {
self.0.delete_local(model_id).map_err(kind_error)
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(kind_error)
}
fn snapshot(&self) -> PyResult<HashMap<String, PyData<Cooldown>>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, cooldown)| (key, PyData(cooldown)))
.collect())
}
}

View file

@ -0,0 +1,182 @@
use std::{collections::HashMap, sync::Arc, time::Duration};
use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage};
use pyo3::prelude::*;
use super::{
PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote,
};
use crate::redis::client::NativeRedisClient;
/// `litellm.redis.stores.router_usage.UsageCounter`.
#[derive(FromPyObject)]
struct PyCounter {
deployment_id: String,
model: Option<String>,
minute: String,
kind: Spelled<Kind>,
}
impl From<PyCounter> for Counter {
fn from(counter: PyCounter) -> Self {
Self {
deployment_id: counter.deployment_id,
model: counter.model,
minute: counter.minute,
kind: counter.kind.0,
}
}
}
fn counters(counters: Vec<PyCounter>) -> Vec<Counter> {
counters.into_iter().map(Counter::from).collect()
}
/// One router's global TPM/RPM counters. The `_sync` methods block with the GIL released, for the
/// Rust router's sync calls, which have no event loop to await on.
#[pyclass(frozen, name = "NativeRouterUsage")]
pub(crate) struct NativeRouterUsage(Arc<RouterUsage>);
#[pymethods]
impl NativeRouterUsage {
#[new]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
let remote = client.map(|client| remote(&client));
let policy = router_tier_policies(py)?.usage;
Ok(Self(Arc::new(RouterUsage::new(remote, policy))))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PyCounter,
value: PyNumber,
ttl_seconds: u64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counter = Counter::from(counter);
redis_async(py, "async_increment", kind_error, async move {
Ok(store
.increment(&counter, value.0, Duration::from_secs(ttl_seconds))
.await
.map(PyNumber))
})
}
/// `_increment_deployment`'s `async_increment_pipeline`, always on its own round trip.
fn increment_deployment<'py>(
&self,
py: Python<'py>,
deployment_id: String,
model: Option<String>,
minute: String,
tokens: PyNumber,
requests: i64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment_pipeline", kind_error, async move {
store
.increment_deployment(
&deployment_id,
model.as_deref(),
&minute,
tokens.0,
requests,
)
.await
})
}
fn get_many<'py>(
&self,
py: Python<'py>,
counters: Vec<PyCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = self::counters(counters);
redis_async(py, "async_batch_get_cache", kind_error, async move {
let read = store.get_many(&counters).await?;
Ok(read
.into_iter()
.map(|count| count.map(PyNumber))
.collect::<Vec<_>>())
})
}
fn increment_sync(
&self,
py: Python<'_>,
counter: PyCounter,
value: PyNumber,
ttl_seconds: u64,
) -> PyResult<Option<PyNumber>> {
let store = self.0.clone();
let counter = Counter::from(counter);
redis_sync(py, "increment_cache", kind_error, async move {
Ok(store
.increment(&counter, value.0, Duration::from_secs(ttl_seconds))
.await
.map(PyNumber))
})
}
fn increment_deployment_sync(
&self,
py: Python<'_>,
deployment_id: String,
model: Option<String>,
minute: String,
tokens: PyNumber,
requests: i64,
) -> PyResult<()> {
let store = self.0.clone();
redis_sync(py, "increment_pipeline", kind_error, async move {
store
.increment_deployment(
&deployment_id,
model.as_deref(),
&minute,
tokens.0,
requests,
)
.await
})
}
fn get_many_sync(
&self,
py: Python<'_>,
counters: Vec<PyCounter>,
) -> PyResult<Vec<Option<PyNumber>>> {
let store = self.0.clone();
let counters = self::counters(counters);
redis_sync(py, "batch_get_cache", kind_error, async move {
let read = store.get_many(&counters).await?;
Ok(read.into_iter().map(|count| count.map(PyNumber)).collect())
})
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(kind_error)
}
fn snapshot(&self) -> PyResult<HashMap<String, PyNumber>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, count)| (key, PyNumber(count)))
.collect())
}
}

View file

@ -0,0 +1,49 @@
use std::sync::Arc;
use litellm_redis_store::sensitive_route_pins::SensitiveRoutePins;
use pyo3::prelude::*;
use super::{duration, kind_error, redis_async, remote};
use crate::redis::client::NativeRedisClient;
#[pyclass(frozen, name = "NativeSensitiveRoutePins")]
pub(crate) struct NativeSensitiveRoutePins(Arc<SensitiveRoutePins>);
#[pymethods]
impl NativeSensitiveRoutePins {
#[new]
fn new(client: Option<PyRef<'_, NativeRedisClient>>) -> Self {
Self(Arc::new(SensitiveRoutePins::new(
client.map(|client| remote(&client)),
)))
}
fn get<'py>(
&self,
py: Python<'py>,
tenant: String,
session_id: String,
default_ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let default_ttl = duration(default_ttl_seconds)?;
redis_async(py, "async_get_cache", kind_error, async move {
store.get(&tenant, &session_id, default_ttl).await
})
}
fn set<'py>(
&self,
py: Python<'py>,
tenant: String,
session_id: String,
model: String,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&tenant, &session_id, &model, ttl).await
})
}
}

View file

@ -0,0 +1,277 @@
use std::{collections::HashMap, sync::Arc};
use litellm_redis_store::{
Number,
spend_counters::{LostSeed, PostCallSettlement, Reading, SpendCounter, SpendCounters},
};
use pyo3::prelude::*;
use super::{
NativePending, PyData, PyNumber, duration, kind_error, pending, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl, unwrap_all,
};
use crate::redis::client::NativeRedisClient;
/// A counter's value and whether Redis answered it, a clean miss included.
#[pyclass(frozen, eq, get_all, name = "SpendRead")]
#[derive(PartialEq)]
pub(crate) struct PySpendRead {
value: Option<PyNumber>,
from_redis: bool,
}
#[pymethods]
impl PySpendRead {
#[new]
fn new(value: Option<PyNumber>, from_redis: bool) -> Self {
Self { value, from_redis }
}
}
/// Whether every increment one `increment_many_post_call` declared has its reply.
#[pyclass(frozen, name = "PostCallSettlement")]
pub(crate) struct PyPostCallSettlement(PostCallSettlement);
#[pymethods]
impl PyPostCallSettlement {
#[getter]
fn settled(&self) -> bool {
self.0.is_settled()
}
}
fn increments(increments: Vec<(PyData<SpendCounter>, PyNumber)>) -> Vec<(SpendCounter, Number)> {
increments
.into_iter()
.map(|(counter, amount)| (counter.0, amount.0))
.collect()
}
fn numbers(values: Option<Vec<Option<Number>>>) -> Option<Vec<Option<PyNumber>>> {
values.map(|values| {
values
.into_iter()
.map(|value| value.map(PyNumber))
.collect()
})
}
/// The proxy's spend counters, Redis first with this worker's memory copy behind it.
#[pyclass(frozen, name = "NativeSpendCounters")]
pub(crate) struct NativeSpendCounters(Arc<SpendCounters>);
#[pymethods]
impl NativeSpendCounters {
#[new]
fn new(client: Option<PyRef<'_, NativeRedisClient>>) -> Self {
Self(Arc::new(SpendCounters::new(
client.map(|client| remote(&client)),
)))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
#[getter]
fn redis_ttl_seconds(&self) -> Option<f64> {
self.0.redis_ttl().map(|ttl| ttl.as_secs_f64())
}
/// The least TTL a Redis write gives a committed spend twin.
fn set_committed_ttl(&self, ttl_seconds: f64) -> PyResult<()> {
self.0.set_committed_ttl(duration(ttl_seconds)?);
Ok(())
}
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
fn read<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_get_cache", kind_error, async move {
let reading = store.read(&counter.0).await?;
Ok(PySpendRead {
value: reading.value().map(PyNumber),
from_redis: matches!(reading, Reading::Redis(_)),
})
})
}
/// One `MGET` on the batch `send` picks; `None` without Redis. It resolves to `None` when the
/// read failed, so each caller reads on its own.
fn read_many(
&self,
counters: Vec<PyData<SpendCounter>>,
send: &NativeSend,
) -> Option<NativePending> {
let read = self.0.read_many(send.send(), &unwrap_all(counters))?;
Some(pending("async_batch_get_cache", kind_error, read, numbers))
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
amount: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment", kind_error, async move {
store.increment(&counter.0, amount.0).await.map(PyNumber)
})
}
fn increment_many<'py>(
&self,
py: Python<'py>,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increments = self::increments(increments);
redis_async(py, "async_increment_pipeline", kind_error, async move {
let totals = store.increment_many(&increments).await?;
Ok(totals.into_iter().map(PyNumber).collect::<Vec<_>>())
})
}
/// Declares the increments on the request's post-call batch; `None` when there is no Redis.
fn increment_many_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> Option<PyPostCallSettlement> {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
.map(PyPostCallSettlement)
}
fn set_max<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_set_max", kind_error, async move {
store
.set_max(&counter.0, spend.0)
.await
.map(|stored| stored.map(PyNumber))
})
}
#[pyo3(signature = (counter, spend, *, increment_if_lost))]
fn seed<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
increment_if_lost: bool,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let lost = if increment_if_lost {
LostSeed::Increment
} else {
LostSeed::Answer
};
redis_async(py, "async_set_cache", kind_error, async move {
store.seed(&counter.0, spend.0, lost).await.map(PyNumber)
})
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&counter.0, spend.0, ttl).await
})
}
fn delete<'py>(
&self,
py: Python<'py>,
counters: Vec<PyData<SpendCounter>>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = unwrap_all(counters);
redis_async(py, "async_delete_cache", kind_error, async move {
store.delete(&counters).await
})
}
fn refresh_ttl<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_refresh_ttl", kind_error, async move {
store.refresh_ttl(&counter.0).await
})
}
fn get_local(&self, counter: PyData<SpendCounter>) -> PyResult<Option<PyNumber>> {
Ok(self
.0
.get_local(&counter.0)
.map_err(kind_error)?
.map(PyNumber))
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set_local(
&self,
counter: PyData<SpendCounter>,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
self.0
.set_local(&counter.0, spend.0, ttl(ttl_seconds)?)
.map_err(kind_error)
}
fn delete_local(&self, counter: PyData<SpendCounter>) -> PyResult<()> {
self.0.delete_local(&counter.0).map_err(kind_error)
}
fn increment_local(
&self,
counter: PyData<SpendCounter>,
amount: PyNumber,
) -> PyResult<PyNumber> {
self.0
.increment_local(&counter.0, amount.0)
.map(PyNumber)
.map_err(kind_error)
}
fn flush_local(&self) -> PyResult<()> {
self.0.flush_local().map_err(kind_error)
}
fn snapshot(&self) -> PyResult<HashMap<String, PyNumber>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, spend)| (key, PyNumber(spend)))
.collect())
}
}

View file

@ -0,0 +1,125 @@
use litellm_python_compat::{BigInt, MAX_DEPTH, Value};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
prelude::*,
types::{
PyBool, PyBytes, PyComplex, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PySet, PyString,
PyTuple,
},
};
/// A plain Python value (what `json.loads` or `ast.literal_eval` returns) across the bridge.
/// Anything else, a live object included, is refused with `TypeError`.
pub(crate) struct PyValue(pub(crate) Value);
impl<'a, 'py> FromPyObject<'a, 'py> for PyValue {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
to_value(&value, 0).map(Self)
}
}
impl<'py> IntoPyObject<'py> for PyValue {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
to_python(py, self.0)
}
}
fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult<Value> {
if depth > MAX_DEPTH {
return Err(PyValueError::new_err("value is nested too deeply"));
}
let nested = |items: Bound<'_, PyAny>| -> PyResult<Vec<Value>> {
items
.try_iter()?
.map(|item| to_value(&item?, depth + 1))
.collect()
};
if value.is_none() {
return Ok(Value::None);
}
if let Ok(flag) = value.cast::<PyBool>() {
return Ok(Value::Bool(flag.is_true()));
}
if value.is_instance_of::<PyInt>() {
return Ok(Value::Int(match value.extract::<i64>() {
Ok(small) => small.into(),
Err(_) => value
.str()?
.to_str()?
.parse::<BigInt>()
.map_err(|error| PyValueError::new_err(error.to_string()))?,
}));
}
if let Ok(number) = value.cast::<PyFloat>() {
return Ok(Value::Float(number.value()));
}
if let Ok(text) = value.cast::<PyString>() {
return Ok(Value::Str(text.to_str()?.to_owned()));
}
if let Ok(bytes) = value.cast::<PyBytes>() {
return Ok(Value::Bytes(bytes.as_bytes().to_vec()));
}
if let Ok(complex) = value.cast::<PyComplex>() {
return Ok(Value::Complex {
re: complex.real(),
im: complex.imag(),
});
}
if let Ok(entries) = value.cast::<PyDict>() {
return entries
.iter()
.map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?)))
.collect::<PyResult<_>>()
.map(Value::Dict);
}
if value.is_instance_of::<PyList>() {
return nested(value.clone()).map(Value::List);
}
if value.is_instance_of::<PyTuple>() {
return nested(value.clone()).map(Value::Tuple);
}
if value.is_instance_of::<PySet>() || value.is_instance_of::<PyFrozenSet>() {
return nested(value.clone()).map(Value::Set);
}
Err(PyTypeError::new_err(format!(
"{} is not a plain data value",
value.get_type().name()?
)))
}
fn to_python(py: Python<'_>, value: Value) -> PyResult<Bound<'_, PyAny>> {
let all = |items: Vec<Value>| -> PyResult<Vec<Bound<'_, PyAny>>> {
items.into_iter().map(|item| to_python(py, item)).collect()
};
Ok(match value {
Value::None => py.None().into_bound(py),
Value::Bool(flag) => PyBool::new(py, flag).to_owned().into_any(),
Value::Int(number) => match i64::try_from(&number) {
Ok(small) => small.into_pyobject(py)?.into_any(),
Err(_) => py
.get_type::<PyInt>()
.call1((number.to_string(),))?
.into_any(),
},
Value::Float(number) => number.into_pyobject(py)?.into_any(),
Value::Complex { re, im } => PyComplex::from_doubles(py, re, im).into_any(),
Value::Str(text) => text.into_pyobject(py)?.into_any(),
Value::Bytes(bytes) => PyBytes::new(py, &bytes).into_any(),
Value::Tuple(items) => PyTuple::new(py, all(items)?)?.into_any(),
Value::List(items) => PyList::new(py, all(items)?)?.into_any(),
Value::Set(items) => PySet::new(py, all(items)?)?.into_any(),
Value::Dict(entries) => {
let dict = PyDict::new(py);
for (key, item) in entries {
dict.set_item(to_python(py, key)?, to_python(py, item)?)?;
}
dict.into_any()
}
})
}

View file

@ -20,10 +20,10 @@ use std::{
use async_trait::async_trait;
use litellm_router_core::{
Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output,
Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, RouterEvent, RouterRequest,
RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness,
build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo,
Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Deployment,
Error, Host, Output, Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache,
RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, SelectStage,
ThreadRandomness, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo,
precall::PreCallFacts,
};
use pyo3::{PyResult, Python, types::PyAnyMethods};
@ -298,8 +298,9 @@ pub(crate) enum HostCallResponse {
PreCall(PreCallFacts),
/// The candidate models a `routing_plugins` select step reported.
Models(Option<Vec<String>>),
/// The ids a `callback_filters` select step left, or the error a filter raised.
Filtered(Result<Vec<String>, AttemptError>),
/// The ids a `callback_filters` select step left with the one a filter prefers, or the error
/// a filter raised.
Filtered(Result<(Vec<String>, Option<String>), AttemptError>),
/// What a pick's select or pre-route step, or a fallback check, raised on the host instead of
/// answering: the walk takes it as the try's failure.
Raised(AttemptError),
@ -352,9 +353,15 @@ impl HostCallResponse {
got: self.kind(),
});
}
if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) =
(self, call)
&& let Some(id) = ids.iter().find(|id| !candidates.contains(id))
if let (
Self::Select(ids) | Self::Filtered(Ok((ids, _))),
HostCall::Select { candidates, .. },
) = (self, call)
&& let Some(id) = {
let known: std::collections::HashSet<&str> =
candidates.iter().map(String::as_str).collect();
ids.iter().find(|id| !known.contains(id.as_str()))
}
{
return Err(DriveError::UnknownCandidate { id: id.clone() });
}
@ -685,14 +692,22 @@ fn shares_ids(candidates: &[Arc<Deployment>]) -> bool {
/// The candidates a host answer names, in its order. An id names every candidate that carries
/// it, so deployments sharing an id stay together through a filter, as Python's list keeps them.
fn answered_by_id(candidates: &[Arc<Deployment>], ids: &[String]) -> Vec<Arc<Deployment>> {
let mut taken = vec![false; candidates.len()];
let mut positions: std::collections::HashMap<&str, Vec<usize>> =
std::collections::HashMap::new();
for (position, deployment) in candidates.iter().enumerate() {
positions
.entry(deployment.id.as_str())
.or_default()
.push(position);
}
let mut kept = Vec::new();
for id in ids {
for (position, deployment) in candidates.iter().enumerate() {
if !taken[position] && deployment.id == *id {
taken[position] = true;
kept.push(deployment.clone());
}
if let Some(found) = positions.remove(id.as_str()) {
kept.extend(
found
.into_iter()
.map(|position| candidates[position].clone()),
);
}
}
kept
@ -890,7 +905,7 @@ impl Host for StepHost {
&self,
ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error> {
) -> Result<CallbackFiltered, Error> {
let (scope, group) = ctx.scope.wire();
let joined = |own: &[String], walked: &[String]| [own, walked].concat();
let call = HostCall::Select {
@ -909,8 +924,11 @@ impl Host for StepHost {
let HostCallResponse::Filtered(result) = call_host(call).await else {
unreachable!("{CHECKED}");
};
let ids = result.map_err(Error::Attempt)?;
Ok(answered_by_id(&candidates, &ids))
let (ids, preferred) = result.map_err(Error::Attempt)?;
Ok(CallbackFiltered {
kept: answered_by_id(&candidates, &ids),
preferred,
})
}
async fn allow_fallback(
@ -1534,6 +1552,26 @@ mod tests {
);
}
#[rstest]
fn an_answer_by_id_keeps_its_order_and_every_candidate_sharing_an_id() {
let candidates: Vec<Arc<Deployment>> = ["a", "b", "a", "c"]
.into_iter()
.map(|id| {
let spec: litellm_router_core::DeploymentSpec =
serde_json::from_value(deployment("g", id)).unwrap();
Arc::new(Deployment::from_spec(spec).unwrap())
})
.collect();
let ids = ["c", "a", "x", "a"].map(String::from);
let kept = answered_by_id(&candidates, &ids);
let positions: Vec<usize> = kept
.iter()
.map(|d| candidates.iter().position(|c| Arc::ptr_eq(c, d)).unwrap())
.collect();
assert_eq!(positions, [3, 0, 2]);
}
#[rstest]
fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() {
let mut second = deployment("primary", "a1");

View file

@ -1181,7 +1181,8 @@ fn host_call_response(
stage: Stage::CallbackFilters { .. },
..
} => {
// The remaining ids, or a failure mapping as an attempt reports one.
// The remaining ids, alone or paired with the id a filter prefers, or a failure
// mapping as an attempt reports one.
let outcome = outcome.ok_or_else(missing)?;
if let Ok(failure) = outcome.cast::<PyDict>() {
return match attempt_outcome(failure)? {
@ -1191,7 +1192,10 @@ fn host_call_response(
)),
};
}
Ok(HostCallResponse::Filtered(Ok(outcome.extract()?)))
if let Ok(answer) = outcome.cast::<PyTuple>() {
return Ok(HostCallResponse::Filtered(Ok(answer.extract()?)));
}
Ok(HostCallResponse::Filtered(Ok((outcome.extract()?, None))))
}
HostCall::Select {
stage: Stage::Strategy { records, .. },

View file

@ -239,9 +239,13 @@ impl PythonBinding for MessagesPythonHost {
fn encode_response(
&mut self,
py: Python<'_>,
response: Box<litellm_llms_types::formats::messages::MessagesResponse>,
response: litellm_inference_messages::MessagesBody,
) -> PyResult<Py<PyAny>> {
public_response(py, ROUTE_HOST_MODULE, response.as_ref())
public_response(py, ROUTE_HOST_MODULE, &response)
}
fn deliver_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
self.cache.store_response(py, response)
}
fn encode_stream_head(

View file

@ -10,7 +10,7 @@
use std::{hint::black_box, time::Duration};
use criterion::{
BatchSize, BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
measurement::WallTime,
};
use litellm_python_compat::{Value, json, literal::literal_eval, pickle, repr::repr};
@ -72,14 +72,9 @@ fn formats(c: &mut Criterion) {
let mut group = c.benchmark_group("cached_completion");
text_formats(&mut group, &text, &value);
binary_formats(&mut group, &value, &pickled);
// `from_json` consumes its input, so each iteration gets a freshly parsed one.
group.throughput(Throughput::Bytes(dumped.len() as u64));
group.bench_function("from_json", |bencher| {
bencher.iter_batched(
|| serde_json::from_str::<serde_json::Value>(&dumped).expect("dumps output parses"),
json::from_json,
BatchSize::SmallInput,
)
group.bench_function("json_loads", |bencher| {
bencher.iter(|| json::loads(black_box(&dumped)))
});
group.finish();
}

View file

@ -6,6 +6,8 @@ use crate::MAX_DEPTH;
pub enum Error {
#[error("malformed Python literal at byte {0}")]
InvalidLiteral(usize),
#[error("malformed JSON at byte {0}")]
InvalidJson(usize),
#[error("unhashable type: '{0}'")]
Unhashable(&'static str),
#[error("value nests deeper than {MAX_DEPTH} levels")]

View file

@ -1,14 +1,19 @@
//! `json.dumps` with CPython's default options, and the JSON value `json.loads` returns.
//! `json.dumps` with CPython's default options, and `json.loads` with its default decoder.
//!
//! Defaults are `ensure_ascii=True`, `allow_nan=True`, separators `(", ", ": ")`, and no key
//! sorting. Python's `json.loads` is mapped by [`from_json`]; a `serde_json` number that does
//! not fit `i64` or `u64` arrives as a float, where Python would keep an `int`.
//! sorting on the way out, and `strict=True` with `NaN`, `Infinity` and `-Infinity` read as floats
//! on the way in.
use std::fmt::Write;
use std::{
collections::{HashMap, hash_map::Entry},
fmt::Write,
str::FromStr,
};
use num_bigint::BigInt;
use serde_json::{Map, Number};
use crate::{Error, Value, repr::float_repr};
use crate::{Error, MAX_DEPTH, Value, repr::float_repr};
/// `json.dumps(value)`.
pub fn dumps(value: &Value) -> Result<String, Error> {
@ -44,6 +49,260 @@ pub fn from_json(value: serde_json::Value) -> Value {
}
}
/// `json.loads(text)`. Integers keep full precision, and a repeated key keeps its first position
/// with its last value, as `dict(pairs)` does. An escape that leaves a lone surrogate fails, since
/// a Rust `String` cannot hold one.
pub fn loads(text: &str) -> Result<Value, Error> {
let mut decoder = Decoder {
text,
bytes: text.as_bytes(),
pos: 0,
depth: 0,
};
decoder.skip_whitespace();
let value = decoder.value()?;
decoder.skip_whitespace();
if decoder.pos != decoder.bytes.len() {
return Err(decoder.error());
}
Ok(value)
}
/// CPython's `json.scanner` and `scanstring`, one pass over the bytes.
struct Decoder<'a> {
text: &'a str,
bytes: &'a [u8],
pos: usize,
depth: usize,
}
impl Decoder<'_> {
fn error(&self) -> Error {
Error::InvalidJson(self.pos)
}
fn peek(&self) -> Option<u8> {
self.bytes.get(self.pos).copied()
}
fn eat(&mut self, token: &str) -> bool {
let found = self.bytes[self.pos..].starts_with(token.as_bytes());
if found {
self.pos += token.len();
}
found
}
fn expect(&mut self, token: &str) -> Result<(), Error> {
if self.eat(token) {
Ok(())
} else {
Err(self.error())
}
}
fn skip_whitespace(&mut self) {
while matches!(self.peek(), Some(b' ' | b'\t' | b'\n' | b'\r')) {
self.pos += 1;
}
}
fn value(&mut self) -> Result<Value, Error> {
match self.peek() {
Some(b'"') => self.string().map(Value::Str),
Some(b'[') => self.nested(Self::array),
Some(b'{') => self.nested(Self::object),
_ if self.eat("null") => Ok(Value::None),
_ if self.eat("true") => Ok(Value::Bool(true)),
_ if self.eat("false") => Ok(Value::Bool(false)),
_ if self.eat("NaN") => Ok(Value::Float(f64::NAN)),
_ if self.eat("Infinity") => Ok(Value::Float(f64::INFINITY)),
_ if self.eat("-Infinity") => Ok(Value::Float(f64::NEG_INFINITY)),
_ => self.number(),
}
}
fn nested(&mut self, parse: fn(&mut Self) -> Result<Value, Error>) -> Result<Value, Error> {
self.depth += 1;
if self.depth > MAX_DEPTH {
return Err(Error::TooDeep);
}
let value = parse(self)?;
self.depth -= 1;
Ok(value)
}
fn array(&mut self) -> Result<Value, Error> {
self.expect("[")?;
self.skip_whitespace();
let mut values = Vec::new();
if self.eat("]") {
return Ok(Value::List(values));
}
loop {
values.push(self.value()?);
self.skip_whitespace();
if self.eat("]") {
return Ok(Value::List(values));
}
self.expect(",")?;
self.skip_whitespace();
}
}
fn object(&mut self) -> Result<Value, Error> {
self.expect("{")?;
self.skip_whitespace();
let mut entries = Vec::new();
let mut positions = HashMap::<_, usize>::new();
if self.eat("}") {
return Ok(Value::Dict(entries));
}
loop {
if self.peek() != Some(b'"') {
return Err(self.error());
}
let key = self.string()?;
self.skip_whitespace();
self.expect(":")?;
self.skip_whitespace();
let value = self.value()?;
match positions.entry(key) {
Entry::Occupied(position) => entries[*position.get()].1 = value,
Entry::Vacant(position) => {
entries.push((Value::Str(position.key().clone()), value));
position.insert(entries.len() - 1);
}
}
self.skip_whitespace();
if self.eat("}") {
return Ok(Value::Dict(entries));
}
self.expect(",")?;
self.skip_whitespace();
}
}
/// `scanstring` with `strict=True`: control characters must be escaped.
fn string(&mut self) -> Result<String, Error> {
self.expect("\"")?;
let mut out = String::new();
loop {
let start = self.pos;
while self
.peek()
.is_some_and(|byte| byte != b'"' && byte != b'\\' && byte >= 0x20)
{
self.pos += 1;
}
out.push_str(&self.text[start..self.pos]);
match self.peek() {
Some(b'"') => {
self.pos += 1;
return Ok(out);
}
Some(b'\\') => {
self.pos += 1;
out.push(self.escape()?);
}
_ => return Err(self.error()),
}
}
}
fn escape(&mut self) -> Result<char, Error> {
let escaped = match self.peek() {
Some(b'"') => '"',
Some(b'\\') => '\\',
Some(b'/') => '/',
Some(b'b') => '\u{8}',
Some(b'f') => '\u{c}',
Some(b'n') => '\n',
Some(b'r') => '\r',
Some(b't') => '\t',
Some(b'u') => {
self.pos += 1;
return self.unicode_escape();
}
_ => return Err(self.error()),
};
self.pos += 1;
Ok(escaped)
}
/// `\uXXXX`, joining a high surrogate with the low one escaped right after it.
fn unicode_escape(&mut self) -> Result<char, Error> {
let unit = self.hex4()?;
let code = if (0xd800..0xdc00).contains(&unit) {
self.expect("\\u")?;
let low = self.hex4()?;
if !(0xdc00..0xe000).contains(&low) {
return Err(self.error());
}
0x10000 + ((unit - 0xd800) << 10) + (low - 0xdc00)
} else {
unit
};
char::from_u32(code).ok_or_else(|| self.error())
}
fn hex4(&mut self) -> Result<u32, Error> {
let digits = self
.text
.get(self.pos..self.pos + 4)
.filter(|digits| digits.bytes().all(|byte| byte.is_ascii_hexdigit()))
.ok_or_else(|| self.error())?;
self.pos += 4;
u32::from_str_radix(digits, 16).map_err(|_| self.error())
}
/// `NUMBER_RE`: `(-?(?:0|[1-9][0-9]*))(\.[0-9]+)?([eE][-+]?[0-9]+)?`, an `int` unless a
/// fraction or exponent matched.
fn number(&mut self) -> Result<Value, Error> {
let start = self.pos;
self.eat("-");
match self.peek() {
Some(b'0') => self.pos += 1,
Some(b'1'..=b'9') => self.digits(),
_ => return Err(Error::InvalidJson(start)),
}
let integer_end = self.pos;
if self.peek() == Some(b'.') && self.bytes.get(self.pos + 1).is_some_and(u8::is_ascii_digit)
{
self.pos += 1;
self.digits();
}
if matches!(self.peek(), Some(b'e' | b'E')) {
let exponent = self.pos;
self.pos += 1;
if matches!(self.peek(), Some(b'+' | b'-')) {
self.pos += 1;
}
if self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
self.digits();
} else {
self.pos = exponent;
}
}
let text = &self.text[start..self.pos];
if self.pos == integer_end {
BigInt::from_str(text)
.map(Value::Int)
.map_err(|_| Error::InvalidJson(start))
} else {
f64::from_str(text)
.map(Value::Float)
.map_err(|_| Error::InvalidJson(start))
}
}
fn digits(&mut self) {
while self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
self.pos += 1;
}
}
}
/// `json.loads(json.dumps(value))` as a `serde_json` value: tuples become arrays and dict
/// keys are coerced to strings as `json.dumps` does. Non-finite floats, which Python writes
/// as `NaN` and `Infinity`, have no `serde_json` form and fail with [`Error::NonFiniteFloat`].
@ -166,3 +425,86 @@ fn write_string(out: &mut String, text: &str) {
}
out.push('"');
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::{dumps, loads};
use crate::{Error, Value, repr::repr};
// Expected values are the repr of CPython 3.14 json.loads(text)
#[rstest]
#[case::nan("NaN", "nan")]
#[case::infinity("Infinity", "inf")]
#[case::negative_infinity("-Infinity", "-inf")]
#[case::non_finite_in_a_list("[NaN, Infinity, -Infinity]", "[nan, inf, -inf]")]
#[case::beyond_u64("18446744073709551616", "18446744073709551616")]
#[case::exponent_overflow("1e400", "inf")]
#[case::upper_exponent("1E+2", "100.0")]
#[case::negative_zero_int("-0", "0")]
#[case::negative_zero_float("-0.0", "-0.0")]
#[case::repeated_key(r#"{"a": 1, "b": 2, "a": 3}"#, "{'a': 3, 'b': 2}")]
#[case::surrogate_pair(r#""😀""#, "'😀'")]
#[case::upper_hex_surrogate_pair(r#""😀""#, "'😀'")]
#[case::escapes(
r#"{"k": [true, false, null, 1.5, "é\/"]}"#,
"{'k': [True, False, None, 1.5, 'é/']}"
)]
#[case::delete_is_not_a_control_character("\"\x7f\"", "'\\x7f'")]
#[case::surrounding_whitespace(" \n[1]\t", "[1]")]
#[case::empty_containers("[{}, []]", "[{}, []]")]
fn loads_reads_like_python(#[case] text: &str, #[case] expected: &str) {
assert_eq!(repr(&loads(text).unwrap()), expected);
}
// CPython 3.14 json.loads raises JSONDecodeError on each
#[rstest]
#[case::lowercase_nan("nan")]
#[case::signed_nan("-NaN")]
#[case::plus_sign("+1")]
#[case::bare_minus("-")]
#[case::leading_zero("01")]
#[case::bare_point("1.")]
#[case::bare_exponent("1e")]
#[case::trailing_data("Infinity1")]
#[case::trailing_comma("[1,]")]
#[case::missing_comma("[1 2]")]
#[case::missing_colon(r#"{"a" 1}"#)]
#[case::unquoted_key("{1: 2}")]
#[case::raw_newline_in_string("\"a\nb\"")]
#[case::unknown_escape(r#""\x41""#)]
#[case::empty("")]
#[case::byte_order_mark("\u{feff}1")]
fn loads_rejects_what_python_rejects(#[case] text: &str) {
assert!(matches!(loads(text), Err(Error::InvalidJson(_))));
}
#[rstest]
fn a_lone_surrogate_is_a_known_limit() {
assert!(matches!(loads(r#""\ud800""#), Err(Error::InvalidJson(_))));
}
#[rstest]
fn loads_reads_back_the_non_finite_floats_dumps_writes() {
let written = dumps(&Value::List(vec![
Value::Float(f64::NAN),
Value::Float(f64::INFINITY),
Value::Float(f64::NEG_INFINITY),
]))
.unwrap();
assert_eq!(repr(&loads(&written).unwrap()), "[nan, inf, -inf]");
}
#[rstest]
fn loads_rejects_nesting_past_max_depth() {
let text = format!(
"{}{}",
"[".repeat(crate::MAX_DEPTH + 1),
"]".repeat(crate::MAX_DEPTH + 1)
);
assert!(matches!(loads(&text), Err(Error::TooDeep)));
}
}

View file

@ -7,9 +7,10 @@
//! |---|---|
//! | [`literal`] | `ast.literal_eval(text)` |
//! | [`repr`] | `repr(value)` and `str(value)` |
//! | [`json`] | `json.dumps(value)`, and `json.loads(json.dumps(value))` as a JSON value |
//! | [`json`] | `json.dumps(value)` and `json.loads(text)` |
//! | [`pickle`] | `pickle.loads(data)` and `pickle.dumps(value)` for plain data |
//! | [`truthy`] | `bool(value)` |
//! | [`number`] | `int(text)`, `float(text)`, and `int(seconds)` as a `Duration` |
//!
//! [`Value`] is the closed data model these formats share. Live Python objects
//! (descriptors, `__bool__`, `__str__`, callbacks) are out of scope: those belong to the
@ -20,10 +21,12 @@
//! (string hashes are randomized per process). Sets keep their literal order.
//! - `str` values are Rust `String`s, so lone surrogates cannot be represented.
//! - [`pickle::loads`] decodes `tuple`, `set`, and `frozenset` as lists.
//! - [`number`] reads ASCII digits only, while Python's `int` and `float` accept any Unicode decimal digit.
mod error;
pub mod json;
pub mod literal;
pub mod number;
pub mod pickle;
pub mod repr;
pub mod truthy;

View file

@ -0,0 +1,121 @@
use std::{str::FromStr, time::Duration};
use num_bigint::BigInt;
/// `int(text)` in base 10, or `None` where Python raises `ValueError`.
pub fn int(text: &str) -> Option<BigInt> {
parse(text)
}
/// `float(text)`, or `None` where Python raises `ValueError`.
pub fn float(text: &str) -> Option<f64> {
parse(text)
}
/// Why [`seconds`] refused a value, as `int(value)` would. Messages are CPython's.
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum SecondsError {
/// `OverflowError`. Finite values past `Duration::MAX` land here too, where Python builds a big int.
#[error("cannot convert float infinity to integer")]
Overflow,
/// `ValueError`.
#[error("cannot convert float NaN to integer")]
NotANumber,
}
/// A float of seconds as a `Duration`, refusing what `int(value)` refuses. Finite negative values clamp to zero.
pub fn seconds(value: f64) -> Result<Duration, SecondsError> {
match value {
value if value.is_nan() => Err(SecondsError::NotANumber),
value if value.is_infinite() => Err(SecondsError::Overflow),
value => Duration::try_from_secs_f64(value.max(0.0)).map_err(|_| SecondsError::Overflow),
}
}
fn parse<N: FromStr>(text: &str) -> Option<N> {
let text = text.trim();
let digit_at = |index: Option<usize>| {
index
.and_then(|index| text.as_bytes().get(index))
.is_some_and(u8::is_ascii_digit)
};
let underscores_between_digits = text
.match_indices('_')
.all(|(index, _)| digit_at(index.checked_sub(1)) && digit_at(Some(index + 1)));
if !underscores_between_digits {
return None;
}
text.replace('_', "").parse().ok()
}
#[cfg(test)]
mod tests {
use num_bigint::BigInt;
use rstest::rstest;
use std::time::Duration;
use super::{SecondsError, float, int, seconds};
// Expected values are CPython 3.14 int(text)
#[rstest]
#[case::trailing_newline("2\n", Some(2))]
#[case::surrounding_spaces_and_underscore(" 6_380 ", Some(6380))]
#[case::plus_sign("+5", Some(5))]
#[case::negative_zero("-0", Some(0))]
#[case::leading_zeros("007", Some(7))]
#[case::unicode_spaces("\u{a0}7\u{2003}", Some(7))]
#[case::beyond_i64("99999999999999999999", Some(99_999_999_999_999_999_999_i128))]
#[case::double_underscore("1__0", None)]
#[case::leading_underscore("_1", None)]
#[case::trailing_underscore("1_", None)]
#[case::decimal_point("5.0", None)]
#[case::information_separator("\x1c7", None)]
#[case::empty("", None)]
#[case::sign_only("+", None)]
#[case::non_ascii_digit_is_a_known_limit("\u{667}", None)]
fn int_parses_like_python(#[case] text: &str, #[case] expected: Option<i128>) {
assert_eq!(int(text), expected.map(BigInt::from));
}
// Expected values are CPython 3.14 float(text)
#[rstest]
#[case::underscore_and_newline(" 1_0.5\n", Some(10.5))]
#[case::exponent("1e3", Some(1000.0))]
#[case::leading_point(".5", Some(0.5))]
#[case::trailing_point("5.", Some(5.0))]
#[case::plus_sign("+1.5", Some(1.5))]
#[case::infinity_word("Infinity", Some(f64::INFINITY))]
#[case::infinity_with_space("infinity ", Some(f64::INFINITY))]
#[case::underscore_before_exponent("1_e3", None)]
#[case::underscore_after_exponent("1e_3", None)]
#[case::underscore_after_point("1._5", None)]
#[case::underscore_after_word("inf_", None)]
#[case::hexadecimal("0x10", None)]
#[case::bare_exponent("1.5e", None)]
fn float_parses_like_python(#[case] text: &str, #[case] expected: Option<f64>) {
assert_eq!(float(text), expected);
}
#[rstest]
fn float_reads_a_signed_nan_like_python() {
assert!(float("-NaN").is_some_and(f64::is_nan));
}
// CPython 3.14: int(inf) and int(-inf) raise OverflowError, int(nan) raises ValueError
#[rstest]
#[case::fraction(1.5, Ok(Duration::from_millis(1500)))]
#[case::zero(0.0, Ok(Duration::ZERO))]
#[case::negative_clamps_to_zero(-3.0, Ok(Duration::ZERO))]
#[case::infinity(f64::INFINITY, Err(SecondsError::Overflow))]
#[case::negative_infinity(f64::NEG_INFINITY, Err(SecondsError::Overflow))]
#[case::past_duration_max(1e20, Err(SecondsError::Overflow))]
#[case::nan(f64::NAN, Err(SecondsError::NotANumber))]
#[case::negative_nan(-f64::NAN, Err(SecondsError::NotANumber))]
fn seconds_refuses_what_int_refuses(
#[case] value: f64,
#[case] expected: Result<Duration, SecondsError>,
) {
assert_eq!(seconds(value), expected);
}
}

View file

@ -76,6 +76,11 @@ const KNOWN: &[(&str, &str, &str)] = &[
"literal_eval",
"deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it",
),
(
"nested_150",
"json.loads",
"deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it",
),
(
"nested_150",
"pickle.loads",
@ -341,3 +346,19 @@ fn to_json_matches_python_json_round_trip(fixtures: &Fixtures) {
}
mismatches.finish();
}
/// CPython gives back the same text from `json.dumps(json.loads(text))` for every text its
/// `json.dumps` wrote here, non-finite floats included, so `loads` must too.
#[rstest]
fn loads_reads_back_what_python_json_dumps_wrote(fixtures: &Fixtures) {
let mut mismatches = Mismatches::default();
for row in &fixtures.rows {
let Some(text) = &row.json else { continue };
let actual = match json::loads(text).and_then(|value| json::dumps(&value)) {
Ok(text) => text,
Err(error) => format!("error: {error}"),
};
mismatches.check(row, "json.loads", text, &actual);
}
mismatches.finish();
}

View file

@ -1,12 +1,12 @@
- The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace
- It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default
- Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route
- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`)
- Every connection sets TCP_NODELAY, like redis-py's `socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)`. Set every redis-rs timeout and retry explicitly, citing the redis-py 5.3.1 or `litellm/` source next to each constant. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's. `Timeouts` takes `socket_timeout` and `socket_connect_timeout`, and the connect timeout falls back to the socket timeout like redis-py. Standalone reconnects on the next command with no retry of the failed one, matching redis-py's `Retry(NoBackoff(), 0)`. redis-rs keeps a socket whose command timed out, so `Renewing` in `connection.rs` swaps in a fresh `ConnectionManager`, or a fresh cluster connection like redis-py's cluster `aclose`, after a timeout. Without it a socket that went silent times out every command until TCP gives up on it. Cluster keeps redis-py's 3 `cluster_error_retry_attempts` and waits a flat 250ms between them, which is redis-py's CLUSTERDOWN sleep. redis-py's jittered per-node reconnect backoff has no redis-rs equivalent. The retry set also differs, redis-rs retries LOADING and READONLY internally, while redis-py 5.3.1 raises them. Do not claim exact retry parity
- Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values
## Errors
- `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it
- Classify in `From<RedisError> for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error
- Classify with `Error::classify(error, deployment)`, never a bare `From<RedisError>`, because redis-py (5.3.1) decides by topology. `Transport` is what Python's `_redis_health_error_types` counts: LOADING everywhere, CLUSTERDOWN and MASTERDOWN only on cluster (only `ClusterParser` maps them), READONLY only under Sentinel (where `SentinelManagedConnection` turns it into `ConnectionError` and rediscovers), plus NOAUTH, WRONGPASS and the `ERR` messages redis-py maps to `ConnectionError`, matched exactly like `parse_error`. TRYAGAIN is `Server`, since the cluster client already retried it. Each code gets its own rstest case built from a parsed RESP error
- A `Server` error never drops the connection
- `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome
- Never swallow an error or turn it into a miss. The caller picks the fallback
@ -17,27 +17,61 @@
- Every public call (each command, `eval`, one `run` of a pipeline) takes one admission through `guarded` / `admit` in `client.rs`. Never call the connection around it
- `Transport` counts as a failure (`Timeout` when `is_timeout()`, else `Hard`). `Server` and `Decode` count as nothing, since Redis answered and the request was bad. A probe that gets one still reopens, through the dropped-probe rule in `circuit-breaker`. A pipeline counts one failure if any op came back `Transport`, else one success
- An open breaker returns `Error::BreakerOpen` without touching the socket
- Python doesn't count MASTERDOWN, TRYAGAIN or READONLY as health failures. Here they are `Transport` and do count
## Namespace
- Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace`
- The rule never changes, because Python and Rust pods share keys in the same Redis
## Commands
- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET
- A command is a value, one constructor in `commands.rs` returning `Command<T>`: its alone and pipelined forms and its reply decode. `RedisClient::send`, `Pipeline::with` and `RedisBatch::add` are the only entry points, and each applies its namespace to the keys once. A new command is one constructor, never a method on an executor. `ping` and `subscribe` stay methods. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET
- A store attaches its value codec with `Command::map`, never a keyspace codec inside this crate
- On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test
- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`)
- Every command is the one `RedisCache` sends, argument for argument, so `MONITOR` on a mixed fleet shows one shape. A difference from main's Python is a defect even when the result is the same
- TTLs are `Duration` and are sent in whole seconds (`SET ... EX`, `EXPIRE`, `TTL`), like redis-py's `int(ttl)`. `set_nx` sends `EX` before `NX`, the order redis-py builds
- Numbers go out as `Amount`, written the way redis-py encodes them: an `int` with `str()` and a `float` with `repr()` from `python-compat`
- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` alone is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. Pipelined it is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when an expiry is given, `KeepExisting` lowered to `Refresh`. `Plan` keeps only the cases whose alone and pipelined forms differ
## Scripts and pipelines
- A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production
- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async`
- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline a script is sent as `EVALSHA`, like `RedisBatch`. A `NOSCRIPT` reply means that op ran nothing, so it alone runs again after the pipeline through `invoke_async` (`EVALSHA`, `SCRIPT LOAD`, `EVALSHA`), like `RedisBatch`'s `_settle_alone`. It then lands after the ops around it, as it does in Python
- `Plan::Script` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice
- `RedisBatch::pipelined` is one `Pipeline` whose ops answer as typed `Declared<T>`s. A store that sends its own pipeline uses it rather than decoding `run`'s raw `Value`s
- A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip
- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order
- Pipeline ops take one key each. A multi-key op would span slots inside one group
- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once
- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key and runs the slots concurrently, returning results in input order. Within a slot, runs of idempotent ops (GET, SET, DEL, EXPIRE) share a request and every script op gets its own, sent in order. redis-rs retries a cluster request whole, so a script op batched with others would be applied twice when only a neighbour failed
- Pipeline ops take one key each. A multi-key op would span slots inside one group. An op can send more than one command for that key (`INCRBYFLOAT` then `EXPIRE`), and settles from the first reply, failing if any of its replies failed
- Command constructors live in `commands.rs` and every executor sends them, so a command is written once
## Tests
## Pub/sub
- Classic `PUBLISH` / `SUBSCRIBE`, never sharded, since redis-py publishes and subscribes through one plain client and cluster broadcasts classic messages to every node. Payloads are raw bytes, and the codec belongs to the store that owns the channel
- Channel names never get the namespace. Python builds `{namespace}:{channel}` itself, so the store does too
- `publish` is a command like any other, guarded by the breaker and reported to events
- `subscribe` returns a `Stream` that opens its own RESP2 `PubSub` connection when first polled, like redis-py's `PubSub`, to the endpoint, the first reachable cluster node, or the master Sentinel names at that moment. It never touches the breaker or events, like Python's subscriber, which reads off its own connection
- After a lost connection or a failed subscribe it waits and resubscribes, 5s doubling to 60s and reset by a successful subscribe, like `AuthCacheInvalidationSubscriber._run`. Messages published while it is down are lost. The connection has no idle ping, so a socket that dies silently is only noticed by TCP
- redis-rs's `PubSub` takes no credentials provider, so a provider is asked for its current credentials once when the subscription connects, which is when redis-py's pubsub connection asks it. The connect is bounded by the connect timeout
## Batches
- `RedisBatch` is litellm's `RedisBatch` (`litellm/caching/redis_batch.py`). It keeps each declared op as data and, on a standalone or Sentinel Redis, sends them later as one `Pipeline`, so it inherits every pipeline rule above. On a cluster it never pipelines and runs each op on its own, concurrently, the way `_settle_alone` calls the `RedisCache` method behind it. An increment there keeps a TTL the key already has, as `async_increment` does, where the pipelined one refreshes it like `_Increment`
- `mget` is one `MGET` of the keys in order with repeats read once, like `_MGet`, and answers one value per key it was given, in that order. Use it instead of one `get` per key whenever Python reads several keys in one call
- It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local
- `RedisBatch::alone` runs every op on its own like the cluster path, so a fresh alone batch is "send it now" and `Declared<T>` is the one reply handle on every route
- Each op returns a `Declared<T>` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once
- An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis
- `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client
- Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes
- Tests that only need the public API live in `tests/<module>.rs`, with client builders in `tests/support`. Inline `#[cfg(test)]` modules are only for private items (`error.rs`, `namespace.rs`)
- Observe connections from the server side (`RedisDeployment::clients`, `cli_on_every_node`), never by reaching into `RedisClient`'s connection
- Every command test runs on each topology through `#[values(Topology::Standalone, Topology::Cluster)]`. A test is single-topology only when the behavior exists on one (a `db` number has no cluster meaning)
- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables
- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables. A socket that never answers again comes from `TcpProxy::blackhole_open_connections`, since `thaw` revives the same socket
- Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline
## TLS, credentials and Sentinel
- `TlsOptions` follows redis-py 5.3.1. It verifies the certificate by default, skips the hostname check by default (`ssl_check_hostname=False`), and trusts the system roots plus `ca_certs`
- rustls only reads unencrypted keys, so a PKCS#8 `ENCRYPTED PRIVATE KEY` is decrypted with `key_password` before it reaches redis-rs. A legacy `Proc-Type: 4,ENCRYPTED` PEM key is refused with `KeyError::LegacyEncryption` rather than decrypted by hand-rolled OpenSSL key derivation
- redis-rs loads the system trust store for every TLS connection and fails if that fails, even with `ca_certs` set. Where the macOS keychain is blocked (agent sandboxes), run the TLS tests with `SSL_CERT_FILE=/etc/ssl/cert.pem`
- `Credentials::Static` are sent at connect, and URL userinfo wins over them, like redis-py's `from_url`
- A url's `?protocol=` is ignored and every connection speaks RESP2, which every reply decode assumes
- `Credentials::Provider` wraps a redis-rs `StreamingCredentialsProvider`. redis-rs takes the first item at connect and re-`AUTH`s live connections on every later item, on each cluster node. Providers push, redis-rs never polls. With a provider, redis-rs overwrites URL userinfo at connect, which matches Python stripping `user:pass@` from the URL
- Sentinel discovers the master on each new connection and drops it on a transport failure or a READONLY reply, to one command or inside a pipeline, like redis-py's `SentinelConnectionPool`
- redis-rs sends `ROLE` to the master a Sentinel names before using it, which redis-py never does. That probe gets the data credentials (a provider's current pair) and TLS, so the Sentinel client is built on each discovery. redis-rs gives the probe no way to skip the hostname check, and drops the client certificate when the certificate is not verified, so a TLS master whose certificate does not name the announced address, or that requires a client certificate while verification is off, fails discovery where Python connects
- Sentinel monitors get `sentinel_password`, the data username when that password is set, and the TLS mode and certificates. Python's `_sentinel_auth_kwargs` forwards every other connection kwarg too, and it sends the username even without a sentinel password
- A `Debug` impl never shows a secret. Follow `.agents/skills/rust-redaction`

View file

@ -6,13 +6,28 @@ license.workspace = true
repository.workspace = true
[dependencies]
tracing.workspace = true
litellm-circuit-breaker.workspace = true
litellm-secret-value.workspace = true
litellm-python-compat.workspace = true
litellm-redis-types.workspace = true
futures-util.workspace = true
redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] }
redis = { workspace = true, features = [
"cluster-async",
"connection-manager",
"sentinel",
"tls-rustls-insecure",
"token-based-authentication",
"tokio-rustls-comp",
] }
pem = "4.0.0"
pkcs8 = { version = "0.10.2", features = ["3des", "encryption", "std"] }
rustls-native-certs.workspace = true
socket2 = "0.6.5"
thiserror.workspace = true
tokio = { workspace = true, features = ["rt", "sync", "time"] }
[dev-dependencies]
litellm-redis-testing.workspace = true
rustls.workspace = true
rstest.workspace = true
tokio.workspace = true

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