scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

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

Browse source
This commit is contained in:
mateo-berri 2026-11-03 21:47:52 -08:00
commit edb3be64e2
99 changed files with 3113 additions and 8470 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

87
litellm-rust/Cargo.lock generated
View file

@ -3569,37 +3569,6 @@ dependencies = [
"uuid",
]
[[package]]
name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"r2d2",
"redis",
"redis-test",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-cache-redis-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-testing",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"tokio",
]
[[package]]
name = "litellm-cache-response"
version = "0.1.0"
@ -3607,11 +3576,8 @@ dependencies = [
"litellm-cache",
"litellm-cache-gcs",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-http",
"py_literal",
"redis",
"redis-test",
"rstest",
"serde",
"serde_json",
@ -3650,24 +3616,6 @@ dependencies = [
"litellm-cache",
]
[[package]]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-testing",
"litellm-error",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"uuid",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -4339,8 +4287,6 @@ dependencies = [
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core-utils",
@ -4369,7 +4315,6 @@ dependencies = [
"prost",
"pyo3",
"pyo3-async-runtimes",
"redis",
"reqwest",
"rstest",
"serde",
@ -5782,17 +5727,6 @@ version = "6.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
[[package]]
name = "r2d2"
version = "0.8.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
dependencies = [
"log",
"parking_lot",
"scheduled-thread-pool",
]
[[package]]
name = "rand"
version = "0.8.7"
@ -5966,18 +5900,6 @@ dependencies = [
"xxhash-rust",
]
[[package]]
name = "redis-test"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
dependencies = [
"rand 0.9.5",
"redis",
"socket2 0.6.5",
"tempfile",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
@ -6454,15 +6376,6 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "scheduled-thread-pool"
version = "0.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
dependencies = [
"parking_lot",
]
[[package]]
name = "schemars"
version = "0.9.0"

3
litellm-rust/Cargo.toml
View file

@ -57,11 +57,9 @@ litellm-db-testing = { path = "crates/db-testing" }
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" }
litellm-cache-response = { path = "crates/cache-response" }
litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
litellm-cache-testing = { path = "crates/cache-testing" }
@ -97,7 +95,6 @@ macro_rules_attribute = "0.2.3"
schemars = "1"
reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] }
redis = "1.8.0"
redis-test = "1.0.4"
testcontainers-modules = "0.15.0"
qdrant-client = { version = "1.19.0", default-features = false }
uuid = { version = "1", features = ["v4"] }

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

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

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

@ -1,404 +0,0 @@
use std::{
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheCodec, Error, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
};
use litellm_cache_redis::{
RedisTopology,
connection::{ConnectionRef, Connections},
};
use sha2::{Digest, Sha256};
use crate::{
RedisSemanticConfig,
index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD},
reply::{bytes_field, first_document, number_field, string_field},
};
struct Inner {
index: Index,
distance_threshold: f64,
clock: fn() -> f64,
}
impl Inner {
fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self {
Self {
index: Index::new(config.index_name),
distance_threshold: 1.0 - f64::from(config.similarity_threshold),
clock,
}
}
fn store(
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
response: Vec<u8>,
prompt: &str,
vector: &[f32],
ttl: Option<Duration>,
) -> Result<(), Error> {
let index = self.index.ensure(connection, vector.len())?;
let entry_id = entry_id(prompt, tag);
let hash_key = format!("{index}:{entry_id}");
redis::cmd("HSET")
.arg(&hash_key)
.arg("entry_id")
.arg(&entry_id)
.arg("prompt")
.arg(prompt)
.arg("response")
.arg(response)
.arg(VECTOR_FIELD)
.arg(vector_buffer(vector))
.arg("inserted_at")
.arg(format!("{}", (self.clock)()))
.arg("updated_at")
.arg(format!("{}", (self.clock)()))
.arg(CACHE_KEY_FIELD)
.arg(tag)
.query::<()>(connection)
.map_err(Error::unavailable)?;
if let Some(ttl) = ttl {
redis::cmd("EXPIRE")
.arg(&hash_key)
.arg(ttl_seconds(ttl))
.query::<()>(connection)
.map_err(Error::unavailable)?;
}
Ok(())
}
fn lookup(
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
vector: &[f32],
) -> Result<SemanticLookup<Vec<u8>>, Error> {
let index = self.index.ensure(connection, vector.len())?;
let query = format!(
"(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]",
escape_tag(tag)
);
let result = redis::cmd("FT.SEARCH")
.arg(&index)
.arg(query)
.arg("RETURN")
.arg(8)
.arg("entry_id")
.arg("prompt")
.arg("response")
.arg("inserted_at")
.arg("updated_at")
.arg("metadata")
.arg(CACHE_KEY_FIELD)
.arg("vector_distance")
.arg("SORTBY")
.arg("vector_distance")
.arg("ASC")
.arg("DIALECT")
.arg(2)
.arg("LIMIT")
.arg(0)
.arg(1)
.arg("PARAMS")
.arg(2)
.arg("vector")
.arg(vector_buffer(vector))
.query::<redis::Value>(connection)
.map_err(Error::unavailable)?;
let Some(fields) = first_document(&result) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
return Ok(SemanticLookup::miss(Some(0.0)));
}
// redisvl's range query only returns entries within the distance threshold, so a
// farther hit reads as no result.
let Some(distance) = number_field(fields, "vector_distance")
.filter(|distance| *distance <= self.distance_threshold)
else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let Some(response) = bytes_field(fields, "response") else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
Ok(SemanticLookup {
value: Some(response),
similarity: Some(1.0 - distance),
})
}
}
/// `RedisSemanticCache`: a redisvl-compatible semantic index on Redis Stack. Values go through
/// the injected codec, so the response layer decides what a cached entry is.
pub struct RedisSemanticCache<E, S, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
codec: S,
inner: Arc<Inner>,
}
impl<E: Embedder, S: CacheCodec> RedisSemanticCache<E, S> {
pub fn new(
url: &str,
embedder: E,
codec: S,
config: RedisSemanticConfig,
) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
codec,
inner: Arc::new(Inner::new(config, timestamp)),
})
}
}
impl<E, S, C> RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(
connection: C,
embedder: E,
codec: S,
config: RedisSemanticConfig,
) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
embedder,
codec,
inner: Arc::new(Inner::new(config, timestamp)),
}
}
pub fn with_clock(self, clock: fn() -> f64) -> Self {
let config = RedisSemanticConfig {
index_name: self.index_name().to_owned(),
similarity_threshold: self.similarity_threshold(),
};
Self {
inner: Arc::new(Inner::new(config, clock)),
..self
}
}
pub fn embedder(&self) -> &E {
&self.embedder
}
pub fn index_name(&self) -> &str {
self.inner.index.name()
}
pub fn similarity_threshold(&self) -> f32 {
(1.0 - self.inner.distance_threshold) as f32
}
fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str {
context.scope.as_deref().unwrap_or(key)
}
fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
Ok(SemanticLookup {
value: lookup
.value
.map(|bytes| self.codec.decode(&bytes))
.transpose()?,
similarity: lookup.similarity,
})
}
}
impl<E, S, C> BaseCache for RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = SemanticCacheContext;
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> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(());
};
let response = self.codec.encode(&value)?;
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context);
self.connections.execute(|connection| {
self.inner
.store(connection, tag, response, &prompt, &vector, context.ttl)
})
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
let Some(prompt) = prompt_from_context(&context) else {
return Ok(());
};
let response = self.codec.encode(&value)?;
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, &context).to_owned();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.store(connection, &tag, response, &prompt, &vector, context.ttl)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.async_get_cache_with_similarity(key, context)
.await
.map(|lookup| lookup.value)
}
}
/// Python stamps a similarity of `0.0` when there is no prompt or no hit in the key's scope.
impl<E, S, C> SemanticCache for RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context);
let lookup = self
.connections
.execute(|connection| self.inner.lookup(connection, tag, &vector))?;
self.decode(lookup)
}
async fn async_get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, context).to_owned();
let inner = Arc::clone(&self.inner);
let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.lookup(connection, &tag, &vector)
})
.await?;
self.decode(lookup)
}
}
fn timestamp() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or_default()
}
fn entry_id(prompt: &str, tag: &str) -> String {
let mut digest = Sha256::new();
digest.update(prompt.as_bytes());
digest.update(CACHE_KEY_FIELD.as_bytes());
digest.update(tag.as_bytes());
format!("{:x}", digest.finalize())
}
fn vector_buffer(vector: &[f32]) -> Vec<u8> {
vector
.iter()
.flat_map(|component| component.to_le_bytes())
.collect()
}
fn escape_tag(value: &str) -> String {
let mut escaped = String::with_capacity(value.len());
for ch in value.chars() {
if matches!(
ch,
',' | '.'
| '<'
| '>'
| '{'
| '}'
| '['
| ']'
| '\\'
| '"'
| '\''
| ':'
| ';'
| '!'
| '@'
| '#'
| '$'
| '%'
| '^'
| '&'
| '*'
| '('
| ')'
| '-'
| '+'
| '='
| '~'
| '|'
| '/'
| ' '
| '?'
) {
escaped.push('\\');
}
escaped.push(ch);
}
escaped
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}

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

@ -1,8 +0,0 @@
/// `RedisSemanticCache.DEFAULT_REDIS_INDEX_NAME`.
pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
#[derive(Clone, Debug)]
pub struct RedisSemanticConfig {
pub index_name: String,
pub similarity_threshold: f32,
}

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

@ -1,209 +0,0 @@
use std::sync::OnceLock;
use litellm_cache::Error;
use litellm_cache_redis::connection::ConnectionRef;
use crate::reply::{number_value, string_value};
pub(crate) const CACHE_KEY_FIELD: &str = "litellm_cache_key";
pub(crate) const VECTOR_FIELD: &str = "prompt_vector";
/// The redisvl `SemanticCache` index, resolved once per cache: the configured name when its
/// schema fits, else `<name>_isolated`, recreated when that one is stale too.
pub(crate) struct Index {
name: String,
resolved: OnceLock<String>,
}
impl Index {
pub(crate) fn new(name: String) -> Self {
Self {
name,
resolved: OnceLock::new(),
}
}
pub(crate) fn name(&self) -> &str {
&self.name
}
pub(crate) fn ensure(
&self,
connection: &mut ConnectionRef<'_>,
dims: usize,
) -> Result<String, Error> {
if let Some(name) = self.resolved.get() {
return Ok(name.clone());
}
let name = match index_compatible(connection, &self.name, dims)? {
Some(true) => self.name.clone(),
Some(false) => self.isolated(connection, dims)?,
None => match create_index(connection, &self.name, dims) {
Ok(()) => self.name.clone(),
Err(error) => match index_compatible(connection, &self.name, dims)? {
Some(true) => self.name.clone(),
Some(false) => self.isolated(connection, dims)?,
None => return Err(error),
},
},
};
let _ = self.resolved.set(name.clone());
Ok(name)
}
fn isolated(&self, connection: &mut ConnectionRef<'_>, dims: usize) -> Result<String, Error> {
let name = format!("{}_isolated", self.name);
match index_compatible(connection, &name, dims)? {
Some(true) => Ok(name),
Some(false) => {
redis::cmd("FT.DROPINDEX")
.arg(&name)
.query::<()>(connection)
.map_err(Error::unavailable)?;
create_index(connection, &name, dims)?;
Ok(name)
}
None => {
create_index(connection, &name, dims)?;
Ok(name)
}
}
}
}
fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
redis::cmd("FT.CREATE")
.arg(name)
.arg("ON")
.arg("HASH")
.arg("PREFIX")
.arg(1)
.arg(name)
.arg("SCORE")
.arg(1.0)
.arg("SCHEMA")
.arg("prompt")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("response")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("inserted_at")
.arg("NUMERIC")
.arg("updated_at")
.arg("NUMERIC")
.arg(VECTOR_FIELD)
.arg("VECTOR")
.arg("FLAT")
.arg(6)
.arg("TYPE")
.arg("FLOAT32")
.arg("DIM")
.arg(dims)
.arg("DISTANCE_METRIC")
.arg("COSINE")
.arg(CACHE_KEY_FIELD)
.arg("TAG")
.arg("SEPARATOR")
.arg(",")
.query::<()>(connection)
.map_err(Error::unavailable)
}
fn index_compatible(
connection: &mut ConnectionRef<'_>,
name: &str,
dims: usize,
) -> Result<Option<bool>, Error> {
let info = match redis::cmd("FT.INFO")
.arg(name)
.query::<redis::Value>(connection)
{
Ok(info) => info,
Err(error) if unknown_index(&error) => return Ok(None),
Err(error) => return Err(Error::unavailable(error)),
};
Ok(Some(schema_compatible(&info, dims)))
}
fn unknown_index(error: &redis::RedisError) -> bool {
error
.code()
.is_some_and(|code| code.eq_ignore_ascii_case("unknown"))
&& error
.detail()
.is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index"))
}
struct Attribute {
name: Option<String>,
field_type: Option<String>,
dim: Option<f64>,
data_type: Option<String>,
distance_metric: Option<String>,
}
fn attribute(value: &redis::Value) -> Option<Attribute> {
let redis::Value::Array(pairs) = value else {
return None;
};
let mut attribute = Attribute {
name: None,
field_type: None,
dim: None,
data_type: None,
distance_metric: None,
};
for pair in pairs.as_chunks::<2>().0 {
match string_value(&pair[0]).as_deref() {
Some("identifier") => attribute.name = string_value(&pair[1]),
Some("type") => attribute.field_type = string_value(&pair[1]),
Some("dim") => attribute.dim = number_value(&pair[1]),
Some("data_type") => attribute.data_type = string_value(&pair[1]),
Some("distance_metric") => attribute.distance_metric = string_value(&pair[1]),
_ => {}
}
}
Some(attribute)
}
fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
let redis::Value::Array(entries) = info else {
return false;
};
let attributes = entries
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
.map(|pair| &pair[1]);
let Some(redis::Value::Array(attributes)) = attributes else {
return false;
};
let fields = attributes.iter().filter_map(attribute).collect::<Vec<_>>();
let has_field = |name: &str, field_type: &str| {
fields.iter().any(|field| {
field.name.as_deref() == Some(name) && field.field_type.as_deref() == Some(field_type)
})
};
has_field("prompt", "TEXT")
&& has_field("response", "TEXT")
&& has_field("inserted_at", "NUMERIC")
&& has_field("updated_at", "NUMERIC")
&& has_field(CACHE_KEY_FIELD, "TAG")
&& fields.iter().any(|field| {
field.name.as_deref() == Some(VECTOR_FIELD)
&& field.field_type.as_deref() == Some("VECTOR")
&& field.dim == Some(dims as f64)
&& field
.data_type
.as_deref()
.is_some_and(|data| data.eq_ignore_ascii_case("float32"))
&& field
.distance_metric
.as_deref()
.is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
})
}

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

@ -1,7 +0,0 @@
mod cache;
mod config;
mod index;
mod reply;
pub use cache::RedisSemanticCache;
pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig};

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

@ -1,57 +0,0 @@
pub(crate) fn string_value(value: &redis::Value) -> Option<String> {
match value {
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
redis::Value::SimpleString(text) => Some(text.clone()),
redis::Value::VerbatimString { text, .. } => Some(text.clone()),
_ => None,
}
}
pub(crate) fn number_value(value: &redis::Value) -> Option<f64> {
match value {
redis::Value::Int(number) => Some(*number as f64),
redis::Value::Double(number) => Some(*number),
_ => string_value(value).and_then(|text| text.parse().ok()),
}
}
pub(crate) fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
let redis::Value::Array(items) = result else {
return None;
};
let [count, _document_id, fields, ..] = items.as_slice() else {
return None;
};
if !matches!(count, redis::Value::Int(count) if *count > 0) {
return None;
}
match fields {
redis::Value::Array(fields) => Some(fields.as_slice()),
_ => None,
}
}
fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
fields
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some(name))
.map(|pair| &pair[1])
}
pub(crate) fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
field_value(fields, name).and_then(string_value)
}
pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
field_value(fields, name).and_then(number_value)
}
pub(crate) fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
match field_value(fields, name)? {
redis::Value::BulkString(bytes) => Some(bytes.clone()),
redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
_ => None,
}
}

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

@ -1,778 +0,0 @@
mod support;
use std::time::Duration;
use litellm_cache::{
BaseCache, Error, JsonCodec, SemanticCacheContext,
semantic::{SemanticCache, SemanticLookup},
};
use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
use redis_test::{MockCmd, MockRedisConnection};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use support::FakeEmbedder;
const INDEX: &str = DEFAULT_INDEX_NAME;
const PROMPT: &str = "hello prompt";
const CLOCK: fn() -> f64 = || 1700000000.5;
const VECTOR: [f32; 3] = [0.1, 0.2, 0.3];
type MockCache = RedisSemanticCache<FakeEmbedder, JsonCodec<Value>, MockRedisConnection>;
#[fixture]
fn config() -> RedisSemanticConfig {
RedisSemanticConfig {
index_name: INDEX.into(),
similarity_threshold: 0.9,
}
}
#[fixture]
fn entry() -> Value {
json!({"timestamp": 1.0, "response": {"answer": "yes"}})
}
#[fixture]
fn context() -> SemanticCacheContext {
messages_context(vec![json!({"role": "user", "content": PROMPT})])
}
fn messages_context(messages: Vec<Value>) -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(Value::Array(messages)),
..Default::default()
}
}
fn cache(commands: Vec<MockCmd>, embedder: FakeEmbedder) -> MockCache {
RedisSemanticCache::with_connection(
MockRedisConnection::new(commands).assert_all_commands_consumed(),
embedder,
JsonCodec::new(),
config(),
)
.with_clock(CLOCK)
}
fn encoded(value: &Value) -> Vec<u8> {
serde_json::to_vec(value).unwrap()
}
fn vector_bytes(vector: &[f32]) -> Vec<u8> {
vector
.iter()
.flat_map(|component| component.to_le_bytes())
.collect()
}
fn entry_id(prompt: &str, tag: &str) -> String {
let mut digest = Sha256::new();
digest.update(prompt.as_bytes());
digest.update(b"litellm_cache_key");
digest.update(tag.as_bytes());
format!("{:x}", digest.finalize())
}
fn s(value: &str) -> redis::Value {
redis::Value::BulkString(value.as_bytes().to_vec())
}
fn unknown_index_error() -> redis::RedisError {
support::server_error("Unknown index name")
}
fn info_missing(index: &str) -> MockCmd {
MockCmd::new(
redis::cmd("FT.INFO").arg(index),
Err::<redis::Value, _>(unknown_index_error()),
)
}
fn info(index: &str, value: redis::Value) -> MockCmd {
MockCmd::new(redis::cmd("FT.INFO").arg(index), Ok(value))
}
fn attribute(name: &str, field_type: &str, extra: Vec<redis::Value>) -> redis::Value {
let mut parts = vec![
s("identifier"),
s(name),
s("attribute"),
s(name),
s("type"),
s(field_type),
];
parts.extend(extra);
redis::Value::Array(parts)
}
fn index_info(attributes: Vec<redis::Value>) -> redis::Value {
redis::Value::Array(vec![
s("index_name"),
s(INDEX),
s("attributes"),
redis::Value::Array(attributes),
])
}
fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value {
attribute(
"prompt_vector",
"VECTOR",
vec![
s("algorithm"),
s("FLAT"),
s("data_type"),
s(data_type),
s("dim"),
redis::Value::Int(dims),
s("distance_metric"),
s(distance_metric),
],
)
}
fn info_with_vector(vector: redis::Value) -> redis::Value {
index_info(vec![
attribute("prompt", "TEXT", vec![]),
attribute("response", "TEXT", vec![]),
attribute("inserted_at", "NUMERIC", vec![]),
attribute("updated_at", "NUMERIC", vec![]),
vector,
attribute("litellm_cache_key", "TAG", vec![]),
])
}
fn compatible_info(dims: i64) -> redis::Value {
info_with_vector(vector_attribute_with(dims, "FLOAT32", "COSINE"))
}
fn unscoped_info(dims: i64) -> redis::Value {
index_info(vec![
attribute("prompt", "TEXT", vec![]),
attribute("response", "TEXT", vec![]),
attribute("inserted_at", "NUMERIC", vec![]),
attribute("updated_at", "NUMERIC", vec![]),
vector_attribute_with(dims, "FLOAT32", "COSINE"),
])
}
fn create_index(name: &str, dims: usize) -> MockCmd {
MockCmd::new(create_index_command(name, dims), Ok("OK"))
}
fn create_index_command(name: &str, dims: usize) -> redis::Cmd {
let mut command = redis::cmd("FT.CREATE");
command
.arg(name)
.arg("ON")
.arg("HASH")
.arg("PREFIX")
.arg(1)
.arg(name)
.arg("SCORE")
.arg(1.0)
.arg("SCHEMA")
.arg("prompt")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("response")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("inserted_at")
.arg("NUMERIC")
.arg("updated_at")
.arg("NUMERIC")
.arg("prompt_vector")
.arg("VECTOR")
.arg("FLAT")
.arg(6)
.arg("TYPE")
.arg("FLOAT32")
.arg("DIM")
.arg(dims)
.arg("DISTANCE_METRIC")
.arg("COSINE")
.arg("litellm_cache_key")
.arg("TAG")
.arg("SEPARATOR")
.arg(",");
command
}
fn hset(index: &str, prompt: &str, tag: &str, vector: &[f32], value: &Value) -> MockCmd {
MockCmd::new(
redis::cmd("HSET")
.arg(format!("{index}:{}", entry_id(prompt, tag)))
.arg("entry_id")
.arg(entry_id(prompt, tag))
.arg("prompt")
.arg(prompt)
.arg("response")
.arg(encoded(value))
.arg("prompt_vector")
.arg(vector_bytes(vector))
.arg("inserted_at")
.arg("1700000000.5")
.arg("updated_at")
.arg("1700000000.5")
.arg("litellm_cache_key")
.arg(tag),
Ok(7),
)
}
fn search(
index: &str,
tag: &str,
vector: &[f32],
reply: redis::RedisResult<redis::Value>,
) -> MockCmd {
let mut command = redis::cmd("FT.SEARCH");
command
.arg(index)
.arg(format!(
"(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]"
))
.arg("RETURN")
.arg(8)
.arg("entry_id")
.arg("prompt")
.arg("response")
.arg("inserted_at")
.arg("updated_at")
.arg("metadata")
.arg("litellm_cache_key")
.arg("vector_distance")
.arg("SORTBY")
.arg("vector_distance")
.arg("ASC")
.arg("DIALECT")
.arg(2)
.arg("LIMIT")
.arg(0)
.arg(1)
.arg("PARAMS")
.arg(2)
.arg("vector")
.arg(vector_bytes(vector));
MockCmd::new(command, reply)
}
fn hit(tag: &str, distance: &str, response: Vec<u8>) -> redis::Value {
redis::Value::Array(vec![
redis::Value::Int(1),
s("litellm_semantic_cache_index:stored-id"),
redis::Value::Array(vec![
s("entry_id"),
s("stored-id"),
s("prompt"),
s(PROMPT),
s("response"),
redis::Value::BulkString(response),
s("inserted_at"),
s("1700000000.5"),
s("updated_at"),
s("1700000000.5"),
s("litellm_cache_key"),
s(tag),
s("vector_distance"),
s(distance),
]),
])
}
fn empty_result() -> redis::Value {
redis::Value::Array(vec![redis::Value::Int(0)])
}
#[rstest]
#[case::creates_index_and_expires(false, Some(Duration::from_secs(5)), Some(5))]
#[case::existing_index_without_ttl(true, None, None)]
#[case::fractional_ttl_rounds_up(true, Some(Duration::from_millis(1500)), Some(2))]
fn store_writes_the_redisvl_hash(
#[case] index_exists: bool,
#[case] ttl: Option<Duration>,
#[case] expire: Option<u64>,
entry: Value,
context: SemanticCacheContext,
) {
let hash_key = format!("{INDEX}:{}", entry_id(PROMPT, "key1"));
let mut commands = if index_exists {
vec![info(INDEX, compatible_info(3))]
} else {
vec![info_missing(INDEX), create_index(INDEX, 3)]
};
commands.push(hset(INDEX, PROMPT, "key1", &VECTOR, &entry));
commands.extend(
expire.map(|seconds| MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(seconds), Ok(1))),
);
let cache = cache(commands, FakeEmbedder::new(&[]));
cache
.set_cache("key1", entry, &SemanticCacheContext { ttl, ..context })
.unwrap();
}
#[rstest]
#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))]
#[case::above_distance_threshold("key1", "0.5", None, Ok(None))]
#[case::other_cache_key("other", "0.05", None, Ok(None))]
#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))]
fn lookup_applies_threshold_scope_and_codec(
#[case] stored_tag: &str,
#[case] distance: &str,
#[case] response: Option<&[u8]>,
#[case] expected: Result<Option<Value>, Error>,
entry: Value,
context: SemanticCacheContext,
) {
let response = response.map_or_else(|| encoded(&entry), <[u8]>::to_vec);
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
search(
INDEX,
"key1",
&VECTOR,
Ok(hit(stored_tag, distance, response)),
),
],
FakeEmbedder::new(&[]),
);
let kind = |result: Result<Option<Value>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected));
}
#[rstest]
#[case::hit(context(), Some(hit("key1", "0.05", encoded(&entry()))), Some(entry()), Some(1.0 - 0.05))]
#[case::beyond_distance_threshold(context(), Some(hit("key1", "0.5", encoded(&entry()))), None, Some(0.0))]
#[case::no_results(context(), Some(empty_result()), None, Some(0.0))]
#[case::other_cache_key(context(), Some(hit("other", "0.05", encoded(&entry()))), None, Some(0.0))]
#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
#[tokio::test]
async fn lookup_reports_python_semantic_similarity(
#[case] context: SemanticCacheContext,
#[case] reply: Option<redis::Value>,
#[case] value: Option<Value>,
#[case] similarity: Option<f64>,
#[values(false, true)] use_async: bool,
) {
let commands = reply.map_or_else(Vec::new, |reply| {
vec![
info(INDEX, compatible_info(3)),
search(INDEX, "key1", &VECTOR, Ok(reply)),
]
});
let cache = cache(commands, FakeEmbedder::new(&[]));
let lookup = if use_async {
cache
.async_get_cache_with_similarity("key1", &context)
.await
} else {
cache.get_cache_with_similarity("key1", &context)
};
assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
}
#[rstest]
#[tokio::test]
async fn missing_prompt_is_a_noop_that_never_embeds(entry: Value) {
let embedder = FakeEmbedder::new(&[]);
let calls = embedder.calls.clone();
let cache = cache(Vec::new(), embedder);
let context = SemanticCacheContext::default();
cache.set_cache("key1", entry.clone(), &context).unwrap();
assert_eq!(cache.get_cache("key1", &context).unwrap(), None);
cache
.async_set_cache("key1", entry, context.clone())
.await
.unwrap();
assert_eq!(cache.async_get_cache("key1", &context).await.unwrap(), None);
assert!(calls.lock().unwrap().is_empty());
}
#[rstest]
fn scope_overrides_key_as_filter_tag(entry: Value, context: SemanticCacheContext) {
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
hset(INDEX, PROMPT, "scope-a", &VECTOR, &entry),
search(
INDEX,
"scope\\-a",
&VECTOR,
Ok(hit("scope-a", "0.05", encoded(&entry))),
),
],
FakeEmbedder::new(&[]),
);
let context = SemanticCacheContext {
scope: Some("scope-a".into()),
..context
};
cache.set_cache("key1", entry.clone(), &context).unwrap();
assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(entry));
}
#[rstest]
#[case::unscoped_schema(unscoped_info(3))]
#[case::wrong_distance_metric(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2")))]
#[case::wrong_data_type(info_with_vector(vector_attribute_with(3, "FLOAT64", "COSINE")))]
fn incompatible_schema_falls_back_to_isolated_index(
#[case] base_info: redis::Value,
entry: Value,
context: SemanticCacheContext,
) {
let isolated = format!("{INDEX}_isolated");
let cache = cache(
vec![
info(INDEX, base_info),
info_missing(&isolated),
create_index(&isolated, 3),
hset(&isolated, PROMPT, "key1", &VECTOR, &entry),
],
FakeEmbedder::new(&[]),
);
cache.set_cache("key1", entry, &context).unwrap();
}
#[rstest]
fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticCacheContext) {
let cache = cache(
vec![
info_missing(INDEX),
MockCmd::new(
create_index_command(INDEX, 3),
Err::<&str, _>(support::server_error("Index already exists")),
),
info(INDEX, compatible_info(3)),
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
],
FakeEmbedder::new(&[]),
);
cache.set_cache("key1", entry, &context).unwrap();
}
#[rstest]
#[case::punctuation_and_spaces("a:b, c|d", "a\\:b\\,\\ c\\|d")]
#[case::braces_and_dots("{x}.y", "\\{x\\}\\.y")]
#[case::plain("key1", "key1")]
fn tag_special_characters_are_escaped_in_search_filter(
#[case] tag: &str,
#[case] escaped: &str,
context: SemanticCacheContext,
) {
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
search(INDEX, escaped, &VECTOR, Ok(empty_result())),
],
FakeEmbedder::new(&[]),
);
assert_eq!(cache.get_cache(tag, &context).unwrap(), None);
}
#[rstest]
#[case::content_parts(
SemanticCacheContext {
messages: Some(json!([
{"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]},
{"role": "assistant", "content": "reply"},
])),
..Default::default()
},
"firstsecondreply"
)]
#[case::responses_string_input(
SemanticCacheContext { input: Some(json!(" plain input ")), ..Default::default() },
"plain input"
)]
#[case::responses_nested_input(
SemanticCacheContext {
input: Some(json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"])),
..Default::default()
},
"nested\ntail"
)]
#[case::responses_output_text(
SemanticCacheContext { input: Some(json!({"output_text": " result text "})), ..Default::default() },
"result text"
)]
#[case::search_results(
messages_context(vec![json!({
"role": "user",
"content": "question",
"search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}],
})]),
"questionsrctfound{\"a\":1}"
)]
#[case::empty_messages_fall_back_to_input(
SemanticCacheContext { messages: Some(json!([])), input: Some(json!("fallback")), ..Default::default() },
"fallback"
)]
fn prompt_extraction_matches_python_message_and_input_shapes(
#[case] context: SemanticCacheContext,
#[case] prompt: &str,
) {
let embedder = FakeEmbedder::new(&[]);
let calls = embedder.calls.clone();
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
search(INDEX, "key1", &VECTOR, Ok(empty_result())),
],
embedder,
);
cache.get_cache("key1", &context).unwrap();
assert_eq!(*calls.lock().unwrap(), vec![(prompt.to_owned(), None)]);
}
#[rstest]
#[case(None)]
#[case(Some(Duration::from_secs(9)))]
fn ttl_passes_through_context_only(#[case] ttl: Option<Duration>) {
let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
assert_eq!(
cache.get_ttl(&SemanticCacheContext {
ttl,
..Default::default()
}),
ttl
);
}
#[rstest]
#[tokio::test]
async fn async_paths_embed_with_metadata_then_run_blocking_redis_work(
entry: Value,
context: SemanticCacheContext,
) {
let embedder = FakeEmbedder::new(&[]);
let calls = embedder.calls.clone();
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
search(
INDEX,
"key1",
&VECTOR,
Ok(hit("key1", "0.05", encoded(&entry))),
),
],
embedder,
);
let context = SemanticCacheContext {
metadata: Some(json!({"tenant": "team"})),
..context
};
cache
.async_set_cache("key1", entry.clone(), context.clone())
.await
.unwrap();
assert_eq!(
cache.async_get_cache("key1", &context).await.unwrap(),
Some(entry)
);
assert_eq!(
*calls.lock().unwrap(),
vec![(PROMPT.to_owned(), context.metadata.clone()); 2]
);
}
#[rstest]
fn accessors_report_the_config(config: RedisSemanticConfig) {
let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
assert_eq!(cache.index_name(), config.index_name);
assert!((cache.similarity_threshold() - config.similarity_threshold).abs() < 1e-6);
}
#[rstest]
fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) {
// Pins parity with Python's `_isolated` + overwrite=True flow.
let prompt = "shared prompt";
let isolated = format!("{INDEX}_isolated");
let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
let vector_a = vec![0.1f32; 8];
let worker_a = cache(
vec![
info_missing(INDEX),
create_index(INDEX, 8),
hset(INDEX, prompt, "key1", &vector_a, &entry),
],
FakeEmbedder::new(&[(prompt, &vector_a)]),
);
worker_a
.set_cache("key1", entry.clone(), &context())
.unwrap();
let vector_b = vec![0.2f32; 4];
let worker_b = cache(
vec![
info(INDEX, compatible_info(8)),
info_missing(&isolated),
create_index(&isolated, 4),
hset(&isolated, prompt, "key1", &vector_b, &entry),
search(
&isolated,
"key1",
&vector_b,
Ok(hit("key1", "0.0", encoded(&entry))),
),
search(
&isolated,
"key1",
&vector_b,
Err(redis::RedisError::from((
redis::ErrorKind::Extension,
"Vector dimension mismatch",
))),
),
],
FakeEmbedder::new(&[(prompt, &vector_b)]),
);
worker_b
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert_eq!(
worker_b.get_cache("key1", &context()).unwrap(),
Some(entry.clone())
);
let vector_c = vec![0.3f32; 16];
let worker_c = cache(
vec![
info(INDEX, compatible_info(8)),
info(&isolated, compatible_info(4)),
MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")),
create_index(&isolated, 16),
hset(&isolated, prompt, "key1", &vector_c, &entry),
],
FakeEmbedder::new(&[(prompt, &vector_c)]),
);
worker_c
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert!(matches!(
worker_b.get_cache("key1", &context()),
Err(Error::Unavailable(Some(_)))
));
}
#[fixture]
fn redis_stack_url() -> Option<String> {
std::env::var("LITELLM_REDIS_STACK_URL").ok()
}
fn live_cache(
url: &str,
index_name: &str,
prompt: &str,
vector: Vec<f32>,
) -> RedisSemanticCache<FakeEmbedder, JsonCodec<Value>> {
RedisSemanticCache::new(
url,
FakeEmbedder::new(&[(prompt, vector.as_slice())]),
JsonCodec::<Value>::new(),
RedisSemanticConfig {
index_name: index_name.to_owned(),
similarity_threshold: 0.9,
},
)
.unwrap()
}
#[rstest]
fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option<String>, entry: Value) {
let Some(url) = redis_stack_url else {
return;
};
// Pins parity with Python's `_isolated` + overwrite=True flow.
let base = format!("rust_semantic_shared_{}", std::process::id());
let isolated = format!("{base}_isolated");
let prompt = "shared live prompt";
let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
let worker_a = live_cache(&url, &base, prompt, vec![0.1f32; 8]);
worker_a
.set_cache("key1", entry.clone(), &context())
.unwrap();
let worker_b = live_cache(&url, &base, prompt, vec![0.2f32; 4]);
worker_b
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert_eq!(
worker_b.get_cache("key1", &context()).unwrap(),
Some(entry.clone())
);
let worker_c = live_cache(&url, &base, prompt, vec![0.3f32; 16]);
worker_c
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert!(matches!(
worker_b.get_cache("key1", &context()),
Err(Error::Unavailable(Some(_)))
));
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
for index in [&base, &isolated] {
let _: Result<(), _> = redis::cmd("FT.DROPINDEX")
.arg(index)
.arg("DD")
.query(&mut connection);
}
}
#[rstest]
fn live_store_lookup_and_ttl_against_redis_stack(redis_stack_url: Option<String>, entry: Value) {
let Some(url) = redis_stack_url else {
return;
};
let prompt = "rust semantic cache live prompt";
let index_name = format!("rust_semantic_test_{}", std::process::id());
let cache = live_cache(&url, &index_name, prompt, vec![0.1, 0.2, 0.3, 0.4]);
let context = SemanticCacheContext {
ttl: Some(Duration::from_secs(120)),
..messages_context(vec![json!({"role": "user", "content": prompt})])
};
cache
.set_cache("live-key", entry.clone(), &context)
.unwrap();
assert_eq!(cache.get_cache("live-key", &context).unwrap(), Some(entry));
assert_eq!(cache.get_cache("other-key", &context).unwrap(), None);
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
let ttl: i64 = redis::Commands::ttl(
&mut connection,
format!("{index_name}:{}", entry_id(prompt, "live-key")),
)
.unwrap();
assert!(
ttl > 0,
"expected stored hash to carry an expiry, got {ttl}"
);
}

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

@ -1,63 +0,0 @@
mod support;
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::FakeSearch;
type Cache = RedisSemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
const PREFIX: &str = "contract:";
#[fixture]
fn cache() -> Cache {
RedisSemanticCache::with_connection(
FakeSearch::default(),
PreparedEmbedding(vec![0.6, 0.8]),
JsonCodec::new(),
RedisSemanticConfig {
index_name: DEFAULT_INDEX_NAME.into(),
similarity_threshold: 0.9,
},
)
}
#[fixture]
fn context() -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
..Default::default()
}
}
#[rstest]
#[tokio::test]
async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) {
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test]
async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) {
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test]
async fn overwrite_replaces(cache: Cache, context: SemanticCacheContext) {
contract::overwrite_replaces(&cache, context, PREFIX, json!(1), json!({"b": 2})).await;
}
#[rstest]
#[tokio::test]
async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) {
contract::pipeline_writes_every_entry(
&cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}

302
litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs
View file

@ -1,302 +0,0 @@
#![allow(dead_code)]
use std::{
collections::{BTreeMap, HashMap},
sync::{Arc, Mutex},
};
use litellm_cache::{Error, semantic::Embedder};
use serde_json::Value;
pub type EmbedCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every
/// prompt with its metadata.
pub struct FakeEmbedder {
vectors: HashMap<String, Vec<f32>>,
pub calls: EmbedCalls,
}
impl FakeEmbedder {
pub fn new(vectors: &[(&str, &[f32])]) -> Self {
Self {
vectors: vectors
.iter()
.map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec()))
.collect(),
calls: EmbedCalls::default(),
}
}
}
impl Embedder for FakeEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.calls
.lock()
.unwrap()
.push((prompt.to_owned(), metadata.cloned()));
Ok(self
.vectors
.get(prompt)
.cloned()
.unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
}
async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.embed(prompt, metadata)
}
}
struct FakeIndex {
prefix: Vec<u8>,
dims: usize,
vector_field: String,
}
#[derive(Default)]
struct SearchState {
indexes: HashMap<String, FakeIndex>,
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
}
/// An in-memory Redis Stack speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache
/// sends, with exact cosine KNN over the hashes under an index prefix.
#[derive(Clone, Default)]
pub struct FakeSearch {
state: Arc<Mutex<SearchState>>,
}
impl FakeSearch {
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
let mut state = self.state.lock().unwrap();
let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned();
match text(0).to_uppercase().as_str() {
"FT.CREATE" => {
let name = text(1);
if state.indexes.contains_key(&name) {
return Err(server_error("Index already exists"));
}
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
let prefix = args[position("PREFIX").unwrap() + 2].clone();
let dims = text(position("DIM").unwrap() + 1).parse().unwrap();
let vector_field = text(position("VECTOR").unwrap() - 1);
state.indexes.insert(
name,
FakeIndex {
prefix,
dims,
vector_field,
},
);
Ok(redis::Value::Okay)
}
"FT.INFO" => {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| server_error("Unknown index name"))?;
Ok(index_info(index))
}
"FT.DROPINDEX" => {
state.indexes.remove(&text(1));
Ok(redis::Value::Okay)
}
"HSET" => {
let hash = state.hashes.entry(args[1].clone()).or_default();
for pair in args[2..].chunks(2) {
hash.insert(
String::from_utf8_lossy(&pair[0]).into_owned(),
pair[1].clone(),
);
}
Ok(redis::Value::Int(((args.len() - 2) / 2) as i64))
}
"EXPIRE" => Ok(redis::Value::Int(i64::from(
state.hashes.contains_key(&args[1]),
))),
"FT.SEARCH" => {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| server_error("no such index"))?;
let query = text(2);
let tag = query_tag(&query);
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
let vector = floats(&args[params + 3]);
let best = state
.hashes
.iter()
.filter(|(key, _)| key.starts_with(&index.prefix))
.filter(|(_, fields)| {
fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes())
})
.filter_map(|(key, fields)| {
let stored = floats(fields.get(&index.vector_field)?);
(stored.len() == index.dims)
.then(|| (key, fields, 1.0 - cosine(&vector, &stored)))
})
.min_by(|left, right| left.2.total_cmp(&right.2));
let Some((key, fields, distance)) = best else {
return Ok(redis::Value::Array(vec![redis::Value::Int(0)]));
};
let mut reply = fields
.iter()
.filter(|(name, _)| **name != index.vector_field)
.flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)])
.collect::<Vec<_>>();
reply.extend([
bulk(b"vector_distance"),
bulk(distance.to_string().as_bytes()),
]);
Ok(redis::Value::Array(vec![
redis::Value::Int(1),
bulk(key),
redis::Value::Array(reply),
]))
}
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
_ => Err(server_error("unsupported command")),
}
}
}
impl redis::ConnectionLike for FakeSearch {
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
let mut commands = parse_commands(command);
self.run(commands.remove(0))
}
fn req_packed_commands(
&mut self,
commands: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
let replies = parse_commands(commands)
.into_iter()
.map(|args| self.run(args))
.collect::<redis::RedisResult<Vec<_>>>()?;
Ok(replies.into_iter().skip(offset).take(count).collect())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}
pub fn server_error(reply: &str) -> redis::RedisError {
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
.expect("error reply")
.extract_error()
.expect_err("server error")
}
fn bulk(bytes: &[u8]) -> redis::Value {
redis::Value::BulkString(bytes.to_vec())
}
fn index_info(index: &FakeIndex) -> redis::Value {
let attribute = |name: &str, field_type: &str| {
redis::Value::Array(vec![
bulk(b"identifier"),
bulk(name.as_bytes()),
bulk(b"type"),
bulk(field_type.as_bytes()),
])
};
redis::Value::Array(vec![
bulk(b"attributes"),
redis::Value::Array(vec![
attribute("prompt", "TEXT"),
attribute("response", "TEXT"),
attribute("inserted_at", "NUMERIC"),
attribute("updated_at", "NUMERIC"),
attribute("litellm_cache_key", "TAG"),
redis::Value::Array(vec![
bulk(b"identifier"),
bulk(index.vector_field.as_bytes()),
bulk(b"type"),
bulk(b"VECTOR"),
bulk(b"dim"),
redis::Value::Int(index.dims as i64),
bulk(b"data_type"),
bulk(b"FLOAT32"),
bulk(b"distance_metric"),
bulk(b"COSINE"),
]),
]),
])
}
/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed.
fn query_tag(query: &str) -> String {
let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len();
let mut tag = String::new();
let mut characters = query[start..].chars();
while let Some(character) = characters.next() {
match character {
'\\' => tag.extend(characters.next()),
'}' => break,
character => tag.push(character),
}
}
tag
}
fn floats(bytes: &[u8]) -> Vec<f32> {
bytes
.as_chunks::<4>()
.0
.iter()
.map(|chunk| f32::from_le_bytes(*chunk))
.collect()
}
fn cosine(left: &[f32], right: &[f32]) -> f64 {
let dot = left
.iter()
.zip(right)
.map(|(left, right)| f64::from(*left) * f64::from(*right))
.sum::<f64>();
let norm = |vector: &[f32]| {
vector
.iter()
.map(|value| f64::from(*value).powi(2))
.sum::<f64>()
.sqrt()
};
dot / (norm(left) * norm(right))
}
/// Splits a packed RESP request into each command's arguments.
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
let line = |bytes: &mut &[u8]| {
let end = bytes
.windows(2)
.position(|window| window == b"\r\n")
.unwrap();
let text = String::from_utf8(bytes[1..end].to_vec()).unwrap();
*bytes = &bytes[end + 2..];
text.parse::<usize>().unwrap()
};
let mut commands = Vec::new();
while !bytes.is_empty() {
let count = line(&mut bytes);
let mut args = Vec::with_capacity(count);
for _ in 0..count {
let length = line(&mut bytes);
args.push(bytes[..length].to_vec());
bytes = &bytes[length + 2..];
}
commands.push(args);
}
commands
}

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

@ -1,20 +0,0 @@
[package]
name = "litellm-cache-redis"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-error.workspace = true
redis = { workspace = true, features = ["cluster", "tls-rustls"] }
r2d2 = "0.8.10"
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
redis-test.workspace = true
rstest.workspace = true
serde_json.workspace = true

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

@ -1,192 +0,0 @@
use std::{
sync::{Arc, OnceLock},
time::Duration,
};
use litellm_cache::{BatchEntry, CacheCodec, Error};
use crate::{
connection::{ConnectionRef, Connections},
topology::RedisTopology,
};
const DEFAULT_TTL: Duration = Duration::from_secs(600);
pub struct RedisCache<S, C = redis::Connection> {
pub(crate) connections: Arc<Connections<C>>,
pub(crate) default_ttl: Duration,
pub(crate) codec: S,
pub(crate) namespace: Option<String>,
pub(crate) topology: RedisTopology,
/// The server's major version, read from `INFO` once, like Python's `redis_version`.
pub(crate) major_version: Arc<OnceLock<u32>>,
}
impl<S: CacheCodec> RedisCache<S> {
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
Self::connect(url, &RedisTopology::Standalone, default_ttl, codec)
}
pub fn connect(
url: &str,
topology: &RedisTopology,
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
let connections = Connections::open(url, topology)?;
Ok(Self {
connections: Arc::new(connections),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
topology: topology.clone(),
major_version: Arc::default(),
})
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
topology: RedisTopology::Standalone,
major_version: Arc::default(),
}
}
pub fn with_namespace(self, namespace: Option<String>) -> Self {
Self {
namespace: namespace.filter(|value| !value.is_empty()),
..self
}
}
pub fn namespace(&self) -> Option<&str> {
self.namespace.as_deref()
}
pub fn topology(&self) -> &RedisTopology {
&self.topology
}
pub(crate) fn namespaced_key(&self, key: &str) -> String {
namespaced_key(self.namespace.as_deref(), key)
}
pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec<String> {
keys.iter().map(|key| self.namespaced_key(key)).collect()
}
/// Whole seconds for `ttl`, falling back to the default TTL like Python's `get_ttl`.
pub(crate) fn ttl_or_default(&self, ttl: Option<Duration>) -> u64 {
ttl_seconds(ttl.unwrap_or(self.default_ttl))
}
pub(crate) fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
self.connections.execute(operation)
}
/// `_parse_redis_major_version`: the major version from `INFO`, or
/// `DEFAULT_REDIS_MAJOR_VERSION` when `INFO` fails or its version does not parse. The first
/// answer is kept, as Python reads `redis_version` once at construction.
pub(crate) async fn major_version(&self) -> u32 {
if let Some(version) = self.major_version.get() {
return *version;
}
let info = self
.run(|connection| connection.node_text(&redis::cmd("INFO")))
.await;
let version = info
.ok()
.and_then(|info| parse_major_version(&info))
.unwrap_or_else(default_major_version);
*self.major_version.get_or_init(|| version)
}
pub(crate) async fn run<T, F>(&self, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
Connections::run_blocking(Arc::clone(&self.connections), operation).await
}
pub(crate) fn namespaced_pattern(&self) -> Result<String, Error> {
let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
let escaped: String = namespace
.chars()
.flat_map(|ch| {
if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
vec!['\\', ch]
} else {
vec![ch]
}
})
.collect();
Ok(format!("{escaped}:*"))
}
pub(crate) fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
_ => Err(Error::InvalidEntry(None)),
}
}
pub(crate) fn decode_batch_response(
&self,
value: redis::Value,
) -> Result<BatchEntry<S::Value>, Error> {
match self.decode_response(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}
}
pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
match namespace {
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
format!("{namespace}:{key}")
}
_ => key.into(),
}
}
pub(crate) fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn parse_major_version(info: &str) -> Option<u32> {
let version = info
.lines()
.find_map(|line| line.trim().strip_prefix("redis_version:"))?
.trim();
match version.split_once('.') {
Some((major, _)) => major.parse().ok(),
None => version.parse::<f64>().ok().map(|major| major as u32),
}
}
fn default_major_version() -> u32 {
std::env::var("DEFAULT_REDIS_MAJOR_VERSION")
.ok()
.and_then(|value| value.parse().ok())
.unwrap_or(7)
}

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

@ -1,109 +0,0 @@
use litellm_cache::{CacheCodec, ClaimCache, Error, ExactCacheContext};
use redis::Commands;
use crate::{cache::RedisCache, connection::ConnectionRef};
const CLAIM_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
"elseif current ~= ARGV[1] then return 0; end; ",
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
);
const CLAIM_ATTEMPTS: usize = 8;
fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
_ => Err(Error::InvalidEntry(None)),
}
}
/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
/// bytes that decision was made on, retried when another claimant wins the race.
fn claim<S: CacheCodec>(
connection: &mut ConnectionRef<'_>,
codec: &S,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
ttl: u64,
) -> Result<S::Value, Error>
where
S::Value: PartialEq,
{
let payload = codec.encode(&candidate)?;
if payload.is_empty() {
return Err(Error::InvalidEntry(None));
}
for _ in 0..CLAIM_ATTEMPTS {
let current = stored_bytes(
connection
.get::<_, redis::Value>(key)
.map_err(Error::unavailable)?,
)?
.filter(|bytes| !bytes.is_empty());
let existing = current
.as_deref()
.and_then(|bytes| codec.decode(bytes).ok())
.filter(|existing| eligible.is_empty() || eligible.contains(existing));
let refresh = existing
.as_ref()
.is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
let write: &[u8] = if existing.is_some() { b"" } else { &payload };
let applied = redis::cmd("EVAL")
.arg(CLAIM_SCRIPT)
.arg(1)
.arg(key)
.arg(current.as_deref().unwrap_or_default())
.arg(ttl)
.arg(write)
.arg(u8::from(refresh))
.query::<bool>(connection)
.map_err(Error::unavailable)?;
if applied {
return Ok(existing.unwrap_or(candidate));
}
}
Err(Error::unavailable(ClaimContended))
}
#[derive(Debug, thiserror::Error)]
#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")]
struct ClaimContended;
impl<S, C> ClaimCache for RedisCache<S, C>
where
S: CacheCodec + Clone + 'static,
S::Value: PartialEq,
C: redis::ConnectionLike + Send + 'static,
{
fn claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
self.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
}
async fn async_claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: Vec<S::Value>,
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
let codec = self.codec.clone();
self.run(move |connection| claim(connection, &codec, &key, candidate, &eligible, ttl))
.await
}
}

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

@ -1,513 +0,0 @@
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::Error;
use redis::{
ConnectionAddr, ConnectionInfo, IntoConnectionInfo,
cluster::{
ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress,
},
cluster_routing::{
MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo,
},
};
use crate::topology::{RedisNode, RedisTopology};
pub(crate) const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
}
impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
match self {
Self::Pool(pool) => {
let mut pooled = pool.get().map_err(Error::unavailable)?;
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
result
}
Self::Cluster(pool) => {
let mut pooled = pool.get().map_err(Error::unavailable)?;
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock()?;
operation(&mut ConnectionRef::Node(&mut *connection))
}
}
}
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.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)?,
)?)),
}
}
/// Closes every idle pooled connection; the next operation opens a fresh one. Connections
/// checked out right now return to the pool, and a caller-owned connection stays open.
pub fn disconnect(&self) {
match self {
Self::Pool(pool) => close_idle(pool),
Self::Cluster(pool) => close_idle(pool),
Self::Fixed(_) => {}
}
}
}
fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
r2d2::Pool::builder()
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(manager)
.map_err(Error::unavailable)
}
fn close_idle<M, T>(pool: &r2d2::Pool<M>)
where
M: r2d2::ManageConnection<Connection = PooledConnection<T>>,
{
let mut idle = Vec::new();
while let Some(mut connection) = pool.try_get() {
connection.failed = true;
idle.push(connection);
}
}
pub(crate) struct PooledConnection<C> {
connection: C,
failed: bool,
}
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
pub(crate) struct ConnectionManager(redis::Client);
impl ConnectionManager {
fn open(url: &str) -> Result<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(Error::unavailable)
}
}
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection<redis::Connection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
#[derive(Debug, thiserror::Error)]
enum TopologyFault {
#[error("Redis Cluster needs at least one startup node")]
NoStartupNodes,
#[error("Redis Cluster nodes need a TCP address")]
NonTcpAddress,
#[error("Redis Cluster answered a fan-out command without a per-node map")]
FanOutReply,
#[error("scan target does not match the connection topology")]
ScanTarget,
}
pub(crate) struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
if startup_nodes.is_empty() {
return Err(Error::unavailable(TopologyFault::NoStartupNodes));
}
let info = url.into_connection_info().map_err(Error::unavailable)?;
let nodes = startup_nodes
.iter()
.map(|node| node_info(&info, node))
.collect::<Result<Vec<_>, _>>()?;
ClusterClientBuilder::new(nodes)
.connection_timeout(REDIS_TIMEOUT)
.response_timeout(REDIS_TIMEOUT)
.build()
.map(Self)
.map_err(Error::unavailable)
}
}
fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo, Error> {
let addr = match info.addr() {
ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port),
ConnectionAddr::TcpTls {
insecure,
tls_params,
..
} => ConnectionAddr::TcpTls {
host: node.host.clone(),
port: node.port,
insecure: *insecure,
tls_params: tls_params.clone(),
},
_ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)),
};
Ok(info.clone().set_addr(addr))
}
impl r2d2::ManageConnection for ClusterConnectionManager {
type Connection = PooledConnection<ClusterConnection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
pub enum ConnectionRef<'a> {
Node(&'a mut dyn redis::ConnectionLike),
Cluster(&'a mut ClusterConnection),
}
impl redis::ConnectionLike for ConnectionRef<'_> {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
match self {
Self::Node(connection) => connection.req_packed_command(cmd),
Self::Cluster(connection) => connection.req_packed_command(cmd),
}
}
fn req_packed_commands(
&mut self,
cmd: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
match self {
Self::Node(connection) => connection.req_packed_commands(cmd, offset, count),
Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count),
}
}
fn get_db(&self) -> i64 {
match self {
Self::Node(connection) => connection.get_db(),
Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection),
}
}
fn supports_pipelining(&self) -> bool {
match self {
Self::Node(connection) => connection.supports_pipelining(),
Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection),
}
}
fn check_connection(&mut self) -> bool {
match self {
Self::Node(connection) => connection.check_connection(),
Self::Cluster(connection) => connection.check_connection(),
}
}
fn is_open(&self) -> bool {
match self {
Self::Node(connection) => connection.is_open(),
Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection),
}
}
}
impl ConnectionRef<'_> {
/// Runs `pipeline` and decodes its non-ignored replies as `T`. A cluster connection refuses
/// `Pipeline::query`, so there a transaction goes to its keys' slot as one MULTI/EXEC and
/// anything else is split per node by `ClusterPipeline`; either way the raw replies are
/// handed back to `pipeline` to decode.
pub(crate) fn query_pipeline<T: redis::FromRedisValue>(
&mut self,
pipeline: &redis::Pipeline,
) -> Result<T, Error> {
match self {
Self::Node(connection) => pipeline.query(*connection),
Self::Cluster(connection) if pipeline.is_transaction() => {
redis::ConnectionLike::req_packed_commands(
*connection,
&pipeline.get_packed_pipeline(),
pipeline.len() + 1,
1,
)
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
}
Self::Cluster(connection) => {
let mut cluster = ClusterPipeline::with_capacity(pipeline.len());
for command in pipeline.cmd_iter() {
cluster.add_command(command.clone());
}
cluster
.query(connection)
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
}
}
.map_err(Error::unavailable)
}
pub(crate) fn scan(
&mut self,
pattern: &str,
count: usize,
mut visit: impl FnMut(&mut Self, Vec<String>) -> Result<bool, Error>,
) -> Result<(), Error> {
let pages = match self {
Self::Node(connection) => {
let page = scan_command(0, pattern, count)
.query::<ScanPage>(*connection)
.map_err(Error::unavailable)?;
vec![(None, page)]
}
Self::Cluster(connection) => connection
.route_command(
&scan_command(0, pattern, count),
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllMasters,
Some(ResponsePolicy::Special),
)),
)
.map_err(Error::unavailable)
.and_then(primary_pages)?
.into_iter()
.map(|(node, page)| (Some(node), page))
.collect(),
};
for (node, (mut cursor, mut keys)) in pages {
loop {
if !visit(self, keys)? {
return Ok(());
}
if cursor == 0 {
break;
}
(cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?;
}
}
Ok(())
}
pub(crate) fn ping(&mut self) -> Result<bool, redis::RedisError> {
let command = redis::cmd("PING");
match self {
Self::Node(connection) => command
.query::<String>(*connection)
.map(|response| response == "PONG"),
Self::Cluster(connection) => connection
.route_command(
&command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::AllSucceeded),
)),
)
.map(|_| true),
}
}
pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
match self {
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
Self::Cluster(connection) => {
let value = connection
.route_command(
command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::Special),
)),
)
.map_err(Error::unavailable)?;
let redis::Value::Map(entries) = value else {
return Err(Error::unavailable(TopologyFault::FanOutReply));
};
let mut replies = entries
.into_iter()
.map(|(node, reply)| {
Ok((
redis::from_redis_value::<String>(node).map_err(Error::unavailable)?,
redis::from_redis_value::<String>(reply).map_err(Error::unavailable)?,
))
})
.collect::<Result<Vec<(String, String)>, Error>>()?;
replies.sort();
Ok(replies
.into_iter()
.map(|(_, reply)| reply)
.collect::<Vec<_>>()
.join("\n"))
}
}
}
pub(crate) fn flushall(&mut self) -> Result<(), Error> {
let command = redis::cmd("FLUSHALL");
match self {
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
Self::Cluster(connection) => connection
.route_command(
&command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllMasters,
Some(ResponsePolicy::AllSucceeded),
)),
)
.map(|_| ())
.map_err(Error::unavailable),
}
}
fn scan_page(
&mut self,
node: Option<&NodeAddress>,
cursor: u64,
pattern: &str,
count: usize,
) -> Result<ScanPage, Error> {
let command = scan_command(cursor, pattern, count);
match (self, node) {
(Self::Node(connection), None) => {
command.query(*connection).map_err(Error::unavailable)
}
(Self::Cluster(connection), Some(node)) => connection
.route_command(
&command,
RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress {
host: node.host().to_string(),
port: node.port(),
}),
)
.map_err(Error::unavailable)
.and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)),
_ => Err(Error::unavailable(TopologyFault::ScanTarget)),
}
}
}
type ScanPage = (u64, Vec<String>);
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
let redis::Value::Map(entries) = value else {
return Err(Error::unavailable(TopologyFault::FanOutReply));
};
entries
.into_iter()
.map(|(node, page)| {
let node = redis::from_redis_value::<String>(node).map_err(Error::unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?;
let page = redis::from_redis_value::<ScanPage>(page).map_err(Error::unavailable)?;
Ok((node, page))
})
.collect()
}
fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
let mut command = redis::cmd("SCAN");
command
.cursor_arg(cursor)
.arg("MATCH")
.arg(pattern)
.arg("COUNT")
.arg(count);
command
}
/// Hands already received pipeline replies to `Pipeline::query`, so it applies its own
/// ignore and error handling to replies a cluster pipeline gathered from several nodes.
struct Replies(Option<Vec<redis::Value>>);
impl redis::ConnectionLike for Replies {
fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult<redis::Value> {
Err((redis::ErrorKind::Client, "replies hold a pipeline only").into())
}
fn req_packed_commands(
&mut self,
_: &[u8],
_: usize,
_: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
self.0
.take()
.ok_or_else(|| (redis::ErrorKind::Client, "replies were already read").into())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}

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

@ -1,205 +0,0 @@
use std::time::Duration;
use litellm_cache::{
BoundedCounterCache, CacheCodec, CountReadCache, CounterCache, Error, ExactCacheContext,
IncrementOperation,
};
use crate::{
cache::{RedisCache, ttl_seconds},
connection::ConnectionRef,
store::mget,
};
const INCREMENT_SCRIPT: &str = concat!(
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
"if redis.call('TTL', KEYS[1]) == -1 then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
);
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return count"
);
const SET_MAX_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]); ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return ARGV[1]; end; return current"
);
impl<S, C> CounterCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
self.execute(|connection| increment(connection, key, amount, ttl, false))
}
/// Python `_incrbyfloat_with_ttl`: without `refresh_ttl` the TTL is set only on a key that
/// has none, in one atomic script; with it, every increment re-arms the TTL.
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
refresh_ttl: bool,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
self.run(move |connection| increment(connection, key, amount, ttl, refresh_ttl))
.await
}
async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error> {
if operations.is_empty() {
return Ok(Vec::new());
}
let mut pipeline = redis::pipe();
for operation in operations {
let key = self.namespaced_key(&operation.key);
pipeline.cmd("INCRBYFLOAT").arg(&key).arg(operation.amount);
if let Some(ttl) = operation.ttl {
pipeline
.cmd("EXPIRE")
.arg(key)
.arg(ttl_seconds(ttl))
.ignore();
}
}
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
}
}
fn increment(
connection: &mut ConnectionRef<'_>,
key: String,
amount: f64,
ttl: u64,
refresh_ttl: bool,
) -> Result<f64, Error> {
if !refresh_ttl {
return redis::cmd("EVAL")
.arg(INCREMENT_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable);
}
connection
.query_pipeline(
redis::pipe()
.cmd("INCRBYFLOAT")
.arg(&key)
.arg(amount)
.cmd("EXPIRE")
.arg(&key)
.arg(ttl)
.ignore(),
)
.map(|(value,)| value)
}
impl<S, C> CountReadCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
let keys = self.namespaced_keys(keys);
self.execute(|connection| mget(connection, keys))?
.into_iter()
.map(count)
.collect()
}
async fn async_batch_get_counts(&self, keys: Vec<String>) -> Result<Vec<Option<i64>>, Error> {
let keys = self.namespaced_keys(&keys);
self.run(move |connection| mget(connection, keys))
.await?
.into_iter()
.map(count)
.collect()
}
}
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
redis::from_redis_value(value).map_err(Error::invalid_entry)
}
impl<S, C> BoundedCounterCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = ttl_seconds(ttl);
self.execute(|connection| increment_with_floor(connection, key, amount, ttl))
}
async fn async_increment_with_floor(
&self,
key: &str,
amount: i64,
ttl: Duration,
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = ttl_seconds(ttl);
self.run(move |connection| increment_with_floor(connection, key, amount, ttl))
.await
}
async fn async_set_max(
&self,
key: &str,
value: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(ttl);
self.run(move |connection| {
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)
.arg(key)
.arg(value)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable)
})
.await
}
}
fn increment_with_floor(
connection: &mut ConnectionRef<'_>,
key: String,
amount: i64,
ttl: u64,
) -> Result<i64, Error> {
redis::cmd("EVAL")
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable)
}

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

@ -1,63 +0,0 @@
use std::time::Duration;
use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache};
use crate::cache::RedisCache;
impl<S, C> TtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
let key = self.namespaced_key(key);
let ttl = self
.run(move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
.map_err(Error::unavailable)
})
.await?;
Ok(u64::try_from(ttl).ok().map(Duration::from_secs))
}
}
impl<S, C> RefreshTtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_refresh_ttl(&self, key: &str, ttl: Option<Duration>) -> Result<bool, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(ttl);
self.run(move |connection| {
redis::cmd("EXPIRE")
.arg(key)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable)
})
.await
}
}
impl<S, C> ScanCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
let pattern = format!("{}*", self.namespaced_key(pattern));
self.run(move |connection| {
let mut matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
Ok(matches.len() < count)
})?;
matches.truncate(count);
Ok(matches)
})
.await
}
}

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

@ -1,15 +0,0 @@
mod cache;
mod claim;
pub mod connection;
mod counter;
mod keys;
mod lifecycle;
mod queue;
mod script;
mod store;
mod topology;
pub use cache::RedisCache;
pub use queue::{RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
pub use script::{RedisArg, RedisScript};
pub use topology::{RedisNode, RedisTopology};

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

@ -1,86 +0,0 @@
use litellm_cache::{
CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache,
DisconnectCache, Error, PingCache,
};
use crate::{cache::RedisCache, topology::RedisTopology};
impl<S, C> PingCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn sync_ping(&self) -> Result<bool, Error> {
self.execute(|connection| connection.ping().map_err(Error::unavailable))
}
async fn ping(&self) -> Result<bool, Error> {
self.run(|connection| connection.ping().map_err(Error::unavailable))
.await
}
}
impl<S, C> ConnectionCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
/// Python `RedisCache.test_connection`, or `RedisClusterCache.test_connection` for a
/// cluster topology, which differs only in its messages.
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
let label = match self.topology {
RedisTopology::Standalone => "Redis",
RedisTopology::Cluster { .. } => "Redis Cluster",
};
let ping = self
.run(|connection| Ok(connection.ping().map_err(Error::unavailable)))
.await
.and_then(|ping| ping)
.map_err(|error| litellm_error::Report::new(&error).to_string());
Ok(match ping {
Ok(true) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: format!("{label} connection test successful"),
error: None,
},
Ok(false) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("{label} ping returned False"),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("{label} connection failed: {error}"),
error: Some(error),
},
})
}
}
impl<S, C> DisconnectCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn disconnect(&self) -> Result<(), Error> {
self.connections.disconnect();
Ok(())
}
}
impl<S, C> ClientInfoCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type ClientList = String;
type Info = String;
fn client_list(&self) -> Result<Self::ClientList, Error> {
self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
}
fn info(&self) -> Result<Self::Info, Error> {
self.execute(|connection| connection.node_text(&redis::cmd("INFO")))
}
}

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

@ -1,226 +0,0 @@
use std::time::Duration;
use litellm_cache::{CacheCodec, Error, PopOperation, PushOperation, QueueCache, SetCache};
use crate::{cache::RedisCache, script::RedisArg};
pub type RedisRpushOperation = PushOperation<RedisArg>;
pub type RedisLpopOperation = PopOperation;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RedisLpopResult {
Missing,
Value(Vec<u8>),
Values(Vec<Vec<u8>>),
}
impl<S, C> SetCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type SetValue = RedisArg;
type SetResult = usize;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(ttl);
let mut pipeline = redis::pipe();
pipeline
.cmd("SADD")
.arg(&key)
.arg(values)
.cmd("EXPIRE")
.arg(&key)
.arg(ttl)
.ignore();
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
.map(|(added,)| added)
}
}
impl<S, C> QueueCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type QueueValue = RedisArg;
type PopResult = RedisLpopResult;
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
self.run(move |connection| {
redis::cmd("RPUSH")
.arg(key)
.arg(values)
.query(connection)
.map_err(Error::unavailable)
})
.await
}
async fn async_rpush_and_trim(
&self,
key: &str,
values: Vec<Self::QueueValue>,
max_len: usize,
) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len);
let mut pipeline = redis::pipe();
pipeline
.atomic()
.cmd("RPUSH")
.arg(&key)
.arg(values)
.cmd("LTRIM")
.arg(&key)
.arg(start)
.arg(-1)
.ignore();
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
.map(|(length,)| length)
}
async fn async_rpush_pipeline(
&self,
operations: Vec<RedisRpushOperation>,
) -> Result<Vec<usize>, Error> {
if operations.is_empty() {
return Ok(Vec::new());
}
let mut pipeline = redis::pipe();
for operation in operations {
if operation.values.is_empty() {
return Err(Error::InvalidEntry(None));
}
pipeline
.cmd("RPUSH")
.arg(self.namespaced_key(&operation.key))
.arg(operation.values);
}
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
}
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
if let Some(count) = count
&& self.major_version().await < 7
{
return self.lpop_one_at_a_time(key, count).await;
}
let command = lpop(self.namespaced_key(key), count);
let value = self
.run(move |connection| {
command
.query::<redis::Value>(connection)
.map_err(Error::unavailable)
})
.await?;
lpop_result(value, count.is_some())
}
async fn async_lpop_pipeline(
&self,
operations: Vec<RedisLpopOperation>,
) -> Result<Vec<Self::PopResult>, Error> {
if operations.is_empty() {
return Ok(Vec::new());
}
if operations.iter().any(|operation| operation.count.is_some())
&& self.major_version().await < 7
{
let mut results = Vec::with_capacity(operations.len());
for operation in &operations {
results.push(self.async_lpop(&operation.key, operation.count).await?);
}
return Ok(results);
}
let multiple = operations
.iter()
.map(|operation| operation.count.is_some())
.collect::<Vec<_>>();
let mut pipeline = redis::pipe();
for operation in operations {
pipeline.add_command(lpop(self.namespaced_key(&operation.key), operation.count));
}
self.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
.await?
.into_iter()
.zip(multiple)
.map(|(value, multiple)| lpop_result(value, multiple))
.collect()
}
}
fn lpop(key: String, count: Option<usize>) -> redis::Cmd {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
command.arg(count);
}
command
}
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
match value {
redis::Value::BulkString(bytes) => Ok(bytes),
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
_ => Err(Error::InvalidEntry(None)),
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
/// `handle_lpop_count_for_older_redis_versions`: `count` single-`LPOP` pipelines, keeping
/// only the values actually popped.
async fn lpop_one_at_a_time(&self, key: &str, count: usize) -> Result<RedisLpopResult, Error> {
let key = self.namespaced_key(key);
let mut values = Vec::new();
for _ in 0..count {
let mut pipeline = redis::pipe();
pipeline.add_command(lpop(key.clone(), None));
let replies = self
.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
.await?;
for reply in replies {
if reply != redis::Value::Nil {
values.push(redis_bytes(reply)?);
}
}
}
Ok(RedisLpopResult::Values(values))
}
}
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
match value {
redis::Value::Nil => Ok(RedisLpopResult::Missing),
redis::Value::Array(values) if multiple => values
.into_iter()
.map(redis_bytes)
.collect::<Result<Vec<_>, _>>()
.map(RedisLpopResult::Values),
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
_ => Err(Error::InvalidEntry(None)),
}
}

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

@ -1,136 +0,0 @@
use std::sync::Arc;
use litellm_cache::{CacheCodec, CacheScript, Error, ScriptCache};
use crate::{
cache::{RedisCache, namespaced_key},
connection::Connections,
};
#[derive(Clone, Debug, PartialEq)]
pub enum RedisArg {
Bytes(Vec<u8>),
Integer(i64),
Float(f64),
}
impl From<&str> for RedisArg {
fn from(value: &str) -> Self {
Self::Bytes(value.as_bytes().to_vec())
}
}
impl From<String> for RedisArg {
fn from(value: String) -> Self {
Self::Bytes(value.into_bytes())
}
}
impl From<Vec<u8>> for RedisArg {
fn from(value: Vec<u8>) -> Self {
Self::Bytes(value)
}
}
impl From<i64> for RedisArg {
fn from(value: i64) -> Self {
Self::Integer(value)
}
}
impl From<f64> for RedisArg {
fn from(value: f64) -> Self {
Self::Float(value)
}
}
impl redis::ToRedisArgs for RedisArg {
fn write_redis_args<W>(&self, out: &mut W)
where
W: ?Sized + redis::RedisWrite,
{
match self {
Self::Bytes(value) => value.write_redis_args(out),
Self::Integer(value) => value.write_redis_args(out),
Self::Float(value) => value.write_redis_args(out),
}
}
}
pub struct RedisScript<C> {
connections: Arc<Connections<C>>,
namespace: Option<String>,
source: String,
}
impl<C> CacheScript for RedisScript<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Argument = RedisArg;
type Output = redis::Value;
async fn invoke(
&self,
keys: Vec<String>,
arguments: Vec<Self::Argument>,
) -> Result<Self::Output, Error> {
let keys = keys
.into_iter()
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
.collect::<Vec<_>>();
let source = self.source.clone();
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
eval(connection, &source, keys, arguments)
})
.await
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub async fn async_eval(
&self,
script: String,
keys: Vec<String>,
arguments: Vec<RedisArg>,
) -> Result<redis::Value, Error> {
let keys = self.namespaced_keys(&keys);
self.run(move |connection| eval(connection, &script, keys, arguments))
.await
}
}
fn eval(
connection: &mut impl redis::ConnectionLike,
script: &str,
keys: Vec<String>,
arguments: Vec<RedisArg>,
) -> Result<redis::Value, Error> {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(Error::unavailable)
}
impl<S, C> ScriptCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Script = RedisScript<C>;
fn async_register_script(&self, source: String) -> Self::Script {
RedisScript {
connections: Arc::clone(&self.connections),
namespace: self.namespace.clone(),
source,
}
}
}

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

@ -1,232 +0,0 @@
use std::time::Duration;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheCodec, DeleteCache, Error,
ExactCacheContext, FlushAllCache, FlushCache, TtlPipelineCache,
};
use redis::Commands;
use crate::{cache::RedisCache, connection::ConnectionRef};
impl<S, C> BaseCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let ttl = self.ttl_or_default(context.ttl);
let key = self.namespaced_key(key);
self.execute(|connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(Error::unavailable)
})
}
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = self.execute(|connection| {
connection
.get::<_, redis::Value>(key)
.map_err(Error::unavailable)
})?;
self.decode_response(value)
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
self.run(move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(Error::unavailable)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = self
.run(move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(Error::unavailable)
})
.await?;
self.decode_response(value)
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, Self::Value)>,
context: ExactCacheContext,
) -> Result<(), Error> {
self.async_set_cache_pipeline_with_ttls(
cache_list
.into_iter()
.map(|(key, value)| (key, value, context.ttl))
.collect(),
)
.await
}
}
impl<S, C> TtlPipelineCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_set_cache_pipeline_with_ttls(
&self,
entries: Vec<(String, Self::Value, Option<Duration>)>,
) -> Result<(), Error> {
if entries.is_empty() {
return Ok(());
}
let mut pipeline = redis::pipe();
for (key, value, ttl) in entries {
pipeline
.cmd("SETEX")
.arg(self.namespaced_key(&key))
.arg(self.ttl_or_default(ttl))
.arg(self.codec.encode(&value)?)
.ignore();
}
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
}
}
impl<S, C> BatchCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_cache(
&self,
keys: &[String],
_: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = self.namespaced_keys(keys);
self.execute(|connection| mget(connection, keys))?
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = self.namespaced_keys(&keys);
self.run(move |connection| mget(connection, keys))
.await?
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
}
pub(crate) fn mget(
connection: &mut ConnectionRef<'_>,
keys: Vec<String>,
) -> Result<Vec<redis::Value>, Error> {
redis::cmd("MGET")
.arg(keys)
.query(connection)
.map_err(Error::unavailable)
}
impl<S, C> DeleteCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
.await
}
}
impl<S, C> BulkDeleteCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
if keys.is_empty() {
return Ok(0);
}
let keys = self.namespaced_keys(&keys);
self.run(move |connection| connection.del(keys).map_err(Error::unavailable))
.await
}
}
impl<S, C> FlushCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
self.execute(|connection| flush_matching(connection, &pattern))
}
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
self.run(move |connection| flush_matching(connection, &pattern))
.await
}
}
impl<S, C> FlushAllCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn flushall(&self) -> Result<(), Error> {
self.execute(|connection| connection.flushall())
}
}
fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
connection.scan(pattern, 1000, |connection, keys| {
if !keys.is_empty() {
connection
.del::<_, usize>(keys)
.map_err(Error::unavailable)?;
}
Ok(true)
})
}

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

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

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

File diff suppressed because it is too large Load diff

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

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

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

@ -1,95 +0,0 @@
//! The shared cache contracts, run against the in-process fake connection and, when
//! `LITELLM_TEST_REDIS_CLUSTER_NODES` is set, against a live Redis Cluster.
mod support;
use std::time::Duration;
use litellm_cache::{ExactCacheContext, JsonCodec};
use litellm_cache_testing as contract;
use rstest::rstest;
use serde_json::json;
use support::{JsonCache, cluster_cache, fake_cache};
const PREFIX: &str = "contract:";
#[derive(Clone, Copy, Debug)]
enum Contract {
HitAndMiss,
SyncAsyncEquivalence,
OverwriteReplaces,
PipelineWritesEveryEntry,
BatchPreservesOrder,
DeleteRemovesKey,
FlushClears,
CounterAccumulates,
}
#[derive(Clone, Copy, Debug)]
enum Server {
Fake,
Cluster,
}
async fn check<C>(contract: Contract, cache: &JsonCache<C>)
where
C: redis::ConnectionLike + Send + 'static,
{
let context = ExactCacheContext::default();
match contract {
Contract::HitAndMiss => {
contract::hit_and_miss(cache, context, PREFIX, json!({"answer": 42})).await
}
Contract::SyncAsyncEquivalence => {
contract::sync_async_equivalence(cache, context, PREFIX, json!("first"), json!([2]))
.await
}
Contract::OverwriteReplaces => {
contract::overwrite_replaces(cache, context, PREFIX, json!(1), json!({"b": 2})).await
}
Contract::PipelineWritesEveryEntry => {
contract::pipeline_writes_every_entry(
cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await
}
Contract::BatchPreservesOrder => {
contract::batch_preserves_order(cache, context, PREFIX, json!("first"), json!(2)).await
}
Contract::DeleteRemovesKey => {
contract::delete_removes_key(cache, context, PREFIX, json!("value")).await
}
Contract::FlushClears => {
contract::flush_clears(cache, context, PREFIX, json!("value")).await
}
Contract::CounterAccumulates => contract::counter_accumulates(cache, context, PREFIX).await,
}
}
#[rstest]
#[case::hit_and_miss(Contract::HitAndMiss)]
#[case::sync_async_equivalence(Contract::SyncAsyncEquivalence)]
#[case::overwrite_replaces(Contract::OverwriteReplaces)]
#[case::pipeline_writes_every_entry(Contract::PipelineWritesEveryEntry)]
#[case::batch_preserves_order(Contract::BatchPreservesOrder)]
#[case::delete_removes_key(Contract::DeleteRemovesKey)]
#[case::flush_clears(Contract::FlushClears)]
#[case::counter_accumulates(Contract::CounterAccumulates)]
#[tokio::test]
async fn redis_satisfies_the_cache_contract(
#[case] contract: Contract,
#[values(Server::Fake, Server::Cluster)] server: Server,
) {
match server {
Server::Fake => check(contract, &fake_cache("contract")).await,
Server::Cluster => {
let label = format!("{contract:?}");
if let Some(cache) = cluster_cache(&label, Duration::from_secs(120), JsonCodec::new()) {
check(contract, &cache).await;
}
}
}
}

231
litellm-rust/crates/cache-redis/tests/support/mod.rs
View file

@ -1,231 +0,0 @@
#![allow(dead_code)]
use std::{
collections::BTreeMap,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{CacheCodec, Error, JsonCodec};
use litellm_cache_redis::{RedisCache, RedisNode, RedisTopology};
/// Encodes a byte behind a tag, so a value written with another tag decodes as invalid.
pub struct TaggedByteCodec(pub u8);
impl CacheCodec for TaggedByteCodec {
type Value = u8;
fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
if *value > 127 {
return Err(Error::InvalidEntry(None));
}
Ok(vec![self.0, *value])
}
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
match bytes {
[tag, value] if *tag == self.0 => Ok(*value),
_ => Err(Error::InvalidEntry(None)),
}
}
}
/// A stateful in-process stand-in for a Redis server that understands the string commands the
/// shared contracts exercise, so they run without a live server. TTLs are accepted and ignored.
#[derive(Default)]
pub struct FakeRedis {
strings: BTreeMap<Vec<u8>, Vec<u8>>,
}
impl FakeRedis {
fn run(&mut self, command: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
let name = String::from_utf8_lossy(&command[0]).to_ascii_uppercase();
let args = &command[1..];
Ok(match name.as_str() {
"PING" => redis::Value::SimpleString("PONG".into()),
"SET" | "SETEX" => {
let value = if name == "SET" { &args[1] } else { &args[2] };
self.strings.insert(args[0].clone(), value.clone());
redis::Value::Okay
}
"GET" => self.get(&args[0]),
"MGET" => redis::Value::Array(args.iter().map(|key| self.get(key)).collect()),
"DEL" => {
let removed = args
.iter()
.filter(|key| self.strings.remove(*key).is_some())
.count();
redis::Value::Int(removed as i64)
}
"SCAN" => {
let pattern = &args[2];
let keys = self
.strings
.keys()
.filter(|key| glob(pattern, key))
.map(|key| redis::Value::BulkString(key.clone()))
.collect();
redis::Value::Array(vec![
redis::Value::BulkString(b"0".to_vec()),
redis::Value::Array(keys),
])
}
"EVAL" if args[0].windows(11).any(|window| window == b"INCRBYFLOAT") => {
self.increment_by_float(&args[2], &args[3])
}
"INCRBYFLOAT" => self.increment_by_float(&args[0], &args[1]),
_ => {
return Err(redis::RedisError::from((
redis::ErrorKind::Client,
"unsupported command",
name,
)));
}
})
}
fn get(&self, key: &[u8]) -> redis::Value {
self.strings.get(key).map_or(redis::Value::Nil, |value| {
redis::Value::BulkString(value.clone())
})
}
fn increment_by_float(&mut self, key: &[u8], amount: &[u8]) -> redis::Value {
let current = self
.strings
.get(key)
.map_or(0.0, |value| parse_float(value));
let total = format!("{}", current + parse_float(amount));
self.strings
.insert(key.to_vec(), total.clone().into_bytes());
redis::Value::BulkString(total.into_bytes())
}
}
fn parse_float(bytes: &[u8]) -> f64 {
std::str::from_utf8(bytes).unwrap().parse().unwrap()
}
/// Redis `MATCH` globbing for `*`, `?` and backslash escapes.
fn glob(pattern: &[u8], key: &[u8]) -> bool {
match pattern.split_first() {
None => key.is_empty(),
Some((b'*', rest)) => (0..=key.len()).any(|skip| glob(rest, &key[skip..])),
Some((b'?', rest)) => !key.is_empty() && glob(rest, &key[1..]),
Some((b'\\', [escaped, rest @ ..])) => {
key.first() == Some(escaped) && glob(rest, &key[1..])
}
Some((literal, rest)) => key.first() == Some(literal) && glob(rest, &key[1..]),
}
}
/// Splits RESP request bytes into the commands they carry.
fn commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
fn line<'a>(bytes: &mut &'a [u8]) -> &'a [u8] {
let end = bytes
.windows(2)
.position(|window| window == b"\r\n")
.unwrap();
let (line, rest) = bytes.split_at(end);
*bytes = &rest[2..];
line
}
fn length(line: &[u8]) -> usize {
std::str::from_utf8(&line[1..]).unwrap().parse().unwrap()
}
let mut commands = Vec::new();
while !bytes.is_empty() {
let count = length(line(&mut bytes));
let command = (0..count)
.map(|_| {
let size = length(line(&mut bytes));
let (argument, rest) = bytes.split_at(size);
bytes = &rest[2..];
argument.to_vec()
})
.collect();
commands.push(command);
}
commands
}
impl redis::ConnectionLike for FakeRedis {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
let command = commands(cmd).into_iter().next().unwrap();
self.run(command)
}
fn req_packed_commands(
&mut self,
cmd: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
let replies = commands(cmd)
.into_iter()
.map(|command| self.run(command))
.collect::<redis::RedisResult<Vec<_>>>()?;
Ok(replies.into_iter().skip(offset).take(count).collect())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}
pub type JsonCache<C = redis::Connection> = RedisCache<JsonCodec<serde_json::Value>, C>;
pub fn fake_cache(namespace: &str) -> JsonCache<FakeRedis> {
RedisCache::with_connection(FakeRedis::default(), None, JsonCodec::new())
.with_namespace(Some(namespace.into()))
}
/// Startup nodes from `LITELLM_TEST_REDIS_CLUSTER_NODES` (`host:port,host:port`); tests that
/// need a live cluster skip when it is unset.
pub fn cluster_topology() -> Option<RedisTopology> {
let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
let startup_nodes = nodes
.split(',')
.map(|node| {
let (host, port) = node.trim().rsplit_once(':').expect("host:port");
RedisNode {
host: host.to_string(),
port: port.parse().expect("port"),
}
})
.collect();
Some(RedisTopology::Cluster { startup_nodes })
}
pub fn cluster_url() -> String {
std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
.unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
}
pub fn unique_namespace(label: &str) -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
format!("cluster-test:{label}:{nanos}")
}
pub fn cluster_cache<S: CacheCodec>(
label: &str,
default_ttl: Duration,
codec: S,
) -> Option<RedisCache<S>> {
let topology = cluster_topology()?;
Some(
RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec)
.expect("cluster connection")
.with_namespace(Some(unique_namespace(label))),
)
}

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

@ -16,9 +16,6 @@ sha2.workspace = true
litellm-cache-gcs.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
redis.workspace = true
redis-test.workspace = true
rstest.workspace = true
tokio.workspace = true
wiremock.workspace = true

72
litellm-rust/crates/cache-response/tests/connection.rs
View file

@ -2,79 +2,11 @@ mod support;
use std::{sync::Arc, time::Duration};
use litellm_cache::CacheConnectionStatus;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{
CacheEntry, ConnectionProbe, ExactResponseCache, ResponseCache, ResponseCacheRequest,
};
use redis_test::MockCmd;
use litellm_cache_response::{CacheEntry, ExactResponseCache, ResponseCache, ResponseCacheRequest};
use rstest::rstest;
use serde_json::json;
use support::{keyed, memory, redis, request};
#[rstest]
#[case::reachable(
Ok("PONG"),
CacheConnectionStatus::Success,
"Redis connection test successful",
false
)]
#[case::unexpected_reply(
Ok("NOPE"),
CacheConnectionStatus::Failed,
"Redis ping returned False",
false
)]
#[case::connection_refused(
Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))),
CacheConnectionStatus::Failed,
"Redis connection failed:",
true
)]
#[tokio::test]
async fn connection_backends_are_reachable_as_a_probe(
#[case] reply: redis::RedisResult<&'static str>,
#[case] status: CacheConnectionStatus,
#[case] message: &str,
#[case] has_error: bool,
) {
let probe: Arc<dyn ConnectionProbe> =
Arc::new(redis(vec![MockCmd::new(redis::cmd("PING"), reply)], None));
let result = probe.test_connection().await.unwrap();
assert_eq!(result.status, status);
assert!(result.message.starts_with(message), "{}", result.message);
assert_eq!(result.error.is_some(), has_error);
}
#[rstest]
#[tokio::test]
async fn one_service_serves_both_the_exact_cache_and_its_probe(request: ResponseCacheRequest) {
let service = Arc::new(redis(
vec![
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(br#"{"timestamp":100.0,"response":{"ok":true}}"#.to_vec()),
),
],
Some("tenant"),
));
let probe: Arc<dyn ConnectionProbe> = service.clone();
let exact: Arc<dyn ExactResponseCache> = service;
assert_eq!(
probe.test_connection().await.unwrap().status,
CacheConnectionStatus::Success
);
assert_eq!(
exact
.async_lookup(&request, Duration::from_secs(100))
.await
.unwrap(),
Some(json!({"ok": true}))
);
}
use support::{keyed, memory, request};
/// The in-memory backend has no `test_connection`, as in Python, and still serves every response
/// operation.

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

@ -17,10 +17,9 @@ use litellm_cache_response::{
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
WriteBuffer, cache_key,
};
use redis_test::MockCmd;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{keyed, memory, redis, request};
use support::{keyed, memory, request};
type Memory = Arc<ResponseCache<InMemoryCache<CacheEntry>>>;
@ -331,79 +330,6 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent(
);
}
#[rstest]
#[case::python_sync_literal(
br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice()
)]
#[case::python_async_json(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice())]
#[tokio::test]
async fn redis_consumer_reads_python_sync_and_async_envelopes(
request: ResponseCacheRequest,
#[case] stored: &[u8],
#[values(false, true)] asynchronous: bool,
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(stored.to_vec()),
)],
Some("tenant"),
);
let now = Duration::from_secs(101);
let found = if asynchronous {
cache.async_lookup(&request, now).await.unwrap()
} else {
cache.lookup(&request, now).unwrap()
};
assert_eq!(found, Some(json!({"ok": true, "text": "cached"})));
}
#[rstest]
#[case::object(
json!({"ok": true, "text": "cached"}),
br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()
)]
#[case::array(json!([1, 2]), br#"{"timestamp":100.0,"response":"[1,2]"}"#.as_slice())]
#[tokio::test]
async fn redis_consumer_writes_python_compatible_json(
request: ResponseCacheRequest,
#[case] response: Value,
#[case] wire: &[u8],
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("SETEX").arg("tenant:key").arg(600).arg(wire),
Ok("OK"),
)],
Some("tenant"),
);
cache
.async_store(&request, response, Duration::from_secs(100))
.await
.unwrap();
}
#[rstest]
#[tokio::test]
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis(
mut request: ResponseCacheRequest,
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(b"invalid".to_vec()),
)],
None,
);
request.controls.no_cache = true;
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
request.controls.no_cache = false;
assert_eq!(
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn captured_service_keeps_the_selected_backend_for_background_writes(
@ -541,27 +467,6 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses(
assert_eq!(memory.lookup(&requests[2], now).unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn batch_lookup_with_no_readable_request_skips_the_backend(
#[values(false, true)] asynchronous: bool,
) {
let cache = redis(Vec::new(), None);
let mut request = keyed("key");
request.controls.no_cache = true;
let requests = [request.clone(), request];
let partial = if asynchronous {
cache
.async_lookup_batch(&requests, Duration::ZERO)
.await
.unwrap()
} else {
cache.lookup_batch(&requests, Duration::ZERO).unwrap()
};
assert_eq!(partial.values, vec![None, None]);
assert_eq!(partial.missing_indices, vec![0, 1]);
}
#[rstest]
#[tokio::test]
async fn deferred_entries_keep_the_time_they_were_produced(

17
litellm-rust/crates/cache-response/tests/support/mod.rs
View file

@ -1,15 +1,9 @@
use std::{sync::Arc, time::Duration};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::RedisCache;
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
};
use redis_test::{MockCmd, MockRedisConnection};
use litellm_cache_response::{CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheRequest};
use rstest::fixture;
pub type MockedRedis = RedisCache<ResponseCacheCodec, MockRedisConnection>;
#[fixture]
pub fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
@ -29,12 +23,3 @@ pub fn keyed(key: &str) -> ResponseCacheRequest {
..Default::default()
})
}
/// A Redis response cache that must receive exactly `commands`, in order.
pub fn redis(commands: Vec<MockCmd>, namespace: Option<&str>) -> ResponseCache<MockedRedis> {
let connection = MockRedisConnection::new(commands).assert_all_commands_consumed();
ResponseCache::new(Arc::new(
RedisCache::with_connection(connection, None, ResponseCacheCodec)
.with_namespace(namespace.map(str::to_owned)),
))
}

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

@ -1,22 +0,0 @@
[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
redis = { workspace = true, features = ["tls-rustls"] }
sha2.workspace = true
thiserror.workspace = true
uuid = { version = "1", features = ["v4"] }
[dev-dependencies]
litellm-error.workspace = true
litellm-cache-testing.workspace = true
redis-test.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true

245
litellm-rust/crates/cache-valkey-semantic/src/cache.rs
View file

@ -1,245 +0,0 @@
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{
BaseCache, CacheCodec, Error, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
};
use litellm_cache_redis::{RedisTopology, connection::Connections};
use crate::{
ValkeySemanticConfig,
index::IndexState,
search::{embedding_bytes, scope_tag, search_document, write_document},
};
/// `ValkeySemanticCache`: a semantic cache on valkey-search's TAG + VECTOR index. Values go
/// through the injected codec, so the response layer decides what a cached entry is.
pub struct ValkeySemanticCache<E, S, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
codec: S,
config: ValkeySemanticConfig,
index_dimension: Arc<Mutex<Option<usize>>>,
}
impl<E: Embedder, S: CacheCodec> ValkeySemanticCache<E, S> {
pub fn new(
url: &str,
embedder: E,
codec: S,
config: ValkeySemanticConfig,
) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
codec,
config,
index_dimension: Arc::new(Mutex::new(None)),
})
}
}
impl<E, S, C> ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(
connection: C,
embedder: E,
codec: S,
config: ValkeySemanticConfig,
) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
embedder,
codec,
config,
index_dimension: Arc::new(Mutex::new(None)),
}
}
pub fn similarity_threshold(&self) -> f64 {
self.config.similarity_threshold
}
pub fn index_name(&self) -> &str {
&self.config.index_name
}
fn index_state(&self) -> IndexState {
IndexState {
name: self.config.index_name.clone(),
prefix: format!("{}:", self.config.index_name),
dimension: Arc::clone(&self.index_dimension),
similarity_threshold: self.config.similarity_threshold,
}
}
fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
Ok(SemanticLookup {
value: lookup
.value
.map(|bytes| self.codec.decode(&bytes))
.transpose()?,
similarity: lookup.similarity,
})
}
}
impl<E, S, C> ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec + Clone,
C: redis::ConnectionLike + Send + 'static,
{
/// The same index and connections behind a different embedder.
pub fn with_embedder<E2: Embedder>(&self, embedder: E2) -> ValkeySemanticCache<E2, S, C> {
ValkeySemanticCache {
connections: Arc::clone(&self.connections),
embedder,
codec: self.codec.clone(),
config: self.config.clone(),
index_dimension: Arc::clone(&self.index_dimension),
}
}
}
impl<E, S, C> BaseCache for ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = SemanticCacheContext;
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> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(());
};
let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let response = self.codec.encode(&value)?;
let index = self.index_state();
self.connections.execute(|connection| {
write_document(
connection,
&index,
&scope_tag(key),
&prompt,
response,
embedding_bytes(&embedding),
self.get_ttl(context),
)
})
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
let Some(prompt) = prompt_from_context(&context) else {
return Ok(());
};
let embedding = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let response = self.codec.encode(&value)?;
let index = self.index_state();
let scope = scope_tag(key);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
write_document(
connection,
&index,
&scope,
&prompt,
response,
embedding_bytes(&embedding),
context.ttl,
)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.async_get_cache_with_similarity(key, context)
.await
.map(|lookup| lookup.value)
}
}
/// Python stamps a similarity of `0.0` when there is no prompt or no document in the key's
/// scope, and the closest document's similarity even when it misses the threshold.
impl<E, S, C> SemanticCache for ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let index = self.index_state();
let lookup = self.connections.execute(|connection| {
search_document(
connection,
&index,
&scope_tag(key),
embedding_bytes(&embedding),
)
})?;
self.decode(lookup)
}
async fn async_get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let embedding = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let index = self.index_state();
let scope = scope_tag(key);
let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
search_document(connection, &index, &scope, embedding_bytes(&embedding))
})
.await?;
self.decode(lookup)
}
}

8
litellm-rust/crates/cache-valkey-semantic/src/config.rs
View file

@ -1,8 +0,0 @@
/// `ValkeySemanticCache.DEFAULT_VALKEY_INDEX_NAME`.
pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
#[derive(Clone, Debug, PartialEq)]
pub struct ValkeySemanticConfig {
pub similarity_threshold: f64,
pub index_name: String,
}

117
litellm-rust/crates/cache-valkey-semantic/src/index.rs
View file

@ -1,117 +0,0 @@
use std::sync::{Arc, Mutex};
use litellm_cache::Error;
use litellm_cache_redis::connection::ConnectionRef;
use crate::search::value_text;
/// The valkey-search index one cache writes to, with the dimension it was last ensured for.
#[derive(Clone)]
pub(crate) struct IndexState {
pub(crate) name: String,
pub(crate) prefix: String,
pub(crate) dimension: Arc<Mutex<Option<usize>>>,
pub(crate) similarity_threshold: f64,
}
/// `_ensure_index_sync` / `_ensure_index_async`: create the TAG + HNSW index once per dimension,
/// and accept an existing index unless it reports a different dimension.
pub(crate) fn ensure_index(
connection: &mut ConnectionRef<'_>,
index: &IndexState,
dimension: usize,
) -> Result<(), Error> {
if index
.dimension
.lock()?
.is_some_and(|existing| existing == dimension)
{
return Ok(());
}
let create = redis::cmd("FT.CREATE")
.arg(&index.name)
.arg("ON")
.arg("HASH")
.arg("PREFIX")
.arg(1)
.arg(&index.prefix)
.arg("SCHEMA")
.arg("litellm_cache_key")
.arg("TAG")
.arg("embedding")
.arg("VECTOR")
.arg("HNSW")
.arg(6)
.arg("TYPE")
.arg("FLOAT32")
.arg("DIM")
.arg(dimension)
.arg("DISTANCE_METRIC")
.arg("COSINE")
.query::<String>(connection)
.map(|_| ());
if let Err(error) = create {
if !index_exists(&error) {
return Err(Error::unavailable(error));
}
let info = redis::cmd("FT.INFO")
.arg(&index.name)
.query::<redis::Value>(connection)
.map_err(Error::unavailable)?;
if let Some(existing) =
index_dimension_from_info(&info).filter(|existing| *existing != dimension)
{
return Err(Error::unavailable(DimensionMismatch {
existing,
dimension,
}));
}
}
*index.dimension.lock()? = Some(dimension);
Ok(())
}
fn index_exists(error: &redis::RedisError) -> bool {
error.code().is_some()
&& error
.detail()
.is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists"))
}
#[derive(Debug, thiserror::Error)]
#[error("index exists with dimension {existing}, not {dimension}")]
struct DimensionMismatch {
existing: usize,
dimension: usize,
}
/// `_extract_index_dim`: flatten each attribute one level and read the value after
/// `dimensions`.
fn index_dimension_from_info(value: &redis::Value) -> Option<usize> {
let redis::Value::Array(values) = value else {
return None;
};
let attributes = values.windows(2).find_map(|pair| {
(value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1])
})?;
let redis::Value::Array(fields) = attributes else {
return None;
};
fields.iter().find_map(|field| {
let redis::Value::Array(values) = field else {
return None;
};
let values = values
.iter()
.flat_map(|value| match value {
redis::Value::Array(values) => values.as_slice(),
_ => std::slice::from_ref(value),
})
.collect::<Vec<_>>();
values.windows(2).find_map(|pair| {
(value_text(pair[0]).as_deref() == Some("dimensions"))
.then(|| value_text(pair[1]).and_then(|value| value.parse().ok()))
.flatten()
})
})
}

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

@ -1,7 +0,0 @@
mod cache;
mod config;
mod index;
mod search;
pub use cache::ValkeySemanticCache;
pub use config::{DEFAULT_INDEX_NAME, ValkeySemanticConfig};

161
litellm-rust/crates/cache-valkey-semantic/src/search.rs
View file

@ -1,161 +0,0 @@
use std::time::Duration;
use litellm_cache::{Error, semantic::SemanticLookup};
use litellm_cache_redis::connection::ConnectionRef;
use sha2::{Digest, Sha256};
use uuid::Uuid;
use crate::index::{IndexState, ensure_index};
/// `_scope_tag`: valkey-search TAG fields cannot match arbitrary keys verbatim, so scopes are
/// the key's lowercase SHA-256.
pub(crate) fn scope_tag(key: &str) -> String {
let digest = Sha256::digest(key.as_bytes());
digest.iter().map(|byte| format!("{byte:02x}")).collect()
}
pub(crate) fn embedding_bytes(embedding: &[f32]) -> Vec<u8> {
embedding
.iter()
.flat_map(|value| value.to_le_bytes())
.collect()
}
/// `HSET` a fresh `<prefix><scope>:<uuid4>` document, then `EXPIRE` it when a TTL is set.
pub(crate) fn write_document(
connection: &mut ConnectionRef<'_>,
index: &IndexState,
scope: &str,
prompt: &str,
response: Vec<u8>,
vector: Vec<u8>,
ttl: Option<Duration>,
) -> Result<(), Error> {
ensure_index(connection, index, vector.len() / size_of::<f32>())?;
let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4());
let mut pipeline = redis::pipe();
pipeline
.cmd("HSET")
.arg(&document)
.arg("litellm_cache_key")
.arg(scope)
.arg("prompt")
.arg(prompt)
.arg("response")
.arg(response)
.arg("embedding")
.arg(vector)
.ignore();
if let Some(ttl) = ttl {
pipeline
.cmd("EXPIRE")
.arg(&document)
.arg(ttl.as_secs())
.ignore();
}
pipeline.query::<()>(connection).map_err(Error::unavailable)
}
/// The KNN-1 search within `scope`: the closest document's similarity, and its stored response
/// when that similarity reaches the threshold. No document reads as a similarity of `0.0`.
pub(crate) fn search_document(
connection: &mut ConnectionRef<'_>,
index: &IndexState,
scope: &str,
vector: Vec<u8>,
) -> Result<SemanticLookup<Vec<u8>>, Error> {
ensure_index(connection, index, vector.len() / size_of::<f32>())?;
let query =
format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]");
let response = redis::cmd("FT.SEARCH")
.arg(&index.name)
.arg(query)
.arg("PARAMS")
.arg(2)
.arg("vec")
.arg(vector)
.arg("RETURN")
.arg(2)
.arg("response")
.arg("vector_distance")
.arg("DIALECT")
.arg(2)
.query::<redis::Value>(connection)
.map_err(Error::unavailable)?;
let Some(fields) = search_fields(response)? else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let field = |name: &str| {
fields
.iter()
.find_map(|(field, value)| (field == name).then(|| value.clone()))
.ok_or(Error::InvalidEntry(None))
};
let response = field("response")?;
let similarity = 1.0 - parse_f64(&field("vector_distance")?)?;
Ok(SemanticLookup {
value: (similarity >= index.similarity_threshold).then_some(response),
similarity: Some(similarity),
})
}
type SearchFields = Vec<(String, Vec<u8>)>;
fn search_fields(value: redis::Value) -> Result<Option<SearchFields>, Error> {
let redis::Value::Array(values) = value else {
return Err(Error::InvalidEntry(None));
};
let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?;
if total <= 0 || values.len() < 3 {
return Ok(None);
}
let redis::Value::Array(fields) = &values[2] else {
return Err(Error::InvalidEntry(None));
};
let (pairs, remainder) = fields.as_chunks::<2>();
if !remainder.is_empty() {
return Err(Error::InvalidEntry(None));
}
pairs
.iter()
.map(|pair| {
Ok((
value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?,
value_bytes(&pair[1])?,
))
})
.collect::<Result<Vec<_>, Error>>()
.map(Some)
}
fn parse_i64(value: &redis::Value) -> Result<i64, Error> {
value_text(value)
.ok_or(Error::InvalidEntry(None))?
.parse()
.map_err(Error::invalid_entry)
}
fn parse_f64(value: &[u8]) -> Result<f64, Error> {
std::str::from_utf8(value)
.map_err(Error::invalid_entry)?
.parse()
.map_err(Error::invalid_entry)
}
pub(crate) fn value_text(value: &redis::Value) -> Option<String> {
match value {
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
redis::Value::SimpleString(value) => Some(value.clone()),
redis::Value::Int(value) => Some(value.to_string()),
_ => None,
}
}
fn value_bytes(value: &redis::Value) -> Result<Vec<u8>, Error> {
match value {
redis::Value::BulkString(bytes) => Ok(bytes.clone()),
redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()),
redis::Value::Int(value) => Ok(value.to_string().into_bytes()),
_ => Err(Error::InvalidEntry(None)),
}
}

430
litellm-rust/crates/cache-valkey-semantic/tests/cache.rs
View file

@ -1,430 +0,0 @@
mod support;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{
BaseCache, Error, JsonCodec, SemanticCacheContext,
semantic::{PreparedEmbedding, SemanticCache, SemanticLookup},
};
use litellm_cache_valkey_semantic::{
DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig,
};
use redis_test::MockRedisConnection;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{EmbedCalls, FakeEmbedder, RecordingConnection};
type Requests = Arc<Mutex<Vec<Vec<u8>>>>;
type RecordingCache = ValkeySemanticCache<FakeEmbedder, JsonCodec<Value>, RecordingConnection>;
const KEY_SCOPE: &str = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683";
struct Recording {
cache: RecordingCache,
requests: Requests,
calls: EmbedCalls,
}
impl Recording {
fn text(&self) -> String {
self.requests
.lock()
.unwrap()
.iter()
.map(|request| String::from_utf8_lossy(request).into_owned())
.collect::<Vec<_>>()
.join("\n")
}
}
fn config() -> ValkeySemanticConfig {
ValkeySemanticConfig {
similarity_threshold: 0.8,
index_name: "test".into(),
}
}
fn recording(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Recording {
let connection = RecordingConnection::new(replies);
let requests = connection.requests();
let embedder = FakeEmbedder::new(&[]);
let calls = Arc::clone(&embedder.calls);
Recording {
cache: ValkeySemanticCache::with_connection(
connection,
embedder,
JsonCodec::new(),
config(),
),
requests,
calls,
}
}
#[fixture]
fn entry() -> Value {
json!({"timestamp": 1.0, "response": {"answer": "ok"}})
}
#[fixture]
fn context() -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "hello"}])),
metadata: Some(json!({"source": "test"})),
..Default::default()
}
}
fn ok() -> redis::RedisResult<redis::Value> {
Ok(redis::Value::SimpleString("OK".into()))
}
fn already_exists() -> redis::RedisResult<redis::Value> {
Err(support::server_error("Index already exists"))
}
fn bulk(value: &[u8]) -> redis::Value {
redis::Value::BulkString(value.to_vec())
}
fn search_hit(response: &[u8], distance: &str) -> redis::Value {
redis::Value::Array(vec![
redis::Value::Int(1),
bulk(b"test:document"),
redis::Value::Array(vec![
bulk(b"response"),
bulk(response),
bulk(b"vector_distance"),
bulk(distance.as_bytes()),
]),
])
}
fn encoded(value: &Value) -> Vec<u8> {
serde_json::to_vec(value).unwrap()
}
/// `FT.INFO` with the vector field's dimension nested one level down, as valkey-search reports.
fn nested_dimension_info(dimension: i64) -> redis::Value {
redis::Value::Array(vec![
redis::Value::SimpleString("attributes".into()),
redis::Value::Array(vec![redis::Value::Array(vec![
redis::Value::SimpleString("embedding".into()),
redis::Value::Array(vec![
redis::Value::SimpleString("dimensions".into()),
redis::Value::Int(dimension),
]),
])]),
])
}
/// `FT.INFO` with `dimensions` as a sibling string of the identifier.
fn flat_dimension_info(dimension: &str) -> redis::Value {
redis::Value::Array(vec![
redis::Value::SimpleString("attributes".into()),
redis::Value::Array(vec![redis::Value::Array(vec![
redis::Value::SimpleString("identifier".into()),
redis::Value::SimpleString("embedding".into()),
redis::Value::Array(vec![
redis::Value::SimpleString("dimensions".into()),
redis::Value::SimpleString(dimension.into()),
]),
])]),
])
}
#[rstest]
#[case::string_content(json!([{"content": "hello"}]), None, Some("hello"))]
#[case::text_parts(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))]
#[case::non_object_parts_are_skipped(json!([{"content": ["raw", {"text": "hello"}]}]), None, Some("hello"))]
#[case::search_results(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))]
#[case::responses_string_input(json!([]), Some(json!(" hello ")), Some("hello"))]
#[case::responses_item_input(json!([]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))]
#[case::blank_input(json!([]), Some(json!(" ")), None)]
fn prompt_shapes_follow_redis_semantic_extraction(
#[case] messages: Value,
#[case] input: Option<Value>,
#[case] expected: Option<&str>,
) {
let recording = recording([ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))]);
let context = SemanticCacheContext {
messages: Some(messages),
input,
..Default::default()
};
assert_eq!(recording.cache.get_cache("key", &context).unwrap(), None);
let prompts = recording
.calls
.lock()
.unwrap()
.iter()
.map(|(prompt, _)| prompt.clone())
.collect::<Vec<_>>();
assert_eq!(
prompts,
expected.into_iter().map(str::to_owned).collect::<Vec<_>>()
);
assert_eq!(
recording.requests.lock().unwrap().is_empty(),
expected.is_none()
);
}
#[rstest]
#[case::key("key", KEY_SCOPE)]
#[case::empty_key("", "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")]
fn documents_are_scoped_by_the_keys_sha256(
#[case] key: &str,
#[case] scope: &str,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([ok()]);
recording.cache.set_cache(key, entry, &context).unwrap();
let text = recording.text();
assert!(text.contains(&format!("test:{scope}:")));
assert!(text.contains(&format!("litellm_cache_key\r\n$64\r\n{scope}\r\n")));
}
#[rstest]
#[case::no_ttl(None, None)]
#[case::whole_seconds(Some(Duration::from_secs(5)), Some("5"))]
#[case::fractional_seconds_truncate(Some(Duration::from_millis(1900)), Some("1"))]
fn set_writes_hset_and_expires_only_with_a_ttl(
#[case] ttl: Option<Duration>,
#[case] expire: Option<&str>,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([ok()]);
recording
.cache
.set_cache(
"key",
entry,
&SemanticCacheContext {
ttl,
..context.clone()
},
)
.unwrap();
let text = recording.text();
assert!(text.contains("FT.CREATE"));
assert!(text.contains("HSET"));
let expire_seconds = text
.split_once("EXPIRE\r\n")
.and_then(|(_, rest)| rest.split("\r\n").nth(3));
assert_eq!(expire_seconds, expire);
assert_eq!(
*recording.calls.lock().unwrap(),
vec![("hello".to_owned(), context.metadata)]
);
}
#[rstest]
fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: SemanticCacheContext) {
let recording = recording([ok()]);
recording
.cache
.set_cache("key", entry.clone(), &context)
.unwrap();
recording.cache.set_cache("key", entry, &context).unwrap();
let text = recording.text();
assert_eq!(text.matches("FT.CREATE").count(), 1);
assert_eq!(text.matches("HSET").count(), 2);
}
#[rstest]
#[case::nested_matching(nested_dimension_info(3), Ok(()))]
#[case::nested_mismatch(
nested_dimension_info(2),
Err("cache is unavailable: index exists with dimension 2, not 3")
)]
#[case::flat_matching(flat_dimension_info("3"), Ok(()))]
#[case::flat_mismatch(
flat_dimension_info("2"),
Err("cache is unavailable: index exists with dimension 2, not 3")
)]
#[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))]
fn existing_index_dimension_must_match_embedding(
#[case] info: redis::Value,
#[case] expected: Result<(), &str>,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([already_exists(), Ok(info)]);
let result = recording.cache.set_cache("key", entry, &context);
assert_eq!(
result.map_err(|error| litellm_error::Report::new(&error).to_string()),
expected.map_err(str::to_owned)
);
}
#[rstest]
#[case::create_failure(Err(redis::RedisError::from((redis::ErrorKind::Io, "boom"))))]
fn index_creation_failures_are_unavailable(
#[case] reply: redis::RedisResult<redis::Value>,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([reply]);
assert!(matches!(
recording.cache.set_cache("key", entry, &context),
Err(Error::Unavailable(Some(_)))
));
}
#[rstest]
#[case::within_threshold(search_hit(&encoded(&entry()), "0.1"), Ok(Some(entry())))]
#[case::at_threshold(search_hit(&encoded(&entry()), "0.2"), Ok(Some(entry())))]
#[case::beyond_threshold(search_hit(&encoded(&entry()), "0.5"), Ok(None))]
#[case::zero_documents(redis::Value::Array(vec![redis::Value::Int(0)]), Ok(None))]
#[case::missing_response(
redis::Value::Array(vec![
redis::Value::Int(1),
bulk(b"document"),
redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]),
]),
Err(Error::InvalidEntry(None))
)]
#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))]
#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))]
fn get_applies_threshold_and_decodes_entry(
#[case] reply: redis::Value,
#[case] expected: Result<Option<Value>, Error>,
context: SemanticCacheContext,
) {
let recording = recording([ok(), Ok(reply)]);
let kind = |result: Result<Option<Value>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(
kind(recording.cache.get_cache("key", &context)),
kind(expected)
);
}
#[rstest]
#[case::hit(context(), Some(search_hit(&encoded(&entry()), "0.1")), Some(entry()), Some(1.0 - 0.1))]
#[case::below_threshold(context(), Some(search_hit(&encoded(&entry()), "0.5")), None, Some(1.0 - 0.5))]
#[case::no_results(context(), Some(redis::Value::Array(vec![redis::Value::Int(0)])), None, Some(0.0))]
#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
#[tokio::test]
async fn lookup_reports_python_semantic_similarity(
#[case] context: SemanticCacheContext,
#[case] reply: Option<redis::Value>,
#[case] value: Option<Value>,
#[case] similarity: Option<f64>,
#[values(false, true)] use_async: bool,
) {
let searched = reply.is_some();
let recording = recording(reply.map_or_else(Vec::new, |reply| vec![ok(), Ok(reply)]));
let lookup = if use_async {
recording
.cache
.async_get_cache_with_similarity("key", &context)
.await
} else {
recording.cache.get_cache_with_similarity("key", &context)
};
assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
assert_eq!(recording.text().contains("FT.SEARCH"), searched);
}
#[rstest]
#[tokio::test]
async fn missing_prompt_does_not_touch_valkey(entry: Value) {
let cache = ValkeySemanticCache::with_connection(
MockRedisConnection::new([]).assert_all_commands_consumed(),
FakeEmbedder::new(&[]),
JsonCodec::new(),
config(),
);
let context = SemanticCacheContext::default();
cache.set_cache("key", entry.clone(), &context).unwrap();
assert_eq!(cache.get_cache("key", &context).unwrap(), None);
cache
.async_set_cache("key", entry, context.clone())
.await
.unwrap();
assert_eq!(cache.async_get_cache("key", &context).await.unwrap(), None);
assert_eq!(cache.get_ttl(&context), None);
}
#[rstest]
fn with_embedder_shares_index_state_and_connections(entry: Value, context: SemanticCacheContext) {
let recording = recording([ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]);
recording
.cache
.set_cache("key", entry.clone(), &context)
.unwrap();
let prepared = recording
.cache
.with_embedder(PreparedEmbedding(vec![0.1, 0.2, 0.3]));
assert_eq!(prepared.get_cache("key", &context).unwrap(), Some(entry));
assert_eq!(recording.text().matches("FT.CREATE").count(), 1);
}
#[rstest]
fn accessors_report_the_config() {
let recording = recording([]);
assert_eq!(recording.cache.index_name(), "test");
assert_eq!(recording.cache.similarity_threshold(), 0.8);
assert_eq!(DEFAULT_INDEX_NAME, "litellm_semantic_cache_index");
}
#[rstest]
#[tokio::test]
async fn async_set_and_get_use_shared_document_helpers(
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([ok(), ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]);
let context = SemanticCacheContext {
ttl: Some(Duration::from_millis(1900)),
..context
};
recording
.cache
.async_set_cache("key", entry.clone(), context.clone())
.await
.unwrap();
assert_eq!(
recording
.cache
.async_get_cache("key", &context)
.await
.unwrap(),
Some(entry)
);
let text = recording.text();
assert!(text.contains("FT.CREATE"));
assert!(text.contains("HSET"));
assert!(text.contains("EXPIRE"));
assert_eq!(
*recording.calls.lock().unwrap(),
vec![("hello".to_owned(), context.metadata.clone()); 2]
);
}

62
litellm-rust/crates/cache-valkey-semantic/tests/contract.rs
View file

@ -1,62 +0,0 @@
//! `overwrite_replaces` does not apply: like Python, every write is a new `<scope>:<uuid4>`
//! document, so a second write with the same prompt adds a tie instead of replacing the first.
mod support;
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
use litellm_cache_testing as contract;
use litellm_cache_valkey_semantic::{
DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig,
};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::FakeSearch;
type Cache = ValkeySemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
const PREFIX: &str = "contract:";
#[fixture]
fn cache() -> Cache {
ValkeySemanticCache::with_connection(
FakeSearch::default(),
PreparedEmbedding(vec![0.6, 0.8]),
JsonCodec::new(),
ValkeySemanticConfig {
similarity_threshold: 0.9,
index_name: DEFAULT_INDEX_NAME.into(),
},
)
}
#[fixture]
fn context() -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
..Default::default()
}
}
#[rstest]
#[tokio::test]
async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) {
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test]
async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) {
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test]
async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) {
contract::pipeline_writes_every_entry(
&cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}

352
litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs
View file

@ -1,352 +0,0 @@
#![allow(dead_code)]
use std::{
collections::{BTreeMap, HashMap, VecDeque},
sync::{Arc, Mutex},
};
use litellm_cache::{Error, semantic::Embedder};
use serde_json::Value;
pub type EmbedCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every
/// prompt with its metadata.
pub struct FakeEmbedder {
vectors: HashMap<String, Vec<f32>>,
pub calls: EmbedCalls,
}
impl FakeEmbedder {
pub fn new(vectors: &[(&str, &[f32])]) -> Self {
Self {
vectors: vectors
.iter()
.map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec()))
.collect(),
calls: EmbedCalls::default(),
}
}
}
impl Embedder for FakeEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.calls
.lock()
.unwrap()
.push((prompt.to_owned(), metadata.cloned()));
Ok(self
.vectors
.get(prompt)
.cloned()
.unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
}
async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.embed(prompt, metadata)
}
}
struct FakeIndex {
prefix: Vec<u8>,
dims: usize,
vector_field: String,
}
#[derive(Default)]
struct SearchState {
indexes: HashMap<String, FakeIndex>,
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
}
/// An in-memory valkey-search speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache
/// sends, with exact cosine KNN over the hashes under an index prefix.
#[derive(Clone, Default)]
pub struct FakeSearch {
state: Arc<Mutex<SearchState>>,
}
impl FakeSearch {
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
let mut state = self.state.lock().unwrap();
let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned();
match text(0).to_uppercase().as_str() {
"FT.CREATE" => {
let name = text(1);
if state.indexes.contains_key(&name) {
return Err(server_error("Index already exists"));
}
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
let prefix = args[position("PREFIX").unwrap() + 2].clone();
let dims = text(position("DIM").unwrap() + 1).parse().unwrap();
let vector_field = text(position("VECTOR").unwrap() - 1);
state.indexes.insert(
name,
FakeIndex {
prefix,
dims,
vector_field,
},
);
Ok(redis::Value::Okay)
}
"FT.INFO" => {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| server_error("Unknown index name"))?;
Ok(index_info(index))
}
"FT.DROPINDEX" => {
state.indexes.remove(&text(1));
Ok(redis::Value::Okay)
}
"HSET" => {
let hash = state.hashes.entry(args[1].clone()).or_default();
for pair in args[2..].chunks(2) {
hash.insert(
String::from_utf8_lossy(&pair[0]).into_owned(),
pair[1].clone(),
);
}
Ok(redis::Value::Int(((args.len() - 2) / 2) as i64))
}
"EXPIRE" => Ok(redis::Value::Int(i64::from(
state.hashes.contains_key(&args[1]),
))),
"FT.SEARCH" => {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| server_error("no such index"))?;
let query = text(2);
let tag = query_tag(&query);
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
let vector = floats(&args[params + 3]);
let best = state
.hashes
.iter()
.filter(|(key, _)| key.starts_with(&index.prefix))
.filter(|(_, fields)| {
fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes())
})
.filter_map(|(key, fields)| {
let stored = floats(fields.get(&index.vector_field)?);
(stored.len() == index.dims)
.then(|| (key, fields, 1.0 - cosine(&vector, &stored)))
})
.min_by(|left, right| left.2.total_cmp(&right.2));
let Some((key, fields, distance)) = best else {
return Ok(redis::Value::Array(vec![redis::Value::Int(0)]));
};
let mut reply = fields
.iter()
.filter(|(name, _)| **name != index.vector_field)
.flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)])
.collect::<Vec<_>>();
reply.extend([
bulk(b"vector_distance"),
bulk(distance.to_string().as_bytes()),
]);
Ok(redis::Value::Array(vec![
redis::Value::Int(1),
bulk(key),
redis::Value::Array(reply),
]))
}
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
_ => Err(server_error("unsupported command")),
}
}
}
impl redis::ConnectionLike for FakeSearch {
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
let mut commands = parse_commands(command);
self.run(commands.remove(0))
}
fn req_packed_commands(
&mut self,
commands: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
let replies = parse_commands(commands)
.into_iter()
.map(|args| self.run(args))
.collect::<redis::RedisResult<Vec<_>>>()?;
Ok(replies.into_iter().skip(offset).take(count).collect())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}
pub fn server_error(reply: &str) -> redis::RedisError {
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
.expect("error reply")
.extract_error()
.expect_err("server error")
}
fn bulk(bytes: &[u8]) -> redis::Value {
redis::Value::BulkString(bytes.to_vec())
}
fn index_info(index: &FakeIndex) -> redis::Value {
redis::Value::Array(vec![
bulk(b"index_name"),
bulk(b"fake"),
bulk(b"attributes"),
redis::Value::Array(vec![
redis::Value::Array(vec![
bulk(b"identifier"),
bulk(b"litellm_cache_key"),
bulk(b"type"),
bulk(b"TAG"),
]),
redis::Value::Array(vec![
bulk(b"identifier"),
bulk(index.vector_field.as_bytes()),
bulk(b"type"),
bulk(b"VECTOR"),
bulk(b"index"),
redis::Value::Array(vec![
bulk(b"dimensions"),
redis::Value::Int(index.dims as i64),
]),
]),
]),
])
}
/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed.
fn query_tag(query: &str) -> String {
let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len();
let mut tag = String::new();
let mut characters = query[start..].chars();
while let Some(character) = characters.next() {
match character {
'\\' => tag.extend(characters.next()),
'}' => break,
character => tag.push(character),
}
}
tag
}
fn floats(bytes: &[u8]) -> Vec<f32> {
bytes
.as_chunks::<4>()
.0
.iter()
.map(|chunk| f32::from_le_bytes(*chunk))
.collect()
}
fn cosine(left: &[f32], right: &[f32]) -> f64 {
let dot = left
.iter()
.zip(right)
.map(|(left, right)| f64::from(*left) * f64::from(*right))
.sum::<f64>();
let norm = |vector: &[f32]| {
vector
.iter()
.map(|value| f64::from(*value).powi(2))
.sum::<f64>()
.sqrt()
};
dot / (norm(left) * norm(right))
}
/// Splits a packed RESP request into each command's arguments.
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
let line = |bytes: &mut &[u8]| {
let end = bytes
.windows(2)
.position(|window| window == b"\r\n")
.unwrap();
let text = String::from_utf8(bytes[1..end].to_vec()).unwrap();
*bytes = &bytes[end + 2..];
text.parse::<usize>().unwrap()
};
let mut commands = Vec::new();
while !bytes.is_empty() {
let count = line(&mut bytes);
let mut args = Vec::with_capacity(count);
for _ in 0..count {
let length = line(&mut bytes);
args.push(bytes[..length].to_vec());
bytes = &bytes[length + 2..];
}
commands.push(args);
}
commands
}
/// Records every packed request and answers from a script, `OK` once the script runs out.
pub struct RecordingConnection {
requests: Arc<Mutex<Vec<Vec<u8>>>>,
replies: Mutex<VecDeque<redis::RedisResult<redis::Value>>>,
}
impl RecordingConnection {
pub fn new(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Self {
Self {
requests: Arc::default(),
replies: Mutex::new(replies.into_iter().collect()),
}
}
pub fn requests(&self) -> Arc<Mutex<Vec<Vec<u8>>>> {
Arc::clone(&self.requests)
}
fn reply(&self) -> redis::RedisResult<redis::Value> {
self.replies
.lock()
.unwrap()
.pop_front()
.unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into())))
}
}
impl redis::ConnectionLike for RecordingConnection {
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
self.requests.lock().unwrap().push(command.to_vec());
self.reply()
}
fn req_packed_commands(
&mut self,
command: &[u8],
_offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
self.requests.lock().unwrap().push(command.to_vec());
(0..count).map(|_| self.reply()).collect()
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}

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

@ -29,8 +29,6 @@ bytes.workspace = true
futures-util.workspace = true
litellm-cache.workspace = true
litellm-error.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
@ -56,7 +54,6 @@ pyo3.workspace = true
prost.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true
redis = { workspace = true, features = ["tls-rustls"] }
serde_json.workspace = true
strum.workspace = true
veil.workspace = true

6
litellm-rust/crates/python-bridge/src/cache/AGENTS.md vendored
View file

@ -1,7 +1,7 @@
# Cache boundary
This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
This folder owns how Rust inference reaches the selected Python cache. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters
`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission and inference protocol composition
`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports
`python/` delegates operations to the selected Python cache without discovering configuration. Routes use only the parent module's exports

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

@ -1,32 +1,15 @@
mod native;
mod python;
mod selection;
pub(crate) use native::NativeCacheHandle;
pub(crate) use python::{CacheCall, PythonCache};
pub(crate) use selection::{Cached, Selection, configure, configured_native};
pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
use litellm_cache::Error;
use pyo3::{
exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
prelude::*,
};
use pyo3::{exceptions::PyRuntimeError, prelude::*};
/// A cache call the bridge rejects itself, before or after the cache ran.
#[derive(Debug, thiserror::Error)]
pub(crate) enum HandleError {
#[error("recreate the v2 cache after fork")]
Forked,
#[error("cache limits must be positive")]
Limits,
#[error("namespace and a positive cache limit are required")]
RedisLimits,
#[error("cache durations must be finite and positive")]
Duration(#[source] Option<std::time::TryFromFloatSecsError>),
#[error("unsupported v2 cache control: {0}")]
Control(String),
#[error("native inference needs a v2 cache, the catalog keeps Python caches on Python")]
NotNative,
#[error("this inference route does not support configured caches")]
Unsupported,
#[error("cache operation without configured cache")]
Unconfigured,
#[error("cache arguments unavailable")]
@ -39,66 +22,10 @@ impl From<HandleError> for PyErr {
fn from(error: HandleError) -> Self {
let message = litellm_error::Report::new(&error).to_string();
match error {
HandleError::Limits
| HandleError::RedisLimits
| HandleError::Duration(_)
| HandleError::Control(_) => PyValueError::new_err(message),
HandleError::Forked
| HandleError::NotNative
HandleError::Unsupported
| HandleError::Unconfigured
| HandleError::MissingArguments
| HandleError::NothingPending => PyRuntimeError::new_err(message),
}
}
}
fn cache_error(error: Error) -> PyErr {
let message = litellm_error::Report::new(&error).to_string();
match error {
Error::InvalidEntry(_) => PyValueError::new_err(message),
Error::UnsupportedOperation => PyNotImplementedError::new_err(message),
Error::Unavailable(_) | Error::UnscopedFlush | Error::MissingPrompt => {
PyRuntimeError::new_err(message)
}
}
}
#[cfg(test)]
mod tests {
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
prelude::*,
};
use rstest::rstest;
use super::{Error, cache_error};
fn refused() -> std::io::Error {
std::io::Error::other("connection refused")
}
#[rstest]
#[case::unavailable(
Error::unavailable(refused()),
false,
"cache is unavailable: connection refused"
)]
#[case::invalid_entry(
Error::invalid_entry(refused()),
true,
"invalid cache entry: connection refused"
)]
fn python_sees_the_class_and_the_whole_chain(
#[case] error: Error,
#[case] caller_fault: bool,
#[case] message: &str,
) {
Python::initialize();
Python::attach(|py| {
let failure = cache_error(error);
assert_eq!(failure.is_instance_of::<PyValueError>(py), caller_fault);
assert_eq!(failure.is_instance_of::<PyRuntimeError>(py), !caller_fault);
assert_eq!(failure.value(py).to_string(), message);
});
}
}

9
litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md vendored
View file

@ -1,9 +0,0 @@
# Native cache bindings
This directory exposes the experimental V2 native cache handles (`NativeCacheHandle`) to Python. Cache algorithms and storage protocols remain in their cache crates
Accept the cache object selected by the parent module. Do not read global `litellm.cache`, decide route admission, or select the Python cache adapter here
Keep Python-facing cache classes and method signatures stable when reorganizing modules. Native internals stay private to this directory unless the shared cache boundary or Python module registration needs them
Verify changes with the V2 cache tests using a freshly built extension. Test cache behavior, not module paths or file structure

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

@ -1,3 +0,0 @@
pub(super) mod v2;
pub(crate) use v2::NativeCacheHandle;

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

@ -1,336 +0,0 @@
use crate::cache::{HandleError, cache_error};
use std::{
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{DeleteCache, DisconnectCache, PingCache};
use litellm_host_python::{from_py, release_gil, to_py};
use serde_json::Value;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec,
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
};
use pyo3::{prelude::*, types::PyDict};
#[pyclass(
frozen,
name = "NativeCacheHandle",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeCacheHandle {
service: Arc<dyn ResponseCacheService>,
backend: Arc<dyn ExactResponseCache>,
storage: Storage,
pid: u32,
}
#[derive(Clone)]
enum Storage {
Memory(Arc<InMemoryCache<CacheEntry>>),
Redis(Arc<RedisCache<ResponseCacheCodec>>),
}
impl NativeCacheHandle {
fn check_process(&self) -> PyResult<()> {
if self.pid != std::process::id() {
return Err(HandleError::Forked.into());
}
Ok(())
}
}
fn request(key: String, ttl: Option<f64>) -> PyResult<ResponseCacheRequest> {
let mut request: ResponseCacheRequest = ResponseCacheRequest::new(CacheKeyInput {
preset: Some(key),
..Default::default()
});
request.context.ttl = ttl.map(duration).transpose()?;
Ok(request)
}
#[pymethods]
impl NativeCacheHandle {
#[staticmethod]
#[pyo3(signature = (*, ttl=600.0, capacity=200, max_entry_bytes=4194304))]
fn memory(ttl: f64, capacity: usize, max_entry_bytes: usize) -> PyResult<Self> {
let ttl = duration(ttl)?;
if capacity == 0 || max_entry_bytes == 0 {
return Err(HandleError::Limits.into());
}
let storage = Arc::new(InMemoryCache::new(Some(capacity), Some(ttl)));
let backend = Arc::new(ResponseCache::new(storage.clone()).with_config(
ResponseCacheConfig {
namespace: "sdk".into(),
max_entry_bytes,
},
));
Ok(Self {
service: backend.clone(),
backend,
storage: Storage::Memory(storage),
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (url, *, namespace, ttl=600.0, max_entry_bytes=4194304))]
fn redis(
py: Python<'_>,
url: &str,
namespace: String,
ttl: f64,
max_entry_bytes: usize,
) -> PyResult<Self> {
let ttl = duration(ttl)?;
if namespace.is_empty() || max_entry_bytes == 0 {
return Err(HandleError::RedisLimits.into());
}
let storage = Arc::new(
release_gil(py, || {
RedisCache::connect(
url,
&RedisTopology::Standalone,
Some(ttl),
ResponseCacheCodec,
)
})
.map_err(cache_error)?
.with_namespace(Some(namespace.clone())),
);
let backend = Arc::new(ResponseCache::new(storage.clone()).with_config(
ResponseCacheConfig {
namespace,
max_entry_bytes,
},
));
Ok(Self {
service: backend.clone(),
backend,
storage: Storage::Redis(storage),
pid: std::process::id(),
})
}
fn get(&self, py: Python<'_>, key: String) -> PyResult<Py<PyAny>> {
self.check_process()?;
let request = request(key, None)?;
let value =
release_gil(py, || self.backend.lookup(&request, now())).map_err(cache_error)?;
to_py(py, &value)
}
#[pyo3(signature = (key, value, *, ttl=None))]
fn set(
&self,
py: Python<'_>,
key: String,
value: &Bound<'_, PyAny>,
ttl: Option<f64>,
) -> PyResult<()> {
self.check_process()?;
let request = request(key, ttl)?;
let value: Value = from_py(value)?;
release_gil(py, || self.backend.store(&request, value, now())).map_err(cache_error)
}
fn async_get<'py>(&self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let request = request(key, None)?;
let backend = self.backend.clone();
crate::execution::run_async(
py,
async move { backend.async_lookup(&request, now()).await },
cache_error,
)
}
#[pyo3(signature = (key, value, *, ttl=None))]
fn async_set<'py>(
&self,
py: Python<'py>,
key: String,
value: &Bound<'_, PyAny>,
ttl: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let request = request(key, ttl)?;
let value: Value = from_py(value)?;
let backend = self.backend.clone();
crate::execution::run_async(
py,
async move { backend.async_store(&request, value, now()).await },
cache_error,
)
}
#[pyo3(signature = (entries, *, ttl=None))]
fn async_set_many<'py>(
&self,
py: Python<'py>,
entries: &Bound<'_, PyAny>,
ttl: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let entries: Vec<(String, Value)> = from_py(entries)?;
let entries = entries
.into_iter()
.map(|(key, value)| Ok((request(key, ttl)?, value)))
.collect::<PyResult<Vec<_>>>()?;
let backend = self.backend.clone();
crate::execution::run_async(
py,
async move { backend.async_store_batch(entries, now()).await },
cache_error,
)
}
fn flush(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
self.check_process()?;
let backend = self.backend.clone();
crate::execution::run_sync(py, async move { backend.async_flush().await }, cache_error)
}
fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let backend = self.backend.clone();
crate::execution::run_async(py, async move { backend.async_flush().await }, cache_error)
}
fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let storage = self.storage.clone();
crate::execution::run_async(
py,
async move {
match storage {
Storage::Memory(_) => Ok(true),
Storage::Redis(cache) => cache.ping().await,
}
},
cache_error,
)
}
fn disconnect<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let storage = self.storage.clone();
crate::execution::run_async(
py,
async move {
match storage {
Storage::Memory(cache) => cache.disconnect().await,
Storage::Redis(cache) => cache.disconnect().await,
}
},
cache_error,
)
}
fn delete<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let storage = self.storage.clone();
crate::execution::run_async(
py,
async move {
for key in keys {
match &storage {
Storage::Memory(cache) => cache.async_delete_cache(&key).await?,
Storage::Redis(cache) => cache.async_delete_cache(&key).await?,
}
}
Ok(())
},
cache_error,
)
}
}
fn duration(seconds: f64) -> Result<Duration, HandleError> {
match Duration::try_from_secs_f64(seconds) {
Ok(duration) if !duration.is_zero() => Ok(duration),
Ok(_) => Err(HandleError::Duration(None)),
Err(error) => Err(HandleError::Duration(Some(error))),
}
}
pub(in crate::cache) fn native_handle<'py>(
configured: &Bound<'py, PyAny>,
) -> PyResult<Option<Bound<'py, PyAny>>> {
Ok(configured
.getattr_opt("cache")?
.map(|backend| backend.getattr_opt("native_handle"))
.transpose()?
.flatten()
.filter(|handle| handle.is_instance_of::<NativeCacheHandle>()))
}
pub(in crate::cache) fn configured(
configured: &Bound<'_, PyAny>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<(
Option<Arc<dyn ResponseCacheService>>,
litellm_cache_response::CacheOptions,
)> {
let handle = native_handle(configured)?.ok_or(HandleError::NotNative)?;
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;
cache.check_process()?;
let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none());
let controls = controls
.as_ref()
.map(|value| value.cast::<PyDict>())
.transpose()?;
if let Some(controls) = controls {
for name in controls.keys() {
let name = name.extract::<String>()?;
if !matches!(
name.as_str(),
"no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache"
) {
return Err(HandleError::Control(name).into());
}
}
}
let boolean = |name: &str| -> PyResult<bool> {
controls
.map(|values| values.get_item(name))
.transpose()?
.flatten()
.map(|value| value.extract())
.transpose()
.map(|value| value.unwrap_or(false))
};
let seconds = |name: &str| -> PyResult<Option<Duration>> {
controls
.map(|values| values.get_item(name))
.transpose()?
.flatten()
.map(|value| Ok(duration(value.extract()?)?))
.transpose()
};
Ok((
Some(cache.service.clone()),
litellm_cache_response::CacheOptions {
policy: litellm_cache_response::CachePolicy {
caching: kwargs
.get_item("caching")?
.filter(|value| !value.is_none())
.map(|value| value.extract())
.transpose()?,
no_cache: boolean("no-cache")?,
no_store: boolean("no-store")?,
ttl: seconds("ttl")?,
max_age: seconds("s-max-age")?.or(seconds("s-maxage")?),
},
scope: litellm_cache_response::CacheScope::Shared,
},
))
}
fn now() -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
}

45
litellm-rust/crates/python-bridge/src/cache/selection.rs vendored
View file

@ -1,14 +1,10 @@
use super::{native, python};
use litellm_cache_response::{
CacheOptions, CachePolicy, CacheScope, ResponseCacheService, ScopedCache,
};
use super::{HandleError, python};
use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache};
use litellm_host::{
machine::{HostServices, MachineFault},
protocol::Protocol,
};
use pyo3::{prelude::*, types::PyDict};
use std::sync::Arc;
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
pub(crate) struct Cached<P>(std::marker::PhantomData<P>);
@ -24,7 +20,6 @@ impl<P: Protocol> Protocol for Cached<P> {
enum Backend {
Disabled,
Native(Arc<dyn ResponseCacheService>),
Python { namespace: String },
}
@ -43,7 +38,6 @@ impl Selection {
{
let service = match self.backend {
Backend::Disabled => None,
Backend::Native(service) => Some(service),
Backend::Python { namespace } => Some(python::service(services, namespace)),
};
(
@ -76,37 +70,21 @@ fn selected_cache<'py>(
Ok(Some(configured))
}
pub(crate) fn configured_native(
pub(crate) fn ensure_cache_disabled(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
call_type: &str,
) -> PyResult<(
Option<Arc<dyn ResponseCacheService>>,
litellm_cache_response::CacheOptions,
)> {
) -> PyResult<()> {
let Some(configured) = selected_cache(py, kwargs, call_type)? else {
return Ok((
None,
litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
));
return Ok(());
};
native_configuration(&configured, kwargs)
}
fn native_configuration(
configured: &Bound<'_, PyAny>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<(Option<Arc<dyn ResponseCacheService>>, CacheOptions)> {
if !configured
if configured
.call_method("should_use_cache", (), Some(kwargs))?
.extract::<bool>()?
{
return Ok((
None,
litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
));
return Err(HandleError::Unsupported.into());
}
native::v2::configured(configured, kwargs)
Ok(())
}
pub(crate) fn configure(
@ -122,13 +100,6 @@ pub(crate) fn configure(
options: CacheOptions::new(CacheScope::Shared),
});
};
if native::v2::native_handle(&cache)?.is_some() {
let (native, options) = native_configuration(&cache, arguments)?;
return Ok(Selection {
backend: native.map_or(Backend::Disabled, Backend::Native),
options,
});
}
let enabled = cache
.call_method("should_use_cache", (), Some(arguments))?
.extract::<bool>()?;

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

@ -65,10 +65,6 @@ mod _native {
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
let py = module.py();
let dict = module.dict();
dict.set_item(
"NativeCacheHandle",
py.get_type::<crate::cache::NativeCacheHandle>(),
)?;
dict.set_item(
"_SecretManagerRuntime",
py.get_type::<crate::secrets::runtime::NativeSecretManager>(),
@ -92,7 +88,6 @@ mod tests {
Python::initialize();
Python::attach(|py| {
let mut expected = vec![
"NativeCacheHandle",
"RustBridgeDeclined",
"RustUpstreamError",
"ForkedAfterNativeRuntimeStarted",

6
litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs
View file

@ -4,7 +4,7 @@ use std::sync::Arc;
use host::ChatCompletionsPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_cache_response::CachePolicy;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions};
@ -57,10 +57,6 @@ impl InferenceRoute for ChatCompletionsRoute {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
self.with_cache(cache)
}
fn machine(
self,
call: <ChatCompletions as Protocol>::Request,

11
litellm-rust/crates/python-bridge/src/routes/inference.rs
View file

@ -1,7 +1,7 @@
use std::sync::Arc;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
use litellm_cache_response::CachePolicy;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host::{call::HostedMachine, protocol::Protocol};
@ -142,7 +142,6 @@ pub(super) trait InferenceRoute: Sized + 'static {
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
) -> Self;
fn with_cache(self, cache: ScopedCache) -> Self;
fn machine(
self,
call: <Self::Protocol as Protocol>::Request,
@ -176,12 +175,8 @@ where
crate::http::resources().auth.clone(),
crate::secrets::source(py)?,
);
let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?;
let route = match cache {
Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)),
None => route,
};
Ok(route.machine(request, cache_options.policy))
crate::cache::ensure_cache_disabled(py, arguments, call_type)?;
Ok(route.machine(request, CachePolicy::default()))
},
host,
hooks,

6
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
View file

@ -5,7 +5,7 @@ use std::sync::Arc;
use host::ResponsesPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_cache_response::CachePolicy;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_inference_responses::{ResponsesRoute, route::Responses};
@ -49,10 +49,6 @@ impl InferenceRoute for ResponsesRoute {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
self.with_cache(cache)
}
fn machine(
self,
call: <Responses as Protocol>::Request,

3
litellm/_v2/AGENTS.md
View file

@ -1,3 +0,0 @@
Everything here is experimental and should not be documented
Use this directory to explore alternative APIs where the Rust migration makes backward compatibility difficult. The gateway can use them for performance, but keep them behind the v2 flag to avoid breaking SDK users

3
litellm/_v2/__init__.py
View file

@ -1,3 +0,0 @@
from litellm._v2.cache import Cache
__all__ = ("Cache",)

11
litellm/_v2/cache/AGENTS.md vendored
View file

@ -1,11 +0,0 @@
# Python v2 cache
Keep `litellm._v2.cache.Cache` import-compatible when reorganizing this package. This package owns Python cache factories and the adapter between the existing `BaseCache` interface and `NativeCacheHandle`
Construct native handles at this boundary and inject the adapter through the existing cache facade's `_backend` parameter. Keep the facade's `type`, namespace, and TTL consistent with the configured native backend
Keep storage implementation in the Rust storage crates and response-cache policy in `litellm-cache-response` and core. Do not duplicate cache-key generation, freshness rules, response encoding, or inference orchestration here
Preserve synchronous and asynchronous cache operations, including TTL forwarding and lifecycle methods. Validate Python values before passing them to typed native interfaces. Keep native extension imports lazy so importing the package does not require loading the extension
Extend the existing v2 cache tests in `tests/test_litellm_rust/test_v2.py` for behavioral changes, following that directory's `AGENTS.md`. Test observable cache behavior rather than package layout or implementation structure

72
litellm/_v2/cache/__init__.py vendored
View file

@ -1,72 +0,0 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
from litellm.caching.base_cache import BaseCache
from litellm.caching.caching import Cache as CacheFacade
from litellm.types.caching import LiteLLMCacheType
if TYPE_CHECKING:
from litellm.rust_bridge._native import NativeCacheHandle
_DURATION: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
class NativeBackend(BaseCache):
def __init__(self, handle: NativeCacheHandle) -> None:
self.native_handle = handle
def get_cache(self, key: str, **kwargs: object) -> object:
return self.native_handle.get(key)
async def async_get_cache(self, key: str, **kwargs: object) -> object:
return await self.native_handle.async_get(key)
def set_cache(self, key: str, value: object, **kwargs: object) -> None:
self.native_handle.set(key, value, ttl=_DURATION.validate_python(kwargs.get("ttl")))
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
await self.native_handle.async_set(key, value, ttl=_DURATION.validate_python(kwargs.get("ttl")))
async def async_set_cache_pipeline(self, cache_list: Sequence[tuple[str, object]], **kwargs: object) -> None:
await self.native_handle.async_set_many(cache_list, ttl=_DURATION.validate_python(kwargs.get("ttl")))
async def batch_cache_write(self, key: str, value: object, **kwargs: object) -> None:
await self.async_set_cache(key, value, **kwargs)
def flush_cache(self) -> None:
self.native_handle.flush()
async def async_flush_cache(self) -> None:
await self.native_handle.async_flush()
async def ping(self) -> bool:
return await self.native_handle.ping()
async def disconnect(self) -> None:
await self.native_handle.disconnect()
async def delete_cache_keys(self, keys: Sequence[str]) -> None:
await self.native_handle.delete(keys)
async def test_connection(self) -> dict[str, str]:
return {"status": "success" if await self.ping() else "failed"}
class Cache:
@staticmethod
def memory(*, ttl: float = 600, capacity: int = 200, max_entry_bytes: int = 4194304) -> CacheFacade:
from litellm.rust_bridge._native import NativeCacheHandle
handle: Final = NativeCacheHandle.memory(ttl=ttl, capacity=capacity, max_entry_bytes=max_entry_bytes)
return CacheFacade(type=LiteLLMCacheType.LOCAL, ttl=ttl, _backend=NativeBackend(handle))
@staticmethod
def redis(url: str, *, namespace: str, ttl: float = 600, max_entry_bytes: int = 4194304) -> CacheFacade:
from litellm.rust_bridge._native import NativeCacheHandle
handle: Final = NativeCacheHandle.redis(url, namespace=namespace, ttl=ttl, max_entry_bytes=max_entry_bytes)
return CacheFacade(type=LiteLLMCacheType.REDIS, namespace=namespace, ttl=ttl, _backend=NativeBackend(handle))

18
litellm/mcp_client/client.py
View file

@ -478,6 +478,13 @@ def _notify_cleanup(callback: Callable[[], None] | None) -> None:
callback()
class _UnvalidatedClientSession(ClientSession):
"""Skips the SDK's output check, which lists the upstream's tools on a session with no schema cached."""
async def validate_tool_result(self, name: str, result: MCPCallToolResult) -> None:
return None
class MCPClient:
"""
MCP Client supporting:
@ -505,12 +512,15 @@ class MCPClient:
elicitation_callback: "ElicitationCallback | None" = None,
logging_callback: Callable | None = None,
protocol_version: MCPUpstreamProtocol = "auto",
skip_tool_output_validation: bool = False,
upstream_transport: httpx2.AsyncBaseTransport | None = None,
):
self.protocol_version: MCPUpstreamProtocol = TypeAdapter[MCPUpstreamProtocol](
MCPUpstreamProtocol
).validate_python(protocol_version)
validate_mcp_protocol_transport(self.protocol_version, transport_type)
# 2026-07-28 lists tools before every call, so its output check adds no list
self.skip_tool_output_validation: Final = skip_tool_output_validation and self.protocol_version != "2026-07-28"
self.server_url: str = server_url
self.transport_type: MCPTransport = transport_type
self.auth_type: MCPAuthType = auth_type
@ -546,7 +556,7 @@ class MCPClient:
"response_hooks": _mcp_response_hooks(),
}
def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool]:
def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool, bool]:
callbacks: Final = self.session_callbacks()
return (
self.transport_type,
@ -555,6 +565,7 @@ class MCPClient:
callbacks["sampling_callback"] is not None,
callbacks["elicitation_callback"] is not None,
callbacks["logging_callback"] is not None,
self.skip_tool_output_validation,
)
def route_session_callbacks(self, callbacks: _SessionCallbackMap) -> None:
@ -776,7 +787,10 @@ class MCPClient:
)
# The SDK drops a response stream that ends without a JSON-RPC reply, so nothing else
# ever fails the request.
session_ctx: Final = ClientSession(
session_type: Final = (
_UnvalidatedClientSession if self.skip_tool_output_validation else ClientSession
)
session_ctx: Final = session_type(
read_stream,
write_stream,
read_timeout_seconds=self.timeout,

4
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -254,7 +254,9 @@ async def anthropic_response(
model_id=model_id,
version=version,
response_cost=(
failure_row_response_cost(base_llm_response_processor.data) if is_agentic_followup_failure(e) else 0
failure_row_response_cost(base_llm_response_processor.data)
if is_agentic_followup_failure(answered)
else 0
),
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
request_data=data,

2
litellm/proxy/mcp_server/mcp_server_manager.py
View file

@ -367,7 +367,7 @@ _UpstreamSessionKey: TypeAlias = tuple[
str,
str,
str,
tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool],
tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool, bool],
]

3
litellm/proxy/mcp_server/upstream.py
View file

@ -231,6 +231,7 @@ async def prepare_upstream_client(
server_url="",
transport_type=transport,
protocol_version=protocol_version,
skip_tool_output_validation=server.skip_tool_output_validation,
auth_type=server.auth_type,
auth_value=auth_value,
timeout=timeout,
@ -265,6 +266,7 @@ async def prepare_upstream_client(
server_url=server.url or "",
transport_type=transport,
protocol_version=protocol_version,
skip_tool_output_validation=server.skip_tool_output_validation,
upstream_transport=upstream_transport,
auth_type=server.auth_type,
timeout=timeout,
@ -308,6 +310,7 @@ async def prepare_upstream_client(
server_url=server.url or "",
transport_type=transport,
protocol_version=protocol_version,
skip_tool_output_validation=server.skip_tool_output_validation,
upstream_transport=upstream_transport,
auth_type=server.auth_type,
auth_value=auth_value,

29
litellm/rust_bridge/_native.pyi
View file

@ -296,32 +296,3 @@ class _SecretManagerRuntime:
timeout: float | httpx.Timeout | None = None,
primary_secret_name: str | None = None,
) -> Future[JsonValue]: ...
@final
class NativeCacheHandle:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
@staticmethod
def memory(
*,
ttl: float = 600.0,
capacity: int = 200,
max_entry_bytes: int = 4194304,
) -> NativeCacheHandle: ...
@staticmethod
def redis(
url: str,
*,
namespace: str,
ttl: float = 600.0,
max_entry_bytes: int = 4194304,
) -> NativeCacheHandle: ...
def get(self, key: str) -> object: ...
def set(self, key: str, value: object, *, ttl: float | None = None) -> None: ...
def async_get(self, key: str) -> Future[object]: ...
def async_set(self, key: str, value: object, *, ttl: float | None = None) -> Future[None]: ...
def async_set_many(self, entries: Sequence[tuple[str, object]], *, ttl: float | None = None) -> Future[None]: ...
def flush(self) -> None: ...
def async_flush(self) -> Future[None]: ...
def ping(self) -> Future[bool]: ...
def disconnect(self) -> Future[None]: ...
def delete(self, keys: Sequence[str]) -> Future[None]: ...

6
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -23,6 +23,8 @@ from litellm.types.mcp import (
def _validate_mcp_protocol_metadata(value: dict[str, object]) -> dict[str, object]:
if "protocol_version" in value:
TypeAdapter[MCPUpstreamProtocol](MCPUpstreamProtocol).validate_python(value["protocol_version"])
if "skip_tool_output_validation" in value:
TypeAdapter[bool](bool).validate_python(value["skip_tool_output_validation"], strict=True)
return value
@ -297,6 +299,10 @@ class MCPServer(LiteLLMBaseModel):
"""
return self.oauth2_flow == "client_credentials"
@property
def skip_tool_output_validation(self) -> bool:
return (self.mcp_info or {}).get("skip_tool_output_validation") is True
@model_validator(mode="after")
def resolve_protocol_version(self) -> Self:
if "protocol_version" not in self.model_fields_set and self.mcp_info is not None:

42
tests/_support/mcp_count_tool_upstream.py Normal file
View file

@ -0,0 +1,42 @@
import json
import queue
from collections.abc import Callable, Mapping
from typing import Final
import httpx2
COUNT_OUTPUT_SCHEMA: Final = {"type": "object", "properties": {"n": {"type": "integer"}}, "required": ["n"]}
def count_tool_upstream(
structured: Mapping[str, object], methods_seen: "queue.SimpleQueue[str]"
) -> Callable[[httpx2.Request], httpx2.Response]:
"""An MCP upstream with one tool, `count`, whose declared output schema wants an integer `n`."""
def respond(request: httpx2.Request) -> httpx2.Response:
if request.method != "POST":
return httpx2.Response(200)
payload: Final = json.loads(request.content)
methods_seen.put(payload["method"])
if "id" not in payload:
return httpx2.Response(202)
results: Final[Mapping[str, Mapping[str, object]]] = {
"initialize": {
"protocolVersion": payload["params"].get("protocolVersion"),
"capabilities": {"tools": {}},
"serverInfo": {"name": "count", "version": "1"},
},
"tools/list": {
"tools": [{"name": "count", "inputSchema": {"type": "object"}, "outputSchema": COUNT_OUTPUT_SCHEMA}]
},
"tools/call": {"content": [{"type": "text", "text": "ok"}], "structuredContent": structured},
}
return httpx2.Response(
200, json={"jsonrpc": "2.0", "id": payload["id"], "result": results.get(payload["method"], {})}
)
return respond
def methods_received(methods_seen: "queue.SimpleQueue[str]") -> tuple[str, ...]:
return tuple(methods_seen.get_nowait() for _ in range(methods_seen.qsize()))

29
tests/e2e/ui/tests/modelHub/modelHub.spec.ts
View file

@ -64,7 +64,7 @@ test.describe("AI Hub (internal admin view)", () => {
test.describe("Public model hub (/ui/model_hub_table)", () => {
// No storageState — the public page is reached anonymously with a `key` query param.
test("Public model_hub_table loads and renders the Model Hub tab", async ({ page }) => {
test("Public model_hub_table loads and renders the Models tab", async ({ page }) => {
// The page expects the proxy key as the `key` query param. Use the master
// key the e2e runner already exports — this matches what the AI Hub copy
// button hands out.
@ -74,13 +74,10 @@ test.describe("Public model hub (/ui/model_hub_table)", () => {
// race can't briefly mask the tab while we're evaluating visibility.
await dismissFeedbackPopup(page);
// Page loads (no auth redirect) and the Model Hub tab is always present.
// Agent Hub and MCP Hub tabs are conditionally rendered only when public
// agents/MCP servers exist, so we don't assert on them in a fresh CI run.
await expect(page.getByRole("tab", { name: "Model Hub" })).toBeVisible({ timeout: 10_000 });
await expect(page.getByRole("tab", { name: /^Models/ })).toBeVisible({ timeout: 10_000 });
});
test("Agent Hub and MCP Hub tabs render their public entries", async ({ page, request }) => {
test("Agents and MCP servers tabs render their public entries", async ({ page, request }) => {
const suffix = `${Date.now()}`;
const agentName = `e2e-public-agent-${suffix}`;
const mcpServerName = `e2e_public_mcp_${suffix}`;
@ -145,18 +142,18 @@ test.describe("Public model hub (/ui/model_hub_table)", () => {
await page.goto(`/ui/model_hub_table?key=${masterKey()}`);
await dismissFeedbackPopup(page);
const agentHubTab = page.getByRole("tab", { name: "Agent Hub" });
await expect(agentHubTab).toBeVisible({ timeout: 15_000 });
await agentHubTab.click();
await expect(page.getByText("Available Agents")).toBeVisible({ timeout: 10_000 });
await expect(page.getByRole("row").filter({ hasText: agentName })).toHaveCount(1, { timeout: 10_000 });
const agentsTab = page.getByRole("tab", { name: /^Agents/ });
await expect(agentsTab).toBeVisible({ timeout: 15_000 });
await agentsTab.click();
await expect(page.getByRole("textbox", { name: "Search agents" })).toBeVisible({ timeout: 10_000 });
await expect(page.getByRole("link").filter({ hasText: agentName })).toHaveCount(1, { timeout: 10_000 });
await expect(page.getByText("E2E public agent").first()).toBeVisible();
const mcpHubTab = page.getByRole("tab", { name: "MCP Hub" });
await expect(mcpHubTab).toBeVisible();
await mcpHubTab.click();
await expect(page.getByText("Available MCP Servers")).toBeVisible({ timeout: 10_000 });
await expect(page.getByRole("row").filter({ hasText: mcpServerName })).toHaveCount(1, { timeout: 10_000 });
const mcpTab = page.getByRole("tab", { name: /^MCP/ });
await expect(mcpTab).toBeVisible();
await mcpTab.click();
await expect(page.getByRole("textbox", { name: "Search MCP servers" })).toBeVisible({ timeout: 10_000 });
await expect(page.getByRole("link").filter({ hasText: mcpServerName })).toHaveCount(1, { timeout: 10_000 });
await expect(page.getByText("E2E public MCP server").first()).toBeVisible();
} finally {
await request.post("/v1/mcp/make_public", { headers: auth, data: { mcp_server_ids: priorMcpIds } });

28
tests/integration/providers/test_responses_lifecycle_wire.py
View file

@ -4,7 +4,7 @@ import json
import os
import threading
import uuid
from collections.abc import Callable
from collections.abc import Callable, Iterator
from types import MappingProxyType
from typing import Final, Literal
from urllib.parse import parse_qsl, urlsplit
@ -12,7 +12,15 @@ from urllib.parse import parse_qsl, urlsplit
import httpx
import openai
import pytest
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from integration._support.client import (
Gateway,
Scenario,
eventually,
gateway_from_environment,
object_value,
string_value,
)
from integration._support.process import owned_proxy
from integration._support.responses_vendor import same_response
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.worker_readiness import wait_for_worker_readiness
@ -26,6 +34,16 @@ _PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "
_ALIASES: Final = ("/v1", "", "/openai/v1")
@pytest.fixture(scope="module")
def cancel_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
directory: Final = tmp_path_factory.mktemp("responses-cancel")
with (
gateway_from_environment() as shared,
owned_proxy(shared, directory, {"PROXY_BATCH_POLLING_ENABLED": "false"}, workers=_PROXY_WORKERS) as owned,
):
yield owned
class _DeleteResponse(BaseModel):
model_config = ConfigDict(extra="forbid")
@ -737,7 +755,7 @@ def test_delete_returns_the_client_held_id_not_the_provider_id(
@pytest.mark.parametrize("prefix", _ALIASES)
def test_cancel_forwards_provider_id_and_returns_cancelled_response(
gateway: Gateway,
cancel_gateway: Gateway,
prefix: str,
) -> None:
raw_id: Final = f"resp_cancel_{uuid.uuid4().hex}"
@ -756,9 +774,9 @@ def test_cancel_forwards_provider_id_and_returns_cancelled_response(
assert request.body in (b"", b"{}")
return Reply(body=_response(raw_id, "cancelled"))
with wire_server(respond) as wire, gateway.scenario() as scenario:
with wire_server(respond) as wire, cancel_gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
client: Final = _client(gateway, gateway.key, prefix)
client: Final = _client(cancel_gateway, cancel_gateway.key, prefix)
created: Final = client.responses.create(model=model, input="cancel lifecycle", background=True)
cancelled: Final = client.responses.cancel(created.id)
assert cancelled.id != raw_id and same_response(cancelled.id, created.id), cancelled.model_dump_json()

16
tests/integration/spend/test_session_grouped_logs_pages.py
View file

@ -319,6 +319,18 @@ class _ExplainedReads:
self.database: Final = database
self.rows_scanned: tuple[int, ...] = ()
@asynccontextmanager
async def tx(self, *, timeout: timedelta) -> AsyncIterator["_ExplainedReads"]:
async with self.database.tx(timeout=timeout) as transaction:
reads: Final = _ExplainedReads(transaction)
try:
yield reads
finally:
self.rows_scanned = (*self.rows_scanned, *reads.rows_scanned)
async def execute_raw(self, query: str, *args: object) -> int:
return await self.database.execute_raw(query, *args)
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
explained: Final = await self.database.query_raw(f"EXPLAIN (ANALYZE, FORMAT JSON) {query}", *args)
plan: Final = _Explained.model_validate(explained[0]["QUERY PLAN"][0]).plan
@ -384,14 +396,14 @@ async def test_a_page_cut_inside_a_tie_reads_on_until_the_tie_is_whole() -> None
@dataclass(frozen=True, slots=True)
class _ExplainedPrisma:
db: _ExplainedReads
read_db: _ExplainedReads
async def test_a_pages_representatives_are_found_by_session_not_by_scanning_the_window() -> None:
async with _spend_logs_copy(*_generated_sessions(rows_per_session=10)) as database:
reads: Final = _ExplainedReads(database)
representatives: Final = await _fetch_session_representatives(
prisma_client=_ExplainedPrisma(reads), # pyright: ignore[reportArgumentType] # only .db is used
prisma_client=_ExplainedPrisma(reads), # pyright: ignore[reportArgumentType] # only .read_db is used
where_clause=_COPY_WINDOW,
sql_params=_COPY_WINDOW_PARAMS,
next_param_index=len(_COPY_WINDOW_PARAMS) + 1,

236
tests/test_litellm_rust/cache/test_python_cache.py vendored
View file

@ -5,13 +5,12 @@ import pytest
from pydantic import TypeAdapter
import litellm
from litellm import _v2
from litellm.caching.caching import Cache, CacheMode
from litellm.caching.caching_handler import (
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
)
from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.hooks.model_max_budget_limiter import (
PROXY_VirtualKeyModelMaxBudgetLimiter,
model_budget_spend_cache_key,
@ -19,10 +18,8 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.utils import InternalUsageCache
from litellm.rust_bridge.catalog import Route
from litellm.types.caching import CachingSupportedCallTypes
from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.routing import native_route
@ -32,27 +29,18 @@ pytestmark = pytest.mark.requires_rust_extension
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "stream", "legacy"),
(
("chat", False, False),
("messages", False, False),
("responses", False, False),
("messages", True, False),
("messages", False, True),
("messages", True, True),
),
"stream",
(False, True),
)
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage(
async def test_legacy_cache_hit_keeps_model_budget_spend_but_accounts_for_usage(
recording_server: RecordingServer,
route: Literal["chat", "messages", "responses"],
stream: bool,
native: bool,
monkeypatch: pytest.MonkeyPatch,
legacy: bool,
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
litellm.cache = Cache() if legacy else _v2.Cache.memory()
litellm.cache = Cache()
counters: Final = litellm.DualCache()
budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters)
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
@ -72,7 +60,7 @@ async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage(
}
spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h")
token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens")
first: Final = await invoke(route, recording_server, options, native=native)
first: Final = await invoke("messages", recording_server, options, native=native)
if stream:
await collect(first)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
@ -81,13 +69,13 @@ async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage(
first_log: Final = TypeAdapter(dict[str, object]).validate_python(first_events[0].kwargs)
first_payload: Final = TypeAdapter(dict[str, object]).validate_python(first_log["standard_logging_object"])
expected_cost: Final = TypeAdapter(float).validate_python(first_log["response_cost"])
usage: Final = RESPONSES_RESPONSE["usage"] if route == "responses" else MESSAGES_RESPONSE["usage"]
usage: Final = MESSAGES_RESPONSE["usage"]
expected_tokens: Final = usage["input_tokens"] + usage["output_tokens"]
assert expected_cost > 0
assert counters.get_cache(spend_key) == pytest.approx(expected_cost)
assert counters.get_cache(token_key) == first_payload["total_tokens"] == expected_tokens
second: Final = await invoke(route, recording_server, options, native=native)
second: Final = await invoke("messages", recording_server, options, native=native)
if stream:
await collect(second)
await drain_logging()
@ -101,7 +89,7 @@ async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage(
assert cached_payload["cache_hit"] is True
assert cached_payload["id"] != first_payload["id"]
assert cached_payload["custom_llm_provider"] == first_payload["custom_llm_provider"]
assert cached_payload["custom_llm_provider"] == ("openai" if route == "responses" else "anthropic"), {
assert cached_payload["custom_llm_provider"] == "anthropic", {
"miss_provider": first_log.get("custom_llm_provider"),
"hit_provider": cached_log.get("custom_llm_provider"),
}
@ -111,110 +99,22 @@ async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage(
@pytest.mark.asyncio
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
@pytest.mark.parametrize("backend", ("disabled", "memory", "redis"))
async def test_response_cache_backend_does_not_control_coordination(
recording_server: RecordingServer,
native: bool,
monkeypatch: pytest.MonkeyPatch,
backend: Literal["disabled", "memory", "redis"],
redis_url: str,
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
litellm.cache = (
None
if backend == "disabled"
else _v2.Cache.memory()
if backend == "memory"
else _v2.Cache.redis(redis_url, namespace="independent-coordination")
)
recording_server.expected_requests = 2 if backend == "disabled" else 1
counters: Final = litellm.DualCache()
budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters)
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
InternalUsageCache(counters), model_group_resolver=lambda model: model
)
key_hash: Final = "b" * 64
identity: Final = UserAPIKeyAuth(api_key=key_hash, rpm_limit=2, tpm_limit=1000, max_parallel_requests=1)
spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h")
request_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "requests")
token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens")
parallel_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "max_parallel_requests")
expected_tokens: Final = MESSAGES_RESPONSE["usage"]["input_tokens"] + MESSAGES_RESPONSE["usage"]["output_tokens"]
recorder: Final = RecordingLogger()
async def request(call_id: str, successes: int) -> object:
data: Final = {
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"litellm_call_id": call_id,
"max_tokens": 32,
"metadata": {
"user_api_key": key_hash,
"model_group": "cached-model",
"user_api_key_model_max_budget": {"cached-model": {"max_budget": 1, "budget_duration": "1h"}},
},
}
await limiter.async_pre_call_hook(identity, counters, data, "acompletion")
assert len(TypeAdapter(dict[str, float]).validate_python(counters.get_cache(parallel_key))) == 1
response: Final = await invoke(
"chat", recording_server, {**data, "callbacks": [budget, limiter, recorder]}, native=native
)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await recorder.wait_for_async("async_log_success_event", count=successes)
return response
await asyncio.create_task(request("cache-miss", 1))
assert counters.get_cache(parallel_key) == {}
assert counters.get_cache(request_key) == 1
assert counters.get_cache(token_key) == expected_tokens
first_events: Final = await recorder.wait_for_async("async_log_success_event")
first_cost: Final = TypeAdapter(float).validate_python(first_events[0].kwargs["response_cost"])
assert first_cost > 0
assert counters.get_cache(spend_key) == pytest.approx(first_cost)
await asyncio.create_task(request("cache-hit", 2))
expected_spend: Final = first_cost * recording_server.expected_requests
assert counters.get_cache(spend_key) == pytest.approx(expected_spend)
assert len(recording_server.requests) == recording_server.expected_requests
assert counters.get_cache(parallel_key) == {}
assert counters.get_cache(request_key) == 2
assert counters.get_cache(token_key) == 2 * expected_tokens
with pytest.raises(litellm.RateLimitError):
await asyncio.create_task(request("over-rpm-limit", 3))
assert len(recording_server.requests) == recording_server.expected_requests
assert counters.get_cache(parallel_key) == {}
assert counters.get_cache(token_key) == 2 * expected_tokens
assert counters.get_cache(spend_key) == pytest.approx(expected_spend)
if litellm.cache is not None:
await litellm.cache.disconnect()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
)
async def test_cache_controls_and_backend_credential_key_semantics(
recording_server: RecordingServer,
route: Literal["chat", "messages", "responses"],
legacy: bool,
) -> None:
recording_server.expected_requests = 3 if legacy else 4
litellm.cache = Cache() if legacy else _v2.Cache.memory()
await invoke(route, recording_server, {"cache": {"no-store": True}})
await invoke(route, recording_server, {})
await invoke(route, recording_server, {})
async def test_legacy_cache_controls_and_backend_credential_key_semantics(recording_server: RecordingServer) -> None:
recording_server.expected_requests = 3
litellm.cache = Cache()
await invoke("messages", recording_server, {"cache": {"no-store": True}})
await invoke("messages", recording_server, {})
await invoke("messages", recording_server, {})
assert len(recording_server.requests) == 2
await invoke(route, recording_server, {"cache": {"no-cache": True}})
await invoke(route, recording_server, {"api_key": "another-key"})
await invoke("messages", recording_server, {"cache": {"no-cache": True}})
await invoke("messages", recording_server, {"api_key": "another-key"})
assert len(recording_server.requests) == recording_server.expected_requests
@pytest.mark.asyncio
@pytest.mark.parametrize("legacy", (False, True))
async def test_v2_messages_replays_a_completed_stream(recording_server: RecordingServer, legacy: bool) -> None:
async def test_legacy_messages_replays_a_completed_stream(recording_server: RecordingServer) -> None:
recording_server.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS)
litellm.cache = Cache() if legacy else _v2.Cache.memory()
litellm.cache = Cache()
recorder: Final = RecordingLogger()
parameters: Final = {
"model": MESSAGES_MODEL,
@ -280,74 +180,30 @@ async def test_rust_messages_uses_a_legacy_cache_without_python_inference(
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ("chat", "messages", "responses"))
@pytest.mark.parametrize("native", (False, True))
@pytest.mark.parametrize("excluded", (None, [], ["embedding"]))
async def test_v2_cache_honors_supported_call_types_for_reads_and_writes(
async def test_legacy_cache_default_off_requires_opt_in_even_for_existing_entries(
recording_server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
route: Literal["chat", "messages", "responses"],
native: bool,
excluded: list[CachingSupportedCallTypes] | None,
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
litellm.cache = _v2.Cache.memory()
call_type: Final[CachingSupportedCallTypes] = (
"acompletion" if route == "chat" else "anthropic_messages" if route == "messages" else "aresponses"
)
recording_server.expected_requests = 4
litellm.cache.supported_call_types = excluded
await invoke(route, recording_server, {}, native=native)
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 2
litellm.cache.supported_call_types = [call_type]
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 3
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 3
litellm.cache.supported_call_types = excluded
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 4
@pytest.mark.asyncio
@pytest.mark.parametrize("native", (False, True))
@pytest.mark.parametrize(
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
)
async def test_v2_default_off_requires_opt_in_even_for_existing_entries(
recording_server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
route: Literal["chat", "messages", "responses"],
native: bool,
legacy: bool,
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
litellm.cache = Cache() if legacy else _v2.Cache.memory()
litellm.cache = Cache()
litellm.cache.mode = CacheMode.default_off
recording_server.expected_requests = 4
await invoke(route, recording_server, {}, native=native)
await invoke(route, recording_server, {}, native=native)
await invoke("messages", recording_server, {}, native=native)
await invoke("messages", recording_server, {}, native=native)
assert len(recording_server.requests) == 2
await invoke(route, recording_server, {"cache": {"use-cache": True}}, native=native)
await invoke("messages", recording_server, {"cache": {"use-cache": True}}, native=native)
assert len(recording_server.requests) == 3
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await invoke(route, recording_server, {"cache": {"use-cache": True}}, native=native)
await invoke("messages", recording_server, {"cache": {"use-cache": True}}, native=native)
assert len(recording_server.requests) == 3
await invoke(route, recording_server, {}, native=native)
await invoke("messages", recording_server, {}, native=native)
assert len(recording_server.requests) == 4
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
)
async def test_cache_lookup_uses_backend_request_callback_semantics(
recording_server: RecordingServer,
route: Literal["chat", "messages", "responses"],
legacy: bool,
) -> None:
async def test_legacy_cache_lookup_uses_backend_request_callback_semantics(recording_server: RecordingServer) -> None:
class Rewrite(RecordingLogger):
temperature = 0.1
@ -356,22 +212,40 @@ async def test_cache_lookup_uses_backend_request_callback_semantics(
super().log_pre_api_call(model, messages, kwargs)
logger: Final = Rewrite()
litellm.cache = Cache() if legacy else _v2.Cache.memory()
recording_server.expected_requests = 1 if legacy else 2
await invoke(route, recording_server, {"callbacks": [logger]})
first_hit: Final = await invoke(route, recording_server, {"callbacks": [logger]})
litellm.cache = Cache()
recording_server.expected_requests = 1
await invoke("messages", recording_server, {"callbacks": [logger]})
first_hit: Final = await invoke("messages", recording_server, {"callbacks": [logger]})
logger.temperature = 0.8
await invoke(route, recording_server, {"callbacks": [logger]})
second_hit: Final = await invoke(route, recording_server, {"callbacks": [logger]})
await invoke("messages", recording_server, {"callbacks": [logger]})
second_hit: Final = await invoke("messages", recording_server, {"callbacks": [logger]})
assert logger.names.count("log_pre_api_call") == 4
assert len(recording_server.requests) == recording_server.expected_requests
assert recording_server.requests[0].body["temperature"] == 0.1
if not legacy:
assert recording_server.requests[1].body["temperature"] == 0.8
assert isinstance(cache_key(first_hit), str)
assert isinstance(cache_key(second_hit), str)
if not legacy:
assert cache_key(first_hit) != cache_key(second_hit)
@pytest.mark.asyncio
async def test_rust_chat_completions_rejects_a_configured_python_cache(
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
recording_server.expected_requests = 0
monkeypatch.setattr(litellm, "cache", Cache())
with pytest.raises(RuntimeError, match="this inference route does not support configured caches"):
await invoke("chat", recording_server, {})
assert recording_server.requests == []
@pytest.mark.asyncio
async def test_rust_responses_rejects_a_configured_python_cache(
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
recording_server.expected_requests = 0
monkeypatch.setattr(litellm, "cache", Cache())
with pytest.raises(RuntimeError, match="this inference route does not support configured caches"):
await invoke("responses", recording_server, {})
assert recording_server.requests == []
@pytest.mark.asyncio

0
tests/test_litellm_rust/cache/v2/__init__.py vendored
View file

207
tests/test_litellm_rust/cache/v2/test_cache.py vendored
View file

@ -1,207 +0,0 @@
import asyncio
from collections.abc import AsyncIterator
from typing import Final, Literal
import pytest
from pydantic import TypeAdapter
import litellm
from litellm import _v2
from litellm._v2.cache import NativeBackend
from litellm.caching.caching_handler import (
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
)
from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.types.utils import ModelResponse
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE, RESPONSES_TEXT
from tests.test_litellm_rust.cache.support import cache_key, invoke, payload
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
pytestmark = pytest.mark.requires_rust_extension
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ("chat", "messages", "responses"))
@pytest.mark.parametrize("backend", ("memory", "redis"))
async def test_v2_cache_skips_provider_and_reports_one_success_per_call(
recording_server: RecordingServer,
route: Literal["chat", "messages", "responses"],
backend: Literal["memory", "redis"],
redis_url: str,
) -> None:
recording_server.expected_requests = 2
litellm.cache = _v2.Cache.memory() if backend == "memory" else _v2.Cache.redis(redis_url, namespace="headers")
recorder: Final = RecordingLogger()
first: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
await recorder.wait_for_async("async_log_success_event")
second: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
assert payload(first) == payload(second)
assert cache_key(first) is None
key: Final = cache_key(second)
assert isinstance(key, str)
assert key == get_hidden_params_dict(second)["cache_key"]
assert len(recording_server.requests) == 1
await drain_logging()
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
assert len(successes) == 2
cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs)
assert cached_log["cache_hit"] is True
assert cached_log["response_cost"] == 0
await litellm.cache.delete_cache_keys([key])
refreshed: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
assert cache_key(refreshed) is None
assert len(recording_server.requests) == 2
assert len(await recorder.wait_for_async("async_log_success_event", count=3)) == 3
await litellm.cache.disconnect()
@pytest.mark.asyncio
async def test_v2_global_cache_leaves_legacy_only_calls_usable() -> None:
litellm.cache = _v2.Cache.memory()
response: Final = await litellm.aembedding(
model="openai/cache-test-embedding",
input=["hello"],
api_key="test-key",
mock_response=[0.25, 0.75],
)
assert response.model_dump(include={"data"}) == {
"data": [{"embedding": [0.25, 0.75], "index": 0, "object": "embedding"}]
}
@pytest.mark.parametrize("route", ("chat", "responses"))
def test_v2_cache_works_through_python_inference(
recording_server: RecordingServer, route: Literal["chat", "messages", "responses"], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
litellm.cache = _v2.Cache.memory()
common: Final = {"api_key": "test-key", "api_base": recording_server.base_url}
if route == "responses":
recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
parameters: Final = {"model": RESPONSES_MODEL, "input": "hello", **common}
first: Final = litellm.responses(**parameters)
second: Final = litellm.responses(**parameters)
assert payload(first) == payload(second)
else:
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
arguments: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32, **common}
initial: Final = litellm.completion(**arguments)
cached: Final = litellm.completion(**arguments)
assert isinstance(initial, ModelResponse) and isinstance(cached, ModelResponse)
assert (
initial.choices[0].message.content
== cached.choices[0].message.content
== MESSAGES_RESPONSE["content"][0]["text"]
)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
async def test_v2_facade_and_backend_share_storage_and_management() -> None:
cache: Final = _v2.Cache.memory()
await cache.async_add_cache({"answer": 7}, cache_key="shared")
assert cache.get_cache(cache_key="shared") == {"answer": 7}
assert await cache.ping() is True
await cache.delete_cache_keys(["shared"])
assert await cache.async_get_cache(cache_key="shared") is None
cache.add_cache({"answer": 8}, cache_key="flush")
backend: Final = cache.cache
assert isinstance(backend, NativeBackend)
backend.flush_cache()
assert cache.get_cache(cache_key="flush") is None
await cache.disconnect()
@pytest.mark.asyncio
@pytest.mark.parametrize("control", ("s-maxage", "s-max-age"))
async def test_v2_native_cache_accepts_existing_freshness_aliases(
recording_server: RecordingServer, control: str
) -> None:
litellm.cache = _v2.Cache.memory()
first: Final = await invoke("responses", recording_server, {})
second: Final = await invoke("responses", recording_server, {"cache": {control: 600}})
assert payload(first) == payload(second)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
async def test_v2_cache_does_not_force_native_responses_streaming(
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
litellm.cache = _v2.Cache.memory()
recording_server.default_response = ResponseSpec(
body=None,
events=(
("response.created", {"type": "response.created", "sequence_number": 0, "response": RESPONSES_RESPONSE}),
(
"response.completed",
{"type": "response.completed", "sequence_number": 1, "response": RESPONSES_RESPONSE},
),
),
)
response: Final = await litellm.aresponses(
model=RESPONSES_MODEL,
input="hello",
stream=True,
caching=False,
api_key="test-key",
api_base=recording_server.base_url,
)
assert isinstance(response, AsyncIterator)
chunks: Final = [chunk async for chunk in response]
assert chunks[-1].type == "response.completed"
assert chunks[-1].response.output[0].content[0].text == RESPONSES_TEXT
@pytest.mark.asyncio
async def test_v2_cache_works_through_python_messages(
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
litellm.cache = _v2.Cache.memory()
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
parameters: Final = {
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"max_tokens": 32,
"api_key": "test-key",
"api_base": recording_server.base_url,
}
first: Final = await litellm.anthropic_messages(**parameters)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
second: Final = await litellm.anthropic_messages(**parameters)
assert payload(first) == payload(second)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_v2_redis_flush_only_removes_its_namespace(
redis_url: str, recording_server: RecordingServer, asynchronous: bool
) -> None:
own: Final = _v2.Cache.redis(redis_url, namespace="flush-own")
other: Final = _v2.Cache.redis(redis_url, namespace="flush-other")
litellm.cache = own
recording_server.expected_requests = 2
await own.async_add_cache({"answer": "own"}, cache_key="shared")
await other.async_add_cache({"answer": "other"}, cache_key="shared")
await invoke("responses", recording_server, {})
hit: Final = await invoke("responses", recording_server, {})
assert isinstance(cache_key(hit), str)
assert await own.async_get_cache(cache_key="shared") == {"answer": "own"}
backend: Final = own.cache
assert isinstance(backend, NativeBackend)
if asynchronous:
await backend.async_flush_cache()
else:
backend.flush_cache()
assert await own.async_get_cache(cache_key="shared") is None
assert await other.async_get_cache(cache_key="shared") == {"answer": "other"}
refreshed: Final = await invoke("responses", recording_server, {})
assert cache_key(refreshed) is None
assert len(recording_server.requests) == 2
await own.disconnect()
await other.disconnect()

54
tests/unit/mcp_client/test_mcp_client.py
View file

@ -3,6 +3,7 @@ import base64
import importlib
import json
import os
import queue
import sys
from collections.abc import AsyncIterator, Callable
from contextlib import AbstractAsyncContextManager, asynccontextmanager
@ -74,6 +75,7 @@ from litellm.proxy.mcp_server.outbound_credentials.httpx_auth import StaticHeade
from litellm.types.mcp import MCPAuth, MCPStdioConfig, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from tests._support.clock_loops import AutojumpClockLoop, ManualClockLoop
from tests._support.mcp_count_tool_upstream import count_tool_upstream, methods_received
_JSONRPC_MESSAGE_ADAPTER: Final = TypeAdapter(JSONRPCMessage)
@ -5250,3 +5252,55 @@ async def test_notification_callback_keeps_stream_errors_separate_from_protocol_
await _receive_session_message(failure, stream_error=stream_error, on_notification=receiver)
assert stream_error.result() is first
receiver.assert_not_awaited()
@pytest.mark.asyncio
async def test_a_skipping_client_never_asks_the_upstream_for_its_tools() -> None:
methods_seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue()
client: Final = _MockTransportClient(
count_tool_upstream({"n": 1}, methods_seen),
server_url="https://upstream.example/mcp",
transport_type=MCPTransport.http,
skip_tool_output_validation=True,
)
result: Final = await client.call_tool(CallToolRequestParams(name="count", arguments={}))
assert isinstance(result, CallToolResult)
assert result.is_error is False
assert "tools/list" not in methods_received(methods_seen)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("skip_validation", "rejected"), [(False, True), (True, False)], ids=["validating", "skipping"]
)
async def test_invalid_structured_output_is_rejected_unless_validation_is_skipped(
skip_validation: bool, rejected: bool
) -> None:
client: Final = _MockTransportClient(
count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue()),
server_url="https://upstream.example/mcp",
transport_type=MCPTransport.http,
skip_tool_output_validation=skip_validation,
)
result: Final = await client.call_tool(CallToolRequestParams(name="count", arguments={}))
assert isinstance(result, CallToolResult)
assert result.is_error is rejected
@pytest.mark.asyncio
async def test_a_skipped_validation_passes_the_upstreams_result_through_unchanged() -> None:
client: Final = _MockTransportClient(
count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue()),
server_url="https://upstream.example/mcp",
transport_type=MCPTransport.http,
skip_tool_output_validation=True,
)
result: Final = await client.call_tool(CallToolRequestParams(name="count", arguments={}))
assert isinstance(result, CallToolResult)
assert result.structured_content == {"n": "not-an-integer"}

192
tests/unit/proxy/mcp_server/test_upstream.py Normal file
View file

@ -0,0 +1,192 @@
import queue
from collections.abc import Iterator, Mapping
from typing import Final
from unittest.mock import MagicMock
import httpx2
import pytest
from mcp.server.context import ServerRequestContext
from mcp.types import CallToolResult
from litellm.proxy.mcp_server.mcp_context import active_mcp_request_ctx_var
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
from litellm.proxy.mcp_server.outbound_credentials.resolver import UpstreamCredentialProvider
from litellm.proxy.mcp_server.upstream import prepare_upstream_client
from litellm.types.mcp import MCPAuth, MCPTransport, MCPUpstreamProtocol
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from tests._support.mcp_count_tool_upstream import count_tool_upstream, methods_received
_HTTP_URL: Final = "https://upstream.example/mcp"
_SERVER_SHAPES: Final = pytest.mark.parametrize(
("server_fields", "per_request_token", "forwarded_token"),
[
({"transport": MCPTransport.http, "url": _HTTP_URL}, None, None),
(
{
"transport": MCPTransport.http,
"url": _HTTP_URL,
"auth_type": MCPAuth.bearer_token,
"authentication_token": "static-token",
},
None,
None,
),
({"transport": MCPTransport.http, "url": _HTTP_URL, "auth_type": MCPAuth.bearer_token}, "caller-key", None),
({"transport": MCPTransport.http, "url": _HTTP_URL, "auth_type": MCPAuth.oauth_delegate}, None, "caller-token"),
],
ids=["no-auth", "static-bearer-token", "per-request-token", "forwarded-caller-token"],
)
@pytest.fixture
def _active_mcp_request_context() -> Iterator[None]:
context: Final = ServerRequestContext(
session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None
)
token: Final = active_mcp_request_ctx_var.set(context)
try:
yield
finally:
active_mcp_request_ctx_var.reset(token)
def _server(server_fields: Mapping[str, object], skip_setting: bool | None) -> MCPServer:
return MCPServer.model_validate(
{
**server_fields,
"server_id": "skip-validation",
"name": "skip_validation",
"mcp_info": {} if skip_setting is None else {"skip_tool_output_validation": skip_setting},
}
)
async def _call_count_tool(
manager: MCPServerManager,
server: MCPServer,
*,
per_request_token: str | None = None,
forwarded_token: str | None = None,
gateway_session_id: str | None = None,
) -> CallToolResult:
result: Final = await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="count",
arguments={},
tasks=[],
mcp_auth_header=per_request_token,
mcp_server_auth_headers=None,
oauth2_headers=None if forwarded_token is None else {"Authorization": f"Bearer {forwarded_token}"},
raw_headers=None if gateway_session_id is None else {"mcp-session-id": gateway_session_id},
proxy_logging_obj=None,
)
assert isinstance(result, CallToolResult)
return result
@pytest.mark.asyncio
@_SERVER_SHAPES
@pytest.mark.parametrize(
("skip_setting", "rejected"), [(None, True), (False, True), (True, False)], ids=["unset", "off", "on"]
)
async def test_the_setting_alone_decides_whether_a_tool_result_is_checked(
server_fields: Mapping[str, object],
per_request_token: str | None,
forwarded_token: str | None,
skip_setting: bool | None,
rejected: bool,
) -> None:
upstream: Final = count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue())
manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(upstream))
result: Final = await _call_count_tool(
manager,
_server(server_fields, skip_setting),
per_request_token=per_request_token,
forwarded_token=forwarded_token,
)
assert result.is_error is rejected
@pytest.mark.asyncio
@_SERVER_SHAPES
async def test_a_skipping_server_is_never_asked_for_its_tools(
server_fields: Mapping[str, object], per_request_token: str | None, forwarded_token: str | None
) -> None:
methods_seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue()
manager: Final = MCPServerManager(
upstream_transport=httpx2.MockTransport(count_tool_upstream({"n": 1}, methods_seen))
)
result: Final = await _call_count_tool(
manager,
_server(server_fields, True),
per_request_token=per_request_token,
forwarded_token=forwarded_token,
)
assert result.is_error is False
assert "tools/list" not in methods_received(methods_seen)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("first_setting", "second_setting", "rejected_in_order"),
[(True, False, (False, True)), (False, True, (True, False))],
ids=["turned-off", "turned-on"],
)
async def test_a_changed_setting_applies_to_the_next_call_in_a_live_gateway_session(
_active_mcp_request_context: None,
first_setting: bool,
second_setting: bool,
rejected_in_order: tuple[bool, bool],
) -> None:
upstream: Final = count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue())
manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(upstream))
manager.track_gateway_session("gateway-1")
fields: Final = {"transport": MCPTransport.http, "url": _HTTP_URL}
try:
first: Final = await _call_count_tool(manager, _server(fields, first_setting), gateway_session_id="gateway-1")
second: Final = await _call_count_tool(manager, _server(fields, second_setting), gateway_session_id="gateway-1")
finally:
await manager.close_all_upstream_sessions()
assert (first.is_error, second.is_error) == rejected_in_order
@pytest.mark.asyncio
async def test_a_stdio_upstream_client_carries_the_setting(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true")
server: Final = _server({"transport": MCPTransport.stdio, "command": "python", "args": ["server.py"]}, True)
client: Final = await prepare_upstream_client(
server, provider=UpstreamCredentialProvider(), root_path="/mcp", protocol_version="auto"
)
assert client.skip_tool_output_validation is True
@pytest.mark.asyncio
@pytest.mark.parametrize(
("protocol_version", "expected"), [("auto", True), ("2025-11-25", True), ("2026-07-28", False)]
)
async def test_upstream_client_ignores_the_setting_where_the_call_already_lists_tools(
protocol_version: MCPUpstreamProtocol, expected: bool
) -> None:
server: Final = MCPServer.model_validate(
{
"server_id": "skip-validation",
"name": "skip_validation",
"url": _HTTP_URL,
"transport": MCPTransport.http,
"protocol_version": protocol_version,
"mcp_info": {"skip_tool_output_validation": True},
}
)
client: Final = await prepare_upstream_client(
server, provider=UpstreamCredentialProvider(), root_path="/mcp", protocol_version=protocol_version
)
assert client.skip_tool_output_validation is expected

33
tests/unit/proxy/test__types.py
View file

@ -520,6 +520,39 @@ def test_mcp_metadata_rejects_unavailable_upstream_protocol(revision):
model.model_validate(payload)
@pytest.mark.parametrize("request_model", (NewMCPServerRequest, UpdateMCPServerRequest))
@pytest.mark.parametrize("value", ["true", "false", 1, 0, None])
def test_mcp_metadata_rejects_a_non_boolean_skip_tool_output_validation(
request_model: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], value: object
) -> None:
payload: Final = {
"server_id": "test",
"transport": "http",
"url": "https://example.com/mcp",
"mcp_info": {"skip_tool_output_validation": value},
}
with pytest.raises(ValidationError):
request_model.model_validate(payload)
@pytest.mark.parametrize("request_model", (NewMCPServerRequest, UpdateMCPServerRequest))
@pytest.mark.parametrize("value", [True, False])
def test_mcp_metadata_accepts_a_boolean_skip_tool_output_validation(
request_model: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], value: bool
) -> None:
parsed: Final = request_model.model_validate(
{
"server_id": "test",
"transport": "http",
"url": "https://example.com/mcp",
"mcp_info": {"skip_tool_output_validation": value},
}
)
assert parsed.mcp_info == {"skip_tool_output_validation": value}
MCP_SERVER_REQUESTS = (NewMCPServerRequest, UpdateMCPServerRequest)
STDIO_SERVER_FIELDS = {"server_id": "stdio-1", "transport": "stdio", "command": "python", "args": ["server.py"]}

2
tests/unit/test_internal_context.py
View file

@ -25,7 +25,7 @@ _REPO: Final = Path(__file__).resolve().parents[2]
_REDIS_PRODUCER_ROOTS: Final = ("litellm", "enterprise")
# The cache implementations and facades: they emit the service events, their callers declare the family.
_CACHE_LAYER_DIRS: Final = ("litellm/caching", "litellm/_v2/cache")
_CACHE_LAYER_DIRS: Final = ("litellm/caching",)
# Helpers that act on a cache handed in by the declaring caller, or forward to the response-cache facade.
_CACHE_PARAMETER_HELPERS: Final = frozenset(
{

10
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -1153,16 +1153,6 @@
"count": 2
}
},
"src/app/model_hub/page.tsx": {
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/app/model_hub_table/page.tsx": {
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/AIHub/ModelHubTable.test.tsx": {
"max-params": {
"count": 1

24
ui/litellm-dashboard/src/app/model_hub/page.tsx
View file

@ -1,27 +1,19 @@
"use client";
import React, { Suspense, useEffect, useState } from "react";
import { useSearchParams } from "next/navigation";
import PublicModelHubPage from "@/components/public_model_hub";
import { Suspense } from "react";
function PublicModelHubContent() {
const searchParams = useSearchParams()!;
const key = searchParams.get("key");
const [accessToken, setAccessToken] = useState<string | null>(null);
import PublicAiHub from "@/components/publicModelHub/PublicAiHub";
useEffect(() => {
if (!key) {
return;
}
setAccessToken(key);
}, [key]);
return <PublicModelHubPage accessToken={accessToken} />;
function PublicAiHubPage() {
const accessToken = useSearchParams()?.get("key") ?? null;
return <PublicAiHub accessToken={accessToken} />;
}
export default function PublicModelHub() {
return (
<Suspense fallback={<div className="flex items-center justify-center min-h-screen">Loading...</div>}>
<PublicModelHubContent />
<Suspense fallback={<div className="flex min-h-screen items-center justify-center">Loading...</div>}>
<PublicAiHubPage />
</Suspense>
);
}

33
ui/litellm-dashboard/src/app/model_hub_table/page.tsx
View file

@ -1,36 +1,19 @@
"use client";
import React, { Suspense, useEffect, useState } from "react";
import { useSearchParams } from "next/navigation";
import ModelHubTable from "@/components/AIHub/ModelHubTable";
import { Suspense } from "react";
function PublicModelHubTableContent() {
const searchParams = useSearchParams()!;
const key = searchParams.get("key");
const [accessToken, setAccessToken] = useState<string | null>(null);
import PublicAiHub from "@/components/publicModelHub/PublicAiHub";
useEffect(() => {
if (!key) {
return;
}
setAccessToken(key);
}, [key]);
return (
<ModelHubTable
accessToken={accessToken}
publicPage={true}
premiumUser={false}
userRole={null}
canOpenPlayground={false}
canAddModels={false}
/>
);
function PublicAiHubPage() {
const accessToken = useSearchParams()?.get("key") ?? null;
return <PublicAiHub accessToken={accessToken} />;
}
export default function PublicModelHubTable() {
return (
<Suspense fallback={<div className="flex items-center justify-center min-h-screen">Loading...</div>}>
<PublicModelHubTableContent />
<Suspense fallback={<div className="flex min-h-screen items-center justify-center">Loading...</div>}>
<PublicAiHubPage />
</Suspense>
);
}

11
ui/litellm-dashboard/src/components/navbar.tsx
View file

@ -61,7 +61,7 @@ const Navbar: React.FC<NavbarProps> = ({
return (
<nav className="sticky top-0 z-chrome border-b border-border bg-card">
<div className="w-full">
<div className="flex h-14 items-center px-4">
<div className={cn("flex h-14 items-center px-4", isPublicPage && "gap-4 sm:px-6")}>
<div className="flex shrink-0 items-center">
{onToggleSidebar && (
<button
@ -82,7 +82,12 @@ const Navbar: React.FC<NavbarProps> = ({
<div className="flex items-center gap-2">
<Link href={uiHref("")} className="flex items-center">
<div className="relative">
<div className="flex h-10 max-w-48 items-center justify-center overflow-hidden">
<div
className={cn(
"flex h-10 max-w-48 items-center justify-center overflow-hidden",
isPublicPage && "h-8 max-w-36 sm:h-10 sm:max-w-48",
)}
>
<img src={imageUrl} alt="LiteLLM Brand" className={cn(NAV_LOGO_CLASS_NAME, "dark:hidden")} />
<img
src={darkImageUrl}
@ -141,7 +146,7 @@ const Navbar: React.FC<NavbarProps> = ({
</nav>
{!hideCommunityLinks && (
<div className="flex shrink-0 items-center border-l border-border pl-4">
<div className="hidden shrink-0 items-center border-l border-border pl-4 sm:flex">
<CommunityEngagementButtons />
</div>
)}

80
ui/litellm-dashboard/src/components/publicModelHub/AgentInspector.tsx Normal file
View file

@ -0,0 +1,80 @@
"use client";
import { ExternalLink } from "lucide-react";
import type { AgentCard } from "@/components/PublicModelHubTableColumns";
import { Button } from "@/components/ui/button";
import { Field, FieldDescription } from "@/components/ui/field";
import { Item, ItemActions, ItemContent, ItemDescription, ItemGroup, ItemTitle } from "@/components/ui/item";
import { agentCapabilities } from "./hubFormat";
import { Badges, CodeBlock, EndpointField, Facts, PaneSection } from "./HubParts";
import { agentSnippets } from "./hubSnippets";
import { InspectorHeader } from "./InspectorHeader";
const modes = (agent: AgentCard) => {
const input = (agent.defaultInputModes ?? []).join(", ");
const output = (agent.defaultOutputModes ?? []).join(", ");
return input || output ? `${input || "-"} / ${output || "-"}` : null;
};
export function AgentInspector({ agent }: { agent: AgentCard }) {
const capabilities = agentCapabilities(agent);
return (
<div className="flex flex-col gap-6 p-6">
<InspectorHeader title={agent.name} copyLabel="Copy agent name" description={agent.description} />
<Facts
groups={[
[
{ label: "Version", value: agent.version },
{ label: "Owner", value: agent.provider?.organization },
{ label: "Protocol", value: agent.protocolVersion && `A2A ${agent.protocolVersion}` },
],
[
{ label: "Input / output", wide: true, value: modes(agent) },
{ label: "Capabilities", wide: true, value: capabilities.length > 0 && <Badges items={capabilities} /> },
],
]}
/>
<EndpointField id={`agent-endpoint-${agent.name}`} value={agent.url} />
{agent.documentationUrl && (
<Button
variant="link"
nativeButton={false}
className="h-auto self-start px-0"
render={<a href={agent.documentationUrl} target="_blank" rel="noopener noreferrer" />}
>
Read the documentation
<ExternalLink />
</Button>
)}
{(agent.skills ?? []).length > 0 && (
<PaneSection title="Skills">
<ItemGroup className="gap-2">
{agent.skills.map((skill) => (
<Item key={skill.id} variant="outline" size="sm" role="listitem">
<ItemContent>
<ItemTitle>{skill.name}</ItemTitle>
<ItemDescription>{skill.description}</ItemDescription>
</ItemContent>
<ItemActions>
<Badges items={skill.tags ?? []} />
</ItemActions>
</Item>
))}
</ItemGroup>
</PaneSection>
)}
<PaneSection title="Call this agent">
{agentSnippets(agent).map((snippet, index) => (
<Field key={snippet.title}>
<FieldDescription>
Step {index + 1}: {snippet.title}
</FieldDescription>
<CodeBlock code={snippet.code} />
</Field>
))}
</PaneSection>
</div>
);
}

91
ui/litellm-dashboard/src/components/publicModelHub/FilterChip.tsx Normal file
View file

@ -0,0 +1,91 @@
"use client";
import { CirclePlus, X } from "lucide-react";
import { Button } from "@/components/ui/button";
import { ButtonGroup } from "@/components/ui/button-group";
import {
Combobox,
ComboboxContent,
ComboboxEmpty,
ComboboxInput,
ComboboxItem,
ComboboxList,
ComboboxTrigger,
} from "@/components/ui/combobox";
import { cn } from "@/lib/cva.config";
import { firstPlus } from "./hubFormat";
import { ProviderAvatar } from "./HubParts";
interface FilterOption {
value: string;
label: string;
provider?: string;
}
export function FilterChip({
label,
noun,
options,
value,
onChange,
}: {
label: string;
noun: string;
options: FilterOption[];
value: string[];
onChange: (value: string[]) => void;
}) {
const byValue = new Map(options.map((option) => [option.value, option]));
const selected = value.map((v) => byValue.get(v)?.label ?? v);
const active = value.length > 0;
return (
<ButtonGroup>
<Combobox<FilterOption, true>
multiple
items={options}
value={options.filter((option) => value.includes(option.value))}
onValueChange={(next) => onChange(next.map((option) => option.value))}
itemToStringLabel={(option) => option.label}
isItemEqualToValue={(a, b) => a.value === b.value}
>
<ComboboxTrigger
render={
<Button
variant="outline"
size="sm"
className={cn("[&>svg:last-child]:hidden", !active && "border-dashed text-muted-foreground")}
/>
}
>
{!active && <CirclePlus />}
{label}
{active && <span className="font-normal text-muted-foreground">{firstPlus(selected)}</span>}
</ComboboxTrigger>
<ComboboxContent className="w-auto min-w-60">
<ComboboxInput showTrigger={false} placeholder={`Search ${noun}`} />
<ComboboxEmpty>No {noun} found</ComboboxEmpty>
<ComboboxList>
{(option: FilterOption) => (
<ComboboxItem key={option.value} value={option}>
{option.provider && <ProviderAvatar provider={option.provider} />}
<span className="whitespace-nowrap">{option.label}</span>
</ComboboxItem>
)}
</ComboboxList>
</ComboboxContent>
</Combobox>
{active && (
<Button
variant="outline"
size="icon-sm"
aria-label={`Clear ${label.toLowerCase()} filter`}
onClick={() => onChange([])}
>
<X />
</Button>
)}
</ButtonGroup>
);
}

92
ui/litellm-dashboard/src/components/publicModelHub/HubEmpty.tsx Normal file
View file

@ -0,0 +1,92 @@
"use client";
import { LayoutGrid, SearchX, TriangleAlert } from "lucide-react";
import * as React from "react";
import { Button } from "@/components/ui/button";
import { Empty, EmptyContent, EmptyDescription, EmptyHeader, EmptyMedia, EmptyTitle } from "@/components/ui/empty";
import { SECTION_TITLE } from "./hubFormat";
const noResultsTitle = (noun: string, query: string, filtered: boolean) => {
if (!filtered) return `No ${noun} published yet`;
return query ? `No ${noun} match "${query}"` : `No ${noun} match these filters`;
};
export function NoResults({
icon,
noun,
query,
onClear,
}: {
icon: React.ReactNode;
noun: string;
query: string;
onClear?: () => void;
}) {
return (
<Empty>
<EmptyHeader>
<EmptyMedia variant="icon">{icon}</EmptyMedia>
<EmptyTitle className={SECTION_TITLE}>{noResultsTitle(noun, query, Boolean(onClear))}</EmptyTitle>
<EmptyDescription>
{onClear ? "Try a shorter search or fewer filters" : "Your admin publishes them from the AI Hub page"}
</EmptyDescription>
</EmptyHeader>
{onClear && (
<EmptyContent>
<Button variant="outline" size="sm" onClick={onClear}>
Clear filters
</Button>
</EmptyContent>
)}
</Empty>
);
}
export function HubNotEnabled() {
return (
<Empty className="rounded-xl border border-solid border-border">
<EmptyHeader>
<EmptyMedia variant="icon">
<LayoutGrid />
</EmptyMedia>
<EmptyTitle className={SECTION_TITLE}>This AI Hub isn&apos;t available</EmptyTitle>
<EmptyDescription>Ask your proxy admin to make it public from their Admin UI</EmptyDescription>
</EmptyHeader>
</Empty>
);
}
export function LoadFailed({ noun, onRetry }: { noun: string; onRetry: () => void }) {
return (
<Empty>
<EmptyHeader>
<EmptyMedia variant="icon">
<TriangleAlert />
</EmptyMedia>
<EmptyTitle className={SECTION_TITLE}>Couldn&apos;t load {noun}</EmptyTitle>
<EmptyDescription>The gateway returned an error. Try again in a moment</EmptyDescription>
</EmptyHeader>
<EmptyContent>
<Button variant="outline" size="sm" onClick={onRetry}>
Try again
</Button>
</EmptyContent>
</Empty>
);
}
export function NotListed({ name }: { name: string }) {
return (
<Empty className="h-full">
<EmptyHeader>
<EmptyMedia variant="icon">
<SearchX />
</EmptyMedia>
<EmptyTitle className={SECTION_TITLE}>&ldquo;{name}&rdquo; isn&apos;t listed here</EmptyTitle>
<EmptyDescription>It may have been renamed or unpublished</EmptyDescription>
</EmptyHeader>
</Empty>
);
}

103
ui/litellm-dashboard/src/components/publicModelHub/HubHeader.tsx Normal file
View file

@ -0,0 +1,103 @@
"use client";
import { ChevronDown, ExternalLink, LayoutGrid } from "lucide-react";
import { Button } from "@/components/ui/button";
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuGroup,
DropdownMenuItem,
DropdownMenuLabel,
DropdownMenuSeparator,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Field, FieldDescription } from "@/components/ui/field";
import { InputGroup, InputGroupAddon, InputGroupInput, InputGroupText } from "@/components/ui/input-group";
import { gatewayUrl, usefulLinkList } from "./hubFormat";
import { CopyIconButton } from "./HubParts";
import { StatusBadge } from "./HubStatus";
import type { HubInfo } from "./useHubData";
export type HubStatus = "loading" | "operational" | "degraded" | "unavailable";
function LinksMenu({ info }: { info: HubInfo | undefined }) {
const links = usefulLinkList(info?.useful_links);
if (links.length === 0) return null;
return (
<DropdownMenu>
<DropdownMenuTrigger render={<Button variant="outline" size="sm" />}>
Links
<ChevronDown className="text-muted-foreground" />
</DropdownMenuTrigger>
<DropdownMenuContent align="start" className="w-auto min-w-(--anchor-width)">
{links.map((link) => (
<DropdownMenuItem key={link.title} render={<a href={link.url} target="_blank" rel="noopener noreferrer" />}>
<ExternalLink />
{link.title}
</DropdownMenuItem>
))}
{info?.litellm_version && (
<>
<DropdownMenuSeparator />
<DropdownMenuGroup>
<DropdownMenuLabel className="text-xs font-normal">LiteLLM v{info.litellm_version}</DropdownMenuLabel>
</DropdownMenuGroup>
</>
)}
</DropdownMenuContent>
</DropdownMenu>
);
}
function BaseUrlField() {
const url = gatewayUrl();
return (
<Field className="w-full gap-1.5 sm:w-[420px]">
<InputGroup className="h-8">
<InputGroupAddon>
<InputGroupText className="text-xs font-medium">Base URL</InputGroupText>
</InputGroupAddon>
<InputGroupInput
readOnly
value={url}
aria-label="Base URL"
aria-describedby="hub-base-url-hint"
className="font-mono text-xs md:text-xs"
/>
<InputGroupAddon align="inline-end">
<CopyIconButton value={url} label="Copy base URL" className="size-6" />
</InputGroupAddon>
</InputGroup>
<FieldDescription id="hub-base-url-hint" className="text-xs text-balance">
Use your LiteLLM key as the API key in any OpenAI-compatible SDK
</FieldDescription>
</Field>
);
}
export function HubHeader({ info, status }: { info: HubInfo | undefined; status: HubStatus }) {
return (
<header className="flex flex-col gap-4">
<div className="flex flex-col gap-1">
<div className="flex items-center gap-2">
<h1 className="flex items-center gap-2 text-xl leading-7 font-semibold">
<LayoutGrid className="size-5" />
AI Hub
</h1>
{status === "operational" && <StatusBadge tone="success" label="Operational" />}
</div>
<p className="text-sm text-muted-foreground">
{info?.custom_docs_description || "Models, agents, MCP servers and skills you can call through this gateway"}
</p>
</div>
{status !== "unavailable" && (
<div className="flex flex-wrap items-start gap-2">
<BaseUrlField />
<LinksMenu info={info} />
</div>
)}
</header>
);
}

234
ui/litellm-dashboard/src/components/publicModelHub/HubMasterDetail.tsx Normal file
View file

@ -0,0 +1,234 @@
"use client";
import * as React from "react";
import { Item, ItemContent, ItemGroup, ItemMedia } from "@/components/ui/item";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet";
import { Skeleton } from "@/components/ui/skeleton";
import { cn } from "@/lib/cva.config";
import { NotListed } from "./HubEmpty";
import { InSheet } from "./InspectorHeader";
import { useOpenHref } from "./useHubUrlState";
import { useIsNarrow } from "./useIsNarrow";
function ListSkeleton() {
return (
<ItemGroup className="gap-0.5 p-3" aria-busy="true">
{[0, 1, 2].map((i) => (
<div key={i} role="listitem">
<Item size="sm">
<ItemMedia variant="icon" className="size-5">
<Skeleton className="size-5 rounded-full" />
</ItemMedia>
<ItemContent className="gap-1.5">
<Skeleton className="h-4 w-32" />
<Skeleton className="h-3 w-48" />
</ItemContent>
</Item>
</div>
))}
</ItemGroup>
);
}
function DetailSkeleton() {
return (
<div className="flex flex-col gap-6 p-6" aria-busy="true">
<Skeleton className="h-5 w-48" />
<div className="grid grid-cols-2 gap-x-6 gap-y-4 sm:grid-cols-4">
{[0, 1, 2, 3].map((i) => (
<div key={i} className="flex flex-col gap-1.5">
<Skeleton className="h-3 w-20" />
<Skeleton className="h-4 w-16" />
</div>
))}
</div>
<Skeleton className="h-40 w-full" />
</div>
);
}
const isPlainClick = (event: React.MouseEvent) => {
const opensElsewhere = event.metaKey || event.ctrlKey || event.shiftKey;
return event.button === 0 && !opensElsewhere;
};
function Rows<T>({
rows,
activeKey,
explicit,
getKey,
onSelect,
renderItem,
}: {
rows: T[];
activeKey: string | undefined;
explicit: boolean;
getKey: (row: T) => string;
onSelect: (key: string) => void;
renderItem: (row: T) => React.ReactNode;
}) {
const hrefFor = useOpenHref();
const follow = (key: string) => (event: React.MouseEvent) => {
if (!isPlainClick(event)) return;
event.preventDefault();
onSelect(key);
};
return (
<ItemGroup className="gap-0.5 p-3">
{rows.map((row) => {
const key = getKey(row);
const active = key === activeKey;
return (
<div key={key} role="listitem">
<Item
size="sm"
render={
<a href={hrefFor(key)} aria-current={active && explicit ? "true" : undefined} onClick={follow(key)} />
}
className={cn(active && (explicit ? "bg-muted" : "lg:bg-muted"))}
>
{renderItem(row)}
</Item>
</div>
);
})}
</ItemGroup>
);
}
const resolveSelection = <T,>(
rows: T[],
selectedKey: string | null,
opened: T | undefined,
getKey: (row: T) => string,
) => {
const chosen = selectedKey === null ? undefined : opened ?? rows.find((row) => getKey(row) === selectedKey);
return { chosen, selected: selectedKey === null ? rows[0] : chosen };
};
function useSheet<T>(
chosen: T | undefined,
missingKey: string | undefined,
getKey: (row: T) => string,
renderDetail: (row: T) => React.ReactNode,
) {
const [content, setContent] = React.useState<{ row?: T; missingKey?: string }>({ row: chosen });
if (chosen !== undefined && chosen !== content.row) setContent({ row: chosen });
if (missingKey !== undefined && missingKey !== content.missingKey) setContent({ missingKey });
const { row } = content;
return {
open: chosen !== undefined || missingKey !== undefined,
title: row ? getKey(row) : content.missingKey ?? "Details",
body: row ? renderDetail(row) : content.missingKey && <NotListed name={content.missingKey} />,
};
}
export function HubMasterDetail<T>({
switcher,
toolbar,
rows,
loading,
selectedKey,
opened,
openPending = false,
getKey,
onSelect,
renderItem,
renderDetail,
empty,
footer,
}: {
switcher?: React.ReactNode;
toolbar: React.ReactNode;
rows: T[];
loading: boolean;
selectedKey: string | null;
opened?: T;
openPending?: boolean;
getKey: (row: T) => string;
onSelect: (key: string | null) => void;
renderItem: (row: T) => React.ReactNode;
renderDetail: (row: T) => React.ReactNode;
empty: React.ReactNode;
footer?: React.ReactNode;
}) {
const narrow = useIsNarrow();
const { chosen, selected } = resolveSelection(rows, selectedKey, opened, getKey);
const isEmpty = !loading && rows.length === 0;
const settled = !loading && !openPending && !isEmpty;
const missingKey = selectedKey !== null && chosen === undefined && settled ? selectedKey : undefined;
const sheet = useSheet(chosen, missingKey, getKey, renderDetail);
const listOnly = isEmpty && chosen === undefined && !openPending;
const detail = () => {
if (loading || (selected === undefined && openPending)) return <DetailSkeleton />;
if (selected === undefined) return selectedKey && <NotListed name={selectedKey} />;
return renderDetail(selected);
};
const list = isEmpty ? (
empty
) : (
<>
<Rows
rows={rows}
activeKey={selected && getKey(selected)}
explicit={chosen !== undefined}
getKey={getKey}
onSelect={onSelect}
renderItem={renderItem}
/>
{footer}
</>
);
return (
<div
className={cn(
"grid min-h-0 grid-cols-1 overflow-hidden rounded-xl border border-border lg:h-[calc(100dvh-18.875rem)] lg:min-h-[560px]",
!listOnly && "lg:grid-cols-[minmax(0,5fr)_minmax(0,6fr)]",
)}
>
<div className={cn("flex min-h-0 flex-col", !listOnly && "lg:border-r lg:border-border")}>
<div className="flex flex-col gap-3 border-b border-border p-3">
{switcher}
<div className="flex flex-col gap-2">{toolbar}</div>
</div>
<ScrollArea className="min-h-0 flex-1">{loading ? <ListSkeleton /> : list}</ScrollArea>
</div>
{!listOnly && <ScrollArea className="hidden min-h-0 lg:block">{detail()}</ScrollArea>}
<DetailSheet open={narrow && sheet.open} title={sheet.title} onClose={() => onSelect(null)}>
{sheet.body}
</DetailSheet>
</div>
);
}
function DetailSheet({
open,
title,
onClose,
children,
}: {
open: boolean;
title: string;
onClose: () => void;
children: React.ReactNode;
}) {
const body = React.useRef<HTMLDivElement>(null);
return (
<Sheet open={open} onOpenChange={(next) => !next && onClose()}>
<SheetContent
side="bottom"
showCloseButton={false}
initialFocus={body}
className="max-h-[85dvh] gap-0 rounded-t-xl p-0 lg:hidden"
>
<SheetTitle className="sr-only">{title}</SheetTitle>
<div ref={body} tabIndex={-1} className="min-h-0 overflow-y-auto outline-none">
<InSheet.Provider value>{children}</InSheet.Provider>
</div>
</SheetContent>
</Sheet>
);
}

187
ui/litellm-dashboard/src/components/publicModelHub/HubParts.tsx Normal file
View file

@ -0,0 +1,187 @@
"use client";
import { Copy, Search, X } from "lucide-react";
import * as React from "react";
import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Field, FieldLabel } from "@/components/ui/field";
import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group";
import { ItemTitle } from "@/components/ui/item";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { cn } from "@/lib/cva.config";
import { copyToClipboard } from "@/utils/dataUtils";
import { providerLogo, providerName, SECTION_TITLE } from "./hubFormat";
export function CopyIconButton({ value, label, className }: { value: string; label: string; className?: string }) {
return (
<Tooltip>
<TooltipTrigger
render={
<Button
variant="ghost"
size="icon-sm"
aria-label={label}
className={cn("text-muted-foreground", className)}
onClick={() => void copyToClipboard(value, "Copied")}
/>
}
>
<Copy />
</TooltipTrigger>
<TooltipContent>{label}</TooltipContent>
</Tooltip>
);
}
export function CodeBlock({ code, label = "Copy code" }: { code: string; label?: string }) {
return (
<div className="relative min-w-0 rounded-lg border border-border bg-muted">
<CopyIconButton value={code} label={label} className="absolute top-1.5 right-1.5 bg-background" />
<pre className="max-h-96 overflow-auto p-3 pr-12 font-mono text-xs leading-5 whitespace-pre text-foreground">
{code}
</pre>
</div>
);
}
export function ProviderAvatar({ provider, className }: { provider: string; className?: string }) {
const logo = providerLogo(provider);
return (
<Avatar className={cn("size-4 bg-white ring-1 ring-border", className)}>
{logo && <AvatarImage src={logo} alt="" className="object-contain p-0.5" />}
<AvatarFallback className="bg-muted text-xs">{providerName(provider).charAt(0)}</AvatarFallback>
</Avatar>
);
}
export function ProviderName({ provider }: { provider: string }) {
return (
<span className="flex min-w-0 items-center gap-1.5">
<ProviderAvatar provider={provider} />
<span className="truncate">{providerName(provider)}</span>
</span>
);
}
export function Badges({ items, mono }: { items: string[]; mono?: boolean }) {
if (items.length === 0) return null;
return (
<div className="flex flex-wrap gap-1">
{items.map((item) => (
<Badge key={item} variant="secondary" className={cn("font-medium", mono && "font-mono font-normal")}>
{item}
</Badge>
))}
</div>
);
}
export function PaneSection({
title,
aside,
children,
}: {
title: string;
aside?: React.ReactNode;
children: React.ReactNode;
}) {
return (
<section className="flex flex-col gap-3">
<div className="flex items-center justify-between gap-2">
<h3 className={SECTION_TITLE}>{title}</h3>
{aside}
</div>
{children}
</section>
);
}
export type Fact = { label: string; value: React.ReactNode; wide?: boolean };
const BLANK: React.ReactNode[] = [null, undefined, false, "", "-"];
const present = (value: React.ReactNode) => !BLANK.includes(value);
export function Facts({ groups }: { groups: Fact[][] }) {
const shown = groups.map((group) => group.filter((fact) => present(fact.value))).filter((group) => group.length);
if (shown.length === 0) return null;
return (
<div className="flex flex-col gap-4">
{shown.map((group) => (
<dl key={group[0].label} className="grid grid-cols-2 gap-x-6 gap-y-4 sm:grid-cols-4">
{group.map((fact) => (
<div key={fact.label} className={cn("flex min-w-0 flex-col gap-1", fact.wide && "col-span-2")}>
<dt className="text-xs font-medium whitespace-nowrap text-muted-foreground">{fact.label}</dt>
<dd className="min-w-0 truncate text-sm tabular-nums">{fact.value}</dd>
</div>
))}
</dl>
))}
</div>
);
}
export function EndpointField({ id, value }: { id: string; value: string }) {
return (
<Field>
<FieldLabel htmlFor={id}>Endpoint</FieldLabel>
<InputGroup>
<InputGroupInput id={id} readOnly value={value} className="font-mono text-xs md:text-xs" />
<InputGroupAddon align="inline-end">
<CopyIconButton value={value} label="Copy endpoint" className="size-6" />
</InputGroupAddon>
</InputGroup>
</Field>
);
}
export function SearchInput({
value,
onChange,
placeholder,
}: {
value: string;
onChange: (value: string) => void;
placeholder: string;
}) {
return (
<InputGroup className="h-8 min-w-0 flex-1">
<InputGroupAddon>
<Search />
</InputGroupAddon>
<InputGroupInput
value={value}
onChange={(event) => onChange(event.target.value)}
placeholder={placeholder}
aria-label={placeholder}
/>
{value && (
<InputGroupAddon align="inline-end">
<InputGroupButton size="icon-xs" aria-label="Clear search" onClick={() => onChange("")}>
<X />
</InputGroupButton>
</InputGroupAddon>
)}
</InputGroup>
);
}
export function RowTitle({ children }: { children: React.ReactNode }) {
return (
<ItemTitle className="w-full min-w-0">
<span className="line-clamp-2 wrap-anywhere">{children}</span>
</ItemTitle>
);
}
export function Joined({ parts }: { parts: (string | null | undefined)[] }) {
return parts.filter(Boolean).map((part, index) => (
<React.Fragment key={part}>
{index > 0 && " · "}
<span className="whitespace-nowrap">{part}</span>
</React.Fragment>
));
}

49
ui/litellm-dashboard/src/components/publicModelHub/HubStatus.tsx Normal file
View file

@ -0,0 +1,49 @@
"use client";
import { CircleCheck, CircleDashed, CircleX } from "lucide-react";
import type { ModelGroupInfo } from "@/components/PublicModelHubTableColumns";
import { Badge } from "@/components/ui/badge";
import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card";
import { cn } from "@/lib/cva.config";
import { healthOf, type HealthTone, relativeTime } from "./hubFormat";
const TONE_ICON = { success: CircleCheck, error: CircleX, neutral: CircleDashed } as const;
const TONE_CLASS = {
success: "border-success/25 bg-success/10 text-success",
error: "border-destructive/25 bg-destructive/10 text-destructive",
neutral: "text-muted-foreground",
} as const;
export function StatusBadge({ tone, label }: { tone: HealthTone; label: string }) {
const Icon = TONE_ICON[tone];
return (
<Badge variant="outline" className={cn("font-medium", TONE_CLASS[tone])}>
<Icon data-icon="inline-start" />
{label}
</Badge>
);
}
export function HealthBadge({ model }: { model: ModelGroupInfo }) {
const badge = <StatusBadge {...healthOf(model)} />;
if (!model.health_status) return badge;
return (
<HoverCard>
<HoverCardTrigger delay={150} render={<span tabIndex={0} className="inline-flex" />}>
{badge}
</HoverCardTrigger>
<HoverCardContent align="end" className="w-56 p-3 text-xs">
<dl className="grid grid-cols-[auto_1fr] gap-x-4 gap-y-1">
<dt className="text-muted-foreground">Response time</dt>
<dd className="tabular-nums">
{model.health_response_time ? `${Math.round(model.health_response_time)} ms` : "-"}
</dd>
<dt className="text-muted-foreground">Last checked</dt>
<dd>{relativeTime(model.health_checked_at)}</dd>
</dl>
</HoverCardContent>
</HoverCard>
);
}

49
ui/litellm-dashboard/src/components/publicModelHub/InspectorHeader.tsx Normal file
View file

@ -0,0 +1,49 @@
"use client";
import { X } from "lucide-react";
import * as React from "react";
import { Button } from "@/components/ui/button";
import { SheetClose } from "@/components/ui/sheet";
import { cn } from "@/lib/cva.config";
import { SECTION_TITLE } from "./hubFormat";
import { CopyIconButton } from "./HubParts";
export const InSheet = React.createContext(false);
export function InspectorHeader({
media,
title,
copyLabel,
status,
description,
}: {
media?: React.ReactNode;
title: string;
copyLabel: string;
status?: React.ReactNode;
description?: React.ReactNode;
}) {
const inSheet = React.useContext(InSheet);
return (
<div className="flex flex-col gap-1.5">
<div className="flex min-w-0 items-center gap-2">
{media}
<h2 className={cn("min-w-0 wrap-break-word", SECTION_TITLE)}>{title}</h2>
<CopyIconButton value={title} label={copyLabel} className="size-7" />
<span className="flex-1" />
{status && <span className="hidden sm:flex">{status}</span>}
{inSheet && (
<SheetClose
render={<Button variant="ghost" size="icon-sm" aria-label="Close" className="text-muted-foreground" />}
>
<X />
</SheetClose>
)}
</div>
{status && <span className="flex sm:hidden">{status}</span>}
{description && <p className="text-sm text-muted-foreground">{description}</p>}
</div>
);
}

29
ui/litellm-dashboard/src/components/publicModelHub/McpInspector.tsx Normal file
View file

@ -0,0 +1,29 @@
"use client";
import type { MCPServerData } from "@/components/PublicModelHubTableColumns";
import { mcpAuthLabel } from "./hubFormat";
import { CodeBlock, EndpointField, Facts, PaneSection } from "./HubParts";
import { mcpEndpoint, mcpSnippet } from "./hubSnippets";
import { InspectorHeader } from "./InspectorHeader";
export function McpInspector({ server }: { server: MCPServerData }) {
return (
<div className="flex flex-col gap-6 p-6">
<InspectorHeader title={server.name} copyLabel="Copy server name" description={server.mcp_info?.description} />
<Facts
groups={[
[
{ label: "Transport", value: server.transport },
{ label: "Sign-in", value: mcpAuthLabel(server.auth_type) },
{ label: "Alias", value: server.alias },
],
]}
/>
<EndpointField id={`mcp-endpoint-${server.server_id}`} value={mcpEndpoint(server)} />
<PaneSection title="Connect a client">
<CodeBlock code={mcpSnippet(server)} />
</PaneSection>
</div>
);
}

164
ui/litellm-dashboard/src/components/publicModelHub/ModelInspector.tsx Normal file
View file

@ -0,0 +1,164 @@
"use client";
import { ChevronDown } from "lucide-react";
import * as React from "react";
import type { ModelGroupInfo } from "@/components/PublicModelHubTableColumns";
import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import {
formatPerMillion,
formatTokens,
modeLabel,
modelCapabilities,
providerName,
wildcardExample,
} from "./hubFormat";
import { Badges, CodeBlock, type Fact, Facts, PaneSection, ProviderAvatar, ProviderName } from "./HubParts";
import { curlSnippet, modelSnippet } from "./hubSnippets";
import { HealthBadge } from "./HubStatus";
import { InspectorHeader } from "./InspectorHeader";
const summary = (model: ModelGroupInfo): string | undefined => {
if (model.description) return model.description;
const providers = (model.providers ?? []).map(providerName).join(", ");
return model.mode && providers ? `${modeLabel(model.mode)} model from ${providers}` : undefined;
};
const tokens = (value: number | undefined) => (value ? `${formatTokens(value)} tokens` : null);
const modelFacts = (model: ModelGroupInfo): Fact[][] => {
const providers = model.providers ?? [];
const tags = model.custom_tags ?? [];
return [
[
{ label: "Mode", value: model.mode ? modeLabel(model.mode) : null },
{ label: "Context", value: tokens(model.max_input_tokens) },
{ label: "Max output", value: tokens(model.max_output_tokens) },
],
[
{ label: "Input / 1M tokens", value: formatPerMillion(model.input_cost_per_token) },
{ label: "Output / 1M tokens", value: formatPerMillion(model.output_cost_per_token) },
],
[
{ label: "Requests per minute", value: model.rpm },
{ label: "Tokens per minute", value: model.tpm ? formatTokens(model.tpm) : null },
],
[
{
label: providers.length > 1 ? "Providers" : "Provider",
wide: true,
value: providers.length > 0 && (
<span className="flex flex-wrap gap-x-4 gap-y-1">
{providers.map((provider) => (
<ProviderName key={provider} provider={provider} />
))}
</span>
),
},
{ label: "Tags", wide: true, value: tags.length > 0 && <Badges items={tags} /> },
],
];
};
function CallThisModel({ model }: { model: ModelGroupInfo }) {
const curl = curlSnippet(model);
if (!curl) {
return (
<PaneSection title="Call this model">
<CodeBlock code={modelSnippet(model)} />
</PaneSection>
);
}
return (
<PaneSection title="Call this model">
<Tabs defaultValue="python" className="gap-3">
<TabsList>
<TabsTrigger value="python">Python</TabsTrigger>
<TabsTrigger value="curl">cURL</TabsTrigger>
</TabsList>
<TabsContent value="python">
<CodeBlock code={modelSnippet(model)} />
</TabsContent>
<TabsContent value="curl">
<CodeBlock code={curl} />
</TabsContent>
</Tabs>
</PaneSection>
);
}
function CustomMetadata({ model }: { model: ModelGroupInfo }) {
const metadata = Object.entries(model.custom_metadata ?? {});
if (metadata.length === 0) return null;
return (
<PaneSection title="Custom metadata">
<dl className="grid grid-cols-[160px_minmax(0,1fr)] gap-x-4 gap-y-2 text-sm">
{metadata.map(([key, value]) => (
<React.Fragment key={key}>
<dt className="text-xs font-medium text-muted-foreground">{key}</dt>
<dd>{String(value)}</dd>
</React.Fragment>
))}
</dl>
</PaneSection>
);
}
function SupportedParams({ params }: { params: string[] }) {
if (params.length === 0) return null;
return (
<Collapsible>
<PaneSection
title="Supported OpenAI parameters"
aside={
<CollapsibleTrigger render={<Button variant="ghost" size="sm" />}>
Show {params.length}
<ChevronDown className="text-muted-foreground" />
</CollapsibleTrigger>
}
>
<CollapsibleContent>
<Badges items={params} mono />
</CollapsibleContent>
</PaneSection>
</Collapsible>
);
}
export function ModelInspector({ model }: { model: ModelGroupInfo }) {
const providers = model.providers ?? [];
const capabilities = modelCapabilities(model);
return (
<div className="flex flex-col gap-6 p-6">
<InspectorHeader
media={<ProviderAvatar provider={providers[0] ?? ""} className="size-6" />}
title={model.model_group}
copyLabel="Copy model name"
status={<HealthBadge model={model} />}
description={summary(model)}
/>
{model.model_group.includes("*") && (
<Alert>
<AlertTitle>Wildcard model</AlertTitle>
<AlertDescription>
Call it as <code className="font-mono">{wildcardExample(model.model_group)}</code>, with any model name this
provider offers
</AlertDescription>
</Alert>
)}
<Facts groups={modelFacts(model)} />
{capabilities.length > 0 && (
<PaneSection title="Capabilities">
<Badges items={capabilities} />
</PaneSection>
)}
<CallThisModel model={model} />
<CustomMetadata model={model} />
<SupportedParams params={model.supported_openai_params ?? []} />
</div>
);
}

238
ui/litellm-dashboard/src/components/publicModelHub/ModelsPane.tsx Normal file
View file

@ -0,0 +1,238 @@
"use client";
import { ArrowDownWideNarrow, ArrowUpNarrowWide, Boxes, ChevronLeft, ChevronRight } from "lucide-react";
import type { ModelGroupInfo } from "@/components/PublicModelHubTableColumns";
import { Button } from "@/components/ui/button";
import { ButtonGroup } from "@/components/ui/button-group";
import { ItemActions, ItemContent, ItemDescription, ItemMedia } from "@/components/ui/item";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Toggle } from "@/components/ui/toggle";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { displayMode } from "@/lib/decisionModels";
import { FilterChip } from "./FilterChip";
import { featureFacetLabel, firstPlus, formatPerMillion, formatTokens, modeLabel, providerName } from "./hubFormat";
import { HubMasterDetail } from "./HubMasterDetail";
import { LoadFailed, NoResults } from "./HubEmpty";
import { Joined, ProviderAvatar, RowTitle, SearchInput } from "./HubParts";
import { type HubData, useOpenedModel } from "./useHubData";
import { ModelInspector } from "./ModelInspector";
const SORTS = [
{ value: "model_group", label: "Name", desc: false },
{ value: "providers", label: "Provider", desc: false },
{ value: "mode", label: "Mode", desc: false },
{ value: "input_cost_per_token", label: "Input price", desc: false },
{ value: "output_cost_per_token", label: "Output price", desc: false },
{ value: "max_input_tokens", label: "Context", desc: true },
{ value: "max_output_tokens", label: "Max output", desc: true },
{ value: "rpm", label: "Rate limit", desc: true },
];
type Models = HubData["models"];
function ModelsPager({ models }: { models: Models }) {
const { pagination, rowCount, rows, onPaginationChange } = models;
if (!rowCount) return null;
const start = pagination.pageIndex * pagination.pageSize + 1;
const end = start + rows.length - 1;
const go = (pageIndex: number) => onPaginationChange({ ...pagination, pageIndex });
return (
<div className="flex items-center justify-between gap-2 px-6 pb-3 text-xs text-muted-foreground tabular-nums">
<span>
{start}–{end} of {rowCount}
</span>
<ButtonGroup>
<Button
variant="outline"
size="icon-sm"
aria-label="Previous page"
disabled={pagination.pageIndex === 0}
onClick={() => go(pagination.pageIndex - 1)}
>
<ChevronLeft />
</Button>
<Button
variant="outline"
size="icon-sm"
aria-label="Next page"
disabled={start + pagination.pageSize > rowCount}
onClick={() => go(pagination.pageIndex + 1)}
>
<ChevronRight />
</Button>
</ButtonGroup>
</div>
);
}
function SortControl({ models }: { models: Models }) {
const field = models.sorting[0]?.id ?? "model_group";
const desc = models.sorting[0]?.desc ?? false;
const pick = (value: string | null) => {
const sort = SORTS.find((option) => option.value === value) ?? SORTS[0];
models.onSortingChange([{ id: sort.value, desc: sort.desc }]);
};
return (
<ButtonGroup>
<Select items={SORTS} value={field} onValueChange={pick}>
<SelectTrigger size="sm" className="w-auto gap-1.5" aria-label="Sort by">
<span className="hidden text-muted-foreground sm:inline">Sort</span>
<SelectValue />
</SelectTrigger>
<SelectContent align="end" className="w-auto min-w-(--anchor-width)">
{SORTS.map((sort) => (
<SelectItem key={sort.value} value={sort.value}>
{sort.label}
</SelectItem>
))}
</SelectContent>
</Select>
<Tooltip>
<TooltipTrigger
render={
<Toggle
variant="outline"
size="sm"
aria-label="Sort descending"
pressed={desc}
onPressedChange={(pressed) => models.onSortingChange([{ id: field, desc: pressed }])}
className="dark:bg-input/30 dark:hover:bg-input/50"
/>
}
>
{desc ? <ArrowDownWideNarrow /> : <ArrowUpNarrowWide />}
</TooltipTrigger>
<TooltipContent>Sort descending</TooltipContent>
</Tooltip>
</ButtonGroup>
);
}
function ModelsToolbar({ hub }: { hub: HubData }) {
const { models, facets } = hub;
return (
<>
<div className="flex items-center gap-2">
<SearchInput value={models.searchValue} onChange={models.onSearchChange} placeholder="Search models" />
<SortControl models={models} />
</div>
<div className="flex flex-wrap items-center gap-2">
<FilterChip
label="Provider"
noun="providers"
options={facets.providers.map((p) => ({ value: p, label: providerName(p), provider: p }))}
value={models.providerValues}
onChange={models.onProvidersChange}
/>
<FilterChip
label="Mode"
noun="modes"
options={[...new Set(facets.modes.map((m) => displayMode(m) ?? m))].map((m) => ({
value: m,
label: modeLabel(m),
}))}
value={models.modeValues}
onChange={models.onModesChange}
/>
<FilterChip
label="Features"
noun="features"
options={facets.features.map((f) => ({ value: f, label: featureFacetLabel(f) }))}
value={models.featureValues}
onChange={models.onFeaturesChange}
/>
</div>
</>
);
}
const price = (cost: number | null | undefined, direction: "input" | "output") => {
if (cost === 0) return `Free ${direction}`;
return cost ? `${formatPerMillion(cost)} / 1M ${direction === "input" ? "in" : "out"}` : null;
};
const prices = (model: ModelGroupInfo) =>
model.input_cost_per_token === 0 && model.output_cost_per_token === 0
? ["Free", null]
: [price(model.input_cost_per_token, "input"), price(model.output_cost_per_token, "output")];
function ModelRow({ model }: { model: ModelGroupInfo }) {
const [input, output] = prices(model);
const priced = Boolean(input || output);
return (
<>
<ItemMedia variant="icon" className="size-5">
<ProviderAvatar provider={model.providers?.[0] ?? ""} className="size-5" />
</ItemMedia>
<ItemContent className="min-w-0 gap-0.5">
<RowTitle>{model.model_group}</RowTitle>
<ItemDescription>
<Joined
parts={[
firstPlus((model.providers ?? []).map(providerName)),
model.mode ? modeLabel(model.mode) : null,
model.max_input_tokens ? `${formatTokens(model.max_input_tokens)} context` : null,
]}
/>
</ItemDescription>
{priced && (
<ItemDescription className="tabular-nums sm:hidden">
<Joined parts={[input, output]} />
</ItemDescription>
)}
</ItemContent>
{priced && (
<ItemActions className="hidden flex-col items-end gap-0.5 text-xs whitespace-nowrap tabular-nums sm:flex">
{input && <span>{input}</span>}
{output && <span className="text-muted-foreground">{output}</span>}
</ItemActions>
)}
</>
);
}
export function ModelsPane({
hub,
selected,
onSelect,
}: {
hub: HubData;
selected: string | null;
onSelect: (key: string | null) => void;
}) {
const { models } = hub;
const opened = useOpenedModel(selected, models.rows, !models.isLoading);
const clear = () => {
models.onSearchChange("");
models.onColumnFiltersChange([]);
};
return (
<HubMasterDetail
rows={models.rows}
loading={models.isLoading}
selectedKey={selected}
opened={opened.model}
openPending={opened.pending}
getKey={(model) => model.model_group}
onSelect={onSelect}
empty={
models.error ? (
<LoadFailed noun="models" onRetry={models.refetch} />
) : (
<NoResults
icon={<Boxes />}
noun="models"
query={models.searchValue}
onClear={models.hasActiveQuery ? clear : undefined}
/>
)
}
toolbar={<ModelsToolbar hub={hub} />}
renderItem={(model) => <ModelRow model={model} />}
renderDetail={(model) => <ModelInspector model={model} />}
footer={<ModelsPager models={models} />}
/>
);
}

342
ui/litellm-dashboard/src/components/publicModelHub/PublicAiHub.integration.test.tsx Normal file
View file

@ -0,0 +1,342 @@
import { fireEvent, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils";
import PublicAiHub from "./PublicAiHub";
const { apiGetMock, realGet, navigation, uiSettings, replaceMock, hubInfo } = vi.hoisted(() => ({
apiGetMock: vi.fn(),
realGet: { current: (..._args: unknown[]): Promise<unknown> => Promise.reject(new Error("unset")) },
navigation: { params: new URLSearchParams() },
uiSettings: { values: {} as Record<string, unknown> },
replaceMock: vi.fn(),
hubInfo: {
docs_title: "",
custom_docs_description: null,
litellm_version: "1.90.0",
useful_links: { "Getting started": { url: "https://docs.example.com", index: 0 } },
},
}));
vi.mock("next/navigation", () => ({
useSearchParams: () => navigation.params,
usePathname: () => "/model_hub_table",
useRouter: () => ({ push: vi.fn(), replace: vi.fn(), refresh: vi.fn() }),
}));
vi.mock("@/components/navbar", () => ({ default: () => <nav aria-label="Top bar" /> }));
vi.mock("@/lib/telemetry/uiEvents", () => ({ recordUiEvent: vi.fn() }));
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
useUISettings: () => ({ data: uiSettings, isLoading: false }),
}));
vi.mock("@/components/networking", async (importOriginal) => {
const actual = await importOriginal<typeof import("@/components/networking")>();
realGet.current = actual.apiClient.get;
return {
...actual,
apiClient: { ...actual.apiClient, get: apiGetMock },
getUiConfig: vi.fn().mockResolvedValue({}),
getProxyBaseUrl: vi.fn(() => "https://gateway.example.com"),
getPublicModelHubInfo: vi.fn().mockResolvedValue(hubInfo),
};
});
const MODEL_HUB_PATH = "/public/v1/model_hub";
const MODELS = [
{
model_group: "claude-sonnet-4-5",
providers: ["anthropic"],
mode: "chat",
input_cost_per_token: 0.000003,
description: "Best for agentic coding",
},
{
model_group: "gpt-5",
providers: ["openai"],
mode: "chat",
input_cost_per_token: 0,
output_cost_per_token: 0,
description: "General purpose",
},
];
const AGENTS = [
{
name: "research-assistant",
description: "Answers research questions",
version: "1.4.0",
url: "https://gateway.example.com/a2a/research",
protocolVersion: "0.3.0",
defaultInputModes: ["text"],
defaultOutputModes: ["text"],
skills: [],
},
];
const HUBS: Record<string, unknown> = {
"/public/agent_hub": AGENTS,
"/public/mcp_hub": [],
"/public/skill_hub": { plugins: [] },
};
const page = (data: unknown[]) => ({
data,
meta: { total_count: data.length, page: 1, page_size: 50, total_pages: 1 },
links: { self: MODEL_HUB_PATH, first: MODEL_HUB_PATH, prev: null, next: null, last: MODEL_HUB_PATH },
});
const OFF_PAGE_MODEL = { model_group: "o1-pro", providers: ["openai"], mode: "chat", description: "Deep reasoning" };
type Query = { query?: Record<string, unknown> };
const respond = (path: string, options?: Query): Promise<unknown> => {
if (path in HUBS) return Promise.resolve(HUBS[path]);
if (path !== MODEL_HUB_PATH) return Promise.resolve({ ...page([]), meta: { page: 1, page_size: 100 } });
const q = String(options?.query?.q ?? "");
if (q === "o1-pro") return Promise.resolve(page([OFF_PAGE_MODEL]));
return Promise.resolve(page(MODELS.filter((model) => model.model_group.includes(q))));
};
const respondWithModels = () => apiGetMock.mockImplementation(respond);
const failModelsWith = (next: (options?: Query) => Promise<unknown>) =>
apiGetMock.mockImplementation((path: string, options?: Query) =>
path === MODEL_HUB_PATH ? next(options) : respond(path, options),
);
const narrowScreen = (narrow: boolean) =>
Object.defineProperty(window, "matchMedia", {
writable: true,
value: (query: string) => ({
matches: narrow,
media: query,
addEventListener: () => {},
removeEventListener: () => {},
}),
});
const tokenExpiringIn = (seconds: number) => {
const encode = (value: object) => btoa(JSON.stringify(value)).replace(/=+$/, "");
return `${encode({ alg: "HS256" })}.${encode({ exp: Math.floor(Date.now() / 1000) + seconds })}.signature`;
};
beforeEach(() => {
narrowScreen(false);
vi.clearAllMocks();
testQueryClient.clear();
uiSettings.values = {};
navigation.params = new URLSearchParams();
document.cookie = "token=; expires=Thu, 01 Jan 1970 00:00:00 GMT; path=/";
respondWithModels();
vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response("{}", { status: 404 }));
Object.defineProperty(window, "location", {
writable: true,
value: { ...window.location, replace: replaceMock },
});
});
describe("PublicAiHub", () => {
it("counts every hub and shows the first model's details", async () => {
renderWithProviders(<PublicAiHub accessToken={null} />);
expect(await screen.findByRole("heading", { level: 2, name: "claude-sonnet-4-5" })).toBeInTheDocument();
expect(screen.getByRole("tab", { name: /^Models\s*2$/ })).toBeInTheDocument();
expect(screen.getByRole("tab", { name: /^Agents\s*1$/ })).toBeInTheDocument();
expect(screen.getByText("Best for agentic coding")).toBeInTheDocument();
expect(screen.getByRole("textbox", { name: "Base URL" })).toHaveValue("https://gateway.example.com");
});
it("opens a model from the list and records it in the URL", async () => {
const onUrlUpdate = vi.fn();
const user = userEvent.setup();
renderWithProviders(<PublicAiHub accessToken={null} />, { onUrlUpdate });
await user.click(await screen.findByRole("link", { name: /gpt-5/ }));
expect(await screen.findByRole("heading", { level: 2, name: "gpt-5" })).toBeInTheDocument();
await waitFor(() =>
expect(onUrlUpdate).toHaveBeenCalledWith(expect.objectContaining({ queryString: "?open=gpt-5" })),
);
});
it("explains an empty search and clears it", async () => {
const user = userEvent.setup();
renderWithProviders(<PublicAiHub accessToken={null} />, { searchParams: { tab: "agents" } });
fireEvent.change(await screen.findByRole("textbox", { name: "Search agents" }), { target: { value: "billing" } });
expect(await screen.findByText('No agents match "billing"')).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Clear filters" }));
expect(await screen.findByRole("link", { name: /research-assistant/ })).toBeInTheDocument();
});
it("lists the admin's links and the gateway version", async () => {
const user = userEvent.setup();
renderWithProviders(<PublicAiHub accessToken={null} />);
await user.click(await screen.findByRole("button", { name: "Links" }));
expect(await screen.findByRole("menuitem", { name: "Getting started" })).toHaveAttribute(
"href",
"https://docs.example.com",
);
expect(screen.getByText("LiteLLM v1.90.0")).toBeInTheDocument();
});
it("keeps agents reachable when only the models call fails", async () => {
const user = userEvent.setup();
failModelsWith(() => Promise.reject(new Error("No models configured")));
renderWithProviders(<PublicAiHub accessToken={null} />);
expect(await screen.findByText("Couldn't load models")).toBeInTheDocument();
expect(screen.getByRole("textbox", { name: "Search models" })).toBeInTheDocument();
await user.click(screen.getByRole("tab", { name: /^Agents\s*1$/ }));
expect(await screen.findByRole("link", { name: /research-assistant/ })).toBeInTheDocument();
});
it("says the hub is unavailable when every hub call fails", async () => {
apiGetMock.mockRejectedValue(new Error("Not found"));
renderWithProviders(<PublicAiHub accessToken={null} />);
expect(await screen.findByText("This AI Hub isn't available")).toBeInTheDocument();
expect(screen.queryByRole("tablist")).not.toBeInTheDocument();
});
it("shows a load error when a hub answers with an HTTP error status", async () => {
const user = userEvent.setup();
vi.mocked(fetch).mockResolvedValue(new Response(JSON.stringify({ detail: "Database error" }), { status: 500 }));
apiGetMock.mockImplementation((path: string, options?: Query) =>
path === "/public/agent_hub" ? realGet.current(path, options) : respond(path, options),
);
renderWithProviders(<PublicAiHub accessToken={null} />);
await user.click(await screen.findByRole("tab", { name: /^Agents/ }));
expect(await screen.findByText("Couldn't load agents")).toBeInTheDocument();
expect(screen.queryByText("No agents published yet")).not.toBeInTheDocument();
});
it("shows loading rather than an empty hub while Try again reloads models", async () => {
const user = userEvent.setup();
let finishRetry: (value: unknown) => void = () => {};
failModelsWith(() => Promise.reject(new Error("Gateway timeout")));
renderWithProviders(<PublicAiHub accessToken={null} />);
await screen.findByText("Couldn't load models");
failModelsWith(() => new Promise((resolve) => (finishRetry = resolve)));
await user.click(screen.getByRole("button", { name: "Try again" }));
expect(screen.queryByText("No models published yet")).not.toBeInTheDocument();
expect(screen.queryByText("Couldn't load models")).not.toBeInTheDocument();
finishRetry(page(MODELS));
expect(await screen.findByRole("heading", { level: 2, name: "claude-sonnet-4-5" })).toBeInTheDocument();
});
it("keeps the open model and its snippet tab while a search hides its row", async () => {
const user = userEvent.setup();
failModelsWith((options) =>
options?.query?.q === "gpt-5" ? new Promise(() => {}) : respond(MODEL_HUB_PATH, options),
);
renderWithProviders(<PublicAiHub accessToken={null} />, { searchParams: { open: "gpt-5" } });
await screen.findByRole("heading", { level: 2, name: "gpt-5" });
await user.click(screen.getByRole("tab", { name: "cURL" }));
fireEvent.change(screen.getByRole("textbox", { name: "Search models" }), { target: { value: "claude" } });
await waitFor(() => expect(screen.queryByRole("link", { name: /gpt-5/ })).not.toBeInTheDocument());
expect(screen.getByRole("heading", { level: 2, name: "gpt-5" })).toBeInTheDocument();
expect(screen.getByRole("tab", { name: "cURL" })).toHaveAttribute("aria-selected", "true");
});
it("keeps the open item's details when a search matches nothing", async () => {
renderWithProviders(<PublicAiHub accessToken={null} />, {
searchParams: { tab: "agents", open: "research-assistant" },
});
fireEvent.change(await screen.findByRole("textbox", { name: "Search agents" }), { target: { value: "billing" } });
expect(await screen.findByText('No agents match "billing"')).toBeInTheDocument();
expect(screen.getByRole("heading", { level: 2, name: "research-assistant" })).toBeInTheDocument();
});
it("marks a free model as free in its row", async () => {
renderWithProviders(<PublicAiHub accessToken={null} />);
const row = await screen.findByRole("link", { name: /gpt-5/ });
expect(within(row).getAllByText("Free").length).toBeGreaterThan(0);
});
it("treats an empty open param as nothing opened", async () => {
renderWithProviders(<PublicAiHub accessToken={null} />, { searchParams: { open: "" } });
expect(await screen.findByRole("heading", { level: 2, name: "claude-sonnet-4-5" })).toBeInTheDocument();
});
it("tells phone visitors when a shared link names something no longer listed", async () => {
narrowScreen(true);
renderWithProviders(<PublicAiHub accessToken={null} />, { searchParams: { open: "retired-model" } });
const sheet = await screen.findByRole("dialog");
expect(within(sheet).getByText("\u201cretired-model\u201d isn't listed here")).toBeInTheDocument();
});
it("keeps the page's other query params in each row's link", async () => {
navigation.params = new URLSearchParams("tab=agents&key=abc");
renderWithProviders(<PublicAiHub accessToken={null} />, { searchParams: { tab: "agents", key: "abc" } });
expect(await screen.findByRole("link", { name: /research-assistant/ })).toHaveAttribute(
"href",
"?tab=agents&key=abc&open=research-assistant",
);
});
it("clears the open item when switching tabs", async () => {
const onUrlUpdate = vi.fn();
const user = userEvent.setup();
renderWithProviders(<PublicAiHub accessToken={null} />, { searchParams: { open: "gpt-5" }, onUrlUpdate });
expect(await screen.findByRole("heading", { level: 2, name: "gpt-5" })).toBeInTheDocument();
await user.click(screen.getByRole("tab", { name: /^Agents/ }));
await waitFor(() =>
expect(onUrlUpdate).toHaveBeenLastCalledWith(expect.objectContaining({ queryString: "?tab=agents" })),
);
});
it("opens a shared link to a model that is not on the first page", async () => {
renderWithProviders(<PublicAiHub accessToken={null} />, { searchParams: { open: "o1-pro" } });
expect(await screen.findByRole("heading", { level: 2, name: "o1-pro" })).toBeInTheDocument();
expect(screen.getByText("Deep reasoning")).toBeInTheDocument();
});
it("says when a shared link names something that is no longer listed", async () => {
renderWithProviders(<PublicAiHub accessToken={null} />, { searchParams: { open: "retired-model" } });
expect(await screen.findByText("\u201cretired-model\u201d isn't listed here")).toBeInTheDocument();
expect(screen.queryByRole("heading", { level: 2, name: "claude-sonnet-4-5" })).not.toBeInTheDocument();
});
it("sends visitors without a session to login when the admin requires it", async () => {
uiSettings.values = { require_auth_for_public_ai_hub: true };
renderWithProviders(<PublicAiHub accessToken={null} />);
await waitFor(() => expect(replaceMock).toHaveBeenCalledWith(expect.stringContaining("login")));
});
it("keeps a signed-in visitor on the hub when the admin requires sign-in", async () => {
uiSettings.values = { require_auth_for_public_ai_hub: true };
document.cookie = `token=${tokenExpiringIn(3600)}; path=/`;
renderWithProviders(<PublicAiHub accessToken={null} />);
expect(await screen.findByRole("heading", { level: 2, name: "claude-sonnet-4-5" })).toBeInTheDocument();
expect(replaceMock).not.toHaveBeenCalled();
});
});

113
ui/litellm-dashboard/src/components/publicModelHub/PublicAiHub.tsx Normal file
View file

@ -0,0 +1,113 @@
"use client";
import type { ReactNode } from "react";
import Navbar from "@/components/navbar";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { TooltipProvider } from "@/components/ui/tooltip";
import { ThemeProvider } from "@/contexts/ThemeContext";
import { HubHeader, type HubStatus } from "./HubHeader";
import { HubNotEnabled } from "./HubEmpty";
import { ModelsPane } from "./ModelsPane";
import { AgentsPane, McpPane, SkillsPane } from "./ResourcePanes";
import { type HubData, useHubData } from "./useHubData";
import { type HubTab, useHubUrlState } from "./useHubUrlState";
import { useRequirePublicHubAuth } from "./useRequirePublicHubAuth";
function TabLabel({ children, count }: { children: ReactNode; count: number | undefined }) {
return (
<>
{children}
<span className="text-xs text-muted-foreground tabular-nums">{count ?? "-"}</span>
</>
);
}
const hubStatus = (hub: HubData): HubStatus => {
const failures = [hub.models.error, hub.agents.error, hub.mcp.error, hub.skills.error].filter(Boolean);
if (failures.length === 4) return "unavailable";
if (hub.models.isLoading) return "loading";
return failures.length > 0 ? "degraded" : "operational";
};
function HubTabs({ hub }: { hub: HubData }) {
const { tab, setTab, open, setOpen } = useHubUrlState();
const paneProps = { selected: open, onSelect: setOpen };
return (
<Tabs value={tab} onValueChange={(next) => setTab(next as HubTab)} className="gap-4">
<TabsList className="w-full sm:w-fit">
<TabsTrigger value="models" className="px-1.5 sm:px-2">
<TabLabel count={hub.modelTotal}>Models</TabLabel>
</TabsTrigger>
<TabsTrigger value="agents" className="px-1.5 sm:px-2">
<TabLabel count={hub.agents.data?.length}>Agents</TabLabel>
</TabsTrigger>
<TabsTrigger value="mcp" className="px-1.5 sm:px-2">
<TabLabel count={hub.mcp.data?.length}>
<span className="sm:hidden">MCP</span>
<span className="hidden sm:inline">MCP servers</span>
</TabLabel>
</TabsTrigger>
<TabsTrigger value="skills" className="px-1.5 sm:px-2">
<TabLabel count={hub.skills.data?.length}>Skills</TabLabel>
</TabsTrigger>
</TabsList>
<TabsContent value="models">
<ModelsPane hub={hub} {...paneProps} />
</TabsContent>
<TabsContent value="agents">
<AgentsPane
rows={hub.agents.data}
loading={hub.agents.isPending}
error={hub.agents.isError}
onRetry={() => void hub.agents.refetch()}
{...paneProps}
/>
</TabsContent>
<TabsContent value="mcp">
<McpPane
rows={hub.mcp.data}
loading={hub.mcp.isPending}
error={hub.mcp.isError}
onRetry={() => void hub.mcp.refetch()}
{...paneProps}
/>
</TabsContent>
<TabsContent value="skills">
<SkillsPane
rows={hub.skills.data}
loading={hub.skills.isPending}
error={hub.skills.isError}
onRetry={() => void hub.skills.refetch()}
{...paneProps}
/>
</TabsContent>
</Tabs>
);
}
function PublicAiHubContent() {
useRequirePublicHubAuth();
const hub = useHubData();
const status = hubStatus(hub);
return (
<div className="flex w-full flex-col gap-6 p-4 sm:p-6">
<HubHeader info={hub.info.data} status={status} />
{status === "unavailable" ? <HubNotEnabled /> : <HubTabs hub={hub} />}
</div>
);
}
export default function PublicAiHub({ accessToken }: { accessToken: string | null }) {
return (
<ThemeProvider accessToken={accessToken}>
<TooltipProvider>
<div className="min-h-screen bg-background">
<Navbar accessToken={accessToken} isPublicPage />
<PublicAiHubContent />
</div>
</TooltipProvider>
</ThemeProvider>
);
}

184
ui/litellm-dashboard/src/components/publicModelHub/ResourcePanes.tsx Normal file
View file

@ -0,0 +1,184 @@
"use client";
import { Bot, Puzzle, Server } from "lucide-react";
import * as React from "react";
import type { Plugin } from "@/components/claude_code_plugins/types";
import type { AgentCard, MCPServerData } from "@/components/PublicModelHubTableColumns";
import { ItemActions, ItemContent, ItemDescription, ItemMedia } from "@/components/ui/item";
import { matchesSearchTerm } from "@/utils/searchUtils";
import { AgentInspector } from "./AgentInspector";
import { FilterChip } from "./FilterChip";
import { mcpAuthLabel } from "./hubFormat";
import { HubMasterDetail } from "./HubMasterDetail";
import { LoadFailed, NoResults } from "./HubEmpty";
import { RowTitle, SearchInput } from "./HubParts";
import { McpInspector } from "./McpInspector";
import { SkillInspector } from "./SkillInspector";
interface ResourcePaneProps<T> {
rows: T[];
loading: boolean;
error: boolean;
onRetry: () => void;
noun: string;
icon: React.ReactNode;
search: (row: T) => (string | null | undefined)[];
filter: { label: string; noun: string; values: (row: T) => string[] };
selected: string | null;
onSelect: (key: string | null) => void;
getKey: (row: T) => string;
renderDetail: (row: T) => React.ReactNode;
title: (row: T) => string;
description: (row: T) => string | null | undefined;
meta: (row: T) => string | null | undefined;
}
function ResourcePane<T>({
rows,
loading,
error,
onRetry,
noun,
icon,
search,
filter,
selected,
onSelect,
getKey,
renderDetail,
title,
description,
meta,
}: ResourcePaneProps<T>) {
const [query, setQuery] = React.useState("");
const [chosen, setChosen] = React.useState<string[]>([]);
const options = [...new Set(rows.flatMap(filter.values))].sort();
const sorted = [...rows].sort((a, b) => title(a).localeCompare(title(b)));
const shown = sorted.filter(
(row) =>
matchesSearchTerm(query, search(row)) &&
(chosen.length === 0 || filter.values(row).some((v) => chosen.includes(v))),
);
const filtered = query.trim() !== "" || chosen.length > 0;
const clear = () => {
setQuery("");
setChosen([]);
};
return (
<HubMasterDetail
rows={shown}
loading={loading}
selectedKey={selected}
opened={rows.find((row) => getKey(row) === selected)}
getKey={getKey}
onSelect={onSelect}
empty={
error ? (
<LoadFailed noun={noun} onRetry={onRetry} />
) : (
<NoResults icon={icon} noun={noun} query={query} onClear={filtered ? clear : undefined} />
)
}
toolbar={
<div className="flex items-center gap-2">
<SearchInput value={query} onChange={setQuery} placeholder={`Search ${noun}`} />
{options.length > 0 && (
<FilterChip
label={filter.label}
noun={filter.noun}
options={options.map((option) => ({ value: option, label: option }))}
value={chosen}
onChange={setChosen}
/>
)}
</div>
}
renderItem={(row) => (
<>
<ItemMedia variant="icon" className="size-5 text-muted-foreground">
{icon}
</ItemMedia>
<ItemContent className="min-w-0 gap-0.5">
<RowTitle>{title(row)}</RowTitle>
{description(row) && <ItemDescription>{description(row)}</ItemDescription>}
</ItemContent>
{meta(row) && <ItemActions className="text-xs text-muted-foreground">{meta(row)}</ItemActions>}
</>
)}
renderDetail={renderDetail}
/>
);
}
interface PaneProps<T> {
rows: T[] | undefined;
loading: boolean;
error: boolean;
onRetry: () => void;
selected: string | null;
onSelect: (key: string | null) => void;
}
export function AgentsPane({ rows, ...props }: PaneProps<AgentCard>) {
return (
<ResourcePane<AgentCard>
{...props}
rows={rows ?? []}
noun="agents"
icon={<Bot />}
search={(agent) => [agent.name, agent.description]}
filter={{
label: "Skill tag",
noun: "skill tags",
values: (agent) => (agent.skills ?? []).flatMap((s) => s.tags ?? []),
}}
getKey={(agent) => agent.name}
title={(agent) => agent.name}
description={(agent) => agent.description}
meta={(agent) => (agent.version ? `v${agent.version}` : undefined)}
renderDetail={(agent) => <AgentInspector agent={agent} />}
/>
);
}
export function McpPane({ rows, ...props }: PaneProps<MCPServerData>) {
return (
<ResourcePane<MCPServerData>
{...props}
rows={rows ?? []}
noun="MCP servers"
icon={<Server />}
search={(server) => [server.name, server.mcp_info?.description]}
filter={{
label: "Transport",
noun: "transports",
values: (server) => (server.transport ? [server.transport] : []),
}}
getKey={(server) => server.name}
title={(server) => server.name}
description={(server) => server.mcp_info?.description}
meta={(server) => (server.auth_type ? mcpAuthLabel(server.auth_type) : undefined)}
renderDetail={(server) => <McpInspector server={server} />}
/>
);
}
export function SkillsPane({ rows, ...props }: PaneProps<Plugin>) {
return (
<ResourcePane<Plugin>
{...props}
rows={rows ?? []}
noun="skills"
icon={<Puzzle />}
search={(skill) => [skill.name, skill.description, skill.namespace, skill.domain, ...(skill.keywords ?? [])]}
filter={{ label: "Domain", noun: "domains", values: (skill) => (skill.domain ? [skill.domain] : []) }}
getKey={(skill) => skill.name}
title={(skill) => skill.name}
description={(skill) => skill.description}
meta={(skill) => skill.domain}
renderDetail={(skill) => <SkillInspector skill={skill} />}
/>
);
}

66
ui/litellm-dashboard/src/components/publicModelHub/SkillInspector.tsx Normal file
View file

@ -0,0 +1,66 @@
"use client";
import { ExternalLink } from "lucide-react";
import { getSourceDisplayText, getSourceLink } from "@/components/claude_code_plugins/helpers";
import type { Plugin } from "@/components/claude_code_plugins/types";
import { Button } from "@/components/ui/button";
import { Field, FieldDescription } from "@/components/ui/field";
import { Badges, CodeBlock, Facts, PaneSection } from "./HubParts";
import { marketplaceCommand, skillInstallCommand } from "./hubSnippets";
import { InspectorHeader } from "./InspectorHeader";
function sourceValue(source: Plugin["source"]) {
const link = getSourceLink(source);
const href = link && source.source === "git-subdir" && source.path ? `${link}/tree/main/${source.path}` : link;
if (!href) return getSourceDisplayText(source) || null;
return (
<Button
variant="link"
nativeButton={false}
className="h-auto max-w-full justify-start px-0"
render={<a href={href} target="_blank" rel="noopener noreferrer" />}
>
<span className="truncate">{href.replace(/^https:\/\//, "")}</span>
<ExternalLink />
</Button>
);
}
export function SkillInspector({ skill }: { skill: Plugin }) {
const keywords = skill.keywords ?? [];
const addedAt = skill.created_at ? new Date(skill.created_at) : null;
const added = addedAt && !Number.isNaN(addedAt.getTime()) ? addedAt.toLocaleDateString() : null;
return (
<div className="flex flex-col gap-6 p-6">
<InspectorHeader title={skill.name} copyLabel="Copy skill name" description={skill.description} />
<Facts
groups={[
[
{ label: "Domain", value: skill.domain },
{ label: "Namespace", value: skill.namespace },
{ label: "Version", value: skill.version },
{ label: "Category", value: skill.category },
],
[
{ label: "Source", wide: true, value: sourceValue(skill.source) },
{ label: "Author", value: skill.author?.name },
{ label: "Added", value: added },
],
[{ label: "Keywords", wide: true, value: keywords.length > 0 && <Badges items={keywords} /> }],
]}
/>
<PaneSection title="Install in Claude Code">
<Field>
<FieldDescription>First time only: add this marketplace</FieldDescription>
<CodeBlock code={marketplaceCommand()} />
</Field>
<Field>
<FieldDescription>Then install the skill</FieldDescription>
<CodeBlock code={skillInstallCommand(skill)} />
</Field>
</PaneSection>
</div>
);
}

138
ui/litellm-dashboard/src/components/publicModelHub/hubFormat.test.ts Normal file
View file

@ -0,0 +1,138 @@
import { describe, expect, it, vi } from "vitest";
import type { AgentCard, ModelGroupInfo } from "@/components/PublicModelHubTableColumns";
import {
agentCapabilities,
capabilityLabel,
featureFacetLabel,
firstPlus,
formatPerMillion,
formatTokens,
healthOf,
mcpAuthLabel,
modeLabel,
modelCapabilities,
relativeTime,
usefulLinkList,
wildcardExample,
} from "./hubFormat";
vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "https://gateway.example.com" }));
const model = (overrides: Partial<ModelGroupInfo>): ModelGroupInfo => ({
model_group: "gpt-5",
providers: ["openai"],
supports_parallel_function_calling: false,
supports_vision: false,
supports_function_calling: false,
...overrides,
});
describe("modeLabel", () => {
it("names known modes in sentence case and falls back for unknown ones", () => {
expect(modeLabel("image_generation")).toBe("Image generation");
expect(modeLabel("audio_speech")).toBe("Text to speech");
expect(modeLabel("custom_thing")).toBe("Custom thing");
expect(modeLabel("evaluation")).toBe("Decisions");
expect(modeLabel(undefined)).toBe("-");
});
});
describe("formatTokens", () => {
it("abbreviates thousands and millions", () => {
expect(formatTokens(512)).toBe("512");
expect(formatTokens(200_000)).toBe("200K");
expect(formatTokens(1_000_000)).toBe("1M");
expect(formatTokens(1_048_576)).toBe("1M");
expect(formatTokens(2_500_000)).toBe("2.5M");
expect(formatTokens(undefined)).toBe("-");
});
});
describe("formatPerMillion", () => {
it("prices per million tokens with at least two decimals and keeps sub-cent prices", () => {
expect(formatPerMillion(0.000003)).toBe("$3.00");
expect(formatPerMillion(0.0000165)).toBe("$16.50");
expect(formatPerMillion(3.75e-8)).toBe("$0.0375");
expect(formatPerMillion(2e-9)).toBe("$0.002");
expect(formatPerMillion(0)).toBe("Free");
expect(formatPerMillion(undefined)).toBe("-");
});
});
describe("capabilities", () => {
it("labels known flags and derives a label for new ones", () => {
expect(capabilityLabel("supports_function_calling")).toBe("Function calling");
expect(capabilityLabel("supports_pdf_input")).toBe("Pdf input");
expect(featureFacetLabel("web_search")).toBe("Web search");
});
it("lists only enabled flags, known ones first in a fixed order", () => {
const flags = model({ supports_pdf_input: true, supports_web_search: true, supports_vision: true });
expect(modelCapabilities(flags)).toEqual(["Vision", "Web search", "Pdf input"]);
});
it("lists an agent's enabled capabilities", () => {
const agent = {
capabilities: { streaming: true, pushNotifications: true, stateTransitionHistory: true, x: false },
};
expect(agentCapabilities(agent as AgentCard)).toEqual([
"Streaming",
"Push notifications",
"State transition history",
]);
});
});
describe("healthOf", () => {
it("maps the health check to a tone and label", () => {
expect(healthOf(model({ health_status: "healthy" }))).toEqual({ tone: "success", label: "Healthy" });
expect(healthOf(model({ health_status: "unhealthy" }))).toEqual({ tone: "error", label: "Unhealthy" });
expect(healthOf(model({}))).toEqual({ tone: "neutral", label: "Not checked" });
});
});
describe("relativeTime", () => {
const now = Date.parse("2026-10-31T12:00:00Z");
it("rounds to minutes, hours or days", () => {
expect(relativeTime("2026-10-31T11:59:45Z", now)).toBe("just now");
expect(relativeTime("2026-10-31T11:45:00Z", now)).toBe("15 min ago");
expect(relativeTime("2026-10-31T09:00:00Z", now)).toBe("3 h ago");
expect(relativeTime("2026-10-30T12:00:00Z", now)).toBe("1 day ago");
expect(relativeTime("2026-10-29T12:00:00Z", now)).toBe("2 days ago");
expect(relativeTime(undefined, now)).toBe("-");
expect(relativeTime("not-a-date", now)).toBe("-");
});
});
describe("small formatters", () => {
it("shows a wildcard model as a placeholder name", () => {
expect(wildcardExample("anthropic/*")).toBe("anthropic/<model-name>");
});
it("collapses a list to its first item and a count", () => {
expect(firstPlus(["Anthropic"])).toBe("Anthropic");
expect(firstPlus(["Anthropic", "Amazon Bedrock", "Vertex AI"])).toBe("Anthropic +2");
expect(firstPlus([])).toBeUndefined();
});
it("names MCP sign-in types", () => {
expect(mcpAuthLabel("oauth2")).toBe("OAuth");
expect(mcpAuthLabel("custom_header")).toBe("custom_header");
expect(mcpAuthLabel(null)).toBe("-");
});
});
describe("usefulLinkList", () => {
it("orders the admin's links by index and accepts the legacy string form", () => {
const links = {
"Status page": { url: "https://status.example.com", index: 2 },
"Getting started": { url: "https://docs.example.com", index: 0 },
Legacy: "https://legacy.example.com",
};
expect(usefulLinkList(links).map((link) => link.title)).toEqual(["Getting started", "Legacy", "Status page"]);
expect(usefulLinkList(undefined)).toEqual([]);
});
});

142
ui/litellm-dashboard/src/components/publicModelHub/hubFormat.ts Normal file
View file

@ -0,0 +1,142 @@
import { getProxyBaseUrl } from "@/components/networking";
import type { AgentCard, ModelGroupInfo } from "@/components/PublicModelHubTableColumns";
import { getProviderLogoAndName } from "@/components/provider_info_helpers";
import { displayMode } from "@/lib/decisionModels";
export const SECTION_TITLE = "text-[15px] leading-[22px] font-semibold";
const MODE_LABELS: Record<string, string> = {
chat: "Chat",
completion: "Completion",
embedding: "Embedding",
rerank: "Rerank",
image_generation: "Image generation",
image_edit: "Image edit",
audio_speech: "Text to speech",
audio_transcription: "Transcription",
video_generation: "Video generation",
realtime: "Realtime",
responses: "Responses",
batch: "Batch",
ocr: "OCR",
decisions: "Decisions",
};
const sentenceCase = (text: string): string => text.replaceAll("_", " ").replace(/^./, (c) => c.toUpperCase());
export const modeLabel = (mode: string | null | undefined): string => {
const shown = displayMode(mode);
return shown ? MODE_LABELS[shown] ?? sentenceCase(shown) : "-";
};
export const formatTokens = (tokens: number | null | undefined): string => {
if (!tokens) return "-";
if (tokens >= 1_000_000) return `${+(tokens / 1_000_000).toFixed(1)}M`;
if (tokens >= 1_000) return `${Math.round(tokens / 1_000)}K`;
return String(tokens);
};
const PER_MILLION_FORMAT: Intl.NumberFormatOptions = {
style: "currency",
currency: "USD",
minimumFractionDigits: 2,
maximumFractionDigits: 4,
useGrouping: false,
};
const PER_MILLION = new Intl.NumberFormat("en-US", PER_MILLION_FORMAT);
export const formatPerMillion = (cost: number | null | undefined): string => {
if (cost === 0) return "Free";
return cost ? PER_MILLION.format(cost * 1_000_000) : "-";
};
export const providerName = (provider: string): string =>
getProviderLogoAndName(provider).displayName.replace(/\s*\(.*\)$/, "");
export const providerLogo = (provider: string): string => getProviderLogoAndName(provider).logo;
const CAPABILITY_LABELS: Record<string, string> = {
supports_vision: "Vision",
supports_function_calling: "Function calling",
supports_reasoning: "Reasoning",
supports_web_search: "Web search",
supports_parallel_function_calling: "Parallel function calling",
supports_url_context: "URL context",
supports_fast_mode: "Fast mode",
};
const CAPABILITY_ORDER = Object.keys(CAPABILITY_LABELS);
const capabilityRank = (key: string): number => {
const index = CAPABILITY_ORDER.indexOf(key);
return index === -1 ? CAPABILITY_ORDER.length : index;
};
export const capabilityLabel = (key: string): string =>
CAPABILITY_LABELS[key] ?? sentenceCase(key.replace(/^supports_/, ""));
export const featureFacetLabel = (feature: string): string => capabilityLabel(`supports_${feature}`);
export const modelCapabilities = (model: ModelGroupInfo): string[] =>
Object.entries(model)
.filter(([key, value]) => key.startsWith("supports_") && value === true)
.map(([key]) => key)
.sort((a, b) => capabilityRank(a) - capabilityRank(b))
.map(capabilityLabel);
export const agentCapabilities = (agent: AgentCard): string[] =>
Object.entries(agent.capabilities ?? {})
.filter(([, on]) => on === true)
.map(([key]) => sentenceCase(key.replace(/([a-z])([A-Z])/g, "$1_$2").toLowerCase()));
export type HealthTone = "success" | "error" | "neutral";
const HEALTH: Record<string, { tone: HealthTone; label: string }> = {
healthy: { tone: "success", label: "Healthy" },
unhealthy: { tone: "error", label: "Unhealthy" },
};
export const healthOf = (model: ModelGroupInfo): { tone: HealthTone; label: string } =>
HEALTH[model.health_status ?? ""] ?? { tone: "neutral", label: "Not checked" };
export const relativeTime = (iso: string | null | undefined, now: number = Date.now()): string => {
const then = iso ? new Date(iso).getTime() : NaN;
if (Number.isNaN(then)) return "-";
const minutes = Math.max(0, Math.round((now - then) / 60_000));
if (minutes < 1) return "just now";
if (minutes < 60) return `${minutes} min ago`;
const hours = Math.round(minutes / 60);
if (hours < 24) return `${hours} h ago`;
const days = Math.round(hours / 24);
return `${days} ${days === 1 ? "day" : "days"} ago`;
};
export const gatewayUrl = (): string => getProxyBaseUrl();
export const wildcardExample = (modelGroup: string): string => modelGroup.replaceAll("*", "<model-name>");
export const firstPlus = (items: string[]): string | undefined =>
items.length > 1 ? `${items[0]} +${items.length - 1}` : items[0];
const MCP_AUTH: Record<string, string> = {
none: "None",
oauth2: "OAuth",
bearer_token: "Bearer token",
api_key: "API key",
};
export const mcpAuthLabel = (authType: string | null | undefined): string =>
authType ? MCP_AUTH[authType] ?? authType : "-";
export const usefulLinkList = (
links: Record<string, string | { url: string; index?: number }> | undefined,
): { title: string; url: string }[] =>
Object.entries(links ?? {})
.map(([title, value]) => ({
title,
url: typeof value === "string" ? value : value.url,
index: typeof value === "string" ? 0 : value.index ?? 0,
}))
.sort((a, b) => a.index - b.index)
.map(({ title, url }) => ({ title, url }));

59
ui/litellm-dashboard/src/components/publicModelHub/hubSnippets.test.ts Normal file
View file

@ -0,0 +1,59 @@
import { afterEach, describe, expect, it, vi } from "vitest";
import type { ModelGroupInfo, MCPServerData } from "@/components/PublicModelHubTableColumns";
import { curlSnippet, marketplaceCommand, mcpEndpoint, mcpSnippet, modelSnippet } from "./hubSnippets";
vi.mock("@/components/networking", () => ({
getProxyBaseUrl: () => "https://gateway.example.com",
getGlobalLitellmHeaderName: () => "Authorization",
}));
const model = (overrides: Partial<ModelGroupInfo>): ModelGroupInfo => ({
model_group: "anthropic/*",
providers: ["anthropic"],
supports_parallel_function_calling: false,
supports_vision: false,
supports_function_calling: false,
...overrides,
});
const server = { server_id: "s1", name: "github", server_name: "github" } as MCPServerData;
afterEach(() => {
vi.unstubAllGlobals();
});
describe("model snippets", () => {
it("points the SDK snippet at the gateway and a concrete model name", () => {
vi.stubGlobal("window", { location: { origin: "http://ui.example.com" } });
const snippet = modelSnippet(model({ mode: "chat" }));
expect(snippet).toContain('base_url="https://gateway.example.com"');
expect(snippet).toContain("anthropic/<model-name>");
expect(snippet).not.toContain("ui.example.com");
});
it("offers cURL only for chat models", () => {
expect(curlSnippet(model({ mode: "chat" }))).toContain("curl https://gateway.example.com/v1/chat/completions");
expect(curlSnippet(model({ mode: undefined }))).not.toBeNull();
expect(curlSnippet(model({ mode: "embedding" }))).toBeNull();
});
});
describe("MCP and skill snippets", () => {
it("builds the MCP endpoint from the gateway and the server name", () => {
expect(mcpEndpoint(server)).toBe("https://gateway.example.com/github/mcp");
expect(mcpSnippet(server)).toContain('"url": "https://gateway.example.com/github/mcp"');
});
it("routes a server without a server name through its listed name", () => {
const aliased = { server_id: "s2", name: "my_alias", server_name: null } as unknown as MCPServerData;
expect(mcpEndpoint(aliased)).toBe("https://gateway.example.com/my_alias/mcp");
});
it("adds the gateway's marketplace", () => {
expect(marketplaceCommand()).toBe(
"/plugin marketplace add https://gateway.example.com/claude-code/marketplace.json",
);
});
});

93
ui/litellm-dashboard/src/components/publicModelHub/hubSnippets.ts Normal file
View file

@ -0,0 +1,93 @@
import { modelUsageExample } from "@/components/AIHub/modelUsageExample";
import type { Plugin } from "@/components/claude_code_plugins/types";
import { formatInstallCommand } from "@/components/claude_code_plugins/helpers";
import { generateCodeSnippet } from "@/components/chat_ui/CodeSnippets";
import { getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping";
import { getGlobalLitellmHeaderName } from "@/components/networking";
import type { AgentCard, MCPServerData, ModelGroupInfo } from "@/components/PublicModelHubTableColumns";
import { isDecisionMode } from "@/lib/decisionModels";
import { gatewayUrl, wildcardExample } from "./hubFormat";
const GREETING = "Hello, how are you?";
export const modelSnippet = (model: ModelGroupInfo): string => {
const name = wildcardExample(model.model_group);
if (isDecisionMode(model.mode)) {
return modelUsageExample(model.mode, gatewayUrl(), name, getGlobalLitellmHeaderName());
}
const params: Parameters<typeof generateCodeSnippet>[0] = {
apiKeySource: "custom",
accessToken: null,
apiKey: "your_api_key",
inputMessage: GREETING,
chatHistory: [{ role: "user", content: GREETING, isImage: false }],
selectedTags: [],
selectedVectorStores: [],
selectedGuardrails: [],
selectedPolicies: [],
selectedMCPServers: [],
endpointType: getEndpointType(model.mode || "chat"),
selectedModel: name,
selectedSdk: "openai",
proxySettings: { PROXY_BASE_URL: gatewayUrl() },
};
return generateCodeSnippet(params);
};
export const curlSnippet = (model: ModelGroupInfo): string | null => {
if (model.mode && model.mode !== "chat") return null;
return `curl ${gatewayUrl()}/v1/chat/completions \\
-H "Authorization: Bearer $LITELLM_API_KEY" \\
-H "Content-Type: application/json" \\
-d '{
"model": "${wildcardExample(model.model_group)}",
"messages": [{"role": "user", "content": "Hello"}]
}'`;
};
export const agentSnippets = (agent: AgentCard): { title: string; code: string }[] => [
{
title: "Retrieve the agent card",
code: `from a2a.client import A2ACardResolver
resolver = A2ACardResolver(httpx_client=httpx_client, base_url="${agent.url}")
card = await resolver.get_agent_card()`,
},
{
title: "Call the agent",
code: `from a2a.client import A2AClient
from a2a.types import MessageSendParams, SendMessageRequest
client = A2AClient(httpx_client=httpx_client, agent_card=card)
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(message={
"role": "user",
"parts": [{"kind": "text", "text": "how much is 10 USD in INR?"}],
"messageId": uuid4().hex,
}),
)
response = await client.send_message(request)`,
},
];
export const mcpEndpoint = (server: MCPServerData): string => `${gatewayUrl()}/${server.name}/mcp`;
export const mcpSnippet = (server: MCPServerData): string => `from fastmcp import Client
client = Client({
"mcpServers": {
"${server.name}": {
"url": "${mcpEndpoint(server)}",
"headers": {"x-litellm-api-key": "Bearer your_api_key"},
}
}
})
async with client:
tools = await client.list_tools()`;
export const marketplaceCommand = (): string => `/plugin marketplace add ${gatewayUrl()}/claude-code/marketplace.json`;
export const skillInstallCommand = (skill: Plugin): string => formatInstallCommand(skill);

88
ui/litellm-dashboard/src/components/publicModelHub/useHubData.ts Normal file
View file

@ -0,0 +1,88 @@
"use client";
import { useQuery } from "@tanstack/react-query";
import { useState } from "react";
import type { ResourceListPage } from "@/app/(dashboard)/hooks/common/useResourceList";
import type { Plugin } from "@/components/claude_code_plugins/types";
import { apiClient, getPublicModelHubInfo, getUiConfig } from "@/components/networking";
import type { AgentCard, MCPServerData, ModelGroupInfo } from "@/components/PublicModelHubTableColumns";
import { usePublicModelHubFacets } from "./usePublicModelHubFacets";
import { PUBLIC_MODEL_HUB_PAGE_SIZE, PUBLIC_MODEL_HUB_PATH, usePublicModelHubList } from "./usePublicModelHubList";
export interface HubInfo {
docs_title: string;
custom_docs_description: string | null;
litellm_version: string;
useful_links: Record<string, string | { url: string; index?: number }>;
}
const asArray = <T>(data: unknown): T[] => (Array.isArray(data) ? (data as T[]) : []);
const UI_CONFIG_QUERY = {
queryKey: ["publicAiHub", "uiConfig"],
queryFn: () => getUiConfig(),
staleTime: Infinity,
retry: false,
};
export function useHubData() {
const config = useQuery(UI_CONFIG_QUERY);
const ready = config.isSuccess || config.isError;
const info = useQuery<HubInfo>({
queryKey: ["publicAiHub", "info"],
queryFn: () => getPublicModelHubInfo(),
enabled: ready,
});
const agents = useQuery<AgentCard[]>({
queryKey: ["publicAiHub", "agents"],
queryFn: async ({ signal }) => asArray<AgentCard>(await apiClient.get("/public/agent_hub", { signal })),
enabled: ready,
});
const mcp = useQuery<MCPServerData[]>({
queryKey: ["publicAiHub", "mcp"],
queryFn: async ({ signal }) => asArray<MCPServerData>(await apiClient.get("/public/mcp_hub", { signal })),
enabled: ready,
});
const skills = useQuery<Plugin[]>({
queryKey: ["publicAiHub", "skills"],
queryFn: async ({ signal }) =>
(await apiClient.get<{ plugins?: Plugin[] }>("/public/skill_hub", { signal })).plugins ?? [],
enabled: ready,
});
const list = usePublicModelHubList(ready);
const listLoading = !ready || list.isLoading;
const models = { ...list, isLoading: listLoading && list.rows.length === 0 };
const facets = usePublicModelHubFacets(ready);
const [modelTotal, setModelTotal] = useState<number | undefined>(undefined);
const settledTotal = !listLoading && !models.hasActiveQuery && !models.error ? models.rowCount : modelTotal;
if (settledTotal !== modelTotal) setModelTotal(settledTotal);
return { info, agents, mcp, skills, models, facets, modelTotal: settledTotal };
}
export type HubData = ReturnType<typeof useHubData>;
const openedModelQuery = (open: string | null, enabled: boolean) => ({
queryKey: ["publicAiHub", "openedModel", open],
queryFn: ({ signal }: { signal: AbortSignal }) =>
apiClient.get<ResourceListPage<ModelGroupInfo>>(PUBLIC_MODEL_HUB_PATH, {
query: { q: open ?? "", page_size: PUBLIC_MODEL_HUB_PAGE_SIZE },
signal,
}),
enabled,
retry: false,
});
export function useOpenedModel(open: string | null, rows: ModelGroupInfo[], enabled: boolean) {
const onPage = rows.find((model) => model.model_group === open);
const [kept, setKept] = useState<ModelGroupInfo | undefined>(undefined);
const known = onPage ?? (kept?.model_group === open ? kept : undefined);
const lookup = useQuery(openedModelQuery(open, enabled && open !== null && known === undefined));
const model = known ?? lookup.data?.data.find((found) => found.model_group === open);
if (model !== undefined && model !== kept) setKept(model);
return { model, pending: model === undefined && lookup.isLoading };
}

32
ui/litellm-dashboard/src/components/publicModelHub/useHubUrlState.ts Normal file
View file

@ -0,0 +1,32 @@
"use client";
import { createSerializer, parseAsString, useQueryState } from "nuqs";
import { useSearchParams } from "next/navigation";
import { useCallback } from "react";
import { useUrlTab } from "@/hooks/useUrlTab";
export const HUB_TABS = ["models", "agents", "mcp", "skills"] as const;
export type HubTab = (typeof HUB_TABS)[number];
const OPEN_PARAM = { open: parseAsString };
const serializeOpen = createSerializer(OPEN_PARAM);
export function useHubUrlState() {
const [tab, setUrlTab] = useUrlTab(HUB_TABS, "models");
const [open, setOpenParam] = useQueryState("open", OPEN_PARAM.open);
const setOpen = useCallback((key: string | null) => void setOpenParam(key), [setOpenParam]);
const setTab = useCallback(
(next: HubTab) => {
setUrlTab(next);
setOpen(null);
},
[setUrlTab, setOpen],
);
return { tab, setTab, open: open || null, setOpen };
}
export function useOpenHref() {
const params = useSearchParams();
return useCallback((key: string) => serializeOpen(params ?? new URLSearchParams(), { open: key }), [params]);
}

16
ui/litellm-dashboard/src/components/publicModelHub/useIsNarrow.ts Normal file
View file

@ -0,0 +1,16 @@
import { useSyncExternalStore } from "react";
const NARROW = "(width < 64rem)";
const subscribe = (onChange: () => void) => {
const query = window.matchMedia(NARROW);
query.addEventListener("change", onChange);
return () => query.removeEventListener("change", onChange);
};
export const useIsNarrow = (): boolean =>
useSyncExternalStore(
subscribe,
() => window.matchMedia(NARROW).matches,
() => false,
);

19
ui/litellm-dashboard/src/components/publicModelHub/useRequirePublicHubAuth.ts Normal file
View file

@ -0,0 +1,19 @@
"use client";
import { useEffect } from "react";
import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
import { getProxyBaseUrl, getUiPageUrl } from "@/components/networking";
import { getCookie } from "@/utils/cookieUtils";
import { checkTokenValidity } from "@/utils/jwtUtils";
import { getLoginUrl } from "@/utils/returnUrlUtils";
export function useRequirePublicHubAuth() {
const { data: uiSettings, isLoading } = useUISettings();
const requireAuth = !isLoading && uiSettings?.values?.require_auth_for_public_ai_hub === true;
useEffect(() => {
if (requireAuth && !checkTokenValidity(getCookie("token"))) {
window.location.replace(getUiPageUrl(getLoginUrl(getProxyBaseUrl())));
}
}, [requireAuth]);
}