scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 98

chore: merge main into MCP ordering fix

Browse source
This commit is contained in:
Joshua Valluru 2026-09-21 14:36:41 -07:00
commit 5ec61e7a0b
244 changed files with 27218 additions and 3024 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.circleci/scripts/run_integration.sh
View file

@ -121,6 +121,10 @@ start_proxy() {
"LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_UPSTREAM_URL/_cost_map"
"MODEL_COST_MAP_MIN_MODEL_COUNT=1"
"MODEL_COST_MAP_MAX_SHRINK_RATIO=0"
"GEMINI_API_BASE=$INTEGRATION_UPSTREAM_URL"
"ANTHROPIC_API_BASE=$INTEGRATION_UPSTREAM_URL"
"GEMINI_API_KEY=sk-scripted-provider"
"ANTHROPIC_API_KEY=sk-scripted-provider"
)
else
cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True")

6
.github/scripts/verify_linux_native_wheel.py vendored
View file

@ -205,7 +205,7 @@ def main(
native_module: Final = load_native_module(native_path)
native_module_loads: Final = native_module is not None
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
native_size_limit: Final = 25_000_000
native_size_limit: Final = 30_000_000
native_size_within_limit: Final = native_member.file_size <= native_size_limit
validations: Final = (
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
@ -222,7 +222,7 @@ def main(
("Python extension entry point is present", extension_entry_point_present),
("Native module loads", native_module_loads),
("Production module omits the panic test hook", panic_test_hook_absent),
("Native extension does not exceed 25 MB", native_size_within_limit),
("Native extension does not exceed 30 MB", native_size_within_limit),
("Wheel contents are valid", not unexpected_members),
)
@ -267,7 +267,7 @@ def main(
),
(
not native_size_within_limit,
f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB",
f"native extension exceeds 30 MB: {native_member.file_size / 1_000_000:.2f} MB",
),
(bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"),
)

3
README.md
View file

@ -307,6 +307,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
| [Deepgram (`deepgram`)](https://docs.litellm.ai/docs/providers/deepgram) | ✅ | ✅ | ✅ | | | ✅ | | | | |
| [DeepInfra (`deepinfra`)](https://docs.litellm.ai/docs/providers/deepinfra) | ✅ | ✅ | ✅ | | | | | | | |
| [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | |
| [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | |
| [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | |
| [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | |
| [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | |
@ -356,7 +357,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

136
litellm-rust/Cargo.lock generated
View file

@ -927,6 +927,12 @@ dependencies = [
"libc",
]
[[package]]
name = "crc16"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "338089f42c427b86394a5ee60ff321da23a5c89c9d89514c829687b26359fcff"
[[package]]
name = "crc32fast"
version = "1.5.1"
@ -2460,8 +2466,8 @@ dependencies = [
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
]
[[package]]
@ -2479,12 +2485,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 +2671,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 +2686,7 @@ dependencies = [
"pyo3",
"pyo3-async-runtimes",
"rstest",
"serde",
"serde_json",
"tokio",
"tokio-tungstenite",
@ -2947,6 +2975,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 +2995,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 +3177,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 +3393,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 +3571,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"
@ -3602,9 +3715,13 @@ checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
dependencies = [
"arcstr",
"combine",
"crc16",
"itoa",
"num-bigint",
"num-bigint 0.5.1",
"percent-encoding",
"rand 0.10.2",
"rustls 0.23.42",
"rustls-native-certs",
"ryu",
"sha1_smol",
"socket2 0.6.5",
@ -4001,6 +4118,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 +5080,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 = ["cluster", "tls-rustls"] }
r2d2 = "0.8.10"
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"
serde_json.workspace = true

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

@ -1,58 +1,201 @@
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;
use crate::topology::RedisTopology;
mod connection;
mod operations;
pub(crate) use connection::ConnectionRef;
use connection::{ClusterConnectionManager, ConnectionManager};
pub use operations::{
RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const KEY_PREFIX: &str = "litellm-cache:";
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
pub struct RedisCache<C = redis::Connection> {
connection: Arc<Mutex<C>>,
default_ttl: Duration,
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>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
}
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))
}
}
impl<C> RedisCache<C>
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::Node(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Cluster(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut ConnectionRef::Node(&mut *connection))
}
}
}
}
pub struct RedisCache<S, C = redis::Connection> {
connections: Arc<Connections<C>>,
default_ttl: Duration,
codec: S,
namespace: Option<String>,
topology: RedisTopology,
}
impl<S: CacheCodec> RedisCache<S> {
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
Self::connect(url, &RedisTopology::Standalone, default_ttl, codec)
}
pub fn connect(
url: &str,
topology: &RedisTopology,
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
let connections = match topology {
RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?),
RedisTopology::Cluster { startup_nodes } => {
Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?)
}
};
Ok(Self {
connections: Arc::new(connections),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
topology: topology.clone(),
})
}
}
fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
r2d2::Pool::builder()
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(manager)
.map_err(|_| Error::Unavailable)
}
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,
topology: RedisTopology::Standalone,
}
}
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
pub fn topology(&self) -> &RedisTopology {
&self.topology
}
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
fn namespaced_key(&self, key: &str) -> String {
namespaced_key(self.namespace.as_deref(), key)
}
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
serde_json::from_slice(&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 flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
connection.scan(pattern, 1000, |connection, keys| {
if !keys.is_empty() {
connection
.del::<_, usize>(keys)
.map_err(|_| Error::Unavailable)?;
}
Ok(true)
})
}
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 +204,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)?;
}
Ok(())
.collect::<Result<Vec<_>, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
if entries.is_empty() {
return Ok(());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = entries
.into_iter()
.map(|(key, payload)| {
let mut command = redis::cmd("SETEX");
command.arg(key).arg(ttl).arg(payload);
command
})
.collect();
connection.pipeline(commands).map(drop)
})
.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 connection.ping() {
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 +623,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 +638,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 +661,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 +680,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,

392
litellm-rust/crates/cache-redis/src/cache/connection.rs vendored Normal file
View file

@ -0,0 +1,392 @@
use std::collections::HashMap;
use litellm_cache::Error;
use redis::{
ConnectionAddr, ConnectionInfo, ConnectionLike, IntoConnectionInfo,
cluster::{ClusterClient, ClusterClientBuilder, ClusterConnection, NodeAddress},
cluster_routing::{
MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo, Slot,
},
};
use super::REDIS_TIMEOUT;
use crate::topology::RedisNode;
pub(super) struct PooledConnection<C> {
pub(super) connection: C,
pub(super) failed: bool,
}
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
pub(super) struct ConnectionManager(redis::Client);
impl ConnectionManager {
pub(super) fn open(url: &str) -> Result<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(|_| Error::Unavailable)
}
}
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection<redis::Connection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
pub(super) struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
if startup_nodes.is_empty() {
return Err(Error::Unavailable);
}
let info = url.into_connection_info().map_err(|_| Error::Unavailable)?;
let nodes = startup_nodes
.iter()
.map(|node| node_info(&info, node))
.collect::<Result<Vec<_>, _>>()?;
ClusterClientBuilder::new(nodes)
.connection_timeout(REDIS_TIMEOUT)
.response_timeout(REDIS_TIMEOUT)
.build()
.map(Self)
.map_err(|_| Error::Unavailable)
}
}
fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo, Error> {
let addr = match info.addr() {
ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port),
ConnectionAddr::TcpTls {
insecure,
tls_params,
..
} => ConnectionAddr::TcpTls {
host: node.host.clone(),
port: node.port,
insecure: *insecure,
tls_params: tls_params.clone(),
},
_ => return Err(Error::Unavailable),
};
Ok(info.clone().set_addr(addr))
}
impl r2d2::ManageConnection for ClusterConnectionManager {
type Connection = PooledConnection<ClusterConnection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
pub(crate) enum ConnectionRef<'a> {
Node(&'a mut dyn redis::ConnectionLike),
Cluster(&'a mut ClusterConnection),
}
impl redis::ConnectionLike for ConnectionRef<'_> {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
match self {
Self::Node(connection) => connection.req_packed_command(cmd),
Self::Cluster(connection) => connection.req_packed_command(cmd),
}
}
fn req_packed_commands(
&mut self,
cmd: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
match self {
Self::Node(connection) => connection.req_packed_commands(cmd, offset, count),
Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count),
}
}
fn get_db(&self) -> i64 {
match self {
Self::Node(connection) => connection.get_db(),
Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection),
}
}
fn supports_pipelining(&self) -> bool {
match self {
Self::Node(connection) => connection.supports_pipelining(),
Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection),
}
}
fn check_connection(&mut self) -> bool {
match self {
Self::Node(connection) => connection.check_connection(),
Self::Cluster(connection) => connection.check_connection(),
}
}
fn is_open(&self) -> bool {
match self {
Self::Node(connection) => connection.is_open(),
Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection),
}
}
}
impl ConnectionRef<'_> {
pub(crate) fn pipeline(
&mut self,
commands: Vec<redis::Cmd>,
) -> Result<Vec<redis::Value>, Error> {
match self {
Self::Node(connection) => {
let mut pipeline = redis::pipe();
for command in &commands {
pipeline.add_command(command.clone());
}
pipeline
.query::<Vec<redis::Value>>(*connection)
.map_err(|_| Error::Unavailable)
}
Self::Cluster(connection) => {
let mut replies: Vec<Option<redis::Value>> = vec![None; commands.len()];
for indices in slot_groups(&commands).into_values() {
let mut pipeline = redis::pipe();
for index in &indices {
pipeline.add_command(commands[*index].clone());
}
let values = connection
.req_packed_commands(&pipeline.get_packed_pipeline(), 0, indices.len())
.map_err(|_| Error::Unavailable)?;
if values.len() != indices.len() {
return Err(Error::Unavailable);
}
for (index, value) in indices.into_iter().zip(values) {
replies[index] = Some(value);
}
}
replies
.into_iter()
.collect::<Option<Vec<_>>>()
.ok_or(Error::Unavailable)
}
}
}
pub(crate) fn scan(
&mut self,
pattern: &str,
count: usize,
mut visit: impl FnMut(&mut Self, Vec<String>) -> Result<bool, Error>,
) -> Result<(), Error> {
let pages = match self {
Self::Node(connection) => {
let page = scan_command(0, pattern, count)
.query::<ScanPage>(*connection)
.map_err(|_| Error::Unavailable)?;
vec![(None, page)]
}
Self::Cluster(connection) => connection
.route_command(
&scan_command(0, pattern, count),
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllMasters,
Some(ResponsePolicy::Special),
)),
)
.map_err(|_| Error::Unavailable)
.and_then(primary_pages)?
.into_iter()
.map(|(node, page)| (Some(node), page))
.collect(),
};
for (node, (mut cursor, mut keys)) in pages {
loop {
if !visit(self, keys)? {
return Ok(());
}
if cursor == 0 {
break;
}
(cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?;
}
}
Ok(())
}
pub(crate) fn ping(&mut self) -> Result<bool, redis::RedisError> {
let command = redis::cmd("PING");
match self {
Self::Node(connection) => command
.query::<String>(*connection)
.map(|response| response == "PONG"),
Self::Cluster(connection) => connection
.route_command(
&command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::AllSucceeded),
)),
)
.map(|_| true),
}
}
pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
match self {
Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable),
Self::Cluster(connection) => {
let value = connection
.route_command(
command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::Special),
)),
)
.map_err(|_| Error::Unavailable)?;
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
};
let mut replies = entries
.into_iter()
.map(|(node, reply)| {
Ok((
redis::from_redis_value::<String>(node)
.map_err(|_| Error::Unavailable)?,
redis::from_redis_value::<String>(reply)
.map_err(|_| Error::Unavailable)?,
))
})
.collect::<Result<Vec<(String, String)>, Error>>()?;
replies.sort();
Ok(replies
.into_iter()
.map(|(_, reply)| reply)
.collect::<Vec<_>>()
.join("\n"))
}
}
}
pub(crate) fn flushall(&mut self) -> Result<(), Error> {
let command = redis::cmd("FLUSHALL");
match self {
Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable),
Self::Cluster(connection) => connection
.route_command(
&command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllMasters,
Some(ResponsePolicy::AllSucceeded),
)),
)
.map(|_| ())
.map_err(|_| Error::Unavailable),
}
}
fn scan_page(
&mut self,
node: Option<&NodeAddress>,
cursor: u64,
pattern: &str,
count: usize,
) -> Result<ScanPage, Error> {
let command = scan_command(cursor, pattern, count);
match (self, node) {
(Self::Node(connection), None) => {
command.query(*connection).map_err(|_| Error::Unavailable)
}
(Self::Cluster(connection), Some(node)) => connection
.route_command(
&command,
RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress {
host: node.host().to_string(),
port: node.port(),
}),
)
.map_err(|_| Error::Unavailable)
.and_then(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable)),
_ => Err(Error::Unavailable),
}
}
}
type ScanPage = (u64, Vec<String>);
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
};
entries
.into_iter()
.map(|(node, page)| {
let node = redis::from_redis_value::<String>(node).map_err(|_| Error::Unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(|_| Error::Unavailable)?;
let page = redis::from_redis_value::<ScanPage>(page).map_err(|_| Error::Unavailable)?;
Ok((node, page))
})
.collect()
}
fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
let mut command = redis::cmd("SCAN");
command
.cursor_arg(cursor)
.arg("MATCH")
.arg(pattern)
.arg("COUNT")
.arg(count);
command
}
fn slot_groups(commands: &[redis::Cmd]) -> HashMap<Slot, Vec<usize>> {
let mut groups: HashMap<Slot, Vec<usize>> = HashMap::new();
for (index, command) in commands.iter().enumerate() {
let key = match command.args_iter().nth(1) {
Some(redis::Arg::Simple(key)) => key,
_ => b"",
};
groups.entry(Slot::for_key(key)).or_default().push(index);
}
groups
}

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

@ -0,0 +1,632 @@
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| connection.ping().map_err(|_| Error::Unavailable))
}
pub async fn ping(&self) -> Result<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
connection.ping().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 matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
Ok(matches.len() < count)
})?;
matches.truncate(count);
Ok(matches)
})
.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 sadd = redis::cmd("SADD");
sadd.arg(&key).arg(values);
let mut expire = redis::cmd("EXPIRE");
expire.arg(&key).arg(ttl);
let replies = connection.pipeline(vec![sadd, expire])?;
replies
.into_iter()
.next()
.map(redis::from_redis_value::<usize>)
.transpose()
.map_err(|_| Error::Unavailable)?
.ok_or(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 commands = operations
.into_iter()
.map(|(key, values)| {
let mut command = redis::cmd("RPUSH");
command.arg(key).arg(values);
command
})
.collect();
connection
.pipeline(commands)?
.into_iter()
.map(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable))
.collect()
})
.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 commands = operations
.into_iter()
.map(|(key, count)| {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
command.arg(count);
}
command
})
.collect();
connection.pipeline(commands)
})
.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| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
}
pub fn info(&self) -> Result<String, Error> {
self.connections
.execute(|connection| connection.node_text(&redis::cmd("INFO")))
}
pub fn flushall(&self) -> Result<(), Error> {
self.connections.execute(|connection| connection.flushall())
}
}
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 commands = Vec::with_capacity(operations.len() * 2);
let mut increments = Vec::with_capacity(operations.len());
for (key, amount, ttl) in operations {
let mut increment = redis::cmd("INCRBYFLOAT");
increment.arg(&key).arg(amount);
increments.push(commands.len());
commands.push(increment);
if let Some(ttl) = ttl {
let mut expire = redis::cmd("EXPIRE");
expire.arg(key).arg(ttl);
commands.push(expire);
}
}
let mut replies = connection.pipeline(commands)?;
increments
.into_iter()
.map(|index| {
redis::from_redis_value(std::mem::take(&mut replies[index]))
.map_err(|_| Error::Unavailable)
})
.collect()
})
.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
);
}

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

@ -0,0 +1,492 @@
//! Contract tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a
//! comma separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, CacheScript, ClaimCache,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
ScriptCache,
};
use litellm_cache_redis::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation,
RedisTopology,
};
use redis::cluster_routing::Slot;
type Cache = RedisCache<JsonCodec<serde_json::Value>>;
fn topology() -> Option<RedisTopology> {
let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
let startup_nodes = nodes
.split(',')
.map(|node| {
let (host, port) = node.trim().rsplit_once(':').expect("host:port");
RedisNode {
host: host.to_string(),
port: port.parse().expect("port"),
}
})
.collect();
Some(RedisTopology::Cluster { startup_nodes })
}
fn namespace(label: &str) -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
format!("cluster-test:{label}:{nanos}")
}
fn cluster_url() -> String {
std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
.unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
}
fn cluster_cache(label: &str) -> Option<Cache> {
let topology = topology()?;
Some(
Cache::connect(
&cluster_url(),
&topology,
Some(Duration::from_secs(120)),
JsonCodec::new(),
)
.expect("cluster connection")
.with_namespace(Some(namespace(label))),
)
}
fn counter_cache(label: &str) -> Option<RedisCache<JsonCodec<f64>>> {
let topology = topology()?;
Some(
RedisCache::connect(
&cluster_url(),
&topology,
Some(Duration::from_secs(60)),
JsonCodec::new(),
)
.expect("cluster connection")
.with_namespace(Some(namespace(label))),
)
}
fn multi_slot_keys(count: usize) -> Vec<String> {
let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
let slots: std::collections::HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
assert!(slots.len() > 1, "keys must span multiple slots");
keys
}
macro_rules! cluster_or_skip {
($label:expr) => {
match cluster_cache($label) {
Some(cache) => cache,
None => return,
}
};
}
#[test]
fn constructor_rejects_clusters_without_startup_nodes() {
let error = Cache::connect(
"redis://127.0.0.1:7000",
&RedisTopology::Cluster {
startup_nodes: Vec::new(),
},
None,
JsonCodec::new(),
)
.err();
assert!(matches!(error, Some(Error::Unavailable)));
}
#[test]
fn constructor_rejects_unix_socket_urls_for_clusters() {
let error = Cache::connect(
"redis+unix:///tmp/redis.sock",
&RedisTopology::Cluster {
startup_nodes: vec![RedisNode {
host: "127.0.0.1".into(),
port: 7000,
}],
},
None,
JsonCodec::new(),
)
.err();
assert!(matches!(error, Some(Error::Unavailable)));
}
#[test]
fn single_key_operations_round_trip_with_ttl_rounding() {
let cache = cluster_or_skip!("single");
let context = ExactCacheContext {
ttl: Some(Duration::from_millis(1500)),
};
let keys = multi_slot_keys(12);
for (index, key) in keys.iter().enumerate() {
cache
.set_cache(key, serde_json::json!({ "index": index }), &context)
.unwrap();
}
for (index, key) in keys.iter().enumerate() {
assert_eq!(
cache.get_cache(key, &context).unwrap(),
Some(serde_json::json!({ "index": index }))
);
}
let runtime = tokio::runtime::Runtime::new().unwrap();
let ttl = runtime.block_on(cache.async_get_ttl(&keys[0])).unwrap();
assert_eq!(ttl, Some(2));
cache.delete_cache(&keys[0]).unwrap();
assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None);
assert!(cache.sync_ping().unwrap());
}
#[tokio::test]
async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries() {
let cache = cluster_or_skip!("batch");
let context = ExactCacheContext::default();
let keys = multi_slot_keys(40);
for (index, key) in keys.iter().enumerate() {
if index % 5 == 0 {
continue;
}
cache
.async_set_cache(key, serde_json::json!(index), context.clone())
.await
.unwrap();
}
let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()])
.unwrap()
.get_connection()
.unwrap();
let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]);
redis::cmd("SET")
.arg(&malformed)
.arg("not json")
.exec(&mut raw)
.unwrap();
let entries = cache
.async_batch_get_cache(keys.clone(), context.clone())
.await
.unwrap();
assert_eq!(entries.len(), keys.len());
for (index, entry) in entries.iter().enumerate() {
let expected = if index == 1 {
BatchEntry::Invalid
} else if index % 5 == 0 {
BatchEntry::Miss
} else {
BatchEntry::Hit(serde_json::json!(index))
};
assert_eq!(*entry, expected, "entry {index}");
}
let sync_entries = cache.batch_get_cache(&keys, &context).unwrap();
assert_eq!(sync_entries, entries);
cache.delete_cache_keys(keys.clone()).await.unwrap();
let entries = cache.async_batch_get_cache(keys, context).await.unwrap();
assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss));
}
#[tokio::test]
async fn pipelines_group_by_slot_and_return_results_in_submission_order() {
let cache = cluster_or_skip!("pipeline");
let keys = multi_slot_keys(30);
let entries = keys
.iter()
.enumerate()
.map(|(index, key)| (key.clone(), serde_json::json!(index)))
.collect();
cache
.async_set_cache_pipeline(entries, ExactCacheContext::default())
.await
.unwrap();
let hits = cache
.async_batch_get_cache(keys.clone(), ExactCacheContext::default())
.await
.unwrap();
assert!(
hits.iter()
.enumerate()
.all(|(index, entry)| *entry == BatchEntry::Hit(serde_json::json!(index)))
);
let queues: Vec<String> = keys.iter().map(|key| format!("queue:{key}")).collect();
let pushed = cache
.async_rpush_pipeline(
queues
.iter()
.enumerate()
.map(|(index, key)| RedisRpushOperation {
key: key.clone(),
values: (0..=index)
.map(|value| RedisArg::Integer(value as i64))
.collect(),
})
.collect(),
)
.await
.unwrap();
assert_eq!(pushed, (1..=keys.len()).collect::<Vec<_>>());
let popped = cache
.async_lpop_pipeline(
queues
.iter()
.enumerate()
.map(|(index, key)| RedisLpopOperation {
key: key.clone(),
count: (index % 2 == 0).then_some(2),
})
.collect(),
)
.await
.unwrap();
for (index, result) in popped.into_iter().enumerate() {
match result {
RedisLpopResult::Value(value) => {
assert_eq!(index % 2, 1, "queue {index}");
assert_eq!(value, b"0");
}
RedisLpopResult::Values(values) => {
assert_eq!(index % 2, 0, "queue {index}");
let expected: Vec<Vec<u8>> = (0..=index)
.take(2)
.map(|value| value.to_string().into_bytes())
.collect();
assert_eq!(values, expected);
}
other => panic!("queue {index}: {other:?}"),
}
}
let counters: Vec<String> = keys.iter().map(|key| format!("counter:{key}")).collect();
let Some(counter) = counter_cache("counter") else {
return;
};
let totals = counter
.async_increment_pipeline(
counters
.iter()
.enumerate()
.map(|(index, key)| IncrementOperation {
key: key.clone(),
amount: index as f64 + 0.5,
ttl: (index % 3 == 0).then_some(Duration::from_secs(30)),
})
.collect(),
)
.await
.unwrap();
let expected: Vec<f64> = (0..keys.len()).map(|index| index as f64 + 0.5).collect();
assert_eq!(totals, expected);
assert_eq!(counter.async_get_ttl(&counters[0]).await.unwrap(), Some(30));
assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None);
counter.async_flush_cache().await.unwrap();
cache.async_flush_cache().await.unwrap();
}
#[tokio::test]
async fn scan_and_scoped_flush_cover_every_primary() {
let cache = cluster_or_skip!("flush");
let other = cluster_or_skip!("other");
let context = ExactCacheContext::default();
let keys = multi_slot_keys(60);
for key in &keys {
cache
.async_set_cache(key, serde_json::json!(true), context.clone())
.await
.unwrap();
other
.async_set_cache(key, serde_json::json!(true), context.clone())
.await
.unwrap();
}
let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap();
scanned.sort();
let mut expected: Vec<String> = keys
.iter()
.map(|key| format!("{}:{key}", cache.namespace().unwrap()))
.collect();
expected.sort();
assert_eq!(scanned, expected);
assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7);
cache.flush_cache().unwrap();
let flushed = cache
.async_batch_get_cache(keys.clone(), context.clone())
.await
.unwrap();
assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss));
let kept = other.async_batch_get_cache(keys, context).await.unwrap();
assert!(
kept.iter()
.all(|entry| *entry == BatchEntry::Hit(serde_json::json!(true)))
);
other.async_flush_cache().await.unwrap();
}
fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> {
let mut connection = startup.get_connection().unwrap();
let nodes: String = redis::cmd("CLUSTER")
.arg("NODES")
.query(&mut connection)
.unwrap();
let mut counts: Vec<(String, u64)> = nodes
.lines()
.map(|line| {
let address = line.split_whitespace().nth(1).unwrap();
let address = address.split('@').next().unwrap();
let mut node = redis::Client::open(format!("redis://{address}"))
.unwrap()
.get_connection()
.unwrap();
let stats: String = redis::cmd("INFO")
.arg("commandstats")
.query(&mut node)
.unwrap();
let calls = stats
.lines()
.find_map(|stat| stat.strip_prefix("cmdstat_ping:calls="))
.and_then(|rest| rest.split(',').next())
.map_or(0, |calls| calls.parse().unwrap());
(address.to_string(), calls)
})
.collect();
counts.sort();
counts
}
#[tokio::test]
async fn ping_reaches_every_node() {
let cache = cluster_or_skip!("ping");
let startup = redis::Client::open(cluster_url()).unwrap();
let before = ping_calls_per_node(&startup);
assert!(before.len() >= 2, "{before:?}");
assert!(cache.ping().await.unwrap());
let after = ping_calls_per_node(&startup);
for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) {
assert!(calls_after > calls_before, "{node} was not pinged");
}
assert!(cache.sync_ping().unwrap());
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
}
#[tokio::test]
async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
let Some(counter) = counter_cache("counter") else {
return;
};
let context = ExactCacheContext::default();
assert_eq!(
counter
.increment_cache("spend", 1.5, context.clone())
.unwrap(),
1.5
);
assert_eq!(
counter
.async_increment("spend", 2.0, context.clone())
.await
.unwrap(),
3.5
);
assert_eq!(
counter
.increment_with_floor("budget", -3, Duration::from_secs(30))
.unwrap(),
0
);
assert_eq!(
counter
.async_increment_with_floor("budget", 7, Duration::from_secs(30))
.await
.unwrap(),
7
);
assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0);
assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0);
counter.flush_cache().unwrap();
let cache = cluster_or_skip!("claim");
let owner = serde_json::json!("owner-a");
let rival = serde_json::json!("owner-b");
assert_eq!(
cache
.claim_cache("lock", owner.clone(), &[], context.clone())
.unwrap(),
owner
);
assert_eq!(
cache
.async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone())
.await
.unwrap(),
owner
);
assert_eq!(
cache
.claim_cache("lock", rival.clone(), &[], context.clone())
.unwrap(),
owner
);
assert_eq!(
cache
.async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone())
.await
.unwrap(),
rival
);
let script = cache
.async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into());
let reply = script
.invoke(
vec!["scripted".into()],
vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)],
)
.await
.unwrap();
assert_eq!(reply, redis::Value::Okay);
assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), Some(5));
let evaluated: redis::Value = cache
.async_eval(
"return redis.call('GET', KEYS[1])".into(),
vec!["scripted".into()],
Vec::new(),
)
.await
.unwrap();
assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec()));
assert_eq!(
cache
.async_set_cache_sadd(
"members",
vec![
RedisArg::Bytes(b"a".to_vec()),
RedisArg::Bytes(b"b".to_vec())
],
Some(Duration::from_secs(9)),
)
.await
.unwrap(),
2
);
assert_eq!(cache.async_get_ttl("members").await.unwrap(), Some(9));
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert!(cache.ping().await.unwrap());
let info = cache.info().unwrap();
assert!(info.matches("redis_version").count() > 1, "{info}");
assert!(cache.client_list().unwrap().contains("id="));
cache.async_flush_cache().await.unwrap();
assert_eq!(cache.async_get_ttl("members").await.unwrap(), None);
assert_eq!(cache.get_cache("lock", &context).unwrap(), None);
}

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)
}

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

@ -0,0 +1,826 @@
use std::time::Duration;
use litellm_cache::CacheType;
use litellm_cache_redis::{RedisNode, RedisTopology};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
prelude::*,
types::{PyAny, PyDict, PyList, 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) topology: RedisTopology,
pub(super) connection: RedisConnectionConfig,
}
struct RedisClientProjection<'py> {
topology: RedisTopology,
host: String,
port: u16,
pool_size: usize,
resolved: Bound<'py, PyDict>,
tls: Option<RedisTlsConfig>,
}
const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
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) if service.topology() != Some(&config.topology) => {
Some("facade and native backend topologies 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, "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 projection = if has_value(&source, "startup_nodes")? {
project_cluster_client(&source, &client)?
} else {
project_standalone_client(&client)?
};
let RedisClientProjection {
topology,
host,
port,
pool_size,
resolved,
tls,
} = match projection {
Ok(projection) => projection,
Err(reason) => return Ok(Err(reason)),
};
if has_value(&resolved, "credential_provider")? {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
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>()?,
topology,
connection: RedisConnectionConfig {
host,
port,
database: optional_i64(&resolved, "db")?.unwrap_or(0),
username: optional_dict_string(&resolved, "username")?,
password: optional_dict_string(&resolved, "password")?,
protocol,
pool_size,
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_standalone_client<'py>(
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
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>()?;
if has_value(&resolved, "redis_connect_func")? {
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));
};
Ok(Ok(RedisClientProjection {
topology: RedisTopology::Standalone,
host: required_string(&resolved, "host")?,
port: port(required_i64(&resolved, "port")?)?,
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
resolved,
tls,
}))
}
#[inline(never)]
fn project_cluster_client<'py>(
source: &Bound<'py, PyDict>,
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
let Some(startup_nodes) = startup_nodes(source)? else {
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
};
if !instance_class_is(client, "redis.cluster", "RedisCluster")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let nodes = client.getattr("nodes_manager")?;
if !class_is(
&nodes.getattr("connection_pool_class")?,
"redis.connection",
"ConnectionPool",
)? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let resolved = nodes.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
if let Some(connect) = resolved.get_item("redis_connect_func")?
&& !connect.is_none()
{
let own_hook = connect
.getattr("__self__")
.is_ok_and(|owner| owner.is(client))
&& connect
.getattr("__func__")
.and_then(|function| Ok(function.is(&client.get_type().getattr("on_connect")?)))
.unwrap_or(false);
if !own_hook {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
}
let tls = if optional_bool(&resolved, "ssl")?.unwrap_or(false) {
Some(project_tls(&resolved)?)
} else {
None
};
let first = &startup_nodes[0];
Ok(Ok(RedisClientProjection {
host: first.host.clone(),
port: first.port,
pool_size: optional_i64(&resolved, "max_connections")?
.map(|value| {
usize::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis pool size"))
})
.transpose()?
.unwrap_or(REDIS_PY_DEFAULT_MAX_CONNECTIONS),
topology: RedisTopology::Cluster { startup_nodes },
resolved,
tls,
}))
}
#[inline(never)]
fn startup_nodes(source: &Bound<'_, PyDict>) -> PyResult<Option<Vec<RedisNode>>> {
let Some(nodes) = source.get_item("startup_nodes")? else {
return Ok(None);
};
let Ok(nodes) = nodes.cast_into::<PyList>() else {
return Ok(None);
};
if nodes.is_empty() {
return Ok(None);
}
let mut parsed = Vec::with_capacity(nodes.len());
for node in nodes.iter() {
let Ok(node) = node.cast_into::<PyDict>() else {
return Ok(None);
};
if node.len() != 2 || !has_value(&node, "host")? || !has_value(&node, "port")? {
return Ok(None);
}
let (Ok(host), Ok(port)) = (
required_string(&node, "host"),
required_i64(&node, "port").and_then(port),
) else {
return Ok(None);
};
parsed.push(RedisNode { host, port });
}
Ok(Some(parsed))
}
#[inline(never)]
fn port(value: i64) -> PyResult<u16> {
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
}
#[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 litellm_cache_redis::{RedisNode, RedisTopology};
use super::{
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig,
RedisProtocol,
};
use crate::cache::native::NativeResponseCache;
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
facade(
py,
&format!(
"RedisCluster = type('RedisCluster', (), {{'__module__': 'redis.cluster', 'on_connect': lambda self, connection: None}})\n\
client = RedisCluster()\n\
client.nodes_manager = SimpleNamespace(connection_pool_class=ConnectionPool, connection_kwargs={{'password': 'secret', 'redis_connect_func': {hook}, 'protocol': 3, 'ssl': True, 'ssl_cert_reqs': 'none'}})\n\
backend = SimpleNamespace(default_ttl=120, namespace='team', redis_flush_size=100, redis_kwargs={{'startup_nodes': {startup_nodes}, 'password': 'secret'}}, redis_client=client)\n\
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=100, semantic_cache_scope='key', cache=backend)"
),
)
}
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");
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {
Python::initialize();
Python::attach(|py| {
let facade = cluster_facade(
py,
"[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]",
"client.on_connect",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("cluster startup nodes should project natively");
};
let CacheBackendConfig::Redis(redis) = &config.backend else {
panic!("expected Redis configuration");
};
let expected = RedisTopology::Cluster {
startup_nodes: vec![
RedisNode {
host: "node-a".into(),
port: 7000,
},
RedisNode {
host: "node-b".into(),
port: 7001,
},
],
};
assert_eq!(redis.topology, expected);
assert_eq!(redis.connection.host, "node-a");
assert_eq!(redis.connection.port, 7000);
assert_eq!(redis.connection.password.as_deref(), Some("secret"));
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(
redis
.connection
.tls
.as_ref()
.unwrap()
.certificate_requirement,
CertificateRequirement::None
);
});
}
#[test]
fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() {
Python::initialize();
Python::attach(|py| {
for (startup_nodes, hook, message) in [
(
"[{'host': 'node-a', 'port': 7000, 'server_type': 'primary'}]",
"client.on_connect",
"native Redis topology is not implemented",
),
(
"[{'host': 'node-a', 'port': 'seven'}]",
"client.on_connect",
"native Redis topology is not implemented",
),
(
"[]",
"client.on_connect",
"native Redis topology is not implemented",
),
(
"[{'host': 'node-a', 'port': 7000}]",
"lambda connection: None",
"native Redis credentials require Python",
),
] {
let facade = cluster_facade(py, startup_nodes, hook);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("{startup_nodes} with {hook} must stay on Python");
};
assert_eq!(reason.message(), message, "{startup_nodes} with {hook}");
}
});
}
}

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

@ -0,0 +1,330 @@
use litellm_cache_redis::RedisTopology;
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: Option<usize>,
attributes: RedisPoolAttributes,
}
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
max_connections: Option<&'static str>,
}
const STANDALONE_POOL: RedisPoolAttributes = RedisPoolAttributes {
pool: "connection_pool",
connection_class: "connection_class",
max_connections: Some("max_connections"),
};
const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
pool: "nodes_manager",
connection_class: "connection_pool_class",
max_connections: None,
};
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>, attributes: RedisPoolAttributes) -> PyResult<Self> {
let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?;
Ok(Self {
reference: pool.clone().unbind(),
connection_class: pool.getattr(attributes.connection_class)?.unbind(),
connection_kwargs: pool
.getattr("connection_kwargs")?
.call_method0("copy")?
.unbind(),
max_connections: Self::max_connections(&pool, &attributes)?,
attributes,
})
}
fn max_connections(
pool: &Bound<'_, PyAny>,
attributes: &RedisPoolAttributes,
) -> PyResult<Option<usize>> {
attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let pool = backend
.getattr("redis_client")?
.getattr(self.attributes.pool)?;
Ok(self.reference.bind(py).is(&pool)
&& self
.connection_class
.bind(py)
.is(&pool.getattr(self.attributes.connection_class)?)
&& self.max_connections == Self::max_connections(&pool, &self.attributes)?
&& 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 cluster = matches!(service.topology(), Some(RedisTopology::Cluster { .. }));
let (module, name, cache_kind) = match (kind, cluster) {
("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"),
("redis", true) => (
"litellm.caching.redis_cluster_cache",
"RedisClusterCache",
"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: match (kind, cluster) {
("redis", false) => Some(RedisPoolGuard::capture(&backend, STANDALONE_POOL)?),
("redis", true) => Some(RedisPoolGuard::capture(&backend, CLUSTER_POOL)?),
_ => None,
},
})
}
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)
}

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

@ -0,0 +1,97 @@
use litellm_cache_redis::{RedisNode, RedisTopology};
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, startup_nodes=None))]
fn redis(
py: Python<'_>,
url: String,
ttl_seconds: f64,
namespace: Option<String>,
startup_nodes: Option<Vec<(String, u16)>>,
) -> PyResult<Self> {
let ttl = Some(duration(ttl_seconds)?);
let topology = match startup_nodes {
None => RedisTopology::Standalone,
Some(nodes) => RedisTopology::Cluster {
startup_nodes: nodes
.into_iter()
.map(|(host, port)| RedisNode { host, port })
.collect(),
},
};
let service = release_gil(py, move || {
NativeResponseCache::redis(&url, &topology, 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()),
}
}

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

@ -0,0 +1,207 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
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,
topology: &RedisTopology,
ttl: Option<Duration>,
namespace: Option<String>,
) -> Result<Self, Error> {
let backend =
RedisCache::connect(url, topology, 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 topology(&self) -> Option<&RedisTopology> {
match self {
Self::Memory(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
}
}
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::*;

32
litellm/__init__.py
View file

@ -689,6 +689,7 @@ recraft_models: Set = set()
cometapi_models: Set = set()
oci_models: Set = set()
vercel_ai_gateway_models: Set = set()
edenai_models: Set = set() # mutable-ok: filled from the price map at import, like the sibling provider sets
volcengine_models: Set = set()
wandb_models: Set = set(WANDB_MODELS)
ovhcloud_models: Set = set()
@ -763,6 +764,8 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None:
openrouter_models.add(key)
elif value.get("litellm_provider") == "vercel_ai_gateway":
vercel_ai_gateway_models.add(key)
elif value.get("litellm_provider") == "edenai":
edenai_models.add(key)
elif value.get("litellm_provider") == "datarobot":
datarobot_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-text-models":
@ -1111,6 +1114,7 @@ model_list = list(
| oci_models
| heroku_models
| vercel_ai_gateway_models
| edenai_models
| volcengine_models
| wandb_models
| ovhcloud_models
@ -1139,6 +1143,7 @@ def _build_models_by_provider() -> dict:
"baseten": baseten_models,
"openrouter": openrouter_models,
"vercel_ai_gateway": vercel_ai_gateway_models,
"edenai": edenai_models,
"datarobot": datarobot_models,
"vertex_ai": vertex_chat_models
| vertex_text_models
@ -1684,6 +1689,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,
@ -2114,6 +2122,30 @@ if TYPE_CHECKING:
from .llms.vercel_ai_gateway.chat.transformation import (
VercelAIGatewayConfig as VercelAIGatewayConfig,
)
from .llms.edenai.chat.transformation import (
EdenAIChatConfig as EdenAIChatConfig,
)
from .llms.edenai.responses.transformation import (
EdenAIResponsesAPIConfig as EdenAIResponsesAPIConfig,
)
from .llms.edenai.messages.transformation import (
EdenAIAnthropicMessagesConfig as EdenAIAnthropicMessagesConfig,
)
from .llms.edenai.embedding.transformation import (
EdenAIEmbeddingConfig as EdenAIEmbeddingConfig,
)
from .llms.edenai.audio_transcription.transformation import (
EdenAIAudioTranscriptionConfig as EdenAIAudioTranscriptionConfig,
)
from .llms.edenai.text_to_speech.transformation import (
EdenAITextToSpeechConfig as EdenAITextToSpeechConfig,
)
from .llms.edenai.image_generation.transformation import (
EdenAIImageGenerationConfig as EdenAIImageGenerationConfig,
)
from .llms.edenai.videos.transformation import (
EdenAIVideoConfig as EdenAIVideoConfig,
)
from .llms.ovhcloud.chat.transformation import (
OVHCloudChatConfig as OVHCloudChatConfig,
)

24
litellm/_lazy_imports_registry.py
View file

@ -176,6 +176,7 @@ LLM_CONFIG_NAMES: Final = (
"BedrockClaudePlatformMessagesConfig",
"AmazonAnthropicClaudeMessagesConfig",
"AmazonMantleMessagesConfig",
"BedrockMantleAnthropicMessagesConfig",
"TogetherAIConfig",
"TogetherAIChatConfig",
"NLPCloudConfig",
@ -326,6 +327,14 @@ LLM_CONFIG_NAMES: Final = (
"InceptionChatConfig",
"HyperbolicChatConfig",
"VercelAIGatewayConfig",
"EdenAIChatConfig",
"EdenAIResponsesAPIConfig",
"EdenAIAnthropicMessagesConfig",
"EdenAIEmbeddingConfig",
"EdenAIAudioTranscriptionConfig",
"EdenAITextToSpeechConfig",
"EdenAIImageGenerationConfig",
"EdenAIVideoConfig",
"OVHCloudChatConfig",
"OVHCloudEmbeddingConfig",
"CometAPIEmbeddingConfig",
@ -746,6 +755,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",
@ -1227,6 +1240,17 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.vercel_ai_gateway.chat.transformation",
"VercelAIGatewayConfig",
),
"EdenAIChatConfig": (".llms.edenai.chat.transformation", "EdenAIChatConfig"),
"EdenAIResponsesAPIConfig": (".llms.edenai.responses.transformation", "EdenAIResponsesAPIConfig"),
"EdenAIAnthropicMessagesConfig": (".llms.edenai.messages.transformation", "EdenAIAnthropicMessagesConfig"),
"EdenAIEmbeddingConfig": (".llms.edenai.embedding.transformation", "EdenAIEmbeddingConfig"),
"EdenAIAudioTranscriptionConfig": (
".llms.edenai.audio_transcription.transformation",
"EdenAIAudioTranscriptionConfig",
),
"EdenAITextToSpeechConfig": (".llms.edenai.text_to_speech.transformation", "EdenAITextToSpeechConfig"),
"EdenAIImageGenerationConfig": (".llms.edenai.image_generation.transformation", "EdenAIImageGenerationConfig"),
"EdenAIVideoConfig": (".llms.edenai.videos.transformation", "EdenAIVideoConfig"),
"OVHCloudChatConfig": (".llms.ovhcloud.chat.transformation", "OVHCloudChatConfig"),
"OVHCloudEmbeddingConfig": (
".llms.ovhcloud.embedding.transformation",

42
litellm/anthropic_beta_headers_config.json
View file

@ -11,6 +11,7 @@
"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",
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": "fast-mode-2026-02-01",
"files-api-2025-04-14": "files-api-2025-04-14",
@ -44,6 +45,7 @@
"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",
"dangerous-tool-use-2026-09-03": null,
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": "files-api-2025-04-14",
@ -76,6 +78,7 @@
"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": null,
"dangerous-tool-use-2026-09-03": null,
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
@ -109,6 +112,7 @@
"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",
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
@ -131,6 +135,42 @@
"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",
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"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",
@ -142,6 +182,7 @@
"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",
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": null,
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
@ -175,6 +216,7 @@
"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",
"dangerous-tool-use-2026-09-03": null,
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": "fast-mode-2026-02-01",
"files-api-2025-04-14": "files-api-2025-04-14",

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

10
litellm/caching/redis_cache.py
View file

@ -80,6 +80,8 @@ class _AsyncRedisCommands(Protocol):
def ttl(self, name: str) -> Awaitable[int]: ...
def expire(self, name: str, time: int) -> Awaitable[bool]: ...
def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ...
def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ...
@ -1948,6 +1950,14 @@ class RedisCache(BaseCache):
_record_swallowed_redis_failure(self._circuit_breaker, e)
return None
@_redis_circuit_breaker_guard
async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool:
"""EXPIRE an existing key without touching its value. False when the key is absent."""
_used_ttl: Final = self.get_ttl(ttl=ttl)
if _used_ttl is None:
return False
return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl)
@_redis_circuit_breaker_guard
async def async_rpush(
self,

4
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
@ -749,6 +750,7 @@ LITELLM_CHAT_PROVIDERS: Final = [
"inception",
"vercel_ai_gateway",
"wandb",
"edenai",
"ovhcloud",
"lemonade",
"docker_model_runner",
@ -924,6 +926,7 @@ openai_compatible_endpoints: Final[list] = [
"https://api.hyperbolic.xyz/v1",
"https://ai-gateway.helicone.ai/",
"https://ai-gateway.vercel.sh/v1",
"https://api.edenai.run/v3",
"https://api.inference.wandb.ai/v1",
"https://api.clarifai.com/v2/ext/openai/v1",
"https://api.libertai.io/v1",
@ -993,6 +996,7 @@ openai_compatible_providers: Final[list] = [
"hyperbolic",
"vercel_ai_gateway",
"aiml",
"edenai",
"wandb",
"cometapi",
"clarifai",

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

1
litellm/images/main.py
View file

@ -388,6 +388,7 @@ def image_generation(
litellm.LlmProviders.DASHSCOPE,
litellm.LlmProviders.QWENCLOUD,
litellm.LlmProviders.QWEN_AI_PLATFORM,
litellm.LlmProviders.EDENAI,
):
if image_generation_config is None:
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")

18
litellm/integrations/otel/model/payloads.py
View file

@ -427,7 +427,7 @@ class LLMCallSpanData:
# plain ``.get`` — no repeated ``isinstance`` guards.
raw_response: Final = payload.get("response")
response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {})
choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response)
choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response) or _ocr_choices(response)
# ``finish_reasons`` is metadata, not content, so derive it from
# ``choices_out`` before gating. The raw message/choice bodies are only
# retained when content capture is enabled (see ``capture_span_content``);
@ -752,6 +752,22 @@ def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]:
return (choice,)
def _ocr_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]:
markdowns: Final = tuple(
text for page in _dicts(response.get("pages")) if (text := as_str(page.get("markdown"))) is not None
)
if not markdowns:
return ()
message: Final[_AssistantMessage] = {
"role": "assistant",
"content": "\n\n".join(markdowns),
"refusal": None,
"tool_calls": None,
}
choice: Final[_Choice] = {"message": message, "finish_reason": None}
return (choice,)
def _responses_parts_text(parts: tuple[Mapping[str, object], ...], part_type: str, field: str) -> str | None:
texts: Final = tuple(
text for part in parts if part.get("type") == part_type if (text := as_str(part.get(field))) is not None

24
litellm/litellm_core_utils/core_helpers.py
View file

@ -4,7 +4,8 @@ import copy
import logging
import re
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
import httpx
from pydantic import TypeAdapter, ValidationError
@ -703,3 +704,24 @@ def redact_nested_match_and_regex_keys(
except Exception:
return payload
return redacted
RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost"
_NO_HEADERS: Final[Mapping[str, object]] = MappingProxyType({})
class _CarriesHiddenParams(Protocol):
_hidden_params: dict[str, object] # mutable-ok: the responses billed here keep hidden params in a plain dict
def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: float | None) -> None:
"""Record a provider-reported cost where the cost calculator looks before the price map."""
if cost is None:
return
hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor
additional_headers: Final[object] = hidden_params.get("additional_headers")
merged: Final[dict[str, object]] = { # mutable-ok: assigned into the plain-dict hidden params
**(additional_headers if isinstance(additional_headers, Mapping) else _NO_HEADERS),
RESPONSE_COST_HEADER: cost,
}
hidden_params["additional_headers"] = merged # rebind-ok: the caller's record is the point

6
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -362,6 +362,9 @@ def get_llm_provider(
elif endpoint == "https://ai-gateway.vercel.sh/v1":
custom_llm_provider = "vercel_ai_gateway"
dynamic_api_key = get_secret_str("VERCEL_AI_GATEWAY_API_KEY")
elif endpoint == "https://api.edenai.run/v3":
custom_llm_provider = "edenai" # rebind-ok: api_base detection resolves the provider in place
dynamic_api_key = get_secret_str("EDENAI_API_KEY")
elif endpoint == "https://api.inference.wandb.ai/v1":
custom_llm_provider = "wandb"
dynamic_api_key = get_secret_str("WANDB_API_KEY")
@ -853,6 +856,9 @@ def _get_openai_compatible_provider_info(
api_base,
dynamic_api_key,
) = litellm.VercelAIGatewayConfig()._get_openai_compatible_provider_info(api_base, api_key)
elif custom_llm_provider == "edenai":
api_base = litellm.EdenAIChatConfig.get_api_base(api_base) # rebind-ok: chain resolves in place
dynamic_api_key = litellm.EdenAIChatConfig.get_api_key(api_key) # rebind-ok: chain resolves in place
elif custom_llm_provider == "aiml":
(
api_base,

7
litellm/litellm_core_utils/litellm_logging.py
View file

@ -69,7 +69,11 @@ from litellm.litellm_core_utils.classifier_logging import (
classifier_input_snapshot,
is_classifier_call,
)
from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name
from litellm.litellm_core_utils.core_helpers import (
is_expected_client_error,
reconstruct_model_name,
set_response_cost_in_hidden_params,
)
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.internal_call_metadata import (
MODEL_ACCESS_GROUP_METADATA_KEY,
@ -3918,6 +3922,7 @@ class Logging(LiteLLMLoggingBaseClass):
):
## return unified Usage object
if isinstance(result.response.usage, ResponseAPIUsage):
set_response_cost_in_hidden_params(result.response, result.response.usage.cost)
transformed_usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
result.response.usage
)

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,

13
litellm/llms/base_llm/videos/transformation.py
View file

@ -272,6 +272,19 @@ class BaseVideoConfig(ABC):
) -> VideoObject:
pass
async def async_transform_video_status_retrieve_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str | None = None,
) -> VideoObject:
"""Async transform video status retrieve response."""
return self.transform_video_status_retrieve_response(
raw_response=raw_response,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
)
def transform_video_create_character_request(
self,
name: 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:

25
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)
@ -530,6 +533,9 @@ class AmazonAnthropicClaudeMessagesConfig(
if anthropic_model_info.is_eager_input_streaming_used(tools):
beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER)
if anthropic_messages_optional_request_params.get("safeguards") is not None:
beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value)
self._filter_context_management_for_bedrock_invoke(
anthropic_messages_request=anthropic_messages_request,
beta_set=beta_set,
@ -546,15 +552,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/custom_httpx/llm_http_handler.py
View file

@ -8881,7 +8881,7 @@ class BaseLLMHTTPHandler:
url=url,
headers=headers,
)
return video_status_provider_config.transform_video_status_retrieve_response(
return await video_status_provider_config.async_transform_video_status_retrieve_response(
raw_response=response,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,

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

91
litellm/llms/edenai/audio_transcription/transformation.py Normal file
View file

@ -0,0 +1,91 @@
"""
Support for OpenAI's `/v1/audio/transcriptions` endpoint on Eden AI, served at `/v3/audio/transcriptions`
with the real per-request cost at the top level of the JSON body.
Docs: https://www.edenai.co/docs/api-reference/audio/audio-transcriptions
"""
from collections.abc import Mapping
from typing import Final
import httpx
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params
from litellm.llms.base_llm.audio_transcription.transformation import AudioTranscriptionRequestData
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.transcriptions.whisper_transformation import OpenAIWhisperAudioTranscriptionConfig
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import FileTypes, TranscriptionResponse
from litellm.utils import convert_to_model_response_object
from ..common_utils import EdenAIException, authorized_headers, endpoint_url, reported_cost
def _form_fields(model: str, optional_params: Mapping[str, object]) -> dict[str, object]: # mutable-ok: httpx form data
"""LiteLLM parks non-OpenAI params, `model` included, under `extra_body` for the OpenAI SDK; a
multipart body carries them as top-level text fields instead."""
extras: Final = optional_params.get("extra_body")
nested: Final = extras.items() if isinstance(extras, Mapping) else ()
fields: Final = (*optional_params.items(), *nested, ("model", model))
return {key: value for key, value in fields if key != "extra_body"} # mutable-ok: httpx form data
class EdenAIAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig):
@property
def has_native_transcription_endpoint(self) -> bool:
return True
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
stream: bool | None = None,
) -> str:
return endpoint_url(api_base, "audio/transcriptions")
def validate_environment(
self,
headers: dict[str, object], # mutable-ok: inherited contract
model: str,
messages: list[AllMessageValues], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, object]: # mutable-ok: inherited contract
return authorized_headers(headers, api_key, model)
def transform_audio_transcription_request(
self,
model: str,
audio_file: FileTypes,
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
) -> AudioTranscriptionRequestData:
"""Eden reports `duration` and `cost` on every body, so the Whisper default of `verbose_json`,
which the gpt-4o-transcribe models reject, is not needed for cost tracking."""
audio: Final = process_audio_file(audio_file)
files: Final = {"file": (audio.filename, audio.file_content, audio.content_type)} # mutable-ok: httpx contract
return AudioTranscriptionRequestData(data=_form_fields(model, optional_params), files=files)
def transform_audio_transcription_response(self, raw_response: httpx.Response) -> TranscriptionResponse:
if "application/json" not in raw_response.headers.get("content-type", ""):
return TranscriptionResponse(text=raw_response.text)
body: Final = raw_response.json()
response: Final[TranscriptionResponse] = convert_to_model_response_object(
response_object=body, model_response_object=TranscriptionResponse(), response_type="audio_transcription"
)
set_response_cost_in_hidden_params(response, reported_cost(body))
return response
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract
) -> BaseLLMException:
return EdenAIException(message=error_message, status_code=status_code, headers=headers)

145
litellm/llms/edenai/chat/transformation.py Normal file
View file

@ -0,0 +1,145 @@
"""
Support for OpenAI's `/v1/chat/completions` endpoint on Eden AI.
Eden AI is an OpenAI-compatible gateway (one key across 1000+ models), so requests go through the
shared HTTP handler untouched. Every Eden response reports the real per-request cost at the top
level of the body; the only translation here lifts that number into LiteLLM's cost tracking.
Docs: https://www.edenai.co/docs
"""
from collections.abc import AsyncIterator, Iterator, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import BaseModel, TypeAdapter
import litellm
from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
from ..common_utils import EdenAIException, reported_cost, resolve_api_base, resolve_api_key
if TYPE_CHECKING:
import tiktoken
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_OPTIONAL_MAPPING: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
class _EdenAIModel(BaseModel):
id: str
class _EdenAIModelCatalog(BaseModel):
data: tuple[_EdenAIModel, ...]
def _stream_options_with_usage(request: Mapping[str, object]) -> Mapping[str, object]:
current: Final = _OPTIONAL_MAPPING.validate_python(request.get("stream_options")) or MappingProxyType({})
return MappingProxyType({**current, "include_usage": True})
class EdenAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler):
def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: inherited contract
parsed: Final = super().chunk_parser(chunk)
cost: Final = reported_cost(chunk)
usage: Final[object] = getattr(parsed, "usage", None)
if cost is not None and isinstance(usage, Usage):
usage.cost = cost
return parsed
class EdenAIChatConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract
reasoning: Final[tuple[str, ...]] = (
("reasoning_effort",)
if litellm.supports_reasoning(model=model, custom_llm_provider=litellm.LlmProviders.EDENAI.value)
else ()
)
return [*super().get_supported_openai_params(model), *reasoning] # mutable-ok: inherited contract
@staticmethod
def get_api_key(api_key: str | None = None) -> str | None:
return resolve_api_key(api_key)
@staticmethod
def get_api_base(api_base: str | None = None) -> str:
return resolve_api_base(api_base)
def transform_request(
self,
model: str,
messages: list[AllMessageValues], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
headers: dict[str, object], # mutable-ok: inherited contract
) -> dict[str, object]: # mutable-ok: inherited contract
request: Final[dict[str, object]] = super().transform_request( # mutable-ok: inherited contract
model, messages, optional_params, litellm_params, headers
)
if not request.get("stream"):
return request
return {**request, "stream_options": dict(_stream_options_with_usage(request))} # mutable-ok: JSON body
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict[str, object], # mutable-ok: inherited contract
messages: list[AllMessageValues], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:
response: Final = super().transform_response(
model=model,
raw_response=raw_response,
model_response=model_response,
logging_obj=logging_obj,
request_data=request_data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
api_key=api_key,
json_mode=json_mode,
)
set_response_cost_in_hidden_params(response, reported_cost(raw_response.content))
return response
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract
) -> BaseLLMException:
return EdenAIException(message=error_message, status_code=status_code, headers=headers)
def get_model_response_iterator(
self,
streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
sync_stream: bool,
json_mode: bool | None = False,
) -> EdenAIChatCompletionStreamingHandler:
return EdenAIChatCompletionStreamingHandler(
streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode
)
def get_models(
self, api_key: str | None = None, api_base: str | None = None
) -> list[str]: # mutable-ok: inherited contract
response: Final = litellm.module_level_client.get(url=f"{self.get_api_base(api_base)}/models")
if not response.is_success:
raise EdenAIException(status_code=response.status_code, message=response.text, headers=response.headers)
catalog: Final = _EdenAIModelCatalog.model_validate(response.json())
return [f"edenai/{model.id}" for model in catalog.data] # mutable-ok: inherited contract

80
litellm/llms/edenai/common_utils.py Normal file
View file

@ -0,0 +1,80 @@
"""
Pieces shared by every Eden AI endpoint: credentials, the exception class, and the per-request
`cost` Eden reports at the top level of each response body, or in a header when the body is binary.
"""
from collections.abc import Container, Mapping
from types import MappingProxyType
from typing import Final
from pydantic import AliasChoices, BaseModel, Field, ValidationError
import litellm
from litellm.exceptions import AuthenticationError
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import LlmProviders
EDENAI_API_BASE: Final = "https://api.edenai.run/v3"
EDENAI_COST_HEADER: Final = "x-edenai-cost"
class EdenAIException(BaseLLMException):
pass
class _EdenAIExtras(BaseModel):
cost: float | None = Field(default=None, validation_alias=AliasChoices("cost", EDENAI_COST_HEADER))
def resolve_api_base(api_base: str | None) -> str:
return api_base or get_secret_str("EDENAI_API_BASE") or EDENAI_API_BASE
def resolve_api_key(api_key: str | None) -> str | None:
return api_key or get_secret_str("EDENAI_API_KEY")
def require_api_key(api_key: str | None, model: str) -> str:
resolved: Final = resolve_api_key(api_key or litellm.api_key)
if resolved is None:
raise AuthenticationError(
message="Missing Eden AI API key: set EDENAI_API_KEY or pass api_key",
llm_provider=LlmProviders.EDENAI.value,
model=model,
)
return resolved
def reported_cost(payload: object) -> float | None:
try:
extras: Final = (
_EdenAIExtras.model_validate_json(payload)
if isinstance(payload, bytes)
else _EdenAIExtras.model_validate(payload)
)
except ValidationError:
return None
return extras.cost
def authorized_headers(
headers: Mapping[str, object], api_key: str | None, model: str
) -> dict[str, object]: # mutable-ok: header contract
return {**headers, "Authorization": f"Bearer {require_api_key(api_key, model)}"} # mutable-ok: header contract
def json_headers(
headers: Mapping[str, object], api_key: str | None, model: str
) -> dict[str, object]: # mutable-ok: header contract
"""The shared HTTP handler sends some JSON bodies as raw content, so the type must be set here."""
authorized: Final = authorized_headers(headers, api_key, model)
return {**authorized, "Content-Type": "application/json"} # mutable-ok: header contract
def endpoint_url(api_base: str | None, path: str) -> str:
return f"{resolve_api_base(api_base).rstrip('/')}/{path}"
def pick(params: Mapping[str, object], keys: Container[str]) -> Mapping[str, object]:
return MappingProxyType({key: value for key, value in params.items() if key in keys})

97
litellm/llms/edenai/embedding/transformation.py Normal file
View file

@ -0,0 +1,97 @@
"""
Support for OpenAI's `/v1/embeddings` endpoint on Eden AI, served at `/v3/embeddings` with the real
per-request cost at the top level of the body.
Docs: https://www.edenai.co/docs/v3/llms/embeddings
"""
from typing import TYPE_CHECKING, Final
import httpx
from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse
from litellm.utils import convert_to_model_response_object
from ..common_utils import EdenAIException, endpoint_url, json_headers, pick, reported_cost
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_SUPPORTED_PARAMS: Final = ("dimensions", "encoding_format", "user")
class EdenAIEmbeddingConfig(BaseEmbeddingConfig):
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract
return list(_SUPPORTED_PARAMS) # mutable-ok: inherited contract
def map_openai_params(
self,
non_default_params: dict[str, object], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: inherited contract
return {**optional_params, **pick(non_default_params, _SUPPORTED_PARAMS)} # mutable-ok: inherited contract
def validate_environment(
self,
headers: dict[str, object], # mutable-ok: inherited contract
model: str,
messages: list[AllMessageValues], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, object]: # mutable-ok: inherited contract
return json_headers(headers, api_key, model)
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
stream: bool | None = None,
) -> str:
return endpoint_url(api_base, "embeddings")
def transform_embedding_request(
self,
model: str,
input: AllEmbeddingInputValues,
optional_params: dict[str, object], # mutable-ok: inherited contract
headers: dict[str, object], # mutable-ok: inherited contract
) -> dict[str, object]: # mutable-ok: inherited contract
return {"model": model, "input": input, **optional_params} # mutable-ok: inherited contract
def transform_embedding_response(
self,
model: str,
raw_response: httpx.Response,
model_response: EmbeddingResponse,
logging_obj: "LiteLLMLoggingObj",
api_key: str | None,
request_data: dict[str, object], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
) -> EmbeddingResponse:
body: Final = raw_response.json()
logging_obj.post_call(original_response=body)
response: Final[EmbeddingResponse] = convert_to_model_response_object(
response_object=body, model_response_object=model_response, response_type="embedding"
)
set_response_cost_in_hidden_params(response, reported_cost(body))
return response
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract
) -> BaseLLMException:
return EdenAIException(message=error_message, status_code=status_code, headers=headers)

115
litellm/llms/edenai/image_generation/transformation.py Normal file
View file

@ -0,0 +1,115 @@
"""
Support for OpenAI's `/v1/images/generations` endpoint on Eden AI, served at `/v3/images/generations`
for every image model in the catalog with the real per-request cost at the top level of the body.
Docs: https://www.edenai.co/docs/v3/llms/image-generation
"""
from typing import TYPE_CHECKING, Final
import httpx
from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.image_generation.transformation import BaseImageGenerationConfig
from litellm.types.llms.openai import AllMessageValues, OpenAIImageGenerationOptionalParams
from litellm.types.utils import ImageResponse
from litellm.utils import convert_to_model_response_object
from ..common_utils import EdenAIException, endpoint_url, json_headers, pick, reported_cost
if TYPE_CHECKING:
import tiktoken
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_SUPPORTED_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = (
"background",
"moderation",
"n",
"output_compression",
"output_format",
"quality",
"response_format",
"size",
"style",
"user",
)
class EdenAIImageGenerationConfig(BaseImageGenerationConfig):
def get_supported_openai_params(
self, model: str
) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: inherited contract
return list(_SUPPORTED_PARAMS) # mutable-ok: inherited contract
def map_openai_params(
self,
non_default_params: dict[str, object], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: inherited contract
return {**optional_params, **pick(non_default_params, _SUPPORTED_PARAMS)} # mutable-ok: inherited contract
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
stream: bool | None = None,
) -> str:
return endpoint_url(api_base, "images/generations")
def validate_environment(
self,
headers: dict[str, object], # mutable-ok: inherited contract
model: str,
messages: list[AllMessageValues], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, object]: # mutable-ok: inherited contract
return json_headers(headers, api_key, model)
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
headers: dict[str, object], # mutable-ok: inherited contract
) -> dict[str, object]: # mutable-ok: inherited contract
return {"model": model, "prompt": prompt, **optional_params} # mutable-ok: inherited contract
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict[str, object], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:
body: Final = raw_response.json()
logging_obj.post_call(original_response=body)
response: Final[ImageResponse] = convert_to_model_response_object(
response_object=body, model_response_object=model_response, response_type="image_generation"
)
set_response_cost_in_hidden_params(response, reported_cost(body))
return response
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract
) -> BaseLLMException:
return EdenAIException(message=error_message, status_code=status_code, headers=headers)

79
litellm/llms/edenai/messages/transformation.py Normal file
View file

@ -0,0 +1,79 @@
"""
Support for Anthropic's `/v1/messages` endpoint on Eden AI.
Eden AI serves the Anthropic Messages API at `/v3/v1/messages` for every model in its catalog, so
the Anthropic payload is forwarded untranslated and the answer comes back in Anthropic's shape with
Eden's per-request `cost` beside it. Eden does not report a cost inside a Messages stream yet, so
streams fall back to the price map.
Docs: https://www.edenai.co/docs/api-reference/anthropic-messages/create-anthropic-message
"""
from typing import TYPE_CHECKING, Final
import httpx
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai_like.json_loader import SimpleProviderConfig
from litellm.llms.openai_like.messages.transformation import JSONProviderAnthropicMessagesConfig
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
from litellm.types.utils import LlmProviders
from ..common_utils import EDENAI_API_BASE, EdenAIException, reported_cost, require_api_key
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_EDENAI_PROVIDER_SPEC: Final[dict[str, str]] = { # mutable-ok: SimpleProviderConfig takes a plain dict
"base_url": EDENAI_API_BASE,
"api_key_env": "EDENAI_API_KEY",
"api_base_env": "EDENAI_API_BASE",
}
_EDENAI_PROVIDER: Final = SimpleProviderConfig(LlmProviders.EDENAI.value, _EDENAI_PROVIDER_SPEC)
class EdenAIAnthropicMessagesConfig(JSONProviderAnthropicMessagesConfig):
def __init__(self) -> None:
super().__init__(_EDENAI_PROVIDER)
def validate_anthropic_messages_environment(
self,
headers: dict[str, str], # mutable-ok: inherited contract
model: str,
messages: list[object], # mutable-ok: inherited contract
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict[str, str], str | None]: # mutable-ok: inherited contract
return super().validate_anthropic_messages_environment(
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=require_api_key(api_key, model),
api_base=api_base,
)
def transform_anthropic_messages_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> AnthropicMessagesResponse:
response: Final = super().transform_anthropic_messages_response(
model=model, raw_response=raw_response, logging_obj=logging_obj
)
cost: Final = reported_cost(response)
if cost is not None:
logging_obj.model_call_details["response_cost"] = cost # rebind-ok: the per-call record spend logging reads
return response
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract
) -> BaseLLMException:
return EdenAIException(message=error_message, status_code=status_code, headers=headers)

80
litellm/llms/edenai/responses/transformation.py Normal file
View file

@ -0,0 +1,80 @@
"""
Support for OpenAI's `/v1/responses` endpoint on Eden AI.
Eden AI serves the Responses API at `/v3/responses` in OpenAI's wire format, so the OpenAI config
does the work; this one points it at Eden and authenticates with the Eden key. Eden reports the
per-request cost on `usage.cost` of every body, the final `response.completed` event included, so
the shared usage-cost lift bills both modes.
Docs: https://www.edenai.co/docs/v3/llms/responses
"""
from typing import TYPE_CHECKING, Final
import httpx
from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from ..common_utils import EdenAIException, authorized_headers, resolve_api_base
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
class EdenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.EDENAI
def validate_environment(
self,
headers: dict[str, object], # mutable-ok: inherited contract
model: str,
litellm_params: GenericLiteLLMParams | None,
) -> dict[str, object]: # mutable-ok: inherited contract
return authorized_headers(headers, litellm_params.api_key if litellm_params else None, model)
def get_complete_url(
self,
api_base: str | None,
litellm_params: dict[str, object], # mutable-ok: inherited contract
) -> str:
return super().get_complete_url(api_base=resolve_api_base(api_base), litellm_params=litellm_params)
def transform_response_api_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> ResponsesAPIResponse:
response: Final = super().transform_response_api_response(
model=model, raw_response=raw_response, logging_obj=logging_obj
)
set_response_cost_in_hidden_params(response, response.usage.cost if response.usage else None)
return response
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract
) -> BaseLLMException:
return EdenAIException(message=error_message, status_code=status_code, headers=headers)
def should_fake_stream(
self,
model: str | None,
stream: bool | None,
custom_llm_provider: str | None = None,
) -> bool:
"""Eden streams every catalog model natively; the base class would fake-stream any model the
price map does not know, which is all of them."""
return False
def supports_native_websocket(self) -> bool:
return False

85
litellm/llms/edenai/text_to_speech/transformation.py Normal file
View file

@ -0,0 +1,85 @@
"""
Support for OpenAI's `/v1/audio/speech` endpoint on Eden AI, served at `/v3/audio/speech`. The answer
is raw audio, so the real per-request cost travels in the `x-edenai-cost` response header.
Docs: https://www.edenai.co/docs/api-reference/audio/audio-speech
"""
from typing import TYPE_CHECKING, Final
import httpx
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig, TextToSpeechRequestData
from litellm.types.llms.openai import HttpxBinaryResponseContent
from ..common_utils import EdenAIException, endpoint_url, json_headers, reported_cost
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_SUPPORTED_PARAMS: Final = ("voice", "response_format", "speed", "instructions")
class EdenAITextToSpeechConfig(BaseTextToSpeechConfig):
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: inherited contract
return list(_SUPPORTED_PARAMS) # mutable-ok: inherited contract
def map_openai_params(
self,
model: str,
optional_params: dict[str, object], # mutable-ok: inherited contract
voice: str | dict[str, object] | None = None, # mutable-ok: inherited contract
drop_params: bool = False,
kwargs: dict[str, object] | None = None, # mutable-ok: inherited contract
) -> tuple[str | None, dict[str, object]]: # mutable-ok: inherited contract
return (voice if isinstance(voice, str) else None), optional_params
def validate_environment(
self,
headers: dict[str, object], # mutable-ok: inherited contract
model: str,
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, object]: # mutable-ok: inherited contract
return json_headers(headers, api_key, model)
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: dict[str, object], # mutable-ok: inherited contract
) -> str:
return endpoint_url(api_base, "audio/speech")
def transform_text_to_speech_request(
self,
model: str,
input: str,
voice: str | None,
optional_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: dict[str, object], # mutable-ok: inherited contract
headers: dict[str, object], # mutable-ok: inherited contract
) -> TextToSpeechRequestData:
fields: Final = (("model", model), ("input", input), ("voice", voice), *optional_params.items())
return TextToSpeechRequestData(
dict_body={key: value for key, value in fields if value is not None} # mutable-ok: TypedDict field
)
def transform_text_to_speech_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> HttpxBinaryResponseContent:
response: Final = HttpxBinaryResponseContent(response=raw_response)
response.set_response_cost(reported_cost(raw_response.headers))
return response
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract
) -> BaseLLMException:
return EdenAIException(message=error_message, status_code=status_code, headers=headers)

146
litellm/llms/edenai/videos/transformation.py Normal file
View file

@ -0,0 +1,146 @@
"""
Support for OpenAI's `/v1/videos` API on Eden AI, served at `/v3/videos`. A job is created, polled and
downloaded through the OpenAI routes; Eden reports `cost` as 0 on the create response and the settled
amount on the status read once the job completes or fails.
Docs: https://www.edenai.co/docs/v3/llms/video-generation
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
import httpx
from httpx._types import RequestFiles
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
from litellm.types.router import GenericLiteLLMParams
from litellm.types.videos.main import VideoObject
from ..common_utils import EdenAIException, authorized_headers, endpoint_url, reported_cost
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
def _usage_with_reported_cost(
usage: Mapping[str, object] | None, body: bytes
) -> dict[str, object]: # mutable-ok: VideoObject.usage is a plain dict field
cost: Final = reported_cost(body)
return { # mutable-ok: VideoObject.usage is a plain dict field
key: value
for key, value in (*(usage.items() if usage else ()), ("provider_reported_cost_usd", cost))
if value is not None
}
class EdenAIVideoConfig(OpenAIVideoConfig):
def validate_environment(
self,
headers: dict[str, object], # mutable-ok: inherited contract
model: str,
api_key: str | None = None,
litellm_params: GenericLiteLLMParams | None = None,
) -> dict[str, object]: # mutable-ok: inherited contract
return authorized_headers(headers, api_key or (litellm_params.api_key if litellm_params else None), model)
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: dict[str, object], # mutable-ok: inherited contract
) -> str:
return endpoint_url(api_base, "videos")
def use_multipart_form_data(self) -> bool:
return False
def transform_video_create_request(
self,
model: str,
prompt: str,
api_base: str,
video_create_optional_request_params: dict[str, object], # mutable-ok: inherited contract
litellm_params: GenericLiteLLMParams,
headers: dict[str, object], # mutable-ok: inherited contract
) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: inherited contract
"""A reference image is a multipart file part, or a JSON `{"file_id"}` / `{"image_url"}` object."""
reference: Final = video_create_optional_request_params.get("input_reference")
if not isinstance(reference, Mapping):
return super().transform_video_create_request(
model=model,
prompt=prompt,
api_base=api_base,
video_create_optional_request_params=video_create_optional_request_params,
litellm_params=litellm_params,
headers=headers,
)
data, files, url = super().transform_video_create_request(
model=model,
prompt=prompt,
api_base=api_base,
video_create_optional_request_params={ # mutable-ok: inherited contract
key: value for key, value in video_create_optional_request_params.items() if key != "input_reference"
},
litellm_params=litellm_params,
headers=headers,
)
return {**data, "input_reference": dict(reference)}, files, url # mutable-ok: JSON body
def transform_video_create_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
custom_llm_provider: str | None = None,
request_data: dict[str, object] | None = None, # mutable-ok: inherited contract
) -> VideoObject:
video: Final = super().transform_video_create_response(
model=model,
raw_response=raw_response,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
request_data=request_data,
)
video.usage = _usage_with_reported_cost(video.usage, raw_response.content)
return video
def transform_video_status_retrieve_response(
self,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
custom_llm_provider: str | None = None,
) -> VideoObject:
raw_response.raise_for_status() # the shared GET helpers return error bodies instead of raising
video: Final = super().transform_video_status_retrieve_response(
raw_response=raw_response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider
)
video.usage = _usage_with_reported_cost(video.usage, raw_response.content)
return video
def transform_video_content_response(
self,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> bytes:
raw_response.raise_for_status() # the shared GET helpers return error bodies instead of raising
return raw_response.content
def transform_video_list_response(
self,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
custom_llm_provider: str | None = None,
) -> dict[str, str]: # mutable-ok: inherited contract
raw_response.raise_for_status() # the shared GET helpers return error bodies instead of raising
return super().transform_video_list_response(
raw_response=raw_response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider
)
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: inherited contract
) -> BaseLLMException:
return EdenAIException(message=error_message, status_code=status_code, headers=headers)

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):

299
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 collections.abc import Callable, 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})
@ -130,12 +163,127 @@ def _response_data(raw_response: httpx.Response) -> Mapping[str, object]:
return TypeAdapter(Mapping[str, object]).validate_python(raw_response.json())
def _response_data_or_none(raw_response: httpx.Response) -> Mapping[str, object] | None:
try:
return _response_data(raw_response)
except ValueError:
return None
def _detail_item_text(item: Mapping[str, object]) -> str | None:
message: Final[object] = item.get("msg")
if not isinstance(message, str):
return None
location: Final[object] = item.get("loc")
if isinstance(location, str) and location:
return f"{location}: {message}"
if isinstance(location, (list, tuple)):
location_parts: Final[tuple[str, ...]] = tuple(part for part in location if isinstance(part, str))
if location_parts:
return f"{'.'.join(location_parts)}: {message}"
return message
def _error_text(response_data: Mapping[str, object]) -> str | None:
detail: Final[object] = response_data.get("detail")
if isinstance(detail, str):
return detail
if isinstance(detail, list):
detail_items: Final[tuple[Mapping[str, object], ...]] = tuple(
item for item in detail if isinstance(item, Mapping)
)
detail_messages: Final[tuple[str, ...]] = tuple(
message for item in detail_items if (message := _detail_item_text(item)) is not None
)
if detail_messages:
return "; ".join(detail_messages)
error: Final[object] = response_data.get("error")
return error if isinstance(error, str) else None
def _result_error(raw_response: httpx.Response) -> str | None:
if raw_response.is_success:
return None
response_data: Final[Mapping[str, object] | None] = _response_data_or_none(raw_response)
error_text: Final[str | None] = _error_text(response_data) if response_data is not None else None
if error_text:
return error_text
response_text: Final[str] = raw_response.text
return response_text or f"fal.ai returned HTTP {raw_response.status_code}"
def _terminal_result_error(raw_response: httpx.Response) -> str | None:
if raw_response.status_code == 429 or raw_response.status_code >= 500:
return None
return _result_error(raw_response)
def _get_fal_ai_async_httpx_client() -> AsyncHTTPHandler:
return get_async_httpx_client(llm_provider=LlmProviders.FAL_AI)
def _response_string(response_data: Mapping[str, object], key: str, default: str = "") -> str:
value: Final[object] = response_data.get(key)
return value if isinstance(value, str) else default
def _result_request(
raw_response: httpx.Response,
response_data: Mapping[str, object],
) -> tuple[str, Mapping[str, str]] | None:
if _response_string(response_data, "status", "IN_QUEUE") != "COMPLETED":
return None
result_url: Final[str] = str(raw_response.request.url).removesuffix("/status")
result_headers: Final[Mapping[str, str]] = MappingProxyType(
{
key: value
for key, value in (
("Authorization", raw_response.request.headers.get("Authorization")),
("Content-Type", raw_response.request.headers.get("Content-Type")),
)
if value is not None
}
)
return result_url, result_headers
def _status_video_object(
response_data: Mapping[str, object],
raw_response: httpx.Response,
custom_llm_provider: str | None,
result_error: str | None,
) -> VideoObject:
raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE")
status: Final[str] = _STATUS_MAP.get(raw_status, "queued")
status_error: Final[str | None] = _error_text(response_data)
error: Final[str | None] = result_error if result_error is not None else status_error
provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER
model_path: Final[str | None] = _model_path_from_request_url(raw_response)
request_id: Final[str] = _response_string(response_data, "request_id") or (
_request_id_from_request_url(raw_response) or ""
)
return VideoObject(
id=encode_video_id_with_provider(request_id, provider, model_path),
object="video",
status="failed" if error else status,
created_at=0,
model=model_path,
error=(
{"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict
),
)
class FalAIVideoConfig(BaseVideoConfig):
def __init__(
self,
sync_client_factory: Callable[[], HTTPHandler] = _get_httpx_client,
async_client_factory: Callable[[], AsyncHTTPHandler] = _get_fal_ai_async_httpx_client,
) -> None:
super().__init__()
self._sync_client_factory: Final = sync_client_factory
self._async_client_factory: Final = async_client_factory
def get_supported_openai_params(self, model: str) -> _SupportedParams:
supported_params: Final[_SupportedParams] = [ # mutable-ok: BaseVideoConfig requires a list
"model",
@ -158,18 +306,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 +347,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 +408,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 +420,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
}
@ -299,25 +460,58 @@ class FalAIVideoConfig(BaseVideoConfig):
custom_llm_provider: str | None = None,
) -> VideoObject:
response_data: Final[Mapping[str, object]] = _response_data(raw_response)
raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE")
status: Final[str] = _STATUS_MAP.get(raw_status, "queued")
error_value: Final[object] = response_data.get("error")
error: Final[str | None] = error_value if isinstance(error_value, str) else None
provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER
model_path: Final[str | None] = _model_path_from_request_url(raw_response)
request_id: Final[str] = _response_string(response_data, "request_id") or (
_request_id_from_request_url(raw_response) or ""
result_error: Final[str | None] = self._fetch_result_error(raw_response, response_data)
return _status_video_object(
response_data=response_data,
raw_response=raw_response,
custom_llm_provider=custom_llm_provider,
result_error=result_error,
)
return VideoObject(
id=encode_video_id_with_provider(request_id, provider, model_path),
object="video",
status="failed" if error else status,
created_at=0,
model=model_path,
error=(
{"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict
),
def _fetch_result_error(
self,
raw_response: httpx.Response,
response_data: Mapping[str, object],
) -> str | None:
result_request: Final[tuple[str, Mapping[str, str]] | None] = _result_request(raw_response, response_data)
if result_request is None:
return None
result_url, result_headers = result_request
result_response: Final[httpx.Response] = self._sync_client_factory().get(
url=result_url,
headers=result_headers,
)
return _terminal_result_error(result_response)
async def async_transform_video_status_retrieve_response(
self,
raw_response: httpx.Response,
logging_obj: object,
custom_llm_provider: str | None = None,
) -> VideoObject:
response_data: Final[Mapping[str, object]] = _response_data(raw_response)
result_error: Final[str | None] = await self._fetch_result_error_async(raw_response, response_data)
return _status_video_object(
response_data=response_data,
raw_response=raw_response,
custom_llm_provider=custom_llm_provider,
result_error=result_error,
)
async def _fetch_result_error_async(
self,
raw_response: httpx.Response,
response_data: Mapping[str, object],
) -> str | None:
result_request: Final[tuple[str, Mapping[str, str]] | None] = _result_request(raw_response, response_data)
if result_request is None:
return None
result_url, result_headers = result_request
result_response: Final[httpx.Response] = await self._async_client_factory().get(
url=result_url,
headers=result_headers,
)
return _terminal_result_error(result_response)
@staticmethod
def _decode_video_id(video_id: str) -> tuple[str, str]:
@ -355,17 +549,23 @@ class FalAIVideoConfig(BaseVideoConfig):
video_url: Final[object] = video_data.get("url")
if isinstance(video_url, str) and video_url:
return video_url
error_message: Final[str | None] = next(
(value for key in ("error", "detail") if isinstance(value := response_data.get(key), str)),
None,
)
error_message: Final[str | None] = _error_text(response_data)
if error_message:
raise ValueError(f"fal.ai video result did not include a video URL: {error_message}")
raise ValueError("fal.ai video result did not include a video URL")
def transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes:
error: Final[str | None] = _result_error(raw_response)
if error is not None:
raise FalAIVideoError(
status_code=raw_response.status_code,
message=error,
headers=dict(raw_response.headers), # mutable-ok: exception headers require a mutable dictionary
request=raw_response.request,
response=raw_response,
)
video_url: Final[str] = self._extract_video_url(_response_data(raw_response))
httpx_client: Final[HTTPHandler] = _get_httpx_client()
httpx_client: Final[HTTPHandler] = self._sync_client_factory()
video_response: Final[httpx.Response] = httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped
video_url
)
@ -373,8 +573,17 @@ class FalAIVideoConfig(BaseVideoConfig):
return video_response.content
async def async_transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes:
error: Final[str | None] = _result_error(raw_response)
if error is not None:
raise FalAIVideoError(
status_code=raw_response.status_code,
message=error,
headers=dict(raw_response.headers), # mutable-ok: exception headers require a mutable dictionary
request=raw_response.request,
response=raw_response,
)
video_url: Final[str] = self._extract_video_url(_response_data(raw_response))
async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client(llm_provider=LlmProviders.FAL_AI)
async_httpx_client: Final[AsyncHTTPHandler] = self._async_client_factory()
video_response: Final[httpx.Response] = await async_httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped
video_url
)

3
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py
View file

@ -108,6 +108,9 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
if anthropic_model_info.is_tool_search_used(tools):
beta_values.add(get_tool_search_beta_header("vertex_ai"))
if optional_params.get("safeguards") is not None:
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value)
if beta_values:
headers["anthropic-beta"] = ",".join(beta_values)

62
litellm/main.py
View file

@ -3568,6 +3568,32 @@ def _complete_vercel_ai_gateway(
return response
def _complete_edenai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult:
api_base: Final = litellm.EdenAIChatConfig.get_api_base(ctx.api_base)
api_key: Final = litellm.EdenAIChatConfig.get_api_key(ctx.api_key or litellm.api_key)
response: Final = base_llm_http_handler.completion(
model=ctx.model,
messages=ctx.messages,
api_base=api_base,
custom_llm_provider="edenai",
model_response=ctx.model_response,
encoding=_get_encoding(),
logging_obj=ctx.logging,
optional_params=ctx.optional_params,
timeout=ctx.timeout,
litellm_params=ctx.litellm_params,
shared_session=ctx.shared_session,
acompletion=ctx.acompletion,
stream=ctx.stream,
api_key=api_key,
headers=ctx.headers or litellm.headers,
client=_dispatch_client_http(ctx),
provider_config=ctx.provider_config,
)
ctx.logging.post_call(input=ctx.messages, api_key=api_key, original_response=response)
return response
def _complete_vertex_ai_beta(
ctx: _CompletionDispatchContext,
) -> _CompletionDispatchResult:
@ -5771,6 +5797,8 @@ def completion(
response = _complete_minimax(_dispatch_ctx)
elif custom_llm_provider == "hosted_vllm":
response = _complete_hosted_vllm(_dispatch_ctx)
elif custom_llm_provider == "edenai":
response = _complete_edenai(_dispatch_ctx) # rebind-ok: dispatch chain binds response per branch
elif (
# A known OpenAI model name only decides the route when nothing else
# resolved a provider. get_llm_provider() already maps these names to
@ -6440,6 +6468,22 @@ def embedding(
litellm_params=litellm_params_dict,
headers=headers or {},
)
elif custom_llm_provider == "edenai":
response = base_llm_http_handler.embedding(
model=model,
input=input,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
logging_obj=logging,
timeout=timeout,
model_response=EmbeddingResponse(),
optional_params=optional_params,
client=client,
aembedding=aembedding,
litellm_params=litellm_params_dict,
headers=headers,
)
elif (
custom_llm_provider == "openai_like"
or custom_llm_provider == "llamafile"
@ -8142,7 +8186,23 @@ def speech(
custom_llm_provider=custom_llm_provider,
)
response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None
if custom_llm_provider == "openai" or (
if custom_llm_provider == "edenai":
litellm_params_dict["api_base"] = api_base
response = base_llm_http_handler.text_to_speech_handler(
model=model,
input=input,
voice=voice if isinstance(voice, str) else None,
text_to_speech_provider_config=text_to_speech_provider_config or litellm.EdenAITextToSpeechConfig(),
text_to_speech_optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client,
_is_async=aspeech or False,
)
elif custom_llm_provider == "openai" or (
custom_llm_provider in litellm.openai_compatible_providers
and custom_llm_provider not in AZURE_OPENAI_AUDIO_PROVIDERS
):

165
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.0741e-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.81482e-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.56175e-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,
@ -76812,5 +76889,65 @@
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"openrouter/xiaomi/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"openrouter/xiaomi/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"openrouter/xiaomi/mimo-v2.6-pro-ultraspeed": {
"cache_read_input_token_cost": 3.6e-08,
"input_cost_per_token": 4.35e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_pdf_input": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
}
}

20
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,
@ -815,6 +815,24 @@
"interactions": true
}
},
"edenai": {
"display_name": "Eden AI (`edenai`)",
"url": "https://docs.litellm.ai/docs/providers/edenai",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true,
"embeddings": true,
"image_generations": true,
"audio_transcriptions": true,
"audio_speech": true,
"moderations": false,
"batches": false,
"rerank": false,
"interactions": false,
"video_generations": true
}
},
"duckduckgo": {
"display_name": "DuckDuckGo (`duckduckgo`)",
"url": "https://docs.litellm.ai/docs/search/duckduckgo",

24
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -2500,9 +2500,8 @@ class MCPServerManager:
# Filter blank scopes (e.g. YAML ``scopes: [""]``) the same way the DB-build path does, so
# an all-blank list normalizes to None rather than a ``("",)`` tuple that skips the
# entra_obo fail-closed scope precondition and POSTs an empty scope to the IdP.
resolved_scopes = self._extract_scopes(server_config.get("scopes")) or (
gated_oauth_metadata.scopes if gated_oauth_metadata else None
)
configured_scopes = self._extract_scopes(server_config.get("scopes"))
resolved_scopes = configured_scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None)
resolved_authorization_url = manual_authorization_url or (
gated_oauth_metadata.authorization_url if gated_oauth_metadata else None
)
@ -2579,6 +2578,7 @@ class MCPServerManager:
client_secret=server_config.get("client_secret", None),
oauth2_flow=self._explicit_oauth2_flow(config_oauth2_flow),
scopes=resolved_scopes,
configured_scopes=tuple(configured_scopes) if configured_scopes else None,
issuer=effective_issuer,
issuer_is_anchored=use_issuer_anchor,
authorization_url=resolved_authorization_url,
@ -3055,6 +3055,18 @@ class MCPServerManager:
if scopes_value is not None:
scopes = self._extract_scopes(scopes_value)
stored_scopes: Final[object] = credentials_dict.get("scopes") if credentials_dict else None
scopes_as_objects: Final = (
cast(Sequence[object], stored_scopes) # cast-ok: list shape validated below
if isinstance(stored_scopes, list)
else ()
)
configured_scopes: Final = (
tuple(scope for scope in scopes_as_objects if isinstance(scope, str))
if scopes_as_objects and all(isinstance(scope, str) and scope for scope in scopes_as_objects)
else None
)
name_for_prefix: Final = mcp_server.alias or mcp_server.server_name or mcp_server.server_id
mcp_info: Final[MCPInfo] = _mcp_info.copy()
@ -3129,6 +3141,7 @@ class MCPServerManager:
client_secret=client_secret_value or getattr(mcp_server, "client_secret", None),
oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)),
scopes=resolved_scopes,
configured_scopes=configured_scopes,
issuer=effective_issuer,
issuer_is_anchored=use_issuer_anchor,
authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None),
@ -7094,6 +7107,11 @@ class MCPServerManager:
spec_path=server.spec_path,
transport=server.transport,
auth_type=server.auth_type,
credentials=(
{"scopes": list(server.configured_scopes)} # mutable-ok: MCPCredentials requires a JSON-array list
if server.configured_scopes
else None
),
created_at=server.created_at,
updated_at=server.updated_at,
teams=[],

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": [
{

1
litellm/proxy/_types.py
View file

@ -641,6 +641,7 @@ class LiteLLMRoutes(enum.Enum):
"/v1/models",
"/sso/get/ui_settings",
"/get/user_banner",
"/get/latest_release_info",
]
# NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend

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

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