scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

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

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 23:51:16 +00:00
commit c31fa2ae7a
51 changed files with 5855 additions and 841 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

195
litellm-rust/Cargo.lock generated
View file

@ -40,6 +40,12 @@ dependencies = [
"cc",
]
[[package]]
name = "allocator-api2"
version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
[[package]]
name = "android_system_properties"
version = "0.1.6"
@ -230,6 +236,7 @@ dependencies = [
"aws-credential-types",
"aws-sigv4",
"aws-smithy-async",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-runtime",
"aws-smithy-runtime-api",
@ -238,7 +245,9 @@ dependencies = [
"bytes",
"bytes-utils",
"fastrand",
"http 0.2.12",
"http 1.4.2",
"http-body 0.4.6",
"http-body 1.1.0",
"percent-encoding",
"pin-project-lite",
@ -272,6 +281,43 @@ dependencies = [
"tracing",
]
[[package]]
name = "aws-sdk-s3"
version = "1.146.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2cd651b4400d4011b8927b83a9552bf90ff11e6e5da0b9f0a7583247aceec971"
dependencies = [
"arc-swap",
"aws-credential-types",
"aws-runtime",
"aws-sigv4",
"aws-smithy-async",
"aws-smithy-checksums",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-json",
"aws-smithy-observability",
"aws-smithy-runtime",
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"aws-smithy-xml 0.62.1",
"aws-types",
"bytes",
"fastrand",
"hex",
"hmac",
"http 0.2.12",
"http 1.4.2",
"http-body 1.1.0",
"lru",
"percent-encoding",
"regex-lite",
"sha2 0.11.0",
"tracing",
"url",
]
[[package]]
name = "aws-sdk-secretsmanager"
version = "1.117.0"
@ -316,7 +362,7 @@ dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"aws-smithy-xml",
"aws-smithy-xml 0.61.1",
"aws-types",
"fastrand",
"http 0.2.12",
@ -332,6 +378,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70"
dependencies = [
"aws-credential-types",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-runtime-api",
"aws-smithy-types",
@ -359,10 +406,31 @@ dependencies = [
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.1"
name = "aws-smithy-checksums"
version = "0.65.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
checksum = "b67ecd999972b58e67cab052f5129906c08c25883bd0788ceefc55ef97d61307"
dependencies = [
"aws-smithy-http",
"aws-smithy-types",
"bytes",
"crc-fast",
"hex",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"md-5",
"pin-project-lite",
"sha1 0.11.0",
"sha2 0.11.0",
"tracing",
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "80c2051c2f1016fb8e6548dd07b8bc2ac9c3fe583721444b92f515e856d31609"
dependencies = [
"aws-smithy-types",
"bytes",
@ -375,6 +443,7 @@ version = "0.64.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-runtime-api",
"aws-smithy-types",
"bytes",
@ -554,6 +623,18 @@ dependencies = [
"xmlparser",
]
[[package]]
name = "aws-smithy-xml"
version = "0.62.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b932c8d6dc127fc980eecd78f8694ae9b9551b69a93a7def2a199c1c0033daf"
dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"xmlparser",
]
[[package]]
name = "aws-types"
version = "1.6.0"
@ -980,6 +1061,16 @@ dependencies = [
"libc",
]
[[package]]
name = "crc-fast"
version = "1.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5"
dependencies = [
"digest 0.10.7",
"spin",
]
[[package]]
name = "crc16"
version = "0.4.0"
@ -1348,7 +1439,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -1925,6 +2016,8 @@ version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
dependencies = [
"allocator-api2",
"equivalent",
"foldhash",
]
@ -2148,7 +2241,7 @@ dependencies = [
"libc",
"percent-encoding",
"pin-project-lite",
"socket2 0.6.5",
"socket2 0.5.10",
"tokio",
"tower-service",
"tracing",
@ -2679,6 +2772,37 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-s3"
version = "0.1.0"
dependencies = [
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-types",
"aws-types",
"litellm-auth-aws",
"litellm-cache",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"tokio",
"uuid",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -2841,6 +2965,7 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
@ -2850,6 +2975,8 @@ dependencies = [
"litellm-cache-redis",
"litellm-cache-redis-semantic",
"litellm-cache-response",
"litellm-cache-s3",
"litellm-cache-valkey-semantic",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -2860,10 +2987,12 @@ dependencies = [
"litellm-types",
"pyo3",
"pyo3-async-runtimes",
"redis",
"rstest",
"serde",
"serde_json",
"serde_with",
"sha2 0.10.9",
"tokio",
"tokio-tungstenite",
]
@ -3091,6 +3220,15 @@ version = "0.4.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
[[package]]
name = "lru"
version = "0.18.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ff9840bcc50b71349309900da0ce7279aa336ae71d73250b07998932c7d97c25"
dependencies = [
"hashbrown 0.17.1",
]
[[package]]
name = "lru-slab"
version = "0.1.2"
@ -3113,6 +3251,16 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
[[package]]
name = "md-5"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98"
dependencies = [
"cfg-if",
"digest 0.11.3",
]
[[package]]
name = "memchr"
version = "2.8.3"
@ -3755,7 +3903,7 @@ dependencies = [
"quinn-udp",
"rustc-hash",
"rustls 0.23.42",
"socket2 0.6.5",
"socket2 0.5.10",
"thiserror 2.0.19",
"tokio",
"tracing",
@ -3794,9 +3942,9 @@ dependencies = [
"cfg_aliases",
"libc",
"once_cell",
"socket2 0.6.5",
"socket2 0.5.10",
"tracing",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4282,7 +4430,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4353,7 +4501,7 @@ dependencies = [
"security-framework",
"security-framework-sys",
"webpki-root-certs",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4626,6 +4774,17 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "sha1"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214"
dependencies = [
"cfg-if",
"cpufeatures 0.3.0",
"digest 0.11.3",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
@ -4742,6 +4901,12 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "spin"
version = "0.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3"
[[package]]
name = "spm_precompiled"
version = "0.1.4"
@ -4908,10 +5073,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"getrandom 0.3.4",
"once_cell",
"rustix",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -5375,7 +5540,7 @@ dependencies = [
"rand 0.8.7",
"rustls 0.23.42",
"rustls-pki-types",
"sha1",
"sha1 0.10.7",
"thiserror 1.0.69",
"utf-8",
]
@ -5784,7 +5949,7 @@ version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]

1
litellm-rust/Cargo.toml
View file

@ -33,6 +33,7 @@ litellm-cache = { path = "crates/cache" }
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-s3 = { path = "crates/cache-s3" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }

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

@ -8,29 +8,26 @@ use litellm_cache::{
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
SemanticCacheContext,
};
use litellm_cache_redis::connection::{ConnectionRef, Connections, ttl_seconds};
use litellm_cache_redis::{
RedisTopology,
connection::{ConnectionRef, Connections},
};
use litellm_cache_response::{CacheEntry, ResponseCacheCodec};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::prompt::prompt_from_context;
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
const CACHE_KEY_FIELD: &str = "litellm_cache_key";
const VECTOR_FIELD: &str = "prompt_vector";
pub trait Embedder: Send + Sync + 'static {
fn embed(
&self,
prompt: &str,
metadata: &serde_json::Map<String, Value>,
) -> Result<Vec<f32>, Error>;
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error>;
fn async_embed(
&self,
prompt: &str,
metadata: &serde_json::Map<String, Value>,
metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
}
@ -212,7 +209,7 @@ pub struct RedisSemanticCache<E: Embedder, C = redis::Connection> {
impl<E: Embedder> RedisSemanticCache<E> {
pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?),
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
inner: Arc::new(Inner::new(config)),
})
@ -277,7 +274,7 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
let Some(prompt) = prompt_from_context(context) else {
return Ok(());
};
let vector = self.embedder.embed(&prompt, &context.metadata)?;
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context).to_string();
self.connections.execute(|connection| {
self.inner
@ -289,7 +286,7 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
let Some(prompt) = prompt_from_context(context) else {
return Ok(None);
};
let vector = self.embedder.embed(&prompt, &context.metadata)?;
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context).to_string();
self.connections
.execute(|connection| self.inner.lookup(connection, &tag, &vector))
@ -306,7 +303,7 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
};
let vector = self
.embedder
.async_embed(&prompt, &context.metadata)
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, &context).to_string();
let inner = Arc::clone(&self.inner);
@ -326,7 +323,7 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
};
let vector = self
.embedder
.async_embed(&prompt, &context.metadata)
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, context).to_string();
let inner = Arc::clone(&self.inner);
@ -613,3 +610,9 @@ fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
_ => None,
}
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}

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

@ -2,8 +2,10 @@ use litellm_cache::SemanticCacheContext;
use serde_json::Value;
pub fn prompt_from_context(context: &SemanticCacheContext) -> Option<String> {
if !context.messages.is_empty() {
return Some(messages_text(&context.messages));
if let Some(messages) = context.messages.as_ref().and_then(Value::as_array)
&& !messages.is_empty()
{
return Some(messages_text(messages));
}
let input = context.input.as_ref()?;
let mut parts = Vec::new();

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

@ -35,7 +35,7 @@ impl FakeEmbedder {
}
impl Embedder for FakeEmbedder {
fn embed(&self, prompt: &str, _: &serde_json::Map<String, Value>) -> Result<Vec<f32>, Error> {
fn embed(&self, prompt: &str, _: Option<&Value>) -> Result<Vec<f32>, Error> {
self.calls.lock().unwrap().push(prompt.to_string());
Ok(self
@ -45,11 +45,7 @@ impl Embedder for FakeEmbedder {
.unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
}
async fn async_embed(
&self,
prompt: &str,
metadata: &serde_json::Map<String, Value>,
) -> Result<Vec<f32>, Error> {
async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.embed(prompt, metadata)
}
}
@ -63,7 +59,7 @@ fn config() -> RedisSemanticConfig {
fn messages_context(messages: Vec<Value>) -> SemanticCacheContext {
SemanticCacheContext {
messages,
messages: Some(Value::Array(messages)),
..Default::default()
}
}

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

@ -40,6 +40,7 @@ const CLAIM_SCRIPT: &str = concat!(
);
const CLAIM_ATTEMPTS: usize = 8;
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
@ -50,21 +51,6 @@ impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
pub fn pooled(url: &str, timeout: Duration, pool_size: u32) -> Result<Self, Error> {
let pool = r2d2::Pool::builder()
.max_size(pool_size)
.min_idle(Some(0))
.connection_timeout(timeout)
.test_on_check_out(false)
.build(ConnectionManager::open(url)?)
.map_err(|_| Error::Unavailable)?;
Ok(Self::Pool(pool))
}
pub fn fixed(connection: C) -> Self {
Self::Fixed(Mutex::new(connection))
}
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
@ -98,6 +84,19 @@ where
.await
.map_err(|_| Error::Unavailable)?
}
pub fn fixed(connection: C) -> Self {
Self::Fixed(Mutex::new(connection))
}
pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
match topology {
RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
ClusterConnectionManager::open(url, startup_nodes)?,
)?)),
}
}
}
pub struct RedisCache<S, C = redis::Connection> {
@ -119,12 +118,7 @@ impl<S: CacheCodec> RedisCache<S> {
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)?)?)
}
};
let connections = Connections::open(url, topology)?;
Ok(Self {
connections: Arc::new(connections),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
@ -224,22 +218,10 @@ where
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl_seconds(ttl)
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
async fn run_blocking<T, F>(connections: Arc<Connections<C>>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
Connections::run_blocking(connections, operation).await
}
}
pub fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
@ -338,7 +320,7 @@ where
if entries.is_empty() {
return Ok(());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = entries
.into_iter()
.map(|(key, payload)| {
@ -357,7 +339,7 @@ where
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Self::run_blocking(Arc::clone(&self.connections), |connection| {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match connection.ping() {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,

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

@ -188,7 +188,7 @@ where
}
pub async fn ping(&self) -> Result<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
Connections::run_blocking(Arc::clone(&self.connections), |connection| {
connection.ping().map_err(|_| Error::Unavailable)
})
.await
@ -208,7 +208,7 @@ where
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| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
@ -231,7 +231,7 @@ where
}
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| {
Connections::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");
@ -279,7 +279,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, values)| {
@ -333,7 +333,7 @@ where
.iter()
.map(|(_, count)| count.is_some())
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, count)| {
@ -426,7 +426,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::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 {

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

@ -2,7 +2,7 @@ mod cache;
mod topology;
pub mod connection {
pub use crate::cache::{ConnectionRef, Connections, ttl_seconds};
pub use crate::cache::{ConnectionRef, Connections};
}
pub use cache::{

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

@ -58,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na
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
Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. 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

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

@ -1,22 +1,21 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error,
ExactCacheContext, FlushCache,
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache,
};
use serde_json::Value;
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
#[derive(Clone)]
pub struct ResponseCacheRequest<C: CacheContext = ExactCacheContext> {
pub struct ResponseCacheRequest<C: CacheContext = litellm_cache::ExactCacheContext> {
pub key: CacheKeyInput,
pub controls: CacheControls,
pub context: C,
pub max_age: Option<Duration>,
}
impl ResponseCacheRequest<ExactCacheContext> {
impl<C: CacheContext + Default> ResponseCacheRequest<C> {
pub fn new(key: CacheKeyInput) -> Self {
Self {
key,
@ -27,17 +26,24 @@ impl ResponseCacheRequest<ExactCacheContext> {
default_on: true,
..Default::default()
},
context: ExactCacheContext::default(),
context: C::default(),
max_age: None,
}
}
}
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>> {
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
where
B::Context: Default + PartialEq,
{
backend: Arc<B>,
}
impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
impl<B> ResponseCache<B>
where
B: BaseCache<Value = CacheEntry>,
B::Context: Default + PartialEq,
{
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
}
@ -46,10 +52,11 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration>
where
B::Context: Default,
{
pub fn backend_arc(&self) -> &Arc<B> {
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration> {
self.backend.get_ttl(&B::Context::default())
}
@ -199,10 +206,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
&self,
entries: Vec<(ResponseCacheRequest<B::Context>, Value)>,
now: Duration,
) -> Result<(), Error>
where
B::Context: PartialEq,
{
) -> Result<(), Error> {
self.async_store_entries(
entries
.into_iter()
@ -217,10 +221,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
pub async fn async_store_entries(
&self,
entries: Vec<(ResponseCacheRequest<B::Context>, Value, Duration)>,
) -> Result<(), Error>
where
B::Context: PartialEq,
{
) -> Result<(), Error> {
let writable = entries
.into_iter()
.filter(|(request, _, _)| request.controls.writes())

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

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-s3"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-auth-aws.workspace = true
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
aws-credential-types = "1.3.0"
aws-smithy-types = "1.6.0"
aws-types = "1.6.0"
tokio.workspace = true
[dev-dependencies]
wiremock = "0.6.5"
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

101
litellm-rust/crates/cache-s3/src/auth.rs Normal file
View file

@ -0,0 +1,101 @@
use aws_credential_types::{
Credentials as AwsCredentials,
provider::{ProvideCredentials, error::CredentialsError, future},
};
use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
#[derive(Clone)]
pub(crate) struct Credentials {
config: AwsAuthConfig,
env: fn(&str) -> Option<String>,
}
impl Credentials {
pub(crate) fn new(config: AwsAuthConfig) -> Self {
Self::with_env(config, |name| std::env::var(name).ok())
}
pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
Self { config, env }
}
}
impl ProvideCredentials for Credentials {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::new(async {
if let (Some(access_key_id), Some(secret_access_key)) = (
self.config.access_key_id.clone(),
self.config.secret_access_key.clone(),
) {
return Ok(AwsCredentials::new(
access_key_id,
secret_access_key,
self.config.session_token.clone(),
None,
"litellm-s3-cache",
));
}
resolve_credentials(self.config.clone(), &self.env)
.await
.map_err(|_| CredentialsError::provider_error("S3 cache authentication failed"))
})
}
}
impl std::fmt::Debug for Credentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Credentials").finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn explicit_keys_ignore_an_ambient_session_token() {
let provider = Credentials::with_env(
AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
|name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()),
);
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "key");
assert_eq!(credentials.secret_access_key(), "secret");
assert_eq!(credentials.session_token(), None);
}
#[tokio::test]
async fn explicit_keys_keep_their_session_token() {
let provider = Credentials::new(AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
session_token: Some("t".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.session_token(), Some("t"));
}
#[tokio::test]
async fn environment_keys_resolve_with_their_session_token() {
let provider = Credentials::with_env(AwsAuthConfig::default(), |name| match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "env-key");
assert_eq!(credentials.secret_access_key(), "env-secret");
assert_eq!(credentials.session_token(), Some("env-token"));
}
}

220
litellm-rust/crates/cache-s3/src/cache.rs Normal file
View file

@ -0,0 +1,220 @@
use std::{
future::Future,
sync::Arc,
time::{Duration, SystemTime},
};
use aws_sdk_s3::{
config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
error::SdkError,
primitives::ByteStream,
};
use aws_smithy_types::{DateTime, date_time::Format};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use crate::auth::Credentials;
pub struct S3Endpoint {
pub url: String,
}
pub struct S3CacheConfig {
pub bucket: String,
pub key_prefix: String,
pub region: String,
pub endpoint: Option<S3Endpoint>,
pub auth: AwsAuthConfig,
}
pub struct S3Cache<C: CacheCodec> {
client: aws_sdk_s3::Client,
codec: C,
runtime: Handle,
bucket: Arc<str>,
key_prefix: Arc<str>,
region: Arc<str>,
endpoint: Option<Arc<str>>,
}
impl<C: CacheCodec> S3Cache<C> {
pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self {
let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
let base = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(config.region.clone()))
.credentials_provider(Credentials::new(config.auth))
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
let builder = match &endpoint_url {
Some(url) => base.endpoint_url(url).force_path_style(true),
None => base,
};
Self {
client: aws_sdk_s3::Client::from_conf(builder.build()),
codec,
runtime,
bucket: config.bucket.into(),
key_prefix: config.key_prefix.into(),
region: config.region.into(),
endpoint: endpoint_url.map(Into::into),
}
}
pub fn bucket(&self) -> &str {
&self.bucket
}
pub fn key_prefix(&self) -> &str {
&self.key_prefix
}
pub fn region(&self) -> &str {
&self.region
}
pub fn endpoint(&self) -> Option<&str> {
self.endpoint.as_deref()
}
pub fn to_s3_key(&self, key: &str) -> String {
format!("{}{}", self.key_prefix, key.replace(':', "/"))
}
fn block_on<F: Future>(&self, future: F) -> F::Output {
if Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
async fn put(
&self,
key: &str,
value: C::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let s3_key = self.to_s3_key(key);
let body = self.codec.encode(&value)?;
let request = self
.client
.put_object()
.bucket(self.bucket.as_ref())
.key(&s3_key)
.body(ByteStream::from(body))
.content_type("application/json")
.content_language("en")
.content_disposition(format!("inline; filename=\"{s3_key}.json\""));
let request = match context.ttl {
Some(ttl) => {
let seconds = ttl.as_secs_f64();
request
.cache_control(format!("immutable, max-age={seconds}, s-maxage={seconds}"))
.expires(DateTime::from(SystemTime::now() + ttl))
}
None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"),
};
request.send().await.map_err(|_| Error::Unavailable)?;
Ok(())
}
async fn get(&self, key: &str) -> Result<Option<C::Value>, Error> {
let output = match self
.client
.get_object()
.bucket(self.bucket.as_ref())
.key(self.to_s3_key(key))
.send()
.await
{
Ok(output) => output,
Err(error) => {
if let SdkError::ServiceError(service) = &error {
let status = error
.raw_response()
.map(|response| response.status().as_u16());
let not_found = service.err().is_no_such_key()
|| service.err().meta().code() == Some("AccessDenied")
|| status == Some(404)
|| status == Some(403);
if not_found {
return Ok(None);
}
}
return Err(Error::Unavailable);
}
};
if let Some(expires) = output.expires_string()
&& let Ok(expires) = DateTime::from_str(expires, Format::HttpDate)
&& expires < DateTime::from(SystemTime::now())
{
return Ok(None);
}
let bytes = output
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?
.into_bytes();
self.codec.decode(&bytes).map(Some)
}
}
impl<C: CacheCodec> BaseCache for S3Cache<C> {
type Value = C::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
self.block_on(self.put(key, value, context))
}
fn get_cache(&self, key: &str, _context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.block_on(self.get(key))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
self.put(key, value, &context).await
}
async fn async_get_cache(
&self,
key: &str,
_context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.get(key).await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<C: CacheCodec> BatchCache for S3Cache<C> {}
impl<C: CacheCodec> FlushCache for S3Cache<C> {
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
}

4
litellm-rust/crates/cache-s3/src/lib.rs Normal file
View file

@ -0,0 +1,4 @@
mod auth;
mod cache;
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

278
litellm-rust/crates/cache-s3/tests/cache.rs Normal file
View file

@ -0,0 +1,278 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
use serde_json::{Value, json};
use tokio::runtime::Handle;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{method, path},
};
fn config(endpoint: String) -> S3CacheConfig {
S3CacheConfig {
bucket: "cache-bucket".to_string(),
key_prefix: "team/".to_string(),
region: "us-east-1".to_string(),
endpoint: Some(S3Endpoint { url: endpoint }),
auth: AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
}
}
fn cache(endpoint: &str) -> S3Cache<JsonCodec<Value>> {
S3Cache::new(
config(endpoint.to_string()),
JsonCodec::<Value>::new(),
Handle::current(),
)
}
async fn mock_server() -> MockServer {
let server = MockServer::start().await;
Mock::given(method("PUT"))
.respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\""))
.mount(&server)
.await;
server
}
fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option<SystemTime> {
use aws_smithy_types::{DateTime, date_time::Format};
headers
.get(name)
.and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok())
.map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos()))
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn set_writes_python_metadata_with_and_without_ttl() {
let server = mock_server().await;
let cache = cache(&server.uri());
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(90)),
};
cache
.set_cache("alpha:beta", json!({"answer": 1}), &context)
.unwrap();
cache
.set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default())
.unwrap();
let requests = server.received_requests().await.unwrap();
let ttl_request = requests
.iter()
.find(|request| request.url.path() == "/cache-bucket/team/alpha/beta")
.expect("ttl write should hit the converted S3 key");
assert_eq!(
ttl_request.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=90, s-maxage=90"
);
assert_eq!(
ttl_request.headers["content-type"].to_str().unwrap(),
"application/json"
);
assert_eq!(
ttl_request.headers["content-language"].to_str().unwrap(),
"en"
);
assert_eq!(
ttl_request.headers["content-disposition"].to_str().unwrap(),
"inline; filename=\"team/alpha/beta.json\""
);
let expires = http_date_from(&ttl_request.headers, "expires").expect("ttl write sets Expires");
let remaining = expires.duration_since(SystemTime::now()).unwrap();
assert!(remaining > Duration::from_secs(60) && remaining <= Duration::from_secs(91));
assert_eq!(
serde_json::from_slice::<Value>(&ttl_request.body).unwrap(),
json!({"answer": 1})
);
let plain = requests
.iter()
.find(|request| request.url.path() == "/cache-bucket/team/plain")
.expect("no-ttl write should hit the converted S3 key");
assert_eq!(
plain.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=31536000, s-maxage=31536000"
);
assert!(plain.headers.get("expires").is_none());
assert_eq!(
plain.headers["content-disposition"].to_str().unwrap(),
"inline; filename=\"team/plain.json\""
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_hit_miss_expired_and_invalid_entries() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/missing"))
.respond_with(
ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/denied"))
.respond_with(
ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/expired"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
.set_body_json(json!({"answer": 4})),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/malformed"))
.respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry"))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
assert_eq!(
cache.get_cache("hit", &context).unwrap(),
Some(json!({"answer": 3}))
);
assert_eq!(cache.get_cache("missing", &context).unwrap(), None);
assert_eq!(cache.get_cache("denied", &context).unwrap(), None);
assert_eq!(cache.get_cache("expired", &context).unwrap(), None);
assert_eq!(
cache.get_cache("malformed", &context),
Err(Error::InvalidEntry)
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
let server = mock_server().await;
for (key, status, body) in [
("first", 200, "{\"answer\": 1}"),
("invalid", 200, "garbage"),
] {
Mock::given(method("GET"))
.and(path(format!("/cache-bucket/team/{key}")))
.respond_with(ResponseTemplate::new(status).set_body_string(body))
.mount(&server)
.await;
}
Mock::given(method("GET"))
.and(path("/cache-bucket/team/miss"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
let keys = vec![
"first".to_string(),
"miss".to_string(),
"invalid".to_string(),
];
let entries = cache.batch_get_cache(&keys, &context).unwrap();
assert_eq!(
entries,
vec![
BatchEntry::Hit(json!({"answer": 1})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unsupported_and_noop_capabilities_match_python() {
let server = mock_server().await;
let cache = cache(&server.uri());
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
cache.flush_cache().unwrap();
cache.disconnect().await.unwrap();
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext {
ttl: Some(Duration::from_secs(45)),
}),
Some(Duration::from_secs(45))
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[test]
fn key_conversion_prefixes_and_splits_colons() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.unwrap();
let _guard = runtime.enter();
let cache = S3Cache::new(
S3CacheConfig {
key_prefix: "team/".to_string(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(cache.bucket(), "cache-bucket");
assert_eq!(cache.key_prefix(), "team/");
assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c");
assert_eq!(cache.to_s3_key("plain"), "team/plain");
let unprefixed = S3Cache::new(
S3CacheConfig {
key_prefix: String::new(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(unprefixed.to_s3_key("a:b"), "a/b");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sync_methods_block_inside_and_outside_the_runtime() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9})))
.mount(&server)
.await;
let uri = server.uri();
let cache = tokio::task::spawn_blocking(move || {
let cache = cache(&uri);
let context = ExactCacheContext::default();
cache
.set_cache("key", json!({"answer": 9}), &context)
.unwrap();
cache.get_cache("key", &context).unwrap()
})
.await
.unwrap();
assert_eq!(cache, Some(json!({"answer": 9})));
}

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

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
serde_json.workspace = true
sha2.workspace = true
tokio.workspace = true
uuid = { version = "1", features = ["v4"] }
[dev-dependencies]
redis-test = "1.0.4"
rstest.workspace = true

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

File diff suppressed because it is too large Load diff

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

@ -35,8 +35,8 @@ impl CacheContext for ExactCacheContext {
#[derive(Clone, Debug, Default, PartialEq)]
pub struct SemanticCacheContext {
pub input: Option<serde_json::Value>,
pub messages: Vec<serde_json::Value>,
pub metadata: serde_json::Map<String, serde_json::Value>,
pub messages: Option<serde_json::Value>,
pub metadata: Option<serde_json::Value>,
pub scope: Option<String>,
pub ttl: Option<Duration>,
}
@ -48,11 +48,8 @@ impl CacheContext for SemanticCacheContext {
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
Self {
input: self.input.clone(),
messages: self.messages.clone(),
metadata: self.metadata.clone(),
scope: self.scope.clone(),
ttl,
..self.clone()
}
}
}
@ -130,3 +127,31 @@ pub trait BaseCache: Send + Sync {
fn test_connection(&self) -> impl Future<Output = Result<CacheConnectionResult, Error>> + Send;
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use serde_json::json;
use super::{CacheContext, SemanticCacheContext};
#[test]
fn semantic_context_with_ttl_only_replaces_ttl() {
let context = SemanticCacheContext {
input: Some(json!({"input": "hello"})),
messages: Some(json!([{"role": "user", "content": "hello"}])),
metadata: Some(json!({"tenant": "team"})),
scope: Some("scope".into()),
ttl: Some(Duration::from_secs(10)),
};
let updated = context.with_ttl(Some(Duration::from_secs(20)));
assert_eq!(updated.ttl, Some(Duration::from_secs(20)));
assert_eq!(updated.input, context.input);
assert_eq!(updated.messages, context.messages);
assert_eq!(updated.metadata, context.metadata);
assert_eq!(updated.scope, context.scope);
}
}

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

@ -131,8 +131,8 @@ fn associated_context_preserves_backend_specific_lookup_inputs() {
fn semantic_context_with_ttl_preserves_lookup_inputs() {
let context = SemanticCacheContext {
input: Some(serde_json::json!("text")),
messages: vec![serde_json::json!({"role": "user", "content": "hi"})],
metadata: serde_json::Map::from_iter([("key".into(), serde_json::json!("value"))]),
messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])),
metadata: Some(serde_json::json!({"key": "value"})),
scope: Some("scope".into()),
ttl: None,
};

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

@ -11,7 +11,7 @@ aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
sse = ["dep:sse-stream"]
[dependencies]
aws-smithy-eventstream = { version = "=0.61.1", optional = true }
aws-smithy-eventstream = { version = "=0.61.4", optional = true }
aws-smithy-types = { version = "1.6.1", optional = true }
bytes = "1"
futures-util.workspace = true

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

@ -34,7 +34,7 @@ tokio = { workspace = true, features = ["sync"] }
url.workspace = true
[dev-dependencies]
aws-smithy-eventstream = "=0.61.1"
aws-smithy-eventstream = "=0.61.4"
aws-smithy-types = "1.6.1"
rstest.workspace = true
tokio.workspace = true

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

@ -24,12 +24,15 @@ litellm-cache.workspace = true
litellm-cache-azure-blob.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-s3.workspace = true
litellm-cache-gcs.workspace = true
litellm-cache-disk.workspace = true
litellm-cache-redis-semantic.workspace = true
litellm-cache-response.workspace = true
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
serde.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-callbacks-legacy-python.workspace = true
litellm-core.workspace = true
litellm-core-utils.workspace = true
@ -41,6 +44,7 @@ litellm-host-python.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
serde_json.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }
@ -50,6 +54,7 @@ serde_with.workspace = true
criterion.workspace = true
futures-util.workspace = true
rstest.workspace = true
sha2.workspace = true
tokio-tungstenite.workspace = true
[[bench]]

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

@ -14,7 +14,6 @@ use super::{
future::{ready_none, ready_value},
native::NativeResponseCache,
request::{now, request, requests},
semantic::{SemanticOperation, drive},
};
pub(super) enum CacheBinding {
@ -57,17 +56,7 @@ impl ResolvedCache {
CacheBinding::Disabled => ready_none(py)?,
CacheBinding::Native(service) => {
let request = request(input)?;
if service.semantic_embedder().is_some() {
return Ok(ExecutionStep::Await(
drive(py, service.clone(), SemanticOperation::Lookup(request))?.unbind(),
));
}
let service = service.clone();
run_async(
py,
async move { service.async_lookup(&request, now()).await },
cache_error,
)?
service.async_lookup_py(py, request)?
}
CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?,
};
@ -185,19 +174,7 @@ impl ResolvedCache {
CacheBinding::Native(service) => {
let request = self::request(request)?;
let response: Value = from_py(response)?;
if service.semantic_embedder().is_some() {
return drive(
py,
service.clone(),
SemanticOperation::Store(request, response),
);
}
let service = service.clone();
run_async(
py,
async move { service.async_store(&request, response, now()).await },
cache_error,
)
service.async_store_py(py, request, response)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store(py, response, callback_kwargs)
@ -253,20 +230,8 @@ impl ResolvedCache {
"batch cache requests and responses must have equal lengths",
));
}
let entries = requests.into_iter().zip(responses).collect::<Vec<_>>();
if service.semantic_embedder().is_some() {
return drive(
py,
service.clone(),
SemanticOperation::StoreBatch(entries.into()),
);
}
let service = service.clone();
run_async(
py,
async move { service.async_store_batch(entries, now()).await },
cache_error,
)
let entries = requests.into_iter().zip(responses).collect();
service.async_store_batch_py(py, entries)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store_batch(py, callback_result, callback_kwargs)

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

@ -1,11 +1,13 @@
use std::{path::PathBuf, time::Duration};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::CacheType;
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
exceptions::{PyAttributeError, PyTypeError, PyValueError},
prelude::*,
types::{PyAny, PyDict, PyList, PyString},
types::{PyAny, PyBool, PyDict, PyList, PyString},
};
use super::{native::NativeResponseCache, request::duration};
@ -115,10 +117,20 @@ struct RedisClientProjection<'py> {
const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct ValkeySemanticCacheConfig {
pub(super) similarity_threshold: f64,
pub(super) index_name: String,
pub(super) embedding_model: String,
pub(super) connection: RedisConnectionConfig,
}
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
S3(Box<S3CacheConfig>),
Gcs(GcsCacheConfig),
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
RedisSemantic(Box<RedisSemanticCacheConfig>),
@ -136,6 +148,9 @@ pub(super) enum UnsupportedCacheConfig {
RedisCredentials,
RedisConnection,
RedisOption,
S3Client,
S3Credentials,
S3Option,
GcsBucket,
DiskStore,
}
@ -148,6 +163,9 @@ impl UnsupportedCacheConfig {
Self::RedisCredentials => "native Redis credentials require Python",
Self::RedisConnection => "native Redis connection type is not implemented",
Self::RedisOption => "native Redis configuration requires Python",
Self::S3Client => "native S3 client type is not implemented",
Self::S3Credentials => "native S3 credentials require Python",
Self::S3Option => "native S3 configuration requires Python",
Self::GcsBucket => "native GCS cache requires a configured bucket name",
Self::DiskStore => "native disk cache requires the built-in diskcache store",
}
@ -192,6 +210,13 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::S3) => match project_s3(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::S3(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Gcs) => match project_gcs(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
@ -199,6 +224,13 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::ValkeySemantic) => match project_valkey_semantic(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::ValkeySemantic(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Disk) => match project_disk(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
@ -218,11 +250,9 @@ impl NativeCacheConfig {
backend: CacheBackendConfig::RedisSemantic(Box::new(backend)),
}))
}),
Some(CacheType::ValkeySemantic | CacheType::S3 | CacheType::QdrantSemantic) | None => {
Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
))
}
Some(CacheType::QdrantSemantic) | None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
)),
}
}
@ -230,12 +260,16 @@ impl NativeCacheConfig {
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO),
CacheBackendConfig::Disk(_)
| CacheBackendConfig::AzureBlob(_)
| CacheBackendConfig::Gcs(_)
| CacheBackendConfig::RedisSemantic(_) => None,
};
if service.default_ttl() != default_ttl {
if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_))
&& service.default_ttl() != default_ttl
{
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
@ -260,6 +294,30 @@ impl NativeCacheConfig {
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
CacheBackendConfig::S3(_) if service.kind() != "s3" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => {
Some("facade and native backend buckets must match")
}
CacheBackendConfig::S3(config)
if service.key_prefix() != Some(config.key_prefix.as_str()) =>
{
Some("facade and native backend key prefixes must match")
}
CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => {
Some("facade and native backend regions must match")
}
CacheBackendConfig::S3(config)
if service.endpoint()
!= config
.endpoint
.as_ref()
.map(|endpoint| endpoint.url.as_str()) =>
{
Some("facade and native backend endpoints must match")
}
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => {
Some("facade and native backend types must match")
}
@ -285,6 +343,16 @@ impl NativeCacheConfig {
Some("facade and native backend credentials must match")
}
CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::ValkeySemantic(config) => {
if service.kind() != "valkey-semantic" {
return Some("facade and native backend types must match");
}
let Some((threshold, index_name)) = service.semantic_config() else {
return Some("facade and native backend types must match");
};
(threshold != config.similarity_threshold || index_name != config.index_name)
.then_some("facade and native semantic settings must match")
}
CacheBackendConfig::Disk(_) if service.kind() != "disk" => {
Some("facade and native backend types must match")
}
@ -497,6 +565,77 @@ fn project_redis(
}))
}
#[inline(never)]
fn project_s3(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<S3CacheConfig, UnsupportedCacheConfig>> {
let client = backend.getattr("s3_client")?;
if !instance_class_is(&client, "botocore.client", "S3")? {
return Ok(Err(UnsupportedCacheConfig::S3Client));
}
let meta = client.getattr("meta")?;
let Some(region) = optional_string(meta.getattr("region_name")?)? else {
return Ok(Err(UnsupportedCacheConfig::S3Option));
};
let Some(endpoint_url) = optional_string(meta.getattr("endpoint_url")?)? else {
return Ok(Err(UnsupportedCacheConfig::S3Option));
};
let client_config = meta.getattr("config")?;
for name in ["s3", "proxies", "client_cert"] {
if optional_attribute(&client_config, name)?.is_some_and(|value| !value.is_none()) {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
}
let signature = match optional_attribute(&client_config, "signature_version")? {
Some(value) => value.extract::<Option<String>>()?,
None => None,
};
if signature.as_deref() != Some("s3v4") {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
let insecure = endpoint_url.starts_with("http://");
let verify = optional_attribute_chain(&client, &["_endpoint", "http_session", "_verify"])?;
let verified = verify
.and_then(|value| value.cast::<PyBool>().ok().map(|value| value.is_true()))
.unwrap_or(false);
if !verified && !insecure {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
let credentials = optional_attribute_chain(&client, &["_request_signer", "_credentials"])?
.ok_or(UnsupportedCacheConfig::S3Credentials);
let credentials = match credentials {
Ok(credentials) if !credentials.is_none() => credentials,
_ => return Ok(Err(UnsupportedCacheConfig::S3Credentials)),
};
let auth = if credentials.getattr("method")?.extract::<String>()?.as_str() == "explicit" {
AwsAuthConfig {
access_key_id: credentials
.getattr("access_key")?
.extract::<Option<String>>()?,
secret_access_key: credentials
.getattr("secret_key")?
.extract::<Option<String>>()?,
session_token: credentials.getattr("token")?.extract::<Option<String>>()?,
region_name: Some(region.clone()),
..Default::default()
}
} else {
AwsAuthConfig {
region_name: Some(region.clone()),
..Default::default()
}
};
let default_endpoint = endpoint_url == format!("https://s3.{region}.amazonaws.com")
|| (region == "us-east-1" && endpoint_url == "https://s3.amazonaws.com");
Ok(Ok(S3CacheConfig {
bucket: backend.getattr("bucket_name")?.extract::<String>()?,
key_prefix: backend.getattr("key_prefix")?.extract::<String>()?,
region,
endpoint: (!default_endpoint).then_some(S3Endpoint { url: endpoint_url }),
auth,
}))
}
#[inline(never)]
fn project_standalone_client<'py>(
client: &Bound<'py, PyAny>,
@ -531,7 +670,7 @@ fn project_standalone_client<'py>(
#[inline(never)]
fn project_cluster_client<'py>(
source: &Bound<'py, PyDict>,
source: &Bound<'_, PyDict>,
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
let Some(startup_nodes) = startup_nodes(source)? else {
@ -619,6 +758,71 @@ fn port(value: i64) -> PyResult<u16> {
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
}
#[inline(never)]
fn project_valkey_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<ValkeySemanticCacheConfig, UnsupportedCacheConfig>> {
let client = backend.getattr("sync_client")?;
let pool = client.getattr("connection_pool")?;
let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
for key in ["credential_provider", "redis_connect_func"] {
if has_value(&resolved, key)? {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
}
if is_tls {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let connection = RedisConnectionConfig {
host: required_string(&resolved, "host")?,
port: u16::try_from(required_i64(&resolved, "port")?)
.map_err(|_| PyValueError::new_err("invalid Redis port"))?,
database: optional_i64(&resolved, "db")?.unwrap_or(0),
username: optional_dict_string(&resolved, "username")?,
password: optional_dict_string(&resolved, "password")?,
protocol: RedisProtocol::Resp2,
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
read_timeout: None,
connect_timeout: None,
socket_keepalive: None,
health_check_interval: Duration::ZERO,
client_name: None,
tls: None,
};
if connection.host.is_empty() {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
Ok(Ok(ValkeySemanticCacheConfig {
similarity_threshold: backend.getattr("similarity_threshold")?.extract()?,
index_name: backend.getattr("index_name")?.extract()?,
embedding_model: backend.getattr("embedding_model")?.extract()?,
connection,
}))
}
#[inline(never)]
fn project_connection_pool<'py>(
pool: &Bound<'py, PyAny>,
) -> PyResult<Result<(Bound<'py, PyDict>, bool), UnsupportedCacheConfig>> {
if !instance_class_is(pool, "redis.connection", "ConnectionPool")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
let connection_class = resolved
.get_item("connection_class")?
.unwrap_or(pool.getattr("connection_class")?);
let is_tls = if class_is(&connection_class, "redis.connection", "Connection")? {
false
} else if class_is(&connection_class, "redis.connection", "SSLConnection")? {
true
} else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
Ok(Ok((resolved, is_tls)))
}
#[inline(never)]
fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
Ok(RedisTlsConfig {
@ -700,6 +904,31 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult<O
}
}
#[inline(never)]
fn optional_attribute<'py>(
value: &Bound<'py, PyAny>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match value.getattr(name) {
Ok(value) => Ok(Some(value)),
Err(error) if error.is_instance_of::<PyAttributeError>(value.py()) => Ok(None),
Err(error) => Err(error),
}
}
#[inline(never)]
fn optional_attribute_chain<'py>(
value: &Bound<'py, PyAny>,
names: &[&str],
) -> PyResult<Option<Bound<'py, PyAny>>> {
names
.iter()
.try_fold(Some(value.clone()), |current, name| match current {
Some(current) => optional_attribute(&current, name),
None => Ok(None),
})
}
#[inline(never)]
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
Ok(value
@ -790,8 +1019,8 @@ mod tests {
use litellm_cache_redis::{RedisNode, RedisTopology};
use super::{
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig,
NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
};
use crate::cache::native::NativeResponseCache;
@ -866,6 +1095,194 @@ mod tests {
});
}
#[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 projects_valkey_semantic_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = Connection\n\
pool.max_connections = 12\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Valkey semantic cache should be supported");
};
let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else {
panic!("expected Valkey semantic configuration");
};
assert_eq!(valkey.similarity_threshold, 0.85);
assert_eq!(valkey.index_name, "semantic_idx");
assert_eq!(valkey.embedding_model, "text-embedding-3-small");
assert_eq!(valkey.connection.host, "cache.internal");
assert_eq!(valkey.connection.port, 6390);
assert_eq!(valkey.connection.database, 2);
assert_eq!(valkey.connection.pool_size, 12);
assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2);
assert!(valkey.connection.tls.is_none());
});
}
#[test]
fn valkey_semantic_tls_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("TLS Valkey semantic cache should stay on Python");
};
assert_eq!(
reason.message(),
"native Redis connection type is not implemented"
);
});
}
#[test]
fn valkey_semantic_dynamic_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = Connection\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'credential_provider': object()}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic Valkey authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[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 projects_gcs_configuration() {
Python::initialize();
@ -931,192 +1348,6 @@ mod tests {
});
}
#[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_builtin_disk_configuration_and_rejects_custom_stores() {
Python::initialize();
Python::attach(|py| {
let root =
std::env::temp_dir().join(format!("litellm-disk-config-{}", std::process::id()));
let directory = root.to_string_lossy();
let disk_facade = facade(
py,
&format!(
"Cache = type('Cache', (), {{'__module__': 'diskcache.core'}})\n\
Disk = type('Disk', (), {{'__module__': 'diskcache.core'}})\n\
store = Cache()\n\
store._disk = Disk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&disk_facade).unwrap()
else {
panic!("disk cache should be supported");
};
let CacheBackendConfig::Disk(disk) = config.backend else {
panic!("expected disk configuration");
};
assert_eq!(disk.directory, root);
let matching = NativeResponseCache::disk(&directory).unwrap();
assert_eq!(
(NativeCacheConfig {
policy: config.policy,
backend: CacheBackendConfig::Disk(disk),
})
.service_mismatch(&matching),
None
);
let other = NativeResponseCache::disk(&root.join("other").to_string_lossy()).unwrap();
let mismatch = NativeCacheConfig {
policy: CachePolicy {
mode: "default-on".into(),
ttl: None,
namespace: None,
supported_call_types: None,
redis_flush_size: None,
semantic_cache_scope: "key".into(),
},
backend: CacheBackendConfig::Disk(DiskCacheConfig {
directory: root.clone(),
}),
};
assert_eq!(
mismatch.service_mismatch(&other),
Some("facade and native backend directories must match")
);
let custom = facade(
py,
&format!(
"CustomCache = type('CustomCache', (), {{'__module__': 'mypkg'}})\n\
CustomDisk = type('CustomDisk', (), {{'__module__': 'mypkg'}})\n\
store = CustomCache()\n\
store._disk = CustomDisk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&custom).unwrap()
else {
panic!("custom disk store must stay on Python");
};
assert_eq!(
reason.message(),
"native disk cache requires the built-in diskcache store"
);
});
}
#[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();

96
litellm-rust/crates/python-bridge/src/cache/embedder.rs vendored
View file

@ -1,10 +1,9 @@
use std::future::Future;
use litellm_cache::Error;
use litellm_cache_redis_semantic::Embedder;
use litellm_host_python::to_py;
use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict};
use serde_json::{Map, Value};
use serde_json::Value;
tokio::task_local! {
static PREPARED_EMBEDDING: Result<Vec<f32>, Error>;
@ -19,11 +18,21 @@ pub(super) fn with_prepared_embedding<F: Future>(
pub(super) struct PythonEmbedder(Py<PyAny>);
impl Clone for PythonEmbedder {
fn clone(&self) -> Self {
Python::attach(|py| Self(self.0.clone_ref(py)))
}
}
impl PythonEmbedder {
pub(super) fn new(object: Py<PyAny>) -> Self {
Self(object)
}
pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self(backend.clone().unbind()))
}
pub(super) fn object(&self) -> &Py<PyAny> {
&self.0
}
@ -32,16 +41,24 @@ impl PythonEmbedder {
visit.call(&self.0)
}
pub(super) fn async_embed_awaitable<'py>(
&self,
py: Python<'py>,
prompt: &str,
metadata: &Option<Value>,
) -> PyResult<Bound<'py, PyAny>> {
let metadata = to_py(py, metadata)?;
self.0
.bind(py)
.call_method1("_get_async_embedding", (prompt, metadata))
}
fn metadata_kwargs<'py>(
py: Python<'py>,
metadata: &Map<String, Value>,
metadata: Option<&Value>,
) -> PyResult<Bound<'py, PyDict>> {
let kwargs = PyDict::new(py);
if metadata.is_empty() {
kwargs.set_item("metadata", py.None())?;
} else {
kwargs.set_item("metadata", to_py(py, metadata)?)?;
}
kwargs.set_item("metadata", to_py(py, &metadata)?)?;
Ok(kwargs)
}
@ -49,7 +66,7 @@ impl PythonEmbedder {
&self,
py: Python<'_>,
prompt: &str,
metadata: &Map<String, Value>,
metadata: Option<&Value>,
) -> PyResult<Py<PyAny>> {
let kwargs = Self::metadata_kwargs(py, metadata)?;
self.0
@ -67,8 +84,33 @@ impl PythonEmbedder {
}
}
impl Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: &Map<String, Value>) -> Result<Vec<f32>, Error> {
impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
let result = Python::attach(|py| -> PyResult<Vec<f64>> {
let metadata = to_py(py, &metadata)?;
self.0
.bind(py)
.call_method1("_get_embedding", (prompt, metadata))?
.extract()
})
.map_err(|_| Error::Unavailable)?;
Ok(result.into_iter().map(|value| value as f32).collect())
}
fn async_embed(
&self,
_prompt: &str,
_metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
let seeded = PREPARED_EMBEDDING
.try_with(Clone::clone)
.unwrap_or(Err(Error::Unavailable));
std::future::ready(seeded)
}
}
impl litellm_cache_redis_semantic::Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
Python::attach(|py| {
let kwargs = Self::metadata_kwargs(py, metadata)?;
Self::extract(self.0.bind(py).call_method(
@ -83,7 +125,7 @@ impl Embedder for PythonEmbedder {
fn async_embed(
&self,
_prompt: &str,
_metadata: &Map<String, Value>,
_metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
let seeded = PREPARED_EMBEDDING
.try_with(Clone::clone)
@ -98,21 +140,19 @@ mod tests {
#[tokio::test]
async fn async_embed_returns_the_seeded_vector_or_unavailable() {
Python::initialize();
let embedder = Python::attach(|py| PythonEmbedder::new(py.None()));
let metadata = Map::new();
let embedder_ref = &embedder;
let metadata_ref = &metadata;
assert_eq!(
with_prepared_embedding(Ok(vec![0.5f32, 0.25]), async move {
embedder_ref.async_embed("prompt", metadata_ref).await
})
.await,
Ok(vec![0.5, 0.25])
);
assert_eq!(
embedder.async_embed("prompt", &metadata).await,
Err(Error::Unavailable)
);
let object = Python::attach(|py| {
Python::initialize();
py.None()
});
let embedder = PythonEmbedder::new(object);
let scoped_embedder = embedder.clone();
let scoped = with_prepared_embedding(Ok(vec![0.25]), async move {
litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None)
.await
});
assert_eq!(scoped.await, Ok(vec![0.25]));
let unscoped =
litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await;
assert_eq!(unscoped, Err(Error::Unavailable));
}
}

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

@ -31,6 +31,7 @@ struct RedisPoolGuard {
connection_class: Py<PyAny>,
connection_kwargs: Py<PyAny>,
max_connections: Option<usize>,
client_name: &'static str,
attributes: RedisPoolAttributes,
}
@ -46,11 +47,18 @@ struct AzureBlobClientGuard {
container_name: String,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
enum ConnectionGuard {
None,
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
S3(S3ClientGuard),
}
#[derive(Clone, Copy)]
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
@ -69,6 +77,8 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
max_connections: None,
};
const VALKEY_POOL: RedisPoolAttributes = STANDALONE_POOL;
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
@ -155,10 +165,9 @@ impl ObjectGuard {
return Ok(false);
}
for (name, value) in &expected.attributes {
if !attributes.get_item(name)?.is(value.bind(py)) {
return Ok(false);
}
if instance.contains(name)? && !self.config_names.contains(&name.as_str()) {
if (instance.contains(name)? && !self.config_names.contains(&name.as_str()))
|| !attributes.get_item(name)?.is(value.bind(py))
{
return Ok(false);
}
}
@ -179,8 +188,12 @@ impl ObjectGuard {
}
impl RedisPoolGuard {
fn capture(backend: &Bound<'_, PyAny>, attributes: RedisPoolAttributes) -> PyResult<Self> {
let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?;
fn capture(
backend: &Bound<'_, PyAny>,
client_name: &'static str,
attributes: RedisPoolAttributes,
) -> PyResult<Self> {
let pool = backend.getattr(client_name)?.getattr(attributes.pool)?;
Ok(Self {
reference: pool.clone().unbind(),
connection_class: pool.getattr(attributes.connection_class)?.unbind(),
@ -188,31 +201,30 @@ impl RedisPoolGuard {
.getattr("connection_kwargs")?
.call_method0("copy")?
.unbind(),
max_connections: Self::max_connections(&pool, &attributes)?,
max_connections: attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()?,
client_name,
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.client_name)?
.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.max_connections
== self
.attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()?
&& self
.connection_kwargs
.bind(py)
@ -274,12 +286,43 @@ impl AzureBlobClientGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl ConnectionGuard {
fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(match (kind, cluster) {
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?),
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"redis_client",
STANDALONE_POOL,
)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"redis_client",
CLUSTER_POOL,
)?),
("valkey-semantic", _) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"sync_client",
VALKEY_POOL,
)?),
("disk", _) => Self::None,
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
("s3", _) => Self::S3(S3ClientGuard::capture(backend)?),
_ => Self::None,
})
}
@ -289,6 +332,7 @@ impl ConnectionGuard {
Self::None => Ok(true),
Self::RedisPool(guard) => guard.matches(py, backend),
Self::AzureBlob(guard) => guard.matches(py, backend),
Self::S3(guard) => guard.matches(py, backend),
}
}
@ -297,6 +341,7 @@ impl ConnectionGuard {
Self::None => Ok(()),
Self::RedisPool(guard) => guard.traverse(visit),
Self::AzureBlob(guard) => guard.traverse(visit),
Self::S3(guard) => guard.traverse(visit),
}
}
}
@ -329,12 +374,18 @@ impl FacadeGuard {
"redis",
),
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
("valkey-semantic", false) => (
"litellm.caching.valkey_semantic_cache",
"ValkeySemanticCache",
"valkey-semantic",
),
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
"azure-blob",
),
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -394,6 +445,11 @@ impl FacadeGuard {
"embedding_timeout",
"_index_name",
"_redis_url",
"similarity_threshold",
"embedding_model",
"index_name",
"embedding_max_input_tokens",
"embedding_timeout",
"bucket_name",
"key_prefix",
"path_service_account",

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

@ -1,5 +1,8 @@
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_redis_semantic::RedisSemanticConfig;
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{
PyTraverseError, PyVisit,
@ -7,8 +10,6 @@ use pyo3::{
prelude::*,
};
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use super::{
cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard,
native::NativeResponseCache, request::duration,
@ -74,6 +75,40 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (bucket, *, region, endpoint_url=None, key_prefix="", access_key_id=None, secret_access_key=None, session_token=None))]
fn s3(
py: Python<'_>,
bucket: String,
region: String,
endpoint_url: Option<String>,
key_prefix: &str,
access_key_id: Option<String>,
secret_access_key: Option<String>,
session_token: Option<String>,
) -> PyResult<Self> {
let config = S3CacheConfig {
bucket,
key_prefix: key_prefix.to_string(),
region: region.clone(),
endpoint: endpoint_url.map(|url| S3Endpoint { url }),
auth: AwsAuthConfig {
access_key_id,
secret_access_key,
session_token,
region_name: Some(region),
..Default::default()
},
};
let service = run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) })?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (bucket_name, *, gcs_path=None, path_service_account=None, endpoint=None, token=None))]
fn gcs(
@ -111,6 +146,29 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[pyo3(signature = (url, similarity_threshold, index_name, embedder))]
fn valkey_semantic(
url: String,
similarity_threshold: f64,
index_name: String,
embedder: &Bound<'_, PyAny>,
) -> PyResult<Self> {
let python_embedder = PythonEmbedder::from_backend(embedder)?;
let service = NativeResponseCache::valkey_semantic(
&url,
similarity_threshold,
index_name,
python_embedder,
)
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (account_url, container))]
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {
@ -164,11 +222,17 @@ impl CacheTestHandle {
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 service = service
.with_scope(
facade
.getattr("semantic_cache_scope")?
.extract::<String>()?,
)
.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
let handle = Py::new(
py,
Self {

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

@ -9,6 +9,7 @@ mod native;
mod request;
mod resolver;
mod semantic;
mod semantic_step;
use litellm_cache::Error;
use pyo3::{

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

File diff suppressed because it is too large Load diff

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

@ -1,11 +1,11 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::{ExactCacheContext, SemanticCacheContext};
use litellm_cache::ExactCacheContext;
use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest};
use litellm_host_python::from_py;
use pyo3::{exceptions::PyValueError, prelude::*};
use serde::Deserialize;
use serde_json::{Map, Value};
use serde_json::Value;
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
@ -14,69 +14,52 @@ struct RequestInput {
controls: Option<CacheControls>,
ttl_seconds: Option<f64>,
max_age_seconds: Option<f64>,
messages: Option<Value>,
input: Option<Value>,
messages: Option<Vec<Value>>,
metadata: Option<Map<String, Value>>,
metadata: Option<Value>,
litellm_metadata: Option<Value>,
litellm_params: Option<Value>,
scope: Option<String>,
}
#[derive(Clone)]
pub(super) struct CacheRequest {
key: CacheKeyInput,
controls: CacheControls,
ttl: Option<Duration>,
max_age: Option<Duration>,
input: Option<Value>,
messages: Vec<Value>,
metadata: Map<String, Value>,
scope: Option<String>,
pub(super) struct NativeRequest {
pub(super) key: CacheKeyInput,
pub(super) controls: CacheControls,
pub(super) ttl: Option<Duration>,
pub(super) max_age: Option<Duration>,
pub(super) messages: Option<Value>,
pub(super) input: Option<Value>,
pub(super) metadata: Option<Value>,
pub(super) litellm_metadata: Option<Value>,
pub(super) litellm_params: Option<Value>,
pub(super) scope: Option<String>,
}
impl CacheRequest {
pub(super) fn exact(&self) -> ResponseCacheRequest<ExactCacheContext> {
let mut request = ResponseCacheRequest::new(self.key.clone());
request.controls = self.controls;
request.context.ttl = self.ttl;
request.max_age = self.max_age;
request
}
pub(super) fn semantic(&self) -> ResponseCacheRequest<SemanticCacheContext> {
ResponseCacheRequest {
key: self.key.clone(),
controls: self.controls,
context: SemanticCacheContext {
input: self.input.clone(),
messages: self.messages.clone(),
metadata: self.metadata.clone(),
scope: self.scope.clone(),
ttl: self.ttl,
},
max_age: self.max_age,
}
}
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<CacheRequest> {
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<NativeRequest> {
let input: RequestInput = from_py(value)?;
request_input(input)
}
fn request_input(input: RequestInput) -> PyResult<CacheRequest> {
let defaults = ResponseCacheRequest::<ExactCacheContext>::new(input.key.clone());
Ok(CacheRequest {
fn request_input(input: RequestInput) -> PyResult<NativeRequest> {
let controls = input.controls.unwrap_or_else(|| {
ResponseCacheRequest::<ExactCacheContext>::new(input.key.clone()).controls
});
Ok(NativeRequest {
key: input.key,
controls: input.controls.unwrap_or(defaults.controls),
controls,
ttl: input.ttl_seconds.map(duration).transpose()?,
max_age: input.max_age_seconds.map(duration).transpose()?,
messages: input.messages,
input: input.input,
messages: input.messages.unwrap_or_default(),
metadata: input.metadata.unwrap_or_default(),
metadata: input.metadata,
litellm_metadata: input.litellm_metadata,
litellm_params: input.litellm_params,
scope: input.scope,
})
}
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<CacheRequest>> {
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<NativeRequest>> {
from_py::<Vec<RequestInput>>(value)?
.into_iter()
.map(request_input)

27
litellm-rust/crates/python-bridge/src/cache/semantic.rs vendored
View file

@ -14,13 +14,13 @@ use super::{
cache_error,
embedder::{PythonEmbedder, with_prepared_embedding},
native::NativeResponseCache,
request::{CacheRequest, now},
request::{NativeRequest, now},
};
pub(super) enum SemanticOperation {
Lookup(CacheRequest),
Store(CacheRequest, Value),
StoreBatch(VecDeque<(CacheRequest, Value)>),
Lookup(NativeRequest),
Store(NativeRequest, Value),
StoreBatch(VecDeque<(NativeRequest, Value)>),
}
enum Phase {
@ -32,7 +32,7 @@ enum Phase {
pub(super) struct SemanticBody {
service: NativeResponseCache,
operation: SemanticOperation,
pending: Option<(CacheRequest, Option<Value>)>,
pending: Option<(NativeRequest, Option<Value>)>,
phase: Phase,
}
@ -101,7 +101,7 @@ impl ExecutionBody for SemanticBody {
let (request, _) = self.pending.as_ref().ok_or_else(|| {
PyRuntimeError::new_err("semantic execution has no pending operation")
})?;
let semantic = request.semantic();
let semantic = NativeResponseCache::redis_semantic_request(request);
let Some(prompt) = prompt_from_context(&semantic.context) else {
return self.backend_step(py, Err(Error::Unavailable));
};
@ -113,7 +113,7 @@ impl ExecutionBody for SemanticBody {
let coroutine = embedder.async_embedding_coroutine(
py,
&prompt,
&semantic.context.metadata,
semantic.context.metadata.as_ref(),
)?;
self.phase = Phase::AwaitingEmbedding;
return Ok(ExecutionStep::Await(coroutine));
@ -160,16 +160,15 @@ impl ExecutionBody for SemanticBody {
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.service.traverse(visit)
if let Some(embedder) = self.service.semantic_embedder() {
embedder.traverse(visit)?;
}
Ok(())
}
}
pub(super) fn drive(
py: Python<'_>,
service: NativeResponseCache,
operation: SemanticOperation,
) -> PyResult<Bound<'_, PyAny>> {
let execution = Py::new(py, Execution::new(SemanticBody::new(service, operation)))?;
pub(super) fn drive(py: Python<'_>, body: SemanticBody) -> PyResult<Bound<'_, PyAny>> {
let execution = Py::new(py, Execution::new(body))?;
py.import("litellm.rust_bridge.lifecycle")?
.getattr("drive")?
.call1((execution,))

249
litellm-rust/crates/python-bridge/src/cache/semantic_step.rs vendored Normal file
View file

@ -0,0 +1,249 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::SemanticCacheContext;
use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest};
use litellm_cache_valkey_semantic::{PreparedEmbedding, ValkeySemanticCache, prompt_from_context};
use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use serde_json::Value;
use super::{cache_error, embedder::PythonEmbedder};
pub(super) enum Op {
Lookup,
Store(Value),
StoreBatch(Vec<Value>),
}
#[derive(Clone, Copy)]
enum State {
Start,
AwaitingEmbedding,
AwaitingStorage,
Done,
}
pub(super) struct SemanticEmbedExecution {
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
op: Op,
now: Option<Duration>,
prepared: Vec<Option<Vec<f32>>>,
index: usize,
state: State,
}
impl SemanticEmbedExecution {
pub(super) fn lookup(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
) -> Self {
Self {
backend,
embedder,
requests: vec![request],
op: Op::Lookup,
now: None,
prepared: vec![None],
index: 0,
state: State::Start,
}
}
pub(super) fn store(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
response: Value,
) -> Self {
Self {
backend,
embedder,
requests: vec![request],
op: Op::Store(response),
now: None,
prepared: vec![None],
index: 0,
state: State::Start,
}
}
pub(super) fn store_batch(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
responses: Vec<Value>,
) -> Self {
Self {
backend,
embedder,
prepared: vec![None; requests.len()],
requests,
op: Op::StoreBatch(responses),
now: None,
index: 0,
state: State::Start,
}
}
fn start(&mut self, py: Python<'_>) -> PyResult<ExecutionStep> {
if self.now.is_none() {
self.now = Some(super::request::now());
}
while self.index < self.requests.len() {
let request = &self.requests[self.index];
let enabled = match &self.op {
Op::Lookup => request.controls.reads(),
Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(),
};
if !enabled {
self.index += 1;
continue;
}
let Some(prompt) = prompt_from_context(&request.context) else {
self.index += 1;
continue;
};
let metadata = request.context.metadata.clone();
let awaitable = self
.embedder
.async_embed_awaitable(py, &prompt, &metadata)?;
self.state = State::AwaitingEmbedding;
return Ok(ExecutionStep::Await(awaitable.unbind()));
}
self.state = State::AwaitingStorage;
self.storage_step(py)
}
fn storage_step(&self, py: Python<'_>) -> PyResult<ExecutionStep> {
let requests = self.requests.clone();
let prepared = self.prepared.clone();
let backend = Arc::clone(&self.backend);
let now = self
.now
.ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?;
let awaitable = match &self.op {
Op::Lookup => {
let Some(request) = requests.into_iter().next() else {
return Err(PyRuntimeError::new_err(
"semantic lookup requires one request",
));
};
match prepared.into_iter().next().flatten() {
Some(values) => {
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = Arc::new(ResponseCache::new(Arc::new(backend)));
run_async(
py,
async move { cache.async_lookup(&request, now).await },
cache_error,
)?
}
None => {
let cache = Arc::new(ResponseCache::new(backend));
run_async(
py,
async move { cache.async_lookup(&request, now).await },
cache_error,
)?
}
}
}
Op::Store(response) => {
let Some(request) = requests.into_iter().next() else {
return Err(PyRuntimeError::new_err(
"semantic store requires one request",
));
};
let response = response.clone();
match prepared.into_iter().next().flatten() {
Some(values) => {
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = Arc::new(ResponseCache::new(Arc::new(backend)));
run_async(
py,
async move { cache.async_store(&request, response, now).await },
cache_error,
)?
}
None => {
let cache = Arc::new(ResponseCache::new(backend));
run_async(
py,
async move { cache.async_store(&request, response, now).await },
cache_error,
)?
}
}
}
Op::StoreBatch(responses) => {
let responses = responses.clone();
run_async(
py,
async move {
for ((request, response), prepared) in
requests.into_iter().zip(responses).zip(prepared)
{
let Some(values) = prepared else {
continue;
};
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = ResponseCache::new(Arc::new(backend));
cache.async_store(&request, response, now).await?;
}
Ok(())
},
cache_error,
)?
}
};
Ok(ExecutionStep::Await(awaitable.unbind()))
}
fn resume_py(
&mut self,
py: Python<'_>,
result: Option<PyResult<Py<PyAny>>>,
) -> PyResult<ExecutionStep> {
match (self.state, result) {
(State::Start, None) => self.start(py),
(State::AwaitingEmbedding, Some(Ok(value))) => {
let values = value.bind(py).extract::<Vec<f64>>()?;
self.prepared[self.index] =
Some(values.into_iter().map(|value| value as f32).collect());
self.index += 1;
self.start(py)
}
(State::AwaitingStorage, Some(Ok(value))) => {
self.state = State::Done;
Ok(ExecutionStep::Return(value))
}
(_, Some(Err(error))) => Err(error),
_ => Err(PyRuntimeError::new_err(
"invalid semantic cache execution state",
)),
}
}
}
impl ExecutionBody for SemanticEmbedExecution {
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| self.resume_py(py, result))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.embedder.traverse(visit)
}
}
pub(super) fn drive_semantic<'py>(
py: Python<'py>,
body: SemanticEmbedExecution,
) -> PyResult<Bound<'py, PyAny>> {
let execution = Py::new(py, Execution::new(body))?;
py.import("litellm.rust_bridge.lifecycle")?
.getattr("drive")?
.call1((execution,))
}

4
litellm/litellm_core_utils/logging_worker.py
View file

@ -484,9 +484,13 @@ class LoggingWorker:
so it correctly handles items that have been dequeued but whose
callback hasn't finished yet — ``queue.empty()`` would return True in
that window and cause us to skip the wait.
``start()`` runs first so a queue left behind by a previous event loop
is carried onto this one and drained here instead of joined forever.
"""
if self._queue is None:
return
self.start()
await self._queue.join()
async def clear_queue(self):

75
litellm/model_prices_and_context_window_backup.json
View file

@ -43037,21 +43037,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 8.92272e-07,
"input_cost_per_token": 8.87226e-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.784544e-06,
"output_cost_per_token": 1.774452e-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.4356e-08,
"cache_read_input_token_cost": 7.39355e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -73794,6 +73794,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_above_128k_tokens": 5e-07,
"litellm_provider": "openrouter",
@ -73815,6 +73816,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6-flash": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 7.5e-08,
"input_cost_per_token_above_128k_tokens": 1e-07,
"litellm_provider": "openrouter",
@ -73855,6 +73857,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-2.0-code": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 5e-07,
"input_cost_per_token_above_128k_tokens": 1e-06,
"litellm_provider": "openrouter",
@ -76876,6 +76879,26 @@
"supports_vision": true,
"supports_web_search": true
},
"moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"global.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
@ -76975,5 +76998,51 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"xiaomi_mimo/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"xiaomi_mimo/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
}
}

11
litellm/proxy/auth/handle_jwt.py
View file

@ -14,7 +14,7 @@ import hashlib
import os
import re
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
@ -1602,6 +1602,7 @@ class JWTAuthManager:
team_object: LiteLLM_TeamTable | None,
route: str,
request_method: str | None = None,
team_allowed_routes: Collection[str] = (),
) -> bool:
normalized_request_method: Final = request_method.upper() if isinstance(request_method, str) else None
if not RouteChecks.is_auth_enforced_pass_through_route(
@ -1610,8 +1611,11 @@ class JWTAuthManager:
):
return True
if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=team_allowed_routes):
return True
# JWT team selection is team-scoped; key metadata is not available here,
# so passthrough access is granted only by the selected team's metadata.
# so beyond the JWT config grant above, only the selected team's metadata grants access.
return RouteChecks.check_passthrough_route_access(
route=route,
user_api_key_dict=UserAPIKeyAuth(team_metadata=(team_object.metadata or {}) if team_object else {}),
@ -1689,6 +1693,7 @@ class JWTAuthManager:
team_object=team_object,
route=route,
request_method=request_method,
team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes,
):
is_allowed = False
denied_auth_enforced_pass_through_route = True
@ -2584,6 +2589,7 @@ class JWTAuthManager:
team_object=team_object,
route=route,
request_method=request_method,
team_allowed_routes=handler.litellm_jwtauth.team_allowed_routes,
):
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
@ -2653,6 +2659,7 @@ class JWTAuthManager:
team_object=team_object,
route=route,
request_method=request_method,
team_allowed_routes=handler.litellm_jwtauth.team_allowed_routes,
):
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
elif team_id is None:

35
litellm/proxy/auth/route_checks.py
View file

@ -278,7 +278,11 @@ class RouteChecks:
route=route,
method=RouteChecks._get_request_method(request=request),
):
RouteChecks._require_auth_pass_through_access(route=route, valid_token=valid_token)
RouteChecks._require_auth_pass_through_access(
route=route,
valid_token=valid_token,
jwt_team_allowed_routes=RouteChecks._jwt_team_allowed_routes(valid_token=valid_token),
)
elif RouteChecks.is_llm_api_route(route=route):
pass
elif RouteChecks.is_info_route(route=route):
@ -689,16 +693,43 @@ class RouteChecks:
),
)
@staticmethod
def jwt_team_routes_grant_pass_through(route: str, team_allowed_routes: Collection[str]) -> bool:
"""
Explicit paths and trailing-wildcard prefixes grant auth=true pass-through. Blanket grants never do:
a named route group like ``openai_routes`` is only ever compared as a path, and an entry that names
no path segment (``*``, ``/*``) is skipped.
"""
return any(
RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route)
for allowed_route in team_allowed_routes
if allowed_route.rstrip("*").strip("/")
)
@staticmethod
def _jwt_team_allowed_routes(valid_token: UserAPIKeyAuth) -> Collection[str]:
"""``team_allowed_routes`` for team tokens built by JWT auth; JWT-mapped virtual keys stay key-scoped."""
if valid_token.jwt_claims is None or valid_token.token is not None or valid_token.team_id is None:
return ()
from litellm.proxy.proxy_server import jwt_handler
return jwt_handler.litellm_jwtauth.team_allowed_routes
@staticmethod
def _require_auth_pass_through_access(
route: str,
valid_token: UserAPIKeyAuth,
jwt_team_allowed_routes: Collection[str] = (),
) -> None:
"""
Require an explicit ``allowed_passthrough_routes`` match for auth=true pass-through.
Require an explicit grant for auth=true pass-through: ``allowed_passthrough_routes`` on the
key or team, or an explicit JWT ``team_allowed_routes`` entry.
"""
if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token):
return
if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=jwt_team_allowed_routes):
return
raise RouteChecks._auth_pass_through_denied_exception(route=route)
@staticmethod

9
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -1078,6 +1078,9 @@ if MCP_AVAILABLE:
return {"servers": registry_servers}
## FastAPI Routes
def _mcp_server_display_order(server: LiteLLM_MCPServerTable) -> tuple[str, str]:
return ((server.server_name or server.alias or server.server_id).lower(), server.server_id)
def _get_user_mcp_management_mode() -> UserMCPManagementMode:
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
@ -1228,10 +1231,12 @@ if MCP_AVAILABLE:
detail="You do not have permission to view MCP servers for this team.",
)
redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id)
redacted_mcp_servers = sorted(
await _get_team_scoped_mcp_server_list(sanitized_team_id), key=_mcp_server_display_order
)
else:
servers: Final = await _resolve_accessible_mcp_servers(user_api_key_dict)
redacted_mcp_servers = _redact_mcp_credentials_list(servers)
redacted_mcp_servers = sorted(_redact_mcp_credentials_list(servers), key=_mcp_server_display_order)
if connected_app_view is True and is_ui_session_credential(user_api_key_dict):
reachable_ids: Final = await _connected_app_reachable_server_ids(user_api_key_dict)

121
litellm/rust_bridge/_native.pyi
View file

@ -1,6 +1,6 @@
from asyncio import Future
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from typing import Literal, Never, final
from typing import Never, final
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
@ -93,6 +93,59 @@ class ResponsesWebSocketConnection:
def recv_text(self) -> Future[str | None]: ...
def close(self) -> Future[None]: ...
@final
class _CacheTestBinding:
@property
def kind(self) -> str: ...
def lookup(
self,
request: object,
*,
callback_kwargs: Mapping[str, object] | Sequence[object] | None = None,
) -> object: ...
def store(
self,
request: object,
response: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> None: ...
def lookup_batch(
self,
requests: Sequence[object],
*,
callback_kwargs: Sequence[object] | None = None,
) -> object: ...
def async_lookup(
self,
request: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[object]: ...
def async_store(
self,
request: object,
response: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[None]: ...
def async_lookup_batch(
self,
requests: Sequence[object],
*,
callback_kwargs: Sequence[object] | None = None,
) -> Future[object]: ...
def async_store_batch(
self,
requests: Sequence[object],
responses: Sequence[object],
*,
callback_result: object = None,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[object]: ...
def async_flush(self) -> Future[None]: ...
def ping(self) -> Future[object]: ...
@final
class _CacheTestHandle:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
@ -109,11 +162,22 @@ class _CacheTestHandle:
*,
ttl_seconds: float = 60.0,
namespace: str | None = None,
startup_nodes: list[tuple[str, int]] | None = None,
startup_nodes: Sequence[tuple[str, int]] | None = None,
) -> _CacheTestHandle: ...
@staticmethod
def disk(directory: str) -> _CacheTestHandle: ...
@staticmethod
def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
@staticmethod
def redis_semantic(backend: object) -> _CacheTestHandle: ...
@staticmethod
def valkey_semantic(
url: str,
similarity_threshold: float,
index_name: str,
embedder: object,
) -> _CacheTestHandle: ...
@staticmethod
def gcs(
bucket_name: str,
*,
@ -123,48 +187,19 @@ class _CacheTestHandle:
token: str | None = None,
) -> _CacheTestHandle: ...
@staticmethod
def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
@staticmethod
def redis_semantic(backend: object) -> _CacheTestHandle: ...
@property
def backend(
self,
) -> Literal["memory", "redis", "gcs", "disk", "azure-blob", "redis_semantic"]: ...
def _bind_facade(self, facade: object) -> None: ...
@final
class _CacheTestBinding:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
@property
def kind(self) -> Literal["disabled", "native", "python_callback"]: ...
def lookup(
self, request: object, *, callback_kwargs: object = None
) -> object: ...
def store(
self, request: object, response: object, *, callback_kwargs: object = None
) -> None: ...
def lookup_batch(
self, requests: object, *, callback_kwargs: object = None
) -> object: ...
def async_lookup(
self, request: object, *, callback_kwargs: object = None
) -> Future[object]: ...
def async_store(
self, request: object, response: object, *, callback_kwargs: object = None
) -> Future[object]: ...
def async_lookup_batch(
self, requests: object, *, callback_kwargs: object = None
) -> Future[object]: ...
def async_store_batch(
self,
requests: object,
responses: object,
def s3(
bucket: str,
*,
callback_result: object = None,
callback_kwargs: object = None,
) -> Future[object]: ...
def async_flush(self) -> Future[object]: ...
def ping(self) -> Future[object]: ...
region: str,
endpoint_url: str | None = None,
key_prefix: str = "",
access_key_id: str | None = None,
secret_access_key: str | None = None,
session_token: str | None = None,
) -> _CacheTestHandle: ...
@property
def backend(self) -> str: ...
def _bind_facade(self, facade: object) -> None: ...
@final
class _CacheTestResolver:

75
model_prices_and_context_window.json
View file

@ -43037,21 +43037,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 8.92272e-07,
"input_cost_per_token": 8.87226e-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.784544e-06,
"output_cost_per_token": 1.774452e-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.4356e-08,
"cache_read_input_token_cost": 7.39355e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -73794,6 +73794,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_above_128k_tokens": 5e-07,
"litellm_provider": "openrouter",
@ -73815,6 +73816,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6-flash": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 7.5e-08,
"input_cost_per_token_above_128k_tokens": 1e-07,
"litellm_provider": "openrouter",
@ -73855,6 +73857,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-2.0-code": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 5e-07,
"input_cost_per_token_above_128k_tokens": 1e-06,
"litellm_provider": "openrouter",
@ -76876,6 +76879,26 @@
"supports_vision": true,
"supports_web_search": true
},
"moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"global.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
@ -76975,5 +76998,51 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"xiaomi_mimo/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"xiaomi_mimo/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
}
}

3
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -89,6 +89,9 @@
- {id: llm.chat_completions.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip"}
- {id: llm.chat_completions.together_ai.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together cost header and spend row match the registry price"}
- {id: llm.chat_completions.together_ai.thinking.nonstream.effort_none_disables, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [effort_none_disables], source: "llm_translation/test_together_ai_e2e.py", rationale: "reasoning_effort=none maps to Together's reasoning disable toggle on hybrid models"}
- {id: llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "Native MiMo v2.6 rows price the cost header and spend row from the cost map"}
- {id: llm.chat_completions.xiaomi_mimo.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo reasoning deltas stream as reasoning_content"}
- {id: llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo tool calls are not dropped"}
- {id: llm.chat_completions.together_ai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: structured_output, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "response_format json_schema reaches Together and constrains the reply"}
- {id: llm.chat_completions.together_ai.prompt_cache_5m.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: prompt_cache_5m, streaming: nonstream, assertions: [cache_hit, cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together prefix-cache reads bill at cache_read_input_token_cost, not full input price"}
- {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"}

1
tests/e2e/coverage_registry/schema.py
View file

@ -58,6 +58,7 @@ LlmRoute = Literal[
"openai",
"together_ai",
"vertex",
"xiaomi_mimo",
]
LlmCapability = Literal[

214
tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py Normal file
View file

@ -0,0 +1,214 @@
"""Live e2e: Xiaomi MiMo v2.6 through the gateway on /chat/completions.
Both native ``xiaomi_mimo/`` v2.6 rows (pro and flash) are registered via
``/model/new`` and driven against Xiaomi's own endpoint. What the gateway owes
us is that the reasoning chain surfaces as ``reasoning_content``, tool calls
survive translation, and the cost header plus spend row follow the proxy's own
cost-map price for the row (read back from ``/model/info``, never pinned here).
Requires XIAOMI_MIMO_API_KEY on the proxy; no skip gate.
"""
from __future__ import annotations
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import StreamingResponse, require_successful_call, unwrap
from lifecycle import ResourceManager
from models import (
ChatBody,
ChatMessage,
ChatResponse,
ChatTool,
ChatToolFunction,
CostMapEntry,
LiteLLMParamsBody,
OutMessage,
SpendLogRow,
)
from passthrough_client import PassthroughClient
from pydantic import BaseModel
pytestmark = pytest.mark.e2e
BACKENDS: Final = ("xiaomi_mimo/mimo-v2.6-pro", "xiaomi_mimo/mimo-v2.6-flash")
ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number."
WEATHER_PROMPT = "What is the weather in Paris? Use the tool."
COUNTING_PROMPT = "Count from 1 to 50, one number per line."
WEATHER_TOOL = ChatTool(
function=ChatToolFunction(
name="get_weather",
description="Get the current weather for a location.",
parameters={
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
)
)
class _WeatherArgs(BaseModel):
location: str
class _StreamDelta(BaseModel):
content: str | None = None
reasoning_content: str | None = None
class _StreamChoice(BaseModel):
delta: _StreamDelta | None = None
class _StreamChunk(BaseModel):
choices: list[_StreamChoice] = []
def _approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
@pytest.fixture(scope="module")
def registry(client: PassthroughClient) -> dict[str, CostMapEntry]:
return client.proxy.model_cost_map()
def _register(client: PassthroughClient, resources: ResourceManager, backend: str) -> tuple[str, str]:
model = f"e2e-xiaomi-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=backend, api_key="os.environ/XIAOMI_MIMO_API_KEY")
)
resources.defer(lambda: client.proxy.delete_model(model_id))
return model, resources.key()
def _message(response: ChatResponse) -> OutMessage:
assert response.choices, f"Xiaomi returned no choices: {response}"
message = response.choices[0].message
assert message is not None, f"Xiaomi choice has no message: {response}"
return message
def _deltas(result: StreamingResponse) -> list[_StreamDelta]:
require_successful_call(result)
assert result.is_streaming, f"response was not streamed: {result.headers}"
assert not result.stream_error, f"stream errored: {result.stream_error}"
assert result.stream_done, f"stream never reached [DONE]: {result.stream_events[-3:]}"
return [
choice.delta
for event in result.stream_events
for choice in _StreamChunk.model_validate_json(event).choices
if choice.delta is not None
]
@pytest.mark.parametrize("backend", BACKENDS)
class TestXiaomiMimoChatCompletions:
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged")
def test_cost_header_and_spend_row_match_the_registry_price(
self,
client: PassthroughClient,
resources: ResourceManager,
registry: dict[str, CostMapEntry],
backend: str,
) -> None:
price = registry.get(backend)
assert price is not None, f"{backend} has no row in the proxy's cost map, so native calls would bill $0"
assert price.litellm_provider == "xiaomi_mimo", f"{backend} is filed under the wrong provider: {price}"
assert price.input_cost_per_token and price.output_cost_per_token, f"{backend} carries no price: {price}"
model, key = _register(client, resources, backend)
result = client.proxy.transport.send(
"/chat/completions",
headers=client.proxy.transport.bearer(key),
json=ChatBody(
model=model,
messages=[ChatMessage(role="user", content=f"{ARITHMETIC_PROMPT} {unique_marker()}")],
max_tokens=1024,
),
)
require_successful_call(result)
response = ChatResponse.model_validate_json(result.body)
message = _message(response)
assert message.content and "43" in message.content, f"answer lost: {message}"
assert message.reasoning_content, f"{backend} reasons, but no reasoning_content came back: {message}"
usage = response.usage
assert usage is not None and usage.prompt_tokens and usage.completion_tokens, (
f"response carries no usage, so the cost cannot be real: {result.body[:300]}"
)
header_cost = result.response_cost
assert header_cost is not None and header_cost > 0, (
f"x-litellm-response-cost header missing or non-positive: {result.headers}"
)
cached = (usage.prompt_tokens_details.cached_tokens or 0) if usage.prompt_tokens_details else 0
expected = (
(usage.prompt_tokens - cached) * price.input_cost_per_token
+ cached * (price.cache_read_input_token_cost or 0.0)
+ usage.completion_tokens * price.output_cost_per_token
)
assert _approx_equal(header_cost, expected), (
f"header cost {header_cost} disagrees with the registry price for {backend} at {usage}: expected {expected}"
)
def _priced(rows: list[SpendLogRow]) -> bool:
return any(row.spend is not None and row.spend > 0 for row in rows)
rows = client.proxy.poll_logs_for_key(key, predicate=_priced)
priced = [row for row in rows if row.spend is not None and row.spend > 0]
assert priced, f"no priced spend row landed for key {key}; got {rows}"
row = priced[0]
assert row.custom_llm_provider == "xiaomi_mimo", f"spend row misattributed: {row}"
assert row.spend is not None and _approx_equal(row.spend, header_cost), (
f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}"
)
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.thinking.stream.works")
def test_reasoning_and_answer_stream_as_deltas(
self, client: PassthroughClient, resources: ResourceManager, backend: str
) -> None:
model, key = _register(client, resources, backend)
deltas = _deltas(
client.proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=COUNTING_PROMPT)],
max_tokens=2048,
stream=True,
),
)
)
reasoning = "".join(delta.reasoning_content or "" for delta in deltas)
content = "".join(delta.content or "" for delta in deltas)
assert reasoning, f"stream carried no reasoning_content deltas: {deltas[:5]}"
assert "50" in content, f"streamed answer lost: {content[:300]!r}"
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works")
def test_tool_call_is_returned(self, client: PassthroughClient, resources: ResourceManager, backend: str) -> None:
model, key = _register(client, resources, backend)
message = _message(
unwrap(
client.proxy.chat(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=WEATHER_PROMPT)],
tools=[WEATHER_TOOL],
max_tokens=1024,
),
)
)
)
assert message.tool_calls, f"{backend} dropped the tool call: {message}"
call = message.tool_calls[0]
assert call.id, f"tool call carries no id, so a tool result cannot answer it: {call}"
assert call.function.name == "get_weather", f"wrong tool called: {call}"
assert call.function.arguments, f"tool call carries no arguments: {call}"
args = _WeatherArgs.model_validate_json(call.function.arguments)
assert "paris" in args.location.lower(), f"tool arguments lost the location: {args}"

47
tests/test_litellm/litellm_core_utils/test_logging_worker.py
View file

@ -180,6 +180,53 @@ class TestLoggingWorker:
assert sorted(fired) == ["first", "second"]
@pytest.mark.parametrize("stranded", ["still_queued", "dequeued_never_started"])
def test_flush_on_new_loop_drains_tasks_stranded_on_previous_loop(self, stranded):
"""
Regression: ``flush()`` from a new event loop used to ``join()`` the queue bound to the
previous loop, whose unfinished counter nothing on the new loop ever decrements. The first
such flush hung until pytest-timeout killed it and every later one raised
``RuntimeError: ... is bound to a different event loop`` from the queue's Event.
"""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
callback = AsyncMock()
async def enqueue_on_first_loop():
if stranded == "still_queued":
worker._ensure_queue()
worker.enqueue(callback())
return
worker.ensure_initialized_and_enqueue(callback())
asyncio.run(enqueue_on_first_loop())
assert worker._queue is not None
expected_shape = (1, 0) if stranded == "still_queued" else (0, 1)
assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape
assert callback.await_count == 0, "precondition: the callback never ran before the first loop closed"
async def flush_twice_on_second_loop():
await asyncio.wait_for(worker.flush(), timeout=5)
await asyncio.wait_for(worker.flush(), timeout=5)
asyncio.run(flush_twice_on_second_loop())
assert callback.await_count == 1
def test_flush_starts_a_worker_when_the_queue_has_none(self):
"""``flush()`` must drain a queue that exists on the current loop without a running worker."""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
callback = AsyncMock()
async def enqueue_then_flush():
worker._ensure_queue()
worker.enqueue(callback())
assert worker._worker_task is None, "precondition: nothing is draining the queue yet"
await asyncio.wait_for(worker.flush(), timeout=3)
asyncio.run(enqueue_then_flush())
assert callback.await_count == 1
def test_flush_on_exit_swallows_cancellation_and_drains_remaining(self):
"""A callback raising CancelledError must not abort the atexit flush of later events."""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)

184
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -291,6 +291,106 @@ async def test_find_team_with_model_access_uses_request_method_for_passthrough_a
assert "allowed_passthrough_routes" in exc_info.value.detail
_AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = {
"test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": {
"endpoint_id": "test-uuid-1",
"path": "/model-host/v1/extractor",
"type": "subpath",
"auth": True,
},
}
@pytest.mark.asyncio
async def test_find_team_with_model_access_team_allowed_routes_wildcard_grants_auth_passthrough():
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_allowed_routes=["openai_routes", "/model-host/*"])
team_without_passthrough_allowlist = LiteLLM_TeamTable(team_id="team-a", models=["all-proxy-models"], metadata={})
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=team_without_passthrough_allowlist,
),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
_AUTH_ENFORCED_MODEL_HOST_ROUTES,
),
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
):
team_id, team_obj = await JWTAuthManager.find_team_with_model_access(
team_ids={"team-a"},
requested_model=None,
route="/model-host/v1/extractor/predict",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
request_method="POST",
)
assert team_id == "team-a"
assert team_obj == team_without_passthrough_allowlist
@pytest.mark.asyncio
async def test_auth_builder_header_team_allows_auth_passthrough_for_team_allowed_routes_wildcard():
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
user_id_jwt_field="sub",
team_allowed_routes=["openai_routes", "/model-host/*"],
),
)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=LiteLLM_TeamTable(team_id="team-2", metadata={}),
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(None, None, None, None, "user-1"),
),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
_AUTH_ENFORCED_MODEL_HOST_ROUTES,
),
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
):
mock_auth_jwt.return_value = {"sub": "user-1", "scope": "", "groups": ["team-1", "team-2"]}
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={},
general_settings={},
route="/model-host/v1/extractor/predict",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_headers={"x-litellm-team-id": "team-2"},
request_method="POST",
)
assert result["team_id"] == "team-2"
@pytest.mark.asyncio
async def test_auth_builder_proxy_admin_user_role():
"""Test that is_proxy_admin is True when user_object.user_role is PROXY_ADMIN"""
@ -6463,6 +6563,90 @@ async def test_auth_builder_db_fallback_enforces_passthrough_route_access():
assert "passthrough route" in exc_info.value.detail
async def _auth_builder_via_db_team_fallback(team_allowed_routes: list[str]):
user_id = "u_passthrough"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["team_no_passthrough"],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_allowed_routes=team_allowed_routes)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id, metadata={})
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock, return_value={"sub": user_id, "scope": ""}),
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
_AUTH_ENFORCED_MODEL_HOST_ROUTES,
),
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
):
return await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={},
general_settings={"enforce_rbac": False},
route="/model-host/v1/extractor/predict",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers=None,
request_method="POST",
)
@pytest.mark.asyncio
async def test_auth_builder_db_fallback_team_allowed_routes_wildcard_grants_auth_passthrough():
result = await _auth_builder_via_db_team_fallback(team_allowed_routes=["openai_routes", "/model-host/*"])
assert result["team_id"] == "team_no_passthrough"
@pytest.mark.asyncio
async def test_auth_builder_db_fallback_route_groups_alone_do_not_grant_auth_passthrough():
with pytest.raises(HTTPException) as exc_info:
await _auth_builder_via_db_team_fallback(team_allowed_routes=["openai_routes", "mapped_pass_through_routes"])
assert exc_info.value.status_code == 403, exc_info.value.detail
assert "allowed_passthrough_routes" in exc_info.value.detail
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_singular_claim_reconciles_memberships():
"""When fallback_to_db_teams is on but the JWT carries a singular team claim

116
tests/test_litellm/proxy/auth/test_route_checks.py
View file

@ -1249,6 +1249,122 @@ def test_non_proxy_admin_allows_auth_pass_through_with_team_allowlist():
)
@pytest.mark.parametrize(
"route, team_allowed_routes, expected",
[
("/model-host/v1/extractor/predict", ["/model-host/*"], True),
("/model-host", ["/model-host/*"], False),
("/model-host/v1/extractor", ["/model-host/v1/extractor"], True),
("/model-host/v1/extractor/predict", ["/model-host/v1/extractor"], False),
("/other/v1/extractor", ["/model-host/*"], False),
("/model-host/v1/extractor", ["openai_routes", "llm_api_routes", "mapped_pass_through_routes"], False),
("/model-host/v1/extractor", ["*"], False),
("/model-host/v1/extractor", ["/*"], False),
("/model-host/v1/extractor", [], False),
],
)
def test_jwt_team_routes_grant_pass_through_only_for_explicit_paths(route, team_allowed_routes, expected):
assert (
RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=team_allowed_routes)
is expected
)
_AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = {
"test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": {
"endpoint_id": "test-uuid-1",
"path": "/model-host/v1/extractor",
"type": "subpath",
"auth": True,
},
}
def _jwt_handler_with_team_allowed_routes(team_allowed_routes: list[str]):
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
jwt_handler: Final = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_allowed_routes=team_allowed_routes)
return jwt_handler
def _check_model_host_route_as(valid_token: UserAPIKeyAuth, team_allowed_routes: list[str]) -> None:
with (
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
_AUTH_ENFORCED_MODEL_HOST_ROUTES,
),
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
patch(
"litellm.proxy.proxy_server.jwt_handler",
_jwt_handler_with_team_allowed_routes(team_allowed_routes),
),
):
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=None,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
route="/model-host/v1/extractor/predict",
request=MagicMock(spec=Request),
valid_token=valid_token,
request_data={},
)
def test_non_proxy_admin_allows_auth_pass_through_for_jwt_team_allowed_routes_wildcard():
jwt_token: Final = UserAPIKeyAuth(
user_id="test_user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
team_id="team-a",
jwt_claims={"sub": "test_user"},
)
_check_model_host_route_as(jwt_token, team_allowed_routes=["openai_routes", "/model-host/*"])
def test_non_proxy_admin_denies_auth_pass_through_for_jwt_when_only_route_groups_configured():
jwt_token: Final = UserAPIKeyAuth(
user_id="test_user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
team_id="team-a",
jwt_claims={"sub": "test_user"},
)
with pytest.raises(HTTPException) as exc_info:
_check_model_host_route_as(jwt_token, team_allowed_routes=["openai_routes", "mapped_pass_through_routes"])
assert exc_info.value.status_code == 403, exc_info.value.detail
assert "allowed_passthrough_routes" in exc_info.value.detail
@pytest.mark.parametrize(
"api_key, team_id, jwt_claims",
[
("sk-test-key", "team-a", None),
("sk-test-key", "team-a", {"sub": "test_user"}),
(None, "team-a", None),
(None, None, {"sub": "test_user"}),
],
ids=["plain_virtual_key", "jwt_mapped_virtual_key", "keyless_non_jwt_caller", "jwt_without_team"],
)
def test_non_proxy_admin_jwt_team_allowed_routes_grant_pass_through_only_to_jwt_team_callers(
api_key, team_id, jwt_claims
):
caller: Final = UserAPIKeyAuth(
api_key=api_key,
user_id="test_user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
team_id=team_id,
jwt_claims=jwt_claims,
)
with pytest.raises(HTTPException) as exc_info:
_check_model_host_route_as(caller, team_allowed_routes=["openai_routes", "/model-host/*"])
assert exc_info.value.status_code == 403, exc_info.value.detail
assert "allowed_passthrough_routes" in exc_info.value.detail
def test_virtual_key_without_llm_api_routes_cannot_access_pass_through():
"""
Test that virtual keys without llm_api_routes permission cannot access registered pass-through endpoints.

56
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -6,7 +6,7 @@ import logging
from contextlib import ExitStack
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import List, Optional, cast
from typing import Final, List, Optional, cast
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -1614,6 +1614,60 @@ class TestTeamScopedMCPServerAccess:
result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id")
assert len(result) == 1
class TestFetchAllMCPServersOrdering:
def test_display_order_is_case_insensitive_name_then_id(self) -> None:
servers: Final = (
LiteLLM_MCPServerTable(server_id="s-2", server_name="GitHub", alias="aaa", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="s-1", alias="github", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="s-0", server_name="Slack", alias="zzz", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="confluence", server_name="", alias="", transport=MCPTransport.http),
)
ordered: Final = sorted(servers, key=mgmt_endpoints._mcp_server_display_order)
assert [s.server_id for s in ordered] == ["confluence", "s-1", "s-2", "s-0"]
@pytest.mark.parametrize("team_id", [None, "team-1"])
@pytest.mark.parametrize("reverse", [False, True])
@pytest.mark.asyncio
async def test_list_is_sorted_by_display_name_regardless_of_resolution_order(
self, team_id: str | None, reverse: bool
) -> None:
mock_user_auth: Final = generate_mock_user_api_key_auth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin_user",
)
servers: Final = (
generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"),
generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"),
generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"),
)
resolved: Final = list(reversed(servers) if reverse else servers)
mock_manager: Final = MagicMock()
mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved)
with (
patch( # test-quality-ok: the route reads a module-global manager with no injection seam
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch( # test-quality-ok: admin view is derived from module-global proxy settings
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
return_value=True,
),
patch( # test-quality-ok: auth contexts need a live prisma client
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
AsyncMock(return_value=[mock_user_auth]),
),
patch( # test-quality-ok: isolate the route's ordering from team database resolution
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list",
AsyncMock(return_value=resolved),
),
):
result: Final = await mgmt_endpoints.fetch_all_mcp_servers(
user_api_key_dict=mock_user_auth, team_id=team_id
)
assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"]
@pytest.mark.asyncio
async def test_restricted_virtual_key_cannot_use_team_id_filter(self):
"""Restricted virtual keys must not bypass access limits via team_id."""

112
tests/test_litellm_rust/support/s3_stub.py Normal file
View file

@ -0,0 +1,112 @@
"""In-process path-style S3 stub for native cache parity tests."""
import threading
from dataclasses import dataclass, field
from email.utils import parsedate_to_datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
from urllib.parse import unquote, urlsplit
_STORED_HEADERS: Final = (
"cache-control",
"content-type",
"content-language",
"content-disposition",
"expires",
)
@dataclass
class S3Object:
body: bytes
headers: dict[str, str] = field(default_factory=dict)
class S3Stub:
"""Minimal path-style S3 endpoint serving PUT and GET object operations."""
def __init__(self) -> None:
self._objects: dict[str, S3Object] = {}
stub: Final = self
class Handler(BaseHTTPRequestHandler):
def _key(self) -> str:
parts: Final = urlsplit(self.path).path.lstrip("/").split("/", 1)
return unquote(parts[1]) if len(parts) == 2 else ""
def _read_body(self) -> bytes:
transfer: Final = self.headers.get("transfer-encoding", "")
if "chunked" not in transfer:
return self.rfile.read(int(self.headers.get("content-length", 0)))
chunks: Final = bytearray()
while True:
size = int(self.rfile.readline().split(b";")[0].strip(), 16)
if size == 0:
while self.rfile.readline().strip():
pass
return bytes(chunks)
chunks.extend(self.rfile.read(size))
self.rfile.readline()
def do_PUT(self) -> None:
body: Final = self._read_body()
headers: Final = {name: self.headers[name] for name in _STORED_HEADERS if name in self.headers}
stub._objects = {**stub._objects, self._key(): S3Object(body=body, headers=headers)}
self.send_response(200)
self.send_header("ETag", '"stub"')
self.send_header("Content-Length", "0")
self.end_headers()
def do_HEAD(self) -> None:
self._object(send_body=False)
def do_GET(self) -> None:
self._object(send_body=True)
def _object(self, send_body: bool) -> None:
entry: Final = stub._objects.get(self._key())
if entry is None:
self.send_response(404)
self.send_header("Content-Type", "application/xml")
body: Final = b'<?xml version="1.0" encoding="UTF-8"?><Error><Code>NoSuchKey</Code></Error>'
self.send_header("Content-Length", str(len(body)))
self.end_headers()
if send_body:
self.wfile.write(body)
return
self.send_response(200)
for name, value in entry.headers.items():
self.send_header(name, value)
self.send_header("ETag", '"stub"')
self.send_header("Content-Length", str(len(entry.body)))
self.end_headers()
if send_body:
self.wfile.write(entry.body)
def log_message(self, format: str, *args: object) -> None:
pass
self._server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
self._worker: Final = threading.Thread(target=self._server.serve_forever, daemon=True)
self._worker.start()
@property
def url(self) -> str:
host, port = self._server.server_address[:2]
return f"http://{host}:{port}"
@property
def objects(self) -> dict[str, S3Object]:
return self._objects
def put_object(self, key: str, body: bytes, headers: dict[str, str] | None = None) -> None:
self._objects = {**self._objects, key: S3Object(body=body, headers=headers or {})}
def expires(self, key: str) -> object:
header: Final = self._objects[key].headers.get("expires")
return parsedate_to_datetime(header) if header else None
def close(self) -> None:
self._server.shutdown()
self._server.server_close()
self._worker.join(timeout=5)

209
tests/test_litellm_rust/test_cache.py
View file

@ -11,12 +11,16 @@ import uuid
import weakref
from collections.abc import Callable, Generator
from contextlib import ExitStack
from datetime import datetime
from pathlib import Path
from types import SimpleNamespace
from typing import Final, Protocol, cast
from unittest.mock import Mock
from urllib.parse import urlparse
from uuid import uuid4
import boto3
import botocore.config
import diskcache
import fakeredis
import pytest
@ -26,17 +30,19 @@ from azure.storage.blob import ContainerClient
import litellm
from litellm.caching.azure_blob_cache import AzureBlobCache
from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache
from litellm.caching.gcs_cache import GCSCache
from litellm.caching.disk_cache import DiskCache
from litellm.caching.gcs_cache import GCSCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.caching.redis_semantic_cache import RedisSemanticCache
from litellm.caching.s3_cache import S3Cache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
from litellm.types.llms.custom_llm import CustomLLMItem
from litellm.types.utils import EmbeddingResponse
from tests.test_litellm_rust.support.fake_gcs import FakeGcs
from tests.test_litellm_rust.support.isolation import rebound
from tests.test_litellm_rust.support.s3_stub import S3Stub
_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name
_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name
@ -436,15 +442,20 @@ def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure
assert handle.backend == "azure-blob"
account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}")
with pytest.raises(TypeError, match="containers must match"):
_native._CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade(
azure_blob_facade
)
_native._CacheTestHandle.azure_blob(
account_url, f"{backend.container_client.container_name}-other"
)._bind_facade(azure_blob_facade)
handle._bind_facade(azure_blob_facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade))
native: Final = resolver.resolve()
assert native.kind == "native"
response: Final = {"choices": [{"text": "caf\u00e9 \u2603"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
response: Final = {
"choices": [{"text": "caf\u00e9 \u2603"}],
"usage": {"total_tokens": 3},
"flag": True,
"empty": None,
}
native.store({**request("sync"), "ttl_seconds": 0.001}, response)
native.store(request("sync"), {"choices": [{"text": "second"}]})
time.sleep(0.01)
@ -499,7 +510,9 @@ async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_pyt
await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2})
time.sleep(0.01)
assert await binding.async_lookup(request("async")) == {"value": 2}
assert await backend.async_get_cache("async") == json.loads(backend.container_client.download_blob("async").readall())
assert await backend.async_get_cache("async") == json.loads(
backend.container_client.download_blob("async").readall()
)
assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2}
await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}])
@ -544,6 +557,8 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None:
assert client.get("second") is not None
await facade.cache.disconnect()
client.close()
async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None:
disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path))
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}}
@ -652,6 +667,183 @@ async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path:
}
@pytest.fixture
def s3_stub() -> Generator[S3Stub]:
stub: Final = S3Stub()
try:
yield stub
finally:
stub.close()
def python_s3(url: str) -> S3Cache:
return S3Cache(
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None:
python_cache: Final = python_s3(s3_stub.url)
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}}
python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90)
python_cache.set_cache("plain", {"timestamp": time.time(), "response": response})
s3_stub.put_object("team/malformed", b"not a cache entry")
s3_stub.put_object(
"team/expired",
json.dumps({"timestamp": time.time(), "response": response}).encode(),
{"expires": "Thu, 01 Jan 1970 00:00:00 GMT"},
)
binding: Final = _native._CacheTestResolver(
SimpleNamespace(
cache=_native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
)
).resolve()
assert binding.lookup(request("sync:key")) == response
assert await binding.async_lookup(request("plain")) == response
assert binding.lookup(request("malformed")) is None
assert binding.lookup(request("expired")) is None
assert binding.lookup(request("absent")) is None
binding.store({**request("native:key"), "ttl_seconds": 90.0}, response)
await binding.async_store(request("no_ttl"), response)
stored: Final = s3_stub.objects["team/native/key"]
assert stored.headers["content-type"] == "application/json"
assert stored.headers["content-language"] == "en"
assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"'
assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90"
expires: Final = cast(datetime, s3_stub.expires("team/native/key"))
remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds()
assert 60 < remaining <= 91
no_ttl: Final = s3_stub.objects["team/no_ttl"]
assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000"
assert "expires" not in no_ttl.headers
assert python_cache.get_cache("native:key")["response"] == response
partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")])
assert partial == {"values": [response, None, None], "missing_indices": [1, 2]}
def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None:
facade: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
handle: Final = _native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
with pytest.raises(TypeError, match="buckets must match"):
_native._CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade)
with pytest.raises(TypeError, match="key prefixes must match"):
_native._CacheTestHandle.s3(
"cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/"
)._bind_facade(facade)
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
binding: Final = resolver.resolve()
assert binding.kind == "native"
handler: Final = Mock()
facade.cache.s3_client.meta.events.register("before-call.s3.*", handler)
binding.store(request("native"), {"answer": 1})
assert binding.lookup(request("native")) == {"answer": 1}
assert handler.call_count == 0
assert "team/native" in s3_stub.objects
with rebound(facade.cache, "bucket_name", "other"):
assert resolver.resolve().kind == "python_callback"
other_client: Final = boto3.client(
"s3",
region_name="us-east-1",
endpoint_url=s3_stub.url,
aws_access_key_id="key",
aws_secret_access_key="secret",
)
with rebound(facade.cache, "s3_client", other_client):
assert resolver.resolve().kind == "python_callback"
class CustomS3Cache(S3Cache):
pass
subclassed: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
subclassed.cache = CustomS3Cache(
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
with pytest.raises(TypeError):
handle._bind_facade(subclassed)
assert _native._CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback"
def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None:
handle: Final = _native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
unverified: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url="https://s3.example.test",
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
s3_verify=False,
)
with pytest.raises(TypeError, match="requires Python"):
handle._bind_facade(unverified)
proxied: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}),
)
with pytest.raises(TypeError, match="requires Python"):
handle._bind_facade(proxied)
async def test_gcs_reads_python_entries_and_writes_python_compatible_objects(
fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch
) -> None:
@ -787,6 +979,7 @@ async def test_gcs_facade_binds_only_exact_matching_configuration(
assert resolver.resolve().kind == "python_callback"
with rebound(facade.cache, "path_service_account", "sa.json"):
assert resolver.resolve().kind == "python_callback"
def no_get_cache(*args: object, **kwargs: object) -> None:
return None
@ -921,7 +1114,9 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n
await binding.async_flush()
remaining: Final = tuple(sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)))
remaining: Final = tuple(
sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node))
)
assert remaining == (), remaining
assert client.get("unscoped") == b"stays"
client.delete("unscoped")

599
tests/test_litellm_rust/test_valkey_semantic_cache_native.py Normal file
View file

@ -0,0 +1,599 @@
import asyncio
import contextvars
import hashlib
import os
import struct
import threading
import time
from collections.abc import Generator, Mapping
from types import SimpleNamespace
from typing import Final, cast
from uuid import uuid4
import pytest
import redis
from litellm.caching.caching import Cache
from litellm.caching.valkey_semantic_cache import ValkeySemanticCache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
pytestmark: Final = pytest.mark.requires_rust_extension
embedding_context: Final = contextvars.ContextVar("embedding_context")
@pytest.fixture
def valkey_url() -> str:
url: Final = os.environ.get("LITELLM_TEST_VALKEY_URL")
if url is None:
pytest.skip("LITELLM_TEST_VALKEY_URL is not set")
return url
@pytest.fixture
def index_name(valkey_url: str) -> Generator[str]:
index: Final = f"litellm_test_{uuid4().hex}"
yield index
client: Final = redis.Redis.from_url(valkey_url)
try:
client.ft(index).dropindex(delete_documents=True)
except redis.ResponseError:
pass
finally:
client.close()
def _request(prompt: str = "semantic cache prompt") -> dict[str, object]:
return {
"key": {"preset": "key"},
"messages": [{"role": "user", "content": prompt}],
}
def _field_request(
prompt: str,
metadata: Mapping[str, object],
*,
namespace: str | None = None,
litellm_metadata: Mapping[str, object] | None = None,
litellm_params: Mapping[str, object] | None = None,
) -> dict[str, object]:
request: Final = {
"key": {
"fields": [
{
"name": "model",
"value": "gpt-4.1",
"api_parameter": True,
"internal_parameter": False,
},
{
"name": "messages",
"value": prompt,
"api_parameter": True,
"internal_parameter": False,
},
],
"namespace": namespace,
},
"messages": [{"role": "user", "content": prompt}],
"metadata": dict(metadata),
}
if litellm_metadata is not None:
request["litellm_metadata"] = dict(litellm_metadata)
if litellm_params is not None:
request["litellm_params"] = dict(litellm_params)
return request
def _facade(
url: str,
index_name: str,
embeddings: Mapping[str, list[float]],
*,
namespace: str | None = None,
) -> Cache:
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url=url,
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
namespace=namespace,
)
vectors: Final = embeddings
def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]:
return vectors[prompt]
async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
return vectors[prompt]
facade.cache._get_embedding = embed
facade.cache._get_async_embedding = async_embedding
return facade
def _backend(
url: str,
index_name: str,
embeddings: Mapping[str, list[float]] | None = None,
) -> ValkeySemanticCache:
vectors: Final = embeddings or {"semantic cache prompt": [1.0, 0.0]}
backend: Final = ValkeySemanticCache(
redis_url=url,
similarity_threshold=0.8,
index_name=index_name,
)
def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]:
return vectors[prompt]
async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
return vectors[prompt]
backend._get_embedding = embed
backend._get_async_embedding = async_embedding
return backend
def test_python_write_native_read(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
response: Final = {"answer": "python"}
backend.set_cache("key", response, messages=_request()["messages"])
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
assert binding.lookup(_request()) == response
def test_native_write_python_read(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
response: Final = {"answer": "native"}
binding.store({**_request(), "ttl_seconds": 2.0}, response)
cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"]))
assert cached["response"] == response
async def test_async_lookup_and_store(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
request: Final = {**_request(), "ttl_seconds": 2.0}
await binding.async_store(request, {"answer": "async"})
assert await binding.async_lookup(request) == {"answer": "async"}
async def test_disabled_cache_controls_skip_async_embedding(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
calls: Final = []
async def fail_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
calls.append(prompt)
raise AssertionError("embedding must not run")
backend._get_async_embedding = fail_embedding
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
controls: Final = {
"supported_call_type": True,
"configured": True,
"native_backend": True,
"default_on": True,
"caching": True,
"no_cache": False,
"no_store": False,
"use_cache": True,
}
no_read_request: Final = {**_request(), "controls": {**controls, "no_cache": True}}
assert await binding.async_lookup(no_read_request) is None
no_write_request: Final = {**_request(), "controls": {**controls, "no_store": True}}
await binding.async_store(no_write_request, {"answer": "blocked"})
assert calls == []
client: Final = redis.Redis.from_url(valkey_url)
assert list(client.scan_iter(f"{index_name}:*")) == []
client.close()
async def test_async_embedding_runs_inline_in_caller_task(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
observed: dict[str, object] = {}
async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
observed["context"] = embedding_context.get("missing")
observed["task"] = asyncio.current_task()
observed["thread"] = threading.get_ident()
embedding_context.set("embedder")
return [1.0, 0.0]
backend._get_async_embedding = async_embedding
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
request: Final = {**_request(), "ttl_seconds": 2.0}
caller_task: Final = asyncio.current_task()
caller_thread: Final = threading.get_ident()
token: Final = embedding_context.set("caller")
try:
await binding.async_store(request, {"answer": "inline"})
assert observed["context"] == "caller"
assert observed["task"] is caller_task
assert observed["thread"] == caller_thread
assert embedding_context.get() == "embedder"
assert await binding.async_lookup(request) == {"answer": "inline"}
finally:
embedding_context.reset(token)
def test_facade_activation_and_mutation_fallback(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url=valkey_url,
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
assert resolver.resolve().kind == "native"
facade.cache.similarity_threshold = 0.7
assert resolver.resolve().kind == "python_callback"
def test_batch_lookup_is_unsupported(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
with pytest.raises(NotImplementedError):
binding.lookup_batch([_request()])
def test_ttl_expiry(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store({**_request(), "ttl_seconds": 1.0}, {"answer": "expires"})
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
assert client.ttl(documents[0]) > 0
time.sleep(1.5)
assert binding.lookup(_request()) is None
def test_no_ttl_is_persistent_and_python_reads_native_value(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
response: Final = {"answer": "persistent"}
binding.store(_request(), response)
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
assert client.ttl(documents[0]) == -1
cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"]))
assert cached["response"] == response
def test_below_threshold_misses_on_native_and_python(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(
valkey_url,
index_name,
{"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]},
)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(_request("prompt A"), {"answer": "A"})
assert binding.lookup(_request("prompt B")) is None
assert backend.get_cache("key", messages=_request("prompt B")["messages"]) is None
def test_malformed_entry_is_a_miss_on_native_and_python(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
client: Final = redis.Redis.from_url(valkey_url)
scope: Final = hashlib.sha256(b"key").hexdigest()
document: Final = f"{index_name}:{scope}:{uuid4().hex}"
client.hset(
document,
mapping={
"litellm_cache_key": scope,
"prompt": "semantic cache prompt",
"response": "not json",
"embedding": struct.pack("<2f", 1.0, 0.0),
},
)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
assert binding.lookup(_request()) is None
assert backend.get_cache("key", messages=_request()["messages"]) is None
def test_mixed_content_parts_match_python_semantic_behavior(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
messages: Final = [{"role": "user", "content": ["raw", {"text": "hello"}]}]
backend.set_cache("key", {"answer": "mixed"}, messages=messages)
assert backend.get_cache("key", messages=messages) is None
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
request: Final = {**_request(), "messages": messages}
binding.store(request, {"answer": "mixed"})
assert binding.lookup(request) is None
client: Final = redis.Redis.from_url(valkey_url)
assert list(client.scan_iter(f"{index_name}:*")) == []
client.close()
async def test_async_store_batch_and_lookup(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(
valkey_url,
index_name,
{"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]},
)
sync_calls: Final = []
async_tasks: Final = []
def sync_embedding(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]:
sync_calls.append(prompt)
return {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}[prompt]
async def async_embedding(
prompt: str,
metadata: dict[str, object] | None = None,
) -> list[float]:
async_tasks.append(asyncio.current_task())
return {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}[prompt]
backend._get_embedding = sync_embedding
backend._get_async_embedding = async_embedding
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
requests: Final = [_request("prompt A"), _request("prompt B")]
responses: Final = [{"answer": "A"}, {"answer": "B"}]
caller_task: Final = asyncio.current_task()
await binding.async_store_batch(requests, responses)
assert sync_calls == []
assert async_tasks
assert all(task is caller_task for task in async_tasks)
assert await binding.async_lookup(requests[0]) == responses[0]
assert await binding.async_lookup(requests[1]) == responses[1]
def test_subclass_backend_falls_back_to_python(
valkey_url: str,
index_name: str,
) -> None:
class Custom(ValkeySemanticCache):
pass
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url=valkey_url,
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
)
facade.cache = Custom(redis_url=valkey_url, similarity_threshold=0.8, index_name=index_name)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
assert resolver.resolve().kind == "python_callback"
def test_field_key_matches_python_semantic_scope(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]})
metadata: Final = {"user_api_key": "k1"}
expected: Final = facade.get_cache_key(
model="gpt-4.1",
messages=[{"role": "user", "content": "semantic cache prompt"}],
metadata=metadata,
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(_field_request("semantic cache prompt", metadata), {"answer": "scoped"})
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
document_parts: Final = documents[0].decode().split(":")
assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest()
client.close()
def test_field_key_reads_all_python_tenant_metadata_sources(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]})
params_metadata: Final = {"user_api_key_team_id": "team-from-params"}
expected: Final = facade.get_cache_key(
model="gpt-4.1",
messages=[{"role": "user", "content": "semantic cache prompt"}],
metadata={},
litellm_params={"metadata": params_metadata},
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(
_field_request(
"semantic cache prompt",
{},
litellm_params={"metadata": params_metadata},
),
{"answer": "params"},
)
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
document_parts: Final = documents[0].decode().split(":")
assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest()
client.close()
assert (
binding.lookup(
_field_request(
"semantic cache prompt",
{},
litellm_metadata={"user_api_key_team_id": "team-from-litellm"},
)
)
is None
)
def test_namespace_isolates_semantic_entries(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(
valkey_url,
index_name,
{"semantic cache prompt": [1.0, 0.0]},
namespace="team-a",
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
team_a: Final = _field_request("semantic cache prompt", {}, namespace="team-a")
team_b: Final = _field_request("semantic cache prompt", {}, namespace="team-b")
binding.store(team_a, {"answer": "team-a"})
assert binding.lookup(team_b) is None
assert binding.lookup(team_a) == {"answer": "team-a"}
cached: Final = cast(
Mapping[str, object],
facade.get_cache(
model="gpt-4.1",
messages=[{"role": "user", "content": "semantic cache prompt"}],
),
)
assert cached == {"answer": "team-a"}
def test_field_key_isolates_tenant_scope(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]})
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(
_field_request("semantic cache prompt", {"user_api_key": "k1"}),
{"answer": "tenant one"},
)
assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k2"})) is None
assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k1"})) == {"answer": "tenant one"}
def test_tls_valkey_facade_falls_back_to_python(
index_name: str,
) -> None:
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url="rediss://127.0.0.1:6390/0",
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
assert resolver.resolve().kind == "python_callback"
async def test_ping_maps_unsupported_native_operation_to_not_implemented(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
with pytest.raises(NotImplementedError):
await binding.ping()

227
ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx
View file

@ -3,7 +3,8 @@ import { render, waitFor, screen, act, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { describe, it, expect, vi, beforeEach } from "vitest";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import MCPServers from "./mcp_servers";
import MCPServers, { compareServers, type SortKey } from "./mcp_servers";
import type { MCPServer } from "@/components/mcp_tools/types";
import * as networking from "@/components/networking";
// Mock the networking module
@ -31,6 +32,98 @@ const createQueryClient = () =>
},
});
describe("compareServers", () => {
const server = (server_id: string, name: string, created_at = ""): MCPServer => ({
server_id,
server_name: name,
created_at,
updated_at: created_at,
created_by: "user",
updated_by: "user",
});
const shuffled = [server("c", "github"), server("a", "slack"), server("b", "Jira")];
it("orders servers without timestamps by name so config.yaml servers render in a stable order", () => {
const byCreated = [...shuffled].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id);
const byUpdated = [...shuffled].sort((a, b) => compareServers(a, b, "updated_desc")).map((s) => s.server_id);
const byHealth = [...shuffled].sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id);
expect(byCreated).toEqual(["c", "b", "a"]);
expect(byUpdated).toEqual(["c", "b", "a"]);
expect(byHealth).toEqual(["c", "b", "a"]);
});
it("keeps newest-first when timestamps differ", () => {
const newest = server("new", "zzz", "2026-02-01T00:00:00Z");
const oldest = server("old", "aaa", "2026-01-01T00:00:00Z");
expect([oldest, newest].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id)).toEqual([
"new",
"old",
]);
});
it.each<SortKey>(["created_desc", "updated_desc", "name_asc", "health"])(
"breaks equal timestamps and names by ID for %s regardless of input order",
(sort) => {
const servers = [
server("b", "GitHub", "2026-01-01T00:00:00Z"),
server("c", "Slack", "2026-01-01T00:00:00Z"),
server("a", "github", "2026-01-01T00:00:00Z"),
];
for (const input of [servers, [...servers].reverse()]) {
expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual(["a", "b", "c"]);
}
},
);
it("uses the display name before alias, then falls back to alias and ID", () => {
const servers: MCPServer[] = [
{ ...server("s-slack", "Slack"), alias: "aaa" },
{ ...server("s-github", ""), server_name: null, alias: "GitHub" },
{ ...server("confluence", ""), alias: "" },
];
for (const input of [servers, [...servers].reverse()]) {
expect([...input].sort((a, b) => compareServers(a, b, "name_asc")).map((s) => s.server_id)).toEqual([
"confluence",
"s-github",
"s-slack",
]);
}
});
it.each<SortKey>(["created_desc", "updated_desc", "health"])(
"keeps timestamped servers before missing timestamps for %s",
(sort) => {
const servers = [
server("config", "aaa"),
server("older", "bbb", "2026-01-01T00:00:00Z"),
server("newer", "zzz", "2026-02-01T00:00:00Z"),
];
for (const input of [servers, [...servers].reverse()]) {
expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual([
"newer",
"older",
"config",
]);
}
},
);
it("sorts health before recency and display name", () => {
const servers: MCPServer[] = [
{ ...server("healthy", "aaa", "2026-03-01T00:00:00Z"), status: "healthy" },
{ ...server("unknown", "bbb", "2026-02-01T00:00:00Z"), status: "unknown" },
{ ...server("unhealthy", "zzz", "2026-01-01T00:00:00Z"), status: "unhealthy" },
];
expect(servers.sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id)).toEqual([
"unhealthy",
"unknown",
"healthy",
]);
});
});
describe("MCPServers", () => {
const defaultProps = {
accessToken: "123",
@ -74,6 +167,134 @@ describe("MCPServers", () => {
const myConnections = await screen.findByRole("link", { name: "My Connections" });
expect(myConnections).toBeVisible();
expect(myConnections).toHaveAttribute("href", "/ui/connect");
for (const name of ["Semantic Filter", "Tool Search", "Network Settings", "Submitted MCPs"]) {
const tab = screen.queryByRole("tab", { name });
if (userRole === "Admin") {
expect(tab).toBeVisible();
} else {
expect(tab).not.toBeInTheDocument();
}
}
expect(
screen.getByRole("button", {
name: userRole === "Admin" ? "+ Add New MCP Server" : "+ Submit MCP Server",
}),
).toBeVisible();
});
it.each(["cancel", "success", "failure", "unnamed"])("preserves delete confirmation on %s", async (outcome) => {
const server: MCPServer = {
created_at: "",
updated_at: "",
server_id: "delete-server",
server_name: outcome === "unnamed" ? null : "Delete fixture",
alias: "delete-alias",
url: outcome === "unnamed" ? null : "https://example.com/mcp",
created_by: "user",
updated_by: "user",
};
vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]);
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
let finishDelete: () => void = () => {};
vi.mocked(networking.deleteMCPServer).mockImplementation(
() =>
new Promise((resolve, reject) => {
finishDelete = () => (outcome === "failure" ? reject(new Error("Delete failed")) : resolve(undefined));
}),
);
render(
<QueryClientProvider client={createQueryClient()}>
<MCPServers {...defaultProps} />
</QueryClientProvider>,
);
await userEvent.click(await screen.findByRole("button", { name: "Server actions" }));
await userEvent.click(await screen.findByRole("menuitem", { name: "Delete" }));
const dialog = await screen.findByRole("alertdialog", { name: "Delete MCP Server?" });
expect(within(dialog).getByText("delete-server")).toBeVisible();
if (outcome === "unnamed") {
expect(within(dialog).queryByText("Name")).not.toBeInTheDocument();
expect(within(dialog).queryByText("URL")).not.toBeInTheDocument();
} else {
expect(within(dialog).getByText("Delete fixture")).toBeVisible();
expect(within(dialog).getByText("https://example.com/mcp")).toBeVisible();
}
if (outcome === "cancel") {
await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" }));
expect(networking.deleteMCPServer).not.toHaveBeenCalled();
} else {
await userEvent.click(within(dialog).getByRole("button", { name: "Delete" }));
expect(within(dialog).getByRole("button", { name: "Deleting..." })).toBeDisabled();
expect(within(dialog).getByRole("button", { name: "Cancel" })).toBeDisabled();
expect(networking.deleteMCPServer).toHaveBeenCalledWith("123", "delete-server");
await act(async () => finishDelete());
}
await waitFor(() => expect(screen.queryByRole("alertdialog")).not.toBeInTheDocument());
});
it("filters servers by access group", async () => {
const server = { created_by: "user", updated_by: "user" };
vi.mocked(networking.fetchMCPServers).mockResolvedValue([
{
...server,
server_id: "string-group",
server_name: "String group",
alias: "string-alias",
mcp_access_groups: ["shared"],
},
{
...server,
server_id: "legacy-group",
server_name: "Legacy group",
alias: "legacy-alias",
mcp_access_groups: ["shared"],
},
{
...server,
server_id: "other-group",
server_name: "Other group",
alias: "other-alias",
mcp_access_groups: ["different"],
},
]);
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
render(
<QueryClientProvider client={createQueryClient()}>
<MCPServers {...defaultProps} />
</QueryClientProvider>,
);
await screen.findByText("String group");
await userEvent.click(screen.getByRole("combobox", { name: "Access Group" }));
await userEvent.click(await screen.findByRole("option", { name: "shared" }));
expect(screen.getByText("String group")).toBeVisible();
expect(screen.getByText("Legacy group")).toBeVisible();
expect(screen.queryByText("Other group")).not.toBeInTheDocument();
});
it.each(["server_name", "alias", "url", "server_id"] as const)("searches by %s case-insensitively", async (field) => {
const server: MCPServer = {
created_at: "",
updated_at: "",
server_id: "search-server",
server_name: "Search fixture",
created_by: "user",
updated_by: "user",
[field]: "Needle",
};
vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]);
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
render(
<QueryClientProvider client={createQueryClient()}>
<MCPServers {...defaultProps} />
</QueryClientProvider>,
);
await screen.findByTestId("mcp-servers-grid");
const search = screen.getByPlaceholderText("Search by name, alias, URL, or ID");
await userEvent.type(search, " NEEDLE ");
expect(screen.getByTestId("mcp-servers-grid")).toBeVisible();
await userEvent.clear(search);
await userEvent.type(search, "no-match");
expect(screen.queryByTestId("mcp-servers-grid")).not.toBeInTheDocument();
expect(screen.getByText("No servers match the current filters or search.")).toBeVisible();
});
it("should render mocked MCP servers data in the table", async () => {
@ -316,9 +537,7 @@ describe("MCPServers", () => {
expect(screen.getByText("Team B Server")).toBeInTheDocument();
expect(screen.getByText("Team A Server 2")).toBeInTheDocument();
// Find the team select by its "Team" label, then the combobox it labels
const teamLabel = screen.getByText("Team");
const teamSelect = within(teamLabel.parentElement!).getByRole("combobox");
const teamSelect = screen.getByRole("combobox", { name: "Team" });
await userEvent.click(teamSelect);

244
ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx
View file

@ -49,7 +49,7 @@ import { cn } from "@/lib/cva.config";
import UserEnvVarsModal from "./UserEnvVarsModal";
import { listMCPUserEnvVarStatus } from "@/components/networking";
type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health";
export type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health";
const SORT_OPTIONS: { value: SortKey; label: string }[] = [
{ value: "created_desc", label: "Recently created" },
@ -64,32 +64,33 @@ const HEALTH_RANK: Record<string, number> = {
healthy: 2,
};
const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => {
const compareByName = (a: MCPServer, b: MCPServer): number => {
const nameA = (a.server_name || a.alias || a.server_id).toLowerCase();
const nameB = (b.server_name || b.alias || b.server_id).toLowerCase();
return nameA.localeCompare(nameB) || a.server_id.localeCompare(b.server_id);
};
const compareByTimestampDesc = (a: string | null | undefined, b: string | null | undefined): number => {
const ta = a ? new Date(a).getTime() : 0;
const tb = b ? new Date(b).getTime() : 0;
return tb - ta;
};
export const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => {
switch (sort) {
case "name_asc": {
const nameA = (a.server_name || a.alias || a.server_id).toLowerCase();
const nameB = (b.server_name || b.alias || b.server_id).toLowerCase();
return nameA.localeCompare(nameB);
}
case "updated_desc": {
const ta = a.updated_at ? new Date(a.updated_at).getTime() : 0;
const tb = b.updated_at ? new Date(b.updated_at).getTime() : 0;
return tb - ta;
}
case "name_asc":
return compareByName(a, b);
case "updated_desc":
return compareByTimestampDesc(a.updated_at, b.updated_at) || compareByName(a, b);
case "health": {
const ra = HEALTH_RANK[a.status ?? "unknown"] ?? 1;
const rb = HEALTH_RANK[b.status ?? "unknown"] ?? 1;
if (ra !== rb) return ra - rb;
const ta = a.created_at ? new Date(a.created_at).getTime() : 0;
const tb = b.created_at ? new Date(b.created_at).getTime() : 0;
return tb - ta;
return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b);
}
case "created_desc":
default: {
const ta = a.created_at ? new Date(a.created_at).getTime() : 0;
const tb = b.created_at ? new Date(b.created_at).getTime() : 0;
return tb - ta;
}
default:
return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b);
}
};
@ -112,6 +113,62 @@ const readToolsOAuthServerId = (): string | null => {
}
};
function DeleteServerDialog({
open,
onOpenChange,
server,
isDeleting,
onConfirm,
}: {
open: boolean;
onOpenChange: (open: boolean) => void;
server: MCPServer | undefined;
isDeleting: boolean;
onConfirm: () => Promise<void>;
}) {
return (
<AlertDialog open={open} onOpenChange={onOpenChange}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Delete MCP Server?</AlertDialogTitle>
</AlertDialogHeader>
<div className="space-y-4">
<p className="text-sm text-muted-foreground">
This action is permanent and cannot be undone. All associated configurations will be removed.
</p>
{server && (
<dl className="mt-3 space-y-1 rounded-lg border border-border bg-muted p-4">
{server.server_name && (
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">Name</dt>
<dd className="text-sm font-semibold">{server.server_name}</dd>
</div>
)}
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">ID</dt>
<dd className="font-mono text-xs">{server.server_id}</dd>
</div>
{server.url && (
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">URL</dt>
<dd className="font-mono text-xs break-all">{server.url}</dd>
</div>
)}
</dl>
)}
</div>
<AlertDialogFooter>
<AlertDialogCancel disabled={isDeleting}>Cancel</AlertDialogCancel>
<Button variant="destructive" disabled={isDeleting} onClick={onConfirm}>
{isDeleting ? "Deleting..." : "Delete"}
</Button>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
);
}
const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, isViewOnly = false }) => {
const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers();
@ -298,16 +355,12 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
}
if (group !== "all") {
filtered = filtered.filter((server) =>
server.mcp_access_groups?.some((g: any) => (typeof g === "string" ? g === group : g && g.name === group)),
server.mcp_access_groups?.some((g: string | { name?: string } | null) =>
typeof g === "string" ? g === group : g?.name === group,
),
);
}
const sorted = [...filtered].sort((a, b) => {
if (!a.created_at && !b.created_at) return 0;
if (!a.created_at) return 1;
if (!b.created_at) return -1;
return new Date(b.created_at).getTime() - new Date(a.created_at).getTime();
});
setFilteredServers(sorted);
setFilteredServers(filtered);
},
[serversWithHealth],
);
@ -338,7 +391,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
const alias = (s.alias || "").toLowerCase();
const url = (s.url || "").toLowerCase();
const id = s.server_id.toLowerCase();
return name.includes(q) || alias.includes(q) || url.includes(q) || id.includes(q);
return [name, alias, url, id].some((value) => value.includes(q));
})
: filteredServers;
return [...matches].sort((a, b) => compareServers(a, b, sortKey));
@ -381,9 +434,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
};
// Find the server to delete from the servers list
const serverToDelete = serverIdToDelete
? (mcpServers || []).find((server) => server.server_id === serverIdToDelete)
: null;
const serverToDelete = mcpServers?.find((server) => server.server_id === serverIdToDelete);
const handleCreateSuccess = (newMcpServer: MCPServer) => {
setFilteredServers((prev) => [...prev, newMcpServer]);
@ -425,45 +476,13 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
return (
<TooltipProvider>
<div className="h-full w-full p-6">
<AlertDialog open={isDeleteModalOpen} onOpenChange={(open) => !open && cancelDelete()}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Delete MCP Server?</AlertDialogTitle>
</AlertDialogHeader>
<div className="space-y-4">
<p className="text-sm text-muted-foreground">
This action is permanent and cannot be undone. All associated configurations will be removed.
</p>
{serverToDelete && (
<dl className="mt-3 space-y-1 rounded-lg border border-border bg-muted p-4">
{serverToDelete.server_name && (
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">Name</dt>
<dd className="text-sm font-semibold">{serverToDelete.server_name}</dd>
</div>
)}
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">ID</dt>
<dd className="font-mono text-xs">{serverToDelete.server_id}</dd>
</div>
{serverToDelete.url && (
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">URL</dt>
<dd className="font-mono text-xs break-all">{serverToDelete.url}</dd>
</div>
)}
</dl>
)}
</div>
<AlertDialogFooter>
<AlertDialogCancel disabled={isDeletingServer}>Cancel</AlertDialogCancel>
<Button variant="destructive" disabled={isDeletingServer} onClick={confirmDelete}>
{isDeletingServer ? "Deleting..." : "Delete"}
</Button>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
<DeleteServerDialog
open={isDeleteModalOpen}
onOpenChange={(open) => !open && cancelDelete()}
server={serverToDelete}
isDeleting={isDeletingServer}
onConfirm={confirmDelete}
/>
<CreateMCPServer
userRole={userRole}
userID={userID}
@ -492,7 +511,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
<Plug />
My Connections
</Link>
{isAdminRole(userRole) && (
{isAdminRole(userRole) ? (
<>
<Button className="shrink-0" variant="secondary" onClick={() => setImportVisible(true)}>
Import from JSON
@ -501,8 +520,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
+ Add New MCP Server
</Button>
</>
)}
{!isAdminRole(userRole) && (
) : (
<Button
className="shrink-0"
onClick={() => {
@ -549,24 +567,23 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
Connect
</TabsTrigger>
{isAdminRole(userRole) && (
<TabsTrigger value="semantic-filter" className="flex-none rounded-none px-4 py-2">
Semantic Filter
</TabsTrigger>
)}
{isAdminRole(userRole) && (
<TabsTrigger value="tool-search" className="flex-none rounded-none px-4 py-2">
Tool Search
</TabsTrigger>
)}
{isAdminRole(userRole) && (
<TabsTrigger value="network-settings" className="flex-none rounded-none px-4 py-2">
Network Settings
</TabsTrigger>
)}
{isAdminRole(userRole) && (
<TabsTrigger value="submitted" className="flex-none rounded-none px-4 py-2">
Submitted MCPs
</TabsTrigger>
<>
<TabsTrigger value="semantic-filter" className="flex-none rounded-none px-4 py-2">
Semantic Filter
</TabsTrigger>
<TabsTrigger value="tool-search" className="flex-none rounded-none px-4 py-2">
Tool Search
</TabsTrigger>
<TabsTrigger value="network-settings" className="flex-none rounded-none px-4 py-2">
Network Settings
</TabsTrigger>
<TabsTrigger value="submitted" className="flex-none rounded-none px-4 py-2">
Submitted MCPs
</TabsTrigger>
</>
)}
{isProxyAdminTierRole(userRole) && (
<TabsTrigger value="connections" className="flex-none rounded-none px-4 py-2">
@ -601,13 +618,11 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
value={selectedTeam}
onValueChange={(v: string | null) => handleTeamChange(v ?? "all")}
>
<SelectTrigger className="w-55">
<SelectTrigger className="w-55" aria-label="Team">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="all">
{isInternalUser ? "All Available Servers" : "All Servers"}
</SelectItem>
<SelectItem value="all">{teamSelectItems.all}</SelectItem>
<SelectItem value="personal">Personal</SelectItem>
{uniqueTeams.map((team) => (
<SelectItem key={team.team_id} value={team.team_id}>
@ -641,7 +656,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
value={selectedMcpAccessGroup}
onValueChange={(v: string | null) => handleMcpAccessGroupChange(v ?? "all")}
>
<SelectTrigger className="w-55">
<SelectTrigger className="w-55" aria-label="Access Group">
<SelectValue />
</SelectTrigger>
<SelectContent>
@ -742,24 +757,23 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
<MCPConnect />
</TabsContent>
{isAdminRole(userRole) && (
<TabsContent value="semantic-filter" keepMounted>
<MCPSemanticFilterSettings accessToken={accessToken} />
</TabsContent>
)}
{isAdminRole(userRole) && (
<TabsContent value="tool-search" keepMounted>
<MCPToolSearchSettings accessToken={accessToken} />
</TabsContent>
)}
{isAdminRole(userRole) && (
<TabsContent value="network-settings" keepMounted>
<MCPNetworkSettings accessToken={accessToken} />
</TabsContent>
)}
{isAdminRole(userRole) && (
<TabsContent value="submitted" keepMounted>
<MCPSubmissionsTab accessToken={accessToken} />
</TabsContent>
<>
<TabsContent value="semantic-filter" keepMounted>
<MCPSemanticFilterSettings accessToken={accessToken} />
</TabsContent>
<TabsContent value="tool-search" keepMounted>
<MCPToolSearchSettings accessToken={accessToken} />
</TabsContent>
<TabsContent value="network-settings" keepMounted>
<MCPNetworkSettings accessToken={accessToken} />
</TabsContent>
<TabsContent value="submitted" keepMounted>
<MCPSubmissionsTab accessToken={accessToken} />
</TabsContent>
</>
)}
{isProxyAdminTierRole(userRole) && (
<TabsContent value="connections">