scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

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

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-21 13:25:45 -07:00
commit ddf6565970
No known key found for this signature in database
145 changed files with 11805 additions and 1042 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -356,7 +356,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
| [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | |
| [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | |
| [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | |
| [Qwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
| [Qianwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
| [QwenCloud (`qwencloud`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
| [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | |
| [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | |

1
backend/routes/allowlist.py
View file

@ -51,6 +51,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/cache_settings",
"/coordination_redis/",
"/cost_tracking",
"/cost_optimization/",
"/cost/",
"/credentials",
"/credential",

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "is_default" BOOLEAN NOT NULL DEFAULT false;

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

@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
is_default Boolean @default(false) // Applied only when no non-default attachment matches
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt

128
litellm-rust/Cargo.lock generated
View file

@ -2460,8 +2460,8 @@ dependencies = [
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
]
[[package]]
@ -2479,12 +2479,29 @@ name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"r2d2",
"redis",
"redis-test",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-cache-response"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-redis",
"py_literal",
"redis",
"redis-test",
"serde",
"serde_json",
"sha2 0.10.9",
"tokio",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -2648,6 +2665,10 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -2659,6 +2680,7 @@ dependencies = [
"pyo3",
"pyo3-async-runtimes",
"rstest",
"serde",
"serde_json",
"tokio",
"tokio-tungstenite",
@ -2947,6 +2969,16 @@ dependencies = [
"minimal-lexical",
]
[[package]]
name = "num-bigint"
version = "0.4.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367"
dependencies = [
"num-integer",
"num-traits",
]
[[package]]
name = "num-bigint"
version = "0.5.1"
@ -2957,6 +2989,15 @@ dependencies = [
"num-traits",
]
[[package]]
name = "num-complex"
version = "0.4.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
dependencies = [
"num-traits",
]
[[package]]
name = "num-conv"
version = "0.2.2"
@ -3130,6 +3171,48 @@ version = "2.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
[[package]]
name = "pest"
version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6d45aeb61b4bf818e12d4205f2466f8c4748f85f4fce0146d1c03d69d753f0ad"
dependencies = [
"memchr",
"ucd-trie",
]
[[package]]
name = "pest_derive"
version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "89cc5a242e25ed4e7704d0be240f2cfbe20a8c27e7e252d94835be93d92dc39f"
dependencies = [
"pest",
"pest_generator",
]
[[package]]
name = "pest_generator"
version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7abf21475cc3820fe4b2ca2dc2142902f67a02189f3b5b3a229f4febc01a43e5"
dependencies = [
"pest",
"pest_meta",
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "pest_meta"
version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adba4db388f687393c18c51348d44a41d870ca9df71a2c98172ea3035dc6936e"
dependencies = [
"pest",
]
[[package]]
name = "pin-project"
version = "1.1.13"
@ -3304,6 +3387,19 @@ dependencies = [
"prost",
]
[[package]]
name = "py_literal"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1"
dependencies = [
"num-bigint 0.4.8",
"num-complex",
"num-traits",
"pest",
"pest_derive",
]
[[package]]
name = "pyo3"
version = "0.29.2"
@ -3469,6 +3565,17 @@ version = "6.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
[[package]]
name = "r2d2"
version = "0.8.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
dependencies = [
"log",
"parking_lot",
"scheduled-thread-pool",
]
[[package]]
name = "rand"
version = "0.8.7"
@ -3603,8 +3710,10 @@ dependencies = [
"arcstr",
"combine",
"itoa",
"num-bigint",
"num-bigint 0.5.1",
"percent-encoding",
"rustls 0.23.42",
"rustls-native-certs",
"ryu",
"sha1_smol",
"socket2 0.6.5",
@ -4001,6 +4110,15 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "scheduled-thread-pool"
version = "0.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
dependencies = [
"parking_lot",
]
[[package]]
name = "schemars"
version = "0.9.0"
@ -4954,6 +5072,12 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "ucd-trie"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
[[package]]
name = "unarray"
version = "0.1.4"

2
litellm-rust/Cargo.toml
View file

@ -28,6 +28,8 @@ litellm-types = { path = "crates/types" }
litellm-core-utils = { path = "crates/core-utils" }
litellm-cache = { path = "crates/cache" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-response = { path = "crates/cache-response" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }

2
litellm-rust/crates/cache-memory/Cargo.toml
View file

@ -7,8 +7,8 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
serde_json.workspace = true
[dev-dependencies]
serde_json.workspace = true
rstest.workspace = true
tokio.workspace = true

296
litellm-rust/crates/cache-memory/src/cache.rs
View file

@ -1,18 +1,20 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap, HashSet},
hash::Hash,
sync::{Arc, Mutex},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
BaseCache, BatchCache, CacheConnectionResult, CacheConnectionStatus, ClaimCache, CounterCache,
DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
type ValueMeasure<V> = Arc<dyn Fn(&V) -> Result<usize, Error> + Send + Sync>;
type ValueValidator<V> = Arc<dyn Fn(&V) -> Result<(), Error> + Send + Sync>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CacheWrite {
@ -33,7 +35,6 @@ pub struct InMemoryCache<V: Clone> {
default_ttl: Duration,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
validate_value: Option<ValueValidator<V>>,
now: Arc<dyn Fn() -> Duration + Send + Sync>,
}
@ -77,7 +78,6 @@ impl<V: Clone> InMemoryCache<V> {
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
max_entry_bytes,
measure_value,
validate_value: None,
now: Arc::new(now),
}
}
@ -91,9 +91,6 @@ impl<V: Clone> InMemoryCache<V> {
if self.max_size_in_memory == 0 {
return Ok(CacheWrite::Disabled);
}
if let Some(validate) = &self.validate_value {
validate(&value)?;
}
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&value)? > limit
{
@ -101,15 +98,13 @@ impl<V: Clone> InMemoryCache<V> {
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now);
let key = key.into();
state.values.insert(key.clone(), value);
Self::evict(&mut state, self.max_size_in_memory, now, &key);
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
let expiration = now + ttl.unwrap_or(self.default_ttl);
state.expirations.insert(key.clone(), expiration);
state.expiration_heap.push(Reverse((expiration, key)));
Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key, value);
Ok(CacheWrite::Stored)
}
@ -126,6 +121,14 @@ impl<V: Clone> InMemoryCache<V> {
Ok(state.values.get(key).cloned())
}
pub fn max_size_in_memory(&self) -> usize {
self.max_size_in_memory
}
pub fn max_entry_bytes(&self) -> Option<usize> {
self.max_entry_bytes
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
Ok(self
.state
@ -136,6 +139,25 @@ impl<V: Clone> InMemoryCache<V> {
.copied())
}
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
self.expires_at(key)
}
pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result<Vec<String>, Error> {
let state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut expirations = state
.expirations
.iter()
.map(|(key, expiration)| (key.clone(), *expiration))
.collect::<Vec<_>>();
expirations.sort_unstable_by_key(|(_, expiration)| *expiration);
Ok(expirations
.into_iter()
.take(count)
.map(|(key, _)| key)
.collect())
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::remove(&mut state, key);
@ -150,7 +172,7 @@ impl<V: Clone> InMemoryCache<V> {
Ok(())
}
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration) {
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: &str) {
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
if state.expirations.get(&key).copied() != Some(expiration) {
state.expiration_heap.pop();
@ -161,6 +183,9 @@ impl<V: Clone> InMemoryCache<V> {
break;
}
}
if state.values.contains_key(key) {
return;
}
while state.values.len() >= capacity {
let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else {
break;
@ -171,84 +196,205 @@ impl<V: Clone> InMemoryCache<V> {
}
}
fn set_expiration(state: &mut CacheState<V>, key: &str, expiration: Duration) {
if state.expirations.get(key).copied() != Some(expiration) {
state.expirations.insert(key.into(), expiration);
state
.expiration_heap
.push(Reverse((expiration, key.into())));
}
}
fn remove(state: &mut CacheState<V>, key: &str) {
state.values.remove(key);
state.expirations.remove(key);
}
}
impl InMemoryCache<CacheEntry> {
pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn response_cache_with_clock(
capacity: usize,
ttl: Duration,
max_entry_bytes: usize,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
let mut cache = Self::with_clock_and_size_measurement(
Some(capacity),
Some(ttl),
Some(max_entry_bytes),
Some(Arc::new(|entry: &CacheEntry| {
serde_json::to_vec(entry)
.map(|bytes| bytes.len())
.map_err(|_| Error::InvalidEntry)
})),
now,
impl<V> ClaimCache for InMemoryCache<V>
where
V: Clone + PartialEq + Send + Sync + 'static,
{
fn claim_cache(
&self,
key: &str,
candidate: V,
eligible: &[V],
context: ExactCacheContext,
) -> Result<V, Error> {
if self.max_size_in_memory == 0 {
return Ok(candidate);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
let existing = state
.values
.get(key)
.filter(|existing| eligible.is_empty() || eligible.contains(existing))
.cloned();
if let Some(existing) = &existing
&& eligible.is_empty()
&& *existing != candidate
{
return Ok(existing.clone());
}
let winner = existing.unwrap_or(candidate);
Self::set_expiration(
&mut state,
key,
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
);
cache.validate_value = Some(Arc::new(|entry: &CacheEntry| {
entry
.timestamp
.is_finite()
.then_some(())
.ok_or(Error::InvalidEntry)
}));
cache
state.values.insert(key.into(), winner.clone());
Ok(winner)
}
}
impl BaseCache for InMemoryCache<CacheEntry> {
type Value = CacheEntry;
impl CounterCache for InMemoryCache<f64> {
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
if self.max_size_in_memory == 0 {
return Ok(amount);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
let value = state.values.get(key).copied().unwrap_or_default() + amount;
if !state.expirations.contains_key(key) {
Self::set_expiration(
&mut state,
key,
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
);
}
state.values.insert(key.into(), value);
Ok(value)
}
}
fn default_ttl(&self) -> Duration {
self.default_ttl
impl InMemoryCache<f64> {
pub async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error> {
operations
.into_iter()
.map(|operation| {
self.increment_cache(
&operation.key,
operation.amount,
ExactCacheContext { ttl: operation.ttl },
)
})
.collect()
}
}
impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
type Value = V;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let ttl = self.get_ttl(&kwargs);
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let ttl = self.get_ttl(context).unwrap_or(self.default_ttl);
self.set_cache(key, value, Some(ttl)).map(|_| ())
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
self.get_cache(key)
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.delete_cache(key)
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
fn flush_cache(&self) -> Result<(), Error> {
self.flush_cache()
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "In-memory cache connection test successful".into(),
error: None,
})
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "In-memory cache connection test successful".into(),
error: None,
})
}
}
impl<V: Clone + Send + Sync + 'static> BatchCache for InMemoryCache<V> {}
impl<V: Clone + Send + Sync + 'static> DeleteCache for InMemoryCache<V> {
fn delete_cache(&self, key: &str) -> Result<(), Error> {
InMemoryCache::delete_cache(self, key)
}
}
impl<V: Clone + Send + Sync + 'static> FlushCache for InMemoryCache<V> {
fn flush_cache(&self) -> Result<(), Error> {
InMemoryCache::flush_cache(self)
}
}
impl<V: Clone + Send + Sync + 'static> TtlCache for InMemoryCache<V> {
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
InMemoryCache::async_get_ttl(self, key).await
}
}
impl<T> SetCache for InMemoryCache<HashSet<T>>
where
T: Clone + Eq + Hash + Send + Sync + 'static,
{
type SetValue = T;
type SetResult = Vec<T>;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
if self.max_size_in_memory == 0 {
return Ok(values);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
let mut stored = state.values.get(key).cloned().unwrap_or_default();
stored.extend(values.iter().cloned());
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&stored)? > limit
{
return Ok(values);
}
if !state.expirations.contains_key(key) {
Self::set_expiration(&mut state, key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key.into(), stored);
Ok(values)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn repeated_increments_keep_one_heap_entry_per_expiration() {
let cache = InMemoryCache::<f64>::new(Some(4), None);
for _ in 0..100 {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1);
}
}

312
litellm-rust/crates/cache-memory/tests/cache.rs
View file

@ -1,8 +1,16 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use std::{
collections::HashSet,
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error};
use litellm_cache::{
BaseCache, CacheBackend, CacheConnectionStatus, ClaimCache, CounterCache, DeleteCache, Error,
ExactCacheContext, IncrementOperation, SetCache, get_cache, set_cache,
};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
@ -84,66 +92,49 @@ fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>
}
#[test]
fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
let disabled = InMemoryCache::<CacheEntry>::response_cache(0, Duration::from_secs(60), 80);
fn disabled_size_limited_and_validated_writes_are_observable() {
let cache = |capacity| {
InMemoryCache::with_clock_and_size_measurement(
Some(capacity),
Some(Duration::from_secs(60)),
Some(4),
Some(Arc::new(|value: &String| {
if value.is_empty() {
return Err(Error::InvalidEntry);
}
Ok(value.len())
})),
|| Duration::from_secs(100),
)
};
let disabled = cache(0);
assert_eq!(
disabled
.set_cache(
"a",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("x")
},
None
)
.unwrap(),
disabled.set_cache("a", "x".into(), None).unwrap(),
CacheWrite::Disabled
);
let cache = InMemoryCache::<CacheEntry>::response_cache(2, Duration::from_secs(60), 80);
let cache = cache(2);
assert_eq!(
cache
.set_cache(
"large",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("x".repeat(100))
},
None
)
.unwrap(),
cache.set_cache("large", "oversized".into(), None).unwrap(),
CacheWrite::TooLarge
);
cache
.set_cache(
"small",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("ok"),
},
None,
)
.unwrap();
assert!(cache.get_cache("small").unwrap().is_some());
assert_eq!(cache.get_cache("large").unwrap(), None);
assert_eq!(
cache
.set_cache(
"invalid",
CacheEntry {
timestamp: f64::NAN,
response: serde_json::json!("bad"),
},
None,
)
.unwrap_err(),
Error::InvalidEntry
cache.set_cache("small", "ok".into(), None).unwrap(),
CacheWrite::Stored
);
assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into()));
assert_eq!(
cache.set_cache("invalid", String::new(), None),
Err(Error::InvalidEntry)
);
assert_eq!(cache.get_cache("invalid").unwrap(), None);
cache.delete_cache("small").unwrap();
cache.flush_cache().unwrap();
assert_eq!(cache.get_cache("small").unwrap(), None);
}
#[tokio::test]
async fn connection_test_matches_python_result_contract() {
let cache = InMemoryCache::<CacheEntry>::default();
let cache = InMemoryCache::<String>::default();
let result = BaseCache::test_connection(&cache).await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "In-memory cache connection test successful");
@ -156,3 +147,222 @@ async fn connection_test_matches_python_result_contract() {
})
);
}
#[tokio::test]
async fn generic_consumers_share_typed_values_and_honor_expiration() {
let clock = clock();
let cache: CacheBackend<InMemoryCache<String>> = Arc::new(cache(clock.clone(), 4));
let reader = Arc::clone(&cache);
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
};
set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap();
assert_eq!(
get_cache(reader.as_ref(), "sync", &context).unwrap(),
Some("first".into())
);
cache
.batch_cache_write("async", "second".into(), context.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(vec![("batch".into(), "third".into())], context.clone())
.await
.unwrap();
drop(cache);
for (key, value) in [("sync", "first"), ("async", "second"), ("batch", "third")] {
assert_eq!(
reader.async_get_cache(key, &context).await.unwrap(),
Some(value.into())
);
}
reader.async_delete_cache("async").await.unwrap();
assert_eq!(
reader.async_get_cache("async", &context).await.unwrap(),
None
);
clock.store(106, Ordering::SeqCst);
assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None);
assert_eq!(
reader.async_get_cache("batch", &context).await.unwrap(),
None
);
}
#[test]
fn claims_are_atomic_and_refresh_eligible_winners() {
let clock = clock();
let cache = InMemoryCache::with_clock(Some(4), Some(Duration::from_secs(60)), {
let clock = clock.clone();
move || Duration::from_secs(clock.load(Ordering::SeqCst))
});
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(10)),
};
assert_eq!(
cache
.claim_cache("affinity", "first".to_string(), &[], context.clone())
.unwrap(),
"first"
);
clock.store(103, Ordering::SeqCst);
assert_eq!(
cache
.claim_cache("affinity", "second".to_string(), &[], context.clone())
.unwrap(),
"first"
);
assert_eq!(
cache.expires_at("affinity").unwrap(),
Some(Duration::from_secs(110))
);
clock.store(105, Ordering::SeqCst);
assert_eq!(
cache
.claim_cache(
"affinity",
"second".to_string(),
&["first".to_string(), "second".to_string()],
context,
)
.unwrap(),
"first"
);
assert_eq!(
cache.expires_at("affinity").unwrap(),
Some(Duration::from_secs(115))
);
}
#[test]
fn counters_increment_under_one_lock() {
let cache = InMemoryCache::<f64>::default();
assert_eq!(
CounterCache::increment_cache(&cache, "counter", 1.5, ExactCacheContext::default())
.unwrap(),
1.5
);
assert_eq!(
CounterCache::increment_cache(&cache, "counter", 2.0, ExactCacheContext::default())
.unwrap(),
3.5
);
}
#[rstest]
fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("hot", "1".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("cold", "2".into(), Some(Duration::from_secs(20)))
.unwrap();
cache.set_cache("cold", "3".into(), None).unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
cache
.claim_cache("cold", "4".into(), &[], ExactCacheContext::default())
.unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
cache.set_cache("new", "5".into(), None).unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), None);
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
assert_eq!(cache.get_cache("new").unwrap(), Some("5".into()));
}
#[test]
fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() {
let cache = InMemoryCache::<f64>::new(Some(2), None);
for key in ["a", "b", "a", "b"] {
cache
.increment_cache(key, 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.get_cache("a").unwrap(), Some(2.0));
assert_eq!(cache.get_cache("b").unwrap(), Some(2.0));
}
#[test]
fn disabled_cache_does_not_retain_claims_or_counters() {
let claims = InMemoryCache::<String>::new(Some(0), None);
assert_eq!(
claims
.claim_cache("key", "first".into(), &[], ExactCacheContext::default())
.unwrap(),
"first"
);
assert_eq!(claims.get_cache("key").unwrap(), None);
let counters = InMemoryCache::<f64>::new(Some(0), None);
assert_eq!(
counters
.increment_cache("key", 2.0, ExactCacheContext::default())
.unwrap(),
2.0
);
assert_eq!(counters.get_cache("key").unwrap(), None);
}
#[tokio::test]
async fn ttl_and_oldest_key_operations_use_the_stored_expirations() {
let clock = Arc::new(AtomicU64::new(100));
let cache = cache(clock, 3);
cache
.set_cache("later", "2".into(), Some(Duration::from_secs(20)))
.unwrap();
cache
.set_cache("first", "1".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.async_get_ttl("first").await.unwrap(),
Some(Duration::from_secs(110))
);
assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]);
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
}
#[tokio::test]
async fn increment_pipeline_preserves_operation_order() {
let cache = InMemoryCache::<f64>::new(Some(3), None);
assert_eq!(
cache
.async_increment_pipeline(vec![
IncrementOperation {
key: "a".into(),
amount: 1.0,
ttl: Some(Duration::from_secs(10)),
},
IncrementOperation {
key: "a".into(),
amount: 2.0,
ttl: Some(Duration::from_secs(20)),
},
])
.await
.unwrap(),
[1.0, 3.0]
);
assert_eq!(cache.get_cache("a").unwrap(), Some(3.0));
}
#[tokio::test]
async fn set_capability_preserves_python_result_and_deduplicates_storage() {
let cache = InMemoryCache::<HashSet<String>>::new(None, None);
let inserted = vec!["a".into(), "a".into(), "b".into()];
assert_eq!(
cache
.async_set_cache_sadd("members", inserted.clone(), None)
.await
.unwrap(),
inserted
);
assert_eq!(
cache.get_cache("members").unwrap(),
Some(HashSet::from(["a".into(), "b".into()]))
);
}

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

@ -7,9 +7,10 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
redis = "1.7.0"
serde_json.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
r2d2 = "0.8.10"
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"
serde_json.workspace = true

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

@ -1,58 +1,243 @@
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
ClaimCache, CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
};
use redis::Commands;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const KEY_PREFIX: &str = "litellm-cache:";
mod operations;
pub struct RedisCache<C = redis::Connection> {
connection: Arc<Mutex<C>>,
default_ttl: Duration,
pub use operations::{
RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
struct PooledConnection {
connection: redis::Connection,
failed: bool,
}
impl RedisCache<redis::Connection> {
pub fn new(url: &str, default_ttl: Option<Duration>) -> Result<Self, Error> {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
Ok(Self::with_connection(connection, default_ttl))
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
struct ConnectionManager(redis::Client);
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection;
type Error = redis::RedisError;
fn connect(&self) -> Result<PooledConnection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut PooledConnection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut PooledConnection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
impl<C> RedisCache<C>
const INCREMENT_SCRIPT: &str = concat!(
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
"if redis.call('TTL', KEYS[1]) == -1 then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
);
const CLAIM_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
"elseif current ~= ARGV[1] then return 0; end; ",
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
);
const CLAIM_ATTEMPTS: usize = 8;
enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Fixed(Mutex<C>),
}
struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike);
impl redis::ConnectionLike for ConnectionRef<'_> {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
self.0.req_packed_command(cmd)
}
fn req_packed_commands(
&mut self,
cmd: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
self.0.req_packed_commands(cmd, offset, count)
}
fn get_db(&self) -> i64 {
self.0.get_db()
}
fn supports_pipelining(&self) -> bool {
self.0.supports_pipelining()
}
fn check_connection(&mut self) -> bool {
self.0.check_connection()
}
fn is_open(&self) -> bool {
self.0.is_open()
}
}
impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn with_connection(connection: C, default_ttl: Option<Duration>) -> Self {
Self {
connection: Arc::new(Mutex::new(connection)),
fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
match self {
Self::Pool(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let result = operation(&mut ConnectionRef(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut ConnectionRef(&mut *connection))
}
}
}
}
pub struct RedisCache<S, C = redis::Connection> {
connections: Arc<Connections<C>>,
default_ttl: Duration,
codec: S,
namespace: Option<String>,
}
impl<S: CacheCodec> RedisCache<S> {
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let pool = r2d2::Pool::builder()
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(ConnectionManager(client))
.map_err(|_| Error::Unavailable)?;
Ok(Self {
connections: Arc::new(Connections::Pool(pool)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
})
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
Self {
connections: Arc::new(Connections::Fixed(Mutex::new(connection))),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
}
}
fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
self.connection.lock().map_err(|_| Error::Unavailable)
pub fn with_namespace(self, namespace: Option<String>) -> Self {
Self {
namespace: namespace.filter(|value| !value.is_empty()),
..self
}
}
fn namespaced_key(key: &str) -> String {
format!("{KEY_PREFIX}{key}")
pub fn namespace(&self) -> Option<&str> {
self.namespace.as_deref()
}
fn namespaced_pattern() -> &'static str {
const PATTERN: &str = "litellm-cache:*";
PATTERN
fn namespaced_key(&self, key: &str) -> String {
namespaced_key(self.namespace.as_deref(), key)
}
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
fn namespaced_pattern(&self) -> Result<String, Error> {
let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
let escaped: String = namespace
.chars()
.flat_map(|ch| {
if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
vec!['\\', ch]
} else {
vec![ch]
}
})
.collect();
Ok(format!("{escaped}:*"))
}
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
let mut cursor = 0u64;
loop {
let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
.cursor_arg(cursor)
.arg("MATCH")
.arg(pattern)
.arg("COUNT")
.arg(1000)
.query(connection)
.map_err(|_| Error::Unavailable)?;
if !keys.is_empty() {
connection
.del::<_, usize>(keys)
.map_err(|_| Error::Unavailable)?;
}
if next_cursor == 0 {
return Ok(());
}
cursor = next_cursor;
}
}
fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
_ => Err(Error::InvalidEntry),
}
}
fn decode_batch_response(&self, value: redis::Value) -> Result<BatchEntry<S::Value>, Error> {
match self.decode_response(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}
fn ttl_seconds(ttl: Duration) -> u64 {
@ -61,196 +246,418 @@ where
.max(1)
}
fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
async fn run_blocking<T, F>(connections: Arc<Connections<C>>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut C) -> Result<T, Error> + Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
Box::pin(async move {
tokio::task::spawn_blocking(move || {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut connection)
})
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
})
}
}
impl<C> BaseCache for RedisCache<C>
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
match namespace {
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
format!("{namespace}:{key}")
}
_ => key.into(),
}
}
impl<S, C> BaseCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = CacheEntry;
type Value = S::Value;
type Context = ExactCacheContext;
fn default_ttl(&self) -> Duration {
self.default_ttl
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let payload = Self::encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
self.connection()?
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
.map_err(|_| Error::Unavailable)
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
self.connection()?
.get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)?
.map(Self::decode)
.transpose()
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.connection()?
.del::<_, ()>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)
}
fn flush_cache(&self) -> Result<(), Error> {
let mut connection = self.connection()?;
let keys = connection
.scan_match(Self::namespaced_pattern())
.map_err(|_| Error::Unavailable)?
.collect::<redis::RedisResult<Vec<String>>>()
.map_err(|_| Error::Unavailable)?;
if keys.is_empty() {
return Ok(());
}
connection
.del::<_, usize>(keys)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
fn async_set_cache<'a>(
&'a self,
key: &'a str,
fn set_cache(
&self,
key: &str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let payload = Self::encode(&value);
let key = Self::namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
context: &ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(context).unwrap_or(self.default_ttl));
let key = self.namespaced_key(key);
self.connections.execute(|connection| {
connection
.set_ex::<_, _, ()>(key, payload?, ttl)
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
})
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
_: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
let key = Self::namespaced_key(key);
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.get::<_, Option<Vec<u8>>>(key)
.map_err(|_| Error::Unavailable)
})
.await?
.map(Self::decode)
.transpose()
})
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = self.connections.execute(|connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
})?;
self.decode_response(value)
}
fn async_set_cache_pipeline<'a>(
&'a self,
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
})
.await?;
self.decode_response(value)
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
context: ExactCacheContext,
) -> Result<(), Error> {
let entries = cache_list
.into_iter()
.map(|(key, value)| {
Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
self.codec
.encode(&value)
.map(|payload| (self.namespaced_key(&key), payload))
})
.collect::<Result<Vec<_>, _>>();
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
for (key, payload) in entries? {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)?;
.collect::<Result<Vec<_>, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, payload) in entries {
pipeline
.cmd("SETEX")
.arg(key)
.arg(ttl)
.arg(payload)
.ignore();
}
Ok(())
pipeline
.query::<()>(connection)
.map_err(|_| Error::Unavailable)
})
.await
}
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
let key = Self::namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Self::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match redis::cmd("PING").query::<String>(connection) {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Redis connection failed: {error}"),
error: Some(error.to_string()),
},
})
})
.await
{
Ok(result) => Ok(result),
Err(error) => Ok(CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Redis connection failed: {error}"),
error: Some(error.to_string()),
}),
}
}
}
impl<S, C> BatchCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_cache(
&self,
keys: &[String],
_: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = self.connections.execute(|connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})?;
values
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
values
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
}
impl<S, C> DeleteCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.connections
.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
.await
}
}
impl<S, C> FlushCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
self.connections
.execute(|connection| Self::flush_matching(connection, &pattern))
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), |connection| {
redis::cmd("PING")
.query::<String>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
})
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Self::flush_matching(connection, &pattern)
})
.await
}
}
impl<S, C> CounterCache for RedisCache<S, C>
where
S: CacheCodec<Value = f64>,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
self.connections
.execute(|connection| increment(connection, key, amount, ttl))
}
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
increment(connection, key, amount, ttl)
})
.await
}
}
fn increment(
connection: &mut ConnectionRef<'_>,
key: String,
amount: f64,
ttl: u64,
) -> Result<f64, Error> {
redis::cmd("EVAL")
.arg(INCREMENT_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
}
fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
_ => Err(Error::InvalidEntry),
}
}
/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
/// bytes that decision was made on, retried when another claimant wins the race.
fn claim<S: CacheCodec>(
connection: &mut ConnectionRef<'_>,
codec: &S,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
ttl: u64,
) -> Result<S::Value, Error>
where
S::Value: PartialEq,
{
let payload = codec.encode(&candidate)?;
if payload.is_empty() {
return Err(Error::InvalidEntry);
}
for _ in 0..CLAIM_ATTEMPTS {
let current = stored_bytes(
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)?,
)?
.filter(|bytes| !bytes.is_empty());
let existing = current
.as_deref()
.and_then(|bytes| codec.decode(bytes).ok())
.filter(|existing| eligible.is_empty() || eligible.contains(existing));
let refresh = existing
.as_ref()
.is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
let write: &[u8] = if existing.is_some() { b"" } else { &payload };
let applied = redis::cmd("EVAL")
.arg(CLAIM_SCRIPT)
.arg(1)
.arg(key)
.arg(current.as_deref().unwrap_or_default())
.arg(ttl)
.arg(write)
.arg(u8::from(refresh))
.query::<bool>(connection)
.map_err(|_| Error::Unavailable)?;
if applied {
return Ok(existing.unwrap_or(candidate));
}
}
Err(Error::Unavailable)
}
impl<S, C> ClaimCache for RedisCache<S, C>
where
S: CacheCodec + Clone + 'static,
S::Value: PartialEq,
C: redis::ConnectionLike + Send + 'static,
{
fn claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
self.connections
.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
}
async fn async_claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: Vec<S::Value>,
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
let codec = self.codec.clone();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
claim(connection, &codec, &key, candidate, &eligible, ttl)
})
.await
}
}
#[cfg(test)]
mod tests {
use super::RedisCache;
use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use std::time::Duration;
fn entry() -> CacheEntry {
CacheEntry {
timestamp: 123.0,
response: json!({"choices": [{"text": "cached"}]}),
}
}
use litellm_cache::{
BaseCache, CacheCodec, DeleteCache, ExactCacheContext, FlushCache, JsonCodec,
};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
#[test]
fn cache_entries_round_trip_through_json() {
let entry = entry();
let encoded = RedisCache::<redis::Connection>::encode(&entry).unwrap();
assert_eq!(
RedisCache::<redis::Connection>::decode(encoded).unwrap(),
entry
);
}
use super::RedisCache;
#[test]
fn invalid_json_is_rejected() {
assert!(RedisCache::<redis::Connection>::decode(b"not json".to_vec()).is_err());
fn entry() -> serde_json::Value {
json!({"deployment": "model-a", "cooldown_seconds": 30})
}
#[test]
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::ZERO),
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_millis(1500)),
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_secs(15)),
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_secs(15)),
15
);
}
@ -258,7 +665,9 @@ mod tests {
#[test]
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
let value = entry();
let payload = RedisCache::<redis::Connection>::encode(&value).unwrap();
let payload = JsonCodec::<serde_json::Value>::new()
.encode(&value)
.unwrap();
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
@ -271,13 +680,17 @@ mod tests {
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
.with_namespace(Some("litellm-cache".into()));
cache
.set_cache("key", value.clone(), CacheKwargs::default())
.set_cache("key", value.clone(), &ExactCacheContext::default())
.unwrap();
assert_eq!(
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(value)
);
cache.delete_cache("key").unwrap();
@ -290,13 +703,17 @@ mod tests {
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("litellm-cache:*"),
.arg("litellm-cache:*")
.arg("COUNT")
.arg(1000),
Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
.with_namespace(Some("litellm-cache".into()));
cache.flush_cache().unwrap();
}
@ -305,7 +722,9 @@ mod tests {
async fn test_connection_runs_ping_off_executor() {
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
.with_namespace(Some("litellm-cache".into()));
assert_eq!(
cache.test_connection().await.unwrap().status,

633
litellm-rust/crates/cache-redis/src/cache/operations.rs vendored Normal file
View file

@ -0,0 +1,633 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
CacheCodec, CacheScript, ClientInfoCache, Error, IncrementOperation, QueueCache, ScanCache,
ScriptCache, SetCache, TtlCache,
};
use redis::Commands;
use super::{ConnectionRef, Connections, RedisCache, namespaced_key};
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return count"
);
const SET_MAX_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]); ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return ARGV[1]; end; return current"
);
#[derive(Clone, Debug, PartialEq)]
pub enum RedisArg {
Bytes(Vec<u8>),
Integer(i64),
Float(f64),
}
impl From<&str> for RedisArg {
fn from(value: &str) -> Self {
Self::Bytes(value.as_bytes().to_vec())
}
}
impl From<String> for RedisArg {
fn from(value: String) -> Self {
Self::Bytes(value.into_bytes())
}
}
impl From<Vec<u8>> for RedisArg {
fn from(value: Vec<u8>) -> Self {
Self::Bytes(value)
}
}
impl From<i64> for RedisArg {
fn from(value: i64) -> Self {
Self::Integer(value)
}
}
impl From<f64> for RedisArg {
fn from(value: f64) -> Self {
Self::Float(value)
}
}
impl redis::ToRedisArgs for RedisArg {
fn write_redis_args<W>(&self, out: &mut W)
where
W: ?Sized + redis::RedisWrite,
{
match self {
Self::Bytes(value) => value.write_redis_args(out),
Self::Integer(value) => value.write_redis_args(out),
Self::Float(value) => value.write_redis_args(out),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisRpushOperation {
pub key: String,
pub values: Vec<RedisArg>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RedisLpopOperation {
pub key: String,
pub count: Option<usize>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RedisLpopResult {
Missing,
Value(Vec<u8>),
Values(Vec<Vec<u8>>),
}
pub struct RedisScript<C> {
connections: Arc<Connections<C>>,
namespace: Option<String>,
source: String,
}
impl<C> CacheScript for RedisScript<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Argument = RedisArg;
type Output = redis::Value;
async fn invoke(
&self,
keys: Vec<String>,
arguments: Vec<Self::Argument>,
) -> Result<Self::Output, Error> {
let keys = keys
.into_iter()
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
.collect::<Vec<_>>();
let connections = Arc::clone(&self.connections);
let source = self.source.clone();
tokio::task::spawn_blocking(move || {
connections.execute(|connection| {
redis::cmd("EVAL")
.arg(source)
.arg(keys.len())
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(|_| Error::Unavailable)
})
})
.await
.map_err(|_| Error::Unavailable)?
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
if keys.is_empty() {
return Ok(0);
}
let keys = keys
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del(keys).map_err(|_| Error::Unavailable)
})
.await
}
pub fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = self.connections.execute(|connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})?;
values.into_iter().map(count).collect()
}
pub async fn async_batch_get_counts(
&self,
keys: Vec<String>,
) -> Result<Vec<Option<i64>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
values.into_iter().map(count).collect()
}
pub fn sync_ping(&self) -> Result<bool, Error> {
self.connections.execute(|connection| {
redis::cmd("PING")
.query::<String>(connection)
.map(|response| response == "PONG")
.map_err(|_| Error::Unavailable)
})
}
pub async fn ping(&self) -> Result<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
redis::cmd("PING")
.query::<String>(connection)
.map(|response| response == "PONG")
.map_err(|_| Error::Unavailable)
})
.await
}
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<i64>, Error> {
let key = self.namespaced_key(key);
let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
Ok((ttl >= 0).then_some(ttl))
}
pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
let pattern = format!("{}*", self.namespaced_key(pattern));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut cursor = 0u64;
let mut matches = Vec::new();
loop {
let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
.cursor_arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(count)
.query(connection)
.map_err(|_| Error::Unavailable)?;
matches.extend(keys);
if matches.len() >= count || next_cursor == 0 {
matches.truncate(count);
return Ok(matches);
}
cursor = next_cursor;
}
})
.await
}
pub async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<RedisArg>,
ttl: Option<Duration>,
) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry);
}
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
pipeline.cmd("SADD").arg(&key).arg(values);
pipeline.cmd("EXPIRE").arg(&key).arg(ttl).ignore();
pipeline
.query::<(usize,)>(connection)
.map(|(added,)| added)
.map_err(|_| Error::Unavailable)
})
.await
}
pub async fn async_rpush(&self, key: &str, values: Vec<RedisArg>) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry);
}
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("RPUSH")
.arg(key)
.arg(values)
.query(connection)
.map_err(|_| Error::Unavailable)
})
.await
}
pub async fn async_rpush_pipeline(
&self,
operations: Vec<RedisRpushOperation>,
) -> Result<Vec<usize>, Error> {
let operations = operations
.into_iter()
.map(|operation| {
if operation.values.is_empty() {
return Err(Error::InvalidEntry);
}
Ok((self.namespaced_key(&operation.key), operation.values))
})
.collect::<Result<Vec<_>, _>>()?;
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, values) in operations {
pipeline.cmd("RPUSH").arg(key).arg(values);
}
pipeline.query(connection).map_err(|_| Error::Unavailable)
})
.await
}
pub async fn async_lpop(
&self,
key: &str,
count: Option<usize>,
) -> Result<RedisLpopResult, Error> {
let key = self.namespaced_key(key);
let multiple = count.is_some();
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
command.arg(count);
}
command
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
lpop_result(value, multiple)
}
pub async fn async_lpop_pipeline(
&self,
operations: Vec<RedisLpopOperation>,
) -> Result<Vec<RedisLpopResult>, Error> {
let operations = operations
.into_iter()
.map(|operation| (self.namespaced_key(&operation.key), operation.count))
.collect::<Vec<_>>();
if operations.is_empty() {
return Ok(Vec::new());
}
let multiple = operations
.iter()
.map(|(_, count)| count.is_some())
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, count) in operations {
let command = pipeline.cmd("LPOP").arg(key);
if let Some(count) = count {
command.arg(count);
}
}
pipeline
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
values
.into_iter()
.zip(multiple)
.map(|(value, multiple)| lpop_result(value, multiple))
.collect()
}
pub async fn async_eval(
&self,
script: String,
keys: Vec<String>,
arguments: Vec<RedisArg>,
) -> Result<redis::Value, Error> {
let keys = keys
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(|_| Error::Unavailable)
})
.await
}
pub fn client_list(&self) -> Result<String, Error> {
self.connections.execute(|connection| {
redis::cmd("CLIENT")
.arg("LIST")
.query(connection)
.map_err(|_| Error::Unavailable)
})
}
pub fn info(&self) -> Result<String, Error> {
self.connections.execute(|connection| {
redis::cmd("INFO")
.query(connection)
.map_err(|_| Error::Unavailable)
})
}
pub fn flushall(&self) -> Result<(), Error> {
self.connections.execute(|connection| {
redis::cmd("FLUSHALL")
.query(connection)
.map_err(|_| Error::Unavailable)
})
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec<Value = f64>,
C: redis::ConnectionLike + Send + 'static,
{
pub fn increment_with_floor(
&self,
key: &str,
amount: i64,
ttl: Duration,
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl);
self.connections
.execute(|connection| increment_with_floor(connection, key, amount, ttl))
}
pub async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error> {
let operations = operations
.into_iter()
.map(|operation| {
(
self.namespaced_key(&operation.key),
operation.amount,
operation.ttl.map(Self::ttl_seconds),
)
})
.collect::<Vec<_>>();
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, amount, ttl) in operations {
pipeline.cmd("INCRBYFLOAT").arg(&key).arg(amount);
if let Some(ttl) = ttl {
pipeline.cmd("EXPIRE").arg(key).arg(ttl).ignore();
}
}
pipeline.query(connection).map_err(|_| Error::Unavailable)
})
.await
}
pub async fn async_increment_with_floor(
&self,
key: &str,
amount: i64,
ttl: Duration,
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
increment_with_floor(connection, key, amount, ttl)
})
.await
}
pub async fn async_set_max(
&self,
key: &str,
value: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)
.arg(key)
.arg(value)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
})
.await
}
}
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
match value {
redis::Value::BulkString(bytes) => Ok(bytes),
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
_ => Err(Error::InvalidEntry),
}
}
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
match value {
redis::Value::Nil => Ok(RedisLpopResult::Missing),
redis::Value::Array(values) if multiple => values
.into_iter()
.map(redis_bytes)
.collect::<Result<Vec<_>, _>>()
.map(RedisLpopResult::Values),
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
_ => Err(Error::InvalidEntry),
}
}
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::Int(value) => Ok(Some(value)),
redis::Value::BulkString(value) => std::str::from_utf8(&value)
.ok()
.and_then(|value| value.parse().ok())
.map(Some)
.ok_or(Error::InvalidEntry),
redis::Value::SimpleString(value) => {
value.parse().map(Some).map_err(|_| Error::InvalidEntry)
}
_ => Err(Error::InvalidEntry),
}
}
fn increment_with_floor(
connection: &mut ConnectionRef<'_>,
key: String,
amount: i64,
ttl: u64,
) -> Result<i64, Error> {
redis::cmd("EVAL")
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
}
impl<S, C> TtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
RedisCache::async_get_ttl(self, key)
.await
.map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64)))
}
}
impl<S, C> ScanCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
RedisCache::async_scan_iter(self, pattern, count).await
}
}
impl<S, C> ClientInfoCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type ClientList = String;
type Info = String;
fn client_list(&self) -> Result<Self::ClientList, Error> {
RedisCache::client_list(self)
}
fn info(&self) -> Result<Self::Info, Error> {
RedisCache::info(self)
}
}
impl<S, C> SetCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type SetValue = RedisArg;
type SetResult = usize;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
RedisCache::async_set_cache_sadd(self, key, values, ttl).await
}
}
impl<S, C> QueueCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type QueueValue = RedisArg;
type PopResult = RedisLpopResult;
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
RedisCache::async_rpush(self, key, values).await
}
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
RedisCache::async_lpop(self, key, count).await
}
}
impl<S, C> ScriptCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Script = RedisScript<C>;
fn async_register_script(&self, source: String) -> Self::Script {
RedisScript {
connections: Arc::clone(&self.connections),
namespace: self.namespace.clone(),
source,
}
}
}

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

@ -1,3 +1,7 @@
mod cache;
mod topology;
pub use cache::RedisCache;
pub use cache::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};
pub use topology::{RedisNode, RedisTopology};

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

@ -0,0 +1,14 @@
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RedisNode {
pub host: String,
pub port: u16,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum RedisTopology {
#[default]
Standalone,
Cluster {
startup_nodes: Vec<RedisNode>,
},
}

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

@ -1,6 +1,703 @@
use litellm_cache_redis::RedisCache;
use std::time::Duration;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheScript, ClaimCache,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
ScriptCache, get_cache, set_cache,
};
use litellm_cache_redis::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation,
};
use redis_test::{MockCmd, MockRedisConnection};
struct TaggedByteCodec(u8);
impl CacheCodec for TaggedByteCodec {
type Value = u8;
fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
if *value > 127 {
return Err(Error::InvalidEntry);
}
Ok(vec![self.0, *value])
}
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
match bytes {
[tag, value] if *tag == self.0 => Ok(*value),
_ => Err(Error::InvalidEntry),
}
}
}
#[test]
fn constructor_rejects_invalid_urls() {
assert!(RedisCache::new("not a redis url", None).is_err());
assert!(RedisCache::new("not a redis url", None, JsonCodec::<String>::new()).is_err());
}
#[test]
fn generic_helpers_use_the_injected_codec_and_ttl() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
.arg("counter")
.arg(2)
.arg([42u8, 7].as_slice()),
Ok("OK"),
),
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
let context = ExactCacheContext {
ttl: Some(Duration::from_millis(1500)),
};
set_cache(&cache, "counter", 7, &context).unwrap();
assert_eq!(get_cache(&cache, "counter", &context).unwrap(), Some(7));
}
#[tokio::test]
async fn async_operations_preserve_codec_ttl_and_missing_values() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
.arg("counter")
.arg(9)
.arg([42u8, 7].as_slice()),
Ok("OK"),
),
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
MockCmd::new(
redis::cmd("SETEX")
.arg("batch")
.arg(2)
.arg([42u8, 8].as_slice()),
Ok("OK"),
),
MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)),
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(
connection,
Some(Duration::from_secs(9)),
TaggedByteCodec(42),
);
let context = ExactCacheContext::default();
cache
.batch_cache_write("counter", 7, context.clone())
.await
.unwrap();
assert_eq!(
cache.async_get_cache("counter", &context).await.unwrap(),
Some(7)
);
cache
.async_set_cache_pipeline(
vec![("batch".into(), 8)],
ExactCacheContext {
ttl: Some(Duration::from_millis(1500)),
},
)
.await
.unwrap();
cache.async_delete_cache("counter").await.unwrap();
assert_eq!(
cache.async_get_cache("counter", &context).await.unwrap(),
None
);
}
#[tokio::test]
async fn codec_errors_propagate_without_writing_partial_batches() {
let connection = MockRedisConnection::new([
MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
let context = ExactCacheContext::default();
assert_eq!(
cache.set_cache("invalid", 255, &context),
Err(Error::InvalidEntry)
);
assert_eq!(
cache.async_set_cache("invalid", 255, context.clone()).await,
Err(Error::InvalidEntry)
);
assert_eq!(
cache
.async_set_cache_pipeline(
vec![("valid".into(), 7), ("invalid".into(), 255)],
context.clone(),
)
.await,
Err(Error::InvalidEntry)
);
assert_eq!(
cache.get_cache("invalid", &context),
Err(Error::InvalidEntry)
);
assert_eq!(
cache.async_get_cache("invalid", &context).await,
Err(Error::InvalidEntry)
);
}
#[test]
fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() {
let connection = MockRedisConnection::new([
MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
.with_namespace(Some("team".into()));
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
None
);
assert_eq!(
cache
.get_cache("team:key", &ExactCacheContext::default())
.unwrap(),
None
);
}
#[test]
fn flush_requires_a_namespace_and_escapes_glob_metacharacters() {
let unscoped = RedisCache::with_connection(
MockRedisConnection::new([]).assert_all_commands_consumed(),
None,
JsonCodec::<String>::new(),
);
assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush));
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("team\\*:*")
.arg("COUNT")
.arg(1000),
Ok(redis_test::redis_value!(["0", ["team*:key"]])),
),
MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let scoped = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
.with_namespace(Some("team*".into()));
scoped.flush_cache().unwrap();
}
#[tokio::test]
async fn connection_failures_use_the_python_result_contract() {
let error = redis::RedisError::from((redis::ErrorKind::Io, "connection refused"));
let connection =
MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Err::<String, _>(error))])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new());
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Failed);
assert!(result.message.starts_with("Redis connection failed:"));
assert!(result.error.is_some());
}
#[tokio::test]
async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() {
let connection = MockRedisConnection::new([MockCmd::new(
redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"),
Ok(vec![
redis::Value::BulkString(vec![42, 7]),
redis::Value::Nil,
redis::Value::BulkString(vec![99, 7]),
]),
)])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
assert_eq!(
cache
.async_batch_get_cache(
vec!["hit".into(), "miss".into(), "invalid".into()],
ExactCacheContext::default(),
)
.await
.unwrap(),
vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid]
);
}
#[tokio::test]
async fn async_flush_deletes_each_scan_page_separately() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("team:*")
.arg("COUNT")
.arg(1000),
Ok(redis_test::redis_value!(["7", ["team:a", "team:b"]])),
),
MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)),
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(7)
.arg("MATCH")
.arg("team:*")
.arg("COUNT")
.arg(1000),
Ok(redis_test::redis_value!(["0", ["team:c"]])),
),
MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
.with_namespace(Some("team".into()));
cache.async_flush_cache().await.unwrap();
}
#[tokio::test]
async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() {
let mut sadd_pipeline = redis::pipe();
sadd_pipeline
.cmd("SADD")
.arg("team:members")
.arg("a")
.arg("b")
.cmd("EXPIRE")
.arg("team:members")
.arg(600u64)
.ignore();
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("MGET").arg("team:count").arg("team:missing"),
Ok(redis_test::redis_value!(["7", nil])),
),
MockCmd::new(
redis::cmd("MGET").arg("team:count").arg("team:missing"),
Ok(redis_test::redis_value!(["7", nil])),
),
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
MockCmd::new(redis::cmd("TTL").arg("team:missing"), Ok(-2i64)),
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("team:job-*")
.arg("COUNT")
.arg(25),
Ok(redis_test::redis_value!(["4", ["team:job-a"]])),
),
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(4)
.arg("MATCH")
.arg("team:job-*")
.arg("COUNT")
.arg(25),
Ok(redis_test::redis_value!(["0", ["team:job-b"]])),
),
MockCmd::new(
redis::cmd("DEL").arg("team:job-a").arg("team:job-b"),
Ok(2u32),
),
MockCmd::with_values(
sadd_pipeline,
Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]),
),
MockCmd::new(
redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"),
Ok(2u32),
),
MockCmd::new(
redis::cmd("LPOP").arg("team:queue").arg(2usize),
Ok(redis_test::redis_value!(["a", "b"])),
),
MockCmd::new(
redis::cmd("EVAL")
.arg("return KEYS[1]")
.arg(1usize)
.arg("team:key"),
Ok("team:key"),
),
MockCmd::new(
redis::cmd("EVAL")
.arg("return KEYS[1]")
.arg(1usize)
.arg("team:key"),
Ok("team:key"),
),
MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")),
MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")),
MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK")),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
.with_namespace(Some("team".into()));
assert_eq!(
cache
.batch_get_counts(&["count".into(), "missing".into()])
.unwrap(),
[Some(7), None]
);
assert_eq!(
cache
.async_batch_get_counts(vec!["count".into(), "missing".into()])
.await
.unwrap(),
[Some(7), None]
);
assert!(cache.sync_ping().unwrap());
assert!(cache.ping().await.unwrap());
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
assert_eq!(
cache.async_scan_iter("job-", 25).await.unwrap(),
["team:job-a", "team:job-b"]
);
assert_eq!(
cache
.delete_cache_keys(vec!["job-a".into(), "job-b".into()])
.await
.unwrap(),
2
);
assert_eq!(
cache
.async_set_cache_sadd("members", vec!["a".into(), "b".into()], None)
.await
.unwrap(),
2
);
assert_eq!(
cache
.async_rpush("queue", vec!["a".into(), "b".into()])
.await
.unwrap(),
2
);
assert_eq!(
cache.async_lpop("queue", Some(2)).await.unwrap(),
RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()])
);
assert_eq!(
cache
.async_eval("return KEYS[1]".into(), vec!["key".into()], Vec::new())
.await
.unwrap(),
redis::Value::BulkString(b"team:key".to_vec())
);
assert_eq!(
cache
.async_register_script("return KEYS[1]".into())
.invoke(vec!["key".into()], Vec::new())
.await
.unwrap(),
redis::Value::BulkString(b"team:key".to_vec())
);
assert_eq!(cache.client_list().unwrap(), "id=1");
assert_eq!(cache.info().unwrap(), "redis_version:7");
cache.flushall().unwrap();
}
#[tokio::test]
async fn direct_redis_pipelines_preserve_operation_order() {
let mut rpush_pipeline = redis::pipe();
rpush_pipeline
.cmd("RPUSH")
.arg("team:a")
.arg("one")
.cmd("RPUSH")
.arg("team:b")
.arg("two");
let mut lpop_pipeline = redis::pipe();
lpop_pipeline
.cmd("LPOP")
.arg("team:a")
.arg(2usize)
.cmd("LPOP")
.arg("team:b");
let connection = MockRedisConnection::new([
MockCmd::with_values(
rpush_pipeline,
Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]),
),
MockCmd::with_values(
lpop_pipeline,
Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]),
),
])
.assert_all_commands_consumed();
let queue = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
.with_namespace(Some("team".into()));
assert_eq!(
queue
.async_rpush_pipeline(vec![
RedisRpushOperation {
key: "a".into(),
values: vec![RedisArg::from("one")],
},
RedisRpushOperation {
key: "b".into(),
values: vec![RedisArg::from("two")],
},
])
.await
.unwrap(),
[1, 2]
);
assert_eq!(
queue
.async_lpop_pipeline(vec![
RedisLpopOperation {
key: "a".into(),
count: Some(2),
},
RedisLpopOperation {
key: "b".into(),
count: None,
},
])
.await
.unwrap(),
[
RedisLpopResult::Values(vec![b"one".to_vec()]),
RedisLpopResult::Missing,
]
);
let mut increment_pipeline = redis::pipe();
increment_pipeline
.cmd("INCRBYFLOAT")
.arg("team:counter")
.arg(1.5f64)
.cmd("EXPIRE")
.arg("team:counter")
.arg(10u64)
.ignore()
.cmd("INCRBYFLOAT")
.arg("team:counter")
.arg(2.0f64);
let connection = MockRedisConnection::new([MockCmd::with_values(
increment_pipeline,
Ok(vec![
redis::Value::BulkString(b"1.5".to_vec()),
redis::Value::Int(1),
redis::Value::BulkString(b"3.5".to_vec()),
]),
)])
.assert_all_commands_consumed();
let counters = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new())
.with_namespace(Some("team".into()));
assert_eq!(
counters
.async_increment_pipeline(vec![
IncrementOperation {
key: "counter".into(),
amount: 1.5,
ttl: Some(Duration::from_secs(10)),
},
IncrementOperation {
key: "counter".into(),
amount: 2.0,
ttl: None,
},
])
.await
.unwrap(),
[1.5, 3.5]
);
}
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return count"
);
const SET_MAX_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]); ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return ARGV[1]; end; return current"
);
#[tokio::test]
async fn counter_repairs_are_atomic_and_use_default_ttl() {
let floor = || {
redis::cmd("EVAL")
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
.arg(1)
.arg("team:counter")
.arg(-2i64)
.arg(30u64)
.clone()
};
let connection = MockRedisConnection::new([
MockCmd::new(floor(), Ok(0i64)),
MockCmd::new(floor(), Ok(0i64)),
MockCmd::new(
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)
.arg("team:counter")
.arg(4.5f64)
.arg(600u64),
Ok("4.5"),
),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new())
.with_namespace(Some("team".into()));
assert_eq!(
cache
.increment_with_floor("counter", -2, Duration::from_secs(30))
.unwrap(),
0
);
assert_eq!(
cache
.async_increment_with_floor("counter", -2, Duration::from_secs(30))
.await
.unwrap(),
0
);
assert_eq!(
cache.async_set_max("counter", 4.5, None).await.unwrap(),
4.5
);
}
const CLAIM_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
"elseif current ~= ARGV[1] then return 0; end; ",
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
);
fn claim_eval(expected: &str, write: &str, refresh: bool) -> redis::Cmd {
let mut cmd = redis::cmd("EVAL");
cmd.arg(CLAIM_SCRIPT)
.arg(1)
.arg("pin")
.arg(expected)
.arg(600)
.arg(write)
.arg(u8::from(refresh));
cmd
}
#[tokio::test]
async fn claims_match_eligible_values_written_by_another_encoder() {
let python_payload = r#"{"model_id": "a", "deployment": "east"}"#;
let stored = serde_json::json!({"deployment": "east", "model_id": "a"});
let candidate = serde_json::json!({"model_id": "b"});
let connection = MockRedisConnection::new([
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)),
MockCmd::new(claim_eval(python_payload, "", true), Ok(1)),
])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
assert_eq!(
cache
.async_claim_cache(
"pin",
candidate,
vec![stored.clone()],
ExactCacheContext::default()
)
.await
.unwrap(),
stored
);
}
#[test]
fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() {
let candidate = serde_json::json!({"model_id": "b"});
let payload = r#"{"model_id":"b"}"#;
let connection = MockRedisConnection::new([
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)),
MockCmd::new(claim_eval("", payload, false), Ok(0)),
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)),
MockCmd::new(claim_eval(r#"{"model_id":"gone"}"#, payload, false), Ok(1)),
])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
assert_eq!(
cache
.claim_cache(
"pin",
candidate.clone(),
&[serde_json::json!({"model_id": "a"})],
ExactCacheContext::default()
)
.unwrap(),
candidate
);
}
#[test]
fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() {
let stored = r#"{"model_id": "a"}"#;
let connection = MockRedisConnection::new([
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)),
MockCmd::new(claim_eval(stored, "", false), Ok(1)),
])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
assert_eq!(
cache
.claim_cache(
"pin",
serde_json::json!({"model_id": "b"}),
&[],
ExactCacheContext::default()
)
.unwrap(),
serde_json::json!({"model_id": "a"})
);
}
#[tokio::test]
async fn async_increment_runs_the_atomic_script() {
let mut eval = redis::cmd("EVAL");
eval.arg(concat!(
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
"if redis.call('TTL', KEYS[1]) == -1 then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
))
.arg(1)
.arg("counter")
.arg(2.5f64)
.arg(600);
let connection =
MockRedisConnection::new([MockCmd::new(eval, Ok("4.5"))]).assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new());
assert_eq!(
cache
.async_increment("counter", 2.5, ExactCacheContext::default())
.await
.unwrap(),
4.5
);
}

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

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-response"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
py_literal = "0.4.0"
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
[dev-dependencies]
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
redis = "1.7.0"
redis-test = "1.0.4"
tokio.workspace = true

61
litellm-rust/crates/cache-response/README.md Normal file
View file

@ -0,0 +1,61 @@
# Response cache foundation
`ResponseCache<B>` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache<Value = CacheEntry>`
## Ownership
`litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations
`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python
The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum, which only dispatches. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host
## Native Rust use
```rust
use std::{sync::Arc, time::Duration};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{CacheKeyInput, ResponseCache, ResponseCacheRequest};
use serde_json::json;
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let request = ResponseCacheRequest::new(CacheKeyInput {
preset: Some("example:key".into()),
..Default::default()
});
let now = Duration::from_secs(100);
cache.store(&request, json!({"answer": 7}), now)?;
assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7})));
```
For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers, including counters and claims, move that blocking work off the executor. The pool skips the checkout PING and instead discards any connection whose command failed
Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it
## Python integration boundary
The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API
Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution
Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend
The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings, including `litellm.default_redis_ttl` and SSL options. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python
Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy
The Redis backend also provides the primitives needed to preserve its direct Python surface later: TLS URLs, ping, bulk delete, counter batches, TTL, scan, set membership, raw queue push and pop, queue and counter pipelines, counter floor and maximum operations, script evaluation, client information, namespaced flush, and full flush. These are backend operations only and are not exported to Python by this PR. Memory provides TTL, oldest-key, and counter-pipeline operations
## Adding another backend
Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python
Verify typed values, TTL precedence, missing entries, serialization failures, namespaces, batch ordering, and sync/async behavior. Run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before enabling a public facade
## Follow-up scope
Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths
Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees

45
litellm-rust/crates/cache-response/src/buffer.rs Normal file
View file

@ -0,0 +1,45 @@
use std::{sync::Mutex, time::Duration};
use litellm_cache::{BaseCache, Error, ExactCacheContext};
use serde_json::Value;
use crate::{CacheEntry, ResponseCache, ResponseCacheRequest};
pub struct WriteBuffer {
flush_size: usize,
entries: Mutex<Vec<(ResponseCacheRequest, Value, Duration)>>,
}
impl WriteBuffer {
pub fn new(flush_size: usize) -> Self {
Self {
flush_size: flush_size.max(1),
entries: Mutex::new(Vec::new()),
}
}
pub async fn async_store<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>>(
&self,
cache: &ResponseCache<B>,
request: &ResponseCacheRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
let pending = {
let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?;
entries.push((request.clone(), response, now));
(entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries))
};
// A failed flush drops its batch, as Python does. Requeueing would grow the
// buffer and re-send an ever larger pipeline on every write during an outage.
match pending {
Some(pending) => cache.async_store_entries(pending).await,
None => Ok(()),
}
}
pub fn clear(&self) -> Result<(), Error> {
self.entries.lock().map_err(|_| Error::Unavailable)?.clear();
Ok(())
}
}

147
litellm-rust/crates/cache-response/src/caching.rs Normal file
View file

@ -0,0 +1,147 @@
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
pub enum CacheMode {
#[default]
#[serde(rename = "default_on")]
DefaultOn,
#[serde(rename = "default_off")]
DefaultOff,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct CacheKeyField {
pub name: String,
pub value: Option<String>,
pub api_parameter: bool,
pub internal_parameter: bool,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
#[serde(default)]
pub struct CacheKeyInput {
pub fields: Vec<CacheKeyField>,
pub preset: Option<String>,
pub namespace: Option<String>,
pub include_provider_parameters: bool,
}
#[derive(Default)]
pub struct CacheKeyContext {
pub model_group: Option<String>,
pub caching_groups: Vec<(Vec<String>, String)>,
pub file_checksum: Option<String>,
pub file_object_name: Option<String>,
pub metadata_file_name: Option<String>,
pub parameters_file_name: Option<String>,
}
impl CacheKeyContext {
pub fn apply(self, input: &mut CacheKeyInput) {
let group = self.model_group.as_ref().and_then(|model| {
self.caching_groups
.iter()
.find(|(models, _)| models.contains(model))
});
for field in &mut input.fields {
match field.name.as_str() {
"model" => {
field.value = group
.map(|(_, formatted)| formatted.clone())
.or_else(|| self.model_group.clone())
.or_else(|| field.value.take())
}
"file" => {
field.value = self
.file_checksum
.clone()
.or_else(|| self.file_object_name.clone())
.or_else(|| self.metadata_file_name.clone())
.or_else(|| self.parameters_file_name.clone())
}
_ => {}
}
}
}
}
pub fn get_cache_key(input: &CacheKeyInput) -> String {
cache_key(input)
}
pub fn cache_key(input: &CacheKeyInput) -> String {
if let Some(preset) = &input.preset {
return preset.clone();
}
let mut digest = Sha256::new();
for field in &input.fields {
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
&& let Some(value) = &field.value
{
digest.update(field.name.as_bytes());
digest.update(b": ");
digest.update(value.as_bytes());
}
}
let hash = format!("{:x}", digest.finalize());
input
.namespace
.as_deref()
.filter(|namespace| !namespace.is_empty())
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
pub struct CacheControls {
pub supported_call_type: bool,
pub configured: bool,
pub native_backend: bool,
pub default_on: bool,
pub caching: Option<bool>,
pub no_cache: bool,
pub no_store: bool,
#[serde(default)]
pub use_cache: bool,
}
impl CacheControls {
pub fn reads(self) -> bool {
self.supported_call_type
&& self.configured
&& self.caching.unwrap_or(true)
&& !self.no_cache
&& (self.default_on || self.use_cache)
}
pub fn writes(self) -> bool {
self.supported_call_type
&& self.configured
&& self.caching.unwrap_or(true)
&& !self.no_store
&& (self.default_on || self.use_cache)
}
}
pub fn should_use_cache(controls: CacheControls) -> bool {
controls.reads() || controls.writes()
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CacheEntry {
#[serde(skip_serializing_if = "Option::is_none")]
pub timestamp: Option<f64>,
pub response: Value,
}
impl CacheEntry {
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
self.timestamp.is_none_or(|timestamp| {
timestamp.is_finite()
&& max_age.is_none_or(|age| now.as_secs_f64() - timestamp <= age.as_secs_f64())
})
}
}

129
litellm-rust/crates/cache-response/src/codec.rs Normal file
View file

@ -0,0 +1,129 @@
use litellm_cache::{CacheCodec, Error};
use serde_json::Value;
use crate::CacheEntry;
#[derive(Clone, Copy, Debug, Default)]
pub struct ResponseCacheCodec;
impl CacheCodec for ResponseCacheCodec {
type Value = CacheEntry;
fn encode(&self, value: &CacheEntry) -> Result<Vec<u8>, Error> {
if value
.timestamp
.is_some_and(|timestamp| !timestamp.is_finite())
{
return Err(Error::InvalidEntry);
}
// Python reads a `response` that is either a dict or a serialized string, so every
// other shape is written serialized. A string on the wire is therefore always a
// serialized response, which keeps string-valued responses unambiguous.
if value.timestamp.is_none() || value.response.is_object() {
return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry);
}
let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?;
serde_json::to_vec(&CacheEntry {
timestamp: value.timestamp,
response: Value::String(response),
})
.map_err(|_| Error::InvalidEntry)
}
fn decode(&self, bytes: &[u8]) -> Result<CacheEntry, Error> {
let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?;
let value = decode_value(text)?;
let Some(timestamp) = value.get("timestamp") else {
return Ok(CacheEntry {
timestamp: None,
response: value,
});
};
let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else {
return Err(Error::InvalidEntry);
};
let response = match value.get("response").ok_or(Error::InvalidEntry)? {
Value::String(text) => decode_value(text)?,
response => response.clone(),
};
Ok(CacheEntry {
timestamp: Some(timestamp),
response,
})
}
}
fn decode_value(text: &str) -> Result<Value, Error> {
if let Ok(value) = serde_json::from_str(text) {
return Ok(value);
}
check_literal_depth(text)?;
let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?;
literal_value(literal, 0)
}
fn literal_value(value: py_literal::Value, depth: usize) -> Result<Value, Error> {
use py_literal::Value as Literal;
if depth > 128 {
return Err(Error::InvalidEntry);
}
match value {
Literal::String(text) => Ok(Value::String(text)),
Literal::Boolean(value) => Ok(Value::Bool(value)),
Literal::None => Ok(Value::Null),
Literal::Integer(value) => {
serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry)
}
Literal::Float(value) => serde_json::Number::from_f64(value)
.map(Value::Number)
.ok_or(Error::InvalidEntry),
Literal::List(values) | Literal::Tuple(values) => values
.into_iter()
.map(|value| literal_value(value, depth + 1))
.collect::<Result<Vec<_>, _>>()
.map(Value::Array),
Literal::Dict(entries) => entries
.into_iter()
.map(|(key, value)| {
let Literal::String(key) = key else {
return Err(Error::InvalidEntry);
};
Ok((key, literal_value(value, depth + 1)?))
})
.collect::<Result<serde_json::Map<_, _>, _>>()
.map(Value::Object),
_ => Err(Error::InvalidEntry),
}
}
fn check_literal_depth(text: &str) -> Result<(), Error> {
let mut quote = None;
let mut escaped = false;
let mut depth = 0usize;
for ch in text.chars() {
if escaped {
escaped = false;
continue;
}
if let Some(delimiter) = quote {
if ch == '\\' {
escaped = true;
} else if ch == delimiter {
quote = None;
}
continue;
}
match ch {
'\'' | '"' => quote = Some(ch),
'[' | '{' | '(' => {
depth += 1;
if depth > 128 {
return Err(Error::InvalidEntry);
}
}
']' | '}' | ')' => depth = depth.saturating_sub(1),
_ => {}
}
}
Ok(())
}

22
litellm-rust/crates/cache-response/src/embedding.rs Normal file
View file

@ -0,0 +1,22 @@
use serde::Serialize;
use serde_json::Value;
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct PartialHits {
pub values: Vec<Option<Value>>,
pub missing_indices: Vec<usize>,
}
impl PartialHits {
pub fn new(values: Vec<Option<Value>>) -> Self {
let missing_indices = values
.iter()
.enumerate()
.filter_map(|(index, value)| value.is_none().then_some(index))
.collect();
Self {
values,
missing_indices,
}
}
}

14
litellm-rust/crates/cache-response/src/lib.rs Normal file
View file

@ -0,0 +1,14 @@
mod buffer;
mod caching;
mod codec;
mod embedding;
mod response;
pub use buffer::WriteBuffer;
pub use caching::{
CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key,
get_cache_key, should_use_cache,
};
pub use codec::ResponseCacheCodec;
pub use embedding::PartialHits;
pub use response::{ResponseCache, ResponseCacheRequest};

280
litellm-rust/crates/cache-response/src/response.rs Normal file
View file

@ -0,0 +1,280 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
};
use serde_json::Value;
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
#[derive(Clone)]
pub struct ResponseCacheRequest {
pub key: CacheKeyInput,
pub controls: CacheControls,
pub context: ExactCacheContext,
pub max_age: Option<Duration>,
}
impl ResponseCacheRequest {
pub fn new(key: CacheKeyInput) -> Self {
Self {
key,
controls: CacheControls {
configured: true,
supported_call_type: true,
native_backend: true,
default_on: true,
..Default::default()
},
context: ExactCacheContext::default(),
max_age: None,
}
}
}
pub struct ResponseCache<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
backend: Arc<B>,
}
impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
}
pub fn backend(&self) -> &B {
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration> {
self.backend.get_ttl(&ExactCacheContext::default())
}
pub async fn async_flush(&self) -> Result<(), Error>
where
B: FlushCache,
{
self.backend.async_flush_cache().await
}
pub async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
self.backend.test_connection().await
}
pub fn lookup(
&self,
request: &ResponseCacheRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
return Ok(None);
}
let entry = match self
.backend
.get_cache(&cache_key(&request.key), &request.context)
{
Ok(entry) => entry,
Err(Error::InvalidEntry) => None,
Err(error) => return Err(error),
};
Ok(Self::fresh_or_miss(entry, now, request.max_age))
}
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
return Ok(None);
}
let entry = match self
.backend
.async_get_cache(&cache_key(&request.key), &request.context)
.await
{
Ok(entry) => entry,
Err(Error::InvalidEntry) => None,
Err(error) => return Err(error),
};
Ok(Self::fresh_or_miss(entry, now, request.max_age))
}
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
now: Duration,
) -> Result<PartialHits, Error>
where
B: BatchCache,
{
let readable = requests
.iter()
.enumerate()
.filter(|(_, request)| request.controls.reads())
.collect::<Vec<_>>();
let keys = readable
.iter()
.map(|(_, request)| cache_key(&request.key))
.collect::<Vec<_>>();
let entries = if let Some((_, request)) = readable.first() {
self.backend.batch_get_cache(&keys, &request.context)?
} else {
Vec::new()
};
Self::partial_hits(requests, readable, entries, now)
}
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
now: Duration,
) -> Result<PartialHits, Error>
where
B: BatchCache,
{
let readable = requests
.iter()
.enumerate()
.filter(|(_, request)| request.controls.reads())
.collect::<Vec<_>>();
let keys = readable
.iter()
.map(|(_, request)| cache_key(&request.key))
.collect::<Vec<_>>();
let entries = if let Some((_, request)) = readable.first() {
self.backend
.async_batch_get_cache(keys, request.context.clone())
.await?
} else {
Vec::new()
};
Self::partial_hits(requests, readable, entries, now)
}
pub fn store(
&self,
request: &ResponseCacheRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
if !request.controls.writes() {
return Ok(());
}
self.backend.set_cache(
&cache_key(&request.key),
CacheEntry {
timestamp: Some(now.as_secs_f64()),
response,
},
&request.context,
)
}
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
if !request.controls.writes() {
return Ok(());
}
self.backend
.async_set_cache(
&cache_key(&request.key),
CacheEntry {
timestamp: Some(now.as_secs_f64()),
response,
},
request.context.clone(),
)
.await
}
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
now: Duration,
) -> Result<(), Error> {
self.async_store_entries(
entries
.into_iter()
.map(|(request, response)| (request, response, now))
.collect(),
)
.await
}
/// Stores entries that each carry the time they were produced, so a deferred write keeps
/// the freshness of its original response.
pub async fn async_store_entries(
&self,
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
) -> Result<(), Error> {
let writable = entries
.into_iter()
.filter(|(request, _, _)| request.controls.writes())
.map(|(request, response, now)| {
(
cache_key(&request.key),
CacheEntry {
timestamp: Some(now.as_secs_f64()),
response,
},
request.context,
)
})
.collect::<Vec<_>>();
let Some((_, _, first_kwargs)) = writable.first() else {
return Ok(());
};
if writable
.iter()
.all(|(_, _, context)| context == first_kwargs)
{
let context = first_kwargs.clone();
let cache_list = writable
.into_iter()
.map(|(key, entry, _)| (key, entry))
.collect();
return self
.backend
.async_set_cache_pipeline(cache_list, context)
.await;
}
for (key, entry, context) in writable {
self.backend.async_set_cache(&key, entry, context).await?;
}
Ok(())
}
fn partial_hits(
requests: &[ResponseCacheRequest],
readable: Vec<(usize, &ResponseCacheRequest)>,
entries: Vec<BatchEntry<CacheEntry>>,
now: Duration,
) -> Result<PartialHits, Error> {
if readable.len() != entries.len() {
return Err(Error::Unavailable);
}
let mut values = vec![None; requests.len()];
for ((index, request), entry) in readable.into_iter().zip(entries) {
let response = match entry {
BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age),
BatchEntry::Miss | BatchEntry::Invalid => None,
};
values[index] = response;
}
Ok(PartialHits::new(values))
}
fn fresh_or_miss(
entry: Option<CacheEntry>,
now: Duration,
max_age: Option<Duration>,
) -> Option<Value> {
entry
.filter(|entry| entry.fresh(now, max_age))
.map(|entry| entry.response)
}
}

90
litellm-rust/crates/cache-response/tests/caching.rs Normal file
View file

@ -0,0 +1,90 @@
use litellm_cache_response::{
CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key,
};
use sha2::{Digest, Sha256};
#[test]
fn keys_match_python_order_groups_files_presets_and_namespaces() {
let mut input = CacheKeyInput {
fields: vec![
CacheKeyField {
name: "model".into(),
value: Some("deployment".into()),
api_parameter: true,
internal_parameter: false,
},
CacheKeyField {
name: "file".into(),
value: None,
api_parameter: true,
internal_parameter: false,
},
],
namespace: Some("team".into()),
..Default::default()
};
CacheKeyContext {
model_group: Some("group".into()),
caching_groups: vec![(vec!["group".into()], "['group']".into())],
file_checksum: Some("checksum".into()),
..Default::default()
}
.apply(&mut input);
assert_eq!(
cache_key(&input),
format!(
"team:{:x}",
Sha256::digest(b"model: ['group']file: checksum")
)
);
input.preset = Some("preset".into());
assert_eq!(get_cache_key(&input), "preset");
}
#[test]
fn cache_controls_honor_default_modes_and_directives() {
let enabled = CacheControls {
supported_call_type: true,
configured: true,
default_on: true,
..Default::default()
};
assert!(enabled.reads());
assert!(enabled.writes());
assert!(
!CacheControls {
default_on: false,
..enabled
}
.reads()
);
assert!(
CacheControls {
default_on: false,
use_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_store: true,
..enabled
}
.writes()
);
assert!(
!CacheControls {
caching: Some(false),
..enabled
}
.writes()
);
}

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

@ -0,0 +1,484 @@
use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use litellm_cache::{BaseCache, CacheCodec, Error};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::RedisCache;
use litellm_cache_response::{
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, WriteBuffer,
};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
Some(8),
Some(Duration::from_secs(600)),
))))
}
fn request() -> ResponseCacheRequest {
ResponseCacheRequest::new(CacheKeyInput {
preset: Some("tenant:key".into()),
..Default::default()
})
}
#[tokio::test]
async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
let clock = Arc::new(AtomicU64::new(100));
let backend = Arc::new(InMemoryCache::with_clock(
Some(8),
Some(Duration::from_secs(600)),
{
let clock = clock.clone();
move || Duration::from_secs(clock.load(Ordering::SeqCst))
},
));
let cache = ResponseCache::new(backend.clone());
let mut request = request();
request.context.ttl = Some(Duration::from_secs(10));
request.max_age = Some(Duration::from_secs(5));
cache
.store(
&request,
json!({"choices": [1], "usage": {"total_tokens": 7}}),
Duration::from_secs(100),
)
.unwrap();
assert_eq!(
backend.expires_at("tenant:key").unwrap(),
Some(Duration::from_secs(110))
);
assert!(
cache
.async_lookup(&request, Duration::from_secs(105))
.await
.unwrap()
.is_some()
);
assert_eq!(
cache.lookup(&request, Duration::from_secs(106)).unwrap(),
None
);
request.max_age = None;
assert_eq!(
cache
.lookup(&request, Duration::from_secs(106))
.unwrap()
.unwrap()["usage"]["total_tokens"],
7
);
clock.store(111, Ordering::SeqCst);
assert_eq!(
cache
.async_lookup(&request, Duration::from_secs(111))
.await
.unwrap(),
None
);
cache
.async_store(&request, json!({"choices": [2]}), Duration::from_secs(111))
.await
.unwrap();
assert_eq!(
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
Some(json!({"choices": [2]}))
);
}
#[tokio::test]
async fn directives_skip_io_and_keep_reads_and_writes_independent() {
let cache = memory();
let mut request = request();
let now = Duration::from_secs(100);
request.controls.no_store = true;
cache
.async_store(&request, json!({"v": 1}), now)
.await
.unwrap();
assert_eq!(cache.lookup(&request, now).unwrap(), None);
request.controls.no_store = false;
request.controls.no_cache = true;
cache.store(&request, json!({"v": 2}), now).unwrap();
assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None);
request.controls.no_cache = false;
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
request.controls.default_on = false;
cache.store(&request, json!({"v": 3}), now).unwrap();
assert_eq!(cache.lookup(&request, now).unwrap(), None);
request.controls.use_cache = true;
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
request.controls.supported_call_type = false;
assert_eq!(cache.lookup(&request, now).unwrap(), None);
}
#[tokio::test]
async fn redis_consumer_reads_python_sync_and_async_envelopes_and_writes_compatible_json() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()),
),
MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()),
),
MockCmd::new(
redis::cmd("SETEX")
.arg("tenant:key")
.arg(600)
.arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()),
Ok("OK"),
),
])
.assert_all_commands_consumed();
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec)
.with_namespace(Some("tenant".into()));
let cache = ResponseCache::new(Arc::new(backend));
let request = request();
let expected = json!({"ok": true, "text": "cached"});
assert_eq!(
cache.lookup(&request, Duration::from_secs(101)).unwrap(),
Some(expected.clone())
);
assert_eq!(
cache
.async_lookup(&request, Duration::from_secs(101))
.await
.unwrap(),
Some(expected.clone())
);
cache
.async_store(&request, expected, Duration::from_secs(100))
.await
.unwrap();
}
#[tokio::test]
async fn captured_service_keeps_the_selected_backend_for_background_writes() {
let original = memory();
let captured = original.clone();
let replacement = memory();
let request = request();
let writer = tokio::spawn({
let request = request.clone();
async move {
captured
.async_store(
&request,
json!({"selected": "original"}),
Duration::from_secs(100),
)
.await
}
});
writer.await.unwrap().unwrap();
assert_eq!(
original.lookup(&request, Duration::from_secs(100)).unwrap(),
Some(json!({"selected":"original"}))
);
assert_eq!(
replacement
.lookup(&request, Duration::from_secs(100))
.unwrap(),
None
);
}
#[test]
fn generated_keys_preserve_namespace_and_explicit_keys() {
let cache = memory();
let key = CacheKeyInput {
fields: vec![CacheKeyField {
name: "model".into(),
value: Some("a".into()),
api_parameter: true,
internal_parameter: false,
}],
namespace: Some("tenant".into()),
..Default::default()
};
let generated = ResponseCacheRequest::new(key.clone());
let explicit = ResponseCacheRequest::new(CacheKeyInput {
preset: Some(litellm_cache_response::cache_key(&key)),
..Default::default()
});
cache
.store(&generated, json!({"value": 7}), Duration::from_secs(100))
.unwrap();
assert_eq!(
cache.lookup(&explicit, Duration::from_secs(100)).unwrap(),
Some(json!({"value":7}))
);
}
#[test]
fn response_codec_accepts_python_literals_without_executing_code() {
let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#;
let entry = ResponseCacheCodec.decode(bytes).unwrap();
assert_eq!(
entry.response,
json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]})
);
for bytes in [
b"__import__('os').system('false')".as_slice(),
b"{'timestamp': 'invalid', 'response': {}}",
b"{'timestamp': 1e9999, 'response': {}}",
] {
assert_eq!(
ResponseCacheCodec.decode(bytes).unwrap_err(),
Error::InvalidEntry
);
}
let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000));
assert_eq!(
ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(),
Error::InvalidEntry
);
assert_eq!(
ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(f64::NAN),
response: json!({})
})
.unwrap_err(),
Error::InvalidEntry
);
}
#[tokio::test]
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() {
let connection = MockRedisConnection::new([MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(b"invalid".to_vec()),
)])
.assert_all_commands_consumed();
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec);
let cache = ResponseCache::new(Arc::new(backend));
let mut request = request();
request.controls.no_cache = true;
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
request.controls.no_cache = false;
assert_eq!(
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
None
);
}
#[test]
fn string_responses_round_trip_through_typed_and_wire_backends() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let now = Duration::from_secs(100);
for response in [json!("hello world"), json!("123"), json!("null")] {
cache.store(&request(), response.clone(), now).unwrap();
assert_eq!(
cache.lookup(&request(), now).unwrap(),
Some(response.clone())
);
let wire = ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(100.0),
response: response.clone(),
})
.unwrap();
assert_eq!(ResponseCacheCodec.decode(&wire).unwrap().response, response);
}
}
#[test]
fn non_object_responses_are_written_as_python_readable_serialized_strings() {
let wire = ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(100.0),
response: json!([1, 2]),
})
.unwrap();
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&wire).unwrap(),
json!({"timestamp": 100.0, "response": "[1,2]"})
);
assert_eq!(
ResponseCacheCodec.decode(&wire).unwrap().response,
json!([1, 2])
);
assert_eq!(
ResponseCacheCodec.decode(br#"{"timestamp": 100.0, "response": "not serialized"}"#),
Err(Error::InvalidEntry)
);
}
#[test]
fn response_entries_preserve_the_existing_json_representation() {
let codec = ResponseCacheCodec;
let entry = CacheEntry {
timestamp: Some(123.0),
response: json!({"choices": [{"text": "cached"}]}),
};
let bytes = codec.encode(&entry).unwrap();
assert_eq!(bytes, serde_json::to_vec(&entry).unwrap());
assert_eq!(codec.decode(&bytes).unwrap(), entry);
}
#[test]
fn response_codec_preserves_values_without_timestamps() {
let codec = ResponseCacheCodec;
let raw = json!({"choices": [{"text": "legacy"}]});
let entry = codec.decode(&serde_json::to_vec(&raw).unwrap()).unwrap();
assert_eq!(entry.timestamp, None);
assert_eq!(entry.response, raw);
let backend = Arc::new(InMemoryCache::default());
BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, &Default::default()).unwrap();
let cache = ResponseCache::new(backend);
assert_eq!(
cache.lookup(&request(), Duration::from_secs(100)).unwrap(),
Some(json!({"choices": [{"text": "legacy"}]}))
);
}
#[tokio::test]
async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() {
let cache = memory();
let requests = ["hit", "miss", "disabled"].map(|key| {
ResponseCacheRequest::new(CacheKeyInput {
preset: Some(key.into()),
..Default::default()
})
});
cache
.store(&requests[0], json!({"value": 1}), Duration::from_secs(100))
.unwrap();
let mut requests = requests.to_vec();
requests[2].controls.caching = Some(false);
let partial = cache
.async_lookup_batch(&requests, Duration::from_secs(100))
.await
.unwrap();
assert_eq!(partial.values, vec![Some(json!({"value": 1})), None, None]);
assert_eq!(partial.missing_indices, vec![1, 2]);
cache
.async_store_batch(
vec![
(requests[1].clone(), json!({"value": 2})),
(requests[2].clone(), json!({"value": 3})),
],
Duration::from_secs(100),
)
.await
.unwrap();
assert_eq!(
cache
.lookup(&requests[1], Duration::from_secs(100))
.unwrap(),
Some(json!({"value": 2}))
);
requests[2].controls.caching = None;
assert_eq!(
cache
.lookup(&requests[2], Duration::from_secs(100))
.unwrap(),
None
);
}
#[tokio::test]
async fn deferred_entries_keep_the_time_they_were_produced() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let mut request = request();
request.max_age = Some(Duration::from_secs(10));
cache
.async_store_entries(vec![(
request.clone(),
json!({"answer": 7}),
Duration::from_secs(100),
)])
.await
.unwrap();
assert_eq!(
cache.lookup(&request, Duration::from_secs(110)).unwrap(),
Some(json!({"answer": 7}))
);
assert_eq!(
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
None
);
}
#[tokio::test]
async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let buffer = WriteBuffer::new(2);
let mut first = request();
first.max_age = Some(Duration::from_secs(10));
let mut second = request();
second.key.preset = Some("tenant:other".into());
buffer
.async_store(
&cache,
&first,
json!({"answer": 7}),
Duration::from_secs(100),
)
.await
.unwrap();
assert_eq!(
cache.lookup(&first, Duration::from_secs(100)).unwrap(),
None
);
buffer
.async_store(
&cache,
&second,
json!({"answer": 8}),
Duration::from_secs(200),
)
.await
.unwrap();
assert_eq!(
cache.lookup(&first, Duration::from_secs(110)).unwrap(),
Some(json!({"answer": 7}))
);
assert_eq!(
cache.lookup(&first, Duration::from_secs(111)).unwrap(),
None
);
assert_eq!(
cache.lookup(&second, Duration::from_secs(200)).unwrap(),
Some(json!({"answer": 8}))
);
}
#[tokio::test]
async fn write_buffer_clear_drops_pending_entries() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let buffer = WriteBuffer::new(2);
let mut other = request();
other.key.preset = Some("tenant:other".into());
let now = Duration::from_secs(100);
buffer
.async_store(&cache, &request(), json!({"answer": 7}), now)
.await
.unwrap();
buffer.clear().unwrap();
buffer
.async_store(&cache, &other, json!({"answer": 8}), now)
.await
.unwrap();
assert_eq!(cache.lookup(&request(), now).unwrap(), None);
assert_eq!(cache.lookup(&other, now).unwrap(), None);
}

2
litellm-rust/crates/cache/Cargo.toml vendored
View file

@ -8,8 +8,8 @@ repository.workspace = true
[dependencies]
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

121
litellm-rust/crates/cache/src/base_cache.rs vendored
View file

@ -1,18 +1,35 @@
use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
use std::{future::Future, time::Duration};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::Error;
pub type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
#[derive(Clone, Debug, PartialEq)]
pub enum BatchEntry<V> {
Hit(V),
Miss,
Invalid,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct CacheKwargs {
pub trait CacheContext: Clone + Send + Sync + 'static {
fn ttl(&self) -> Option<Duration>;
fn with_ttl(&self, ttl: Option<Duration>) -> Self;
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ExactCacheContext {
pub ttl: Option<Duration>,
pub extras: Map<String, Value>,
}
impl CacheContext for ExactCacheContext {
fn ttl(&self) -> Option<Duration> {
self.ttl
}
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
Self { ttl }
}
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
@ -32,67 +49,59 @@ pub struct CacheConnectionResult {
pub trait BaseCache: Send + Sync {
type Value: Clone + Send + Sync + 'static;
type Context: CacheContext;
fn default_ttl(&self) -> Duration {
Duration::from_secs(60)
}
fn get_ttl(&self, context: &Self::Context) -> Option<Duration>;
fn get_ttl(&self, kwargs: &CacheKwargs) -> Duration {
kwargs.ttl.unwrap_or_else(|| self.default_ttl())
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error>;
fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<Self::Value>, Error>;
fn async_set_cache<'a>(
&'a self,
key: &'a str,
fn set_cache(
&self,
key: &str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
Box::pin(async move { self.set_cache(key, value, kwargs) })
context: &Self::Context,
) -> Result<(), Error>;
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error>;
fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> impl Future<Output = Result<(), Error>> + Send {
async move { self.set_cache(key, value, &context) }
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
kwargs: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
Box::pin(async move { self.get_cache(key, kwargs) })
fn async_get_cache(
&self,
key: &str,
context: &Self::Context,
) -> impl Future<Output = Result<Option<Self::Value>, Error>> + Send {
async move { self.get_cache(key, context) }
}
fn async_set_cache_pipeline<'a>(
&'a self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
Box::pin(async move {
for (key, value) in cache_list {
self.set_cache(&key, value, kwargs.clone())?;
fn async_set_cache_pipeline(
&self,
entries: Vec<(String, Self::Value)>,
context: Self::Context,
) -> impl Future<Output = Result<(), Error>> + Send {
async move {
for (key, value) in entries {
self.async_set_cache(&key, value, context.clone()).await?;
}
Ok(())
})
}
}
fn batch_cache_write<'a>(
&'a self,
key: &'a str,
fn batch_cache_write(
&self,
key: &str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
self.async_set_cache(key, value, kwargs)
context: Self::Context,
) -> impl Future<Output = Result<(), Error>> + Send {
self.async_set_cache(key, value, context)
}
fn delete_cache(&self, key: &str) -> Result<(), Error>;
fn disconnect(&self) -> impl Future<Output = Result<(), Error>> + Send;
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
Box::pin(async move { self.delete_cache(key) })
}
fn flush_cache(&self) -> Result<(), Error>;
fn disconnect(&self) -> CacheFuture<'_, ()>;
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult>;
fn test_connection(&self) -> impl Future<Output = Result<CacheConnectionResult, Error>> + Send;
}

85
litellm-rust/crates/cache/src/cache_type.rs vendored Normal file
View file

@ -0,0 +1,85 @@
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)]
pub enum CacheType {
#[serde(rename = "local")]
Local,
#[serde(rename = "redis")]
Redis,
#[serde(rename = "redis-semantic")]
RedisSemantic,
#[serde(rename = "valkey-semantic")]
ValkeySemantic,
#[serde(rename = "s3")]
S3,
#[serde(rename = "disk")]
Disk,
#[serde(rename = "qdrant-semantic")]
QdrantSemantic,
#[serde(rename = "azure-blob")]
AzureBlob,
#[serde(rename = "gcs")]
Gcs,
}
impl CacheType {
pub const ALL: [Self; 9] = [
Self::Local,
Self::Redis,
Self::RedisSemantic,
Self::ValkeySemantic,
Self::S3,
Self::Disk,
Self::QdrantSemantic,
Self::AzureBlob,
Self::Gcs,
];
pub const fn as_python_name(self) -> &'static str {
match self {
Self::Local => "local",
Self::Redis => "redis",
Self::RedisSemantic => "redis-semantic",
Self::ValkeySemantic => "valkey-semantic",
Self::S3 => "s3",
Self::Disk => "disk",
Self::QdrantSemantic => "qdrant-semantic",
Self::AzureBlob => "azure-blob",
Self::Gcs => "gcs",
}
}
pub fn from_python_name(value: &str) -> Option<Self> {
Self::ALL
.into_iter()
.find(|cache_type| cache_type.as_python_name() == value)
}
}
#[cfg(test)]
mod tests {
use super::CacheType;
#[test]
fn every_python_cache_type_has_one_round_trip_identity() {
let names = CacheType::ALL.map(CacheType::as_python_name);
assert_eq!(
names,
[
"local",
"redis",
"redis-semantic",
"valkey-semantic",
"s3",
"disk",
"qdrant-semantic",
"azure-blob",
"gcs",
]
);
assert_eq!(
names.map(CacheType::from_python_name),
CacheType::ALL.map(Some)
);
}
}

167
litellm-rust/crates/cache/src/caching.rs vendored
View file

@ -1,166 +1,23 @@
use std::sync::Arc;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::{BaseCache, CacheKwargs, Error};
pub use crate::BaseCache as Cache;
use crate::{BaseCache, Error};
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
pub enum CacheMode {
#[default]
#[serde(rename = "default_on")]
DefaultOn,
#[serde(rename = "default_off")]
DefaultOff,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct CacheKeyField {
pub name: String,
pub value: Option<String>,
pub api_parameter: bool,
pub internal_parameter: bool,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
pub struct CacheKeyInput {
pub fields: Vec<CacheKeyField>,
pub preset: Option<String>,
pub namespace: Option<String>,
pub include_provider_parameters: bool,
}
#[derive(Default)]
pub struct CacheKeyContext {
pub model_group: Option<String>,
pub caching_groups: Vec<(Vec<String>, String)>,
pub file_checksum: Option<String>,
pub file_object_name: Option<String>,
pub metadata_file_name: Option<String>,
pub parameters_file_name: Option<String>,
}
impl CacheKeyContext {
pub fn apply(self, input: &mut CacheKeyInput) {
let group = self.model_group.as_ref().and_then(|model| {
self.caching_groups
.iter()
.find(|(models, _)| models.contains(model))
});
for field in &mut input.fields {
match field.name.as_str() {
"model" => {
field.value = group
.map(|(_, formatted)| formatted.clone())
.or_else(|| self.model_group.clone())
.or_else(|| field.value.take())
}
"file" => {
field.value = self
.file_checksum
.clone()
.or_else(|| self.file_object_name.clone())
.or_else(|| self.metadata_file_name.clone())
.or_else(|| self.parameters_file_name.clone())
}
_ => {}
}
}
}
}
pub fn get_cache_key(input: &CacheKeyInput) -> String {
cache_key(input)
}
pub fn cache_key(input: &CacheKeyInput) -> String {
if let Some(preset) = &input.preset {
return preset.clone();
}
let mut digest = Sha256::new();
for field in &input.fields {
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
&& let Some(value) = &field.value
{
digest.update(field.name.as_bytes());
digest.update(b": ");
digest.update(value.as_bytes());
}
}
let hash = format!("{:x}", digest.finalize());
input
.namespace
.as_deref()
.filter(|namespace| !namespace.is_empty())
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
pub struct CacheControls {
pub supported_call_type: bool,
pub configured: bool,
pub native_backend: bool,
pub default_on: bool,
pub caching: Option<bool>,
pub no_cache: bool,
pub no_store: bool,
#[serde(default)]
pub use_cache: bool,
}
impl CacheControls {
pub fn reads(self) -> bool {
self.supported_call_type
&& self.configured
&& self.caching.unwrap_or(true)
&& !self.no_cache
&& (self.default_on || self.use_cache)
}
pub fn writes(self) -> bool {
self.supported_call_type
&& self.configured
&& !self.no_store
&& (self.default_on || self.use_cache)
}
}
pub fn should_use_cache(controls: CacheControls) -> bool {
controls.reads() || controls.writes()
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CacheEntry {
pub timestamp: f64,
pub response: Value,
}
impl CacheEntry {
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
self.timestamp.is_finite()
&& max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64())
}
}
pub fn get_cache(
cache: &dyn BaseCache<Value = CacheEntry>,
pub fn get_cache<B: BaseCache>(
cache: &B,
key: &str,
kwargs: &CacheKwargs,
) -> Result<Option<CacheEntry>, Error> {
cache.get_cache(key, kwargs)
context: &B::Context,
) -> Result<Option<B::Value>, Error> {
cache.get_cache(key, context)
}
pub fn set_cache(
cache: &dyn BaseCache<Value = CacheEntry>,
pub fn set_cache<B: BaseCache>(
cache: &B,
key: &str,
entry: CacheEntry,
kwargs: CacheKwargs,
value: B::Value,
context: &B::Context,
) -> Result<(), Error> {
cache.set_cache(key, entry, kwargs)
cache.set_cache(key, value, context)
}
pub type CacheBackend = Arc<dyn BaseCache<Value = CacheEntry>>;
pub type CacheBackend<B> = Arc<B>;

169
litellm-rust/crates/cache/src/capabilities.rs vendored Normal file
View file

@ -0,0 +1,169 @@
use std::{future::Future, time::Duration};
use crate::{BaseCache, BatchEntry, Error};
#[derive(Clone, Debug, PartialEq)]
pub struct IncrementOperation {
pub key: String,
pub amount: f64,
pub ttl: Option<Duration>,
}
pub trait BatchCache: BaseCache {
fn batch_get_cache(
&self,
keys: &[String],
context: &Self::Context,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
keys.iter()
.map(|key| match self.get_cache(key, context) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
}
fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: Self::Context,
) -> impl Future<Output = Result<Vec<BatchEntry<Self::Value>>, Error>> + Send {
async move {
let mut entries = Vec::with_capacity(keys.len());
for key in keys {
entries.push(match self.async_get_cache(&key, &context).await {
Ok(Some(value)) => BatchEntry::Hit(value),
Ok(None) => BatchEntry::Miss,
Err(Error::InvalidEntry) => BatchEntry::Invalid,
Err(error) => return Err(error),
});
}
Ok(entries)
}
}
}
pub trait DeleteCache: BaseCache {
fn delete_cache(&self, key: &str) -> Result<(), Error>;
fn async_delete_cache(&self, key: &str) -> impl Future<Output = Result<(), Error>> + Send {
async move { self.delete_cache(key) }
}
}
pub trait FlushCache: BaseCache {
fn flush_cache(&self) -> Result<(), Error>;
fn async_flush_cache(&self) -> impl Future<Output = Result<(), Error>> + Send {
async move { self.flush_cache() }
}
}
pub trait CounterCache: BaseCache<Value = f64> {
fn increment_cache(&self, key: &str, amount: f64, context: Self::Context)
-> Result<f64, Error>;
fn async_increment(
&self,
key: &str,
amount: f64,
context: Self::Context,
) -> impl Future<Output = Result<f64, Error>> + Send {
async move { self.increment_cache(key, amount, context) }
}
}
pub trait ClaimCache: BaseCache
where
Self::Value: PartialEq,
{
fn claim_cache(
&self,
key: &str,
candidate: Self::Value,
eligible: &[Self::Value],
context: Self::Context,
) -> Result<Self::Value, Error>;
fn async_claim_cache(
&self,
key: &str,
candidate: Self::Value,
eligible: Vec<Self::Value>,
context: Self::Context,
) -> impl Future<Output = Result<Self::Value, Error>> + Send {
async move { self.claim_cache(key, candidate, &eligible, context) }
}
}
pub trait TtlCache: BaseCache {
fn async_get_ttl(
&self,
key: &str,
) -> impl Future<Output = Result<Option<Duration>, Error>> + Send;
}
pub trait SetCache: BaseCache {
type SetValue: Clone + Send + Sync + 'static;
type SetResult: Send + Sync + 'static;
fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> impl Future<Output = Result<Self::SetResult, Error>> + Send;
}
pub trait QueueCache: BaseCache {
type QueueValue: Clone + Send + Sync + 'static;
type PopResult: Send + Sync + 'static;
fn async_rpush(
&self,
key: &str,
values: Vec<Self::QueueValue>,
) -> impl Future<Output = Result<usize, Error>> + Send;
fn async_lpop(
&self,
key: &str,
count: Option<usize>,
) -> impl Future<Output = Result<Self::PopResult, Error>> + Send;
}
pub trait ScanCache: BaseCache {
fn async_scan_iter(
&self,
pattern: &str,
count: usize,
) -> impl Future<Output = Result<Vec<String>, Error>> + Send;
}
pub trait ClientInfoCache: BaseCache {
type ClientList: Send + Sync + 'static;
type Info: Send + Sync + 'static;
fn client_list(&self) -> Result<Self::ClientList, Error>;
fn info(&self) -> Result<Self::Info, Error>;
}
pub trait CacheScript: Send + Sync + 'static {
type Argument: Clone + Send + Sync + 'static;
type Output: Send + Sync + 'static;
fn invoke(
&self,
keys: Vec<String>,
arguments: Vec<Self::Argument>,
) -> impl Future<Output = Result<Self::Output, Error>> + Send;
}
pub trait ScriptCache: BaseCache {
type Script: CacheScript;
fn async_register_script(&self, source: String) -> Self::Script;
}

50
litellm-rust/crates/cache/src/codec.rs vendored Normal file
View file

@ -0,0 +1,50 @@
use std::marker::PhantomData;
use serde::{Serialize, de::DeserializeOwned};
use crate::Error;
pub trait CacheCodec: Send + Sync {
type Value: Clone + Send + Sync + 'static;
fn encode(&self, value: &Self::Value) -> Result<Vec<u8>, Error>;
fn decode(&self, bytes: &[u8]) -> Result<Self::Value, Error>;
}
pub struct JsonCodec<V>(PhantomData<fn() -> V>);
impl<V> Clone for JsonCodec<V> {
fn clone(&self) -> Self {
*self
}
}
impl<V> Copy for JsonCodec<V> {}
impl<V> Default for JsonCodec<V> {
fn default() -> Self {
Self::new()
}
}
impl<V> JsonCodec<V> {
pub const fn new() -> Self {
Self(PhantomData)
}
}
impl<V> CacheCodec for JsonCodec<V>
where
V: Clone + Send + Sync + Serialize + DeserializeOwned + 'static,
{
type Value = V;
fn encode(&self, value: &Self::Value) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
}
fn decode(&self, bytes: &[u8]) -> Result<Self::Value, Error> {
serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry)
}
}

390
litellm-rust/crates/cache/src/dual.rs vendored Normal file
View file

@ -0,0 +1,390 @@
use std::{sync::Arc, time::Duration};
use crate::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ClaimCache,
CounterCache, DeleteCache, Error, FlushCache,
};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum ReadPolicy {
#[default]
LocalThenRemote,
LocalOnly,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum WritePolicy {
#[default]
Both,
LocalOnly,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum RemoteFailurePolicy {
#[default]
Propagate,
UseLocal,
}
pub struct DualCache<L1, L2> {
l1: Arc<L1>,
l2: Arc<L2>,
read_policy: ReadPolicy,
write_policy: WritePolicy,
remote_failure_policy: RemoteFailurePolicy,
promotion_ttl: Option<Duration>,
}
impl<L1, L2> DualCache<L1, L2> {
pub fn new(l1: Arc<L1>, l2: Arc<L2>) -> Self {
Self {
l1,
l2,
read_policy: ReadPolicy::default(),
write_policy: WritePolicy::default(),
remote_failure_policy: RemoteFailurePolicy::default(),
promotion_ttl: None,
}
}
pub fn with_read_policy(self, read_policy: ReadPolicy) -> Self {
Self {
read_policy,
..self
}
}
pub fn with_write_policy(self, write_policy: WritePolicy) -> Self {
Self {
write_policy,
..self
}
}
pub fn with_remote_failure_policy(self, remote_failure_policy: RemoteFailurePolicy) -> Self {
Self {
remote_failure_policy,
..self
}
}
pub fn with_promotion_ttl(self, promotion_ttl: Duration) -> Self {
Self {
promotion_ttl: Some(promotion_ttl),
..self
}
}
fn reads_remote(&self) -> bool {
self.read_policy == ReadPolicy::LocalThenRemote
}
fn writes_remote(&self) -> bool {
self.write_policy == WritePolicy::Both
}
fn remote<T>(&self, result: Result<T, Error>) -> Result<Option<T>, Error> {
match result {
Ok(value) => Ok(Some(value)),
Err(Error::Unavailable)
if self.remote_failure_policy == RemoteFailurePolicy::UseLocal =>
{
Ok(None)
}
Err(error) => Err(error),
}
}
fn promotion_context<C: CacheContext>(&self, context: &C) -> C {
context.with_ttl(self.promotion_ttl.or(context.ttl()))
}
}
impl<V, C, L1, L2> DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: BaseCache<Value = V, Context = C>,
L2: BaseCache<Value = V, Context = C>,
{
fn missing(entries: &[BatchEntry<V>]) -> Vec<usize> {
entries
.iter()
.enumerate()
.filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index))
.collect()
}
fn merge_batch(
&self,
keys: &[String],
context: &C,
mut entries: Vec<BatchEntry<V>>,
missing: Vec<usize>,
remote: Vec<BatchEntry<V>>,
) -> Result<Vec<BatchEntry<V>>, Error> {
if missing.len() != remote.len() {
return Err(Error::Unavailable);
}
for (index, entry) in missing.into_iter().zip(remote) {
if let BatchEntry::Hit(value) = &entry {
let promotion_context = self.promotion_context(context);
self.l1
.set_cache(&keys[index], value.clone(), &promotion_context)?;
}
entries[index] = entry;
}
Ok(entries)
}
}
impl<V, C, L1, L2> BaseCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: BaseCache<Value = V, Context = C>,
L2: BaseCache<Value = V, Context = C>,
{
type Value = V;
type Context = C;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
self.l2.get_ttl(context)
}
fn set_cache(&self, key: &str, value: V, context: &C) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.set_cache(key, value.clone(), context))?;
}
self.l1.set_cache(key, value, context)
}
fn get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
if let Some(value) = self.l1.get_cache(key, context)? {
return Ok(Some(value));
}
if !self.reads_remote() {
return Ok(None);
}
let value = self.remote(self.l2.get_cache(key, context))?.flatten();
if let Some(value) = &value {
let promotion_context = self.promotion_context(context);
self.l1.set_cache(key, value.clone(), &promotion_context)?;
}
Ok(value)
}
async fn async_set_cache(&self, key: &str, value: V, context: C) -> Result<(), Error> {
if self.writes_remote() {
self.remote(
self.l2
.async_set_cache(key, value.clone(), context.clone())
.await,
)?;
}
self.l1.async_set_cache(key, value, context).await
}
async fn async_get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
if let Some(value) = self.l1.async_get_cache(key, context).await? {
return Ok(Some(value));
}
if !self.reads_remote() {
return Ok(None);
}
let value = self
.remote(self.l2.async_get_cache(key, context).await)?
.flatten();
if let Some(value) = &value {
self.l1
.async_set_cache(key, value.clone(), self.promotion_context(context))
.await?;
}
Ok(value)
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, V)>,
context: C,
) -> Result<(), Error> {
if self.writes_remote() {
self.remote(
self.l2
.async_set_cache_pipeline(entries.clone(), context.clone())
.await,
)?;
}
self.l1.async_set_cache_pipeline(entries, context).await
}
async fn disconnect(&self) -> Result<(), Error> {
self.l2.disconnect().await?;
self.l1.disconnect().await
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
self.l2.test_connection().await
}
}
impl<V, C, L1, L2> BatchCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: BatchCache<Value = V, Context = C>,
L2: BatchCache<Value = V, Context = C>,
{
fn batch_get_cache(&self, keys: &[String], context: &C) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self.l1.batch_get_cache(keys, context)?;
let missing = Self::missing(&entries);
if missing.is_empty() || !self.reads_remote() {
return Ok(entries);
}
let remote_keys = missing
.iter()
.map(|index| keys[*index].clone())
.collect::<Vec<_>>();
match self.remote(self.l2.batch_get_cache(&remote_keys, context))? {
Some(remote) => self.merge_batch(keys, context, entries, missing, remote),
None => Ok(entries),
}
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: C,
) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self
.l1
.async_batch_get_cache(keys.clone(), context.clone())
.await?;
let missing = Self::missing(&entries);
if missing.is_empty() || !self.reads_remote() {
return Ok(entries);
}
let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect();
match self.remote(
self.l2
.async_batch_get_cache(remote_keys, context.clone())
.await,
)? {
Some(remote) => self.merge_batch(&keys, &context, entries, missing, remote),
None => Ok(entries),
}
}
}
impl<V, C, L1, L2> DeleteCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: DeleteCache<Value = V, Context = C>,
L2: DeleteCache<Value = V, Context = C>,
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.delete_cache(key))?;
}
self.l1.delete_cache(key)
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.async_delete_cache(key).await)?;
}
self.l1.async_delete_cache(key).await
}
}
impl<V, C, L1, L2> FlushCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: FlushCache<Value = V, Context = C>,
L2: FlushCache<Value = V, Context = C>,
{
fn flush_cache(&self) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.flush_cache())?;
}
self.l1.flush_cache()
}
async fn async_flush_cache(&self) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.async_flush_cache().await)?;
}
self.l1.async_flush_cache().await
}
}
impl<C, L1, L2> CounterCache for DualCache<L1, L2>
where
C: CacheContext,
L1: BaseCache<Value = f64, Context = C>,
L2: CounterCache<Context = C>,
{
fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
let value = self.l2.increment_cache(key, amount, context.clone())?;
self.l1.set_cache(key, value, &context)?;
Ok(value)
}
async fn async_increment(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
let value = self
.l2
.async_increment(key, amount, context.clone())
.await?;
self.l1.async_set_cache(key, value, context).await?;
Ok(value)
}
}
impl<V, C, L1, L2> ClaimCache for DualCache<L1, L2>
where
V: Clone + PartialEq + Send + Sync + 'static,
C: CacheContext,
L1: ClaimCache<Value = V, Context = C>,
L2: ClaimCache<Value = V, Context = C>,
{
fn claim_cache(&self, key: &str, candidate: V, eligible: &[V], context: C) -> Result<V, Error> {
match self.remote(
self.l2
.claim_cache(key, candidate.clone(), eligible, context.clone()),
)? {
Some(winner) => {
self.l1.set_cache(key, winner.clone(), &context)?;
Ok(winner)
}
None => self.l1.claim_cache(key, candidate, eligible, context),
}
}
async fn async_claim_cache(
&self,
key: &str,
candidate: V,
eligible: Vec<V>,
context: C,
) -> Result<V, Error> {
match self.remote(
self.l2
.async_claim_cache(key, candidate.clone(), eligible.clone(), context.clone())
.await,
)? {
Some(winner) => {
self.l1
.async_set_cache(key, winner.clone(), context)
.await?;
Ok(winner)
}
None => {
self.l1
.async_claim_cache(key, candidate, eligible, context)
.await
}
}
}
}

2
litellm-rust/crates/cache/src/error.rs vendored
View file

@ -4,4 +4,6 @@ pub enum Error {
Unavailable,
#[error("invalid cache entry")]
InvalidEntry,
#[error("flushing Redis requires an explicit namespace")]
UnscopedFlush,
}

17
litellm-rust/crates/cache/src/lib.rs vendored
View file

@ -1,12 +1,21 @@
mod base_cache;
mod cache_type;
mod caching;
mod capabilities;
mod codec;
mod dual;
mod error;
pub use base_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs,
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext,
ExactCacheContext,
};
pub use caching::{
Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput,
CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache,
pub use cache_type::CacheType;
pub use caching::{Cache, CacheBackend, get_cache, set_cache};
pub use capabilities::{
BatchCache, CacheScript, ClaimCache, ClientInfoCache, CounterCache, DeleteCache, FlushCache,
IncrementOperation, QueueCache, ScanCache, ScriptCache, SetCache, TtlCache,
};
pub use codec::{CacheCodec, JsonCodec};
pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy};
pub use error::Error;

208
litellm-rust/crates/cache/tests/caching.rs vendored
View file

@ -1,42 +1,97 @@
use std::{sync::Mutex, time::Duration};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext,
CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key,
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache,
};
use sha2::{Digest, Sha256};
use std::time::Duration;
struct TestCache {
default_ttl: Duration,
writes: Mutex<Vec<(String, String, ExactCacheContext)>>,
}
#[derive(Clone)]
struct SemanticContext {
ttl: Option<Duration>,
query: String,
}
impl CacheContext for SemanticContext {
fn ttl(&self) -> Option<Duration> {
self.ttl
}
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
Self {
ttl,
query: self.query.clone(),
}
}
}
struct SemanticCache;
impl BaseCache for SemanticCache {
type Value = String;
type Context = SemanticContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> {
Ok(())
}
fn get_cache(&self, _: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
Ok((context.query == "matching prompt").then(|| "semantic hit".into()))
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
unreachable!()
}
}
impl BaseCache for TestCache {
type Value = CacheEntry;
type Value = String;
type Context = ExactCacheContext;
fn default_ttl(&self) -> Duration {
self.default_ttl
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> {
fn set_cache(&self, _: &str, _: Self::Value, _: &ExactCacheContext) -> Result<(), Error> {
Err(Error::Unavailable)
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: ExactCacheContext,
) -> Result<(), Error> {
if key == "unavailable" {
return Err(Error::Unavailable);
}
self.writes
.lock()
.unwrap()
.push((key.into(), value, context));
Ok(())
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
Ok(None)
}
fn delete_cache(&self, _: &str) -> Result<(), Error> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
unreachable!()
}
}
@ -45,95 +100,64 @@ impl BaseCache for TestCache {
fn ttl_uses_default_and_allows_per_call_override() {
let cache = TestCache {
default_ttl: Duration::from_secs(60),
writes: Mutex::default(),
};
assert_eq!(
cache.get_ttl(&CacheKwargs::default()),
Duration::from_secs(60)
cache.get_ttl(&ExactCacheContext::default()),
Some(Duration::from_secs(60))
);
assert_eq!(
cache.get_ttl(&CacheKwargs {
cache.get_ttl(&ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
..Default::default()
}),
Duration::from_secs(5)
Some(Duration::from_secs(5))
);
}
#[test]
fn keys_match_python_order_groups_files_presets_and_namespaces() {
let mut input = CacheKeyInput {
fields: vec![
CacheKeyField {
name: "model".into(),
value: Some("deployment".into()),
api_parameter: true,
internal_parameter: false,
},
CacheKeyField {
name: "file".into(),
value: None,
api_parameter: true,
internal_parameter: false,
},
],
namespace: Some("team".into()),
..Default::default()
fn associated_context_preserves_backend_specific_lookup_inputs() {
let context = SemanticContext {
ttl: None,
query: "matching prompt".into(),
};
CacheKeyContext {
model_group: Some("group".into()),
caching_groups: vec![(vec!["group".into()], "['group']".into())],
file_checksum: Some("checksum".into()),
..Default::default()
}
.apply(&mut input);
assert_eq!(
cache_key(&input),
format!(
"team:{:x}",
Sha256::digest(b"model: ['group']file: checksum")
)
get_cache(&SemanticCache, "shared-key", &context).unwrap(),
Some("semantic hit".into())
);
input.preset = Some("preset".into());
assert_eq!(get_cache_key(&input), "preset");
}
#[test]
fn cache_controls_honor_default_modes_and_directives() {
let enabled = CacheControls {
supported_call_type: true,
configured: true,
default_on: true,
..Default::default()
#[tokio::test]
async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
let cache = TestCache {
default_ttl: Duration::from_secs(60),
writes: Mutex::default(),
};
assert!(enabled.reads());
assert!(enabled.writes());
assert!(
!CacheControls {
default_on: false,
..enabled
}
.reads()
let entry = String::from("cached");
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
};
cache
.batch_cache_write("single", entry.clone(), context.clone())
.await
.unwrap();
assert_eq!(
cache
.async_set_cache_pipeline(
vec![
("first".into(), entry.clone()),
("unavailable".into(), entry.clone()),
("skipped".into(), entry.clone()),
],
context.clone(),
)
.await,
Err(Error::Unavailable)
);
assert!(
CacheControls {
default_on: false,
use_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_store: true,
..enabled
}
.writes()
assert_eq!(
*cache.writes.lock().unwrap(),
vec![
("single".into(), entry.clone(), context.clone()),
("first".into(), entry, context),
]
);
}

41
litellm-rust/crates/cache/tests/codec.rs vendored Normal file
View file

@ -0,0 +1,41 @@
use std::collections::BTreeMap;
use litellm_cache::{CacheCodec, Error, JsonCodec};
use serde::{Deserialize, Serialize};
use serde_json::json;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
struct RoutingState {
deployment: String,
cooldown_seconds: u64,
}
#[test]
fn json_codec_round_trips_typed_domain_values() {
let codec = JsonCodec::<RoutingState>::new();
let value = RoutingState {
deployment: "deployment-a".into(),
cooldown_seconds: 30,
};
let bytes = codec.encode(&value).unwrap();
assert_eq!(codec.decode(&bytes).unwrap(), value);
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&bytes).unwrap(),
json!({"deployment": "deployment-a", "cooldown_seconds": 30})
);
}
#[test]
fn json_codec_rejects_malformed_and_wrongly_typed_entries() {
let codec = JsonCodec::<RoutingState>::new();
for bytes in [b"not json".as_slice(), br#"{"deployment":12}"#.as_slice()] {
assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry);
}
}
#[test]
fn json_codec_propagates_encoding_errors() {
let codec = JsonCodec::<BTreeMap<(u8, u8), String>>::new();
let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]);
assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry);
}

385
litellm-rust/crates/cache/tests/dual.rs vendored Normal file
View file

@ -0,0 +1,385 @@
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{
BaseCache, BatchCache, CacheConnectionResult, ClaimCache, CounterCache, DeleteCache, DualCache,
Error, ExactCacheContext, FlushCache, ReadPolicy, RemoteFailurePolicy, WritePolicy,
};
struct TestCache<V> {
value: Mutex<Option<V>>,
fail: bool,
}
impl<V> TestCache<V> {
fn new(value: Option<V>, fail: bool) -> Self {
Self {
value: Mutex::new(value),
fail,
}
}
}
impl<V> BaseCache for TestCache<V>
where
V: Clone + Send + Sync + 'static,
{
type Value = V;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(Duration::from_secs(60)))
}
fn set_cache(&self, _: &str, value: V, _: &ExactCacheContext) -> Result<(), Error> {
*self.value.lock().unwrap() = Some(value);
Ok(())
}
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<V>, Error> {
Ok(self.value.lock().unwrap().clone())
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
unreachable!()
}
}
impl<V> BatchCache for TestCache<V> where V: Clone + Send + Sync + 'static {}
impl<V> DeleteCache for TestCache<V>
where
V: Clone + Send + Sync + 'static,
{
fn delete_cache(&self, _: &str) -> Result<(), Error> {
*self.value.lock().unwrap() = None;
Ok(())
}
}
impl<V> FlushCache for TestCache<V>
where
V: Clone + Send + Sync + 'static,
{
fn flush_cache(&self) -> Result<(), Error> {
*self.value.lock().unwrap() = None;
Ok(())
}
}
impl CounterCache for TestCache<f64> {
fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result<f64, Error> {
if self.fail {
return Err(Error::Unavailable);
}
let mut value = self.value.lock().unwrap();
let incremented = value.unwrap_or_default() + amount;
*value = Some(incremented);
Ok(incremented)
}
}
impl<V> ClaimCache for TestCache<V>
where
V: Clone + PartialEq + Send + Sync + 'static,
{
fn claim_cache(
&self,
_: &str,
candidate: V,
eligible: &[V],
_: ExactCacheContext,
) -> Result<V, Error> {
if self.fail {
return Err(Error::Unavailable);
}
let mut value = self.value.lock().unwrap();
let winner = match value.as_ref() {
Some(existing) if eligible.is_empty() || eligible.contains(existing) => {
existing.clone()
}
_ => candidate,
};
*value = Some(winner.clone());
Ok(winner)
}
}
#[test]
fn failed_l2_increment_leaves_l1_unchanged() {
let l1 = Arc::new(TestCache::new(Some(10.0), false));
let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true)));
assert_eq!(
cache.increment_cache("counter", 2.0, ExactCacheContext::default()),
Err(Error::Unavailable)
);
assert_eq!(
l1.get_cache("counter", &ExactCacheContext::default())
.unwrap(),
Some(10.0)
);
}
#[test]
fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() {
let l1 = Arc::new(TestCache::new(Some("first".to_string()), false));
let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true)))
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(
cache
.claim_cache(
"affinity",
"second".into(),
&["first".into(), "second".into()],
ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
},
)
.unwrap(),
"first"
);
}
struct SyncPanics(TestCache<String>);
impl BaseCache for SyncPanics {
type Value = String;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
self.0.get_ttl(context)
}
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
panic!("sync L2 write on an async path")
}
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
panic!("sync L2 read on an async path")
}
async fn async_set_cache(
&self,
key: &str,
value: String,
context: ExactCacheContext,
) -> Result<(), Error> {
self.0.set_cache(key, value, &context)
}
async fn async_get_cache(
&self,
key: &str,
context: &ExactCacheContext,
) -> Result<Option<String>, Error> {
self.0.get_cache(key, context)
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, String)>,
context: ExactCacheContext,
) -> Result<(), Error> {
for (key, value) in cache_list {
self.0.set_cache(&key, value, &context)?;
}
Ok(())
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
unreachable!()
}
}
impl BatchCache for SyncPanics {
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: ExactCacheContext,
) -> Result<Vec<litellm_cache::BatchEntry<String>>, Error> {
assert_eq!(keys, ["missing"]);
Ok(vec![match self.0.get_cache("missing", &context)? {
Some(value) => litellm_cache::BatchEntry::Hit(value),
None => litellm_cache::BatchEntry::Miss,
}])
}
}
impl DeleteCache for SyncPanics {
fn delete_cache(&self, _: &str) -> Result<(), Error> {
panic!("sync L2 delete on an async path")
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
self.0.delete_cache(key)
}
}
impl FlushCache for SyncPanics {
fn flush_cache(&self) -> Result<(), Error> {
panic!("sync L2 flush on an async path")
}
}
#[tokio::test]
async fn async_operations_use_the_async_l2_methods() {
let l1 = Arc::new(TestCache::new(None, false));
let cache = DualCache::new(
l1.clone(),
Arc::new(SyncPanics(TestCache::new(
Some("remote".to_string()),
false,
))),
);
let context = ExactCacheContext::default();
assert_eq!(
cache.async_get_cache("missing", &context).await.unwrap(),
Some("remote".into())
);
assert_eq!(
l1.get_cache("missing", &context).unwrap(),
Some("remote".into())
);
l1.delete_cache("missing").unwrap();
assert_eq!(
cache
.async_batch_get_cache(vec!["missing".into()], context.clone())
.await
.unwrap(),
[litellm_cache::BatchEntry::Hit("remote".to_string())]
);
cache
.async_set_cache("missing", "written".into(), context.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(vec![("missing".into(), "piped".into())], context.clone())
.await
.unwrap();
cache.async_delete_cache("missing").await.unwrap();
assert_eq!(
cache.async_get_cache("missing", &context).await.unwrap(),
None
);
}
struct Unavailable;
impl BaseCache for Unavailable {
type Value = String;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
Err(Error::Unavailable)
}
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
Err(Error::Unavailable)
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
unreachable!()
}
}
impl BatchCache for Unavailable {}
impl DeleteCache for Unavailable {
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Err(Error::Unavailable)
}
}
impl FlushCache for Unavailable {
fn flush_cache(&self) -> Result<(), Error> {
Err(Error::Unavailable)
}
}
impl ClaimCache for Unavailable {
fn claim_cache(
&self,
_: &str,
_: String,
_: &[String],
_: ExactCacheContext,
) -> Result<String, Error> {
Err(Error::InvalidEntry)
}
}
#[test]
fn remote_failure_policy_selects_propagation_or_the_local_tier() {
let context = ExactCacheContext::default();
let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable));
assert_eq!(
strict.set_cache("key", "value".into(), &context),
Err(Error::Unavailable)
);
assert_eq!(strict.get_cache("key", &context), Err(Error::Unavailable));
let l1 = Arc::new(TestCache::new(None, false));
let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable))
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(degraded.get_cache("key", &context), Ok(None));
degraded.set_cache("key", "value".into(), &context).unwrap();
assert_eq!(
degraded.get_cache("key", &context),
Ok(Some("value".into()))
);
degraded.delete_cache("key").unwrap();
assert_eq!(l1.get_cache("key", &context), Ok(None));
}
#[test]
fn claim_fallback_does_not_hide_non_availability_errors() {
let cache = DualCache::new(
Arc::new(TestCache::new(Some("first".to_string()), false)),
Arc::new(Unavailable),
)
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(
cache.claim_cache(
"affinity",
"second".into(),
&[],
ExactCacheContext::default()
),
Err(Error::InvalidEntry)
);
}
#[test]
fn local_only_policies_never_touch_l2() {
let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false));
let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone())
.with_read_policy(ReadPolicy::LocalOnly)
.with_write_policy(WritePolicy::LocalOnly);
let context = ExactCacheContext::default();
assert_eq!(cache.get_cache("key", &context), Ok(None));
cache.set_cache("key", "local".into(), &context).unwrap();
assert_eq!(l2.get_cache("key", &context), Ok(Some("remote".into())));
}

5
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -20,6 +20,11 @@ tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
bytes.workspace = true
litellm-cache.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
litellm-callbacks-legacy-python.workspace = true
litellm-core.workspace = true

291
litellm-rust/crates/python-bridge/src/cache/binding.rs vendored Normal file
View file

@ -0,0 +1,291 @@
use litellm_cache_response::PartialHits;
use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py};
use pyo3::{
PyTraverseError, PyVisit,
exceptions::{PyRuntimeError, PyValueError},
prelude::*,
types::PyDict,
};
use serde_json::Value;
use super::{
cache_error,
callback::PythonCallback,
future::{ready_none, ready_value},
native::NativeResponseCache,
request::{now, request, requests},
};
pub(super) enum CacheBinding {
Disabled,
Native(NativeResponseCache),
PythonCallback(PythonCallback),
}
#[pyclass(frozen, name = "_CacheTestBinding")]
pub(crate) struct ResolvedCache {
binding: CacheBinding,
pid: u32,
}
impl ResolvedCache {
pub(super) fn new(binding: CacheBinding) -> Self {
Self {
binding,
pid: std::process::id(),
}
}
fn check_process(&self) -> PyResult<()> {
if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() {
return Err(PyRuntimeError::new_err(
"native cache bindings must be resolved again after fork",
));
}
Ok(())
}
pub(crate) fn lookup_step(
&self,
py: Python<'_>,
input: &Bound<'_, PyAny>,
kwargs: Option<&Bound<'_, PyDict>>,
) -> PyResult<ExecutionStep> {
self.check_process()?;
let awaitable = match &self.binding {
CacheBinding::Disabled => ready_none(py)?,
CacheBinding::Native(service) => {
let request = request(input)?;
let service = service.clone();
run_async(
py,
async move { service.async_lookup(&request, now()).await },
cache_error,
)?
}
CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?,
};
Ok(ExecutionStep::Await(awaitable.unbind()))
}
}
#[pymethods]
impl ResolvedCache {
#[getter]
fn kind(&self) -> &'static str {
match self.binding {
CacheBinding::Disabled => "disabled",
CacheBinding::Native(_) => "native",
CacheBinding::PythonCallback(_) => "python_callback",
}
}
#[pyo3(signature = (request, *, callback_kwargs=None))]
fn lookup(
&self,
py: Python<'_>,
request: &Bound<'_, PyAny>,
callback_kwargs: Option<&Bound<'_, PyDict>>,
) -> PyResult<Py<PyAny>> {
self.check_process()?;
match &self.binding {
CacheBinding::Disabled => Ok(py.None()),
CacheBinding::Native(service) => {
let request = self::request(request)?;
let service = service.clone();
let response = release_gil(py, move || service.lookup(&request, now()))
.map_err(cache_error)?;
to_py(py, &response)
}
CacheBinding::PythonCallback(callback) => {
callback.lookup(py, callback_kwargs).map(Bound::unbind)
}
}
}
#[pyo3(signature = (request, response, *, callback_kwargs=None))]
fn store(
&self,
py: Python<'_>,
request: &Bound<'_, PyAny>,
response: &Bound<'_, PyAny>,
callback_kwargs: Option<&Bound<'_, PyDict>>,
) -> PyResult<()> {
self.check_process()?;
match &self.binding {
CacheBinding::Disabled => Ok(()),
CacheBinding::Native(service) => {
let request = self::request(request)?;
let response: Value = from_py(response)?;
let service = service.clone();
release_gil(py, move || service.store(&request, response, now()))
.map_err(cache_error)
}
CacheBinding::PythonCallback(callback) => callback.store(py, response, callback_kwargs),
}
}
#[pyo3(signature = (requests, *, callback_kwargs=None))]
fn lookup_batch(
&self,
py: Python<'_>,
requests: &Bound<'_, PyAny>,
callback_kwargs: Option<&Bound<'_, PyAny>>,
) -> PyResult<Py<PyAny>> {
self.check_process()?;
match &self.binding {
CacheBinding::Disabled => {
let requests = self::requests(requests)?;
to_py(py, &PartialHits::new(vec![None; requests.len()]))
}
CacheBinding::Native(service) => {
let requests = self::requests(requests)?;
let service = service.clone();
let response = release_gil(py, move || service.lookup_batch(&requests, now()))
.map_err(cache_error)?;
to_py(py, &response)
}
CacheBinding::PythonCallback(callback) => callback
.lookup_batch(py, requests, callback_kwargs)
.map(Bound::unbind),
}
}
#[pyo3(signature = (request, *, callback_kwargs=None))]
fn async_lookup<'py>(
&self,
py: Python<'py>,
request: &Bound<'py, PyAny>,
callback_kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)?
else {
unreachable!()
};
Ok(awaitable.into_bound(py))
}
#[pyo3(signature = (request, response, *, callback_kwargs=None))]
fn async_store<'py>(
&self,
py: Python<'py>,
request: &Bound<'py, PyAny>,
response: &Bound<'py, PyAny>,
callback_kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
match &self.binding {
CacheBinding::Disabled => ready_none(py),
CacheBinding::Native(service) => {
let request = self::request(request)?;
let response: Value = from_py(response)?;
let service = service.clone();
run_async(
py,
async move { service.async_store(&request, response, now()).await },
cache_error,
)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store(py, response, callback_kwargs)
}
}
}
#[pyo3(signature = (requests, *, callback_kwargs=None))]
fn async_lookup_batch<'py>(
&self,
py: Python<'py>,
requests: &Bound<'py, PyAny>,
callback_kwargs: Option<&Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
match &self.binding {
CacheBinding::Disabled => {
let requests = self::requests(requests)?;
ready_value(py, &PartialHits::new(vec![None; requests.len()]))
}
CacheBinding::Native(service) => {
let requests = self::requests(requests)?;
let service = service.clone();
run_async(
py,
async move { service.async_lookup_batch(&requests, now()).await },
cache_error,
)
}
CacheBinding::PythonCallback(callback) => {
callback.async_lookup_batch(py, requests, callback_kwargs)
}
}
}
#[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))]
fn async_store_batch<'py>(
&self,
py: Python<'py>,
requests: &Bound<'py, PyAny>,
responses: &Bound<'py, PyAny>,
callback_result: Option<&Bound<'py, PyAny>>,
callback_kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
match &self.binding {
CacheBinding::Disabled => ready_none(py),
CacheBinding::Native(service) => {
let requests = self::requests(requests)?;
let responses: Vec<Value> = from_py(responses)?;
if requests.len() != responses.len() {
return Err(PyValueError::new_err(
"batch cache requests and responses must have equal lengths",
));
}
let entries = requests.into_iter().zip(responses).collect();
let service = service.clone();
run_async(
py,
async move { service.async_store_batch(entries, now()).await },
cache_error,
)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store_batch(py, callback_result, callback_kwargs)
}
}
}
fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
match &self.binding {
CacheBinding::Disabled => ready_none(py),
CacheBinding::Native(service) => {
let service = service.clone();
run_async(py, async move { service.async_flush().await }, cache_error)
}
CacheBinding::PythonCallback(callback) => callback.async_flush(py),
}
}
fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
match &self.binding {
CacheBinding::Disabled => ready_none(py),
CacheBinding::Native(service) => {
let service = service.clone();
run_async(
py,
async move { service.test_connection().await },
cache_error,
)
}
CacheBinding::PythonCallback(callback) => callback.ping(py),
}
}
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
if let CacheBinding::PythonCallback(callback) = &self.binding {
callback.traverse(&visit)?;
}
Ok(())
}
}

162
litellm-rust/crates/python-bridge/src/cache/callback.rs vendored Normal file
View file

@ -0,0 +1,162 @@
use pyo3::{
PyTraverseError, PyVisit,
exceptions::{PyTypeError, PyValueError},
prelude::*,
types::{PyDict, PyList, PyTuple},
};
use super::future::ready_none;
pub(super) struct PythonCallback(Py<PyAny>);
impl PythonCallback {
pub(super) fn new(object: Py<PyAny>) -> Self {
Self(object)
}
pub(super) fn lookup<'py>(
&self,
py: Python<'py>,
kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
self.0
.bind(py)
.call_method("get_cache", (), Some(callback_kwargs(kwargs)?))
}
pub(super) fn async_lookup<'py>(
&self,
py: Python<'py>,
kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
self.0
.bind(py)
.call_method("async_get_cache", (), Some(callback_kwargs(kwargs)?))
}
pub(super) fn store(
&self,
py: Python<'_>,
response: &Bound<'_, PyAny>,
kwargs: Option<&Bound<'_, PyDict>>,
) -> PyResult<()> {
self.0
.bind(py)
.call_method("add_cache", (response,), Some(callback_kwargs(kwargs)?))
.map(|_| ())
}
pub(super) fn async_store<'py>(
&self,
py: Python<'py>,
response: &Bound<'py, PyAny>,
kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
self.0.bind(py).call_method(
"async_add_cache",
(response,),
Some(callback_kwargs(kwargs)?),
)
}
pub(super) fn lookup_batch<'py>(
&self,
py: Python<'py>,
requests: &Bound<'py, PyAny>,
kwargs: Option<&Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyAny>> {
let results = PyList::empty(py);
for kwargs in batch_callback_kwargs(requests, kwargs)? {
results.append(
self.0
.bind(py)
.call_method("get_cache", (), Some(&kwargs))?,
)?;
}
Ok(results.into_any())
}
pub(super) fn async_lookup_batch<'py>(
&self,
py: Python<'py>,
requests: &Bound<'py, PyAny>,
kwargs: Option<&Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyAny>> {
let awaitables = batch_callback_kwargs(requests, kwargs)?
.iter()
.map(|kwargs| {
self.0
.bind(py)
.call_method("async_get_cache", (), Some(kwargs))
})
.collect::<PyResult<Vec<_>>>()?;
py.import("asyncio")?
.call_method1("gather", PyTuple::new(py, awaitables)?)
}
pub(super) fn async_store_batch<'py>(
&self,
py: Python<'py>,
result: Option<&Bound<'py, PyAny>>,
kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
let result = result.ok_or_else(|| {
PyTypeError::new_err("Python cache callbacks require their original callback_result")
})?;
self.0.bind(py).call_method(
"async_add_cache_pipeline",
(result,),
Some(callback_kwargs(kwargs)?),
)
}
pub(super) fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let object = self.0.bind(py);
let backend = match object.getattr_opt("cache")? {
Some(backend) if !backend.is_none() => backend,
_ => object.clone(),
};
if backend.hasattr("async_flush_cache")? {
return backend.call_method0("async_flush_cache");
}
backend.call_method0("flush_cache")?;
ready_none(py)
}
pub(super) fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.0.bind(py).call_method0("ping")
}
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
}
fn callback_kwargs<'a, 'py>(
kwargs: Option<&'a Bound<'py, PyDict>>,
) -> PyResult<&'a Bound<'py, PyDict>> {
kwargs.ok_or_else(|| {
PyTypeError::new_err("Python cache callbacks require their original callback_kwargs")
})
}
fn batch_callback_kwargs<'py>(
requests: &Bound<'py, PyAny>,
kwargs: Option<&Bound<'py, PyAny>>,
) -> PyResult<Vec<Bound<'py, PyDict>>> {
let kwargs = kwargs
.ok_or_else(|| {
PyTypeError::new_err(
"Python cache callbacks require one original callback_kwargs mapping per request",
)
})?
.try_iter()?
.map(|item| Ok(item?.cast_into::<PyDict>()?))
.collect::<PyResult<Vec<_>>>()?;
if kwargs.len() != requests.len()? {
return Err(PyValueError::new_err(
"batch cache requests and callback_kwargs must have equal lengths",
));
}
Ok(kwargs)
}

594
litellm-rust/crates/python-bridge/src/cache/config.rs vendored Normal file
View file

@ -0,0 +1,594 @@
use std::time::Duration;
use litellm_cache::CacheType;
use pyo3::{
exceptions::{PyTypeError, PyValueError},
prelude::*,
types::{PyAny, PyDict, PyString},
};
use super::{native::NativeResponseCache, request::duration};
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct CachePolicy {
pub(super) mode: String,
pub(super) ttl: Option<Duration>,
pub(super) namespace: Option<String>,
pub(super) supported_call_types: Option<Vec<String>>,
pub(super) redis_flush_size: Option<usize>,
pub(super) semantic_cache_scope: String,
}
pub(super) struct MemoryCacheConfig {
pub(super) default_ttl: Duration,
pub(super) capacity: usize,
pub(super) max_entry_bytes: usize,
}
#[derive(Debug, PartialEq)]
pub(super) enum RedisProtocol {
Resp2,
Resp3,
}
#[derive(Debug, PartialEq)]
pub(super) enum CertificateRequirement {
None,
Optional,
Required,
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct RedisTlsConfig {
pub(super) certificate_requirement: CertificateRequirement,
pub(super) check_hostname: bool,
pub(super) ca_certificate: Option<String>,
pub(super) ca_data: Option<String>,
pub(super) client_certificate: Option<String>,
pub(super) client_key: Option<String>,
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct RedisConnectionConfig {
pub(super) host: String,
pub(super) port: u16,
pub(super) database: i64,
pub(super) username: Option<String>,
pub(super) password: Option<String>,
pub(super) protocol: RedisProtocol,
pub(super) pool_size: usize,
pub(super) read_timeout: Option<Duration>,
pub(super) connect_timeout: Option<Duration>,
pub(super) socket_keepalive: Option<bool>,
pub(super) health_check_interval: Duration,
pub(super) client_name: Option<String>,
pub(super) tls: Option<RedisTlsConfig>,
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct RedisCacheConfig {
pub(super) default_ttl: Duration,
pub(super) namespace: Option<String>,
pub(super) flush_size: usize,
pub(super) connection: RedisConnectionConfig,
}
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct NativeCacheConfig {
pub(super) policy: CachePolicy,
pub(super) backend: CacheBackendConfig,
}
pub(super) enum UnsupportedCacheConfig {
Backend,
RedisTopology,
RedisCredentials,
RedisConnection,
RedisOption,
}
impl UnsupportedCacheConfig {
pub(super) fn message(&self) -> &'static str {
match self {
Self::Backend => "native cache backend is not implemented",
Self::RedisTopology => "native Redis topology is not implemented",
Self::RedisCredentials => "native Redis credentials require Python",
Self::RedisConnection => "native Redis connection type is not implemented",
Self::RedisOption => "native Redis configuration requires Python",
}
}
}
pub(super) enum CacheConfigProjection {
Native(Box<NativeCacheConfig>),
Unsupported(UnsupportedCacheConfig),
}
impl NativeCacheConfig {
#[inline(never)]
pub(super) fn project(facade: &Bound<'_, PyAny>) -> PyResult<CacheConfigProjection> {
let backend_name = facade.getattr("type")?.extract::<String>()?;
let policy = CachePolicy {
mode: facade.getattr("mode")?.extract::<String>()?,
ttl: optional_duration(facade.getattr("ttl")?)?,
namespace: optional_string(facade.getattr("namespace")?)?,
supported_call_types: facade
.getattr("supported_call_types")?
.extract::<Option<Vec<String>>>()?,
redis_flush_size: facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
semantic_cache_scope: facade
.getattr("semantic_cache_scope")?
.extract::<String>()?,
};
let backend = facade.getattr("cache")?;
match CacheType::from_python_name(&backend_name) {
Some(CacheType::Local) => project_memory(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Memory(backend),
}))
}),
Some(CacheType::Redis) => match project_redis(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Redis(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(
CacheType::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic
| CacheType::AzureBlob
| CacheType::Gcs,
)
| None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
)),
}
}
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
if service.default_ttl()
!= Some(match &self.backend {
CacheBackendConfig::Memory(config) => config.default_ttl,
CacheBackendConfig::Redis(config) => config.default_ttl,
})
{
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
CacheBackendConfig::Memory(config) if service.kind() != "memory" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => {
Some("facade and native backend capacities must match")
}
CacheBackendConfig::Memory(config)
if service.max_entry_bytes() != Some(config.max_entry_bytes) =>
{
Some("facade and native backend item limits must match")
}
CacheBackendConfig::Memory(_) => None,
CacheBackendConfig::Redis(_) if service.kind() != "redis" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
}
}
}
#[inline(never)]
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
Ok(MemoryCacheConfig {
default_ttl: duration(backend.getattr("default_ttl")?.extract::<f64>()?)?,
capacity: backend.getattr("max_size_in_memory")?.extract::<usize>()?,
max_entry_bytes: max_size_kib
.checked_mul(1024)
.ok_or_else(|| PyValueError::new_err("memory cache item limit is too large"))?,
})
}
#[inline(never)]
fn project_redis(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<RedisCacheConfig, UnsupportedCacheConfig>> {
let source = backend.getattr("redis_kwargs")?.cast_into::<PyDict>()?;
if has_value(&source, "startup_nodes")? {
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
}
if has_value(&source, "sentinel_nodes")? {
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
}
for key in ["credential_provider", "redis_connect_func"] {
if has_value(&source, key)? {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
}
if has_value(&source, "connection_pool")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
for key in [
"retry",
"retry_on_error",
"socket_keepalive_options",
"unix_socket_path",
"cache",
"cache_config",
"event_dispatcher",
"ssl_ca_path",
"ssl_password",
"ssl_min_version",
"ssl_ciphers",
"ssl_validate_ocsp",
"ssl_validate_ocsp_stapled",
"ssl_ocsp_context",
"ssl_ocsp_expected_cert",
] {
if has_value(&source, key)? {
return Ok(Err(UnsupportedCacheConfig::RedisOption));
}
}
for key in ["retry_on_timeout", "single_connection_client"] {
if optional_coerced_bool(&source, key)?.unwrap_or(false) {
return Ok(Err(UnsupportedCacheConfig::RedisOption));
}
}
let client = backend.getattr("redis_client")?;
let pool = client.getattr("connection_pool")?;
if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
for key in ["credential_provider", "redis_connect_func"] {
if has_value(&resolved, key)? {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
}
let connection_class = resolved
.get_item("connection_class")?
.unwrap_or(pool.getattr("connection_class")?);
let tls = if class_is(&connection_class, "redis.connection", "Connection")? {
None
} else if class_is(&connection_class, "redis.connection", "SSLConnection")? {
Some(project_tls(&resolved)?)
} else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) {
2 => RedisProtocol::Resp2,
3 => RedisProtocol::Resp3,
_ => return Err(PyValueError::new_err("unsupported Redis protocol version")),
};
let health_check_interval =
duration(optional_f64(&resolved, "health_check_interval")?.unwrap_or(0.0))?;
Ok(Ok(RedisCacheConfig {
default_ttl: duration(backend.getattr("default_ttl")?.extract::<f64>()?)?,
namespace: optional_attribute_string(backend, "namespace")?,
flush_size: backend.getattr("redis_flush_size")?.extract::<usize>()?,
connection: RedisConnectionConfig {
host: required_string(&resolved, "host")?,
port: u16::try_from(required_i64(&resolved, "port")?)
.map_err(|_| PyValueError::new_err("invalid Redis port"))?,
database: optional_i64(&resolved, "db")?.unwrap_or(0),
username: optional_dict_string(&resolved, "username")?,
password: optional_dict_string(&resolved, "password")?,
protocol,
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
read_timeout: optional_dict_duration(&resolved, "socket_timeout")?,
connect_timeout: optional_dict_duration(&resolved, "socket_connect_timeout")?,
socket_keepalive: optional_bool(&resolved, "socket_keepalive")?,
health_check_interval,
client_name: optional_dict_string(&resolved, "client_name")?,
tls,
},
}))
}
#[inline(never)]
fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
Ok(RedisTlsConfig {
certificate_requirement: certificate_requirement(values)?,
check_hostname: optional_bool(values, "ssl_check_hostname")?.unwrap_or(false),
ca_certificate: optional_dict_string(values, "ssl_ca_certs")?,
ca_data: optional_dict_string(values, "ssl_ca_data")?,
client_certificate: optional_dict_string(values, "ssl_certfile")?,
client_key: optional_dict_string(values, "ssl_keyfile")?,
})
}
#[inline(never)]
fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult<CertificateRequirement> {
let Some(value) = values.get_item("ssl_cert_reqs")? else {
return Ok(CertificateRequirement::Required);
};
if value.is_none() {
return Ok(CertificateRequirement::Required);
}
if let Ok(number) = value.extract::<i32>() {
return match number {
0 => Ok(CertificateRequirement::None),
1 => Ok(CertificateRequirement::Optional),
2 => Ok(CertificateRequirement::Required),
_ => Err(PyValueError::new_err(
"invalid Redis TLS certificate requirement",
)),
};
}
let text = value.str()?;
let text = text.to_str()?;
if text.eq_ignore_ascii_case("none") || text.eq_ignore_ascii_case("cert_none") {
return Ok(CertificateRequirement::None);
}
if text.eq_ignore_ascii_case("optional") || text.eq_ignore_ascii_case("cert_optional") {
return Ok(CertificateRequirement::Optional);
}
if text.eq_ignore_ascii_case("required") || text.eq_ignore_ascii_case("cert_required") {
return Ok(CertificateRequirement::Required);
}
Err(PyValueError::new_err(
"invalid Redis TLS certificate requirement",
))
}
#[inline(never)]
fn instance_class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult<bool> {
class_is(value.get_type().as_any(), module, name)
}
#[inline(never)]
fn class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult<bool> {
Ok(value
.getattr("__module__")?
.cast_into::<PyString>()?
.to_str()?
== module
&& value
.getattr("__qualname__")?
.cast_into::<PyString>()?
.to_str()?
== name)
}
#[inline(never)]
fn optional_duration(value: Bound<'_, PyAny>) -> PyResult<Option<Duration>> {
value.extract::<Option<f64>>()?.map(duration).transpose()
}
#[inline(never)]
fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<String>> {
match value.getattr(name) {
Ok(value) => optional_string(value),
Err(error) if error.is_instance_of::<pyo3::exceptions::PyAttributeError>(value.py()) => {
Ok(None)
}
Err(error) => Err(error),
}
}
#[inline(never)]
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
Ok(value
.extract::<Option<String>>()?
.filter(|value| !value.is_empty()))
}
#[inline(never)]
fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult<bool> {
Ok(values.get_item(key)?.is_some_and(|value| !value.is_none()))
}
#[inline(never)]
fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult<String> {
values
.get_item(key)?
.ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))?
.extract::<String>()
}
#[inline(never)]
fn required_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<i64> {
values
.get_item(key)?
.ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))?
.extract::<i64>()
}
#[inline(never)]
fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<String>> {
match values.get_item(key)? {
Some(value) if !value.is_none() => optional_string(value),
_ => Ok(None),
}
}
#[inline(never)]
fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<f64>> {
match values.get_item(key)? {
Some(value) => value.extract::<Option<f64>>(),
None => Ok(None),
}
}
#[inline(never)]
fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<i64>> {
match values.get_item(key)? {
Some(value) => value.extract::<Option<i64>>(),
None => Ok(None),
}
}
#[inline(never)]
fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<bool>> {
match values.get_item(key)? {
Some(value) => value.extract::<Option<bool>>(),
None => Ok(None),
}
}
#[inline(never)]
fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<bool>> {
let Some(value) = values.get_item(key)? else {
return Ok(None);
};
if value.is_none() {
return Ok(None);
}
if let Ok(text) = value.extract::<String>() {
return Ok(Some(
text == "1" || text.eq_ignore_ascii_case("true") || text.eq_ignore_ascii_case("yes"),
));
}
value.extract::<bool>().map(Some)
}
#[inline(never)]
fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<Duration>> {
optional_f64(values, key)?.map(duration).transpose()
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
use pyo3::{prelude::*, types::PyDict};
use super::{
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig,
RedisProtocol,
};
use crate::cache::native::NativeResponseCache;
fn facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> {
let locals = PyDict::new(py);
py.run(
&CString::new(format!(
"from types import SimpleNamespace\n\
ConnectionPool = type('ConnectionPool', (), {{'__module__': 'redis.connection'}})\n\
Connection = type('Connection', (), {{'__module__': 'redis.connection'}})\n\
SSLConnection = type('SSLConnection', (), {{'__module__': 'redis.connection'}})\n\
{body}"
))
.unwrap(),
None,
Some(&locals),
)
.unwrap();
locals.get_item("facade").unwrap().unwrap()
}
#[test]
fn projects_effective_memory_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(default_ttl=913, max_size_in_memory=37, max_size_per_item=8)\n\
facade = SimpleNamespace(type='local', mode='default-on', ttl=11.5, namespace=None, supported_call_types=['completion'], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("memory cache should be supported");
};
assert_eq!(
config.policy.ttl.unwrap(),
std::time::Duration::from_secs_f64(11.5)
);
let CacheBackendConfig::Memory(memory) = config.backend else {
panic!("expected memory configuration");
};
assert_eq!(memory.default_ttl, std::time::Duration::from_secs(913));
assert_eq!(memory.capacity, 37);
assert_eq!(memory.max_entry_bytes, 8192);
let matching =
NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8192);
let mismatched =
NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8191);
let matching_config = NativeCacheConfig {
policy: config.policy,
backend: CacheBackendConfig::Memory(memory),
};
assert_eq!(matching_config.service_mismatch(&matching), None);
assert_eq!(
matching_config.service_mismatch(&mismatched),
Some("facade and native backend item limits must match")
);
});
}
#[test]
fn projects_resolved_redis_tls_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.max_connections = 29\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\
facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Redis cache should be supported");
};
let CacheBackendConfig::Redis(redis) = config.backend else {
panic!("expected Redis configuration");
};
assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777));
assert_eq!(redis.namespace.as_deref(), Some("team"));
assert_eq!(redis.flush_size, 31);
assert_eq!(redis.connection.host, "cache.internal");
assert_eq!(redis.connection.port, 6380);
assert_eq!(redis.connection.database, 4);
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(redis.connection.pool_size, 29);
let tls = redis.connection.tls.unwrap();
assert_eq!(
tls.certificate_requirement,
CertificateRequirement::Optional
);
assert!(tls.check_hostname);
assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem"));
assert_eq!(tls.ca_data.as_deref(), Some("CA DATA"));
assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem"));
assert_eq!(tls.client_key.as_deref(), Some("/client.key"));
});
}
#[test]
fn dynamic_redis_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
}

293
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored Normal file
View file

@ -0,0 +1,293 @@
use litellm_host_python::from_py;
use pyo3::{
PyTraverseError, PyVisit,
exceptions::PyTypeError,
prelude::*,
types::{PyDict, PyTuple, PyType},
};
use serde_json::Value;
use super::{
config::{CacheConfigProjection, NativeCacheConfig},
handle::CacheTestHandle,
native::NativeResponseCache,
};
struct ClassGuard {
class: Py<PyType>,
attributes: Vec<(String, Py<PyAny>)>,
}
struct ObjectGuard {
reference: Py<PyAny>,
classes: Vec<ClassGuard>,
config_names: &'static [&'static str],
config: Vec<Value>,
}
struct RedisPoolGuard {
reference: Py<PyAny>,
connection_class: Py<PyAny>,
connection_kwargs: Py<PyAny>,
max_connections: usize,
}
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
redis_pool: Option<RedisPoolGuard>,
}
impl ObjectGuard {
fn capture(
py: Python<'_>,
object: &Bound<'_, PyAny>,
config_names: &'static [&'static str],
) -> PyResult<Self> {
let classes = object
.get_type()
.getattr("__mro__")?
.cast_into::<PyTuple>()?
.iter()
.map(|class| {
let class = class.cast_into::<PyType>()?;
let attributes = class
.getattr("__dict__")?
.call_method0("items")?
.try_iter()?
.map(|item| item?.extract::<(String, Py<PyAny>)>())
.collect::<PyResult<Vec<_>>>()?;
Ok(ClassGuard {
class: class.unbind(),
attributes,
})
})
.collect::<PyResult<Vec<_>>>()?;
let guard = Self {
reference: py
.import("weakref")?
.getattr("ref")?
.call1((object,))?
.unbind(),
classes,
config_names,
config: Self::config(object, config_names)?,
};
if !guard.matches(py, object)? {
return Err(PyTypeError::new_err(
"native facade registration requires unmodified built-in methods",
));
}
Ok(guard)
}
fn config(object: &Bound<'_, PyAny>, names: &[&str]) -> PyResult<Vec<Value>> {
names
.iter()
.map(|name| match object.getattr(*name) {
Ok(value) => from_py(&value),
Err(error)
if error.is_instance_of::<pyo3::exceptions::PyAttributeError>(object.py()) =>
{
Ok(Value::Null)
}
Err(error) => Err(error),
})
.collect()
}
fn matches(&self, py: Python<'_>, object: &Bound<'_, PyAny>) -> PyResult<bool> {
if !self.reference.bind(py).call0()?.is(object) {
return Ok(false);
}
let mro = object
.get_type()
.getattr("__mro__")?
.cast_into::<PyTuple>()?;
if mro.len() != self.classes.len() {
return Ok(false);
}
let instance = object.getattr("__dict__")?.cast_into::<PyDict>()?;
for (class, expected) in mro.iter().zip(&self.classes) {
if !class.is(expected.class.bind(py)) {
return Ok(false);
}
let attributes = class.getattr("__dict__")?;
if attributes.len()? != expected.attributes.len() {
return Ok(false);
}
for (name, value) in &expected.attributes {
if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) {
return Ok(false);
}
}
}
Ok(Self::config(object, self.config_names)? == self.config)
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)?;
for class in &self.classes {
visit.call(&class.class)?;
for (_, value) in &class.attributes {
visit.call(value)?;
}
}
Ok(())
}
}
impl RedisPoolGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let pool = backend
.getattr("redis_client")?
.getattr("connection_pool")?;
Ok(Self {
reference: pool.clone().unbind(),
connection_class: pool.getattr("connection_class")?.unbind(),
connection_kwargs: pool
.getattr("connection_kwargs")?
.call_method0("copy")?
.unbind(),
max_connections: pool.getattr("max_connections")?.extract::<usize>()?,
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let pool = backend
.getattr("redis_client")?
.getattr("connection_pool")?;
Ok(self.reference.bind(py).is(&pool)
&& self
.connection_class
.bind(py)
.is(&pool.getattr("connection_class")?)
&& self.max_connections == pool.getattr("max_connections")?.extract::<usize>()?
&& self
.connection_kwargs
.bind(py)
.eq(pool.getattr("connection_kwargs")?)?)
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)?;
visit.call(&self.connection_class)?;
visit.call(&self.connection_kwargs)
}
}
impl FacadeGuard {
pub(super) fn capture(
py: Python<'_>,
facade: &Bound<'_, PyAny>,
service: &NativeResponseCache,
) -> PyResult<Self> {
let kind = service.kind();
let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?;
if !facade.get_type().is(&cache_type) {
return Err(PyTypeError::new_err(
"only exact built-in Cache facades can be registered",
));
}
let (module, name, cache_kind) = match kind {
"memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
"redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
if facade.getattr("type")?.extract::<String>()? != cache_kind
|| !backend.get_type().is(&py.import(module)?.getattr(name)?)
{
return Err(PyTypeError::new_err(
"facade and native backend types must match",
));
}
let config = match NativeCacheConfig::project(facade)? {
CacheConfigProjection::Native(config) => *config,
CacheConfigProjection::Unsupported(reason) => {
return Err(PyTypeError::new_err(reason.message()));
}
};
if let Some(message) = config.service_mismatch(service) {
return Err(PyTypeError::new_err(message));
}
Ok(Self {
outer: ObjectGuard::capture(
py,
facade,
&[
"type",
"mode",
"ttl",
"namespace",
"supported_call_types",
"redis_flush_size",
"semantic_cache_scope",
],
)?,
backend: ObjectGuard::capture(
py,
&backend,
&[
"namespace",
"default_ttl",
"max_size_in_memory",
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
],
)?,
redis_pool: (kind == "redis")
.then(|| RedisPoolGuard::capture(&backend))
.transpose()?,
})
}
fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<bool> {
if !self.outer.matches(py, facade)? {
return Ok(false);
}
let backend = facade.getattr("cache")?;
if !self.backend.matches(py, &backend)? {
return Ok(false);
}
match &self.redis_pool {
Some(guard) => guard.matches(py, &backend),
None => Ok(true),
}
}
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
self.outer.traverse(&visit)?;
self.backend.traverse(&visit)?;
if let Some(guard) = &self.redis_pool {
guard.traverse(&visit)?;
}
Ok(())
}
}
pub(super) fn resolve(
py: Python<'_>,
facade: &Bound<'_, PyAny>,
) -> PyResult<Option<NativeResponseCache>> {
let Ok(dict) = facade
.getattr("__dict__")
.and_then(|dict| dict.cast_into::<PyDict>().map_err(Into::into))
else {
return Ok(None);
};
let Some(handle) = dict.get_item("_native_cache_handle")? else {
return Ok(None);
};
let Ok(handle) = handle.extract::<PyRef<'_, CacheTestHandle>>() else {
return Ok(None);
};
let Some(guard) = &handle.guard else {
return Ok(None);
};
if !guard.matches(py, facade).unwrap_or(false) {
return Ok(None);
}
handle.service().map(Some)
}

18
litellm-rust/crates/python-bridge/src/cache/future.rs vendored Normal file
View file

@ -0,0 +1,18 @@
use litellm_host_python::to_py;
use pyo3::prelude::*;
pub(super) fn ready_none(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
ready_value(py, &())
}
pub(super) fn ready_value<'py, T: serde::Serialize>(
py: Python<'py>,
value: &T,
) -> PyResult<Bound<'py, PyAny>> {
let future = py
.import("asyncio")?
.call_method0("get_running_loop")?
.call_method0("create_future")?;
future.call_method1("set_result", (to_py(py, value)?,))?;
Ok(future)
}

84
litellm-rust/crates/python-bridge/src/cache/handle.rs vendored Normal file
View file

@ -0,0 +1,84 @@
use litellm_host_python::release_gil;
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
#[pyclass(frozen, name = "_CacheTestHandle")]
pub(crate) struct CacheTestHandle {
service: NativeResponseCache,
pub(super) guard: Option<FacadeGuard>,
pid: u32,
}
impl CacheTestHandle {
pub(super) fn service(&self) -> PyResult<NativeResponseCache> {
if self.pid != std::process::id() {
return Err(PyRuntimeError::new_err(
"native cache handles must be recreated after fork",
));
}
Ok(self.service.clone())
}
}
#[pymethods]
impl CacheTestHandle {
#[staticmethod]
#[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))]
fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult<Self> {
Ok(Self {
service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes),
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (url, *, ttl_seconds=60.0, namespace=None))]
fn redis(
py: Python<'_>,
url: String,
ttl_seconds: f64,
namespace: Option<String>,
) -> PyResult<Self> {
let ttl = Some(duration(ttl_seconds)?);
let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace))
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()
}
fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
let service = self.service()?;
let guard = FacadeGuard::capture(py, facade, &service)?;
let service = service.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
let handle = Py::new(
py,
Self {
service,
guard: Some(guard),
pid: self.pid,
},
)?;
facade.setattr("_native_cache_handle", handle)
}
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
if let Some(guard) = &self.guard {
guard.traverse(visit)?;
}
Ok(())
}
}

26
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored Normal file
View file

@ -0,0 +1,26 @@
mod binding;
mod callback;
mod config;
mod facade;
mod future;
mod handle;
mod native;
mod request;
mod resolver;
use litellm_cache::Error;
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
prelude::*,
};
pub(crate) use self::{
binding::ResolvedCache, handle::CacheTestHandle, resolver::CacheTestResolver,
};
fn cache_error(error: Error) -> PyErr {
match error {
Error::InvalidEntry => PyValueError::new_err(error.to_string()),
_ => PyRuntimeError::new_err(error.to_string()),
}
}

198
litellm-rust/crates/python-bridge/src/cache/native.rs vendored Normal file
View file

@ -0,0 +1,198 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::RedisCache;
use litellm_cache_response::{
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
};
use serde_json::Value;
#[derive(Clone)]
pub(super) enum NativeResponseCache {
Memory(Arc<ResponseCache<InMemoryCache<CacheEntry>>>),
Redis {
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
buffer: Option<Arc<WriteBuffer>>,
},
}
impl NativeResponseCache {
pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
Self::Memory(Arc::new(ResponseCache::new(Arc::new(
InMemoryCache::with_clock_and_size_measurement(
Some(capacity),
Some(ttl),
Some(max_entry_bytes),
Some(Arc::new(|entry| {
ResponseCacheCodec.encode(entry).map(|bytes| bytes.len())
})),
super::request::now,
),
))))
}
pub fn redis(
url: &str,
ttl: Option<Duration>,
namespace: Option<String>,
) -> Result<Self, Error> {
let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace);
Ok(Self::Redis {
cache: Arc::new(ResponseCache::new(Arc::new(backend))),
buffer: None,
})
}
}
impl NativeResponseCache {
pub fn kind(&self) -> &'static str {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
}
}
pub fn default_ttl(&self) -> Option<Duration> {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
}
}
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } => None,
}
}
pub fn with_redis_flush_size(self, flush_size: Option<usize>) -> Self {
match self {
Self::Redis { cache, .. } => Self::Redis {
cache,
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
},
memory => memory,
}
}
pub fn lookup(
&self,
request: &ResponseCacheRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
match self {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
}
}
pub fn store(
&self,
request: &ResponseCacheRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
}
}
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
now: Duration,
) -> Result<PartialHits, Error> {
match self {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
}
}
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
match self {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
}
}
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_store(request, response, now).await,
Self::Redis {
cache,
buffer: None,
} => cache.async_store(request, response, now).await,
Self::Redis {
cache,
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
}
}
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
now: Duration,
) -> Result<PartialHits, Error> {
match self {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
}
}
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
}
}
pub async fn async_flush(&self) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_flush().await,
Self::Redis { cache, buffer } => {
if let Some(buffer) = buffer {
buffer.clear()?;
}
cache.async_flush().await
}
}
}
pub async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match self {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
}
}
}

48
litellm-rust/crates/python-bridge/src/cache/request.rs vendored Normal file
View file

@ -0,0 +1,48 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest};
use litellm_host_python::from_py;
use pyo3::{exceptions::PyValueError, prelude::*};
use serde::Deserialize;
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RequestInput {
key: CacheKeyInput,
controls: Option<CacheControls>,
ttl_seconds: Option<f64>,
max_age_seconds: Option<f64>,
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<ResponseCacheRequest> {
let input: RequestInput = from_py(value)?;
request_input(input)
}
fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest> {
let mut request = ResponseCacheRequest::new(input.key);
if let Some(controls) = input.controls {
request.controls = controls;
}
request.context.ttl = input.ttl_seconds.map(duration).transpose()?;
request.max_age = input.max_age_seconds.map(duration).transpose()?;
Ok(request)
}
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<ResponseCacheRequest>> {
from_py::<Vec<RequestInput>>(value)?
.into_iter()
.map(request_input)
.collect()
}
pub(super) fn duration(seconds: f64) -> PyResult<Duration> {
Duration::try_from_secs_f64(seconds)
.map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative"))
}
pub(super) fn now() -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
}

39
litellm-rust/crates/python-bridge/src/cache/resolver.rs vendored Normal file
View file

@ -0,0 +1,39 @@
use pyo3::{PyTraverseError, PyVisit, prelude::*};
use super::{
binding::{CacheBinding, ResolvedCache},
callback::PythonCallback,
facade,
handle::CacheTestHandle,
};
#[pyclass(frozen, name = "_CacheTestResolver")]
pub(crate) struct CacheTestResolver {
namespace: Py<PyAny>,
}
#[pymethods]
impl CacheTestResolver {
#[new]
fn new(namespace: Py<PyAny>) -> Self {
Self { namespace }
}
pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult<ResolvedCache> {
let object = self.namespace.bind(py).getattr("cache")?;
let binding = if object.is_none() {
CacheBinding::Disabled
} else if let Ok(handle) = object.extract::<PyRef<'_, CacheTestHandle>>() {
CacheBinding::Native(handle.service()?)
} else if let Some(service) = facade::resolve(py, &object)? {
CacheBinding::Native(service)
} else {
CacheBinding::PythonCallback(PythonCallback::new(object.unbind()))
};
Ok(ResolvedCache::new(binding))
}
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.namespace)
}
}

12
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -1,3 +1,4 @@
mod cache;
mod credentials;
mod diagnostics;
mod errors;
@ -9,6 +10,7 @@ mod token_counter;
#[pymodule(gil_used = true)]
mod _native {
use crate::cache::{CacheTestHandle, CacheTestResolver, ResolvedCache};
#[cfg(feature = "panic-test")]
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
@ -32,6 +34,16 @@ mod _native {
use crate::token_counter::TokenCounter;
#[pymodule_export]
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
use pyo3::{prelude::*, types::PyModule};
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
let py = module.py();
let dict = module.dict();
dict.set_item("_CacheTestHandle", py.get_type::<CacheTestHandle>())?;
dict.set_item("_CacheTestResolver", py.get_type::<CacheTestResolver>())?;
dict.set_item("_CacheTestBinding", py.get_type::<ResolvedCache>())
}
}
use pyo3::prelude::*;

3
litellm/__init__.py
View file

@ -1684,6 +1684,9 @@ if TYPE_CHECKING:
from .llms.bedrock.messages.mantle_transformation import (
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
)
from .llms.bedrock_mantle.messages.transformation import (
BedrockMantleAnthropicMessagesConfig as BedrockMantleAnthropicMessagesConfig,
)
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
from .llms.together_ai.chat.transformation import (
TogetherAIChatConfig as TogetherAIChatConfig,

5
litellm/_lazy_imports_registry.py
View file

@ -176,6 +176,7 @@ LLM_CONFIG_NAMES: Final = (
"BedrockClaudePlatformMessagesConfig",
"AmazonAnthropicClaudeMessagesConfig",
"AmazonMantleMessagesConfig",
"BedrockMantleAnthropicMessagesConfig",
"TogetherAIConfig",
"TogetherAIChatConfig",
"NLPCloudConfig",
@ -746,6 +747,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.bedrock.messages.mantle_transformation",
"AmazonMantleMessagesConfig",
),
"BedrockMantleAnthropicMessagesConfig": (
".llms.bedrock_mantle.messages.transformation",
"BedrockMantleAnthropicMessagesConfig",
),
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
"TogetherAIChatConfig": (
".llms.together_ai.chat.transformation",

35
litellm/anthropic_beta_headers_config.json
View file

@ -131,6 +131,41 @@
"web-fetch-2025-09-10": null,
"web-search-2025-03-05": null
},
"bedrock_mantle": {
"advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19",
"advisor-tool-2026-03-01": null,
"bash_20241022": null,
"bash_20250124": null,
"claude-code-20250219": "claude-code-20250219",
"code-execution-2025-08-25": null,
"compact-2026-01-12": "compact-2026-01-12",
"computer-use-2025-01-24": "computer-use-2025-01-24",
"computer-use-2025-11-24": "computer-use-2025-11-24",
"context-1m-2025-08-07": "context-1m-2025-08-07",
"context-management-2025-06-27": "context-management-2025-06-27",
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
"mcp-client-2025-04-04": null,
"mcp-client-2025-11-20": null,
"mcp-servers-2025-12-04": null,
"output-128k-2025-02-19": "output-128k-2025-02-19",
"per-turn-control-2026-07-01": "per-turn-control-2026-07-01",
"prompt-caching-scope-2026-01-05": null,
"skills-2025-10-02": null,
"structured-output-2024-03-01": null,
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"tool-examples-2025-10-29": "tool-examples-2025-10-29",
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
"web-fetch-2025-09-10": null,
"web-search-2025-03-05": "web-search-2025-03-05"
},
"vertex_ai": {
"advisor-tool-2026-03-01": null,
"advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19",

2
litellm/anthropic_beta_headers_manager.py
View file

@ -334,7 +334,7 @@ def update_headers_with_filtered_beta(
Updated headers dict
"""
existing_beta: Final = headers.get("anthropic-beta")
if not existing_beta:
if existing_beta is None:
return headers
# Parse existing beta headers

1
litellm/constants.py
View file

@ -402,6 +402,7 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = (
if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None
else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
)
PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20
DEFAULT_TRIM_RATIO: Final = float(
os.getenv("DEFAULT_TRIM_RATIO", 0.75)
) # default ratio of tokens to trim from the end of a prompt

60
litellm/experimental_mcp_client/client.py
View file

@ -7,12 +7,13 @@ import base64
import hashlib
import json
import os
from collections.abc import Awaitable, Callable, Generator
from collections.abc import Awaitable, Callable, Generator, Sequence
from contextlib import AbstractAsyncContextManager
from functools import partial
from types import MappingProxyType
from typing import Any, Final, TypeAlias, TypeVar
import anyio
import httpx2
from httpx2._client import UseClientDefault
from httpx2._types import AuthTypes
@ -38,6 +39,8 @@ from mcp.types import (
ListPromptsResult,
ListResourcesResult,
ListResourceTemplatesResult,
PaginatedRequestParams,
PaginatedResult,
Prompt,
ResourceTemplate,
ServerNotification,
@ -49,7 +52,12 @@ from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
from litellm._logging import verbose_logger
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT
from litellm.constants import (
MCP_CLIENT_TIMEOUT,
MCP_NPM_CACHE_DIR,
MCP_TOOL_LISTING_MAX_PAGES,
MCP_TOOL_LISTING_TIMEOUT,
)
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response
@ -147,6 +155,8 @@ def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None:
TSessionResult = TypeVar("TSessionResult")
_ListPage = TypeVar("_ListPage", bound=PaginatedResult)
_ListItem = TypeVar("_ListItem")
class _MCPHTTPClient(httpx2.AsyncClient):
@ -793,6 +803,33 @@ class MCPClient:
# Return a default error result instead of raising
return self.error_tool_result(e)
async def _list_optional_pages(
self,
fetch_page: Callable[[PaginatedRequestParams | None], Awaitable[_ListPage]],
items_of: Callable[[_ListPage], Sequence[_ListItem]],
) -> list[_ListItem]: # mutable-ok: existing list discovery API
items: Final[list[_ListItem]] = [] # mutable-ok: bounded iterative page accumulation
cursors: Final[set[str]] = set() # mutable-ok: constant-time detection of cursor cycles
cursor: str | None = None # rebind-ok: iterative traversal avoids recursion at the existing page cap
with anyio.fail_after(max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)):
for page_index in range(MCP_TOOL_LISTING_MAX_PAGES):
try:
page = await fetch_page( # rebind-ok: each SDK page replaces the previous one
None if cursor is None else PaginatedRequestParams(cursor=cursor)
)
except MCPError as error:
if page_index > 0 and error.error.code == METHOD_NOT_FOUND:
raise RuntimeError("MCP list operation became unavailable during pagination") from error
raise
items.extend(items_of(page))
if not page.next_cursor:
return items
if page.next_cursor in cursors:
raise RuntimeError("MCP list pagination repeated a cursor")
cursors.add(page.next_cursor)
cursor = page.next_cursor
raise RuntimeError(f"MCP list pagination exceeded {MCP_TOOL_LISTING_MAX_PAGES} pages")
async def list_prompts(self, *, raise_on_error: bool = False) -> list[Prompt]:
"""List available prompts from the server."""
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
@ -802,7 +839,11 @@ class MCPClient:
if capabilities is not None and capabilities.prompts is None:
return ListPromptsResult(prompts=[])
try:
return await session.list_prompts()
return ListPromptsResult(
prompts=await self._list_optional_pages(
lambda params: session.list_prompts(params=params), lambda page: page.prompts
)
)
except MCPError as error:
if error.error.code != METHOD_NOT_FOUND:
raise
@ -892,7 +933,11 @@ class MCPClient:
if capabilities is not None and capabilities.resources is None:
return ListResourcesResult(resources=[])
try:
return await session.list_resources()
return ListResourcesResult(
resources=await self._list_optional_pages(
lambda params: session.list_resources(params=params), lambda page: page.resources
)
)
except MCPError as error:
if error.error.code != METHOD_NOT_FOUND:
raise
@ -941,7 +986,12 @@ class MCPClient:
if capabilities is not None and capabilities.resources is None:
return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload
try:
return await session.list_resource_templates()
return ListResourceTemplatesResult(
resource_templates=await self._list_optional_pages(
lambda params: session.list_resource_templates(params=params),
lambda page: page.resource_templates,
)
)
except MCPError as error:
if error.error.code != METHOD_NOT_FOUND:
raise

11
litellm/litellm_core_utils/token_counter.py
View file

@ -46,6 +46,8 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionDocumentObject,
ChatCompletionNamedToolChoiceParam,
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
ChatCompletionToolParam,
OpenAIMessageContentListBlock,
)
@ -854,6 +856,8 @@ def _count_content_list(
content_list: str
| Iterable[
OpenAIMessageContentListBlock
| ChatCompletionThinkingBlock
| ChatCompletionRedactedThinkingBlock
| AnthropicMessagesTextParam
| AnthropicMessagesImageParam
| AnthropicMessagesDocumentParam
@ -898,9 +902,9 @@ def _count_content_list(
use_default_image_token_count,
default_token_count,
)
elif c["type"] == "thinking":
elif c["type"] in ("thinking", "redacted_thinking"):
# Claude extended thinking content block
# Count the thinking text and skip signature (opaque signature blob)
# Count the thinking text and skip the opaque blobs (signature, redacted data)
thinking_text = str(c.get("thinking", ""))
if thinking_text:
num_tokens += count_function(thinking_text)
@ -920,7 +924,8 @@ def _count_content_list(
raise ValueError(
f"Invalid content item type: {content_type}. "
f"Expected str or dict with 'type' field "
f"(text, image_url, image, document, file, tool_use, tool_result, thinking, tool_reference)."
f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, "
f"tool_reference)."
)
return num_tokens
except Exception as e:

7
litellm/llms/anthropic/experimental_pass_through/messages/handler.py
View file

@ -651,6 +651,11 @@ def anthropic_messages_handler(
"display": "summarized",
}
resolved_api_base: Final = (
dynamic_api_base
if dynamic_api_base is not None and anthropic_messages_provider_config.uses_get_llm_provider_api_base()
else api_base
)
return base_llm_http_handler.anthropic_messages_handler(
model=model,
messages=strip_provider_specific_fields_from_anthropic_messages(messages),
@ -662,7 +667,7 @@ def anthropic_messages_handler(
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
api_key=api_key,
api_base=api_base,
api_base=resolved_api_base,
stream=stream,
kwargs=kwargs,
)

3
litellm/llms/base_llm/anthropic_messages/transformation.py
View file

@ -128,6 +128,9 @@ class BaseAnthropicMessagesConfig(ABC):
"""
return True
def uses_get_llm_provider_api_base(self) -> bool:
return False
def get_async_streaming_response_iterator(
self,
model: str,

51
litellm/llms/bedrock/batches/transformation.py
View file

@ -33,6 +33,7 @@ from ..base_aws_llm import BaseAWSLLM
from ..common_utils import (
CommonBatchFilesUtils,
merge_bedrock_aws_request_params,
resolve_s3_bucket_owner,
resolve_s3_encryption_key_id,
)
@ -51,6 +52,26 @@ _S3_BATCH_FILE_UUID_SUFFIX_PATTERN: Final = re.compile(
_BEDROCK_TAGS_ADAPTER: Final[TypeAdapter[list[BedrockTag]]] = TypeAdapter(list[BedrockTag])
def _build_s3_input_config(s3_uri: str, s3_bucket_owner: str | None) -> BedrockS3InputDataConfig:
if s3_bucket_owner is None:
return BedrockS3InputDataConfig(s3Uri=s3_uri)
return BedrockS3InputDataConfig(s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner)
def _build_s3_output_config(
s3_uri: str, s3_bucket_owner: str | None, s3_encryption_key_id: str | None
) -> BedrockS3OutputDataConfig:
if s3_bucket_owner is None:
if s3_encryption_key_id is None:
return BedrockS3OutputDataConfig(s3Uri=s3_uri)
return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3EncryptionKeyId=s3_encryption_key_id)
if s3_encryption_key_id is None:
return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner)
return BedrockS3OutputDataConfig(
s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner, s3EncryptionKeyId=s3_encryption_key_id
)
def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]:
try:
return _BEDROCK_TAGS_ADAPTER.validate_python(raw_tags, strict=True)
@ -214,25 +235,23 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
job_name: Final = self.common_utils.generate_unique_job_name(model, prefix="litellm")
output_key: Final = f"litellm-batch-outputs/{job_name}/"
# Build input data config
input_data_config: Final[BedrockInputDataConfig] = {
"s3InputDataConfig": BedrockS3InputDataConfig(s3Uri=f"s3://{input_bucket}/{input_key}")
}
# Build output data config
s3_output_config: Final[BedrockS3OutputDataConfig] = BedrockS3OutputDataConfig(
s3Uri=f"s3://{output_bucket}/{output_key}"
)
# Add optional KMS encryption key ID if provided
s3_encryption_key_id = resolve_s3_encryption_key_id(
s3_bucket_owner: Final = resolve_s3_bucket_owner(litellm_params=litellm_params, optional_params=optional_params)
s3_encryption_key_id: Final = resolve_s3_encryption_key_id(
litellm_params=litellm_params,
optional_params=optional_params,
)
if s3_encryption_key_id:
s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id
output_data_config: Final[BedrockOutputDataConfig] = {"s3OutputDataConfig": s3_output_config}
input_data_config: Final[BedrockInputDataConfig] = {
"s3InputDataConfig": _build_s3_input_config(
s3_uri=f"s3://{input_bucket}/{input_key}", s3_bucket_owner=s3_bucket_owner
)
}
output_data_config: Final[BedrockOutputDataConfig] = {
"s3OutputDataConfig": _build_s3_output_config(
s3_uri=f"s3://{output_bucket}/{output_key}",
s3_bucket_owner=s3_bucket_owner,
s3_encryption_key_id=s3_encryption_key_id,
)
}
# Create Bedrock batch request with proper typing
bedrock_request: Final[BedrockCreateBatchRequest] = {

26
litellm/llms/bedrock/common_utils.py
View file

@ -1555,11 +1555,33 @@ def resolve_s3_encryption_key_id(
Precedence: `s3_encryption_key_id` in litellm_params, then optional_params
(client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var.
"""
return _resolve_s3_setting("s3_encryption_key_id", "AWS_S3_ENCRYPTION_KEY_ID", litellm_params, optional_params)
def resolve_s3_bucket_owner(
litellm_params: Mapping[str, object],
optional_params: Mapping[str, object] | None = None,
) -> str | None:
"""
Resolve the AWS account id that owns the S3 buckets used by Bedrock batch jobs.
Precedence: `s3_bucket_owner` in litellm_params, then optional_params
(client-side / request params), then the AWS_S3_BUCKET_OWNER env var.
"""
return _resolve_s3_setting("s3_bucket_owner", "AWS_S3_BUCKET_OWNER", litellm_params, optional_params)
def _resolve_s3_setting(
param_name: str,
env_var: str,
litellm_params: Mapping[str, object],
optional_params: Mapping[str, object] | None,
) -> str | None:
candidates: Final = tuple(
source.get("s3_encryption_key_id") for source in (litellm_params, optional_params) if source is not None
source.get(param_name) for source in (litellm_params, optional_params) if source is not None
)
explicit: Final = next((value for value in candidates if isinstance(value, str) and value), None)
return explicit or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID")
return explicit or get_secret_str(env_var)
class CommonBatchFilesUtils:

22
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -1,4 +1,4 @@
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
@ -445,13 +445,16 @@ class AmazonAnthropicClaudeMessagesConfig(
# Bedrock InvokeModel DOES support ``clear_tool_uses_20250919`` under the
# ``context-management-2025-06-27`` beta. AWS docs:
# https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-tool-use.md
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: dict[str, str] = {
"compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
"clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
}
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType(
{
"compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
"clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
}
)
@staticmethod
@classmethod
def _filter_context_management_for_bedrock_invoke(
cls,
anthropic_messages_request: dict,
beta_set: set,
) -> None:
@ -481,7 +484,7 @@ class AmazonAnthropicClaudeMessagesConfig(
anthropic_messages_request.pop("context_management", None)
return
supported: Final = AmazonAnthropicClaudeMessagesConfig._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS
supported: Final = cls._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS
retained_edits: Final = [e for e in edits if isinstance(e, dict) and e.get("type") in supported]
if not retained_edits:
anthropic_messages_request.pop("context_management", None)
@ -546,15 +549,16 @@ class AmazonAnthropicClaudeMessagesConfig(
if "tool-search-tool-2025-10-19" in beta_set:
beta_set.add("tool-examples-2025-10-29")
beta_provider: Final = self.custom_llm_provider or "bedrock"
filtered_betas: Final = sorted(
filter_and_transform_beta_headers(
beta_headers=list(beta_set),
provider="bedrock",
provider=beta_provider,
)
)
dropped_user_betas: Final = sorted(
b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider="bedrock")
b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider=beta_provider)
)
if dropped_user_betas:
verbose_logger.warning(

0
litellm/llms/bedrock_mantle/messages/__init__.py Normal file
View file

127
litellm/llms/bedrock_mantle/messages/transformation.py Normal file
View file

@ -0,0 +1,127 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
DEFAULT_ANTHROPIC_API_VERSION,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import MANTLE_MESSAGES_PATH
from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig
from litellm.llms.bedrock_mantle.common_utils import (
MANTLE_HOST_RE,
BedrockMantleAuthMixin,
resolve_mantle_region,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.router import GenericLiteLLMParams
_BASE_SUFFIXES_TO_STRIP: Final = (
MANTLE_MESSAGES_PATH,
"/v1/messages",
"/messages",
"/anthropic/v1",
"/openai/v1",
"/v1",
)
_BODY_FIELDS_MANTLE_READS_FROM_HEADERS: Final = frozenset({"anthropic_version", "anthropic_beta"})
_ANTHROPIC_BETAS: Final = TypeAdapter(tuple[str, ...])
_MANTLE_REQUEST: Final = TypeAdapter(dict[str, object])
def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mapping[str, object]) -> str:
region: Final = resolve_mantle_region(MappingProxyType({**litellm_params, "api_base": api_base}))
configured: Final = (
api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") or f"https://bedrock-mantle.{region}.api.aws"
).rstrip("/")
stripped: Final = next(
(configured[: -len(suffix)] for suffix in _BASE_SUFFIXES_TO_STRIP if configured.endswith(suffix)),
configured,
)
host: Final = f"https://bedrock-mantle.{region}.api.aws" if MANTLE_HOST_RE.match(stripped) else stripped
return f"{host}{MANTLE_MESSAGES_PATH}"
class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleMessagesConfig):
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType(
{
**AmazonMantleMessagesConfig._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS,
"clear_thinking_20251015": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
}
)
def __init__(self, aws_signer: BaseAWSLLM | None = None) -> None:
AmazonMantleMessagesConfig.__init__(self)
self._aws_signer = aws_signer or self
@property
def custom_llm_provider(self) -> str | None:
return "bedrock_mantle"
def uses_get_llm_provider_api_base(self) -> bool:
return True
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict,
litellm_params: dict,
stream: bool | None = None,
) -> str:
return build_mantle_native_messages_url(api_base=api_base, litellm_params=litellm_params)
def validate_anthropic_messages_environment(
self,
headers: dict,
model: str,
messages: list[dict],
optional_params: dict,
litellm_params: dict,
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict, str | None]:
merged_headers, resolved_api_base = super().validate_anthropic_messages_environment(
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=api_key,
api_base=api_base,
)
if any(name.lower() == "anthropic-version" for name in merged_headers):
return merged_headers, resolved_api_base
return { # mutable-ok: the base class contract returns a dict the handler signs into in place
**merged_headers,
"anthropic-version": DEFAULT_ANTHROPIC_API_VERSION,
}, resolved_api_base
def transform_anthropic_messages_request(
self,
model: str,
messages: list[dict],
anthropic_messages_optional_request_params: dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> dict:
request: Final = _MANTLE_REQUEST.validate_python(
super().transform_anthropic_messages_request(
model=model,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
litellm_params=litellm_params,
headers=headers,
),
)
betas: Final = request.get("anthropic_beta")
if betas is not None:
header_betas: Final = ",".join(_ANTHROPIC_BETAS.validate_python(betas))
headers["anthropic-beta"] = header_betas # rebind-ok: the handler signs and sends this same dict
return { # mutable-ok: the base class contract returns the dict the handler serializes as the body
key: value for key, value in request.items() if key not in _BODY_FIELDS_MANTLE_READS_FROM_HEADERS
}

2
litellm/llms/dashscope/common_utils.py
View file

@ -103,7 +103,7 @@ def missing_dashscope_family_key_message(custom_llm_provider: str) -> str:
)
if custom_llm_provider == "qwen_ai_platform":
return (
"Missing API key for Qwen AI Platform. Set QWEN_AI_PLATFORM_API_KEY or "
"Missing API key for Qianwen AI Platform. Set QWEN_AI_PLATFORM_API_KEY or "
"DASHSCOPE_API_KEY environment variable or pass api_key parameter."
)
return "Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter."

2
litellm/llms/dashscope/qwen_ai_platform.py
View file

@ -23,7 +23,7 @@ def _require_qwen_ai_platform_api_key(api_key: str | None) -> str:
resolved: Final = _resolve_qwen_ai_platform_api_key(api_key)
if resolved is None:
raise ValueError(
"Qwen AI Platform API key is required. Set 'QWEN_AI_PLATFORM_API_KEY' or 'DASHSCOPE_API_KEY' env var "
"Qianwen AI Platform API key is required. Set 'QWEN_AI_PLATFORM_API_KEY' or 'DASHSCOPE_API_KEY' env var "
"or pass api_key explicitly."
)
return resolved

126
litellm/llms/fal_ai/cost_calculator.py
View file

@ -1,12 +1,16 @@
from collections.abc import Mapping
from math import ceil
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter
import litellm
from litellm.types.utils import ImageResponse
from litellm.types.utils import ImageObject, ImageResponse
FAL_KEYED_PRICING_DEFAULT_QUALITY: Final[str] = "high"
FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = "1024-x-768"
FAL_PIXELS_PER_MEGAPIXEL: Final[int] = 1_048_576
FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType(
{
"square_hd": "1024-x-1024",
@ -18,14 +22,17 @@ FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType(
}
)
_OBJECT_MAP: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _keyed_size(optional_params: Mapping[str, object]) -> str | None:
image_size: Final = optional_params.get("image_size")
if image_size is None or image_size == "auto":
return FAL_TEXT_TO_IMAGE_DEFAULT_SIZE
if isinstance(image_size, Mapping):
width: Final = image_size.get("width")
height: Final = image_size.get("height")
image_size_map: Final = _OBJECT_MAP.validate_python(image_size)
width: Final = image_size_map.get("width")
height: Final = image_size_map.get("height")
if isinstance(width, int) and isinstance(height, int):
return f"{width}-x-{height}"
return None
@ -34,21 +41,71 @@ def _keyed_size(optional_params: Mapping[str, object]) -> str | None:
return None
def _keyed_cost_per_image(model: str, optional_params: Mapping[str, object] | None) -> float | None:
if optional_params is None:
def _image_dimensions(image: object) -> tuple[int, int] | None:
if not isinstance(image, ImageObject):
return None
size: Final = _keyed_size(optional_params)
if size is None:
raw_provider_specific_fields: Final = image.provider_specific_fields
if not isinstance(raw_provider_specific_fields, Mapping):
return None
provider_specific_fields: Final = _OBJECT_MAP.validate_python(raw_provider_specific_fields)
width: Final = provider_specific_fields.get("width")
height: Final = provider_specific_fields.get("height")
if type(width) is not int or width <= 0 or type(height) is not int or height <= 0:
return None
return width, height
def _response_size(image: object) -> str | None:
dimensions: Final = _image_dimensions(image)
if dimensions is None:
return None
width, height = dimensions
return f"{width}-x-{height}"
def _keyed_quality(optional_params: Mapping[str, object]) -> str:
raw_quality: Final = optional_params.get("quality")
quality: Final = (
raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY
)
keyed_entry: Final = litellm.model_cost.get(f"fal_ai/{quality}/{size}/{model}")
if keyed_entry is None:
return raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY
def _keyed_cost_per_image(
model: str,
image: object,
optional_params: Mapping[str, object],
) -> float | None:
quality: Final = _keyed_quality(optional_params)
request_size: Final = _keyed_size(optional_params) or FAL_TEXT_TO_IMAGE_DEFAULT_SIZE
sizes: Final = (_response_size(image), request_size, FAL_TEXT_TO_IMAGE_DEFAULT_SIZE)
for size in sizes:
if size is None:
continue
keyed_entry = _entry(f"fal_ai/{quality}/{size}/{model}")
if keyed_entry is None:
continue
keyed_cost = keyed_entry.get("output_cost_per_image")
if isinstance(keyed_cost, (int, float)):
return float(keyed_cost)
return None
def _flat_cost_per_image(
image: object,
output_cost_per_image: float,
output_cost_per_pixel: float | None,
) -> float:
dimensions: Final = _image_dimensions(image)
if dimensions is None or output_cost_per_pixel is None:
return output_cost_per_image
width, height = dimensions
megapixels: Final = ceil(width * height / FAL_PIXELS_PER_MEGAPIXEL)
return output_cost_per_pixel * FAL_PIXELS_PER_MEGAPIXEL * megapixels
def _entry(key: str) -> Mapping[str, object] | None:
raw_entry: Final[object] = litellm.model_cost.get(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # global catalog is untyped
if not isinstance(raw_entry, Mapping):
return None
keyed_cost: Final = keyed_entry.get("output_cost_per_image")
return float(keyed_cost) if isinstance(keyed_cost, (int, float)) else None
return _OBJECT_MAP.validate_python(raw_entry)
def cost_calculator(
@ -61,15 +118,36 @@ def cost_calculator(
"""
if not isinstance(image_response, ImageResponse):
raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}")
# the proxy cost path passes the provider-prefixed model name
model = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/")
num_images: Final[int] = len(image_response.data) if image_response.data else 0
keyed_cost_per_image: Final = _keyed_cost_per_image(model=model, optional_params=optional_params)
if keyed_cost_per_image is not None:
return keyed_cost_per_image * num_images
_model_info: Final = litellm.get_model_info(
model=model,
normalized_model: Final = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/")
params: Final[Mapping[str, object]] = optional_params or MappingProxyType({})
images: Final = tuple(image_response.data or ())
keyed_costs: Final = tuple(
_keyed_cost_per_image(
model=normalized_model,
image=image,
optional_params=params,
)
for image in images
)
if all(cost is not None for cost in keyed_costs):
return sum(cost for cost in keyed_costs if cost is not None)
model_info: Final = litellm.get_model_info(
model=normalized_model,
custom_llm_provider=litellm.LlmProviders.FAL_AI.value,
)
output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0
return output_cost_per_image * num_images
raw_output_cost_per_image: Final = model_info.get("output_cost_per_image")
output_cost_per_image: Final = (
float(raw_output_cost_per_image) if isinstance(raw_output_cost_per_image, (int, float)) else 0.0
)
raw_output_cost_per_pixel: Final = model_info.get("output_cost_per_pixel")
output_cost_per_pixel: Final = (
float(raw_output_cost_per_pixel) if isinstance(raw_output_cost_per_pixel, (int, float)) else None
)
return sum(
_flat_cost_per_image(
image=image,
output_cost_per_image=output_cost_per_image,
output_cost_per_pixel=output_cost_per_pixel,
)
for image in images
)

23
litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py
View file

@ -3,9 +3,9 @@ from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
from litellm.types.utils import ImageObject, ImageResponse
from litellm.types.utils import ImageResponse
from .transformation import FalAIBaseConfig
from .transformation import FalAIBaseConfig, fal_images_to_image_objects
if TYPE_CHECKING:
import tiktoken
@ -229,25 +229,8 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig):
if not model_response.data:
model_response.data = []
# Handle Flux Pro v1.1-ultra response format
images: Final = response_data.get("images", [])
if isinstance(images, list):
for image_data in images:
if isinstance(image_data, dict):
model_response.data.append(
ImageObject(
url=image_data.get("url", None),
b64_json=None, # Flux Pro returns URLs only
)
)
elif isinstance(image_data, str):
# If images is just a list of URLs
model_response.data.append(
ImageObject(
url=image_data,
b64_json=None,
)
)
model_response.data.extend(fal_images_to_image_objects(images))
# Add additional metadata from Flux Pro response
if hasattr(model_response, "_hidden_params"):

2
litellm/llms/fal_ai/image_generation/gpt_image_2_transformation.py
View file

@ -51,7 +51,7 @@ def supported_gpt_image_qualities(
and "-x-" in parts[2]
and "/".join(parts[3:]) == qualified_endpoint
)
return qualities | {"auto"} if qualities else frozenset()
return qualities | frozenset({"auto"}) if qualities else frozenset()
def map_gpt_image_quality(

41
litellm/llms/fal_ai/image_generation/transformation.py
View file

@ -1,6 +1,9 @@
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
@ -22,16 +25,40 @@ else:
LiteLLMLoggingObj = Any
class FalImageProviderSpecificFields(TypedDict, total=False):
width: ReadOnly[int]
height: ReadOnly[int]
content_type: ReadOnly[str]
_FAL_IMAGE_DATA: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def fal_images_to_image_objects(images: object) -> tuple[ImageObject, ...]:
if not isinstance(images, list):
return ()
return tuple(
ImageObject(url=image_data.get("url", None), b64_json=image_data.get("b64_json", None))
if isinstance(image_data, dict)
else ImageObject(url=image_data, b64_json=None)
for image_data in images
if isinstance(image_data, (dict, str))
)
def to_image_object(image_data: object) -> ImageObject:
if isinstance(image_data, Mapping):
image_map: Final = _FAL_IMAGE_DATA.validate_python(image_data)
url: Final = image_map.get("url")
b64_json: Final = image_map.get("b64_json")
width: Final = image_map.get("width")
height: Final = image_map.get("height")
content_type: Final = image_map.get("content_type")
provider_specific_fields: Final[FalImageProviderSpecificFields] = {
**({"width": width} if isinstance(width, int) and type(width) is int and width > 0 else {}),
**({"height": height} if isinstance(height, int) and type(height) is int and height > 0 else {}),
**({"content_type": content_type} if isinstance(content_type, str) else {}),
}
return ImageObject(
url=url if isinstance(url, str) else None,
b64_json=b64_json if isinstance(b64_json, str) else None,
provider_specific_fields=provider_specific_fields or None,
)
return ImageObject(url=image_data if isinstance(image_data, str) else None, b64_json=None)
return tuple(to_image_object(image_data) for image_data in images if isinstance(image_data, (Mapping, str)))
class FalAIBaseConfig(BaseImageGenerationConfig):

88
litellm/llms/fal_ai/videos/transformation.py
View file

@ -1,6 +1,8 @@
import math
import sys
import time
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, TypeAlias
@ -36,11 +38,33 @@ class FalAIVideoError(BaseLLMException):
_ALLOWED_ASPECT_RATIOS: Final[frozenset[str]] = frozenset({"auto", "16:9", "9:16", "1:1", "4:3", "3:4", "21:9"})
_ALLOWED_RESOLUTIONS: Final[frozenset[str]] = frozenset({"480p", "720p", "1080p", "4k"})
_RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = (
(480, "480p"),
(720, "720p"),
(1080, "1080p"),
@dataclass(frozen=True, slots=True)
class _ModelProfile:
resolutions: frozenset[str]
resolution_tiers: tuple[tuple[int, str], ...]
default_resolution: str
integer_duration: bool
reference_key: str
reference_as_list: bool
_SEEDANCE_PROFILE: Final[_ModelProfile] = _ModelProfile(
resolutions=frozenset({"480p", "720p", "1080p", "4k"}),
resolution_tiers=((480, "480p"), (720, "720p"), (1080, "1080p"), (sys.maxsize, "4k")),
default_resolution="720p",
integer_duration=False,
reference_key="image_url",
reference_as_list=False,
)
_H3_PROFILE: Final[_ModelProfile] = _ModelProfile(
resolutions=frozenset({"480P", "768P", "2K", "4K"}),
resolution_tiers=((480, "480P"), (768, "768P"), (1440, "2K"), (sys.maxsize, "4K")),
default_resolution="2K",
integer_duration=True,
reference_key="reference_image_urls",
reference_as_list=True,
)
_QUEUE_NAMESPACES: Final[frozenset[str]] = frozenset(("workflows", "comfy"))
_STATUS_MAP: Final[Mapping[str, str]] = MappingProxyType(
@ -75,8 +99,12 @@ def _duration_value(value: object) -> str | None:
return None
def _resolution_for_short_side(short_side: int) -> str:
return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k")
def _profile_for_model(model: str) -> _ModelProfile:
return _H3_PROFILE if model.startswith("minimax/h3/") else _SEEDANCE_PROFILE
def _resolution_for_short_side(short_side: int, profile: _ModelProfile) -> str:
return next(resolution for threshold, resolution in profile.resolution_tiers if short_side <= threshold)
def _model_path_from_request_url(raw_response: httpx.Response) -> str | None:
@ -97,14 +125,19 @@ def _request_id_from_request_url(raw_response: httpx.Response) -> str | None:
return segments[request_id_index] if len(segments) > request_id_index else None
def _size_params(size: object) -> Mapping[str, str]:
def _size_params(size: object, profile: _ModelProfile) -> Mapping[str, str]:
if not isinstance(size, str):
return MappingProxyType({})
if size in _ALLOWED_RESOLUTIONS:
return MappingProxyType({"resolution": size})
if size.count("x") != 1:
normalized_size: Final[str] = size.lower()
canonical_resolution: Final[str | None] = next(
(resolution for resolution in profile.resolutions if resolution.lower() == normalized_size),
None,
)
if canonical_resolution is not None:
return MappingProxyType({"resolution": canonical_resolution})
if normalized_size.count("x") != 1:
return MappingProxyType({})
width_text, height_text = size.split("x")
width_text, height_text = normalized_size.split("x")
if not (width_text.isdigit() and height_text.isdigit()):
return MappingProxyType({})
width: Final[int] = int(width_text)
@ -113,7 +146,7 @@ def _size_params(size: object) -> Mapping[str, str]:
return MappingProxyType({})
reduced_gcd: Final[int] = math.gcd(width, height)
aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}"
resolution: Final[str] = _resolution_for_short_side(min(width, height))
resolution: Final[str] = _resolution_for_short_side(min(width, height), profile)
if aspect_ratio in _ALLOWED_ASPECT_RATIOS:
return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio})
return MappingProxyType({"resolution": resolution})
@ -158,18 +191,27 @@ class FalAIVideoConfig(BaseVideoConfig):
input_reference: Final[object] = video_create_optional_params.get("input_reference")
if "input_reference" in video_create_optional_params and not isinstance(input_reference, str):
raise ValueError("fal.ai needs a public image URL for input_reference")
input_reference_params: Final[Mapping[str, str]] = (
profile: Final[_ModelProfile] = _profile_for_model(model)
input_reference_params: Final[Mapping[str, object]] = (
MappingProxyType({})
if not isinstance(input_reference, str)
else MappingProxyType({"image_url": input_reference})
else MappingProxyType(
{
profile.reference_key: (
[input_reference] # mutable-ok: fal.ai expects a list for H3 references
if profile.reference_as_list
else input_reference
),
}
)
)
duration_params: Final[Mapping[str, str]] = (
duration_params: Final[Mapping[str, object]] = (
MappingProxyType({})
if "seconds" not in video_create_optional_params
else self._duration_params(video_create_optional_params["seconds"])
else self._duration_params(video_create_optional_params["seconds"], profile)
)
size_params: Final[Mapping[str, str]] = (
_size_params(video_create_optional_params["size"])
_size_params(video_create_optional_params["size"], profile)
if "size" in video_create_optional_params
else MappingProxyType({})
)
@ -190,11 +232,11 @@ class FalAIVideoConfig(BaseVideoConfig):
return mapped_params
@staticmethod
def _duration_params(seconds: object) -> Mapping[str, str]:
def _duration_params(seconds: object, profile: _ModelProfile) -> Mapping[str, object]:
duration: Final[str | None] = _duration_value(seconds)
if duration is None:
raise ValueError("fal.ai seconds must be a numeric value")
return MappingProxyType({"duration": duration})
return MappingProxyType({"duration": int(duration) if profile.integer_duration else duration})
def validate_environment(
self,
@ -251,6 +293,7 @@ class FalAIVideoConfig(BaseVideoConfig):
request_data: Mapping[str, object] | None = None,
) -> VideoObject:
response_data: Final[Mapping[str, object]] = _response_data(raw_response)
profile: Final[_ModelProfile] = _profile_for_model(model)
request_params: Final[Mapping[str, object]] = request_data or MappingProxyType({})
request_id: Final[str] = _response_string(response_data, "request_id")
provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER
@ -262,7 +305,10 @@ class FalAIVideoConfig(BaseVideoConfig):
key: value
for key, value in (
("duration_seconds", duration),
("video_resolution", resolution if isinstance(resolution, str) else "720p"),
(
"video_resolution",
resolution if isinstance(resolution, str) else profile.default_resolution,
),
)
if value is not None
}

105
litellm/model_prices_and_context_window_backup.json
View file

@ -22889,6 +22889,45 @@
"video"
]
},
"fal_ai/minimax/h3/text-to-video": {
"litellm_provider": "fal_ai",
"mode": "video_generation",
"output_cost_per_second": 0.13,
"output_cost_per_second_480p": 0.05,
"output_cost_per_second_768p": 0.06,
"output_cost_per_second_2k": 0.13,
"output_cost_per_second_4k": 0.16,
"source": "https://fal.ai/models/minimax/h3/text-to-video",
"supported_endpoints": [
"/v1/videos"
],
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
]
},
"fal_ai/minimax/h3/reference-to-video": {
"litellm_provider": "fal_ai",
"mode": "video_generation",
"output_cost_per_second": 0.13,
"output_cost_per_second_480p": 0.05,
"output_cost_per_second_768p": 0.06,
"output_cost_per_second_2k": 0.13,
"output_cost_per_second_4k": 0.16,
"source": "https://fal.ai/models/minimax/h3/reference-to-video",
"supported_endpoints": [
"/v1/videos"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
},
"fal_ai/bytedance/seedance-2.0/text-to-video": {
"litellm_provider": "fal_ai",
"mode": "video_generation",
@ -24917,10 +24956,11 @@
"fal_ai/fal-ai/flux/dev": {
"litellm_provider": "fal_ai",
"metadata": {
"notes": "fal bills FLUX.1 [dev] at $0.025 per megapixel, rounding each image up to the nearest megapixel. Every named fal image_size (including the landscape_4_3 default) rounds up to 1 megapixel, so this flat per-image price is exact for them"
"notes": "fal bills FLUX.1 [dev] at $0.025 per megapixel, rounding each image up to the nearest megapixel. The per-pixel rate is used when Fal reports the output size, and the flat per-image price is the fallback when dimensions are unavailable"
},
"mode": "image_generation",
"output_cost_per_image": 0.025,
"output_cost_per_pixel": 2.384185791015625e-08,
"source": "https://fal.ai/models/fal-ai/flux/dev",
"supported_endpoints": [
"/v1/images/generations"
@ -42971,21 +43011,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 9.22722e-07,
"input_cost_per_token": 9.15936e-07,
"input_cost_per_token_cache_hit": 4.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.845444e-06,
"output_cost_per_token": 1.831872e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 7.68935e-08,
"cache_read_input_token_cost": 7.6328e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -54441,16 +54481,19 @@
"zai.glm-4.7": {
"input_cost_per_token": 6e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"max_input_tokens": 203000,
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 2.2e-06,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_response_schema": true,
"supports_vision": false
},
"zai.glm-5": {
"input_cost_per_token": 1e-06,
@ -54465,21 +54508,27 @@
"supports_native_structured_output": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_response_schema": true,
"supports_vision": false
},
"zai.glm-4.7-flash": {
"input_cost_per_token": 7e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"max_input_tokens": 203000,
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 4e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_response_schema": true,
"supports_vision": false
},
"zai/glm-5": {
"cache_creation_input_token_cost": 0,
@ -60558,6 +60607,34 @@
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/anthropic.claude-haiku-4-5": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_1hr": 2e-06,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "bedrock_mantle",
"supports_tool_search": true,
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 5e-06,
"source": "https://aws.amazon.com/about-aws/whats-new/2025/10/claude-4-5-haiku-anthropic-amazon-bedrock",
"supports_assistant_prefill": true,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_native_structured_output": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 4096,
"input_cost_per_token_batches": 5e-07,
"output_cost_per_token_batches": 2.5e-06
},
"us.xai.grok-4.6": {
"input_cost_per_token": 2.2e-06,
"output_cost_per_token": 6.6e-06,
@ -76803,7 +76880,7 @@
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.65e-05,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3.html",
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_caching": true,

2
litellm/provider_endpoints_support_backup.json
View file

@ -691,7 +691,7 @@
}
},
"qwen_ai_platform": {
"display_name": "Qwen AI Platform (`qwen_ai_platform`)",
"display_name": "Qianwen AI Platform (`qwen_ai_platform`)",
"url": "https://docs.litellm.ai/docs/providers/qwencloud",
"endpoints": {
"chat_completions": true,

4
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -3865,7 +3865,7 @@ if MCP_AVAILABLE:
try:
data: Final = json.loads(body)
return isinstance(data, dict) and data.get("method") == "initialize"
except (json.JSONDecodeError, TypeError):
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
return False
def _extract_initialize_client_info(body: bytes) -> Implementation | None:
@ -4791,7 +4791,7 @@ if MCP_AVAILABLE:
"MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock",
_peeked.get("id"),
)
except (json.JSONDecodeError, TypeError):
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
# Peek cap truncated the body, so it can't be fully parsed.
# Scan the top-level keys (depth-aware) instead of a flat
# substring search: a response's result payload may nest a

18
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -34982,6 +34982,12 @@
"PolicyAttachmentCreateRequest": {
"description": "Request body for creating a policy attachment.",
"properties": {
"default": {
"default": false,
"description": "Apply this attachment only when no non-default attachment matches the request.",
"title": "Default",
"type": "boolean"
},
"keys": {
"anyOf": [
{
@ -35113,6 +35119,12 @@
"description": "Who created the attachment.",
"title": "Created By"
},
"default": {
"default": false,
"description": "Apply this attachment only when no non-default attachment matches the request.",
"title": "Default",
"type": "boolean"
},
"definition_location": {
"default": "db",
"description": "Where this attachment is defined: 'db' (database) or 'config' (config.yaml).",
@ -37141,6 +37153,12 @@
"PolicyAttachmentCreateRequest": {
"description": "Request body for creating a policy attachment.",
"properties": {
"default": {
"default": false,
"description": "Apply this attachment only when no non-default attachment matches the request.",
"title": "Default",
"type": "boolean"
},
"keys": {
"anyOf": [
{

11
litellm/proxy/litellm_pre_call_utils.py
View file

@ -3216,7 +3216,9 @@ def _match_and_track_policies(
attachment_registry: Final = (
attachment_registry_override if attachment_registry_override is not None else get_attachment_registry()
)
matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(context)
matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(
context, PolicyMatcher.policy_applies(context, policies_override)
)
matching_policy_names: Final = [m["policy_name"] for m in matches_with_reasons]
policy_reasons: Final = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons}
@ -3418,7 +3420,12 @@ async def add_guardrails_from_policy_engine(
_ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join(
(LlmProviders.ANTHROPIC.value, LlmProviders.BEDROCK.value, LlmProviders.VERTEX_AI.value)
(
LlmProviders.ANTHROPIC.value,
LlmProviders.BEDROCK.value,
LlmProviders.BEDROCK_MANTLE.value,
LlmProviders.VERTEX_AI.value,
)
)
_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = LlmProviders.ANTHROPIC.value

184
litellm/proxy/management_endpoints/prompt_caching_requests.py Normal file
View file

@ -0,0 +1,184 @@
from collections.abc import Callable, Mapping
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, Json, TypeAdapter
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth, user_api_key_has_admin_view
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.spend_tracking.savings import (
extract_cache_creation_tokens,
extract_cache_read_tokens,
marks_gateway_injection,
prompt_caching_savings_for_request,
)
from litellm.proxy.spend_tracking.spend_tracking_utils import (
_query_raw_rows, # pyright: ignore[reportPrivateUsage] # existing typed spend-query adapter; rows validated below
)
from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY
from litellm.types.management_endpoints.prompt_caching_requests import (
PromptCachingRequest,
PromptCachingRequestCursor,
PromptCachingRequestFilter,
PromptCachingRequestsResponse,
)
if TYPE_CHECKING:
from litellm.router import Router
router: Final = APIRouter()
def _numeric_token_sql(path: str) -> str:
value: Final = f"metadata #> '{{usage_object,{path}}}'"
return (
f"CASE WHEN jsonb_typeof({value}) = 'number' THEN ({value} #>> '{{}}')::numeric "
f"WHEN {value} = 'true'::jsonb THEN 1 WHEN {value} = 'false'::jsonb THEN 0 END"
)
def _cache_tokens_sql(*paths: str) -> str:
candidates: Final = ", ".join(f"NULLIF(({_numeric_token_sql(path)}), 0)" for path in paths)
return f"TRUNC(COALESCE({candidates}, 0))"
_CACHE_READ_SQL: Final = _cache_tokens_sql("cache_read_input_tokens", "prompt_tokens_details,cached_tokens")
_CACHE_CREATION_SQL: Final = _cache_tokens_sql(
"cache_creation_input_tokens",
"prompt_tokens_details,cache_write_tokens",
"prompt_tokens_details,cache_creation_tokens",
)
_GATEWAY_INJECTED_SQL: Final = (
f"(jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string' "
f"AND (metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = '' "
f"OR metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = model_id))"
)
_FILTER_SQL: Final = MappingProxyType(
{
"all": f"({_GATEWAY_INJECTED_SQL} OR {_CACHE_READ_SQL} > 0 OR {_CACHE_CREATION_SQL} > 0)",
"injected": _GATEWAY_INJECTED_SQL,
"hits": f"{_CACHE_READ_SQL} > 0",
}
)
def prompt_caching_requests_sql(filter: PromptCachingRequestFilter) -> str:
return f"""
SELECT request_id, "startTime" AS start_time, "endTime" AS end_time,
model, model_id, custom_llm_provider, spend,
CASE WHEN jsonb_typeof(metadata->'usage_object') = 'object'
THEN metadata->'usage_object' END AS usage_object,
CASE WHEN jsonb_typeof(metadata->'cost_breakdown') = 'object'
THEN metadata->'cost_breakdown' END AS cost_breakdown,
CASE WHEN jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string'
THEN metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' END AS gateway_marker
FROM "LiteLLM_SpendLogs"
WHERE "startTime" >= ($1::text::timestamptz AT TIME ZONE 'UTC')
AND "startTime" <= ($2::text::timestamptz AT TIME ZONE 'UTC')
AND COALESCE(LOWER(cache_hit), 'false') != 'true'
AND {_FILTER_SQL[filter]}
AND ($4::text::timestamptz IS NULL OR
("startTime", request_id) < (($4::text::timestamptz AT TIME ZONE 'UTC'), $5::text))
ORDER BY "startTime" DESC, request_id DESC
LIMIT $3::integer
"""
class _PromptCachingRow(BaseModel):
request_id: str
start_time: datetime
end_time: datetime
model: str
model_id: str | None
custom_llm_provider: str | None
spend: float
usage_object: Json[Mapping[str, object]] | Mapping[str, object] | None
cost_breakdown: Json[Mapping[str, object]] | Mapping[str, object] | None
gateway_marker: str | None
_REQUEST_ROWS: Final = TypeAdapter(tuple[_PromptCachingRow, ...])
def _request_result(row: _PromptCachingRow, llm_router: "Callable[[], Router | None]") -> PromptCachingRequest:
return PromptCachingRequest(
request_id=row.request_id,
start_time=row.start_time.replace(tzinfo=timezone.utc) if row.start_time.tzinfo is None else row.start_time,
model=row.model,
gateway_injected=marks_gateway_injection(
MappingProxyType({GATEWAY_INJECTED_CACHE_METADATA_KEY: row.gateway_marker}), row.model_id
),
cache_read_tokens=extract_cache_read_tokens(row.usage_object),
cache_creation_tokens=extract_cache_creation_tokens(row.usage_object),
spend=row.spend,
net_savings=prompt_caching_savings_for_request(
model=row.model,
custom_llm_provider=row.custom_llm_provider,
usage_object=row.usage_object,
model_id=row.model_id,
llm_router=llm_router,
cost_breakdown=row.cost_breakdown,
billed_at=row.end_time,
),
)
@router.get(
"/cost_optimization/prompt_caching/requests",
tags=["Cost Optimization"], # mutable-ok: FastAPI's route API requires a list
response_model=PromptCachingRequestsResponse,
)
async def get_prompt_caching_requests(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
start_date: datetime,
end_date: datetime,
page_size: Annotated[int, Query(ge=1, le=100)] = 50,
filter: PromptCachingRequestFilter = "all",
cursor_start_time: datetime | None = None,
cursor_request_id: Annotated[str | None, Query(min_length=1)] = None,
) -> PromptCachingRequestsResponse:
from litellm.proxy.proxy_server import llm_router, prisma_client
if not user_api_key_has_admin_view(user_api_key_dict):
raise HTTPException(status_code=403, detail="Only proxy admin roles can view prompt caching requests")
if (cursor_start_time is None) != (cursor_request_id is None):
raise HTTPException(status_code=400, detail="cursor_start_time and cursor_request_id must be provided together")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
start: Final = start_date.replace(tzinfo=timezone.utc) if start_date.tzinfo is None else start_date
end: Final = end_date.replace(tzinfo=timezone.utc) if end_date.tzinfo is None else end_date
if end < start:
raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date")
cursor_time: Final = (
cursor_start_time.replace(tzinfo=timezone.utc)
if cursor_start_time is not None and cursor_start_time.tzinfo is None
else cursor_start_time
)
rows: Final = _REQUEST_ROWS.validate_python(
await _query_raw_rows(
prisma_client,
prompt_caching_requests_sql(filter),
start.isoformat(),
end.isoformat(),
page_size + 1,
cursor_time.isoformat() if cursor_time is not None else None,
cursor_request_id,
)
or ()
)
def current_router() -> "Router | None":
return llm_router
requests: Final = tuple(_request_result(row, current_router) for row in rows[:page_size])
has_more: Final = len(rows) > page_size
return PromptCachingRequestsResponse(
requests=requests,
page_size=page_size,
has_more=has_more,
next_cursor=PromptCachingRequestCursor(start_time=requests[-1].start_time, request_id=requests[-1].request_id)
if has_more
else None,
)

42
litellm/proxy/policy_engine/attachment_registry.py
View file

@ -5,6 +5,7 @@ Attachments define WHERE policies apply, separate from the policy definitions.
This allows the same policy to be attached to multiple scopes.
"""
from collections.abc import Callable
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypedDict
@ -119,35 +120,49 @@ class AttachmentRegistry:
models=attachment_data.get("models"),
tags=attachment_data.get("tags"),
priority=attachment_data.get("priority"),
default=attachment_data.get("default", False),
)
def get_attached_policies(self, context: PolicyMatchContext) -> list[str]:
def get_attached_policies(
self,
context: PolicyMatchContext,
policy_applies: Callable[[str], bool] | None = None,
) -> list[str]:
"""
Get list of policy names attached to the given context.
Args:
context: The request context to match against
policy_applies: Optional predicate; attachments whose policy does not apply are ignored
Returns:
List of policy names that are attached to matching scopes
"""
return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)]
return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context, policy_applies)]
def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[PolicyAttachmentMatch]:
def get_attached_policies_with_reasons(
self,
context: PolicyMatchContext,
policy_applies: Callable[[str], bool] | None = None,
) -> list[PolicyAttachmentMatch]:
"""
Get list of policy names and match reasons for the given context.
Returns a list of dicts with 'policy_name' and 'matched_via' keys.
The 'matched_via' describes which dimension caused the match.
Attachments whose policy fails `policy_applies` are dropped before defaults are considered.
"""
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
in_scope: Final = tuple(
attachment
for attachment in self._attachments
if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
and (policy_applies is None or policy_applies(attachment.policy))
)
non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default)
matching_attachments: Final = sorted(
(
attachment
for attachment in self._attachments
if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
),
non_default or tuple(attachment for attachment in in_scope if attachment.default),
key=_attachment_sort_key,
)
broadest_attachment_by_policy: Final = MappingProxyType(
@ -169,6 +184,11 @@ class AttachmentRegistry:
@staticmethod
def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str:
"""Describe why an attachment matched the context."""
reason: Final = AttachmentRegistry._describe_scope_match(attachment, context)
return f"default:{reason}" if attachment.default else reason
@staticmethod
def _describe_scope_match(attachment: PolicyAttachment, context: PolicyMatchContext) -> str:
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
if attachment.is_global():
@ -324,6 +344,7 @@ class AttachmentRegistry:
"models": attachment_request.models or [],
"tags": attachment_request.tags or [],
"priority": attachment_request.priority,
"is_default": attachment_request.default,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
"created_by": created_by,
@ -340,6 +361,7 @@ class AttachmentRegistry:
models=attachment_request.models,
tags=attachment_request.tags,
priority=attachment_request.priority,
default=attachment_request.default,
)
self.add_attachment(attachment)
@ -352,6 +374,7 @@ class AttachmentRegistry:
models=created_attachment.models or [],
tags=created_attachment.tags or [],
priority=created_attachment.priority,
default=created_attachment.is_default,
created_at=created_attachment.created_at,
updated_at=created_attachment.updated_at,
created_by=created_attachment.created_by,
@ -429,6 +452,7 @@ class AttachmentRegistry:
models=attachment.models or [],
tags=attachment.tags or [],
priority=attachment.priority,
default=attachment.is_default,
created_at=attachment.created_at,
updated_at=attachment.updated_at,
created_by=attachment.created_by,
@ -468,6 +492,7 @@ class AttachmentRegistry:
models=a.models or [],
tags=a.tags or [],
priority=a.priority,
default=a.is_default,
created_at=a.created_at,
updated_at=a.updated_at,
created_by=a.created_by,
@ -502,6 +527,7 @@ class AttachmentRegistry:
models=(attachment_response.models if attachment_response.models else None),
tags=attachment_response.tags if attachment_response.tags else None,
priority=attachment_response.priority,
default=attachment_response.default,
)
for attachment_response in attachments
]

1
litellm/proxy/policy_engine/policy_endpoints.py
View file

@ -61,6 +61,7 @@ def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment)
models=attachment.models or [],
tags=attachment.tags or [],
priority=attachment.priority,
default=attachment.default,
definition_location="config",
)

36
litellm/proxy/policy_engine/policy_matcher.py
View file

@ -7,6 +7,7 @@ apply to a given request based on team alias, key alias, and model.
Policies are matched via policy_attachments which define WHERE each policy applies.
"""
from collections.abc import Callable, Sequence
from typing import Final
from litellm._logging import verbose_proxy_logger
@ -113,7 +114,7 @@ class PolicyMatcher:
verbose_proxy_logger.debug("AttachmentRegistry not initialized, returning empty list")
return []
return registry.get_attached_policies(context)
return registry.get_attached_policies(context, PolicyMatcher.policy_applies(context))
@staticmethod
def get_matching_policies_from_registry(
@ -130,9 +131,31 @@ class PolicyMatcher:
"""
return PolicyMatcher.get_matching_policies(context=context)
@staticmethod
def policy_applies(
context: PolicyMatchContext,
policies: dict[str, Policy] | None = None,
) -> Callable[[str], bool]:
"""Predicate telling whether a policy exists and its condition matches the context."""
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
return lambda policy_name: bool(
PolicyMatcher.get_policies_with_matching_conditions(
policy_names=(policy_name,),
context=context,
policies=resolved,
)
)
@staticmethod
def _registry_policies() -> dict[str, Policy]:
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
registry: Final = get_policy_registry()
return registry.get_all_policies() if registry.is_initialized() else {}
@staticmethod
def get_policies_with_matching_conditions(
policy_names: list[str],
policy_names: Sequence[str],
context: PolicyMatchContext,
policies: dict[str, Policy] | None = None,
) -> list[str]:
@ -152,17 +175,12 @@ class PolicyMatcher:
List of policy names whose conditions match the context
"""
from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
if policies is None:
registry: Final = get_policy_registry()
if not registry.is_initialized():
return []
policies = registry.get_all_policies()
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
matching_policies: Final = []
for policy_name in policy_names:
policy = policies.get(policy_name)
policy = resolved.get(policy_name)
if policy is None:
continue
# Policy matches if it has no condition OR condition evaluates to True

4
litellm/proxy/policy_engine/policy_resolve_endpoints.py
View file

@ -265,7 +265,9 @@ async def resolve_policies_for_context(
)
# Get matching policies with reasons
match_results: Final = get_attachment_registry().get_attached_policies_with_reasons(context=context)
match_results: Final = get_attachment_registry().get_attached_policies_with_reasons(
context=context, policy_applies=PolicyMatcher.policy_applies(context)
)
if not match_results:
return PolicyResolveResponse(

4
litellm/proxy/policy_engine/response_retrieval.py
View file

@ -84,7 +84,9 @@ def _retrieval_context(
def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]:
matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context)
matches: Final = get_attachment_registry().get_attached_policies_with_reasons(
context, PolicyMatcher.policy_applies(context)
)
if not matches:
return (), MappingProxyType({})
applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions(

4
litellm/proxy/proxy_server.py
View file

@ -608,6 +608,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
from litellm.proxy.management_endpoints.organization_endpoints import (
router as organization_router,
)
from litellm.proxy.management_endpoints.prompt_caching_requests import (
router as prompt_caching_requests_router,
)
from litellm.proxy.management_endpoints.router_settings_endpoints import (
router as router_settings_router,
)
@ -19315,6 +19318,7 @@ app.include_router(workflow_management_router)
app.include_router(memory_router)
app.include_router(plugin_router)
app.include_router(cost_tracking_settings_router)
app.include_router(prompt_caching_requests_router)
app.include_router(router_settings_router)
app.include_router(fallback_management_router)
app.include_router(cache_settings_router)

6
litellm/proxy/public_endpoints/provider_create_fields.json
View file

@ -1129,12 +1129,12 @@
},
{
"provider": "Qwen_AI_Platform",
"provider_display_name": "Qwen AI Platform",
"provider_display_name": "Qianwen AI Platform",
"litellm_provider": "qwen_ai_platform",
"credential_fields": [
{
"key": "api_key",
"label": "Qwen AI Platform API Key",
"label": "Qianwen AI Platform API Key",
"placeholder": null,
"tooltip": null,
"required": true,
@ -1146,7 +1146,7 @@
"key": "api_base",
"label": "API Base",
"placeholder": "https://dashscope.aliyuncs.com/compatible-mode/v1",
"tooltip": "The base URL for Qwen AI Platform. Defaults to https://dashscope.aliyuncs.com/compatible-mode/v1 if not specified.",
"tooltip": "The base URL for Qianwen AI Platform. Defaults to https://dashscope.aliyuncs.com/compatible-mode/v1 if not specified.",
"required": true,
"field_type": "text",
"options": null,

1
litellm/proxy/schema.prisma
View file

@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
is_default Boolean @default(false) // Applied only when no non-default attachment matches
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt

75
litellm/proxy/spend_tracking/savings.py
View file

@ -578,6 +578,56 @@ def autorouter_savings_for_logging_payload(
)
def _request_savings_pricing(
model: str | None,
custom_llm_provider: str | None,
model_id: str | None,
llm_router: "Callable[[], Router | None] | None",
) -> tuple[str | None, ModelInfo | None]:
router_instance: Final = llm_router() if llm_router else None
identity: Final = _resolve_model(model, custom_llm_provider)
pricing: Final = _effective_model_info(router_instance, model_id, model or "") or (
_model_info(identity) if identity else None
)
return identity.provider if identity else custom_llm_provider, pricing
def _prompt_caching_savings(
pricing: ModelInfo | None,
provider: str | None,
usage_object: Mapping[str, object] | None,
cost_breakdown: Mapping[str, object] | None,
billed_at: datetime | str | None,
) -> float | None:
usage: Final = _usage_from_spend_log(usage_object)
if pricing is None or usage is None:
return None
basis: Final = _pricing_basis(cost_breakdown)
result: Final = calculate_prompt_caching_savings(
model_info=pricing,
usage=usage,
custom_llm_provider=provider,
service_tier=basis.service_tier,
data_residency=basis.data_residency,
vertex_location=basis.vertex_location,
billed_at=_coerce_billed_at(billed_at),
)
return result if isfinite(result) else None
def prompt_caching_savings_for_request(
model: str | None,
custom_llm_provider: str | None,
usage_object: Mapping[str, object] | None,
model_id: str | None = None,
llm_router: "Callable[[], Router | None] | None" = None,
cost_breakdown: Mapping[str, object] | None = None,
billed_at: datetime | str | None = None,
) -> float | None:
request_pricing: Final = _request_savings_pricing(model, custom_llm_provider, model_id, llm_router)
return _prompt_caching_savings(request_pricing[1], request_pricing[0], usage_object, cost_breakdown, billed_at)
def compute_savings_spend(
model: str | None,
custom_llm_provider: str | None,
@ -639,29 +689,12 @@ def compute_savings_spend(
# Deployment rates when the request came through one, public rates otherwise --
# `_effective_model_info` merges a deployment's configured prices over the built-in
# map, so a negotiated price is not silently replaced by the list rate.
router_instance: Router | None = llm_router() if llm_router else None
identity: Final = _resolve_model(model, custom_llm_provider)
pricing: Final = _effective_model_info(router_instance, model_id, model or "") or (
_model_info(identity) if identity else None
)
request_pricing: Final = _request_savings_pricing(model, custom_llm_provider, model_id, llm_router)
provider: Final = request_pricing[0]
pricing: Final = request_pricing[1]
input_cost: Final = (_get_cost_per_unit(pricing, "input_cost_per_token") or 0.0) if pricing else 0.0
compression: Final = max(compression_saved_tokens, 0) * input_cost
usage: Final = _usage_from_spend_log(usage_object)
basis: Final = _pricing_basis(cost_breakdown)
billed_at_datetime: Final = _coerce_billed_at(billed_at)
prompt_caching: Final = (
calculate_prompt_caching_savings(
model_info=pricing,
usage=usage,
custom_llm_provider=identity.provider if identity else custom_llm_provider,
service_tier=basis.service_tier,
data_residency=basis.data_residency,
vertex_location=basis.vertex_location,
billed_at=billed_at_datetime,
)
if pricing is not None and usage is not None
else 0.0
)
prompt_caching: Final = _prompt_caching_savings(pricing, provider, usage_object, cost_breakdown, billed_at) or 0.0
gateway_injected_caching: Final = prompt_caching if gateway_injected_cache else 0.0
# The figure the logging path recorded wins, before the usage gate on purpose: a row

260
litellm/router_utils/prompt_caching_cache.py
View file

@ -4,12 +4,19 @@ Wrapper around router cache. Meant to store model id when prompt caching support
import hashlib
import json
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from typing import TYPE_CHECKING, Any, Final, cast
from pydantic import JsonValue, TypeAdapter
from pydantic_core import to_jsonable_python
from typing_extensions import TypedDict
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import PROMPT_CACHE_LOOKBACK_POSITIONS
from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
if TYPE_CHECKING:
@ -28,27 +35,102 @@ class PromptCachingCacheValue(TypedDict):
model_id: str
PROMPT_CACHE_PIN_TTL_SECONDS: Final = 300
_TOOL_RUN_BLOCK_TYPES: Final = frozenset({"tool_use", "tool_result"})
_PREFIX_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...])
_TOOLS_ADAPTER: Final = TypeAdapter(tuple[JsonValue, ...])
_PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None)
@dataclass(frozen=True, slots=True)
class PrefixPosition:
cache_key: str
position: int
def _sorted_pairs(pairs: Iterable[tuple[str, JsonValue]]) -> tuple[tuple[str, JsonValue], ...]:
return tuple(sorted(pairs, key=lambda pair: pair[0]))
def _canonical_bytes(value: object) -> bytes:
return json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
def _block_unit(
envelope: tuple[tuple[str, JsonValue], ...], message_run_type: str | None, block: JsonValue
) -> tuple[bytes, str | None]:
if not isinstance(block, dict):
return _canonical_bytes((envelope, block)), message_run_type
block_type: Final = block.get("type")
block_run_type: Final = block_type if isinstance(block_type, str) and block_type in _TOOL_RUN_BLOCK_TYPES else None
stripped: Final = _sorted_pairs(item for item in block.items() if item[0] != "cache_control")
return _canonical_bytes((envelope, stripped)), message_run_type or block_run_type
def _message_units(message: Mapping[str, JsonValue]) -> tuple[tuple[bytes, str | None], ...]:
envelope: Final = _sorted_pairs(item for item in message.items() if item[0] not in ("content", "cache_control"))
message_run_type: Final = "tool_result" if message.get("role") == "tool" else None
content: Final = message.get("content")
if isinstance(content, list) and content:
return tuple(_block_unit(envelope, message_run_type, block) for block in content)
if isinstance(content, str) and content:
return ((_canonical_bytes((envelope, (("text", content), ("type", "text")))), message_run_type),)
return ((_canonical_bytes((envelope, None)), message_run_type),)
def _chain_digest(digest: bytes, unit: bytes) -> bytes:
return hashlib.sha256(digest + unit).digest()
def _seed(tools: Sequence[ChatCompletionToolParam] | None) -> bytes:
if tools is None:
return hashlib.sha256(b"").digest()
return hashlib.sha256(
_canonical_bytes(
_TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True, bytes_mode="base64"))
)
).digest()
def _positions_of(
prefix: tuple[Mapping[str, JsonValue], ...], tools: Sequence[ChatCompletionToolParam] | None
) -> tuple[PrefixPosition, ...]:
units: Final = tuple(unit for message in prefix for unit in _message_units(message))
digests: Final = tuple(accumulate((unit_bytes for unit_bytes, _ in units), _chain_digest, initial=_seed(tools)))[1:]
run_types: Final = tuple(run_type for _, run_type in units)
positions: Final = accumulate(
0 if run_type is not None and run_type == previous else 1
for run_type, previous in zip(run_types, (None, *run_types[:-1]))
)
return tuple(
PrefixPosition(cache_key=f"deployment:{digest.hex()}:prompt_caching", position=position)
for digest, position in zip(digests, positions)
)
def _lookback_keys(positions: tuple[PrefixPosition, ...]) -> tuple[str, ...]:
if not positions:
return ()
oldest_probed_position: Final = positions[-1].position - PROMPT_CACHE_LOOKBACK_POSITIONS
return tuple(entry.cache_key for entry in reversed(positions) if entry.position > oldest_probed_position)
def _pinned_value(value: JsonValue) -> PromptCachingCacheValue | None:
if not isinstance(value, dict):
return None
model_id: Final = value.get("model_id")
return PromptCachingCacheValue(model_id=model_id) if isinstance(model_id, str) else None
def _first_pin(values: tuple[JsonValue, ...] | None) -> PromptCachingCacheValue | None:
if values is None:
return None
return next((pin for pin in map(_pinned_value, values) if pin is not None), None)
class PromptCachingCache:
def __init__(self, cache: DualCache):
self.cache = cache
self.in_memory_cache = InMemoryCache()
@staticmethod
def serialize_object(obj: Any) -> object:
"""Helper function to serialize Pydantic objects, dictionaries, or fallback to string."""
if hasattr(obj, "dict"):
# If the object is a Pydantic model, use its `dict()` method
return obj.dict()
elif isinstance(obj, dict):
# If the object is a dictionary, serialize it with sorted keys
return json.dumps(obj, sort_keys=True, separators=(",", ":")) # Standardize serialization
elif isinstance(obj, list):
# Serialize lists by ensuring each element is handled properly
return [PromptCachingCache.serialize_object(item) for item in obj]
elif isinstance(obj, (int, float, bool)):
return obj # Keep primitive types as-is
return str(obj)
@staticmethod
def extract_cacheable_prefix(
@ -140,114 +222,116 @@ class PromptCachingCache:
return cacheable_prefix
@staticmethod
def get_prompt_caching_cache_key(
def prefix_positions(
messages: list[AllMessageValues] | None,
tools: list[ChatCompletionToolParam] | None,
) -> str | None:
if messages is None and tools is None:
return None
tools: Sequence[ChatCompletionToolParam] | None,
) -> tuple[PrefixPosition, ...]:
"""
One cache key per content block of the cacheable prefix, oldest block first.
# Extract cacheable prefix from messages (only include up to last cache_control block)
cacheable_messages = None
if messages is not None:
cacheable_messages = PromptCachingCache.extract_cacheable_prefix(messages)
# If no cacheable prefix found, return None (can't cache)
if not cacheable_messages:
return None
Each key hashes the prefix content up to and including that block, with cache_control markers
left out, so the key of a block is the same whichever turn's breakpoint the prefix ends at.
String content hashes like a single text block, which is how the provider treats it and how
Claude Code re-sends a previously marked message. `position` counts a run of consecutive
tool_use (or tool_result) blocks as one, matching the provider's lookback window.
# Use serialize_object for consistent and stable serialization
data_to_hash: Final = {}
if cacheable_messages is not None:
serialized_messages: Final = PromptCachingCache.serialize_object(cacheable_messages)
data_to_hash["messages"] = serialized_messages
if tools is not None:
serialized_tools: Final = PromptCachingCache.serialize_object(tools)
data_to_hash["tools"] = serialized_tools
# Combine serialized data into a single string
data_to_hash_str: Final = json.dumps(
data_to_hash,
sort_keys=True,
separators=(",", ":"),
The prefix is hashed in the shape the success event sees it, with long base64 data URIs
already replaced by their size placeholder, so a request carrying the raw image bytes
derives the same keys the write side stored.
"""
if not messages:
return ()
return _positions_of(
_PREFIX_ADAPTER.validate_python(
to_jsonable_python(
truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)),
serialize_unknown=True,
bytes_mode="base64",
)
),
tools,
)
# Create a hash of the serialized data for a stable cache key
hashed_data: Final = hashlib.sha256(data_to_hash_str.encode()).hexdigest()
return f"deployment:{hashed_data}:prompt_caching"
@staticmethod
async def async_prefix_positions(
messages: list[AllMessageValues] | None,
tools: Sequence[ChatCompletionToolParam] | None,
) -> tuple[PrefixPosition, ...]:
if not messages:
return ()
return await offload_token_count(PromptCachingCache.prefix_positions)(messages, tools)
@staticmethod
def get_prompt_caching_cache_key(
messages: list[AllMessageValues] | None,
tools: Sequence[ChatCompletionToolParam] | None,
) -> str | None:
positions: Final = PromptCachingCache.prefix_positions(messages, tools)
return positions[-1].cache_key if positions else None
def add_model_id(
self,
model_id: str,
messages: list[AllMessageValues] | None,
tools: list[ChatCompletionToolParam] | None,
tools: Sequence[ChatCompletionToolParam] | None,
) -> None:
if messages is None and tools is None:
return
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
# If no cacheable prefix found, don't cache (can't generate cache key)
if cache_key is None:
return
self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300)
return
self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=PROMPT_CACHE_PIN_TTL_SECONDS)
async def async_add_model_id(
self,
model_id: str,
messages: list[AllMessageValues] | None,
tools: list[ChatCompletionToolParam] | None,
tools: Sequence[ChatCompletionToolParam] | None,
) -> None:
if messages is None and tools is None:
return
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
# If no cacheable prefix found, don't cache (can't generate cache key)
if cache_key is None:
positions: Final = await PromptCachingCache.async_prefix_positions(messages, tools)
if not positions:
return
await self.cache.async_set_cache(
cache_key,
positions[-1].cache_key,
PromptCachingCacheValue(model_id=model_id),
ttl=300, # store for 5 minutes
ttl=PROMPT_CACHE_PIN_TTL_SECONDS,
)
return
async def async_get_model_id(
self,
messages: list[AllMessageValues] | None,
tools: list[ChatCompletionToolParam] | None,
tools: Sequence[ChatCompletionToolParam] | None,
) -> PromptCachingCacheValue | None:
"""
Get model ID from cache using the cacheable prefix.
The cache key is based on the cacheable prefix (everything up to and including
the last cache_control block), so requests with the same cacheable prefix but
different user messages will have the same cache key.
Find the deployment that last served this prefix, walking back from the breakpoint the
same way the provider cache does, so a breakpoint that moved forward since the last
turn still lands on the deployment whose cache holds the earlier prefix.
"""
if messages is None and tools is None:
cache_keys: Final = _lookback_keys(await PromptCachingCache.async_prefix_positions(messages, tools))
if not cache_keys:
return None
# Generate cache key using cacheable prefix
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
if cache_key is None:
return None
# Perform cache lookup
cache_result: Final = await self.cache.async_get_cache(key=cache_key)
return cache_result
return _first_pin(
_PINS_ADAPTER.validate_python(
await self.cache.async_batch_get_cache(
keys=list(cache_keys), # mutable-ok: DualCache.async_batch_get_cache only takes a list
)
)
)
def get_model_id(
self,
messages: list[AllMessageValues] | None,
tools: list[ChatCompletionToolParam] | None,
tools: Sequence[ChatCompletionToolParam] | None,
) -> PromptCachingCacheValue | None:
if messages is None and tools is None:
cache_keys: Final = _lookback_keys(PromptCachingCache.prefix_positions(messages, tools))
if not cache_keys:
return None
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
# If no cacheable prefix found, return None (can't cache)
if cache_key is None:
return None
return self.cache.get_cache(cache_key)
return _first_pin(
_PINS_ADAPTER.validate_python(
self.cache.batch_get_cache(
keys=list(cache_keys), # mutable-ok: DualCache.batch_get_cache only takes a list
)
)
)

4
litellm/types/llms/bedrock.py
View file

@ -4,7 +4,7 @@ from enum import Enum
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict
from typing_extensions import ReadOnly, Required, TypedDict, override
from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, override
from .openai import ChatCompletionToolCallChunk
@ -1082,6 +1082,7 @@ class BedrockS3InputDataConfig(TypedDict):
"""S3 input data configuration for Bedrock batch jobs."""
s3Uri: str
s3BucketOwner: NotRequired[ReadOnly[str]]
class BedrockInputDataConfig(TypedDict):
@ -1095,6 +1096,7 @@ class BedrockS3OutputDataConfig(TypedDict, total=False):
s3Uri: str
s3EncryptionKeyId: str | None
s3BucketOwner: ReadOnly[str]
class BedrockOutputDataConfig(TypedDict):

35
litellm/types/management_endpoints/prompt_caching_requests.py Normal file
View file

@ -0,0 +1,35 @@
from datetime import datetime
from typing import Literal, TypeAlias
from pydantic import BaseModel, ConfigDict
PromptCachingRequestFilter: TypeAlias = Literal["all", "injected", "hits"]
class PromptCachingRequest(BaseModel):
model_config = ConfigDict(frozen=True)
request_id: str
start_time: datetime
model: str
gateway_injected: bool
cache_read_tokens: int
cache_creation_tokens: int
spend: float
net_savings: float | None
class PromptCachingRequestCursor(BaseModel):
model_config = ConfigDict(frozen=True)
start_time: datetime
request_id: str
class PromptCachingRequestsResponse(BaseModel):
model_config = ConfigDict(frozen=True)
requests: tuple[PromptCachingRequest, ...]
page_size: int
has_more: bool
next_cursor: PromptCachingRequestCursor | None

4
litellm/types/proxy/policy_engine/policy_types.py
View file

@ -294,6 +294,10 @@ class PolicyAttachment(BaseModel):
le=2147483647,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
default: bool = Field(
default=False,
description="Apply this attachment only when no non-default attachment matches the request.",
)
model_config = ConfigDict(extra="forbid")

8
litellm/types/proxy/policy_engine/resolver_types.py
View file

@ -311,6 +311,10 @@ class PolicyAttachmentCreateRequest(BaseModel):
le=2147483647,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
default: bool = Field(
default=False,
description="Apply this attachment only when no non-default attachment matches the request.",
)
class PolicyAttachmentDBResponse(BaseModel):
@ -327,6 +331,10 @@ class PolicyAttachmentDBResponse(BaseModel):
default=None,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
default: bool = Field(
default=False,
description="Apply this attachment only when no non-default attachment matches the request.",
)
created_at: datetime | None = Field(default=None, description="When the attachment was created.")
updated_at: datetime | None = Field(default=None, description="When the attachment was last updated.")
created_by: str | None = Field(default=None, description="Who created the attachment.")

3
litellm/types/router.py
View file

@ -306,6 +306,7 @@ class CredentialLiteLLMParams(BaseModel):
s3_endpoint_url: str | None = None
s3_region_name: str | None = None
s3_encryption_key_id: str | None = None
s3_bucket_owner: str | None = None
aws_batch_role_arn: str | None = None
s3_output_bucket_name: str | None = None
bedrock_tags: list | None = None
@ -539,6 +540,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
output_cost_per_second: float | None
output_cost_per_second_480p: ReadOnly[float | None]
output_cost_per_second_720p: ReadOnly[float | None]
output_cost_per_second_768p: ReadOnly[float | None]
output_cost_per_second_2k: ReadOnly[float | None]
output_cost_per_second_1080p: float | None
output_cost_per_second_4k: ReadOnly[float | None]
num_retries: int | None

10
litellm/types/utils.py
View file

@ -315,6 +315,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
output_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x output
output_cost_per_character_above_128k_tokens: float | None # only for vertex ai models
output_cost_per_image: float | None
output_cost_per_pixel: ReadOnly[float | None]
output_cost_per_image_token: float | None
output_cost_per_video_token: float | None # for gemini omni models with video output
output_vector_size: int | None
@ -329,6 +330,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
) # video_generation tier: key output_cost_per_second_<resolution> (e.g. 1080p, 720p)
output_cost_per_second_480p: ReadOnly[float | None]
output_cost_per_second_720p: ReadOnly[float | None]
output_cost_per_second_768p: ReadOnly[float | None]
output_cost_per_second_2k: ReadOnly[float | None]
output_cost_per_second_4k: ReadOnly[float | None]
ocr_cost_per_page: float | None # for OCR models
ocr_cost_per_page_batches: ReadOnly[float | None]
@ -2551,6 +2554,10 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
model_config = ConfigDict(extra="allow", protected_namespaces=())
@field_serializer("data")
def _serialize_image_data(self, data: Sequence[OpenAIImage] | None) -> Sequence[Mapping[str, object]] | None:
return None if data is None else [image.model_dump() for image in data]
def __init__(
self,
created: int | None = None,
@ -3610,6 +3617,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
output_cost_per_second_1080p: float | None = None
output_cost_per_second_480p: float | None = None
output_cost_per_second_720p: float | None = None
output_cost_per_second_768p: float | None = None
output_cost_per_second_2k: float | None = None
output_cost_per_second_4k: float | None = None
input_cost_per_pixel: float | None = None
output_cost_per_pixel: float | None = None
@ -3830,6 +3839,7 @@ bedrock_batch_litellm_params: Final = (
"s3_region_name",
"s3_endpoint_url",
"s3_output_bucket_name",
"s3_bucket_owner",
"bedrock_tags",
)

34
litellm/utils.py
View file

@ -5624,6 +5624,12 @@ def _get_model_info_from_generalization(
return None
def _strip_mantle_region_prefix(model: str) -> str:
from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix
return split_mantle_region_prefix(model)[1]
def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> PotentialModelNamesAndCustomLLMProvider:
if custom_llm_provider is None:
# Get custom_llm_provider
@ -5656,20 +5662,30 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P
split_model = strip_bedrock_routing_prefix(split_model)
region_free_split_model: Final = (
_strip_mantle_region_prefix(split_model) if custom_llm_provider == "bedrock_mantle" else split_model
)
region_free_combined_stripped_model_name: Final = (
f"bedrock_mantle/{_strip_model_name(model=region_free_split_model, custom_llm_provider=custom_llm_provider)}"
if custom_llm_provider == "bedrock_mantle"
else combined_stripped_model_name
)
provider_model_info: Final = (
ProviderConfigManager.get_provider_model_info(model=split_model, provider=LlmProviders(custom_llm_provider))
ProviderConfigManager.get_provider_model_info(
model=region_free_split_model, provider=LlmProviders(custom_llm_provider)
)
if custom_llm_provider in LlmProvidersSet
else None
)
provider_cost_key: Final = (
provider_model_info.get_model_cost_key(split_model) if provider_model_info is not None else None
provider_model_info.get_model_cost_key(region_free_split_model) if provider_model_info is not None else None
)
return PotentialModelNamesAndCustomLLMProvider(
split_model=split_model,
split_model=region_free_split_model,
combined_model_name=combined_model_name,
stripped_model_name=stripped_model_name,
combined_stripped_model_name=combined_stripped_model_name,
combined_stripped_model_name=region_free_combined_stripped_model_name,
provider_prefixed_model_name=provider_cost_key or provider_prefixed_model_name,
custom_llm_provider=cast(str, custom_llm_provider),
)
@ -6084,9 +6100,12 @@ def _get_model_info_helper(
output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None),
output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None),
output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None),
output_cost_per_second_768p=_model_info.get("output_cost_per_second_768p", None),
output_cost_per_second_2k=_model_info.get("output_cost_per_second_2k", None),
output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None),
output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None),
output_cost_per_image=_model_info.get("output_cost_per_image", None),
output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None),
output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None),
output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None),
output_vector_size=_model_info.get("output_vector_size", None),
@ -8681,6 +8700,13 @@ class ProviderConfigManager:
from litellm.llms.bedrock.common_utils import BedrockModelInfo
return BedrockModelInfo.get_bedrock_provider_config_for_messages_api(model)
elif litellm.LlmProviders.BEDROCK_MANTLE == provider:
if "claude" in model_lower:
from litellm.llms.bedrock_mantle.messages.transformation import (
BedrockMantleAnthropicMessagesConfig,
)
return BedrockMantleAnthropicMessagesConfig()
elif litellm.LlmProviders.VERTEX_AI == provider:
if "claude" in model_lower:
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import (

105
model_prices_and_context_window.json
View file

@ -22889,6 +22889,45 @@
"video"
]
},
"fal_ai/minimax/h3/text-to-video": {
"litellm_provider": "fal_ai",
"mode": "video_generation",
"output_cost_per_second": 0.13,
"output_cost_per_second_480p": 0.05,
"output_cost_per_second_768p": 0.06,
"output_cost_per_second_2k": 0.13,
"output_cost_per_second_4k": 0.16,
"source": "https://fal.ai/models/minimax/h3/text-to-video",
"supported_endpoints": [
"/v1/videos"
],
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
]
},
"fal_ai/minimax/h3/reference-to-video": {
"litellm_provider": "fal_ai",
"mode": "video_generation",
"output_cost_per_second": 0.13,
"output_cost_per_second_480p": 0.05,
"output_cost_per_second_768p": 0.06,
"output_cost_per_second_2k": 0.13,
"output_cost_per_second_4k": 0.16,
"source": "https://fal.ai/models/minimax/h3/reference-to-video",
"supported_endpoints": [
"/v1/videos"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
},
"fal_ai/bytedance/seedance-2.0/text-to-video": {
"litellm_provider": "fal_ai",
"mode": "video_generation",
@ -24917,10 +24956,11 @@
"fal_ai/fal-ai/flux/dev": {
"litellm_provider": "fal_ai",
"metadata": {
"notes": "fal bills FLUX.1 [dev] at $0.025 per megapixel, rounding each image up to the nearest megapixel. Every named fal image_size (including the landscape_4_3 default) rounds up to 1 megapixel, so this flat per-image price is exact for them"
"notes": "fal bills FLUX.1 [dev] at $0.025 per megapixel, rounding each image up to the nearest megapixel. The per-pixel rate is used when Fal reports the output size, and the flat per-image price is the fallback when dimensions are unavailable"
},
"mode": "image_generation",
"output_cost_per_image": 0.025,
"output_cost_per_pixel": 2.384185791015625e-08,
"source": "https://fal.ai/models/fal-ai/flux/dev",
"supported_endpoints": [
"/v1/images/generations"
@ -42971,21 +43011,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 9.22722e-07,
"input_cost_per_token": 9.15936e-07,
"input_cost_per_token_cache_hit": 4.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.845444e-06,
"output_cost_per_token": 1.831872e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 7.68935e-08,
"cache_read_input_token_cost": 7.6328e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -54441,16 +54481,19 @@
"zai.glm-4.7": {
"input_cost_per_token": 6e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"max_input_tokens": 203000,
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 2.2e-06,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_response_schema": true,
"supports_vision": false
},
"zai.glm-5": {
"input_cost_per_token": 1e-06,
@ -54465,21 +54508,27 @@
"supports_native_structured_output": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_response_schema": true,
"supports_vision": false
},
"zai.glm-4.7-flash": {
"input_cost_per_token": 7e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"max_input_tokens": 203000,
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 4e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_response_schema": true,
"supports_vision": false
},
"zai/glm-5": {
"cache_creation_input_token_cost": 0,
@ -60558,6 +60607,34 @@
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/anthropic.claude-haiku-4-5": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_1hr": 2e-06,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "bedrock_mantle",
"supports_tool_search": true,
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 5e-06,
"source": "https://aws.amazon.com/about-aws/whats-new/2025/10/claude-4-5-haiku-anthropic-amazon-bedrock",
"supports_assistant_prefill": true,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_native_structured_output": true,
"supports_parallel_tool_use_config": true,
"prompt_cache_min_tokens": 4096,
"input_cost_per_token_batches": 5e-07,
"output_cost_per_token_batches": 2.5e-06
},
"us.xai.grok-4.6": {
"input_cost_per_token": 2.2e-06,
"output_cost_per_token": 6.6e-06,
@ -76803,7 +76880,7 @@
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.65e-05,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k3.html",
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_caching": true,

8
model_prices_and_context_window.schema.json
View file

@ -607,6 +607,10 @@
"type": "number",
"minimum": 0
},
"output_cost_per_second_2k": {
"type": "number",
"minimum": 0
},
"output_cost_per_second_480p": {
"type": "number",
"minimum": 0
@ -619,6 +623,10 @@
"type": "number",
"minimum": 0
},
"output_cost_per_second_768p": {
"type": "number",
"minimum": 0
},
"output_cost_per_token": {
"type": "number",
"minimum": 0,

2
provider_endpoints_support.json
View file

@ -744,7 +744,7 @@
}
},
"qwen_ai_platform": {
"display_name": "Qwen AI Platform (`qwen_ai_platform`)",
"display_name": "Qianwen AI Platform (`qwen_ai_platform`)",
"url": "https://docs.litellm.ai/docs/providers/qwencloud",
"endpoints": {
"chat_completions": true,

1
schema.prisma
View file

@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
is_default Boolean @default(false) // Applied only when no non-default attachment matches
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt

3
tests/integration/contracts.json
View file

@ -175,6 +175,9 @@
"tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row": [
"other.provider_wire.fal_ai.image_edit_json_data_urls_and_keyed_pricing"
],
"tests/integration/providers/test_fal_ai_video_wire.py::test_fal_h3_video_create_uses_canonical_body_and_status_path": [
"other.provider_wire.fal_ai.video_queue_create_status_and_content_download"
],
"tests/integration/mcp/test_mcp_lifecycle.py::test_saved_headers_reach_real_mcp_tool_and_survive_unrelated_edit": [
"mcp.call_tool.saved_headers.reach_actual_transport"
],

67
tests/integration/providers/test_fal_ai_image_wire.py
View file

@ -23,14 +23,14 @@ _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]])
def _catalog_cost(key: str) -> float:
def _catalog_cost(key: str, field: str = "output_cost_per_image") -> float:
cost_map: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())
cost_value: Final = cost_map[key]["output_cost_per_image"]
cost_value: Final = cost_map[key][field]
assert isinstance(cost_value, (int, float))
return float(cost_value)
def _image_response(urls: tuple[str, ...], prompt: str) -> bytes:
def _image_response(images: tuple[tuple[str, int, int], ...], prompt: str) -> bytes:
return json.dumps(
{
"images": [
@ -39,10 +39,10 @@ def _image_response(urls: tuple[str, ...], prompt: str) -> bytes:
"content_type": "image/png",
"file_name": url.rsplit("/", 1)[-1],
"file_size": 123456,
"width": 1024,
"height": 768,
"width": width,
"height": height,
}
for url in urls
for url, width, height in images
],
"timings": {"inference": 2.1},
"seed": 1234567,
@ -69,9 +69,9 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro
body: Final = _JSON_OBJECT.validate_json(request.body)
if body.get("quality") == "high":
assert body == {"prompt": _PROMPT, "quality": "high", "image_size": {"width": 1024, "height": 1536}}
return Reply(body=_image_response((f"{wire_url}/files/high.png",), _PROMPT))
return Reply(body=_image_response(((f"{wire_url}/files/high.png", 1024, 1536),), _PROMPT))
assert body == {"prompt": _PROMPT, "quality": "low"}
return Reply(body=_image_response((f"{wire_url}/files/low.png",), _PROMPT))
return Reply(body=_image_response(((f"{wire_url}/files/low.png", 1024, 1536),), _PROMPT))
with wire_server(respond) as wire, gateway.scenario() as scenario:
wire_url: Final = wire.url
@ -85,7 +85,14 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro
)
assert high_response.status_code == 200, high_response.text
high_payload: Final = _JSON_OBJECT.validate_json(high_response.content)
assert high_payload["data"] == [{"url": f"{wire.url}/files/high.png", "b64_json": None, "revised_prompt": None}]
assert high_payload["data"] == [
{
"url": f"{wire.url}/files/high.png",
"b64_json": None,
"revised_prompt": None,
"provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"},
}
]
high_cost: Final = _response_cost(high_response)
assert high_cost == _approx(_catalog_cost("fal_ai/high/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image"))
@ -96,9 +103,16 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro
)
assert low_response.status_code == 200, low_response.text
low_payload: Final = _JSON_OBJECT.validate_json(low_response.content)
assert low_payload["data"] == [{"url": f"{wire.url}/files/low.png", "b64_json": None, "revised_prompt": None}]
assert low_payload["data"] == [
{
"url": f"{wire.url}/files/low.png",
"b64_json": None,
"revised_prompt": None,
"provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"},
}
]
low_cost: Final = _response_cost(low_response)
assert low_cost == _approx(_catalog_cost("fal_ai/low/1024-x-768/openai/gpt-image-2.5/flare/text-to-image"))
assert low_cost == _approx(_catalog_cost("fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image"))
assert high_cost != low_cost
assert [(request.method, request.target) for request in wire.drain()] == [
("POST", "/openai/gpt-image-2.5/flare/text-to-image"),
@ -119,7 +133,7 @@ def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gate
}
return Reply(
body=_image_response(
(f"{wire_url}/files/flux-1.png", f"{wire_url}/files/flux-2.png"),
((f"{wire_url}/files/flux-1.png", 1024, 1024), (f"{wire_url}/files/flux-2.png", 1920, 1080)),
_PROMPT,
)
)
@ -135,11 +149,21 @@ def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gate
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["data"] == [
{"url": f"{wire.url}/files/flux-1.png", "b64_json": None, "revised_prompt": None},
{"url": f"{wire.url}/files/flux-2.png", "b64_json": None, "revised_prompt": None},
{
"url": f"{wire.url}/files/flux-1.png",
"b64_json": None,
"revised_prompt": None,
"provider_specific_fields": {"width": 1024, "height": 1024, "content_type": "image/png"},
},
{
"url": f"{wire.url}/files/flux-2.png",
"b64_json": None,
"revised_prompt": None,
"provider_specific_fields": {"width": 1920, "height": 1080, "content_type": "image/png"},
},
]
cost: Final = _response_cost(response)
assert cost == _approx(2 * _catalog_cost("fal_ai/fal-ai/flux/dev"))
assert cost == _approx(3 * _catalog_cost("fal_ai/fal-ai/flux/dev", "output_cost_per_pixel") * 1_048_576)
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/fal-ai/flux/dev")]
@ -155,7 +179,7 @@ def test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row(
"image_urls": ["data:image/png;base64," + base64.b64encode(_PNG_BYTES).decode()],
"quality": "low",
}
return Reply(body=_image_response((f"{wire_url}/files/edit.png",), _PROMPT))
return Reply(body=_image_response(((f"{wire_url}/files/edit.png", 1024, 1536),), _PROMPT))
with wire_server(respond) as wire, gateway.scenario() as scenario:
wire_url: Final = wire.url
@ -168,9 +192,16 @@ def test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row(
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["data"] == [{"url": f"{wire.url}/files/edit.png", "b64_json": None, "revised_prompt": None}]
assert payload["data"] == [
{
"url": f"{wire.url}/files/edit.png",
"b64_json": None,
"revised_prompt": None,
"provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"},
}
]
cost: Final = _response_cost(response)
assert cost == _approx(_catalog_cost("fal_ai/low/1024-x-768/openai/gpt-image-2.5/flare/edit"))
assert cost == _approx(_catalog_cost("fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/edit"))
assert [(request.method, request.target) for request in wire.drain()] == [
("POST", "/openai/gpt-image-2.5/flare/edit")
]

54
tests/integration/providers/test_fal_ai_video_wire.py
View file

@ -7,6 +7,7 @@ from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
_MODEL: Final = "bytedance/seedance-2.5/text-to-video"
_H3_MODEL: Final = "minimax/h3/text-to-video"
_MP4: Final = b"\x00\x00\x00\x18ftypmp42" + uuid.uuid4().bytes * 4
@ -67,3 +68,56 @@ def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway:
("GET", f"/bytedance/seedance-2.5/requests/{request_id}"),
("GET", f"/files/{request_id}.mp4"),
]
@pytest.mark.covers("other.provider_wire.fal_ai.video_queue_create_status_and_content_download")
def test_fal_h3_video_create_uses_canonical_body_and_status_path(gateway: Gateway) -> None:
request_id: Final = "fal-h3-req-" + uuid.uuid4().hex
def respond(request: Request) -> Reply:
assert request.headers["authorization"] == "Key synthetic-fal-key"
if request.method == "POST":
assert request.target == f"/{_H3_MODEL}"
assert json.loads(request.body) == {
"prompt": "a cat playing volleyball on a beach",
"duration": 6,
"resolution": "2K",
}
return Reply(
body=json.dumps({"status": "IN_QUEUE", "request_id": request_id, "queue_position": 0}).encode()
)
assert request.method == "GET"
if request.target == f"/minimax/h3/requests/{request_id}/status":
return Reply(body=json.dumps({"status": "COMPLETED", "request_id": request_id}).encode())
assert request.target == f"/minimax/h3/requests/{request_id}"
return Reply(body=json.dumps({"video": {"url": f"{wire_url}/files/{request_id}.mp4"}}).encode())
with wire_server(respond) as wire, gateway.scenario() as scenario:
wire_url: Final = wire.url
model: Final = scenario.model(
model=f"fal_ai/{_H3_MODEL}",
api_base=wire.url,
api_key="synthetic-fal-key",
)
created: Final = gateway.post(
"/v1/videos",
{
"model": model,
"prompt": "a cat playing volleyball on a beach",
"seconds": 6,
"size": "2k",
},
)
assert created["status"] == "queued"
video_id: Final = created["id"]
status: Final = gateway.get(f"/v1/videos/{video_id}")
assert status["status"] == "completed"
content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content")
assert content.status_code == 200, content.text
assert content.content == _MP4
assert [(request.method, request.target) for request in wire.drain()] == [
("POST", f"/{_H3_MODEL}"),
("GET", f"/minimax/h3/requests/{request_id}/status"),
("GET", f"/minimax/h3/requests/{request_id}"),
("GET", f"/files/{request_id}.mp4"),
]

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