Merge branch 'litellm_alerts_pr2_final_validation' of /Users/mfkhalil/Code/litellm-activity-before into litellm_alerts_reconcile_03

This commit is contained in:
moe-berri 2026-11-07 11:31:47 +08:00
commit bc6caf55c3
843 changed files with 49716 additions and 14134 deletions

View file

@ -1036,9 +1036,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'))"
@ -1295,7 +1295,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'
@ -1498,9 +1498,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

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

@ -4381,6 +4381,7 @@ dependencies = [
"criterion",
"fancy-regex 0.19.2",
"futures-util",
"indexmap 2.14.0",
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
@ -5773,6 +5774,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",

View file

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

@ -42,9 +42,24 @@ mod _native {
#[pymodule_export]
use crate::redis::client::native_redis_breaker_states;
#[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::{NativePending, NativeRedisStoreError};
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, PySpendRead};
#[pymodule_export]
use crate::redis::store::{
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
};
#[pymodule_export]
use crate::router::{RouterCall, RouterCoreHandle};
#[pymodule_export]
@ -123,8 +138,22 @@ mod tests {
"NativeRedisClient",
"NativeRequestRedis",
"NativeSend",
"NativeAuthCache",
"NativeRateLimits",
"WindowCount",
"AtomicApplied",
"AtomicOverLimit",
"GuardedReply",
"GaugesAcquired",
"GaugeOverLimit",
"NativePending",
"NativeRouterCooldowns",
"NativeRouterUsage",
"NativeSensitiveRoutePins",
"NativeRedisStoreError",
"NativeSpendCounters",
"SpendRead",
"LocalSizes",
"native_redis_breaker_states",
"NativeClickHouseSpendConfig",
"NativeClickHouseSpendStorage",

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

@ -18,7 +18,7 @@ use pyo3::{
exceptions::{PyOverflowError, PyTypeError, PyValueError},
prelude::*,
sync::MutexExt,
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple},
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple},
};
use self::credentials::PythonCredentialSource;
@ -40,16 +40,18 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
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 kwargs.cast::<PyDict>()?.iter() {
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 = kwargs
.cast::<PyDict>()?
let credential_provider = source
.get_item("credential_provider")?
.filter(|provider| !provider.is_none() && !is_plain(provider))
.map(Bound::unbind);
@ -65,27 +67,44 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
}
}
/// `litellm.constants`' `REDIS_CIRCUIT_BREAKER_*`, which Python reads from the environment once.
pub(crate) struct RedisBreaker(BreakerConfig);
/// `litellm.redis.store_settings.RedisBreakerSettings`, which `RedisCache` fills from `litellm.constants`.
#[derive(FromPyObject)]
pub(crate) struct RedisBreaker {
#[pyo3(attribute)]
failure_threshold: u32,
#[pyo3(attribute)]
recovery_timeout: f64,
#[pyo3(attribute)]
timeout_min_duration: f64,
}
impl<'a, 'py> FromPyObject<'a, 'py> for RedisBreaker {
type Error = PyErr;
fn extract(breaker: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let (failure_threshold, recovery_timeout, timeout_min_duration): (u32, f64, f64) =
breaker.extract()?;
impl RedisBreaker {
fn config(&self) -> PyResult<BreakerConfig> {
let seconds = |seconds: f64| {
Duration::try_from_secs_f64(seconds)
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))
.map_err(|error| PyValueError::new_err(error.to_string()))
};
Ok(Self(BreakerConfig {
failure_threshold,
recovery_timeout: seconds(recovery_timeout)?,
timeout_min_duration: seconds(timeout_min_duration)?,
}))
Ok(BreakerConfig {
failure_threshold: self.failure_threshold,
recovery_timeout: seconds(self.recovery_timeout)?,
timeout_min_duration: seconds(self.timeout_min_duration)?,
})
}
}
/// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built.
#[derive(FromPyObject)]
pub(crate) struct StoreSettings {
#[pyo3(attribute)]
redis_kwargs: RedisSpec,
#[pyo3(attribute)]
namespace: Option<String>,
#[pyo3(attribute)]
default_ttl: f64,
#[pyo3(attribute)]
breaker: Option<RedisBreaker>,
}
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
value.is_instance_of::<PyString>()
|| value.is_instance_of::<PyBool>()
@ -119,7 +138,7 @@ fn shared_client(
py: Python<'_>,
spec: RedisSpec,
namespace: Option<String>,
breaker: Option<RedisBreaker>,
breaker: Option<BreakerConfig>,
) -> PyResult<Option<RedisClient>> {
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
let RedisSpec {
@ -176,7 +195,7 @@ fn shared_client(
}
let connected = {
let _runtime = runtime.enter();
litellm_redis_connect::connect(resolved.clone(), host, breaker.map(|breaker| breaker.0))
litellm_redis_connect::connect(resolved.clone(), host, breaker)
};
let client = match connected {
Ok(client) => client,
@ -205,15 +224,15 @@ pub(crate) struct NativeRedisClient {
impl NativeRedisClient {
/// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory.
#[staticmethod]
#[pyo3(signature = (redis_kwargs, *, default_ttl_seconds, namespace = None, breaker = None))]
fn connect(
py: Python<'_>,
redis_kwargs: RedisSpec,
default_ttl_seconds: f64,
namespace: Option<String>,
breaker: Option<RedisBreaker>,
) -> PyResult<Option<Self>> {
let default_ttl = number::seconds(default_ttl_seconds).map_err(seconds_error)?;
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(RedisBreaker::config).transpose()?;
Ok(
shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self {
client,

View file

@ -1,11 +1,3 @@
#[expect(
dead_code,
reason = "the native stores that use these helpers land in the next PR"
)]
pub(crate) mod client;
mod events;
#[expect(
dead_code,
reason = "the native stores that use these helpers land in the next PR"
)]
pub(crate) mod store;

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

@ -1,4 +1,11 @@
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;
@ -159,6 +166,7 @@ where
}
/// 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 {
@ -188,6 +196,21 @@ impl<'py> IntoPyObject<'py> for PyNumber {
}
}
/// 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);

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,591 @@
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 {
#[pyo3(attribute)]
descriptor_key: String,
#[pyo3(attribute)]
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 {
#[pyo3(attribute)]
scope: PyScope,
#[pyo3(attribute)]
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 {
#[pyo3(attribute)]
suffix: String,
#[pyo3(attribute)]
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 {
#[pyo3(attribute)]
counter: PyCounter,
#[pyo3(attribute)]
amount: PyNumber,
#[pyo3(attribute)]
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 {
#[pyo3(attribute)]
counter: PyCounter,
#[pyo3(attribute)]
limit: i64,
#[pyo3(attribute)]
increment: i64,
#[pyo3(attribute)]
ttl: i64,
#[pyo3(attribute)]
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,183 @@
use std::sync::Arc;
use litellm_redis_store::{
Send,
router_cooldowns::{Cooldown, RouterCooldowns},
};
use pyo3::{
prelude::*,
types::{PyDict, 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<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, cooldown) in self.0.snapshot().map_err(kind_error)? {
snapshot.set_item(key, PyData(cooldown))?;
}
Ok(snapshot)
}
}

View file

@ -0,0 +1,184 @@
use std::{sync::Arc, time::Duration};
use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage};
use pyo3::{prelude::*, types::PyDict};
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 {
#[pyo3(attribute)]
deployment_id: String,
#[pyo3(attribute)]
model: Option<String>,
#[pyo3(attribute)]
minute: String,
#[pyo3(attribute)]
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<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, count) in self.0.snapshot().map_err(kind_error)? {
snapshot.set_item(key, PyNumber(count))?;
}
Ok(snapshot)
}
}

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,256 @@
use std::sync::Arc;
use litellm_redis_store::{
Number,
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
};
use pyo3::{prelude::*, types::PyDict};
use super::{
NativePending, PyData, PyNumber, 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 }
}
}
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())
}
/// `("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; `False` when there is no Redis.
fn increment_many_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> bool {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
}
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<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, spend) in self.0.snapshot().map_err(kind_error)? {
snapshot.set_item(key, PyNumber(spend))?;
}
Ok(snapshot)
}
}

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,8 +353,10 @@ impl HostCallResponse {
got: self.kind(),
});
}
if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) =
(self, call)
if let (
Self::Select(ids) | Self::Filtered(Ok((ids, _))),
HostCall::Select { candidates, .. },
) = (self, call)
&& let Some(id) = ids.iter().find(|id| !candidates.contains(id))
{
return Err(DriveError::UnknownCandidate { id: id.clone() });
@ -890,7 +893,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 +912,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(

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

@ -1,11 +1,11 @@
- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (Python-facing values with their `FromValue` / `IntoValue`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`)
- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`)
- An op declared on a `Send` returns `Pending<T>`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on
- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s
- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec
- A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc<RemoteSlot>`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient`
- Every family works with no Redis, answering from memory. `attach` adds the Redis tier at most once and keeps memory, as `DualCache.attach_redis_cache` does, so a host builds the store at startup and never rebuilds it to connect
- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not
- Key suffixes and wire spellings derive from strum on the family's enums, and a type whose Python side is plain data implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`, so the bridge never names a field or a spelling
- Key suffixes and wire spellings derive from strum on the family's enums. A value the family stores in Redis implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`. An argument or reply is a typed record, never a positional tuple: a frozen Python dataclass read in `python-bridge` by a `#[derive(FromPyObject)]` struct with `#[pyo3(attribute)]` fields and converted into the family's type with `From`, and a reply returned as a frozen `#[pyclass]`. The bridge reads an enum through its strum spelling, never a hand-written match
- Wire contracts are hand-written cases with expected values, inlined as the same literals in this crate's test and in `tests/unit/redis/stores/`, checked here against real Redis and there against Python's own codec. Change both sides together
- `rate_limits.rs` keeps the v3 limiter's Lua byte for byte (`rate_limits/lua.rs`) so pods on either stack share one script SHA; the admission decisions stay in the Python limiter
- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. Each op exists once and takes a `Send` (`Now`, `PreCall`, `PostCall`) picked by the caller, and the store calls `send.batch(&client)`. A store never encodes a phase in a method name. A post-call method exists only where a follow-up runs after the write settles (`token_increment_post_call`), as a thin wrapper over the shared op. `Send::Now` is an alone batch, so a script call outside a request goes out alone, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here

View file

@ -0,0 +1,8 @@
/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every
/// other object under a prefixed key.
pub fn is_key_object(key: &str) -> bool {
key.len() == 64
&& key
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}

View file

@ -0,0 +1,258 @@
mod keys;
mod types;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_python_compat::{Value, json};
use litellm_redis_client::{Declared, RedisBatch, commands};
pub use self::keys::is_key_object;
pub use self::types::Settings;
use crate::{
Error, Number, Pending, RequestRedis, Send,
dual_cache::DualCache,
memory::{Clock, lock, system_clock},
remote::{Remote, RemoteSlot},
request::ready,
};
/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it.
const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])";
/// The proxy's auth object cache with `enable_redis_auth_cache` on: `UserApiKeyCache`'s two
/// `DualCache` partitions, so churn in management objects cannot evict user-key objects, in front of
/// one Redis shared by every pod. Keys are the logical cache keys `UserApiKeyCache` callers already
/// use; values are what `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`.
pub struct AuthCache {
key_objects: DualCache<Value>,
objects: DualCache<Value>,
redis: Arc<RemoteSlot>,
settings: Mutex<Settings>,
}
impl AuthCache {
pub fn new(remote: Option<Remote>, settings: Settings) -> Self {
Self::with_clock(remote, settings, system_clock())
}
pub fn with_clock(remote: Option<Remote>, settings: Settings, now: Clock) -> Self {
let redis = Arc::new(RemoteSlot::new(remote));
let partition = || DualCache::new(redis.clone(), settings.memory(), now.clone());
Self {
key_objects: partition(),
objects: partition(),
redis,
settings: Mutex::new(settings),
}
}
pub fn has_redis(&self) -> bool {
self.redis.get().is_some()
}
pub fn settings(&self) -> Settings {
*lock(&self.settings)
}
/// `UserApiKeyCache.update_cache_ttl` and `update_in_memory_max_size` on both partitions, which
/// keep what they hold.
pub fn reconfigure(&self, settings: Settings) {
*lock(&self.settings) = settings;
self.key_objects.reconfigure(settings.memory());
self.objects.reconfigure(settings.memory());
}
pub fn get_local(&self, key: &str) -> Result<Option<Value>, Error> {
Ok(self.partition(key).memory().get(key)?)
}
/// A memory write with `ttl`, or `default_in_memory_ttl` when it names none.
pub fn set_local(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
let partition = self.partition(key);
Ok(partition
.memory()
.set(key, value, ttl.or(partition.default_in_memory_ttl()))?)
}
pub fn delete_local(&self, key: &str) -> Result<(), Error> {
Ok(self.partition(key).memory().delete(key)?)
}
pub fn flush_local(&self) -> Result<(), Error> {
self.key_objects.memory().flush()?;
Ok(self.objects.memory().flush()?)
}
/// Entries held by the management-object and the user-key partitions.
pub fn len_local(&self) -> Result<(usize, usize), Error> {
Ok((
self.objects.memory().len()?,
self.key_objects.memory().len()?,
))
}
/// `DualCache.async_get_cache`. A key `request` already read as missing from Redis is not read
/// again.
pub async fn get(
&self,
request: Option<&RequestRedis>,
key: &str,
) -> Result<Option<Value>, Error> {
self.partition(key).get(request, key, None).await
}
/// `DualCache.async_set_cache`, writing `json.dumps(value)`.
pub async fn set(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
self.partition(key).set(key, value, ttl).await
}
/// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`.
pub async fn set_str(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<(), Error> {
self.partition(key).set_str(key, value, ttl).await
}
/// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks.
pub fn declare_set(
&self,
send: Send<'_>,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<Pending<()>, Error> {
self.partition(key).declare_set(send, key, value, ttl)
}
/// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks.
pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result<Pending<()>, Error> {
self.partition(key).declare_delete(send, key)
}
/// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then
/// the rest in another, each with one TTL.
pub async fn set_many(
&self,
entries: Vec<(String, Value)>,
ttl: Option<Duration>,
) -> Result<(), Error> {
let (key_objects, others): (Vec<_>, Vec<_>) =
entries.into_iter().partition(|(key, _)| is_key_object(key));
self.key_objects.set_pipeline(key_objects, ttl).await?;
self.objects.set_pipeline(others, ttl).await
}
/// The auth prefetch write-back: memory first, an entry without a TTL taking
/// `default_in_memory_ttl`, then `RedisCache.async_set_cache_pipeline_with_ttls`, which encodes
/// every value before it sends anything and gives an entry without a TTL the connection's
/// default. On a request batch the `SET`s are declared and nothing is left to wait for.
/// [`Send::Now`] is one pipeline of its own rather than one round trip per entry.
pub fn write_back(
&self,
send: Send<'_>,
entries: Vec<(String, Value, Option<Duration>)>,
) -> Result<Pending<()>, Error> {
let encoded = self
.redis
.get()
.map(|_| {
entries
.iter()
.map(|(_, value, _)| json::dumps(value))
.collect::<Result<Vec<_>, _>>()
})
.transpose();
let writes = entries
.into_iter()
.map(|(key, value, ttl)| {
self.set_local(&key, value, ttl)?;
Ok((key, ttl))
})
.collect::<Result<Vec<_>, Error>>()?;
let (Some(remote), Some(encoded)) = (self.redis.get(), encoded?) else {
return Ok(ready(()));
};
let sets = writes.iter().zip(encoded).map(|((key, ttl), value)| {
commands::set(key, value, ttl.unwrap_or(remote.default_ttl))
});
let Send::Now = send else {
let batch = send.batch(&remote.client);
sets.for_each(|set| drop(batch.add(set)));
return Ok(ready(()));
};
let pipeline = RedisBatch::pipelined(remote.client.clone());
let written: Vec<_> = sets.map(|set| pipeline.add(set)).collect();
Ok(Box::pin(async move {
Declared::all(written)
.await
.into_iter()
.try_for_each(|reply| Ok(reply?))
}))
}
/// `UserApiKeyCache.async_delete_cache_keys`: the key-object partition and the rest as two
/// concurrent `DualCache.async_delete_cache_keys`. Both run even when one fails.
pub async fn delete(&self, keys: &[String]) -> Result<(), Error> {
let (key_objects, others): (Vec<String>, Vec<String>) =
keys.iter().cloned().partition(|key| is_key_object(key));
let (key_objects, others) = tokio::join!(
self.key_objects.delete_keys(&key_objects),
self.objects.delete_keys(&others)
);
key_objects.and(others)
}
/// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on
/// the batch `send` picks, noting the keys Redis lacks as the request's misses. `None` without
/// Redis. Any failed read or undecodable value fails the whole read, as a failed `MGET` does.
pub fn get_many(&self, send: Send<'_>, keys: &[String]) -> Option<Pending<Vec<Option<Value>>>> {
let client = self.redis.client()?;
let misses = send.misses(client);
let read = send.batch(client).add(commands::mget::<Vec<u8>>(
&keys.iter().map(String::as_str).collect::<Vec<_>>(),
));
let keys = keys.to_vec();
Some(Box::pin(async move {
let raw = read.resolve().await?;
if let Some(misses) = misses {
misses.note(
keys.into_iter()
.zip(&raw)
.filter(|(_, raw)| raw.is_none())
.map(|(key, _)| key),
);
}
raw.into_iter().map(|raw| Ok(types::decode(raw)?)).collect()
}))
}
/// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds.
/// `None` without Redis or for an absent entry.
pub async fn increment_if_present(
&self,
key: &str,
amount: Number,
) -> Result<Option<f64>, Error> {
let Some(client) = self.redis.client() else {
return Ok(None);
};
let amount = litellm_redis_client::Amount::from(amount);
Ok(client
.send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount))
.await?)
}
fn partition(&self, key: &str) -> &DualCache<Value> {
if is_key_object(key) {
&self.key_objects
} else {
&self.objects
}
}
}

View file

@ -0,0 +1,35 @@
use std::time::Duration;
use litellm_python_compat::Value;
use crate::{DecodeError, codec::python_value, dual_cache::MemorySettings};
/// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Settings {
pub memory_entries: usize,
pub memory_ttl: Option<Duration>,
}
impl Default for Settings {
/// `InMemoryCache`'s 200 entries.
fn default() -> Self {
Self {
memory_entries: 200,
memory_ttl: None,
}
}
}
impl Settings {
pub(super) fn memory(self) -> MemorySettings {
MemorySettings {
entries: Some(self.memory_entries),
ttl: self.memory_ttl,
}
}
}
pub(super) fn decode(raw: Option<Vec<u8>>) -> Result<Option<Value>, DecodeError> {
raw.as_deref().map(python_value).transpose()
}

View file

@ -1,23 +1,17 @@
#![expect(
dead_code,
reason = "the store families built on these tiers and codecs land in the next PR"
)]
#![cfg_attr(
not(test),
expect(
unused_imports,
reason = "the store families built on these tiers and codecs land in the next PR"
)
)]
pub mod auth_cache;
mod codec;
mod dual_cache;
mod error;
mod failure_log;
mod memory;
mod number;
pub mod rate_limits;
mod remote;
mod request;
pub mod router_cooldowns;
pub mod router_usage;
pub mod sensitive_route_pins;
pub mod spend_counters;
mod value;
pub use codec::DecodeError;

View file

@ -0,0 +1,74 @@
/// A rate-limit descriptor, `{descriptor:value}`, the hash tag every key of its window shares.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Scope {
pub descriptor: String,
pub value: String,
}
impl Scope {
pub(super) fn key(&self, suffix: &str) -> String {
format!("{{{}:{}}}:{suffix}", self.descriptor, self.value)
}
pub(super) fn window_key(&self) -> String {
self.key(WINDOW)
}
pub(super) fn gauge_key(&self) -> String {
self.key(GAUGE)
}
}
const WINDOW: &str = "window";
const GAUGE: &str = "max_parallel_requests";
/// A counter's key suffix, which is also how the Python limiter names it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)]
pub enum Kind {
#[strum(serialize = "requests")]
Requests,
#[strum(serialize = "tokens")]
Tokens,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Counter {
pub scope: Scope,
pub kind: Kind,
}
impl Counter {
pub(super) fn key(&self) -> String {
self.scope.key(self.kind.into())
}
}
/// One value of the limiter's keyspace, as both tiers address it.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum Entry {
Window(Scope),
Counter(Counter),
Gauge(Scope),
}
impl Entry {
/// The entry the Python limiter names by its key suffix: `window`, `max_parallel_requests`, or a
/// counter's `Kind`.
pub fn new(suffix: &str, scope: Scope) -> Result<Self, strum::ParseError> {
match suffix {
WINDOW => Ok(Self::Window(scope)),
GAUGE => Ok(Self::Gauge(scope)),
kind => kind
.parse()
.map(|kind| Self::Counter(Counter { scope, kind })),
}
}
pub(super) fn key(&self) -> String {
match self {
Self::Window(scope) => scope.window_key(),
Self::Counter(counter) => counter.key(),
Self::Gauge(scope) => scope.gauge_key(),
}
}
}

View file

@ -0,0 +1,260 @@
//! The limiter's Lua, byte for byte as `main`'s Python registers it, so both share one script SHA.
pub(super) const BATCH_RATE_LIMITER: &str = r#"
local results = {}
local now = tonumber(ARGV[1])
local window_size = tonumber(ARGV[2])
-- Process each window/counter pair
for i = 1, #KEYS, 2 do
local window_key = KEYS[i]
local counter_key = KEYS[i + 1]
local increment_value = 1
-- Check if window exists and is valid
local window_start = redis.call('GET', window_key)
if not window_start or (now - tonumber(window_start)) >= window_size then
-- Reset window and counter
local prefix = string.sub(window_key, 1, -(#':window') - 1)
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, increment_value)
redis.call('EXPIRE', window_key, window_size)
redis.call('EXPIRE', counter_key, window_size)
table.insert(results, tostring(now)) -- window_start
table.insert(results, increment_value) -- counter
else
local counter = redis.call('INCR', counter_key)
-- This happens when window_key exists but counter_key doesn't (e.g., tokens key
-- created after requests key when both share the same window_key)
local current_ttl = redis.call('TTL', counter_key)
if current_ttl == -1 then
redis.call('EXPIRE', counter_key, window_size)
end
table.insert(results, window_start) -- window_start
table.insert(results, counter) -- counter
end
end
return results
"#;
pub(super) const BATCH_COUNTER_READ: &str = r#"return redis.call('MGET', unpack(KEYS))"#;
pub(super) const CHECK_AND_INCREMENT_BY_N: &str = r#"
-- Atomic check-and-increment-by-N across one or more descriptors.
-- All-or-nothing: if any descriptor would exceed its limit, no counter is
-- modified.
--
-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied
-- timestamp so that window resets are deterministic across replicas with
-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the
-- TOCTOU window across multi-replica deployments.
--
-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor.
-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]:
-- ARGV[(i-1)*4 + 1] = limit
-- ARGV[(i-1)*4 + 2] = increment
-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets)
-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length)
--
-- Return on success:
-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... }
-- Return on over-limit: { 1, descriptor_index, current_counter, limit }
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
local descriptor_count = #KEYS / 2
local reset_windows = {}
-- Pass 1: read state, validate. Abort without writing if any over limit.
local descriptor_state = {}
for i = 1, descriptor_count do
local window_key = KEYS[(i - 1) * 2 + 1]
local counter_key = KEYS[(i - 1) * 2 + 2]
local arg_base = (i - 1) * 4 + 1
local limit = tonumber(ARGV[arg_base])
local increment = tonumber(ARGV[arg_base + 1])
local window_size = tonumber(ARGV[arg_base + 3])
local window_start = redis.call('GET', window_key)
local window_expired = (not window_start) or
((now - tonumber(window_start)) >= window_size)
local current_counter
if window_expired then
current_counter = 0
else
current_counter = tonumber(redis.call('GET', counter_key) or 0)
end
local blocked
if increment > 0 then
blocked = current_counter + increment > limit
else
blocked = current_counter >= limit
end
if blocked then
return { 1, i, current_counter, limit }
end
descriptor_state[i] = { window_expired, current_counter, window_start }
end
-- Pass 2: all checks passed. Apply increments.
local results = { 0 }
for i = 1, descriptor_count do
local window_key = KEYS[(i - 1) * 2 + 1]
local counter_key = KEYS[(i - 1) * 2 + 2]
local arg_base = (i - 1) * 4 + 1
local increment = tonumber(ARGV[arg_base + 1])
local ttl = tonumber(ARGV[arg_base + 2])
local window_size = tonumber(ARGV[arg_base + 3])
local window_expired = descriptor_state[i][1]
local active_window_start
if window_expired then
active_window_start = now
if not reset_windows[window_key] then
local prefix = string.sub(window_key, 1, -(#':window') - 1)
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
reset_windows[window_key] = true
end
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, increment)
redis.call('EXPIRE', window_key, window_size)
if ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, increment)
else
active_window_start = tonumber(descriptor_state[i][3])
local new_counter = redis.call('INCRBY', counter_key, increment)
local current_ttl = redis.call('TTL', counter_key)
if current_ttl == -1 and ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, new_counter)
end
table.insert(results, active_window_start)
end
return results
"#;
pub(super) const WINDOW_GUARDED_TOKEN_INCREMENT: &str = r#"
local results = {}
for i = 1, #KEYS, 2 do
local window_key = KEYS[i]
local counter_key = KEYS[i + 1]
local arg_base = ((i - 1) / 2) * 3 + 1
local expected_window_start = ARGV[arg_base]
local increment = tonumber(ARGV[arg_base + 1])
local ttl = tonumber(ARGV[arg_base + 2])
local active_window_start = redis.call('GET', window_key)
if active_window_start and active_window_start == expected_window_start then
local new_counter = redis.call('INCRBY', counter_key, increment)
local current_ttl = redis.call('TTL', counter_key)
if current_ttl == -1 and ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, 1)
table.insert(results, new_counter)
else
table.insert(results, 0)
table.insert(results, tonumber(redis.call('GET', counter_key) or 0))
end
end
return results
"#;
pub(super) const PARALLEL_ACQUIRE: &str = r#"
-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge.
-- Each gauge key is a sorted set of per-request slot ids scored by acquire
-- time (Redis server clock). In-flight requests are counted by ZCARD after
-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM
-- counters the gauge is never reset while requests are in flight, a
-- rejected request never occupies a slot, and a slot leaked by a crashed
-- worker self-heals after the slot TTL even under continuous traffic.
--
-- KEYS: one gauge zset key per descriptor.
-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id).
-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }.
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
for i = 1, #KEYS do
local limit = tonumber(ARGV[(i - 1) * 3 + 1])
local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2])
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl)
local in_flight = redis.call('ZCARD', KEYS[i])
if in_flight + 1 > limit then
return { 1, i, in_flight, limit }
end
end
local results = { 0 }
for i = 1, #KEYS do
local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2])
local slot_id = ARGV[(i - 1) * 3 + 3]
redis.call('ZADD', KEYS[i], now, slot_id)
redis.call('EXPIRE', KEYS[i], slot_ttl)
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"#;
pub(super) const PARALLEL_RELEASE: &str = r#"
-- Release one slot per gauge key by removing this request's slot id.
-- ZREM of an absent member (or key) is a no-op, so a release without a
-- matching acquire (proxy-side rejection, double-fired callback, slot
-- already expired) can never free a slot owned by another request.
-- KEYS: gauge zset keys. ARGV: per-key slot_id.
-- Returns the remaining in-flight count per key.
local results = {}
for i = 1, #KEYS do
redis.call('ZREM', KEYS[i], ARGV[i])
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"#;
pub(super) const PARALLEL_COUNT: &str = r#"
-- Read the current in-flight count per gauge key (prunes expired slots
-- first so leaked slots do not inflate the reading).
-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds.
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
local results = {}
for i = 1, #KEYS do
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i]))
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"#;
pub(super) const TOKEN_INCREMENT: &str = r#"
local results = {}
-- Process each key/increment_value/ttl triplet
for i = 1, #KEYS do
local key = KEYS[i]
local increment_value = tonumber(ARGV[i * 2 - 1])
local ttl_seconds = tonumber(ARGV[i * 2])
-- Increment the value
local new_value = redis.call('INCRBYFLOAT', key, increment_value)
-- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL
-- ttl_seconds can be 0 (no TTL) or positive (set TTL)
if ttl_seconds and ttl_seconds > 0 then
local current_ttl = redis.call('TTL', key)
if current_ttl == -1 then
redis.call('EXPIRE', key, ttl_seconds)
end
end
table.insert(results, new_value)
end
return results
"#;

View file

@ -0,0 +1,304 @@
mod keys;
mod lua;
mod scripts;
mod types;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::{Declared, RedisBatch, RedisClient, Slot};
pub use self::keys::{Counter, Entry, Kind, Scope};
pub use self::types::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply,
Increment, Local, WindowCount,
};
use crate::{
Error, Number, Pending, RequestRedis, Send,
dual_cache::{DualCache, MemorySettings, PipelineIncrement},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
request::pending,
};
/// The v3 rate limiter's keyspace: its memory tier, its Lua scripts and the plain counter writes
/// `DualCache` made for it on `main`. The admission logic stays with the caller. A script needs
/// Redis, so without one it declares nothing and answers `None`.
pub struct RateLimits {
cache: DualCache<Local>,
}
impl RateLimits {
/// `backfill_ttl` is the `default_in_memory_ttl` of the `DualCache` this tier replaces. An
/// increment without a TTL takes the remote's default.
pub fn new(remote: Option<Remote>, backfill_ttl: Option<Duration>) -> Self {
Self::with_clock(remote, backfill_ttl, system_clock())
}
pub fn with_clock(remote: Option<Remote>, backfill_ttl: Option<Duration>, now: Clock) -> Self {
Self {
cache: DualCache::new(
Arc::new(RemoteSlot::new(remote)),
MemorySettings {
entries: None,
ttl: backfill_ttl,
},
now,
),
}
}
pub fn has_redis(&self) -> bool {
self.cache.redis().get().is_some()
}
/// `_group_keys_by_hash_tag`: indices of `scopes` sharing a cluster slot, in first-seen order,
/// or one group when the client is not a cluster and grouping is not forced.
pub fn slot_groups(&self, scopes: &[Scope], force: bool) -> Vec<Vec<usize>> {
let by_slot = force
|| self
.cache
.redis()
.client()
.is_some_and(RedisClient::is_cluster);
if !by_slot {
return vec![(0..scopes.len()).collect()];
}
scopes
.iter()
.enumerate()
.fold(
Vec::<(Slot, Vec<usize>)>::new(),
|mut groups, (index, scope)| {
let slot = Slot::for_key(scope.window_key());
match groups.iter_mut().find(|(seen, _)| *seen == slot) {
Some((_, members)) => members.push(index),
None => groups.push((slot, vec![index])),
}
groups
},
)
.into_iter()
.map(|(_, members)| members)
.collect()
}
pub fn get_local(&self, entry: &Entry) -> Result<Option<Local>, Error> {
Ok(self.cache.memory().get(&entry.key())?)
}
pub fn set_local(
&self,
entry: &Entry,
value: Local,
ttl: Option<Duration>,
) -> Result<(), Error> {
Ok(self.cache.memory().set(&entry.key(), value, ttl)?)
}
/// `DualCache.async_get_cache`, backfilling memory for `backfill_ttl` or the store's default.
/// A key `request` already read as missing from Redis is not read again.
pub async fn get(
&self,
request: Option<&RequestRedis>,
entry: &Entry,
backfill_ttl: Option<Duration>,
) -> Result<Option<Local>, Error> {
self.cache.get(request, &entry.key(), backfill_ttl).await
}
pub fn sliding_window(
&self,
send: Send<'_>,
counters: &[Counter],
now: i64,
window_size: i64,
) -> Option<Pending<Vec<WindowCount>>> {
Some(pending(scripts::sliding_window(
&self.batch(send)?,
counters,
now,
window_size,
)))
}
pub fn read_counters(
&self,
send: Send<'_>,
counters: &[Counter],
) -> Option<Pending<Vec<Option<String>>>> {
Some(pending(scripts::read_counters(
&self.batch(send)?,
counters,
)))
}
pub fn check_and_increment(
&self,
send: Send<'_>,
counters: &[AtomicCounter],
) -> Option<Pending<AtomicReply>> {
Some(pending(scripts::check_and_increment(
&self.batch(send)?,
counters,
)))
}
pub fn guarded_increment(
&self,
send: Send<'_>,
increment: &GuardedIncrement,
) -> Option<Pending<GuardedReply>> {
Some(pending(scripts::guarded_increment(
&self.batch(send)?,
increment,
)))
}
pub fn parallel_acquire(
&self,
send: Send<'_>,
gauges: &[Gauge],
slot_ttl: Duration,
slot_id: &str,
) -> Option<Pending<AcquireReply>> {
Some(pending(scripts::parallel_acquire(
&self.batch(send)?,
gauges,
slot_ttl,
slot_id,
)))
}
pub fn parallel_release(
&self,
send: Send<'_>,
gauges: &[Scope],
slot_id: &str,
) -> Option<Pending<Vec<i64>>> {
Some(pending(self.declare_release(send, gauges, slot_id)?))
}
pub fn parallel_count(
&self,
send: Send<'_>,
gauges: &[Scope],
slot_ttl: Duration,
) -> Option<Pending<Vec<i64>>> {
Some(pending(scripts::parallel_count(
&self.batch(send)?,
gauges,
slot_ttl,
)))
}
pub fn token_increment(
&self,
send: Send<'_>,
increments: &[Increment],
) -> Option<Pending<Vec<f64>>> {
Some(pending(self.declare_token_increment(send, increments)?))
}
/// `TOKEN_INCREMENT_SCRIPT` on the request's post-call batch. A failed script falls back to
/// `increment_many` for these increments once the batch has gone out. `false` without Redis.
pub fn token_increment_post_call(
self: &Arc<Self>,
request: &RequestRedis,
increments: Vec<Increment>,
) -> bool {
let Some(declared) = self.declare_token_increment(Send::PostCall(request), &increments)
else {
return false;
};
let store = Arc::clone(self);
request.spawn(async move {
let Err(error) = declared.settled().await else {
return;
};
crate::failure_log::redis_failure!(
&error,
"TTL preservation failed, falling back to regular pipeline"
);
store.increment_many(&increments).await;
});
true
}
/// `PARALLEL_RELEASE_SCRIPT` on the request's post-call batch; the caller has already freed
/// the slot in memory. `false` without Redis.
pub fn parallel_release_post_call(
&self,
request: &RequestRedis,
gauges: &[Scope],
slot_id: &str,
) -> bool {
let Some(declared) = self.declare_release(Send::PostCall(request), gauges, slot_id) else {
return false;
};
request.spawn(async move {
if let Err(error) = declared.settled().await {
crate::failure_log::redis_failure!(
&error,
"parallel_release_script failed, the slot stays released in memory only"
);
}
});
true
}
/// `DualCache.async_increment_cache`. A Redis failure answers with the memory count, and an
/// entry memory cannot add to with `None`.
pub async fn increment(&self, increment: &Increment) -> Option<Number> {
self.cache
.increment(&increment.counter.key(), increment.amount, increment.ttl)
.await
}
/// `DualCache.async_increment_cache_pipeline`, whose failures only log.
pub async fn increment_many(&self, increments: &[Increment]) {
self.cache.increment_pipeline(&pipeline(increments)).await;
}
/// `DualCache.async_increment_cache_pipeline_post_call`: memory now, the Redis increments on
/// the request's post-call batch.
pub fn increment_many_post_call(&self, request: &RequestRedis, increments: &[Increment]) {
self.cache
.increment_pipeline_post_call(request, &pipeline(increments));
}
fn declare_release(
&self,
send: Send<'_>,
gauges: &[Scope],
slot_id: &str,
) -> Option<Declared<Vec<i64>>> {
Some(scripts::parallel_release(
&self.batch(send)?,
gauges,
slot_id,
))
}
fn declare_token_increment(
&self,
send: Send<'_>,
increments: &[Increment],
) -> Option<Declared<Vec<f64>>> {
Some(scripts::token_increment(&self.batch(send)?, increments))
}
fn batch(&self, send: Send<'_>) -> Option<RedisBatch> {
self.cache.redis().batch(send)
}
}
fn pipeline(increments: &[Increment]) -> Vec<PipelineIncrement> {
increments
.iter()
.map(|increment| PipelineIncrement {
key: increment.counter.key(),
amount: increment.amount,
ttl: increment.ttl,
})
.collect()
}

View file

@ -0,0 +1,303 @@
use std::{sync::LazyLock, time::Duration};
use litellm_python_compat::repr::float_repr;
use litellm_redis_client::{
Command, Declared, Error, RedisBatch, Script, Value,
commands::{self, unexpected_reply},
};
use super::{
keys::{Counter, Scope},
lua,
types::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply,
Increment, WindowCount,
},
};
use crate::Number;
static BATCH_RATE_LIMITER: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::BATCH_RATE_LIMITER));
static BATCH_COUNTER_READ: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::BATCH_COUNTER_READ));
static CHECK_AND_INCREMENT_BY_N: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::CHECK_AND_INCREMENT_BY_N));
static WINDOW_GUARDED_TOKEN_INCREMENT: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::WINDOW_GUARDED_TOKEN_INCREMENT));
static PARALLEL_ACQUIRE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_ACQUIRE));
static PARALLEL_RELEASE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_RELEASE));
static PARALLEL_COUNT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_COUNT));
static TOKEN_INCREMENT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::TOKEN_INCREMENT));
pub(super) fn sliding_window(
batch: &RedisBatch,
counters: &[Counter],
now: i64,
window_size: i64,
) -> Declared<Vec<WindowCount>> {
let keys = counters
.iter()
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
.collect::<Vec<_>>();
batch.add(
eval(
&BATCH_RATE_LIMITER,
&keys,
vec![now.to_string(), window_size.to_string()],
)
.map(window_counts),
)
}
/// `BATCH_COUNTER_READ_SCRIPT` over each counter's window and counter keys, in that order.
pub(super) fn read_counters(
batch: &RedisBatch,
counters: &[Counter],
) -> Declared<Vec<Option<String>>> {
let keys = counters
.iter()
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
.collect::<Vec<_>>();
batch.add(eval(&BATCH_COUNTER_READ, &keys, Vec::new()).map(optional_texts))
}
pub(super) fn check_and_increment(
batch: &RedisBatch,
counters: &[AtomicCounter],
) -> Declared<AtomicReply> {
let keys = counters
.iter()
.flat_map(|counter| [counter.counter.scope.window_key(), counter.counter.key()])
.collect::<Vec<_>>();
let args = counters
.iter()
.flat_map(|counter| {
[
counter.limit,
counter.increment,
counter.ttl_seconds,
counter.window_size_seconds,
]
.map(|arg| arg.to_string())
})
.collect();
batch.add(eval(&CHECK_AND_INCREMENT_BY_N, &keys, args).map(atomic_reply))
}
pub(super) fn guarded_increment(
batch: &RedisBatch,
increment: &GuardedIncrement,
) -> Declared<GuardedReply> {
let keys = [
increment.counter.scope.window_key(),
increment.counter.key(),
];
let args = vec![
increment.expected_window_start.clone(),
python_arg(increment.amount),
ttl_arg(increment.ttl),
];
batch.add(eval(&WINDOW_GUARDED_TOKEN_INCREMENT, &keys, args).map(guarded_reply))
}
pub(super) fn parallel_acquire(
batch: &RedisBatch,
gauges: &[Gauge],
slot_ttl: Duration,
slot_id: &str,
) -> Declared<AcquireReply> {
let keys = gauges
.iter()
.map(|gauge| gauge.scope.gauge_key())
.collect::<Vec<_>>();
let args = gauges
.iter()
.flat_map(|gauge| {
[
gauge.limit.to_string(),
slot_ttl.as_secs().to_string(),
slot_id.to_owned(),
]
})
.collect();
batch.add(eval(&PARALLEL_ACQUIRE, &keys, args).map(acquire_reply))
}
pub(super) fn parallel_release(
batch: &RedisBatch,
gauges: &[Scope],
slot_id: &str,
) -> Declared<Vec<i64>> {
let keys = gauges.iter().map(Scope::gauge_key).collect::<Vec<_>>();
let args = gauges.iter().map(|_| slot_id.to_owned()).collect();
batch.add(eval(&PARALLEL_RELEASE, &keys, args).map(integers))
}
pub(super) fn parallel_count(
batch: &RedisBatch,
gauges: &[Scope],
slot_ttl: Duration,
) -> Declared<Vec<i64>> {
let keys = gauges.iter().map(Scope::gauge_key).collect::<Vec<_>>();
let args = gauges
.iter()
.map(|_| slot_ttl.as_secs().to_string())
.collect();
batch.add(eval(&PARALLEL_COUNT, &keys, args).map(integers))
}
pub(super) fn token_increment(batch: &RedisBatch, increments: &[Increment]) -> Declared<Vec<f64>> {
let keys = increments
.iter()
.map(|increment| increment.counter.key())
.collect::<Vec<_>>();
let args = increments
.iter()
.flat_map(|increment| [python_arg(increment.amount), ttl_arg(increment.ttl)])
.collect();
batch.add(eval(&TOKEN_INCREMENT, &keys, args).map(floats))
}
/// `RedisBatch`'s `_Increment`: `INCRBYFLOAT`, then `EXPIRE` when the increment carries a TTL. On
/// its own it keeps a TTL the key already has, like `async_increment`.
fn eval(script: &Script, keys: &[String], args: Vec<String>) -> Command<Value> {
let keys = keys.iter().map(String::as_str).collect::<Vec<_>>();
commands::eval(script, &keys, args)
}
/// The bytes redis-py sends for a Python `int` or `float` argument.
fn python_arg(amount: Number) -> String {
match amount {
Number::Int(value) => value.to_string(),
Number::Float(value) => float_repr(value),
}
}
/// `ttl or 0`: the scripts read a zero TTL as "leave the expiry alone".
fn ttl_arg(ttl: Option<Duration>) -> String {
ttl.map_or(0, |ttl| ttl.as_secs()).to_string()
}
fn items(reply: Value) -> Result<Vec<Value>, Error> {
match reply {
Value::Array(items) => Ok(items),
_ => Err(unexpected_reply("rate limiter script reply is not a list")),
}
}
fn integer(value: &Value) -> Result<i64, Error> {
match value {
Value::Int(value) => Ok(*value),
_ => Err(unexpected_reply(
"rate limiter script reply holds a non-integer where a count belongs",
)),
}
}
fn text(value: &Value) -> Result<Option<String>, Error> {
match value {
Value::Nil => Ok(None),
Value::BulkString(bytes) => String::from_utf8(bytes.clone()).map(Some).map_err(|_| {
unexpected_reply("rate limiter script reply holds a value that is not UTF-8")
}),
Value::SimpleString(text) => Ok(Some(text.clone())),
Value::Int(value) => Ok(Some(value.to_string())),
_ => Err(unexpected_reply(
"rate limiter script reply holds a nested value",
)),
}
}
fn float(reply: Value) -> Result<f64, Error> {
text(&reply)?
.and_then(|text| text.parse().ok())
.ok_or_else(|| unexpected_reply("rate limiter script reply is not a float"))
}
fn index(raw: &Value) -> Result<usize, Error> {
usize::try_from(integer(raw)? - 1)
.map_err(|_| unexpected_reply("rate limiter script reply names an index below 1"))
}
fn window_counts(reply: Value) -> Result<Vec<WindowCount>, Error> {
items(reply)?
.chunks(2)
.map(|pair| match pair {
[window_start, count] => Ok(WindowCount {
window_start: text(window_start)?.ok_or_else(|| {
unexpected_reply("rate limiter script reply holds a nil window start")
})?,
count: integer(count)?,
}),
_ => Err(unexpected_reply(
"rate limiter script reply has an odd length",
)),
})
.collect()
}
fn optional_texts(reply: Value) -> Result<Vec<Option<String>>, Error> {
items(reply)?.iter().map(text).collect()
}
fn integers(reply: Value) -> Result<Vec<i64>, Error> {
items(reply)?.iter().map(integer).collect()
}
fn floats(reply: Value) -> Result<Vec<f64>, Error> {
items(reply)?.into_iter().map(float).collect()
}
fn atomic_reply(reply: Value) -> Result<AtomicReply, Error> {
let items = items(reply)?;
match items.split_first() {
Some((Value::Int(0), rest)) => rest
.chunks(2)
.map(|pair| match pair {
[count, window_start] => Ok(Applied {
count: integer(count)?,
window_start: integer(window_start)?,
}),
_ => Err(unexpected_reply(
"rate limiter script reply has an odd length",
)),
})
.collect::<Result<_, _>>()
.map(AtomicReply::Applied),
Some((Value::Int(1), [index_value, current, limit])) => Ok(AtomicReply::OverLimit {
index: index(index_value)?,
current: integer(current)?,
limit: integer(limit)?,
}),
_ => Err(unexpected_reply("rate limiter script reply has no status")),
}
}
fn guarded_reply(reply: Value) -> Result<GuardedReply, Error> {
match items(reply)?.as_slice() {
[applied, count] => Ok(GuardedReply {
applied: integer(applied)? == 1,
count: integer(count)?,
}),
_ => Err(unexpected_reply(
"rate limiter script reply is not an (applied, count) pair",
)),
}
}
fn acquire_reply(reply: Value) -> Result<AcquireReply, Error> {
let items = items(reply)?;
match items.split_first() {
Some((Value::Int(0), rest)) => rest
.iter()
.map(integer)
.collect::<Result<_, _>>()
.map(AcquireReply::Acquired),
Some((Value::Int(1), [index_value, in_flight, limit])) => Ok(AcquireReply::OverLimit {
index: index(index_value)?,
in_flight: integer(in_flight)?,
limit: integer(limit)?,
}),
_ => Err(unexpected_reply("rate limiter script reply has no status")),
}
}

View file

@ -0,0 +1,117 @@
use std::time::Duration;
use litellm_python_compat::Value;
use super::keys::{Counter, Scope};
use crate::{DecodeError, FromValue, Number};
/// `BATCH_RATE_LIMITER_SCRIPT`'s pair for one counter: the window start it saw or opened and the
/// counter after its `+1`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WindowCount {
pub window_start: String,
pub count: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct AtomicCounter {
pub counter: Counter,
pub limit: i64,
pub increment: i64,
pub ttl_seconds: i64,
pub window_size_seconds: i64,
}
/// `CHECK_AND_INCREMENT_BY_N_SCRIPT`'s reply. `index` is 0-based where the Lua's is 1-based.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AtomicReply {
Applied(Vec<Applied>),
OverLimit {
index: usize,
current: i64,
limit: i64,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Applied {
pub count: i64,
pub window_start: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct GuardedIncrement {
pub counter: Counter,
pub expected_window_start: String,
pub amount: Number,
pub ttl: Option<Duration>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct GuardedReply {
pub applied: bool,
pub count: i64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Gauge {
pub scope: Scope,
pub limit: i64,
}
/// `PARALLEL_ACQUIRE_SCRIPT`'s reply. `index` is 0-based where the Lua's is 1-based.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AcquireReply {
Acquired(Vec<i64>),
OverLimit {
index: usize,
in_flight: i64,
limit: i64,
},
}
/// What the memory tier holds: counters and mirrored gauge counts, window starts, and the per-slot
/// registry a gauge keeps while Redis is unreachable.
#[derive(Clone, Debug, PartialEq)]
pub enum Local {
Number(Number),
Text(String),
Slots(Vec<(String, f64)>),
}
impl From<Number> for Local {
fn from(number: Number) -> Self {
Self::Number(number)
}
}
impl crate::memory::Counter for Local {
fn as_number(&self) -> Option<Number> {
match self {
Self::Number(number) => Some(*number),
Self::Text(_) | Self::Slots(_) => None,
}
}
}
/// A `RedisPipelineIncrementOperation`.
#[derive(Clone, Debug, PartialEq)]
pub struct Increment {
pub counter: Counter,
pub amount: Number,
pub ttl: Option<Duration>,
}
/// A rate-limit value as `RedisCache._get_cache_logic` reads it.
impl FromValue for Local {
fn from_value(value: &Value) -> Result<Self, DecodeError> {
match value {
Value::Str(text) => Ok(Self::Text(text.clone())),
value => Number::from_value(value)
.map(Local::Number)
.ok_or(DecodeError::Shape(
"a rate-limit value is a number or a string",
)),
}
}
}

View file

@ -0,0 +1,17 @@
pub(super) fn cooldown(model_id: &str) -> String {
format!("deployment:{model_id}:cooldown")
}
pub(super) fn cooldowns(model_ids: &[String]) -> Vec<String> {
model_ids
.iter()
.map(|model_id| cooldown(model_id))
.collect()
}
pub(super) fn allowed_fails(model_id: &str, suffix: Option<&str>) -> String {
match suffix.filter(|suffix| !suffix.is_empty()) {
Some(suffix) => format!("deployment:{model_id}:allowed_fails:{suffix}"),
None => format!("deployment:{model_id}:allowed_fails"),
}
}

View file

@ -0,0 +1,205 @@
mod keys;
mod types;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::{Amount, Command, Expiry, commands};
pub use self::types::Cooldown;
use crate::{
Error, Number, Pending, Send,
dual_cache::{BatchReads, DualCache, TierPolicy},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
};
const MAX_CORRECTED_TTL: Duration = Duration::from_secs(60);
const CORRECTION_SLACK: Duration = Duration::from_secs(5);
/// One router's cooldowns and allowed-fails counters: memory per router, Redis shared by every pod.
pub struct RouterCooldowns {
cooldowns: DualCache<Cooldown, BatchReads>,
fails: DualCache<i64>,
redis: Arc<RemoteSlot>,
}
impl RouterCooldowns {
pub fn new(remote: Option<Remote>, policy: TierPolicy) -> Self {
Self::with_clock(remote, policy, system_clock())
}
pub fn with_clock(remote: Option<Remote>, policy: TierPolicy, now: Clock) -> Self {
let redis = Arc::new(RemoteSlot::new(remote));
Self {
fails: DualCache::new(redis.clone(), policy.memory(), now.clone()),
cooldowns: DualCache::batch_reading(redis.clone(), policy, now),
redis,
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.redis.attach(remote)
}
pub fn has_redis(&self) -> bool {
self.redis.get().is_some()
}
/// `DualCache.async_set_cache`: memory first, then a `SET` of `json.dumps(value)` with
/// `EX int(ttl)` whose failure only logs.
pub async fn put(
&self,
model_id: &str,
cooldown: &Cooldown,
ttl: Duration,
) -> Result<(), Error> {
self.cooldowns
.set(&keys::cooldown(model_id), cooldown.clone(), Some(ttl))
.await
}
/// `DualCache.set_cache`: memory first, then a `SET` of `str(value)` with `EX int(ttl)` whose
/// failure only logs.
pub async fn put_str(
&self,
model_id: &str,
cooldown: &Cooldown,
ttl: Duration,
) -> Result<(), Error> {
self.cooldowns
.set_str(&keys::cooldown(model_id), cooldown.clone(), Some(ttl))
.await
}
/// `DualCache.async_increment_cache`: the fleet-wide fail count through `INCRBYFLOAT`. A Redis
/// failure falls back to this router's own count.
pub async fn increment_allowed_fails(
&self,
model_id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<i64, Error> {
let key = keys::allowed_fails(model_id, suffix);
let command =
commands::incr_by_float(&key, Amount::Int(1), Some(Expiry::KeepExisting(ttl)));
self.count_fail(&key, ttl, command, Number::Float).await
}
/// `DualCache.increment_cache`: the fleet-wide fail count through `INCRBY`. A Redis failure
/// falls back to this router's own count.
pub async fn increment_allowed_fails_int(
&self,
model_id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<i64, Error> {
let key = keys::allowed_fails(model_id, suffix);
let command = commands::incr_by(&key, 1, ttl);
self.count_fail(&key, ttl, command, Number::Int).await
}
async fn count_fail<T>(
&self,
key: &str,
ttl: Duration,
command: Command<T>,
total: fn(T) -> Number,
) -> Result<i64, Error> {
let local = self
.fails
.memory()
.update(key, Some(ttl), |count| Ok(count.unwrap_or(0) + 1))?;
let Some(client) = self.redis.client() else {
return Ok(local);
};
match client.send(command).await {
Ok(count) => Ok(total(count).truncate()),
Err(error) => {
crate::failure_log::redis_failure!(
&error,
"allowed_fails counter fell back to this worker's in-memory count"
);
Ok(local)
}
}
}
/// The deployments still cooling down, read on the batch `send` picks. An entry whose cooldown
/// already ended is dropped from memory, and one memory would keep past its end is shortened, at
/// most to 60 s.
pub fn active(
self: &Arc<Self>,
send: Send<'_>,
model_ids: &[String],
) -> Result<Pending<Vec<(String, Cooldown)>>, Error> {
let read = self
.cooldowns
.declare_batch_get(send, &keys::cooldowns(model_ids))?;
let model_ids = model_ids.to_vec();
let store = Arc::clone(self);
Ok(Box::pin(async move {
let read = store.cooldowns.resolve_batch_get(read).await?;
let now = store.cooldowns.now();
let active = model_ids
.into_iter()
.zip(read)
.filter_map(|(model_id, cooldown)| cooldown.map(|cooldown| (model_id, cooldown)))
.map(|(model_id, cooldown)| {
Ok(store
.correct(&model_id, &cooldown, now)?
.then_some((model_id, cooldown)))
})
.collect::<Result<Vec<_>, Error>>()?;
Ok(active.into_iter().flatten().collect())
}))
}
/// The shortest stored cooldown, ended ones included, as `main`'s `get_min_cooldown` reads it.
pub async fn minimum(&self, model_ids: &[String]) -> Result<Option<Number>, Error> {
let read = self
.cooldowns
.batch_get(&keys::cooldowns(model_ids))
.await?;
Ok(Number::min(
read.into_iter()
.flatten()
.map(|cooldown| cooldown.cooldown_time),
))
}
pub fn delete_local(&self, model_id: &str) -> Result<(), Error> {
Ok(self.cooldowns.memory().delete(&keys::cooldown(model_id))?)
}
/// Drops this router's allowed-fails counts, leaving its cooldowns in place.
pub fn flush_local_counters(&self) -> Result<(), Error> {
Ok(self.fails.memory().flush()?)
}
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
/// in-memory router state.
pub fn snapshot(&self) -> Result<Vec<(String, Cooldown)>, Error> {
Ok(self.cooldowns.memory().entries()?)
}
fn correct(&self, model_id: &str, cooldown: &Cooldown, now: Duration) -> Result<bool, Error> {
let key = keys::cooldown(model_id);
let memory = self.cooldowns.memory();
let remaining = cooldown.remaining(now);
if remaining <= 0.0 {
memory.delete(&key)?;
return Ok(false);
}
let remaining = Duration::from_secs_f64(remaining);
if memory.expires_at(&key)? > Some(now + remaining + CORRECTION_SLACK) {
memory.delete(&key)?;
memory.set(
&key,
cooldown.clone(),
Some(remaining.min(MAX_CORRECTED_TTL)),
)?;
}
Ok(true)
}
}

View file

@ -0,0 +1,61 @@
use std::time::Duration;
use litellm_python_compat::Value;
use crate::{DecodeError, FromValue, IntoValue, Number};
#[derive(Clone, Debug, PartialEq)]
pub struct Cooldown {
pub exception_received: String,
pub status_code: String,
pub timestamp: Number,
pub cooldown_time: Number,
}
impl Cooldown {
pub(super) fn remaining(&self, now: Duration) -> f64 {
self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64()
}
}
/// The `CooldownCacheValue` dict the router reads and writes.
impl IntoValue for Cooldown {
fn into_value(self) -> Value {
let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value);
Value::Dict(vec![
field("exception_received", Value::Str(self.exception_received)),
field("status_code", Value::Str(self.status_code)),
field("timestamp", self.timestamp.to_value()),
field("cooldown_time", self.cooldown_time.to_value()),
])
}
}
impl FromValue for Cooldown {
fn from_value(value: &Value) -> Result<Self, DecodeError> {
let Value::Dict(fields) = value else {
return Err(DecodeError::Shape("a cooldown is a dict"));
};
let field = |name: &str| {
fields
.iter()
.find(|(key, _)| matches!(key, Value::Str(key) if key == name))
.map(|(_, value)| value)
};
let text = |name: &'static str| match field(name) {
Some(Value::Str(text)) => Ok(text.clone()),
_ => Err(DecodeError::Shape(name)),
};
let number = |name: &'static str| {
field(name)
.and_then(Number::from_value)
.ok_or(DecodeError::Shape(name))
};
Ok(Self {
exception_received: text("exception_received")?,
status_code: text("status_code")?,
timestamp: number("timestamp")?,
cooldown_time: number("cooldown_time")?,
})
}
}

View file

@ -0,0 +1,28 @@
/// A counter's key segment, which is also how the Python router names it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)]
pub enum Kind {
#[strum(serialize = "tpm")]
Tpm,
#[strum(serialize = "rpm")]
Rpm,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Counter {
pub deployment_id: String,
pub model: Option<String>,
pub minute: String,
pub kind: Kind,
}
impl Counter {
/// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`.
pub(super) fn key(&self) -> String {
let kind: &str = self.kind.into();
let model = self.model.as_deref().unwrap_or("None");
format!(
"global_router:{}:{model}:{kind}:{}",
self.deployment_id, self.minute
)
}
}

View file

@ -0,0 +1,110 @@
mod keys;
use std::{sync::Arc, time::Duration};
pub use self::keys::{Counter, Kind};
use crate::{
Error, Number,
dual_cache::{BatchReads, DualCache, PipelineIncrement, TierPolicy},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
};
/// `RoutingArgs.ttl`, the window of a per-minute counter.
pub const COUNTER_TTL: Duration = Duration::from_secs(60);
/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod.
pub struct RouterUsage {
usage: DualCache<Number, BatchReads>,
}
impl RouterUsage {
pub fn new(remote: Option<Remote>, policy: TierPolicy) -> Self {
Self::with_clock(remote, policy, system_clock())
}
pub fn with_clock(remote: Option<Remote>, policy: TierPolicy, now: Clock) -> Self {
Self {
usage: DualCache::batch_reading(Arc::new(RemoteSlot::new(remote)), policy, now),
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.usage.redis().attach(remote)
}
pub fn has_redis(&self) -> bool {
self.usage.redis().get().is_some()
}
/// `DualCache.async_increment_cache`. A Redis failure answers with the memory count, and an
/// entry memory cannot add to with `None`.
pub async fn increment(
&self,
counter: &Counter,
value: Number,
ttl: Duration,
) -> Option<Number> {
self.usage.increment(&counter.key(), value, Some(ttl)).await
}
/// One deployment's tokens and requests, `async_increment_pipeline` on its own round trip, with
/// the totals mirrored back into memory as `int`s.
pub async fn increment_deployment(
&self,
deployment_id: &str,
model: Option<&str>,
minute: &str,
tokens: Number,
requests: i64,
) -> Result<(), Error> {
let increment = |kind, amount| PipelineIncrement {
key: Counter {
deployment_id: deployment_id.to_owned(),
model: model.map(str::to_owned),
minute: minute.to_owned(),
kind,
}
.key(),
amount,
ttl: Some(COUNTER_TTL),
};
let increments = [
increment(Kind::Tpm, tokens),
increment(Kind::Rpm, Number::Int(requests)),
];
let Some(totals) = self.usage.increment_pipeline(&increments).await else {
return Ok(());
};
if self.usage.redis().get().is_none() {
return Ok(());
}
increments
.iter()
.zip(totals)
.try_for_each(|(increment, total)| {
self.usage.memory().set(
&increment.key,
Number::Int(total.truncate()),
Some(COUNTER_TTL),
)
})?;
Ok(())
}
pub async fn get_many(&self, counters: &[Counter]) -> Result<Vec<Option<Number>>, Error> {
let keys = counters.iter().map(Counter::key).collect::<Vec<_>>();
self.usage.batch_get(&keys).await
}
pub fn flush_local_counters(&self) -> Result<(), Error> {
Ok(self.usage.memory().flush()?)
}
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
/// in-memory router state.
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, Error> {
Ok(self.usage.memory().entries()?)
}
}

View file

@ -0,0 +1,3 @@
pub(super) fn pin(tenant: &str, session_id: &str) -> String {
format!("{{sensitive_route:{tenant}:{session_id}}}:model")
}

View file

@ -0,0 +1,111 @@
mod keys;
use std::{sync::Arc, time::Duration};
use litellm_python_compat::{Value, json};
use litellm_redis_client::{Ttl, commands};
use crate::{
DecodeError, Error,
codec::python_value,
dual_cache::{DualCache, MemorySettings},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
value::text,
};
/// The model each session is pinned to: memory per worker, Redis shared by every pod. Like
/// `SensitiveDataRoutingHandler` on `main`, it reaches past `DualCache` to Redis first.
pub struct SensitiveRoutePins {
pins: DualCache<String>,
}
impl SensitiveRoutePins {
pub fn new(remote: Option<Remote>) -> Self {
Self::with_clock(remote, system_clock())
}
pub fn with_clock(remote: Option<Remote>, now: Clock) -> Self {
Self {
pins: DualCache::new(
Arc::new(RemoteSlot::new(remote)),
MemorySettings::default(),
now,
),
}
}
/// Redis first, kept in memory for the key's remaining TTL. A Redis miss or failure serves memory.
pub async fn get(
&self,
tenant: &str,
session_id: &str,
default_ttl: Duration,
) -> Result<Option<String>, Error> {
let key = keys::pin(tenant, session_id);
if let Some(model) = self.read_remote(&key, default_ttl).await? {
return Ok(Some(model));
}
Ok(self.pins.memory().get(&key)?)
}
/// `SET ... EX` whose failure only logs, then memory.
pub async fn set(
&self,
tenant: &str,
session_id: &str,
model: &str,
ttl: Duration,
) -> Result<(), Error> {
let key = keys::pin(tenant, session_id);
if let Some(client) = self.pins.redis().client() {
let encoded = json::dumps(&Value::Str(model.to_owned()))?;
if let Err(error) = client.send(commands::set(&key, encoded, ttl)).await {
crate::failure_log::redis_failure!(
&error,
"writing a sensitive route pin to Redis failed"
);
}
}
self.pins.memory().set(&key, model.to_owned(), Some(ttl))?;
Ok(())
}
/// `GET`, then on a hit its own `TTL` trip, as `RedisCache.async_get_cache` and `async_get_ttl`
/// make them. A key without a TTL, or a failed `TTL`, keeps `default_ttl`.
async fn read_remote(&self, key: &str, default_ttl: Duration) -> Result<Option<String>, Error> {
let Some(client) = self.pins.redis().client() else {
return Ok(None);
};
let read = client
.send(commands::get::<Vec<u8>>(key))
.await
.map_err(Error::from)
.and_then(|raw| raw.as_deref().map(decode).transpose().map_err(Error::from));
let model = match read {
Ok(Some(model)) => model,
Ok(None) => return Ok(None),
Err(error) => {
crate::failure_log::redis_failure!(
&error,
"reading a sensitive route pin from Redis failed, serving memory"
);
return Ok(None);
}
};
let ttl = match client.send(commands::ttl(key)).await {
Ok(Ttl::Expires(remaining)) => remaining,
Ok(Ttl::Missing | Ttl::Persistent) => default_ttl,
Err(error) => {
tracing::debug!(%error, "Redis TTL Error");
default_ttl
}
};
self.pins.memory().set(key, model.clone(), Some(ttl))?;
Ok(Some(model))
}
}
fn decode(raw: &[u8]) -> Result<String, DecodeError> {
text(&python_value(raw)?, "a pinned model name")
}

View file

@ -0,0 +1,85 @@
/// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum SpendCounter {
Key {
token: String,
window: Option<String>,
},
Team {
team_id: String,
window: Option<String>,
},
TeamMember {
user_id: String,
team_id: String,
},
User(String),
EndUser(String),
Org(String),
Tag(String),
Project(String),
ModelAccessGroup(String),
ShadowEval(String),
}
#[derive(Debug, thiserror::Error)]
#[error("{0:?} is not a spend counter")]
pub struct NotASpendCounter(pub String);
impl SpendCounter {
pub fn key(&self) -> String {
let windowed = |scope: &str, id: &str, window: &Option<String>| match window {
Some(window) => format!("spend:{scope}:{id}:window:{window}"),
None => format!("spend:{scope}:{id}"),
};
match self {
Self::Key { token, window } => windowed("key", token, window),
Self::Team { team_id, window } => windowed("team", team_id, window),
Self::TeamMember { user_id, team_id } => {
format!("spend:team_member:{user_id}:{team_id}")
}
Self::User(id) => format!("spend:user:{id}"),
Self::EndUser(id) => format!("spend:end_user:{id}"),
Self::Org(id) => format!("spend:org:{id}"),
Self::Tag(name) => format!("spend:tag:{name}"),
Self::Project(id) => format!("spend:project:{id}"),
Self::ModelAccessGroup(name) => format!("spend:model_access_group:{name}"),
Self::ShadowEval(job_id) => format!("spend:shadow_eval:{job_id}"),
}
}
pub fn parse(key: &str) -> Result<Self, NotASpendCounter> {
let invalid = || NotASpendCounter(key.to_owned());
let rest = key.strip_prefix("spend:").ok_or_else(invalid)?;
let (scope, id) = rest.split_once(':').ok_or_else(invalid)?;
let windowed = || match id.rsplit_once(":window:") {
Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())),
_ => (id.to_owned(), None),
};
Ok(match scope {
"key" => {
let (token, window) = windowed();
Self::Key { token, window }
}
"team" => {
let (team_id, window) = windowed();
Self::Team { team_id, window }
}
"team_member" => {
let (user_id, team_id) = id.rsplit_once(':').ok_or_else(invalid)?;
Self::TeamMember {
user_id: user_id.to_owned(),
team_id: team_id.to_owned(),
}
}
"user" => Self::User(id.to_owned()),
"end_user" => Self::EndUser(id.to_owned()),
"org" => Self::Org(id.to_owned()),
"tag" => Self::Tag(id.to_owned()),
"project" => Self::Project(id.to_owned()),
"model_access_group" => Self::ModelAccessGroup(id.to_owned()),
"shadow_eval" => Self::ShadowEval(id.to_owned()),
_ => return Err(invalid()),
})
}
}

View file

@ -0,0 +1,380 @@
mod keys;
mod types;
use std::{sync::Arc, time::Duration};
use litellm_python_compat::repr::to_str;
use litellm_redis_client::{Declared, Expiry, RedisBatch, commands};
pub use self::keys::{NotASpendCounter, SpendCounter};
pub use self::types::{LostSeed, Reading};
use crate::{
Error, Number, Pending, RequestRedis, Send,
dual_cache::{DualCache, MemorySettings},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
};
/// `RedisCache.async_set_max`, byte for byte, sent as `EVAL` like it: a counter only ever rises to
/// the authoritative spend.
const SET_MAX: &str = concat!(
"local cur = redis.call('GET', KEYS[1]) ",
"if cur == false or tonumber(cur) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]) ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end ",
"return ARGV[1] end ",
"return cur"
);
/// The proxy's spend counters: Redis shared by every pod, and this pod's memory copy that answers
/// only when Redis cannot. Like `spend_counter_cache`'s callers on `main`, most ops reach past
/// `DualCache` to its Redis and its memory.
pub struct SpendCounters {
cache: DualCache<Number>,
}
impl SpendCounters {
pub fn new(remote: Option<Remote>) -> Self {
Self::with_clock(remote, system_clock())
}
pub fn with_clock(remote: Option<Remote>, now: Clock) -> Self {
Self {
cache: DualCache::new(
Arc::new(RemoteSlot::new(remote)),
MemorySettings::default(),
now,
),
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.cache.redis().attach(remote)
}
pub fn has_redis(&self) -> bool {
self.cache.redis().get().is_some()
}
/// The TTL every Redis write gives a counter, `None` without Redis.
pub fn redis_ttl(&self) -> Option<Duration> {
self.cache.redis().get().map(|remote| remote.default_ttl)
}
/// `read_spend_counter_cache_value`: Redis first. `RedisCache.async_get_cache` swallows an unreachable
/// `GET` as a miss, so an open breaker, a transport error, or a value `float()` refuses reads memory instead.
pub async fn read(&self, counter: &SpendCounter) -> Result<Reading, Error> {
let Some(client) = self.cache.redis().client() else {
return Ok(Reading::Memory(self.get_local(counter)?));
};
match client.send(commands::get::<Vec<u8>>(&counter.key())).await {
Ok(raw) => match types::decode(raw) {
Ok(value) => Ok(Reading::Redis(value)),
Err(error) => {
tracing::debug!(%error, "spend counter Redis value did not decode, falling back to memory");
Ok(Reading::Memory(self.get_local(counter)?))
}
},
Err(litellm_redis_client::Error::BreakerOpen) => {
Ok(Reading::Memory(self.get_local(counter)?))
}
Err(error) => {
crate::failure_log::redis_failure!(&error, "spend counter Redis read failed");
Ok(Reading::Memory(self.get_local(counter)?))
}
}
}
/// One `MGET` for many counters on the batch `send` picks; `None` without Redis. The read
/// resolves to `None` when it failed, so each caller reads on its own.
pub fn read_many(
&self,
send: Send<'_>,
counters: &[SpendCounter],
) -> Option<Pending<Option<Vec<Option<Number>>>>> {
let batch = self.cache.redis().batch(send)?;
let keys: Vec<String> = counters.iter().map(SpendCounter::key).collect();
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
let read = batch.add(commands::mget::<Vec<u8>>(&keys));
Some(Box::pin(async move {
let decoded = read
.resolve()
.await
.map_err(Error::from)
.and_then(|replies| {
replies
.into_iter()
.map(|raw| Ok(types::decode(raw)?))
.collect::<Result<Vec<_>, Error>>()
});
Ok(decoded
.inspect_err(|error| {
tracing::debug!(%error, "spend counter batch read failed, falling back to per-key reads");
})
.ok())
}))
}
/// `increment_spend_counter_cache`: `INCRBYFLOAT` refreshing the TTL, mirrored into memory. A
/// failed increment may or may not have landed, so the counter is dropped everywhere before the
/// error returns. Without Redis it is a memory increment.
pub async fn increment(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
let Some(remote) = self.cache.redis().get() else {
return self.increment_local(counter, amount);
};
let key = counter.key();
let increment = commands::incr_by_float(
&key,
amount.into(),
Some(Expiry::Refresh(remote.default_ttl)),
);
match remote.client.send(increment).await {
Ok(total) => {
let total = Number::Float(total);
self.cache.memory().set(&key, total, None)?;
Ok(total)
}
Err(error) => {
self.invalidate(std::slice::from_ref(counter)).await;
Err(error.into())
}
}
}
/// `run_spend_counter_pipeline`: every increment in one round trip, totals in order. The caller
/// decides what a failure means for counters whose increment may or may not have landed.
pub async fn increment_many(
&self,
increments: &[(SpendCounter, Number)],
) -> Result<Vec<Number>, Error> {
let Some(remote) = self.cache.redis().get() else {
return increments
.iter()
.map(|(counter, amount)| self.increment_local(counter, *amount))
.collect();
};
let pipeline = RedisBatch::pipelined(remote.client.clone());
let totals = increments.iter().map(|(counter, amount)| {
pipeline.add(commands::incr_by_float(
&counter.key(),
(*amount).into(),
Some(Expiry::Refresh(remote.default_ttl)),
))
});
let totals = Declared::all(totals.collect::<Vec<_>>()).await;
increments
.iter()
.zip(totals)
.map(|((counter, _), total)| {
let total = Number::Float(total?);
self.cache.memory().set(&counter.key(), total, None)?;
Ok(total)
})
.collect()
}
/// `increment_many` on the request's post-call batch. Each total reaches memory once the batch
/// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis.
pub fn increment_many_post_call(
self: &Arc<Self>,
request: &RequestRedis,
increments: &[(SpendCounter, Number)],
) -> bool {
let Some(remote) = self.cache.redis().get() else {
return false;
};
let batch = Send::PostCall(request).batch(&remote.client);
let declared = increments
.iter()
.map(|(counter, amount)| {
let total = batch.add(commands::incr_by_float(
&counter.key(),
(*amount).into(),
Some(Expiry::KeepExisting(remote.default_ttl)),
));
(counter.clone(), total)
})
.collect::<Vec<_>>();
let store = Arc::clone(self);
request.spawn(async move {
for (counter, total) in declared {
match total.settled().await {
Ok(total) => {
if let Err(error) = store.cache.memory().set(&counter.key(), Number::Float(total), None) {
tracing::warn!(%error, "mirroring a spend counter total into memory failed");
}
}
Err(error) => {
let key = counter.key();
crate::failure_log::redis_failure!(
&error,
&format!("spend counter {key} increment did not land in the post-call pipeline; invalidating it"),
);
store.invalidate(std::slice::from_ref(&counter)).await;
}
}
}
});
true
}
/// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either.
/// `None` without Redis, else what Redis holds afterwards.
pub async fn set_max(
&self,
counter: &SpendCounter,
spend: Number,
) -> Result<Option<Number>, Error> {
let key = counter.key();
{
let _guard = self.cache.memory().guard();
if self
.cache
.memory()
.get(&key)?
.is_none_or(|cached| cached.as_f64() < spend.as_f64())
{
self.cache.memory().set(&key, spend, None)?;
}
}
let Some(remote) = self.cache.redis().get() else {
return Ok(None);
};
let args = (
to_str(&spend.to_value()),
remote.default_ttl.as_secs().to_string(),
);
let stored: Option<f64> = remote
.client
.send(commands::eval_source(SET_MAX, &[&key], args))
.await?;
Ok(stored.map(Number::Float))
}
/// `SpendCounterReseed.coalesced(_window)`: seed a cold counter with `SET NX` so one pod wins,
/// answer with whatever the counter holds, and mirror it into memory. Without Redis, memory
/// keeps the larger of the seed and what it held.
pub async fn seed(
&self,
counter: &SpendCounter,
spend: Number,
lost: LostSeed,
) -> Result<Number, Error> {
let key = counter.key();
let Some(Remote {
client,
default_ttl,
}) = self.cache.redis().get()
else {
return Ok(self.cache.memory().update(&key, None, |cached| {
Ok(match cached {
Some(cached) if cached.as_f64() > spend.as_f64() => {
Number::Float(cached.as_f64())
}
_ => spend,
})
})?);
};
let current = if client
.send(commands::set_nx(&key, spend.to_json(), *default_ttl))
.await?
{
spend
} else {
match types::decode(client.send(commands::get::<Vec<u8>>(&key)).await?)? {
Some(winner) => Number::Float(winner.as_f64()),
None => match lost {
LostSeed::Answer => spend,
LostSeed::Increment => {
let increment = commands::incr_by_float(
&key,
spend.into(),
Some(Expiry::KeepExisting(*default_ttl)),
);
Number::Float(client.send(increment).await?)
}
},
}
};
self.cache.memory().set(&key, current, None)?;
Ok(current)
}
/// A spend reset: memory, then Redis, each for `ttl` or its own default.
pub async fn set(
&self,
counter: &SpendCounter,
spend: Number,
ttl: Option<Duration>,
) -> Result<(), Error> {
let key = counter.key();
self.cache.memory().set(&key, spend, ttl)?;
let Some(remote) = self.cache.redis().get() else {
return Ok(());
};
let ttl = ttl.unwrap_or(remote.default_ttl);
remote
.client
.send(commands::set(&key, spend.to_json(), ttl))
.await?;
Ok(())
}
/// Drop counters from memory, then from Redis in one round trip.
pub async fn delete(&self, counters: &[SpendCounter]) -> Result<(), Error> {
let keys: Vec<String> = counters.iter().map(SpendCounter::key).collect();
self.cache.delete_keys(&keys).await
}
/// `refresh_spend_counter_ttl`: `false` without Redis or for a counter Redis no longer holds.
pub async fn refresh_ttl(&self, counter: &SpendCounter) -> Result<bool, Error> {
let Some(remote) = self.cache.redis().get() else {
return Ok(false);
};
Ok(remote
.client
.send(commands::expire(&counter.key(), remote.default_ttl))
.await?)
}
pub fn get_local(&self, counter: &SpendCounter) -> Result<Option<Number>, Error> {
Ok(self.cache.memory().get(&counter.key())?)
}
pub fn set_local(
&self,
counter: &SpendCounter,
spend: Number,
ttl: Option<Duration>,
) -> Result<(), Error> {
Ok(self.cache.memory().set(&counter.key(), spend, ttl)?)
}
pub fn delete_local(&self, counter: &SpendCounter) -> Result<(), Error> {
Ok(self.cache.memory().delete(&counter.key())?)
}
/// `DualCache.async_increment_cache(local_only=True)`, which passes `InMemoryCache` no TTL.
pub fn increment_local(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
Ok(self
.cache
.memory()
.increment(&counter.key(), amount, None)?)
}
pub fn flush_local(&self) -> Result<(), Error> {
Ok(self.cache.memory().flush()?)
}
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, Error> {
Ok(self.cache.memory().entries()?)
}
/// `invalidate_spend_counter`: memory always, Redis best effort.
pub async fn invalidate(&self, counters: &[SpendCounter]) {
if let Err(error) = self.delete(counters).await {
tracing::debug!(%error, "dropping spend counters from Redis failed");
}
}
}

View file

@ -0,0 +1,48 @@
use litellm_python_compat::Value;
use super::keys::SpendCounter;
use crate::{DecodeError, FromValue, Number, codec};
/// The `spend:*` key Python names a counter by.
impl FromValue for SpendCounter {
fn from_value(value: &Value) -> Result<Self, DecodeError> {
let Value::Str(key) = value else {
return Err(DecodeError::Shape("a spend counter is a spend:* key"));
};
Self::parse(key).map_err(|_| DecodeError::Shape("a spend counter is a spend:* key"))
}
}
/// Where a read was answered, named the way the Python store reports it. Redis is cross-pod truth,
/// including a clean miss; memory only holds this pod's writes and answers only when there is no
/// Redis or it failed.
#[derive(Clone, Copy, Debug, PartialEq, strum::IntoStaticStr)]
pub enum Reading {
#[strum(serialize = "redis")]
Redis(Option<Number>),
#[strum(serialize = "memory")]
Memory(Option<Number>),
}
impl Reading {
pub fn value(self) -> Option<Number> {
match self {
Self::Redis(value) | Self::Memory(value) => value,
}
}
}
/// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LostSeed {
/// `SpendCounterReseed.coalesced`: answer with the seed without writing it.
Answer,
/// `SpendCounterReseed.coalesced_window`: `INCRBYFLOAT` the seed, keeping any TTL.
Increment,
}
pub(super) fn decode(raw: Option<Vec<u8>>) -> Result<Option<Number>, DecodeError> {
raw.as_deref()
.map(|raw| codec::number(raw, "a spend counter is a number"))
.transpose()
}

View file

@ -2,12 +2,12 @@ use litellm_python_compat::Value;
use crate::{DecodeError, Number};
/// A store type read from the plain Python data its host sends, so the host never names a field.
/// A store value read from the plain Python data it is held as.
pub trait FromValue: Sized {
fn from_value(value: &Value) -> Result<Self, DecodeError>;
}
/// A store type written as the plain Python data its host hands back.
/// A store value written as the plain Python data it is held as.
pub trait IntoValue {
fn into_value(self) -> Value;
}
@ -36,20 +36,6 @@ impl IntoValue for Number {
}
}
/// The fields of a Python tuple of exactly `N` items.
pub(crate) fn fields<'a, const N: usize>(
value: &'a Value,
shape: &'static str,
) -> Result<&'a [Value; N], DecodeError> {
match value {
Value::Tuple(items) | Value::List(items) => items
.as_slice()
.try_into()
.map_err(|_| DecodeError::Shape(shape)),
_ => Err(DecodeError::Shape(shape)),
}
}
pub(crate) fn text(value: &Value, shape: &'static str) -> Result<String, DecodeError> {
match value {
Value::Str(text) => Ok(text.clone()),
@ -57,23 +43,6 @@ pub(crate) fn text(value: &Value, shape: &'static str) -> Result<String, DecodeE
}
}
pub(crate) fn optional_text(
value: &Value,
shape: &'static str,
) -> Result<Option<String>, DecodeError> {
match value {
Value::None => Ok(None),
value => text(value, shape).map(Some),
}
}
pub(crate) fn integer(value: &Value, shape: &'static str) -> Result<i64, DecodeError> {
match value {
Value::Int(value) => i64::try_from(value).map_err(|_| DecodeError::Shape(shape)),
_ => Err(DecodeError::Shape(shape)),
}
}
pub(crate) fn number(value: &Value, shape: &'static str) -> Result<Number, DecodeError> {
Number::from_value(value).ok_or(DecodeError::Shape(shape))
}

View file

@ -0,0 +1,741 @@
mod support;
use std::time::Duration;
use litellm_python_compat::Value;
use litellm_redis_client::commands;
use litellm_redis_client::{Operation, RedisClient, Ttl};
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, Remote, RequestRedis,
auth_cache::{self, AuthCache, Settings},
router_cooldowns::RouterCooldowns,
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within};
use tokio::runtime::Handle;
const LONG_DEADLINE: Duration = Duration::from_secs(60);
const KEY_OBJECT: &str = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b";
const TEAM_JSON: &str = concat!(
r#"{"team_id": "team-\u00e9", "models": ["gpt-5", "claude-opus-5"], "max_budget": null, "#,
r#""spend": 0.30000000000000004, "blocked": false, "tpm_limit": 10000000000, "#,
r#""metadata": {"tags": [], "quoted": "it's \"x\""}}"#,
);
const TEAM_STR: &str = concat!(
"{'team_id': 'team-é', 'models': ['gpt-5', 'claude-opus-5'], 'max_budget': None, ",
"'spend': 0.30000000000000004, 'blocked': False, 'tpm_limit': 10000000000, ",
r#"'metadata': {'tags': [], 'quoted': 'it\'s "x"'}}"#,
);
fn dict<const N: usize>(entries: [(&str, Value); N]) -> Value {
Value::Dict(
entries
.into_iter()
.map(|(key, value)| (Value::from(key), value))
.collect(),
)
}
fn strings(items: &[&str]) -> Value {
Value::List(items.iter().copied().map(Value::from).collect())
}
fn team_record() -> Value {
dict([
("team_id", Value::from("team-é")),
("models", strings(&["gpt-5", "claude-opus-5"])),
("max_budget", Value::None),
("spend", Value::Float(0.30000000000000004)),
("blocked", Value::Bool(false)),
("tpm_limit", Value::from(10000000000)),
(
"metadata",
dict([
("tags", strings(&[])),
("quoted", Value::from("it's \"x\"")),
]),
),
])
}
/// The writes `tests/unit/redis/stores/test_auth_cache.py` checks against Python's own
/// `json.dumps` and `str`.
fn written() -> [(Value, &'static str, &'static str); 7] {
[
(team_record(), TEAM_JSON, TEAM_STR),
(
dict([
("token", Value::from(KEY_OBJECT)),
("spend", Value::Float(1e-05)),
("rpm_limit", Value::from(5)),
("expires", Value::from("2026-11-03T00:00:00Z")),
]),
concat!(
r#"{"token": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", "spend": 1e-05, "#,
r#""rpm_limit": 5, "expires": "2026-11-03T00:00:00Z"}"#,
),
concat!(
"{'token': '88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b', 'spend': 1e-05, ",
"'rpm_limit': 5, 'expires': '2026-11-03T00:00:00Z'}",
),
),
(
Value::from("__no_team_membership__"),
"\"__no_team_membership__\"",
"__no_team_membership__",
),
(
strings(&["mcp-a", "mcp-b"]),
r#"["mcp-a", "mcp-b"]"#,
"['mcp-a', 'mcp-b']",
),
(Value::Float(12.5), "12.5", "12.5"),
(Value::from(7), "7", "7"),
(Value::Bool(true), "true", "True"),
]
}
/// The reads `tests/unit/redis/stores/test_auth_cache.py` checks against Python's own
/// `RedisCache._get_cache_logic`.
fn read_back() -> [(&'static str, Value); 5] {
[
(TEAM_JSON, team_record()),
(TEAM_STR, team_record()),
("['mcp-a', 'mcp-b']", strings(&["mcp-a", "mcp-b"])),
(
"{'user_id': 'u1', 'teams': ['t1'], 'active': True, 'budget': None}",
dict([
("user_id", Value::from("u1")),
("teams", strings(&["t1"])),
("active", Value::Bool(true)),
("budget", Value::None),
]),
),
("\"__NO_MAPPING__\"", Value::from("__NO_MAPPING__")),
]
}
fn settings(memory_ttl: Option<u64>) -> Settings {
Settings {
memory_ttl: memory_ttl.map(Duration::from_secs),
..Settings::default()
}
}
fn store(client: Option<RedisClient>) -> AuthCache {
AuthCache::new(client.map(remote), settings(Some(60)))
}
fn team(id: &str) -> Value {
Value::Dict(vec![(Value::from("team_id"), Value::from(id))])
}
async fn redis_ttl(raw: &RedisClient, key: &str) -> Duration {
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
remaining
}
#[rstest]
#[case(KEY_OBJECT, true)]
#[case(
"88DC28D0F030C55ED4AB77ED8FAF098196CB1C05DF778539800C9F1243FE6B4B",
false
)]
#[case(
"88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4",
false
)]
#[case(
"team_id:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
false
)]
#[case("team_membership:user-1:team-1", false)]
#[case("end_user_restricted_registry", false)]
fn only_a_bare_lowercase_sha256_digest_is_a_user_key_object(
#[case] key: &str,
#[case] key_object: bool,
) {
assert_eq!(auth_cache::is_key_object(key), key_object, "{key}");
}
#[rstest]
#[tokio::test]
async fn writes_mains_json_and_str_bytes(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
for (index, (value, json, printed)) in written().into_iter().enumerate() {
let key = format!("team_id:t{index}");
store.set(&key, value.clone(), None).await.unwrap();
assert_eq!(
raw.send(commands::get::<String>(&key))
.await
.unwrap()
.as_deref(),
Some(json)
);
store.set_str(&key, value, None).await.unwrap();
assert_eq!(
raw.send(commands::get::<String>(&key))
.await
.unwrap()
.as_deref(),
Some(printed)
);
}
}
#[rstest]
#[tokio::test]
async fn reads_every_encoding_main_writes_and_backfills_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
for (index, (bytes, expected)) in read_back().into_iter().enumerate() {
let key = format!("user_id:u{index}");
raw.send(commands::set(&key, bytes, Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(store.get(None, &key).await.unwrap(), Some(expected.clone()));
raw.send(commands::del(&[&key])).await.unwrap();
assert_eq!(
store.get(None, &key).await.unwrap(),
Some(expected),
"{key}"
);
}
assert_eq!(store.get(None, "user_id:absent").await.unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn ttls_follow_dual_cache_and_the_write_back_rule(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = AuthCache::new(
Some(Remote {
client: connect(&deployment),
default_ttl: Duration::from_secs(45),
}),
Settings {
memory_ttl: Some(Duration::from_secs(20)),
..Settings::default()
},
);
store
.set(
"team_id:given",
team("a"),
Some(Duration::from_secs_f64(30.9)),
)
.await
.unwrap();
store.set("team_id:default", team("b"), None).await.unwrap();
store
.write_back(
RedisSend::Now,
vec![
("user-1".to_owned(), team("c"), None),
("user-2".to_owned(), team("d"), Some(Duration::from_secs(5))),
],
)
.unwrap()
.await
.unwrap();
assert!(within(redis_ttl(&raw, "team_id:given").await, 30));
assert!(within(redis_ttl(&raw, "team_id:default").await, 20));
assert!(within(redis_ttl(&raw, "user-1").await, 45));
assert!(within(redis_ttl(&raw, "user-2").await, 5));
assert_eq!(store.get_local("user-1").unwrap(), Some(team("c")));
}
#[tokio::test]
async fn user_key_objects_survive_churn_in_the_management_object_partition() {
let store = AuthCache::new(
None,
Settings {
memory_entries: 2,
..Settings::default()
},
);
store.set(KEY_OBJECT, team("key"), None).await.unwrap();
for index in 0..5 {
store
.set(&format!("team_id:t{index}"), team("t"), None)
.await
.unwrap();
}
assert_eq!(
store.get(None, KEY_OBJECT).await.unwrap(),
Some(team("key"))
);
assert_eq!(store.len_local().unwrap(), (2, 1));
}
#[rstest]
#[tokio::test]
async fn a_value_json_refuses_stays_in_memory_and_never_reaches_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let members = Value::Set(vec![Value::from("a")]);
store
.set("team_id:set", members.clone(), None)
.await
.unwrap();
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
assert_eq!(
raw.send(commands::get::<String>("team_id:set"))
.await
.unwrap(),
None
);
}
// RedisCache._pipeline_helper runs json.dumps inside the pipeline's try, so one value it refuses
// aborts every SET of that pipeline, while DualCache has already written them all to memory
#[tokio::test]
async fn a_value_json_refuses_keeps_its_whole_pipeline_out_of_redis() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let members = Value::Set(vec![Value::from("a")]);
store
.set_many(
vec![
("team_id:a".to_owned(), team("a")),
("team_id:set".to_owned(), members.clone()),
],
None,
)
.await
.unwrap();
assert_eq!(store.get_local("team_id:a").unwrap(), Some(team("a")));
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
assert_eq!(
raw.send(commands::get::<String>("team_id:a"))
.await
.unwrap(),
None
);
}
// auth_object_prefetch writes every entry to memory before
// RedisCache.async_set_cache_pipeline_with_ttls encodes them, and that encodes all before sending
#[tokio::test]
async fn a_write_back_value_json_refuses_fails_after_memory_and_before_redis() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let members = Value::Set(vec![Value::from("a")]);
let written = store.write_back(
RedisSend::Now,
vec![
("team_id:set".to_owned(), members.clone(), None),
("team_id:a".to_owned(), team("a"), None),
],
);
assert!(written.is_err());
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
assert_eq!(store.get_local("team_id:a").unwrap(), Some(team("a")));
assert_eq!(
raw.send(commands::get::<String>("team_id:a"))
.await
.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn pre_call_writes_and_reads_ride_one_pipeline_with_the_cooldown_read(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let cooldowns = std::sync::Arc::new(RouterCooldowns::new(
Some(remote(client.clone())),
COOLDOWN_POLICY,
));
let store = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
raw.send(commands::set(
KEY_OBJECT,
"{\"token\": \"t\"}",
Duration::from_secs(60),
))
.await
.unwrap();
raw.send(commands::set(
"team_alias:alpha",
"\"stale\"",
Duration::from_secs(60),
))
.await
.unwrap();
let reads = store
.get_many(
RedisSend::PreCall(&request),
&[KEY_OBJECT.to_owned(), "end_user_id:absent".to_owned()],
)
.unwrap();
let active = cooldowns
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
.unwrap();
let written = store
.declare_set(RedisSend::PreCall(&request), "team_id:t1", team("t1"), None)
.unwrap();
let deleted = store
.declare_delete(RedisSend::PreCall(&request), "team_alias:alpha")
.unwrap();
assert!(trips.seen().is_empty());
assert_eq!(store.get_local("team_id:t1").unwrap(), Some(team("t1")));
request.flush_pre().await;
let mut seen = trips.seen();
seen.sort_by_key(|operation| format!("{operation:?}"));
let expected = match topology {
Topology::Standalone => vec![Operation::Pipeline { ops: 4 }],
Topology::Cluster => vec![
Operation::Command("DEL"),
Operation::Command("MGET"),
Operation::Command("MGET"),
Operation::Command("SET"),
],
};
assert_eq!(seen, expected);
assert_eq!(
reads.await.unwrap(),
vec![
Some(Value::Dict(vec![(Value::from("token"), Value::from("t"))])),
None
]
);
assert_eq!(active.await.unwrap(), vec![]);
written.await.unwrap();
deleted.await.unwrap();
assert_eq!(
raw.send(commands::get::<String>("team_id:t1"))
.await
.unwrap(),
Some("{\"team_id\": \"t1\"}".to_owned())
);
assert!(within(redis_ttl(&raw, "team_id:t1").await, 60));
assert_eq!(
raw.send(commands::get::<String>("team_alias:alpha"))
.await
.unwrap(),
None
);
assert_eq!(trips.seen().len(), expected.len());
}
#[rstest]
#[tokio::test]
async fn a_write_back_on_the_request_waits_for_the_next_flush(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let store = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
store
.write_back(
RedisSend::PreCall(&request),
vec![(
"user-1".to_owned(),
team("u"),
Some(Duration::from_secs(42)),
)],
)
.unwrap()
.await
.unwrap();
assert!(trips.seen().is_empty());
assert_eq!(store.get_local("user-1").unwrap(), Some(team("u")));
request.flush_pre().await;
let write = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
Topology::Cluster => Operation::Command("SET"),
};
assert_eq!(trips.seen(), [write]);
assert!(within(redis_ttl(&raw, "user-1").await, 42));
}
#[rstest]
#[tokio::test]
async fn an_eviction_broadcast_drops_only_this_pods_copy_and_the_next_read_refills_from_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let this_pod = store(Some(connect(&deployment)));
let other_pod = store(Some(connect(&deployment)));
this_pod.set("team_id:t1", team("old"), None).await.unwrap();
other_pod
.set("team_id:t1", team("new"), None)
.await
.unwrap();
assert_eq!(
this_pod.get(None, "team_id:t1").await.unwrap(),
Some(team("old"))
);
this_pod.delete_local("team_id:t1").unwrap();
assert_eq!(
this_pod.get(None, "team_id:t1").await.unwrap(),
Some(team("new"))
);
assert!(
raw.send(commands::get::<String>("team_id:t1"))
.await
.unwrap()
.is_some()
);
}
// RedisBatch.read_as_missing: once the request's MGET found a key absent, or the request deleted it,
// a GET of it is answered without a round trip until the request sets it again
#[rstest]
#[tokio::test]
async fn a_key_the_request_read_as_missing_is_not_read_again(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let other_pod_writes = |key: &'static str| {
raw.send(commands::set(
key,
TEAM_JSON.to_owned(),
Duration::from_secs(60),
))
};
let read = store
.get_many(RedisSend::PreCall(&request), &["team_id:read".to_owned()])
.unwrap();
assert_eq!(read.await.unwrap(), vec![None]);
store
.declare_delete(RedisSend::PreCall(&request), "team_id:deleted")
.unwrap()
.await
.unwrap();
other_pod_writes("team_id:read").await.unwrap();
other_pod_writes("team_id:deleted").await.unwrap();
assert_eq!(
store.get(Some(&request), "team_id:read").await.unwrap(),
None
);
assert_eq!(
store.get(Some(&request), "team_id:deleted").await.unwrap(),
None
);
assert_eq!(
store.get(None, "team_id:read").await.unwrap(),
Some(team_record())
);
store
.declare_set(
RedisSend::PreCall(&request),
"team_id:deleted",
team("t"),
None,
)
.unwrap()
.await
.unwrap();
store.delete_local("team_id:deleted").unwrap();
assert_eq!(
store.get(Some(&request), "team_id:deleted").await.unwrap(),
Some(team("t"))
);
}
#[rstest]
#[tokio::test]
async fn delete_drops_memory_and_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let keys = [KEY_OBJECT.to_owned(), "team_id:t1".to_owned()];
store
.set_many(
keys.iter().map(|key| (key.clone(), team("t"))).collect(),
None,
)
.await
.unwrap();
assert!(
raw.send(commands::get::<String>("team_id:t1"))
.await
.unwrap()
.is_some()
);
store.delete(&keys).await.unwrap();
for key in &keys {
assert_eq!(store.get_local(key).unwrap(), None);
assert_eq!(raw.send(commands::get::<String>(key)).await.unwrap(), None);
}
}
#[rstest]
#[tokio::test]
async fn increment_if_present_only_raises_an_entry_redis_holds(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
assert_eq!(
store
.increment_if_present("litellm-proxy-budget:spend", Number::Float(2.5))
.await
.unwrap(),
None
);
assert_eq!(
raw.send(commands::get::<String>("litellm-proxy-budget:spend"))
.await
.unwrap(),
None
);
store
.set("litellm-proxy-budget:spend", Value::Float(1.5), None)
.await
.unwrap();
assert_eq!(
store
.increment_if_present("litellm-proxy-budget:spend", Number::Float(2.5))
.await
.unwrap(),
Some(4.0)
);
}
#[tokio::test]
async fn without_redis_every_operation_is_memory_only() {
let store = store(None);
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
store.set("team_id:t1", team("t"), None).await.unwrap();
store
.declare_set(RedisSend::PreCall(&request), "team_id:t2", team("t"), None)
.unwrap()
.await
.unwrap();
assert_eq!(store.get_local("team_id:t2").unwrap(), Some(team("t")));
store
.declare_delete(RedisSend::PreCall(&request), "team_id:t2")
.unwrap()
.await
.unwrap();
assert_eq!(store.get_local("team_id:t2").unwrap(), None);
assert!(
store
.get_many(RedisSend::PreCall(&request), &["team_id:t1".to_owned()])
.is_none()
);
assert!(
store
.get_many(RedisSend::Now, &["team_id:t1".to_owned()])
.is_none()
);
assert_eq!(
store
.increment_if_present("k", Number::Float(1.0))
.await
.unwrap(),
None
);
assert_eq!(
store.get(None, "team_id:t1").await.unwrap(),
Some(team("t"))
);
assert_eq!(store.get_local("team_id:t2").unwrap(), None);
}
// UserApiKeyCache writes the key-object partition, then the rest, and deletes both partitions
// concurrently, while async_increment_if_present sends its script as a plain EVAL
#[rstest]
#[tokio::test]
async fn the_auth_cache_sends_the_commands_user_api_key_cache_sends(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let (client, trips) = connect_observed(&deployment);
let store = store(Some(client));
deployment
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
.await
.unwrap();
store
.set_many(
vec![
(KEY_OBJECT.to_owned(), team("t0")),
("team_id:t1".to_owned(), team("t1")),
],
None,
)
.await
.unwrap();
let written = trips.seen();
store
.delete(&[KEY_OBJECT.to_owned(), "team_id:t1".to_owned()])
.await
.unwrap();
store
.increment_if_present("litellm-proxy-budget:spend", Number::Int(5))
.await
.unwrap();
assert_eq!(written, vec![Operation::Pipeline { ops: 1 }; 2]);
assert_eq!(
trips.seen()[written.len()..written.len() + 2],
[Operation::Command("DEL"), Operation::Command("DEL")]
);
assert_eq!(support::calls(&deployment, "eval").await, 1);
assert_eq!(support::calls(&deployment, "evalsha").await, 0);
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,512 @@
mod support;
use std::time::Duration;
use litellm_redis_client::commands;
use litellm_redis_client::{Operation, Ttl};
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, RequestRedis,
router_cooldowns::{Cooldown, RouterCooldowns},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{ALLOWED_FAILS, COOLDOWN_POLICY as POLICY, connect, cooldowns, remote, within};
use tokio::runtime::Handle;
async fn read_active(
store: &std::sync::Arc<RouterCooldowns>,
model_ids: &[String],
) -> Vec<(String, Cooldown)> {
store
.active(RedisSend::Now, model_ids)
.unwrap()
.await
.unwrap()
}
/// A clock fixed just after the fixture timestamps, so every fixture cooldown is still active.
fn fixture_clock() -> litellm_redis_store::Clock {
std::sync::Arc::new(|| Duration::from_secs_f64(1_730_000_000.75))
}
#[rstest]
#[tokio::test]
async fn put_writes_the_bytes_mains_set_path_writes_with_whole_second_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
#[values("json", "str_dict")] format: &str,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = std::sync::Arc::new(RouterCooldowns::new(
Some(remote(connect(&deployment))),
POLICY,
));
for case in cooldowns() {
let ttl = Duration::from_secs_f64(case.ttl_seconds);
let (written_by_put, expected) = match format {
"json" => (
store.put(case.model_id, &case.cooldown, ttl).await,
case.json,
),
_ => (
store.put_str(case.model_id, &case.cooldown, ttl).await,
case.str_dict,
),
};
written_by_put.unwrap();
let key = case.key;
let written: Option<String> = raw.send(commands::get(key)).await.unwrap();
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(written.as_deref(), Some(expected));
assert!(within(remaining, case.redis_ttl_seconds), "{remaining:?}");
}
}
#[rstest]
#[tokio::test]
async fn reads_cooldowns_python_wrote_as_str_dict_or_json(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
#[values("str_dict", "json")] format: &str,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(connect(&deployment))),
POLICY,
fixture_clock(),
));
for case in cooldowns() {
let bytes = match format {
"json" => case.json,
_ => case.str_dict,
};
raw.send(commands::set(case.key, bytes, Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(
read_active(&store, &[case.model_id.to_owned()]).await,
vec![(case.model_id.to_owned(), case.cooldown)]
);
}
}
#[rstest]
#[tokio::test]
async fn an_ended_cooldown_is_not_active_but_still_counts_toward_the_minimum(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let late = std::sync::Arc::new(|| Duration::from_secs(1_730_000_100));
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(connect(&deployment))),
POLICY,
late,
));
let models = cooldowns()
.iter()
.map(|case| case.model_id.to_owned())
.collect::<Vec<_>>();
for case in cooldowns() {
raw.send(commands::set(
case.key,
case.str_dict,
Duration::from_secs(60),
))
.await
.unwrap();
}
assert_eq!(
store.minimum(&models).await.unwrap(),
Some(Number::Float(0.5))
);
assert_eq!(read_active(&store, &models).await, vec![]);
}
#[rstest]
#[tokio::test]
async fn allowed_fails_count_across_async_and_sync_routers_and_keep_the_first_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let first = std::sync::Arc::new(RouterCooldowns::new(
Some(remote(connect(&deployment))),
POLICY,
));
let second = std::sync::Arc::new(RouterCooldowns::new(
Some(remote(connect(&deployment))),
POLICY,
));
for case in ALLOWED_FAILS {
let key = case.key;
raw.send(commands::del(&[key])).await.unwrap();
let suffix = case.suffix;
let counts = [
first
.increment_allowed_fails("dep-a", suffix, Duration::from_secs_f64(30.9))
.await
.unwrap(),
second
.increment_allowed_fails_int("dep-a", suffix, Duration::from_secs(120))
.await
.unwrap(),
];
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(counts, [1, 2], "{key}");
assert!(within(remaining, 30), "{remaining:?}");
}
}
#[rstest]
#[case::async_increment_cache_sends_incrbyfloat("INCRBYFLOAT")]
#[case::increment_cache_sends_incrby("INCRBY")]
#[tokio::test]
async fn allowed_fails_send_the_increment_mains_path_sends(
#[case] command: &'static str,
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let (client, trips) = support::connect_observed(&deployment);
let store = std::sync::Arc::new(RouterCooldowns::new(Some(remote(client)), POLICY));
let ttl = Duration::from_secs(60);
let count = match command {
"INCRBYFLOAT" => store.increment_allowed_fails("dep-a", None, ttl).await,
_ => store.increment_allowed_fails_int("dep-a", None, ttl).await,
};
assert_eq!(count.unwrap(), 1);
assert_eq!(trips.seen().first(), Some(&Operation::Command(command)));
}
#[rstest]
#[tokio::test]
async fn python_written_fail_counts_continue_in_rust(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = std::sync::Arc::new(RouterCooldowns::new(
Some(remote(connect(&deployment))),
POLICY,
));
raw.send(commands::incr_by(
"deployment:dep-a:allowed_fails",
4,
Duration::from_secs(60),
))
.await
.unwrap();
assert_eq!(
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
5
);
}
#[rstest]
#[tokio::test]
async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() {
let store = std::sync::Arc::new(RouterCooldowns::with_clock(None, POLICY, fixture_clock()));
let [case, _] = cooldowns();
store
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
.await
.unwrap();
let fails = [
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
];
store.flush_local_counters().unwrap();
assert_eq!(
read_active(&store, &[case.model_id.to_owned()]).await,
vec![(case.model_id.to_owned(), case.cooldown)]
);
assert_eq!(fails, [1, 2]);
assert_eq!(
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
1
);
}
#[rstest]
#[tokio::test]
async fn a_redis_outage_serves_memory_and_counts_fails_locally() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(connect(&deployment))),
POLICY,
fixture_clock(),
));
let [case, _] = cooldowns();
store
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
.await
.unwrap();
deployment.stop().await.unwrap();
let active = read_active(
&store,
&[case.model_id.to_owned(), "dep-missing".to_owned()],
)
.await;
let fails = store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap();
assert_eq!(active, vec![(case.model_id.to_owned(), case.cooldown)]);
assert_eq!(fails, 1);
}
struct StepClock(std::sync::Mutex<Duration>);
impl StepClock {
fn at(seconds: f64) -> std::sync::Arc<Self> {
std::sync::Arc::new(Self(std::sync::Mutex::new(Duration::from_secs_f64(
seconds,
))))
}
fn advance(&self, by: Duration) {
*self.0.lock().unwrap() += by;
}
fn clock(self: &std::sync::Arc<Self>) -> litellm_redis_store::Clock {
let this = self.clone();
std::sync::Arc::new(move || *this.0.lock().unwrap())
}
}
#[rstest]
#[tokio::test]
async fn a_miss_is_read_from_redis_again_only_after_the_one_second_refresh(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let time = StepClock::at(1_730_000_000.75);
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(connect(&deployment))),
POLICY,
time.clock(),
));
let [case, _] = cooldowns();
let before = read_active(&store, &[case.model_id.to_owned()]).await;
raw.send(commands::set(
case.key,
case.str_dict,
Duration::from_secs(60),
))
.await
.unwrap();
let throttled = read_active(&store, &[case.model_id.to_owned()]).await;
time.advance(Duration::from_secs(1));
let refreshed = read_active(&store, &[case.model_id.to_owned()]).await;
assert_eq!(before, vec![]);
assert_eq!(throttled, vec![]);
assert_eq!(
refreshed,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}
#[rstest]
#[tokio::test]
async fn a_breaker_refusal_does_not_hold_back_the_next_read() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let breaker = litellm_circuit_breaker::Breaker::new(
litellm_circuit_breaker::Config {
failure_threshold: 1,
recovery_timeout: Duration::from_millis(200),
..litellm_circuit_breaker::Config::default()
},
Box::new(|_| {}),
);
let time = StepClock::at(1_730_000_000.75);
let guarded = support::connect_guarded(&deployment, breaker);
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(guarded.clone())),
POLICY,
time.clock(),
));
let [case, _] = cooldowns();
deployment.stop().await.unwrap();
read_active(&store, &[case.model_id.to_owned()]).await;
time.advance(Duration::from_secs(1));
let refused = read_active(&store, &[case.model_id.to_owned()]).await;
deployment.resume().await.unwrap();
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
raw.send(commands::set(
case.key,
case.str_dict,
Duration::from_secs(60),
))
.await
.unwrap();
let recovered = read_active(&store, &[case.model_id.to_owned()]).await;
assert_eq!(refused, vec![]);
assert_eq!(
recovered,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}
#[rstest]
#[tokio::test]
async fn a_declared_active_read_rides_the_request_batch_and_answers_what_active_does(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = support::connect_observed(&deployment);
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(client)),
POLICY,
fixture_clock(),
));
let models = cooldowns().map(|case| case.model_id.to_owned());
for case in cooldowns() {
raw.send(commands::set(
case.key,
case.str_dict,
Duration::from_secs(60),
))
.await
.unwrap();
}
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let declared = store.active(RedisSend::PreCall(&request), &models).unwrap();
assert!(trips.seen().is_empty());
request.flush_pre().await;
let active = declared.await.unwrap();
let one_mget = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
Topology::Cluster => Operation::Command("MGET"),
};
assert_eq!(trips.seen(), [one_mget]);
assert_eq!(
active,
cooldowns()
.map(|case| (case.model_id.to_owned(), case.cooldown))
.to_vec()
);
assert_eq!(read_active(&store, &models).await, active);
assert_eq!(trips.seen().len(), 1, "the declared read backfilled memory");
}
#[rstest]
#[tokio::test]
async fn a_cooldown_recorded_after_the_declaration_wins_over_the_redis_miss(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(connect(&deployment))),
POLICY,
fixture_clock(),
));
let [case, _] = cooldowns();
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let declared = store
.active(RedisSend::PreCall(&request), &[case.model_id.to_owned()])
.unwrap();
request.flush_pre().await;
store
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
.await
.unwrap();
assert_eq!(
declared.await.unwrap(),
vec![(case.model_id.to_owned(), case.cooldown)]
);
}
#[rstest]
#[tokio::test]
async fn a_breaker_refusal_of_a_declared_read_does_not_hold_back_the_next_read() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let breaker = litellm_circuit_breaker::Breaker::new(
litellm_circuit_breaker::Config {
failure_threshold: 1,
recovery_timeout: Duration::from_millis(200),
..litellm_circuit_breaker::Config::default()
},
Box::new(|_| {}),
);
let guarded = support::connect_guarded(&deployment, breaker);
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(guarded.clone())),
POLICY,
fixture_clock(),
));
let [case, _] = cooldowns();
let models = [case.model_id.to_owned()];
deployment.stop().await.unwrap();
let opened = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let failed = store
.active(RedisSend::PreCall(&opened), &["dep-other".to_owned()])
.unwrap();
opened.flush_pre().await;
failed.await.unwrap();
let refused_request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let refused = store
.active(RedisSend::PreCall(&refused_request), &models)
.unwrap();
refused_request.flush_pre().await;
let refused = refused.await.unwrap();
deployment.resume().await.unwrap();
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
raw.send(commands::set(
case.key,
case.str_dict,
Duration::from_secs(60),
))
.await
.unwrap();
assert_eq!(refused, vec![]);
assert_eq!(
read_active(&store, &models).await,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}

View file

@ -0,0 +1,329 @@
mod support;
use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use litellm_redis_client::commands;
use litellm_redis_client::{Amount, Expiry, Operation, Ttl};
use litellm_redis_store::{
Number, TierPolicy,
router_usage::{COUNTER_TTL, Counter, Kind, RouterUsage},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{USAGE_POLICY as POLICY, connect, remote, usage, within};
fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] {
[Kind::Tpm, Kind::Rpm].map(|kind| Counter {
deployment_id: "dep".to_owned(),
model: model.map(str::to_owned),
minute: "12-34".to_owned(),
kind,
})
}
#[rstest]
#[case::tpm(Kind::Tpm, "tpm")]
#[case::rpm(Kind::Rpm, "rpm")]
fn a_counter_kind_round_trips_through_its_key_segment(#[case] kind: Kind, #[case] segment: &str) {
assert_eq!(<&str>::from(kind), segment);
assert_eq!(segment.parse::<Kind>().unwrap(), kind);
}
#[rstest]
#[case::unknown("tokens")]
#[case::wrong_case("TPM")]
fn other_text_is_not_a_counter_kind(#[case] segment: &str) {
assert!(segment.parse::<Kind>().is_err());
}
// DualCache._backfill_kwargs: a router's batch read keeps a Redis value in memory for
// litellm.default_in_memory_ttl when one is set, else for InMemoryCache's own 600 s
#[rstest]
#[case::configured(Some(Duration::from_secs(5)), None)]
#[case::unset(None, Some(Number::Int(3)))]
#[tokio::test]
async fn a_batch_read_backfills_memory_for_default_in_memory_ttl(
#[case] default_in_memory_ttl: Option<Duration>,
#[case] after_six_seconds: Option<Number>,
) {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let seconds = Arc::new(AtomicU64::new(1_000));
let now = {
let seconds = seconds.clone();
Arc::new(move || Duration::from_secs(seconds.load(Ordering::SeqCst)))
};
let policy = TierPolicy {
refresh_after: Duration::from_secs(3600),
default_in_memory_ttl,
..POLICY
};
let store = RouterUsage::with_clock(Some(remote(connect(&deployment))), policy, now);
let [case, _] = usage();
raw.send(commands::incr_by_float(
case.key,
Amount::Float(3.0),
Some(Expiry::Refresh(COUNTER_TTL)),
))
.await
.unwrap();
store
.get_many(std::slice::from_ref(&case.counter))
.await
.unwrap();
raw.send(commands::del(&[case.key])).await.unwrap();
seconds.fetch_add(6, Ordering::SeqCst);
let read = store.get_many(&[case.counter]).await.unwrap();
assert_eq!(read, vec![after_six_seconds]);
}
#[rstest]
#[tokio::test]
async fn reads_counters_python_incremented_under_mains_keys(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let cases = usage();
for (case, delta) in cases.iter().zip([7.75, 3.0]) {
raw.send(commands::incr_by_float(
case.key,
Amount::Float(delta),
Some(Expiry::Refresh(COUNTER_TTL)),
))
.await
.unwrap();
}
assert_eq!(
store
.get_many(&cases.map(|case| case.counter))
.await
.unwrap(),
vec![Some(Number::Float(7.75)), Some(Number::Int(3))]
);
}
#[rstest]
#[tokio::test]
async fn deployment_increments_write_incrbyfloat_totals_with_a_refreshed_minute_window(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let writer = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let reader = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let [case, _] = usage();
let model = case.counter.model.as_deref();
let [tpm_key, rpm_key] =
["tpm", "rpm"].map(|kind| case.key.replace(":tpm:", &format!(":{kind}:")));
writer
.increment_deployment("dep", model, "12-34", Number::Float(5.5), 1)
.await
.unwrap();
raw.send(commands::expire(&tpm_key, Duration::from_secs(5)))
.await
.unwrap();
writer
.increment_deployment("dep", model, "12-34", Number::Int(2), 1)
.await
.unwrap();
let written: [Option<String>; 2] = [
raw.send(commands::get(&tpm_key)).await.unwrap(),
raw.send(commands::get(&rpm_key)).await.unwrap(),
];
let Ttl::Expires(remaining) = raw.send(commands::ttl(&tpm_key)).await.unwrap() else {
panic!("{tpm_key} has no TTL");
};
assert_eq!(written, [Some("7.5".to_owned()), Some("2".to_owned())]);
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
reader.get_many(&tpm_and_rpm(model)).await.unwrap(),
vec![Some(Number::Float(7.5)), Some(Number::Int(2))]
);
assert_eq!(
writer.get_many(&tpm_and_rpm(model)).await.unwrap(),
vec![Some(Number::Int(7)), Some(Number::Int(2))]
);
}
#[rstest]
#[tokio::test]
async fn a_request_increment_keeps_the_window_python_started(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let [_, case] = usage();
let key = case.key;
raw.send(commands::incr_by_float(
key,
Amount::Float(1.0),
Some(Expiry::KeepExisting(Duration::from_secs(20))),
))
.await
.unwrap();
let total = store
.increment(&case.counter, Number::Int(1), COUNTER_TTL)
.await
.unwrap();
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(total, Number::Float(2.0));
assert!(within(remaining, 20), "{remaining:?}");
}
#[rstest]
#[tokio::test]
async fn without_redis_counters_accumulate_in_memory_until_flushed() {
let store = RouterUsage::new(None, POLICY);
let counters = tpm_and_rpm(Some("model"));
store
.increment_deployment("dep", Some("model"), "12-34", Number::Int(2), 1)
.await
.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.5), 1)
.await
.unwrap();
let counted = store.get_many(&counters).await.unwrap();
store.flush_local_counters().unwrap();
assert_eq!(
counted,
vec![Some(Number::Float(7.5)), Some(Number::Int(2))]
);
assert_eq!(store.get_many(&counters).await.unwrap(), vec![None, None]);
}
#[rstest]
#[tokio::test]
async fn a_late_attach_keeps_memory_and_later_attaches_are_ignored(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let attached = RedisDeployment::start(topology).await.unwrap();
let ignored = RedisDeployment::start(topology).await.unwrap();
let store = RouterUsage::new(None, POLICY);
let [case, _] = usage();
store
.increment(&case.counter, Number::Int(2), COUNTER_TTL)
.await
.unwrap();
assert!(store.attach(remote(connect(&attached))));
assert!(!store.attach(remote(connect(&ignored))));
let total = store
.increment(&case.counter, Number::Int(3), COUNTER_TTL)
.await
.unwrap();
assert_eq!(total, Number::Float(3.0));
assert_eq!(
store
.get_many(std::slice::from_ref(&case.counter))
.await
.unwrap(),
vec![Some(Number::Int(5))]
);
let ignored_total: Option<String> = connect(&ignored)
.send(commands::get(case.key))
.await
.unwrap();
assert_eq!(ignored_total, None);
}
#[rstest]
#[tokio::test]
async fn a_failed_pipeline_keeps_the_memory_counts_as_integers() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
deployment.stop().await.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.75), 1)
.await
.unwrap();
assert_eq!(
store.get_many(&tpm_and_rpm(Some("model"))).await.unwrap(),
vec![Some(Number::Int(5)), Some(Number::Int(1))]
);
}
// main's _increment_deployment sends async_increment_pipeline on its own round trip, never on the
// request's post-call pipeline
#[rstest]
#[tokio::test]
async fn a_deployment_increment_is_its_own_pipeline_and_mirrors_the_redis_totals(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = support::connect_observed(&deployment);
let store = RouterUsage::new(Some(remote(client)), POLICY);
let counters = tpm_and_rpm(Some("model"));
let [tpm_key, rpm_key] =
["tpm", "rpm"].map(|kind| format!("global_router:dep:model:{kind}:12-34"));
raw.send(commands::incr_by_float(
&tpm_key,
Amount::Float(10.0),
Some(Expiry::Refresh(COUNTER_TTL)),
))
.await
.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.5), 1)
.await
.unwrap();
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
assert_eq!(
store.get_many(&counters).await.unwrap(),
vec![Some(Number::Int(15)), Some(Number::Int(1))]
);
let written: [Option<String>; 2] = [
raw.send(commands::get(&tpm_key)).await.unwrap(),
raw.send(commands::get(&rpm_key)).await.unwrap(),
];
assert_eq!(written, [Some("15.5".to_owned()), Some("1".to_owned())]);
let Ttl::Expires(remaining) = raw.send(commands::ttl(&rpm_key)).await.unwrap() else {
panic!("{rpm_key} has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
}
#[rstest]
#[tokio::test]
async fn a_failed_deployment_pipeline_keeps_the_memory_counts_as_integers() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let counters = tpm_and_rpm(Some("model"));
deployment.stop().await.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.75), 1)
.await
.unwrap();
assert_eq!(
store.get_many(&counters).await.unwrap(),
[Some(Number::Int(5)), Some(Number::Int(1))]
);
}

View file

@ -0,0 +1,134 @@
mod support;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_redis_client::{Ttl, commands};
use litellm_redis_store::{Clock, sensitive_route_pins::SensitiveRoutePins};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, remote, within};
const PIN_KEY: &str = "{sensitive_route:tenant-a:session-1}:model";
const HOUR: Duration = Duration::from_secs(3600);
struct StepClock(Mutex<Duration>);
impl StepClock {
fn start() -> Arc<Self> {
Arc::new(Self(Mutex::new(Duration::from_secs(1_000))))
}
fn advance(&self, by: Duration) {
*self.0.lock().unwrap() += by;
}
fn clock(self: &Arc<Self>) -> Clock {
let this = self.clone();
Arc::new(move || *this.0.lock().unwrap())
}
}
#[rstest]
#[tokio::test]
async fn a_pin_one_pod_writes_is_read_by_another_under_mains_key_and_encoding(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let writer = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
let reader = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
writer
.set("tenant-a", "session-1", "on-prem-model", HOUR)
.await
.unwrap();
let written: Option<String> = raw.send(commands::get(PIN_KEY)).await.unwrap();
let Ttl::Expires(remaining) = raw.send(commands::ttl(PIN_KEY)).await.unwrap() else {
panic!("{PIN_KEY} has no TTL");
};
assert_eq!(written.as_deref(), Some("\"on-prem-model\""));
assert!(within(remaining, 3600), "{remaining:?}");
assert_eq!(
reader.get("tenant-a", "session-1", HOUR).await.unwrap(),
Some("on-prem-model".to_owned())
);
assert_eq!(
reader.get("tenant-b", "session-1", HOUR).await.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn a_pin_read_from_redis_stays_in_memory_for_its_remaining_ttl_only() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let time = StepClock::start();
let reader = SensitiveRoutePins::with_clock(Some(remote(connect(&deployment))), time.clock());
raw.send(commands::set(
PIN_KEY,
"\"on-prem-model\"",
Duration::from_secs(30),
))
.await
.unwrap();
let read = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
deployment.stop().await.unwrap();
time.advance(Duration::from_secs(20));
let before_expiry = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
time.advance(Duration::from_secs(15));
let after_expiry = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
assert_eq!(read, Some("on-prem-model".to_owned()));
assert_eq!(before_expiry, read);
assert_eq!(after_expiry, None);
}
#[rstest]
#[tokio::test]
async fn a_redis_outage_keeps_pinning_sessions_in_memory() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
store
.set("tenant-a", "session-1", "on-prem-model", HOUR)
.await
.unwrap();
deployment.stop().await.unwrap();
let pinned_before = store.get("tenant-a", "session-1", HOUR).await.unwrap();
store
.set("tenant-a", "session-2", "other-model", HOUR)
.await
.unwrap();
let pinned_during = store.get("tenant-a", "session-2", HOUR).await.unwrap();
assert_eq!(pinned_before, Some("on-prem-model".to_owned()));
assert_eq!(pinned_during, Some("other-model".to_owned()));
}
#[tokio::test]
async fn without_redis_pins_live_in_memory_until_they_expire() {
let time = StepClock::start();
let store = SensitiveRoutePins::with_clock(None, time.clock());
store
.set(
"tenant-a",
"session-1",
"on-prem-model",
Duration::from_secs(10),
)
.await
.unwrap();
let pinned = store.get("tenant-a", "session-1", HOUR).await.unwrap();
time.advance(Duration::from_secs(11));
let expired = store.get("tenant-a", "session-1", HOUR).await.unwrap();
assert_eq!(pinned, Some("on-prem-model".to_owned()));
assert_eq!(expired, None);
}

View file

@ -0,0 +1,648 @@
mod support;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::commands;
use litellm_redis_client::{Operation, RedisClient, Ttl};
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, RequestRedis,
router_cooldowns::RouterCooldowns,
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within};
use tokio::runtime::Handle;
const TTL: Duration = Duration::from_secs(60);
const LONG_DEADLINE: Duration = Duration::from_secs(60);
const READS: [(&str, Number); 4] = [
("7.25", Number::Float(7.25)),
("3", Number::Int(3)),
("0.0", Number::Float(0.0)),
("1e-05", Number::Float(1e-05)),
];
const SETS: [(Number, &str); 4] = [
(Number::Float(0.0), "0.0"),
(Number::Int(12), "12"),
(Number::Float(1e-05), "1e-05"),
(Number::Float(0.30000000000000004), "0.30000000000000004"),
];
const SET_MAXES: [(Number, &str); 3] = [
(Number::Float(2.5), "2.5"),
(Number::Int(100), "100"),
(Number::Float(1e-05), "1e-05"),
];
fn team(id: &str) -> SpendCounter {
SpendCounter::Team {
team_id: id.to_owned(),
window: None,
}
}
fn store(client: Option<RedisClient>) -> Arc<SpendCounters> {
Arc::new(SpendCounters::new(client.map(remote)))
}
/// Mirrored by `tests/unit/redis/stores/test_spend_counters.py`, which checks the same keys and
/// bytes against Python's own key builders and `RedisCache` codec.
#[rstest]
#[case::redis(Reading::Redis(Some(Number::Int(3))), "redis", Some(Number::Int(3)))]
#[case::redis_miss(Reading::Redis(None), "redis", None)]
#[case::memory(
Reading::Memory(Some(Number::Float(2.5))),
"memory",
Some(Number::Float(2.5))
)]
fn a_reading_names_where_it_was_answered(
#[case] reading: Reading,
#[case] source: &str,
#[case] value: Option<Number>,
) {
assert_eq!(<&str>::from(reading), source);
assert_eq!(reading.value(), value);
}
#[rstest]
#[case(
SpendCounter::Key {
token: "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b".to_owned(),
window: None,
},
"spend:key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
)]
#[case(
SpendCounter::Key {
token: "88dc28d0".to_owned(),
window: Some("30d".to_owned()),
},
"spend:key:88dc28d0:window:30d"
)]
#[case(
SpendCounter::Team {
team_id: "team-a".to_owned(),
window: None,
},
"spend:team:team-a"
)]
#[case(
SpendCounter::Team {
team_id: "team-a".to_owned(),
window: Some("1h".to_owned()),
},
"spend:team:team-a:window:1h"
)]
#[case(
SpendCounter::TeamMember {
user_id: "user:with:colons".to_owned(),
team_id: "team-a".to_owned(),
},
"spend:team_member:user:with:colons:team-a"
)]
#[case(SpendCounter::User("user@example.com".to_owned()), "spend:user:user@example.com")]
#[case(SpendCounter::EndUser("customer é 1".to_owned()), "spend:end_user:customer é 1")]
#[case(SpendCounter::Org("org-1".to_owned()), "spend:org:org-1")]
#[case(SpendCounter::Tag("prod:eu".to_owned()), "spend:tag:prod:eu")]
#[case(SpendCounter::Project("proj-1".to_owned()), "spend:project:proj-1")]
#[case(
SpendCounter::ModelAccessGroup("gpt-group".to_owned()),
"spend:model_access_group:gpt-group"
)]
#[case(SpendCounter::ShadowEval("job-7".to_owned()), "spend:shadow_eval:job-7")]
fn every_counter_formats_and_parses_back_to_mains_key(
#[case] counter: SpendCounter,
#[case] key: &str,
) {
assert_eq!(counter.key(), key);
assert_eq!(SpendCounter::parse(key).unwrap(), counter, "{key}");
}
#[rstest]
#[case("spend_db_floor:spend:key:abc")]
#[case("virtual_key_spend:abc:gpt:1d")]
#[case("spend:unknown:abc")]
#[case("spend:team_member:no-team")]
#[case("spend:")]
fn keys_outside_the_spend_counter_keyspace_are_refused(#[case] key: &str) {
assert!(SpendCounter::parse(key).is_err(), "{key}");
}
#[rstest]
#[tokio::test]
async fn reads_values_python_wrote_and_a_clean_miss_is_a_redis_answer(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counters = (0..READS.len())
.map(|index| team(&format!("t{index}")))
.collect::<Vec<_>>();
for ((bytes, _), counter) in READS.iter().zip(&counters) {
raw.send(commands::set(&counter.key(), *bytes, TTL))
.await
.unwrap();
}
let expected = READS
.iter()
.map(|(_, spend)| Some(*spend))
.collect::<Vec<_>>();
for (counter, value) in counters.iter().zip(&expected) {
assert_eq!(store.read(counter).await.unwrap(), Reading::Redis(*value));
}
assert_eq!(
store
.read_many(RedisSend::Now, &counters)
.unwrap()
.await
.unwrap(),
Some(expected)
);
assert_eq!(
store.read(&team("absent")).await.unwrap(),
Reading::Redis(None)
);
}
#[rstest]
#[tokio::test]
async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counter = team("t");
for (spend, bytes) in SETS {
store.set(&counter, spend, None).await.unwrap();
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some(bytes.to_owned())
);
}
for (spend, bytes) in SET_MAXES {
raw.send(commands::del(&[&counter.key()])).await.unwrap();
let repaired = store.set_max(&counter, spend).await.unwrap();
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some(bytes.to_owned())
);
assert_eq!(repaired, Some(Number::Float(spend.as_f64())));
}
store
.set(&counter, Number::Int(5), Some(Duration::from_secs(30)))
.await
.unwrap();
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
panic!("a reset counter has no TTL");
};
assert!(within(remaining, 30), "{remaining:?}");
}
#[rstest]
#[tokio::test]
async fn set_max_only_ever_raises_a_counter(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counter = SpendCounter::Key {
token: "hashed".to_owned(),
window: None,
};
store.increment(&counter, Number::Float(9.5)).await.unwrap();
let lower = store.set_max(&counter, Number::Float(4.0)).await.unwrap();
let higher = store.set_max(&counter, Number::Float(12.0)).await.unwrap();
assert_eq!(lower, Some(Number::Float(9.5)));
assert_eq!(higher, Some(Number::Float(12.0)));
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some("12.0".to_owned())
);
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
panic!("a repaired counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
store.get_local(&counter).unwrap(),
Some(Number::Float(12.0))
);
}
#[rstest]
#[tokio::test]
async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let other_pod = store(Some(connect(&deployment)));
let store = store(Some(client));
let [first, second] = [team("a"), team("b")];
other_pod
.increment(&first, Number::Float(1.25))
.await
.unwrap();
raw.send(commands::expire(&first.key(), Duration::from_secs(5)))
.await
.unwrap();
let totals = store
.increment_many(&[
(first.clone(), Number::Float(0.5)),
(second.clone(), Number::Int(2)),
])
.await
.unwrap();
assert_eq!(totals, [Number::Float(1.75), Number::Float(2.0)]);
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
let Ttl::Expires(remaining) = raw.send(commands::ttl(&first.key())).await.unwrap() else {
panic!("an incremented counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
raw.send(commands::get::<String>(&first.key()))
.await
.unwrap(),
Some("1.75".to_owned())
);
assert_eq!(store.get_local(&first).unwrap(), Some(Number::Float(1.75)));
assert_eq!(store.get_local(&second).unwrap(), Some(Number::Float(2.0)));
}
#[rstest]
#[tokio::test]
async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_and_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
let declared = std::thread::spawn({
let (store, request, counter) = (Arc::clone(&store), request.clone(), counter.clone());
move || store.increment_many_post_call(&request, &[(counter, Number::Float(3.5))])
})
.join()
.expect("a Python thread has no Tokio context, and declaring there must not panic");
request.flush_post().await;
assert!(declared);
let mirrored =
litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await;
assert_eq!(mirrored, Number::Float(3.5));
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some("3.5".to_owned())
);
}
#[rstest]
#[tokio::test]
async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let store = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))]));
request.flush_pre().await;
assert!(trips.seen().is_empty());
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
None
);
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(100)).await;
let increment = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
Topology::Cluster => Operation::Command("INCRBYFLOAT"),
};
assert_eq!(trips.seen(), [increment]);
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some("3.5".to_owned())
);
assert_eq!(store.get_local(&counter).unwrap(), Some(Number::Float(3.5)));
}
#[rstest]
#[tokio::test]
async fn the_budget_read_rides_the_same_pre_call_pipeline_as_the_cooldown_read(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let cooldowns = std::sync::Arc::new(RouterCooldowns::new(
Some(remote(client.clone())),
COOLDOWN_POLICY,
));
let spend = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counters = [
SpendCounter::Key {
token: "hashed".to_owned(),
window: None,
},
team("t"),
];
raw.send(commands::set(&counters[0].key(), "4.5", TTL))
.await
.unwrap();
let active = cooldowns
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
.unwrap();
let budget = spend
.read_many(RedisSend::PreCall(&request), &counters)
.unwrap();
request.flush_pre().await;
let expected = match topology {
Topology::Standalone => vec![Operation::Pipeline { ops: 2 }],
Topology::Cluster => vec![Operation::Command("MGET"); 2],
};
assert_eq!(trips.seen(), expected);
assert_eq!(active.await.unwrap(), vec![]);
assert_eq!(
budget.await.unwrap(),
Some(vec![Some(Number::Float(4.5)), None])
);
assert_eq!(trips.seen().len(), expected.len());
}
#[rstest]
#[tokio::test]
async fn a_seed_lets_the_first_pod_win_and_others_answer_with_its_value(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let first = store(Some(connect(&deployment)));
let second = store(Some(connect(&deployment)));
let counter = SpendCounter::Org("org".to_owned());
let won = first
.seed(&counter, Number::Float(10.0), LostSeed::Answer)
.await
.unwrap();
let lost = second
.seed(&counter, Number::Float(99.0), LostSeed::Increment)
.await
.unwrap();
assert_eq!([won, lost], [Number::Float(10.0), Number::Float(10.0)]);
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some("10.0".to_owned())
);
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
panic!("a seeded counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
second.get_local(&counter).unwrap(),
Some(Number::Float(10.0))
);
}
#[rstest]
#[tokio::test]
async fn delete_drops_memory_and_redis_and_refresh_ttl_reports_whether_redis_held_it(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counters = [team("a"), team("b")];
for counter in &counters {
store.increment(counter, Number::Int(1)).await.unwrap();
}
raw.send(commands::expire(&counters[0].key(), Duration::from_secs(5)))
.await
.unwrap();
assert!(store.refresh_ttl(&counters[0]).await.unwrap());
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counters[0].key())).await.unwrap() else {
panic!("a refreshed counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
store.delete(&counters).await.unwrap();
for counter in &counters {
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
None
);
assert_eq!(store.get_local(counter).unwrap(), None);
}
assert!(!store.refresh_ttl(&counters[0]).await.unwrap());
}
// main's read_spend_counter_cache_value: RedisCache.async_get_cache swallows a failed GET and
// returns None, an authoritative miss. Only the breaker guard's refusal raises into the fallback
// that reads this pod's memory
#[rstest]
#[tokio::test]
async fn a_failed_read_serves_the_local_counter_and_a_failed_increment_drops_the_counter() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = store(Some(connect(&deployment)));
let [kept, failed] = [team("kept"), team("failed")];
store.increment(&kept, Number::Float(2.0)).await.unwrap();
store.increment(&failed, Number::Float(3.0)).await.unwrap();
deployment.stop().await.unwrap();
let read = store.read(&kept).await.unwrap();
let increment = store.increment(&failed, Number::Float(1.0)).await;
assert_eq!(read, Reading::Memory(Some(Number::Float(2.0))));
assert!(increment.is_err());
assert_eq!(store.get_local(&failed).unwrap(), None);
assert_eq!(
store
.read_many(RedisSend::Now, &[kept])
.unwrap()
.await
.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn a_read_with_an_open_breaker_serves_the_local_counter() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let breaker = litellm_circuit_breaker::Breaker::new(
litellm_circuit_breaker::Config {
failure_threshold: 1,
recovery_timeout: Duration::from_secs(60),
..litellm_circuit_breaker::Config::default()
},
Box::new(|_| {}),
);
let store = store(Some(support::connect_guarded(&deployment, breaker)));
let counter = team("t");
store.increment(&counter, Number::Float(2.0)).await.unwrap();
deployment.stop().await.unwrap();
let failed = store.read(&counter).await.unwrap();
let refused = store.read(&counter).await.unwrap();
assert_eq!(failed, Reading::Memory(Some(Number::Float(2.0))));
assert_eq!(refused, Reading::Memory(Some(Number::Float(2.0))));
}
// increment_in_memory calls async_increment_cache(local_only=True), which hands InMemoryCache no
// TTL, so the entry takes InMemoryCache's own default rather than spend_counter_cache's 60 s
#[test]
fn a_local_increment_outlives_the_default_in_memory_ttl() {
let seconds = Arc::new(std::sync::atomic::AtomicU64::new(1_000));
let now = {
let seconds = seconds.clone();
Arc::new(move || Duration::from_secs(seconds.load(std::sync::atomic::Ordering::SeqCst)))
};
let store = SpendCounters::with_clock(None, now);
let counter = team("t");
store.increment_local(&counter, Number::Int(1)).unwrap();
seconds.fetch_add(61, std::sync::atomic::Ordering::SeqCst);
assert_eq!(store.get_local(&counter).unwrap(), Some(Number::Int(1)));
}
#[rstest]
#[tokio::test]
async fn a_failed_post_call_increment_drops_the_counter_from_memory() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
store.increment(&counter, Number::Float(1.0)).await.unwrap();
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]));
deployment.stop().await.unwrap();
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(store.get_local(&counter).unwrap(), None);
}
#[tokio::test]
async fn without_redis_counters_live_in_memory() {
let store = store(None);
let counter = team("t");
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let totals = store
.increment_many(&[
(counter.clone(), Number::Int(2)),
(counter.clone(), Number::Float(0.5)),
])
.await
.unwrap();
let seeded = store
.seed(&counter, Number::Float(1.0), LostSeed::Answer)
.await
.unwrap();
let raised = store.set_max(&counter, Number::Float(7.0)).await.unwrap();
assert_eq!(totals, [Number::Int(2), Number::Float(2.5)]);
assert_eq!(seeded, Number::Float(2.5));
assert_eq!(raised, None);
assert_eq!(
store.read(&counter).await.unwrap(),
Reading::Memory(Some(Number::Float(7.0)))
);
assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))]));
assert!(
store
.read_many(RedisSend::PreCall(&request), &[counter])
.is_none()
);
}
// main's RedisCache: async_increment(refresh_ttl=True), MGET for a batched read, async_set_max as
// a plain EVAL, SET EX NX for a seed and one multi-key DEL per chunk
#[rstest]
#[tokio::test]
async fn spend_counters_send_the_commands_main_sends(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let store = store(Some(connect(&deployment)));
let counters = [team("a"), team("b")];
deployment
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
.await
.unwrap();
store
.increment(&counters[0], Number::Float(1.5))
.await
.unwrap();
store
.read_many(RedisSend::Now, &counters)
.unwrap()
.await
.unwrap()
.unwrap();
store
.set_max(&counters[0], Number::Float(9.0))
.await
.unwrap();
store
.seed(&counters[1], Number::Float(2.0), LostSeed::Answer)
.await
.unwrap();
store.delete(&counters).await.unwrap();
// commandstats also counts the GET, SET and EXPIRE the set_max script runs inside Redis
let calls = async |command| support::calls(&deployment, command).await;
assert_eq!(calls("incrbyfloat").await, 1);
assert_eq!(calls("expire").await, 2);
assert_eq!(calls("eval").await, 1);
assert_eq!(calls("evalsha").await, 0);
assert_eq!(calls("set").await, 2);
assert!(calls("mget").await > 0);
assert_eq!(calls("get").await, 1);
assert!(calls("del").await > 0);
}

View file

@ -10,9 +10,27 @@ use litellm_redis_client::{
CallEvent, Cluster, Credentials, Namespace, Node, Operation, RedisClient, RedisConfig,
RedisEvents, RedisUrl, Target, Timeouts, current_tag,
};
use litellm_redis_store::Remote;
use litellm_redis_store::{
Number, Remote, TierPolicy,
router_cooldowns::Cooldown,
router_usage::{Counter, Kind},
};
use litellm_redis_testing::RedisDeployment;
/// `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const COOLDOWN_POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(1),
tracked_reads: 1000,
default_in_memory_ttl: None,
};
/// `DualCache`'s default `redis_batch_cache_expiry` and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const USAGE_POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(10),
tracked_reads: 1000,
default_in_memory_ttl: None,
};
pub fn connect(deployment: &RedisDeployment) -> RedisClient {
connect_with(deployment, None, None)
}
@ -116,6 +134,103 @@ fn connect_with(
.unwrap()
}
/// What main's `RedisCache` writes for a cooldown: `str(value)` with `EX int(ttl)`, read back with
/// `json.loads` then `ast.literal_eval` (`litellm/caching/redis_cache.py` `set_cache` and
/// `_get_cache_logic`). Keys come from main's `CooldownCache`. The same literals sit in
/// `tests/unit/redis/stores/test_router_cooldowns.py`, checked against main's codec.
pub struct CooldownCase {
pub model_id: &'static str,
pub cooldown: Cooldown,
pub ttl_seconds: f64,
pub key: &'static str,
pub redis_ttl_seconds: u64,
pub str_dict: &'static str,
pub json: &'static str,
}
pub fn cooldowns() -> [CooldownCase; 2] {
[
CooldownCase {
model_id: "dep-a",
cooldown: Cooldown {
exception_received: "RateLimitError: limited é 'quoted'".to_owned(),
status_code: "429".to_owned(),
timestamp: Number::Float(1_730_000_000.25),
cooldown_time: Number::Int(60),
},
ttl_seconds: 90.5,
key: "deployment:dep-a:cooldown",
redis_ttl_seconds: 90,
str_dict: "{'exception_received': \"RateLimitError: limited é 'quoted'\", 'status_code': '429', 'timestamp': 1730000000.25, 'cooldown_time': 60}",
json: "{\"exception_received\": \"RateLimitError: limited \\u00e9 'quoted'\", \"status_code\": \"429\", \"timestamp\": 1730000000.25, \"cooldown_time\": 60}",
},
CooldownCase {
model_id: "dep-b",
cooldown: Cooldown {
exception_received: "Timeout".to_owned(),
status_code: "408".to_owned(),
timestamp: Number::Float(1_730_000_000.5),
cooldown_time: Number::Float(0.5),
},
ttl_seconds: 30.0,
key: "deployment:dep-b:cooldown",
redis_ttl_seconds: 30,
str_dict: "{'exception_received': 'Timeout', 'status_code': '408', 'timestamp': 1730000000.5, 'cooldown_time': 0.5}",
json: "{\"exception_received\": \"Timeout\", \"status_code\": \"408\", \"timestamp\": 1730000000.5, \"cooldown_time\": 0.5}",
},
]
}
/// main's allowed-fails keys for `dep-a`: an empty suffix is the same key as none.
pub struct AllowedFailsCase {
pub suffix: Option<&'static str>,
pub key: &'static str,
}
pub const ALLOWED_FAILS: [AllowedFailsCase; 3] = [
AllowedFailsCase {
suffix: None,
key: "deployment:dep-a:allowed_fails",
},
AllowedFailsCase {
suffix: Some(""),
key: "deployment:dep-a:allowed_fails",
},
AllowedFailsCase {
suffix: Some("timeout"),
key: "deployment:dep-a:allowed_fails:timeout",
},
];
/// main's `RouterCacheEnum` usage keys, including the `None` a missing model prints into the key.
pub struct UsageCase {
pub counter: Counter,
pub key: &'static str,
}
pub fn usage() -> [UsageCase; 2] {
[
UsageCase {
counter: Counter {
deployment_id: "dep".to_owned(),
model: Some("anthropic/claude-sonnet-5-5".to_owned()),
minute: "12-34".to_owned(),
kind: Kind::Tpm,
},
key: "global_router:dep:anthropic/claude-sonnet-5-5:tpm:12-34",
},
UsageCase {
counter: Counter {
deployment_id: "dep".to_owned(),
model: None,
minute: "09-05".to_owned(),
kind: Kind::Rpm,
},
key: "global_router:dep:None:rpm:09-05",
},
]
}
/// A TTL Redis reports for a key written `whole_seconds` ago at most a few seconds back.
pub fn within(ttl: Duration, whole_seconds: u64) -> bool {
ttl <= Duration::from_secs(whole_seconds)

View file

@ -284,6 +284,7 @@ impl Error {
strum::EnumString,
strum::IntoStaticStr,
strum::VariantArray,
serde_with::SerializeDisplay,
)]
pub enum FailureKind {
#[strum(serialize = "RateLimitError")]
@ -471,6 +472,17 @@ mod tests {
assert_eq!(<&'static str>::from(kind), wire);
}
#[rstest]
#[case::rate_limit(FailureKind::RateLimit, "RateLimitError")]
#[case::timeout(FailureKind::Timeout, "Timeout")]
#[case::api_connection(FailureKind::ApiConnection, "APIConnectionError")]
fn failure_kind_serializes_as_its_exception_name(
#[case] kind: FailureKind,
#[case] wire: &str,
) {
assert_eq!(serde_json::to_value(kind).unwrap(), serde_json::json!(wire));
}
#[rstest]
#[case::empty("")]
#[case::lowercase("ratelimiterror")]

View file

@ -342,7 +342,10 @@ pub struct Output {
}
/// How a provider blocked a response the call otherwise served.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)]
#[derive(
Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, strum::IntoStaticStr, strum::EnumString,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum Blocked {
/// A chat response whose first choice was content-filtered (or that has no choices).
@ -402,8 +405,16 @@ pub enum SelectStage {
OverrideCheck(RoutingStrategy),
}
/// Routing events, delivered through [`Host::observe`].
#[derive(Clone, Debug, PartialEq)]
#[derive(Debug)]
pub struct CallbackFiltered {
pub kept: Vec<Arc<Deployment>>,
pub preferred: Option<String>,
}
/// Routing events, delivered through [`Host::observe`]. They serialize as JSON objects tagged by
/// `event` (the variant name), durations in seconds.
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
#[serde(tag = "event")]
pub enum RouterEvent {
/// A pick for `model` (the level's group) begins: its route phase opens.
PickStarted {
@ -413,6 +424,7 @@ pub enum RouterEvent {
model_group: String,
deployment_id: String,
/// How long the pick took (`_acompletion`'s `async_get_available_deployment` service event).
#[serde(serialize_with = "seconds")]
pick: Duration,
},
Succeeded {
@ -425,11 +437,13 @@ pub enum RouterEvent {
},
CooldownSet {
deployment_id: String,
#[serde(serialize_with = "seconds")]
cooldown: Duration,
},
Retry {
model_group: String,
retry: u32,
#[serde(serialize_with = "seconds")]
delay: Duration,
},
Fallback {
@ -584,6 +598,13 @@ impl PickError {
}
}
/// A pick's error serializes as its message.
impl serde::Serialize for PickError {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.collect_str(&self.0)
}
}
impl PartialEq for PickError {
fn eq(&self, other: &Self) -> bool {
format!("{:?}", self.0) == format!("{:?}", other.0)
@ -591,7 +612,8 @@ impl PartialEq for PickError {
}
/// Which error a [`RouterEvent::RetryLogged`] records.
#[derive(Clone, Debug, PartialEq)]
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
#[serde(tag = "cause", content = "error", rename_all = "snake_case")]
pub enum RetryCause {
/// The last attempt's failure.
Attempt,
@ -698,14 +720,19 @@ pub trait Host: Send + Sync {
}
/// `async_callback_filter_deployments`: the host's routing callbacks narrow `candidates` in
/// turn. A filter that raises fails the try with [`Error::Attempt`], which the call retries and
/// falls back from like any attempt failure. A host without filters keeps every candidate.
/// turn, and may name one deployment they prefer, which the pick takes only when the filters
/// after the callbacks keep it. A filter that raises fails the try with [`Error::Attempt`],
/// which the call retries and falls back from like any attempt failure. A host without
/// filters keeps every candidate and prefers none.
async fn callback_filters(
&self,
_ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error> {
Ok(candidates)
) -> Result<CallbackFiltered, Error> {
Ok(CallbackFiltered {
kept: candidates,
preferred: None,
})
}
/// The values of `refs`, one per ref, in order: what the deployment build substitutes for
/// each `os.environ/` reference it found in entry `index` ([`crate::build`]). Asked once per
@ -745,3 +772,83 @@ pub trait Host: Send + Sync {
Ok(Value::Null)
}
}
fn seconds<S: serde::Serializer>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_f64(duration.as_secs_f64())
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
#[rstest]
#[case::selected(
RouterEvent::Selected {
model_group: "gpt-4o".to_owned(),
deployment_id: "azure-eastus".to_owned(),
pick: Duration::from_millis(1500),
},
json!({"event": "Selected", "model_group": "gpt-4o", "deployment_id": "azure-eastus", "pick": 1.5})
)]
#[case::attempt_failed(
RouterEvent::AttemptFailed {
deployment_id: "openai".to_owned(),
kind: FailureKind::RateLimit,
status: Some(429),
},
json!({"event": "AttemptFailed", "deployment_id": "openai", "kind": "RateLimitError", "status": 429})
)]
#[case::cooldown_set(
RouterEvent::CooldownSet { deployment_id: "openai".to_owned(), cooldown: Duration::from_secs(30) },
json!({"event": "CooldownSet", "deployment_id": "openai", "cooldown": 30.0})
)]
#[case::retry(
RouterEvent::Retry { model_group: "gpt-4o".to_owned(), retry: 1, delay: Duration::ZERO },
json!({"event": "Retry", "model_group": "gpt-4o", "retry": 1, "delay": 0.0})
)]
#[case::blocked(
RouterEvent::Blocked { level: 0, model: "claude".to_owned(), blocked: Blocked::ContentFilter },
json!({"event": "Blocked", "level": 0, "model": "claude", "blocked": "content_filter"})
)]
#[case::retry_logged_after_an_attempt(
RouterEvent::RetryLogged {
scope: 0,
level: 0,
model: "gpt-4o".to_owned(),
attempted_retries: None,
cause: RetryCause::Attempt,
},
json!({
"event": "RetryLogged", "scope": 0, "level": 0, "model": "gpt-4o",
"attempted_retries": null, "cause": {"cause": "attempt"},
})
)]
#[case::retry_logged_after_a_pick(
RouterEvent::RetryLogged {
scope: 1,
level: 1,
model: "backup".to_owned(),
attempted_retries: Some(2),
cause: RetryCause::Pick(PickError(Arc::new(Error::UnknownModel { model: "backup".to_owned() }))),
},
json!({
"event": "RetryLogged", "scope": 1, "level": 1, "model": "backup", "attempted_retries": 2,
"cause": {"cause": "pick", "error": "no deployment serves model `backup`"},
})
)]
#[case::mock_failed(
RouterEvent::MockFailed {
mock: MockFailure::RateLimit,
level: 0,
model: "gpt-4o".to_owned(),
num_retries: None,
},
json!({"event": "MockFailed", "mock": "rate_limit", "level": 0, "model": "gpt-4o", "num_retries": null})
)]
fn an_event_serializes_tagged_by_its_name(#[case] event: RouterEvent, #[case] expected: Value) {
assert_eq!(serde_json::to_value(&event).unwrap(), expected);
}
}

View file

@ -124,8 +124,8 @@ pub use crate::{
error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported},
health::{Clock, Health, Outcome, SystemClock},
host::{
Attempt, Blocked, BuildAsk, BuildPath, Built, Host, Output, Provided, RetryCause,
RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage,
Attempt, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Host, Output, Provided,
RetryCause, RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage,
},
outcome::{OutcomeCallbacks, OutcomeWrite},
prerouting::{PreRoute, PreRouted},

View file

@ -7,7 +7,15 @@ use crate::error::{AttemptError, FailureKind};
/// Which mock failure a request asks for.
#[derive(
Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::IntoStaticStr, strum::EnumString,
Clone,
Copy,
Debug,
PartialEq,
Eq,
Deserialize,
serde::Serialize,
strum::IntoStaticStr,
strum::EnumString,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]

View file

@ -456,7 +456,7 @@ impl Core {
let request_order = request.target_order.filter(|_| !inputs.fallback_hop);
// `async_callback_filter_deployments` runs after the cooldown and blocked filters on every
// async candidate list but a deployment named by id and a listed pass-through pick.
if settings.callback_filters
let preferred = if settings.callback_filters
&& !inputs.count_only
&& !resolution.kind.is_single()
&& request.pass_through != Some(PassThrough::Listed)
@ -472,8 +472,12 @@ impl Core {
ctx.excluded_ids = &inputs.excluded_ids;
ctx.skipped_ids = &inputs.skipped_ids;
ctx.target_order = inputs.target_order.or(request_order);
healthy = self.host.callback_filters(&ctx, healthy).await?;
}
let filtered = self.host.callback_filters(&ctx, healthy).await?;
healthy = filtered.kept;
filtered.preferred
} else {
None
};
if inputs.count_only {
return Ok(Candidates {
name: name.to_owned(),
@ -648,8 +652,13 @@ impl Core {
healthy.retain(|d| !d.cannot_run_batches());
}
if let Some(pinned) = request.pinned_deployment_id.as_deref()
&& healthy.iter().any(|d| d.id == pinned)
if let Some(pinned) = [
request.pinned_deployment_id.as_deref(),
preferred.as_deref(),
]
.into_iter()
.flatten()
.find(|id| healthy.iter().any(|d| d.id == *id))
{
healthy.retain(|d| d.id == pinned);
}

View file

@ -17,7 +17,7 @@ use litellm_router_core::{
};
use parking_lot::Mutex;
use rstest::{fixture, rstest};
use serde_json::json;
use serde_json::{Value, json};
#[derive(Default)]
struct ScriptedHost {
@ -1809,3 +1809,55 @@ fn request_for(model: &str) -> RouterRequest {
..RouterRequest::default()
}
}
#[rstest]
#[tokio::test]
async fn every_event_of_a_retried_and_fallen_back_call_serializes_tagged_by_its_name(
config: RouterConfig,
) {
let host = Arc::new(ScriptedHost::failing(
"primary",
vec![failure(FailureKind::RateLimit, 429); 3],
));
let router = router(
config,
json!({"num_retries": 1, "fallbacks": [{"primary": ["backup"]}]}),
host.clone(),
);
router
.call(Op::from("acompletion"), request())
.await
.unwrap();
let tags: Value = host
.events
.lock()
.iter()
.map(|event| serde_json::to_value(event).unwrap()["event"].clone())
.collect();
assert_eq!(
tags,
json!([
"RetriesOpened",
"PickStarted",
"Selected",
"AttemptFailed",
"CooldownSet",
"RetryBudget",
"RetryLogged",
"Retry",
"PickStarted",
"Selected",
"AttemptFailed",
"CooldownSet",
"RetryLogged",
"FallingBack",
"RetryLogged",
"Fallback",
"RetriesOpened",
"PickStarted",
"Selected",
"Succeeded",
"FallbackDone",
])
);
}

View file

@ -105,8 +105,15 @@ def current_service_target() -> str | None:
return _service_target.get()
def with_service_target(target: str) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
"""Run every call of the decorated function, coroutine functions included, under ``service_target(target)``."""
def with_service_target(
target: str, *, preserve_existing: bool = False
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
"""Run calls under the target, or the inherited target when ``preserve_existing`` is true."""
def call_target() -> str:
if preserve_existing:
return current_service_target() or target
return target
def decorate(fn: Callable[_P, _R]) -> Callable[_P, _R]:
if inspect.iscoroutinefunction(fn):
@ -116,14 +123,14 @@ def with_service_target(target: str) -> Callable[[Callable[_P, _R]], Callable[_P
@wraps(fn)
async def run_async(*args: _P.args, **kwargs: _P.kwargs) -> object:
with service_target(target):
with service_target(call_target()):
return await awaitable_fn(*args, **kwargs)
return cast("Callable[_P, _R]", run_async) # cast-ok: same coroutine-returning signature as ``fn``
@wraps(fn)
def run(*args: _P.args, **kwargs: _P.kwargs) -> _R:
with service_target(target):
with service_target(call_target()):
return fn(*args, **kwargs)
return run

View file

@ -518,7 +518,7 @@ def _or_secrets(value: object, read: Callable[[str], object], *names: str) -> ob
return reduce(lambda found, name: found or read(name), names, value)
def _merge_redis_kwargs(
def merge_redis_kwargs(
env_overrides: Mapping[str, object],
) -> dict[str, object]: # mutable-ok: _get_redis_client_logic swaps settings for client objects in place
overrides: Final = {
@ -580,7 +580,7 @@ def _get_redis_client_logic(**env_overrides):
"""
Common functionality across sync + async redis client implementations
"""
redis_kwargs: Final = _merge_redis_kwargs(env_overrides)
redis_kwargs: Final = merge_redis_kwargs(env_overrides)
_startup_nodes: Final = redis_kwargs.get("startup_nodes") or None

View file

@ -4,13 +4,14 @@ Custom A2A Card Resolver for LiteLLM.
Extends the A2A SDK's card resolver to support multiple well-known paths.
"""
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable
from litellm._logging import verbose_logger
from litellm.a2a_protocol.exceptions import A2AAgentCardDiscoveryError
from litellm.constants import LOCALHOST_URL_PATTERNS
from litellm.litellm_core_utils.core_helpers import has_attributes
if TYPE_CHECKING:
from a2a.types import AgentCard
@ -130,15 +131,34 @@ def normalize_agent_card_interfaces(agent_card: "AgentCard", protocol_version: s
return normalized
class _UrlCarrier(Protocol):
url: str
class _InterfacesCarrier(Protocol):
@property
def supported_interfaces(self) -> Sequence[_UrlCarrier]: ...
def _url_of(value: object) -> str | None:
return value.url if has_attributes(value, _UrlCarrier, "url") else None
def _supported_interfaces(agent_card: object) -> Sequence[_UrlCarrier] | None:
if has_attributes(agent_card, _InterfacesCarrier, "supported_interfaces"):
return agent_card.supported_interfaces
return None
def get_agent_card_url(agent_card: "AgentCard") -> str | None:
"""Return the agent endpoint URL from the resolved SDK card."""
url: Final = getattr(agent_card, "url", None)
url: Final = _url_of(agent_card)
if url:
return url
interfaces: Final = getattr(agent_card, "supported_interfaces", None)
interfaces: Final = _supported_interfaces(agent_card)
if interfaces:
return getattr(interfaces[0], "url", None)
return _url_of(interfaces[0])
return None
@ -168,16 +188,16 @@ def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard":
Returns:
The agent card with the URL fixed if necessary
"""
card_url: Final = getattr(agent_card, "url", None)
card_url: Final = _url_of(agent_card)
if card_url and is_localhost_or_internal_url(card_url):
# Normalize base_url to ensure it ends with /
fixed_url: Final = base_url.rstrip("/") + "/"
agent_card.url = fixed_url
interfaces: Final = getattr(agent_card, "supported_interfaces", None)
interfaces: Final = _supported_interfaces(agent_card)
if interfaces:
interface_url: Final = getattr(interfaces[0], "url", None)
interface_url: Final = _url_of(interfaces[0])
if interface_url and is_localhost_or_internal_url(interface_url):
interfaces[0].url = base_url.rstrip("/") + "/"
@ -194,6 +214,9 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):
- /agentCard/v1.0
"""
if TYPE_CHECKING:
base_url: str
async def get_agent_card(
self,
relative_card_path: str | None = None,

View file

@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils.
Maps A2A SDK exceptions to LiteLLM A2A exception types.
"""
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_logger
from litellm.a2a_protocol.card_resolver import (
@ -21,6 +21,7 @@ from litellm.constants import CONNECTION_ERROR_PATTERNS
if TYPE_CHECKING:
from a2a.client import Client as A2AClientType
from a2a.types import AgentCard
_A2A_KIND_DIALECT_HINT: Final = (
@ -165,7 +166,7 @@ def map_a2a_exception(
async def handle_a2a_localhost_retry(
error: A2ALocalhostURLError,
agent_card: Any,
agent_card: "AgentCard | None",
a2a_client: "A2AClientType",
is_streaming: bool = False,
) -> "A2AClientType":

View file

@ -10,7 +10,7 @@ A2A Streaming Events (in order):
4. Status update (kind: "status-update") - Final status "completed" with final=true
"""
from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
from collections.abc import AsyncIterator, Mapping
from types import MappingProxyType
from typing import Any, Final
@ -126,10 +126,7 @@ class A2ACompletionBridgeHandler:
@staticmethod
async def _acompletion(completion_params: Mapping[str, object]) -> ModelResponse | CustomStreamWrapper:
acompletion_fn: Final[Callable[..., Coroutine[object, object, ModelResponse | CustomStreamWrapper]]] = vars(
litellm
)["acompletion"]
return await acompletion_fn(**completion_params)
return await litellm.acompletion(**completion_params)
@staticmethod
async def handle_non_streaming(

View file

@ -14,7 +14,7 @@ import datetime
import uuid
from collections.abc import AsyncIterator, Coroutine, Mapping
from types import MappingProxyType, ModuleType
from typing import TYPE_CHECKING, Final, Optional, cast
from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
from typing_extensions import TypedDict, Unpack
@ -24,6 +24,7 @@ from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
from litellm.a2a_protocol.utils import A2ARequestUtils
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -36,7 +37,6 @@ if TYPE_CHECKING:
from a2a.client import Client as A2AClientType
from a2a.client import ClientCallContext as A2ACallContextType
from a2a.compat.v0_3.types import (
AgentCard,
Message,
SendMessageRequest,
SendMessageResponse,
@ -45,6 +45,7 @@ if TYPE_CHECKING:
SendStreamingMessageSuccessResponse,
Task,
)
from a2a.types import AgentCard
from a2a.types.a2a_pb2 import SendMessageRequest as CoreSendMessageRequest
from a2a.types.a2a_pb2 import StreamResponse as CoreStreamResponse
@ -95,6 +96,12 @@ class _SendMessageOptions(TypedDict, total=False, extra_items=object):
pass
class _StreamingRequestConverter(Protocol):
def to_core_send_message_request(
self, compat_req: "SendStreamingMessageRequest", /
) -> "CoreSendMessageRequest": ...
def _set_usage_on_logging_obj(
kwargs: Mapping[str, object],
prompt_tokens: int,
@ -375,7 +382,9 @@ async def _stream_messages(
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
)
pb_request: Final[CoreSendMessageRequest] = _a2a_conversions.to_core_send_message_request(request)
pb_request: Final[CoreSendMessageRequest] = require_attributes(
_a2a_conversions, _StreamingRequestConverter, "to_core_send_message_request"
).to_core_send_message_request(request)
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
compat_chunk = _to_compat_stream_response(event, request_id=request.id)
yield SendStreamingMessageResponse(root=compat_chunk)

View file

@ -427,6 +427,7 @@ def update_headers_with_filtered_beta(
provider: str,
api_base: str | None = None,
unlisted: UnlistedBetas = "drop",
keep: frozenset[str] = frozenset(),
) -> dict:
"""
Update headers dict by filtering and transforming anthropic-beta header values.
@ -437,6 +438,7 @@ def update_headers_with_filtered_beta(
provider: Provider name
api_base: Deployment api_base, used to pick the allowlist row (Foundry hosts use ``azure_ai``)
unlisted: Whether a beta missing from the row is dropped or forwarded as sent
keep: Betas the routed model accepts, forwarded as sent whatever the row says
Returns:
Updated headers dict
@ -446,13 +448,20 @@ def update_headers_with_filtered_beta(
return headers
# Parse existing beta headers
beta_values: Final = [b.strip() for b in existing_beta.split(",") if b.strip()]
beta_values: Final[tuple[str, ...]] = tuple(b.strip() for b in existing_beta.split(",") if b.strip())
# Filter and transform based on provider
filtered_beta_values: Final = filter_and_transform_beta_headers(
beta_headers=beta_values,
provider=resolve_beta_provider(provider, api_base),
unlisted=unlisted,
filtered_beta_values: Final = tuple(
dict.fromkeys(
[
*filter_and_transform_beta_headers(
beta_headers=[b for b in beta_values if b not in keep],
provider=resolve_beta_provider(provider, api_base),
unlisted=unlisted,
),
*(b for b in beta_values if b in keep),
]
)
)
# Update or remove the header

View file

@ -13,10 +13,10 @@ import json
import logging
import time
import traceback
from collections.abc import Generator, Mapping
from collections.abc import Awaitable, Callable, Generator, Mapping
from contextlib import contextmanager
from enum import Enum
from typing import Any, Final, Literal, cast
from typing import Any, Final, Literal, Protocol, cast
from pydantic import BaseModel, ConfigDict, TypeAdapter
@ -25,6 +25,7 @@ from litellm._internal_context import current_service_target, service_target
from litellm._logging import verbose_logger
from litellm.constants import CACHED_STREAMING_CHUNK_DELAY
from litellm.integrations.otel.runtime import phase_span
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.litellm_core_utils.model_param_helper import (
ModelParamHelper,
embeds_thought_signature_in_tool_call_id,
@ -52,6 +53,16 @@ _REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter(
)
class _SupportsAsyncPing(Protocol):
@property
def ping(self) -> Callable[[], Awaitable[object]] | None: ...
class _SupportsAsyncDeleteCacheKeys(Protocol):
@property
def delete_cache_keys(self) -> Callable[[object], Awaitable[object]] | None: ...
def _keyed_provider_params_in_extra_body(kwargs: Mapping[str, object]) -> str:
raw_extra_body: Final = kwargs.get("extra_body")
if not isinstance(raw_extra_body, dict):
@ -1032,14 +1043,16 @@ class Cache:
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
await self.cache.batch_cache_write(cache_key, cached_data, **kwargs)
async def ping(self):
cache_ping: Final = getattr(self.cache, "ping")
async def ping(self) -> object:
cache_ping: Final = require_attributes(self.cache, _SupportsAsyncPing, "ping").ping
if cache_ping:
return await cache_ping()
return None
async def delete_cache_keys(self, keys):
cache_delete_cache_keys: Final = getattr(self.cache, "delete_cache_keys")
async def delete_cache_keys(self, keys: object) -> object:
cache_delete_cache_keys: Final = require_attributes(
self.cache, _SupportsAsyncDeleteCacheKeys, "delete_cache_keys"
).delete_cache_keys
if cache_delete_cache_keys:
return await cache_delete_cache_keys(keys)
return None

View file

@ -22,6 +22,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable,
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError
from typing_extensions import TypedDict
import litellm
from litellm._internal_context import is_internal_call, post_response_phase
@ -90,6 +91,10 @@ class CachingHandlerResponse(LiteLLMBaseModel):
in_memory_cache_obj: Final = InMemoryCache()
class _CacheCallKwargs(TypedDict, total=False, extra_items=object):
pass
def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]:
"""
The caching handler is stored on the Logging object
@ -455,7 +460,7 @@ class LLMCachingHandler:
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
args = args or ()
# Now that we confirmed caching will happen, prepare kwargs
new_kwargs: Final = kwargs.copy()
new_kwargs: Final = dict[str, object](kwargs)
new_kwargs.update(
convert_args_to_kwargs(
self.original_function,
@ -470,7 +475,7 @@ class LLMCachingHandler:
print_verbose("Checking Sync Cache")
logging_obj.reset_caching_details()
with response_cache_phase("get"):
cached_result = litellm.cache.get_cache(**new_kwargs)
cached_result = litellm.cache.get_cache(**_CacheCallKwargs(**new_kwargs))
if is_response_without_output(cached_result):
verbose_logger.debug("LiteLLM Cache: cached response has no output, treating it as a miss")
return CachingHandlerResponse(cached_result=None)
@ -480,7 +485,7 @@ class LLMCachingHandler:
pass
else:
call_type = original_function.__name__
cached_result = self._convert_cached_result_to_model_response(
converted_result: Final = self._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=call_type,
kwargs=kwargs,
@ -508,16 +513,16 @@ class LLMCachingHandler:
logging_obj=logging_obj,
model=f"{custom_llm_provider}/{model}",
kwargs=kwargs,
cached_result=cached_result,
cached_result=converted_result,
is_async=False,
custom_llm_provider=custom_llm_provider,
)
if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks(
cached_result=cached_result
cached_result=converted_result
):
logging_obj.handle_sync_cache_hit_success_callbacks(
result=cached_result,
result=converted_result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
@ -527,8 +532,8 @@ class LLMCachingHandler:
or self.request_kwargs.get("cache_key")
or litellm.cache.get_cache_key(**self.request_kwargs)
)
_set_cached_hidden_param(cached_result, "cache_key", cache_key)
return CachingHandlerResponse(cached_result=cached_result)
_set_cached_hidden_param(converted_result, "cache_key", cache_key)
return CachingHandlerResponse(cached_result=converted_result)
return CachingHandlerResponse(cached_result=cached_result)
_sync_get_cache = sync_get_cache
@ -950,22 +955,14 @@ class LLMCachingHandler:
def _convert_cached_result_to_model_response(
self,
cached_result: Any,
cached_result: object,
call_type: str,
kwargs: dict[str, object],
logging_obj: LiteLLMLoggingObj,
model: str,
args: tuple[object, ...],
custom_llm_provider: str | None = None,
) -> (
ModelResponse
| TextCompletionResponse
| EmbeddingResponse
| RerankResponse
| TranscriptionResponse
| CustomStreamWrapper
| None
):
) -> object:
"""
Internal method to process the cached result
@ -1141,7 +1138,7 @@ class LLMCachingHandler:
self,
result: object,
original_function: Callable,
kwargs: dict[str, Any],
kwargs: Mapping[str, object],
args: tuple[object, ...] | None = None,
):
"""
@ -1169,7 +1166,7 @@ class LLMCachingHandler:
return
cache: Final = litellm.cache
new_kwargs: Final = kwargs.copy()
new_kwargs: Final = dict(kwargs)
new_kwargs.update(
convert_args_to_kwargs(
original_function,
@ -1206,7 +1203,7 @@ class LLMCachingHandler:
)
)
else:
create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs))
create_cache_write_task(lambda: cache.async_add_cache(result, **_CacheCallKwargs(**new_kwargs)))
def sync_set_cache(
self,

View file

@ -12,14 +12,14 @@ import asyncio
import itertools
import logging
import time
from collections.abc import Mapping, Sequence
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from threading import Lock
from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS
from litellm.types.caching import RedisPipelineIncrementOperation
from .base_cache import BaseCache
@ -91,7 +91,9 @@ class DualCache(BaseCache):
self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size)
self._last_redis_batch_access_time_lock = Lock()
self.redis_batch_cache_expiry = (
default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10
default_redis_batch_cache_expiry
or litellm.default_redis_batch_cache_expiry
or DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS
)
self.default_in_memory_ttl = default_in_memory_ttl or litellm.default_in_memory_ttl
self.default_redis_ttl = default_redis_ttl or litellm.default_redis_ttl
@ -497,7 +499,7 @@ class DualCache(BaseCache):
batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache)
return None if batch is None else await self._set_on_batch(batch, key, value, ttl)
async def async_delete_cache_pre_call(self, key: str) -> BatchResult[None] | None:
async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None:
"""Memory now, the Redis DEL on the request's pipeline; None when no pipeline is open, so the caller
takes its direct path."""
batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache)

View file

@ -622,9 +622,11 @@ class RequestRedisBatches:
self._deadline_flush = asyncio.ensure_future(self.flush_post_call())
async def flush_all(self) -> None:
"""Batch whatever is still declared (write-backs nobody awaits) before the request scope closes."""
"""Batch whatever is still declared (write-backs nobody awaits) before the request scope closes. The pre-call
flush hooks are skipped: they declare reads for later readers, and none are left."""
await asyncio.gather(
*(batch.flush() for batch in self._batches.values() if batch.pending), self.flush_native_pre()
*(batch.flush() for batch in self._batches.values() if batch.pending),
*(() if self._native is None else (self._native.flush_pre(),)),
)
async def flush_post_call(self) -> None:

View file

@ -40,9 +40,9 @@ from litellm.constants import (
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
REDIS_TIMEOUT_LOG_INTERVAL,
)
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, has_attributes
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
from litellm.redis.store_settings import RedisStoreSettings
from litellm.redis.store_settings import RedisStoreSettings, breaker_settings
from litellm.types.caching import (
RedisPipelineIncrementOperation,
RedisPipelineLpopOperation,
@ -97,6 +97,22 @@ class _AsyncRedisCommands(Protocol):
def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
class _AsyncLuaScript(Protocol):
def __call__(
self, keys: Sequence[str], args: Sequence[str | bytes | int | float], client: object
) -> Awaitable[object]: ...
class _SupportsRegisterScript(Protocol):
def register_script(self, script: str, /) -> _AsyncLuaScript: ...
class _SupportsScriptLoad(Protocol):
def script_load(self, script: str, /) -> object: ...
def evalsha(self, sha: object, numkeys: int, /, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
_GENERIC_CALLER_MODULES: Final = frozenset(
{
__name__,
@ -661,7 +677,7 @@ class RedisCache(BaseCache):
):
from litellm._service_logger import ServiceLogging
from .._redis import get_redis_client, get_redis_connection_pool
from .._redis import get_redis_client, get_redis_connection_pool, merge_redis_kwargs
redis_kwargs: Final = {}
if host is not None:
@ -715,9 +731,10 @@ class RedisCache(BaseCache):
else:
super().__init__() # defaults to 60s
self.store_settings = RedisStoreSettings(
redis_kwargs=MappingProxyType(dict(redis_kwargs)),
redis_kwargs=MappingProxyType(merge_redis_kwargs(redis_kwargs)),
namespace=namespace,
default_ttl=self.default_ttl,
breaker=breaker_settings(),
)
def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:
@ -843,6 +860,9 @@ class RedisCache(BaseCache):
def _async_commands(self) -> _AsyncRedisCommands:
return self.init_async_client()
def _async_client_object(self) -> object:
return self.init_async_client()
def check_and_fix_namespace(self, key: str) -> str:
"""
Make sure each key starts with the given namespace
@ -1116,8 +1136,8 @@ class RedisCache(BaseCache):
Kept separate from async_register_script so each loop caches its own
executor; see that method for why the binding must be per loop.
"""
_redis_client: Final[Any] = self.init_async_client()
if hasattr(_redis_client, "register_script"):
_redis_client: Final = self._async_client_object()
if has_attributes(_redis_client, _SupportsRegisterScript, "register_script"):
registered_script: Final = _redis_client.register_script(script)
async def standalone_executor(
@ -1130,7 +1150,7 @@ class RedisCache(BaseCache):
return standalone_executor
if hasattr(_redis_client, "script_load"):
if has_attributes(_redis_client, _SupportsScriptLoad, "script_load"):
script_sha: Final = _redis_client.script_load(script)
async def cluster_executor(

View file

@ -43,6 +43,7 @@ if TYPE_CHECKING:
_CACHE_RESPONSE_ADAPTER: Final[TypeAdapter[str | bytes | None]] = TypeAdapter(
str | bytes | None, config=ConfigDict(strict=True, hide_input_in_errors=True)
)
_DECODED_CACHE_VALUE: Final = TypeAdapter(object)
class RedisSemanticCache(BaseCache):
@ -407,10 +408,10 @@ class RedisSemanticCache(BaseCache):
# Convert string representation to Python object
try:
return json.loads(response_text)
return _DECODED_CACHE_VALUE.validate_python(json.loads(response_text))
except json.JSONDecodeError:
try:
return ast.literal_eval(response_text)
return _DECODED_CACHE_VALUE.validate_python(ast.literal_eval(response_text))
except (ValueError, SyntaxError) as e:
print_verbose(f"Error parsing cached response: {e}")
return None

View file

@ -39,7 +39,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import (
parse_chat_allowed_tools_tool_choice,
responses_allowed_tools_tool_choice,
)
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, iter_items
from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params
from litellm.litellm_core_utils.prompt_templates.prompt_cache_breakpoint_utils import (
prompt_cache_breakpoint_with_readable_mode,
@ -54,6 +54,7 @@ from litellm.llms.base_llm.bridges.completion_transformation import (
CompletionTransformationBridge,
)
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.responses.additional_tools import additional_tools_in_input
from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool
from litellm.responses.sse_output_recovery import (
parse_sse_json_chunk,
@ -348,6 +349,9 @@ _TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_IT
_TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items"
_DECODED_JSON: Final = TypeAdapter(object)
def _json_field(value: object, key: str) -> object:
return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None # cast-ok: untyped JSON
@ -426,6 +430,17 @@ def replayed_tool_search_tool_names(request: object) -> tuple[str, ...]:
return tuple(chain.from_iterable(_tool_definition_name(tool) for tool in replayed_tool_search_tools(request)))
def _top_level_namespaces(request: object) -> Iterator[object]:
tools: Final = _json_list(_json_field(request, "tools")) or ()
return (tool for tool in tools if _json_field(tool, "type") == "namespace")
def nested_tool_member_names(request: object) -> tuple[str, ...]:
tools: Final = (*_top_level_namespaces(request), *additional_tools_in_input(_json_field(request, "input")))
members: Final = chain.from_iterable(_namespace_members(tool) for tool in tools)
return tuple(chain.from_iterable(_tool_definition_name(member) for member in members))
def _loads_any_tool(message: object, names: frozenset[str]) -> bool:
loaded_names: Final = chain.from_iterable(_tool_definition_name(tool) for tool in _message_loaded_tools(message))
return any(name in names for name in loaded_names)
@ -845,8 +860,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Handle message items with output_text content
if item_type == "message":
content_list: Final = item.get("content", [])
for content_item in content_list:
content_list: Final = _DECODED_JSON.validate_python(item.get("content", []))
for content_item in iter_items(content_list):
if isinstance(content_item, dict):
content_type = content_item.get("type")
if content_type == "output_text":

View file

@ -159,6 +159,7 @@ DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3)
DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1)
DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5)
DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = get_env_float("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", 1.0)
DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS: Final = 10
DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLING_RETRIES", 5)
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)
DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250)

View file

@ -8,7 +8,7 @@ from datetime import datetime
from functools import lru_cache
from itertools import accumulate
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, cast
from httpx import Response
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
@ -24,7 +24,12 @@ from litellm.constants import (
REALTIME_METERED_USAGE_KEY,
)
from litellm.litellm_core_utils.audio_utils.utils import normalized_audio_duration_seconds
from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, billed_usage_cost
from litellm.litellm_core_utils.core_helpers import (
as_str_mapping,
attribute_or_none,
billed_usage_cost,
require_attributes,
)
from litellm.litellm_core_utils.cost_discount import resolve_cost_discount
from litellm.litellm_core_utils.decisions_pricing import (
decisions_pricing_fields,
@ -1188,6 +1193,10 @@ def _model_contains_known_llm_provider(model: str) -> bool:
return _provider_prefix in LlmProvidersSet
class _SupportsGet(Protocol):
def get(self, key: str, /) -> object: ...
def _get_response_model(completion_response: object) -> str | None:
"""
Extract the model name from a completion response object.
@ -1219,7 +1228,7 @@ _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = {
}
def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None:
def _map_traffic_type_to_service_tier(traffic_type: object) -> str | None:
"""
Map a Gemini usageMetadata.trafficType value to a LiteLLM service_tier string.
@ -1273,11 +1282,13 @@ def _served_service_tier(completion_response: object, usage_object: Usage | None
usage_tier: Final = _extract_service_tier(usage_object)
if isinstance(usage_tier, str):
return usage_tier
hidden_params: Final = getattr(completion_response, "_hidden_params", None)
hidden_params: Final[object] = getattr(completion_response, "_hidden_params", None)
if hidden_params is None:
return None
provider_specific: Final = hidden_params.get("provider_specific_fields") or {}
raw_traffic_type: Final = provider_specific.get("traffic_type")
provider_specific: Final[object] = (
require_attributes(hidden_params, _SupportsGet, "get").get("provider_specific_fields") or {}
)
raw_traffic_type: Final = require_attributes(provider_specific, _SupportsGet, "get").get("traffic_type")
if not raw_traffic_type:
return None
return _map_traffic_type_to_service_tier(raw_traffic_type) or "default"
@ -1341,8 +1352,16 @@ def get_usage_object(
return None
class _DumpedUsageFields(TypedDict, total=False, extra_items=object):
pass
def _dumped_usage_fields(usage_model: BaseModel) -> Mapping[str, object]:
return usage_model.model_dump()
def _usage_from_model(usage_model: BaseModel) -> Usage:
return Usage(**usage_model.model_dump())
return Usage(**_DumpedUsageFields(**_dumped_usage_fields(usage_model)))
def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper:

View file

@ -12,10 +12,11 @@
import datetime
import enum
from collections.abc import Mapping, Sequence
from typing import Any, Final
from typing import Any, Final, Protocol, SupportsIndex, SupportsInt, runtime_checkable
import httpx
import openai
from typing_extensions import Buffer
import litellm
from litellm.types.utils import LiteLLMCommonStrings
@ -1250,6 +1251,19 @@ class BlockedPiiEntityError(Exception):
super().__init__(self.message)
@runtime_checkable
class _SupportsTrunc(Protocol):
def __trunc__(self) -> int: ...
def _status_code_int(value: object) -> int:
if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex, _SupportsTrunc)):
return int(value)
raise TypeError(
f"int() argument must be a string, a bytes-like object or a real number, not '{type(value).__name__}'"
)
class MidStreamFallbackError(ServiceUnavailableError):
def __init__(
self,
@ -1265,8 +1279,8 @@ class MidStreamFallbackError(ServiceUnavailableError):
is_pre_first_chunk: bool = False,
emitted_disqualifying_content: bool = False,
):
original_status: Final = getattr(original_exception, "status_code", None)
self.status_code = int(original_status) if original_status is not None else 503
original_status: Final[object] = getattr(original_exception, "status_code", None)
self.status_code = _status_code_int(original_status) if original_status is not None else 503
self.message = f"litellm.MidStreamFallbackError: {message}"
self.model = model
self.llm_provider = llm_provider
@ -1309,7 +1323,7 @@ class MidStreamFallbackError(ServiceUnavailableError):
)
# Restore the propagated status and original response/request objects
self.status_code = int(original_status) if original_status is not None else 503
self.status_code = _status_code_int(original_status) if original_status is not None else 503
self.response = _saved_response
self.request = _saved_request
self.message = _saved_message

View file

@ -14,10 +14,11 @@ import contextlib
import itertools
import json
import logging
import math
import secrets
from collections.abc import AsyncIterable, AsyncIterator, Mapping
from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType, ModuleType
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
import httpx
@ -36,6 +37,7 @@ from litellm.constants import (
from litellm.harness.context import GatewayTarget
from litellm.harness.errors import HarnessError, HarnessInstallFailed
from litellm.harness.types import Harness, Usage
from litellm.litellm_core_utils.core_helpers import has_attributes
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -45,7 +47,7 @@ if TYPE_CHECKING:
from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from uvicorn import Server
from uvicorn import Config, Server
verbose_logger: Final = logging.getLogger("LiteLLM")
_JSON_VALUE_ADAPTER: Final = TypeAdapter(object, config=ConfigDict(strict=True))
@ -110,9 +112,22 @@ class _ResponsesModule(Protocol):
StreamingResponse: type[StreamingResponse]
class _UvicornModule(Protocol):
Config: type[Config]
Server: type[Server]
class _CapturesSignals(Protocol):
capture_signals: Callable[[], contextlib.AbstractContextManager[None]]
class _InstallsSignalHandlers(Protocol):
install_signal_handlers: Callable[[], None]
@dataclass(frozen=True)
class _ServerDeps:
uvicorn: ModuleType
uvicorn: _UvicornModule
applications: _ApplicationsModule
routing: _RoutingModule
responses: _ResponsesModule
@ -438,6 +453,14 @@ def _noop() -> None:
return None
def _signal_capture(server: object) -> _CapturesSignals | None:
return server if has_attributes(server, _CapturesSignals, "capture_signals") else None
def _signal_handler_installer(server: object) -> _InstallsSignalHandlers | None:
return server if has_attributes(server, _InstallsSignalHandlers, "install_signal_handlers") else None
class ModelEndpoint:
"""Local HTTP endpoint for one harness session. Use as an async context manager."""
@ -521,7 +544,7 @@ class ModelEndpoint:
await asyncio.sleep(DEFAULT_POLLING_INTERVAL)
def _build_server(self, deps: _ServerDeps) -> Server:
config = deps.uvicorn.Config(
config: Final = deps.uvicorn.Config(
self._build_app(deps),
host=HARNESS_ENDPOINT_HOST,
port=0,
@ -529,14 +552,16 @@ class ModelEndpoint:
log_level="warning",
access_log=False,
lifespan="off",
timeout_graceful_shutdown=HARNESS_PROCESS_KILL_GRACE_SECONDS,
timeout_graceful_shutdown=math.ceil(HARNESS_PROCESS_KILL_GRACE_SECONDS),
)
server = deps.uvicorn.Server(config)
server: Final = deps.uvicorn.Server(config)
# Never touch the host process's signal handlers.
if hasattr(server, "capture_signals"):
server.capture_signals = contextlib.nullcontext
if hasattr(server, "install_signal_handlers"):
server.install_signal_handlers = _noop
signal_capture: Final = _signal_capture(server)
if signal_capture is not None:
signal_capture.capture_signals = contextlib.nullcontext
handler_installer: Final = _signal_handler_installer(server)
if handler_installer is not None:
handler_installer.install_signal_handlers = _noop
return server
def _build_app(self, deps: _ServerDeps) -> Starlette:

View file

@ -6,12 +6,25 @@ from pathlib import Path
from typing import Any, Final, Protocol, cast, get_type_hints, runtime_checkable
import litellm
from litellm.litellm_core_utils.core_helpers import has_attributes
from litellm.litellm_core_utils.token_counter import get_image_type
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.types.files import FILE_MIME_TYPES, FileType
from litellm.types.images.main import ImageEditOptionalRequestParams
class _ReadsBytes(Protocol):
def read(self, size: int, /) -> bytes: ...
class _Tells(Protocol):
def tell(self) -> int: ...
class _Seeks(Protocol):
def seek(self, offset: int, /) -> object: ...
class ImageEditRequestUtils:
@staticmethod
def get_optional_params_image_edit(
@ -114,12 +127,12 @@ class ImageEditRequestUtils:
bytes_data = image_data[:100]
else:
# For other types, try to read if possible
if hasattr(image_data, "read"):
current_pos = getattr(image_data, "tell", lambda: 0)()
if hasattr(image_data, "seek"):
if has_attributes(image_data, _ReadsBytes, "read"):
current_pos = image_data.tell() if has_attributes(image_data, _Tells, "tell") else 0
if has_attributes(image_data, _Seeks, "seek"):
image_data.seek(0)
bytes_data = image_data.read(100)
if hasattr(image_data, "seek"):
if has_attributes(image_data, _Seeks, "seek"):
image_data.seek(current_pos)
else:
return FILE_MIME_TYPES[FileType.PNG] # Default fallback

View file

@ -53,11 +53,13 @@ async def fetch_user_spend_rows(
today: datetime.date,
baseline_days: int,
) -> tuple[UserSpendRow, ...]:
from litellm.proxy.utils import query_raw_rows
today_str: Final = today.strftime("%Y-%m-%d")
month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d")
baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d")
return USER_SPEND_ROWS_ADAPTER.validate_python(
await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str)
await query_raw_rows(prisma_client.db, USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str)
)

View file

@ -7,11 +7,12 @@ import json
import os
import random
import types
from collections.abc import Mapping
from typing import Final
from collections.abc import Mapping, Sequence
from typing import Final, Protocol
import httpx
from pydantic import BaseModel
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -29,6 +30,22 @@ from litellm.types.integrations.argilla import (
from litellm.types.utils import StandardLoggingPayload
class _ArgillaDataset(TypedDict):
id: ReadOnly[str]
class _ArgillaDatasetsBody(TypedDict, total=False):
items: ReadOnly[Sequence[_ArgillaDataset]]
class _ArgillaDatasetsResponse(Protocol):
def json(self) -> _ArgillaDatasetsBody: ...
def _datasets_body(response: _ArgillaDatasetsResponse) -> _ArgillaDatasetsBody:
return response.json()
def is_serializable(value):
non_serializable_types: Final = (
types.CoroutineType,
@ -102,7 +119,7 @@ class ArgillaLogger(CustomBatchLogger):
url=f"{_credentials_base_url}/api/v1/me/datasets?name={_credentials_dataset_name}",
headers={"X-Argilla-Api-Key": _credentials_api_key},
)
json_response: Final = dataset_response.json()
json_response: Final = _datasets_body(dataset_response)
if (
"items" in json_response
and isinstance(json_response["items"], list)

View file

@ -1,6 +1,6 @@
import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import ConfigDict, TypeAdapter
from typing_extensions import ReadOnly, TypedDict, override
@ -10,10 +10,12 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor
BaseLLMObsOTELAttributes,
safe_set_attribute,
)
from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes, subscript
from litellm.litellm_core_utils.redact_messages import (
should_redact_message_logging,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall, StandardLoggingPayload
if TYPE_CHECKING:
@ -33,32 +35,40 @@ from litellm.integrations._types.open_inference import (
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
class _SupportsGetWithDefault(Protocol):
def get(self, key: str, default: object = None, /) -> object: ...
class ArizeOTELAttributes(BaseLLMObsOTELAttributes):
@staticmethod
@override
def set_messages(span: "Span", kwargs: dict[str, Any]):
def set_messages(span: "Span", kwargs: Mapping[str, object]):
messages: Final = kwargs.get("messages")
# for /chat/completions
# https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions
if messages:
last_message: Final = messages[-1]
last_message: Final = subscript(messages, -1)
safe_set_attribute(
span,
SpanAttributes.INPUT_VALUE,
last_message.get("content", ""),
require_attributes(last_message, _SupportsGetWithDefault, "get").get("content", ""),
)
# LLM_INPUT_MESSAGES shows up under `input_messages` tab on the span page.
for idx, msg in enumerate(messages):
for idx, msg in enumerate(iter_items(messages)):
prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}"
# Set the role per message.
safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", msg.get("role"))
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
require_attributes(msg, _SupportsGetWithDefault, "get").get("role"),
)
# Set the content per message.
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
msg.get("content", ""),
require_attributes(msg, _SupportsGetWithDefault, "get").get("content", ""),
)
# Additive: emit structured tool_calls / multimodal content
@ -663,12 +673,11 @@ def _to_plain_dict(value):
def _get_tool_calls(message) -> list | None:
"""Return ``message.tool_calls`` only when it's a non-empty list.
Works for dicts and Pydantic message objects via ``_safe_get``.
"""Return the message's tool calls, Claude's server tool calls first, or
``None`` when there are none. Works for dicts and Pydantic message objects.
"""
tool_calls: Final = _safe_get(message, "tool_calls")
return tool_calls if isinstance(tool_calls, list) and tool_calls else None
tool_calls: Final = tool_calls_with_server_tool_calls(message)
return list(tool_calls) if tool_calls else None
class _NormalizedToolCallFunction(TypedDict):

View file

@ -18,10 +18,11 @@ import time
import traceback
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final, TypeVar
from typing import Final, Protocol, TypeVar
from urllib.parse import urlparse
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.integrations.batch_utils import (
@ -52,6 +53,19 @@ MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType(
)
class _OAuthTokenBody(TypedDict, total=False):
access_token: ReadOnly[str]
expires_in: ReadOnly[int]
class _OAuthTokenResponse(Protocol):
def json(self) -> _OAuthTokenBody: ...
def _oauth_token_body(response: _OAuthTokenResponse) -> _OAuthTokenBody:
return response.json()
class AzureSentinelLogger(CustomBatchLogger):
"""
Logger that sends LiteLLM logs to Azure Sentinel via Azure Monitor Logs Ingestion API
@ -226,7 +240,7 @@ class AzureSentinelLogger(CustomBatchLogger):
if response.status_code != 200:
raise Exception(f"Failed to get OAuth2 token: {response.status_code} - {response.text}")
token_response: Final = response.json()
token_response: Final = _oauth_token_body(response)
self.oauth_token = token_response.get("access_token")
expires_in: Final = token_response.get("expires_in", 3600)

View file

@ -5,7 +5,7 @@ BitBucket API client for fetching .prompt files from BitBucket repositories.
import base64
import urllib.parse
from collections.abc import Mapping
from typing import Final, TypedDict
from typing import Final, Protocol, TypedDict
from typing_extensions import NotRequired, ReadOnly
@ -36,6 +36,30 @@ class BitBucketFileMetadata(TypedDict):
last_modified: ReadOnly[str | None]
class _SrcListingResponse(Protocol):
def json(self) -> BitBucketSrcListing: ...
class _RepositoryInfoResponse(Protocol):
def json(self) -> Mapping[str, object]: ...
class _BranchListingResponse(Protocol):
def json(self) -> BitBucketBranchListing: ...
def _src_listing(response: _SrcListingResponse) -> BitBucketSrcListing:
return response.json()
def _repository_info(response: _RepositoryInfoResponse) -> Mapping[str, object]:
return response.json()
def _branch_listing(response: _BranchListingResponse) -> BitBucketBranchListing:
return response.json()
def _sanitize_file_path(file_path: str) -> str:
"""Reject path traversal and URL-encode each path segment."""
if "#" in file_path or "?" in file_path:
@ -166,7 +190,7 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final[BitBucketSrcListing] = response.json()
data: Final = _src_listing(response)
return [
file_path
for item in data.get("values", [])
@ -201,7 +225,7 @@ class BitBucketClient:
try:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
return response.json()
return _repository_info(response)
except Exception as e:
raise Exception(f"Failed to get repository info: {e}")
@ -231,7 +255,7 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final[BitBucketBranchListing] = response.json()
data: Final = _branch_listing(response)
return data.get("values", [])
except Exception as e:
raise Exception(f"Failed to get branches: {e}")

View file

@ -4,12 +4,13 @@ Fetches .prompt files from BitBucket repositories and provides team-based access
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol
from jinja2 import DictLoader, select_autoescape
from jinja2.sandbox import ImmutableSandboxedEnvironment
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.litellm_core_utils.core_helpers import require_attributes
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -26,6 +27,10 @@ from litellm.types.utils import StandardCallbackDynamicParams
from .bitbucket_client import BitBucketClient
class _SupportsGetWithDefault(Protocol):
def get(self, key: str, default: object = None, /) -> object: ...
class BitBucketPromptTemplate:
"""
Represents a prompt template loaded from BitBucket.
@ -35,7 +40,7 @@ class BitBucketPromptTemplate:
self,
template_id: str,
content: str,
metadata: dict[str, Any],
metadata: Mapping[str, object],
model: str | None = None,
):
self.template_id = template_id
@ -44,7 +49,9 @@ class BitBucketPromptTemplate:
self.model = model or metadata.get("model")
self.temperature = metadata.get("temperature")
self.max_tokens = metadata.get("max_tokens")
self.input_schema = metadata.get("input", {}).get("schema", {})
self.input_schema = require_attributes(metadata.get("input", {}), _SupportsGetWithDefault, "get").get(
"schema", {}
)
self.optional_params = {k: v for k, v in metadata.items() if k not in ["model", "input", "content"]}
def __repr__(self):

View file

@ -7,10 +7,10 @@ so `response_id` is always the raw provider response id (cache-hit suffix stripp
import json
import re
from collections.abc import Iterator, Mapping
from collections.abc import Iterable, Iterator, Mapping
from math import isfinite
from types import MappingProxyType
from typing import Any, Final
from typing import Any, Final, Protocol
from pydantic import JsonValue, TypeAdapter, ValidationError
@ -19,6 +19,7 @@ from litellm._logging import verbose_logger
from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger
from litellm.integrations.clickhouse.context import is_lens_analysis
from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
@ -40,6 +41,14 @@ _INTERNAL_METADATA_KEYS: Final = frozenset(
)
class _SupportsGet(Protocol):
def get(self, key: str, /) -> object: ...
class _SupportsItems(Protocol):
def items(self) -> Iterable[tuple[object, object]]: ...
def _request_metadata_fields(value: object) -> Iterator[tuple[str, JsonValue]]:
if value is None:
return
@ -85,22 +94,22 @@ def _json_mapping(value: Mapping[str, object]) -> str:
return _json(dict(value))
def _find_traceparent(metadata: Mapping[str, Any], kwargs: Mapping[str, Any]) -> tuple[str, str]:
def _find_traceparent(metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> tuple[str, str]:
custom_headers = metadata.get("requester_custom_headers") or MappingProxyType({})
proxy_request = (kwargs.get("litellm_params") or MappingProxyType({})).get(
proxy_request = require_attributes(kwargs.get("litellm_params") or MappingProxyType({}), _SupportsGet, "get").get(
"proxy_server_request"
) or MappingProxyType({})
request_headers = proxy_request.get("headers") or MappingProxyType({})
request_headers = require_attributes(proxy_request, _SupportsGet, "get").get("headers") or MappingProxyType({})
for headers in (custom_headers, request_headers):
for name, value in headers.items():
for name, value in require_attributes(headers, _SupportsItems, "items").items():
if str(name).lower() == "traceparent":
return parse_traceparent(value)
return "", ""
def _cache_tokens(usage: Mapping[str, Any]) -> tuple[int, int]:
def _cache_tokens(usage: Mapping[str, object]) -> tuple[int, int]:
"""(cache_read, cache_write) from a Usage dict: OpenAI prompt_tokens_details first, Anthropic fields as fallback."""
details = usage.get("prompt_tokens_details") or MappingProxyType({})
details = require_attributes(usage.get("prompt_tokens_details") or MappingProxyType({}), _SupportsGet, "get")
cache_read = _int(details.get("cached_tokens")) or _int(usage.get("cache_read_input_tokens"))
cache_write = (
_int(details.get("cache_write_tokens"))
@ -116,10 +125,17 @@ def _request_tags(value: object) -> list[str]:
return [str(tag) for tag in value]
def _session_id(payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, Any]) -> str:
def _session_id(payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, object]) -> str:
"""Mirrors proxy `_get_session_id_for_spend_log`: explicit session id, else the payload trace id."""
request_metadata = (kwargs.get("litellm_params") or MappingProxyType({})).get("metadata") or MappingProxyType({})
return str(payload.get("session_id") or request_metadata.get("session_id") or payload.get("trace_id") or "")
request_metadata = require_attributes(
kwargs.get("litellm_params") or MappingProxyType({}), _SupportsGet, "get"
).get("metadata") or MappingProxyType({})
return str(
payload.get("session_id")
or require_attributes(request_metadata, _SupportsGet, "get").get("session_id")
or payload.get("trace_id")
or ""
)
def _is_trace_ingest(payload: StandardLoggingPayload) -> bool:
@ -128,7 +144,7 @@ def _is_trace_ingest(payload: StandardLoggingPayload) -> bool:
def spend_log_row_from_payload(
payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, Any]
payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, object]
) -> SpendLogRecord:
metadata: Mapping[str, Any] = payload.get("metadata") or MappingProxyType({})
hidden_params: Mapping[str, Any] = payload.get("hidden_params") or MappingProxyType({})

View file

@ -1,16 +1,17 @@
import contextvars
import copy
import hashlib
import json
import os
import secrets
from collections.abc import Iterable, Mapping, Sequence
from datetime import datetime
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast, get_args
import httpx
from pydantic import TypeAdapter
from pydantic import BaseModel, TypeAdapter
from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger
@ -20,6 +21,7 @@ from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
redact_nested_match_and_regex_keys,
require_attributes,
)
from litellm.secret_managers.main import str_to_bool
from litellm.types.guardrails import (
@ -73,6 +75,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions"
_PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16)
_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16)
_POST_CALL_CHECKED_TOKEN: Final = secrets.token_hex(16)
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
@ -167,18 +170,22 @@ def _strict_guardrail_modes_enabled() -> bool:
return True if parsed is None else parsed
def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None:
class _SupportsGet(Protocol):
def get(self, key: str, /) -> object: ...
def get_session_id_from_request_data(request_data: Mapping[str, object]) -> str | None:
"""Extract session_id from request data (litellm_session_id or metadata)."""
session_id = request_data.get("litellm_session_id")
if session_id:
return str(session_id)
metadata: Final = request_data.get("metadata") or {}
metadata: Final = require_attributes(request_data.get("metadata") or {}, _SupportsGet, "get")
session_id = metadata.get("session_id")
if session_id:
return str(session_id)
litellm_metadata: Final = request_data.get("litellm_metadata") or {}
litellm_metadata: Final = require_attributes(request_data.get("litellm_metadata") or {}, _SupportsGet, "get")
session_id = litellm_metadata.get("session_id")
if session_id:
return str(session_id)
@ -281,6 +288,95 @@ def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> b
return False
def _record_guardrail_marker(
data: dict[str, object], # mutable-ok: the request payload the marker is recorded on
marker: str | None,
) -> None:
if marker is None:
return
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key)
if isinstance(meta, dict):
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
if isinstance(executed, list):
if marker not in executed:
executed.append(marker)
else:
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
return
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
_RESPONSE_METADATA_KEYS: Final = frozenset(
{
"id",
"created",
"created_at",
"model",
"usage",
"usage_info",
"system_fingerprint",
"_hidden_params",
"_response_ms",
}
)
def _response_fingerprint(response: object) -> str | None:
"""Digest of everything but metadata set between layers, so an unknown field counts as content."""
try:
fields: Final[Mapping[str, object]] = (
response.model_dump(mode="json")
if isinstance(response, BaseModel)
else _REQUEST_FIELDS.validate_python(response)
)
content: Final = {key: value for key, value in fields.items() if key not in _RESPONSE_METADATA_KEYS}
if not content:
return None
return hashlib.sha256(
json.dumps(content, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()
).hexdigest()
except (TypeError, ValueError):
return None
_RESPONSE_GUARDRAIL_CHECKS_KEY: Final = "guardrail_checks"
def _str_keyed_dict(value: object) -> Mapping[str, object] | None:
if not isinstance(value, dict):
return None
return cast("Mapping[str, object]", value) # cast-ok: response bodies and _hidden_params are keyed by field name
def _response_hidden_params(response: object) -> Mapping[str, object] | None:
as_dict: Final = _str_keyed_dict(response)
return _str_keyed_dict(
as_dict.get("_hidden_params") if as_dict is not None else getattr(response, "_hidden_params", None)
)
def _write_response_marker(response: object, marker: str | None) -> None:
"""Record a check on the response itself, so it survives the router copying request metadata per hop."""
if marker is None:
return
hidden: Final = _response_hidden_params(response)
if not isinstance(hidden, dict):
return
hidden[_RESPONSE_GUARDRAIL_CHECKS_KEY] = tuple(
_str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY)) | {marker}
)
def _response_carries_marker(response: object, marker: str | None) -> bool:
hidden: Final = _response_hidden_params(response)
return (
marker is not None
and hidden is not None
and marker in _str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY))
)
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
@ -291,6 +387,7 @@ def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuard
class CustomGuardrail(CustomLogger):
mcp_block_behavior: Literal["tool_result", "raise"] | None = None
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
use_native_during_call_hook: ClassVar[bool] = False
@ -912,20 +1009,7 @@ class CustomGuardrail(CustomLogger):
With ``raised``, records instead that the hook raised as a policy pipeline
step the pipeline then moved past. The deployment-level hook still runs it.
"""
marker: Final = self._pre_call_marker(raised=raised)
if marker is None:
return
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key)
if isinstance(meta, dict):
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
if isinstance(executed, list):
if marker not in executed:
executed.append(marker)
else:
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
return
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
_record_guardrail_marker(data, self._pre_call_marker(raised=raised))
def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool:
"""Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided."""
@ -1123,10 +1207,26 @@ class CustomGuardrail(CustomLogger):
cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed
)
if not self._is_valid_response_type(result):
return None
returned: Final = result if self._is_valid_response_type(result) else None
if target is self:
return returned
checked: Final = returned if returned is not None else response
marker: Final = self._post_call_check_marker(checked)
_record_guardrail_marker(
cast("dict[str, object]", request_data), # cast-ok: proxy request payloads are str-keyed
marker,
)
_write_response_marker(checked, marker)
return returned
return result
def _post_call_check_marker(self, response: object) -> str | None:
fingerprint: Final = _response_fingerprint(response)
return self._marker(f"{_POST_CALL_CHECKED_TOKEN}:{fingerprint}") if fingerprint is not None else None
def post_call_already_checked(self, data: Mapping[str, object], response: object) -> bool:
"""Whether a deployment-level hook of this request already ran this guardrail on this exact response."""
marker: Final = self._post_call_check_marker(response)
return _request_carries_marker(data, marker) or _response_carries_marker(response, marker)
async def async_logging_hook(
self,

View file

@ -44,6 +44,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -315,7 +316,10 @@ def _to_dd_tool_calls(message: Mapping[str, object]) -> tuple[ToolCall, ...]:
tool_id=_text_field(tool_call, "id"),
type=_text_field(tool_call, "type", "function"),
)
for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ())
for tool_call in (
*anthropic_server_tool_call_params(message),
*(raw_tool_calls if isinstance(raw_tool_calls, list) else ()),
)
if isinstance(tool_call, dict)
for function in [_mapping_field(tool_call, "function")]
)

View file

@ -3,17 +3,29 @@ Functions for sending Email Alerts
"""
import os
from typing import Final
from collections.abc import Awaitable, Mapping, Sequence
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.repositories.team_repository import TeamRepository
from litellm.types.integrations.slack_alerting import WebhookEvent
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
# we use this for the email header, please send a test email if you change this. verify it looks good on email
LITELLM_LOGO_URL: Final = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
LITELLM_SUPPORT_CONTACT: Final = "support@berri.ai"
def _query_raw_rows(
client: "PrismaClient", query: str, *args: object
) -> Awaitable[Sequence[Mapping[str, object]] | None]:
from litellm.proxy.utils import query_raw_rows
return query_raw_rows(client.db, query, *args)
async def get_all_team_member_emails(team_id: str | None = None) -> list:
verbose_logger.debug("Email Alerting: Getting all team members for team_id=%s", team_id)
if team_id is None:
@ -51,17 +63,18 @@ async def get_all_team_member_emails(team_id: str | None = None) -> list:
WHERE user_id = ANY($1::TEXT[]);
"""
_result: Final = await prisma_client.db.query_raw(sql_query, _team_member_user_ids)
_result: Final = await _query_raw_rows(prisma_client, sql_query, _team_member_user_ids)
verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", _result)
if _result is None:
return []
emails: Final = []
for user in _result:
if user and isinstance(user, dict) and user.get("user_email", None) is not None:
emails.append(user.get("user_email"))
emails: Final = [
user.get("user_email")
for user in _result
if user and isinstance(user, dict) and user.get("user_email", None) is not None
]
return emails

View file

@ -123,7 +123,9 @@ class FocusLiteLLMDatabase:
ORDER BY ordinal_position;
"""
try:
columns_response: Final = await client.db.query_raw(info_query)
from litellm.proxy.utils import query_raw_rows
columns_response: Final = await query_raw_rows(client.db, info_query)
return {"columns": columns_response, "table_name": "LiteLLM_DailyUserSpend"}
except Exception as exc:
raise RuntimeError(f"Error getting table info: {exc}") from exc

View file

@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import require_attributes, subscript
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
get_content_from_model_response,
@ -60,6 +61,10 @@ class _JsonResponse(Protocol):
def json(self) -> object: ...
class _SupportsGet(Protocol):
def get(self, key: str, /) -> object: ...
def _login_access_token(response: _GalileoLoginResponse) -> str:
"""Read the bearer token out of a Galileo login response body."""
return response.json()["access_token"]
@ -463,13 +468,13 @@ class GalileoObserve(CustomLogger):
pass
@staticmethod
def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, object]:
optional_params: Final[Mapping[str, object]] = kwargs.get("optional_params", {}) or {}
def _build_prompt(kwargs: Mapping[str, object]) -> dict[str, object]:
optional_params: Final = require_attributes(kwargs.get("optional_params", {}) or {}, _SupportsGet, "get")
prompt: Final[dict[str, object]] = {"messages": kwargs.get("messages")}
if optional_params.get("functions") is not None:
prompt["functions"] = optional_params["functions"]
prompt["functions"] = subscript(optional_params, "functions")
if optional_params.get("tools") is not None:
prompt["tools"] = optional_params["tools"]
prompt["tools"] = subscript(optional_params, "tools")
return prompt
@staticmethod

View file

@ -5,7 +5,7 @@ import os
import time
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm._logging import verbose_logger
from litellm._uuid import uuid
@ -16,7 +16,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
from litellm.litellm_core_utils.cloud_storage_security import (
sanitize_cloud_object_component,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
from litellm.types.integrations.gcs_bucket import *
@ -28,6 +28,10 @@ else:
VertexBase = Any
class _SupportsGetWithDefault(Protocol):
def get(self, key: str, default: str, /) -> object: ...
class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
def __init__(self, bucket_name: str | None = None) -> None:
from litellm.proxy.proxy_server import premium_user
@ -298,7 +302,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
await self._send_individual_logs(items_to_process)
def _get_object_name(
self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: Any
self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: object
) -> str:
"""
Get the object name to use for the current payload
@ -313,7 +317,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
return f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}"
if logging_payload.get("error_str", None) is not None:
return self._generate_failure_object_name(request_date_str=current_date)
raw_response_id: Final[object] = response_obj.get("id", "")
raw_response_id: Final[object] = require_attributes(response_obj, _SupportsGetWithDefault, "get").get("id", "")
response_id: Final = raw_response_id if isinstance(raw_response_id, str) else ""
return self._generate_success_object_name(
request_date_str=current_date,
@ -360,7 +364,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
def _generate_success_object_name(
self,
request_date_str: str,
response_id: str,
response_id: object,
) -> str:
return f"{request_date_str}/{response_id}"

View file

@ -132,7 +132,7 @@ class GenericPromptManager(CustomPromptManagement):
async def async_fetch_prompt_from_api(
self, prompt_id: str | None, prompt_spec: PromptSpec | None
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Fetch a prompt from the API asynchronously.
"""

View file

@ -10,6 +10,23 @@ from litellm.integrations.helicone_mock_client import (
create_mock_helicone_client,
should_use_helicone_mock,
)
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
from litellm.types.utils import ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall
def _tool_use_block(
tool_call: ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall,
) -> dict[str, object]:
match tool_call:
case ChatCompletionMessageCustomToolCall(custom=custom):
return {"type": "tool_use", "id": tool_call.id, "name": custom.name, "input": custom.input}
case _:
return {
"type": "tool_use",
"id": tool_call.id,
"name": tool_call.function.name,
"input": tool_call.function.arguments,
}
class HeliconeLogger:
@ -57,23 +74,8 @@ class HeliconeLogger:
message = choice["message"]
content = []
if "tool_calls" in message and message["tool_calls"]:
for tool_call in message["tool_calls"]:
function = tool_call.get("function")
custom = tool_call.get("custom")
if not function and not custom:
continue
name, tool_input = (
(function["name"], function["arguments"]) if function else (custom["name"], custom["input"])
)
content.append(
{
"type": "tool_use",
"id": tool_call["id"],
"name": name,
"input": tool_input,
}
)
if tool_calls := tool_calls_with_server_tool_calls(message):
content = [_tool_use_block(tool_call) for tool_call in tool_calls]
elif "content" in message and message["content"]:
content = [{"type": "text", "text": message["content"]}]

View file

@ -7,10 +7,12 @@ https://humanloop.com/
from typing import Final, cast
import httpx
from pydantic import TypeAdapter
from typing_extensions import TypedDict
import litellm
from litellm.caching import DualCache
from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, decoded_json_object, subscript
from litellm.llms.custom_httpx.http_handler import get_httpx_client
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
@ -19,6 +21,8 @@ from litellm.types.utils import StandardCallbackDynamicParams
from .custom_logger import CustomLogger
_DECODED_JSON: Final = TypeAdapter(object)
class PromptManagementClient(TypedDict):
prompt_id: str
@ -79,18 +83,20 @@ class HumanLoopPromptManager(DualCache):
raise Exception(f"Error getting prompt from Humanloop: {e.response.text}")
json_response: Final = response.json()
template_message: Final = json_response["template"]
if isinstance(template_message, dict):
template_messages = [template_message]
elif isinstance(template_message, list):
template_messages = template_message
decoded_response: Final = _DECODED_JSON.validate_python(json_response)
template_message: Final = subscript(decoded_response, "template")
template_mapping: Final = as_str_mapping(template_message)
template_list: Final = as_object_list(template_message)
if template_mapping is not None:
template_messages = [template_mapping]
elif template_list is not None:
template_messages = template_list
else:
raise ValueError(f"Invalid template message type: {type(template_message)}")
template_model: Final = json_response["model"]
optional_params: Final = {}
for k, v in json_response.items():
if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS:
optional_params[k] = v
optional_params: Final = {
k: v for k, v in decoded_json_object(decoded_response).items() if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS
}
return PromptManagementClient(
prompt_id=humanloop_prompt_id,
prompt_template=cast(list[AllMessageValues], template_messages),

View file

@ -3,7 +3,7 @@
import json
import os
from typing import Final, Literal
from typing import Final, Literal, Protocol
import httpx
@ -11,6 +11,7 @@ import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.llms.custom_httpx.http_handler import (
HTTPHandler,
get_async_httpx_client,
@ -18,6 +19,10 @@ from litellm.llms.custom_httpx.http_handler import (
)
class _SupportsGet(Protocol):
def get(self, key: str, default: object = None, /) -> object: ...
def get_utc_datetime():
import datetime as dt
from datetime import datetime
@ -70,10 +75,11 @@ class LagoLogger(CustomLogger):
if (
isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse)
) and hasattr(response_obj, "usage"):
response_usage: Final = require_attributes(response_obj["usage"], _SupportsGet, "get")
usage = {
"prompt_tokens": response_obj["usage"].get("prompt_tokens", 0),
"completion_tokens": response_obj["usage"].get("completion_tokens", 0),
"total_tokens": response_obj["usage"].get("total_tokens"),
"prompt_tokens": response_usage.get("prompt_tokens", 0),
"completion_tokens": response_usage.get("completion_tokens", 0),
"total_tokens": response_usage.get("total_tokens"),
}
litellm_params: Final = kwargs.get("litellm_params", {}) or {}

View file

@ -7,6 +7,7 @@ from itertools import chain
from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, runtime_checkable
from opentelemetry import trace
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_logger
@ -26,12 +27,15 @@ from litellm.integrations.opentelemetry import (
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
from litellm.secret_managers.main import str_to_bool
from litellm.types.integrations.langfuse_otel import (
LangfuseSpanAttributes,
)
from litellm.types.utils import StandardCallbackDynamicParams
_DECODED_JSON: Final = TypeAdapter(object)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -605,8 +609,8 @@ def _extract_choices_output(response_obj: _Gettable) -> str | None:
message: Final = _as_gettable(first_choice.get("message", {}))
if message is None:
return None
tool_calls: Final = message.get("tool_calls")
if isinstance(tool_calls, list) and tool_calls:
tool_calls: Final = tool_calls_with_server_tool_calls(message)
if tool_calls:
transformed_tool_calls: Final = [
entry for tool_call in tool_calls if (entry := _transformed_tool_call(response_obj, tool_call)) is not None
]
@ -642,7 +646,7 @@ def _tool_call_arguments(arguments: object) -> object:
if not isinstance(arguments, str):
return arguments
try:
return json.loads(arguments)
return _DECODED_JSON.validate_python(json.loads(arguments))
except json.JSONDecodeError:
return {}

View file

@ -7,7 +7,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764
import json
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import BaseModel
from typing_extensions import override
@ -16,6 +16,7 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor
BaseLLMObsOTELAttributes,
safe_set_attribute,
)
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse
from litellm.types.utils import (
EmbeddingResponse,
@ -30,6 +31,10 @@ if TYPE_CHECKING:
from opentelemetry.trace import Span
class _SupportsGet(Protocol):
def get(self, key: str, /) -> object: ...
def serialize_langfuse_datetime(value: datetime) -> str:
return json.dumps(value.astimezone().isoformat())
@ -87,11 +92,11 @@ def get_output_content_by_type(
class LangfuseLLMObsOTELAttributes(BaseLLMObsOTELAttributes):
@staticmethod
@override
def set_messages(span: "Span", kwargs: dict[str, Any]):
def set_messages(span: "Span", kwargs: Mapping[str, object]):
prompt: Final = {"messages": kwargs.get("messages")}
optional_params: Final = kwargs.get("optional_params", {})
functions: Final = optional_params.get("functions")
tools: Final = optional_params.get("tools")
functions: Final = require_attributes(optional_params, _SupportsGet, "get").get("functions")
tools: Final = require_attributes(optional_params, _SupportsGet, "get").get("tools")
if functions is not None:
prompt["functions"] = functions
if tools is not None:

View file

@ -6,6 +6,7 @@ from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final
import packaging
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
if TYPE_CHECKING:
from openai.types.chat.chat_completion_message_custom_tool_call import Custom
@ -67,8 +68,8 @@ def parse_messages(input):
}
# Only add tool_calls and function_call to res if they are set
if message.get("tool_calls"):
serialized["tool_calls"] = parse_tool_calls(message.get("tool_calls"))
if tool_calls := tool_calls_with_server_tool_calls(message):
serialized["tool_calls"] = parse_tool_calls(tool_calls)
return serialized

View file

@ -1,9 +1,11 @@
import json
import threading
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import attribute_or_none, iter_items, require_attributes
if TYPE_CHECKING:
from litellm.types.utils import StandardLoggingPayload
@ -11,6 +13,25 @@ else:
StandardLoggingPayload = Any
class _ChoiceWithMessage(Protocol):
@property
def message(self) -> object: ...
class _ExcludeNoneDumpable(Protocol):
def model_dump(self, *, exclude_none: bool) -> object: ...
def _dumped_choice_message(choice: object) -> object:
message: Final = require_attributes(choice, _ChoiceWithMessage, "message").message
return require_attributes(message, _ExcludeNoneDumpable, "model_dump").model_dump(exclude_none=True)
def _output_messages(response_obj: object) -> Sequence[object]:
choices: Final = attribute_or_none(response_obj, "choices") if hasattr(response_obj, "choices") else ()
return [_dumped_choice_message(c) for c in iter_items(choices)]
class MlflowLogger(CustomLogger):
def __init__(self):
from mlflow.tracking import MlflowClient
@ -62,7 +83,7 @@ class MlflowLogger(CustomLogger):
inputs: Final = self._construct_input(kwargs)
input_messages: Final = inputs.get("messages", [])
output_messages: Final = [c.message.model_dump(exclude_none=True) for c in getattr(response_obj, "choices", [])]
output_messages: Final = _output_messages(response_obj)
if messages := [*input_messages, *output_messages]:
set_span_chat_messages(span, messages)
if tools := inputs.get("tools"):

View file

@ -56,6 +56,7 @@ from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import has_attributes
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
from litellm.types.integrations.newrelic import NewRelicInitParams
from litellm.types.utils import Message, ModelResponse, StandardLoggingPayload
@ -479,11 +480,11 @@ class NewRelicLogger(CustomLogger):
content: Final = message.get("content")
# Handle tool calls
if message.get("tool_calls"):
if tool_calls := tool_calls_with_server_tool_calls(message):
try:
return json.dumps(message["tool_calls"])
return json.dumps([tool_call.model_dump() for tool_call in tool_calls])
except Exception:
return str(message["tool_calls"])
return str(tool_calls)
# Handle None or missing content
if content is None:

View file

@ -4,13 +4,13 @@
import json
import os
from collections.abc import Mapping
from typing import Final
from typing import Final, Protocol
import httpx
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
from litellm.llms.custom_httpx.http_handler import (
HTTPHandler,
@ -19,6 +19,10 @@ from litellm.llms.custom_httpx.http_handler import (
)
class _SupportsGet(Protocol):
def get(self, key: str, default: object = None, /) -> object: ...
def get_utc_datetime():
import datetime as dt
from datetime import datetime
@ -60,10 +64,11 @@ class OpenMeterLogger(CustomLogger):
if (
isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse)
) and hasattr(response_obj, "usage"):
response_usage: Final = require_attributes(response_obj["usage"], _SupportsGet, "get")
usage = {
"prompt_tokens": response_obj["usage"].get("prompt_tokens", 0),
"completion_tokens": response_obj["usage"].get("completion_tokens", 0),
"total_tokens": response_obj["usage"].get("total_tokens"),
"prompt_tokens": response_usage.get("prompt_tokens", 0),
"completion_tokens": response_usage.get("completion_tokens", 0),
"total_tokens": response_usage.get("total_tokens"),
}
# OPENMETER_TRUST_REQUEST_USER (default "true"): when set to "false",

View file

@ -6,7 +6,7 @@ from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
from pydantic import ConfigDict, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
@ -35,6 +35,7 @@ from litellm.integrations.otel.model.metadata import flatten_metadata
from litellm.integrations.otel.model.semconv import LiteLLM, Metric
from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls
from litellm.integrations.otel.routing import routing_decision_attributes
from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.secret_redaction import redact_string
@ -42,11 +43,11 @@ from litellm.litellm_core_utils.service_tier_utils import (
get_requested_service_tier,
get_served_service_tier,
)
from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params, tool_calls_with_server_tool_calls
from litellm.secret_managers.main import get_secret_bool, str_to_bool
from litellm.types.integrations.otel_span_attributes import SpanAttributes
from litellm.types.services import ServiceLoggerPayload
from litellm.types.utils import (
ChatCompletionMessageToolCall,
CostBreakdown,
Function,
LLMResponseTypes,
@ -129,6 +130,22 @@ def _text_as_stored_on_span(span: Span, key: str, text: str) -> str:
_MAX_DYNAMIC_TRACER_PROVIDERS: Final = 256
class _SupportsGet(Protocol):
def get(self, key: str, /) -> object: ...
def _response_field(response_obj: object, key: str) -> object:
if not response_obj:
return None
return require_attributes(response_obj, _SupportsGet, "get").get(key)
def _billed_usage(response_obj: object, call_type: str | None, litellm_params: Mapping[str, object] | None) -> object:
if not response_obj or is_unbilled_non_inference_call_from_params(call_type, litellm_params, response_obj):
return None
return _response_field(response_obj, "usage")
class _CacheHitView(TypedDict, total=False):
cache_hit: ReadOnly[object]
saved_cache_cost: ReadOnly[object]
@ -2406,17 +2423,17 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
@staticmethod
def _tool_calls_kv_pair(
tool_calls: list[ChatCompletionMessageToolCall],
tool_calls: Iterable[object],
) -> dict[str, object]:
kv_pairs: Final[dict[str, object]] = {}
for idx, tool_call in enumerate(tool_calls):
_function = tool_call.get("function")
_function: object = require_attributes(tool_call, _SupportsGet, "get").get("function")
if not _function:
continue
keys = Function.__annotations__.keys()
for key in keys:
_value = _function.get(key)
_value = require_attributes(_function, _SupportsGet, "get").get(key)
if _value:
kv_pairs[f"{SpanAttributes.LLM_COMPLETIONS.value}.{idx}.function_call.{key}"] = _value
@ -2583,7 +2600,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
# but Embeddings and Image-gen responses do not. Fall back to
# the litellm call ID so every call type can be correlated
# across LiteLLM UI, Phoenix traces, and provider logs (Issue #8).
response_id = (response_obj.get("id") if response_obj else None) or standard_logging_payload.get("id")
response_id = _response_field(response_obj, "id") or standard_logging_payload.get("id")
if response_id:
self.safe_set_attribute(
span=span,
@ -2600,42 +2617,37 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
)
# The model used to generate the response.
if response_obj and response_obj.get("model"):
response_model: Final = _response_field(response_obj, "model")
if response_model:
self.safe_set_attribute(
span=span,
key=SpanAttributes.LLM_RESPONSE_MODEL.value,
value=response_obj.get("model"),
value=response_model,
)
self._set_service_tier_attributes(span=span, standard_logging_payload=standard_logging_payload)
usage: Final = (
response_obj.get("usage")
if response_obj
and not is_unbilled_non_inference_call_from_params(
kwargs.get("call_type"), litellm_params, response_obj
)
else None
)
usage: Final = _billed_usage(response_obj, kwargs.get("call_type"), litellm_params)
if usage:
usage_fields: Final = require_attributes(usage, _SupportsGet, "get")
self.safe_set_attribute(
span=span,
key=SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS.value,
value=usage.get("total_tokens"),
value=usage_fields.get("total_tokens"),
)
# The number of tokens used in the LLM response (completion).
self.safe_set_attribute(
span=span,
key=SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS.value,
value=usage.get("completion_tokens"),
value=usage_fields.get("completion_tokens"),
)
# The number of tokens used in the LLM prompt.
self.safe_set_attribute(
span=span,
key=SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS.value,
value=usage.get("prompt_tokens"),
value=usage_fields.get("prompt_tokens"),
)
########################################################################
@ -2711,7 +2723,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
########## LLM Response Attributes ##########
#############################################
if response_obj is not None:
if response_obj.get("choices"):
response_view: Final = require_attributes(response_obj, _SupportsGet, "get")
if response_view.get("choices"):
transformed_choices: Final = self._transform_choices_to_otel_semantic_conventions(
response_obj.get("choices")
)
@ -2721,10 +2734,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
value=safe_dumps(transformed_choices),
)
finish_reasons: Final = []
for idx, choice in enumerate(response_obj.get("choices")):
if choice.get("finish_reason"):
finish_reasons.append(choice.get("finish_reason"))
choice_views: Final = [
require_attributes(choice, _SupportsGet, "get")
for choice in iter_items(response_view.get("choices"))
]
finish_reasons: Final = [
choice.get("finish_reason") for choice in choice_views if choice.get("finish_reason")
]
if finish_reasons:
self.safe_set_attribute(
@ -2733,10 +2749,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
value=safe_dumps(finish_reasons),
)
for idx, choice in enumerate(response_obj.get("choices")):
for choice in choice_views:
if choice.get("finish_reason"):
message = choice.get("message")
tool_calls = message.get("tool_calls")
tool_calls = tool_calls_with_server_tool_calls(choice.get("message"))
if tool_calls:
kv_pairs = OpenTelemetry._tool_calls_kv_pair(tool_calls)
for key, value in kv_pairs.items():
@ -2746,13 +2761,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
value=value,
)
elif response_obj.get("output"):
elif response_view.get("output"):
# Responses API: ResponsesAPIResponse has an "output"
# list instead of "choices". Each item with
# type="message" contains a "content" list of
# OutputText objects (type="output_text").
output_items: Final = response_obj.get("output")
output_messages: Final = self._transform_responses_api_output_to_otel(output_items)
output_items: Final = response_view.get("output")
output_messages: Final = self._transform_responses_api_output_to_otel(iter_items(output_items))
if output_messages:
self.safe_set_attribute(
span=span,
@ -2766,7 +2781,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
# ChatCompletionMessageToolCall format expected by
# _tool_calls_kv_pair.
tool_calls = []
for out_item in output_items:
for out_item in iter_items(output_items):
item_d = self._to_dict(out_item)
if item_d and item_d.get("type") == "function_call":
tool_calls.append(
@ -2787,7 +2802,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
)
# Extract finish reason from ResponsesAPIResponse.status
status: Final = response_obj.get("status")
status: Final = response_view.get("status")
if status:
self.safe_set_attribute(
span=span,
@ -2856,7 +2871,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
transformed_msg = {"role": role, "parts": parts}
if "id" in msg:
transformed_msg["id"] = msg["id"]
if "tool_calls" in msg:
if server_tool_calls := anthropic_server_tool_call_params(cast(object, msg)):
transformed_msg["tool_calls"] = [*server_tool_calls, *(msg.get("tool_calls") or [])]
elif "tool_calls" in msg:
transformed_msg["tool_calls"] = msg["tool_calls"]
if "tool_call_id" in msg:
transformed_msg["tool_call_id"] = msg["tool_call_id"]
@ -2905,7 +2922,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
return obj.model_dump()
return None
def _transform_responses_api_output_to_otel(self, output: list) -> list[dict]:
def _transform_responses_api_output_to_otel(self, output: Iterable[object]) -> list[dict]:
"""
Transform Responses API output items into OTEL GenAI 1.38 format.

View file

@ -1,4 +1,5 @@
from abc import ABC
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
if TYPE_CHECKING:
@ -7,7 +8,7 @@ if TYPE_CHECKING:
class BaseLLMObsOTELAttributes(ABC):
@staticmethod
def set_messages(span: "Span", kwargs: dict[str, Any]):
def set_messages(span: "Span", kwargs: Mapping[str, object]):
pass
@staticmethod

View file

@ -1034,10 +1034,12 @@ def emit_guardrail_span(entry: "StandardLoggingGuardrailInformation") -> None:
pass
def seed_request_identity(user_api_key_dict: object, model: str | None = None) -> None:
def seed_request_identity(user_api_key_dict: object, model: str | None = None, request_data: object = None) -> None:
logger: Final = _registered_v2_logger()
if logger is not None:
logger.seed_request_identity(user_api_key_dict, model=model)
logger.seed_request_identity(
user_api_key_dict, model=model, request_metadata=metadata_from_request_data(request_data)
)
def fan_out_provider() -> ApiTracerProvider:

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