Merge pull request #42316 from BerriAI/litellm_valkey_semantic_native_cache

This commit is contained in:
yujonglee 2026-09-21 16:28:12 -07:00 • committed by GitHub
commit 0b8faae494
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 3243 additions and 602 deletions

View file

@ -2772,6 +2772,22 @@ dependencies = [
"wiremock",
]
[[package]]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"tokio",
"uuid",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -2944,6 +2960,7 @@ dependencies = [
"litellm-cache-redis",
"litellm-cache-response",
"litellm-cache-s3",
"litellm-cache-valkey-semantic",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -2954,10 +2971,12 @@ dependencies = [
"litellm-types",
"pyo3",
"pyo3-async-runtimes",
"redis",
"rstest",
"serde",
"serde_json",
"serde_with",
"sha2 0.10.9",
"tokio",
"tokio-tungstenite",
]

View file

@ -14,7 +14,7 @@ use crate::topology::RedisTopology;
mod connection;
mod operations;
pub(crate) use connection::ConnectionRef;
pub use connection::ConnectionRef;
use connection::{ClusterConnectionManager, ConnectionManager};
pub use operations::{
@ -40,7 +40,8 @@ const CLAIM_SCRIPT: &str = concat!(
);
const CLAIM_ATTEMPTS: usize = 8;
enum Connections<C> {
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
@ -50,7 +51,7 @@ impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn execute<T>(
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
@ -73,6 +74,29 @@ where
}
}
}
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)?,
)?)),
}
}
}
pub struct RedisCache<S, C = redis::Connection> {
@ -94,12 +118,7 @@ impl<S: CacheCodec> RedisCache<S> {
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
let connections = match topology {
RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?),
RedisTopology::Cluster { startup_nodes } => {
Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?)
}
};
let connections = Connections::open(url, topology)?;
Ok(Self {
connections: Arc::new(connections),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
@ -127,7 +146,7 @@ where
{
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
Self {
connections: Arc::new(Connections::Fixed(Mutex::new(connection))),
connections: Arc::new(Connections::fixed(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
@ -203,16 +222,6 @@ where
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
async fn run_blocking<T, F>(connections: Arc<Connections<C>>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
}
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
@ -271,7 +280,7 @@ where
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
@ -285,7 +294,7 @@ where
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
@ -311,7 +320,7 @@ where
if entries.is_empty() {
return Ok(());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = entries
.into_iter()
.map(|(key, payload)| {
@ -330,7 +339,7 @@ where
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Self::run_blocking(Arc::clone(&self.connections), |connection| {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match connection.ping() {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
@ -391,7 +400,7 @@ where
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
@ -418,7 +427,7 @@ where
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
.await
@ -438,7 +447,7 @@ where
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
Self::flush_matching(connection, &pattern)
})
.await
@ -470,7 +479,7 @@ where
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment(connection, key, amount, ttl)
})
.await
@ -581,7 +590,7 @@ where
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
let codec = self.codec.clone();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
claim(connection, &codec, &key, candidate, &eligible, ttl)
})
.await

View file

@ -117,7 +117,7 @@ impl r2d2::ManageConnection for ClusterConnectionManager {
}
}
pub(crate) enum ConnectionRef<'a> {
pub enum ConnectionRef<'a> {
Node(&'a mut dyn redis::ConnectionLike),
Cluster(&'a mut ClusterConnection),
}

View file

@ -144,7 +144,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del(keys).map_err(|_| Error::Unavailable)
})
.await
@ -172,7 +172,7 @@ where
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
@ -188,7 +188,7 @@ where
}
pub async fn ping(&self) -> Result<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
Connections::run_blocking(Arc::clone(&self.connections), |connection| {
connection.ping().map_err(|_| Error::Unavailable)
})
.await
@ -196,7 +196,7 @@ where
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<i64>, Error> {
let key = self.namespaced_key(key);
let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
@ -208,7 +208,7 @@ where
pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
let pattern = format!("{}*", self.namespaced_key(pattern));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
@ -231,7 +231,7 @@ where
}
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut sadd = redis::cmd("SADD");
sadd.arg(&key).arg(values);
let mut expire = redis::cmd("EXPIRE");
@ -253,7 +253,7 @@ where
return Err(Error::InvalidEntry);
}
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("RPUSH")
.arg(key)
.arg(values)
@ -279,7 +279,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, values)| {
@ -304,7 +304,7 @@ where
) -> Result<RedisLpopResult, Error> {
let key = self.namespaced_key(key);
let multiple = count.is_some();
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
@ -333,7 +333,7 @@ where
.iter()
.map(|(_, count)| count.is_some())
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, count)| {
@ -365,7 +365,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
@ -426,7 +426,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut commands = Vec::with_capacity(operations.len() * 2);
let mut increments = Vec::with_capacity(operations.len());
for (key, amount, ttl) in operations {
@ -460,7 +460,7 @@ where
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment_with_floor(connection, key, amount, ttl)
})
.await
@ -474,7 +474,7 @@ where
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)

View file

@ -1,6 +1,10 @@
mod cache;
mod topology;
pub mod connection {
pub use crate::cache::{ConnectionRef, Connections};
}
pub use cache::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};

View file

@ -58,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na
Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths
Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees
Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees

View file

@ -1,21 +1,21 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache,
};
use serde_json::Value;
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
#[derive(Clone)]
pub struct ResponseCacheRequest {
pub struct ResponseCacheRequest<C: CacheContext = litellm_cache::ExactCacheContext> {
pub key: CacheKeyInput,
pub controls: CacheControls,
pub context: ExactCacheContext,
pub context: C,
pub max_age: Option<Duration>,
}
impl ResponseCacheRequest {
impl<C: CacheContext + Default> ResponseCacheRequest<C> {
pub fn new(key: CacheKeyInput) -> Self {
Self {
key,
@ -26,17 +26,24 @@ impl ResponseCacheRequest {
default_on: true,
..Default::default()
},
context: ExactCacheContext::default(),
context: C::default(),
max_age: None,
}
}
}
pub struct ResponseCache<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
where
B::Context: Default + PartialEq,
{
backend: Arc<B>,
}
impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
impl<B> ResponseCache<B>
where
B: BaseCache<Value = CacheEntry>,
B::Context: Default + PartialEq,
{
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
}
@ -45,8 +52,12 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
&self.backend
}
pub fn backend_arc(&self) -> &Arc<B> {
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration> {
self.backend.get_ttl(&ExactCacheContext::default())
self.backend.get_ttl(&B::Context::default())
}
pub async fn async_flush(&self) -> Result<(), Error>
@ -62,7 +73,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
@ -81,7 +92,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
@ -101,7 +112,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<B::Context>],
now: Duration,
) -> Result<PartialHits, Error>
where
@ -126,7 +137,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<B::Context>],
now: Duration,
) -> Result<PartialHits, Error>
where
@ -153,7 +164,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
@ -172,7 +183,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
@ -193,7 +204,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
entries: Vec<(ResponseCacheRequest<B::Context>, Value)>,
now: Duration,
) -> Result<(), Error> {
self.async_store_entries(
@ -209,7 +220,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
/// the freshness of its original response.
pub async fn async_store_entries(
&self,
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
entries: Vec<(ResponseCacheRequest<B::Context>, Value, Duration)>,
) -> Result<(), Error> {
let writable = entries
.into_iter()
@ -249,8 +260,8 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
}
fn partial_hits(
requests: &[ResponseCacheRequest],
readable: Vec<(usize, &ResponseCacheRequest)>,
requests: &[ResponseCacheRequest<B::Context>],
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,
entries: Vec<BatchEntry<CacheEntry>>,
now: Duration,
) -> Result<PartialHits, Error> {

View file

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

File diff suppressed because it is too large Load diff

View file

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

View file

@ -8,7 +8,7 @@ mod error;
pub use base_cache::{
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext,
ExactCacheContext,
ExactCacheContext, SemanticCacheContext,
};
pub use cache_type::CacheType;
pub use caching::{Cache, CacheBackend, get_cache, set_cache};

View file

@ -28,6 +28,7 @@ litellm-cache-s3.workspace = true
litellm-cache-gcs.workspace = true
litellm-cache-disk.workspace = true
litellm-cache-response.workspace = true
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
serde.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
@ -42,6 +43,7 @@ litellm-host-python.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
serde_json.workspace = true
tokio = { workspace = true, features = ["sync"] }
@ -51,6 +53,7 @@ serde_with.workspace = true
criterion.workspace = true
futures-util.workspace = true
rstest.workspace = true
sha2.workspace = true
tokio-tungstenite.workspace = true
[[bench]]

View file

@ -56,12 +56,7 @@ impl ResolvedCache {
CacheBinding::Disabled => ready_none(py)?,
CacheBinding::Native(service) => {
let request = request(input)?;
let service = service.clone();
run_async(
py,
async move { service.async_lookup(&request, now()).await },
cache_error,
)?
service.async_lookup_py(py, request)?
}
CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?,
};
@ -179,12 +174,7 @@ impl ResolvedCache {
CacheBinding::Native(service) => {
let request = self::request(request)?;
let response: Value = from_py(response)?;
let service = service.clone();
run_async(
py,
async move { service.async_store(&request, response, now()).await },
cache_error,
)
service.async_store_py(py, request, response)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store(py, response, callback_kwargs)
@ -241,12 +231,7 @@ impl ResolvedCache {
));
}
let entries = requests.into_iter().zip(responses).collect();
let service = service.clone();
run_async(
py,
async move { service.async_store_batch(entries, now()).await },
cache_error,
)
service.async_store_batch_py(py, entries)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store_batch(py, callback_result, callback_kwargs)

View file

@ -88,11 +88,6 @@ pub(super) struct GcsCacheConfig {
pub(super) path_service_account: Option<String>,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
struct RedisClientProjection<'py> {
topology: RedisTopology,
host: String,
@ -104,11 +99,25 @@ struct RedisClientProjection<'py> {
const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct ValkeySemanticCacheConfig {
pub(super) similarity_threshold: f64,
pub(super) index_name: String,
pub(super) embedding_model: String,
pub(super) connection: RedisConnectionConfig,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
S3(Box<S3CacheConfig>),
Gcs(GcsCacheConfig),
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
}
@ -201,6 +210,13 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::ValkeySemantic) => match project_valkey_semantic(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::ValkeySemantic(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Disk) => match project_disk(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
@ -214,12 +230,9 @@ impl NativeCacheConfig {
backend: CacheBackendConfig::AzureBlob(backend),
}))
}),
Some(
CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::QdrantSemantic,
)
| None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
)),
Some(CacheType::RedisSemantic | CacheType::QdrantSemantic) | None => Ok(
CacheConfigProjection::Unsupported(UnsupportedCacheConfig::Backend),
),
}
}
@ -228,11 +241,14 @@ impl NativeCacheConfig {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO),
CacheBackendConfig::Disk(_)
| CacheBackendConfig::AzureBlob(_)
| CacheBackendConfig::Gcs(_) => None,
};
if service.default_ttl() != default_ttl {
if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_))
&& service.default_ttl() != default_ttl
{
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
@ -306,6 +322,16 @@ impl NativeCacheConfig {
Some("facade and native backend credentials must match")
}
CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::ValkeySemantic(config) => {
if service.kind() != "valkey-semantic" {
return Some("facade and native backend types must match");
}
let Some((threshold, index_name)) = service.semantic_config() else {
return Some("facade and native backend types must match");
};
(threshold != config.similarity_threshold || index_name != config.index_name)
.then_some("facade and native semantic settings must match")
}
CacheBackendConfig::Disk(_) if service.kind() != "disk" => {
Some("facade and native backend types must match")
}
@ -588,7 +614,7 @@ fn project_standalone_client<'py>(
#[inline(never)]
fn project_cluster_client<'py>(
source: &Bound<'py, PyDict>,
source: &Bound<'_, PyDict>,
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
let Some(startup_nodes) = startup_nodes(source)? else {
@ -676,6 +702,71 @@ fn port(value: i64) -> PyResult<u16> {
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
}
#[inline(never)]
fn project_valkey_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<ValkeySemanticCacheConfig, UnsupportedCacheConfig>> {
let client = backend.getattr("sync_client")?;
let pool = client.getattr("connection_pool")?;
let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
for key in ["credential_provider", "redis_connect_func"] {
if has_value(&resolved, key)? {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
}
if is_tls {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let connection = RedisConnectionConfig {
host: required_string(&resolved, "host")?,
port: u16::try_from(required_i64(&resolved, "port")?)
.map_err(|_| PyValueError::new_err("invalid Redis port"))?,
database: optional_i64(&resolved, "db")?.unwrap_or(0),
username: optional_dict_string(&resolved, "username")?,
password: optional_dict_string(&resolved, "password")?,
protocol: RedisProtocol::Resp2,
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
read_timeout: None,
connect_timeout: None,
socket_keepalive: None,
health_check_interval: Duration::ZERO,
client_name: None,
tls: None,
};
if connection.host.is_empty() {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
Ok(Ok(ValkeySemanticCacheConfig {
similarity_threshold: backend.getattr("similarity_threshold")?.extract()?,
index_name: backend.getattr("index_name")?.extract()?,
embedding_model: backend.getattr("embedding_model")?.extract()?,
connection,
}))
}
#[inline(never)]
fn project_connection_pool<'py>(
pool: &Bound<'py, PyAny>,
) -> PyResult<Result<(Bound<'py, PyDict>, bool), UnsupportedCacheConfig>> {
if !instance_class_is(pool, "redis.connection", "ConnectionPool")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
let connection_class = resolved
.get_item("connection_class")?
.unwrap_or(pool.getattr("connection_class")?);
let is_tls = if class_is(&connection_class, "redis.connection", "Connection")? {
false
} else if class_is(&connection_class, "redis.connection", "SSLConnection")? {
true
} else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
Ok(Ok((resolved, is_tls)))
}
#[inline(never)]
fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
Ok(RedisTlsConfig {
@ -869,16 +960,13 @@ mod tests {
use pyo3::{prelude::*, types::PyDict};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::run_sync_value;
use litellm_cache_redis::{RedisNode, RedisTopology};
use super::{
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig,
NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
};
use crate::cache::native::NativeResponseCache;
use litellm_cache_redis::{RedisNode, RedisTopology};
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
facade(
@ -951,6 +1039,194 @@ mod tests {
});
}
#[test]
fn projects_resolved_redis_tls_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.max_connections = 29\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\
facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Redis cache should be supported");
};
let CacheBackendConfig::Redis(redis) = config.backend else {
panic!("expected Redis configuration");
};
assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777));
assert_eq!(redis.namespace.as_deref(), Some("team"));
assert_eq!(redis.flush_size, 31);
assert_eq!(redis.connection.host, "cache.internal");
assert_eq!(redis.connection.port, 6380);
assert_eq!(redis.connection.database, 4);
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(redis.connection.pool_size, 29);
let tls = redis.connection.tls.unwrap();
assert_eq!(
tls.certificate_requirement,
CertificateRequirement::Optional
);
assert!(tls.check_hostname);
assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem"));
assert_eq!(tls.ca_data.as_deref(), Some("CA DATA"));
assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem"));
assert_eq!(tls.client_key.as_deref(), Some("/client.key"));
});
}
#[test]
fn projects_valkey_semantic_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = Connection\n\
pool.max_connections = 12\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Valkey semantic cache should be supported");
};
let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else {
panic!("expected Valkey semantic configuration");
};
assert_eq!(valkey.similarity_threshold, 0.85);
assert_eq!(valkey.index_name, "semantic_idx");
assert_eq!(valkey.embedding_model, "text-embedding-3-small");
assert_eq!(valkey.connection.host, "cache.internal");
assert_eq!(valkey.connection.port, 6390);
assert_eq!(valkey.connection.database, 2);
assert_eq!(valkey.connection.pool_size, 12);
assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2);
assert!(valkey.connection.tls.is_none());
});
}
#[test]
fn valkey_semantic_tls_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("TLS Valkey semantic cache should stay on Python");
};
assert_eq!(
reason.message(),
"native Redis connection type is not implemented"
);
});
}
#[test]
fn valkey_semantic_dynamic_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = Connection\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'credential_provider': object()}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic Valkey authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn dynamic_redis_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {
Python::initialize();
Python::attach(|py| {
let facade = cluster_facade(
py,
"[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]",
"client.on_connect",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("cluster startup nodes should project natively");
};
let CacheBackendConfig::Redis(redis) = &config.backend else {
panic!("expected Redis configuration");
};
let expected = RedisTopology::Cluster {
startup_nodes: vec![
RedisNode {
host: "node-a".into(),
port: 7000,
},
RedisNode {
host: "node-b".into(),
port: 7001,
},
],
};
assert_eq!(redis.topology, expected);
assert_eq!(redis.connection.host, "node-a");
assert_eq!(redis.connection.port, 7000);
assert_eq!(redis.connection.password.as_deref(), Some("secret"));
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(
redis
.connection
.tls
.as_ref()
.unwrap()
.certificate_requirement,
CertificateRequirement::None
);
});
}
#[test]
fn projects_gcs_configuration() {
Python::initialize();
@ -1016,375 +1292,6 @@ mod tests {
});
}
#[test]
fn projects_resolved_redis_tls_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.max_connections = 29\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\
facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Redis cache should be supported");
};
let CacheBackendConfig::Redis(redis) = config.backend else {
panic!("expected Redis configuration");
};
assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777));
assert_eq!(redis.namespace.as_deref(), Some("team"));
assert_eq!(redis.flush_size, 31);
assert_eq!(redis.connection.host, "cache.internal");
assert_eq!(redis.connection.port, 6380);
assert_eq!(redis.connection.database, 4);
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(redis.connection.pool_size, 29);
let tls = redis.connection.tls.unwrap();
assert_eq!(
tls.certificate_requirement,
CertificateRequirement::Optional
);
assert!(tls.check_hostname);
assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem"));
assert_eq!(tls.ca_data.as_deref(), Some("CA DATA"));
assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem"));
assert_eq!(tls.client_key.as_deref(), Some("/client.key"));
});
}
#[test]
fn dynamic_redis_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn projects_builtin_disk_configuration_and_rejects_custom_stores() {
Python::initialize();
Python::attach(|py| {
let root =
std::env::temp_dir().join(format!("litellm-disk-config-{}", std::process::id()));
let directory = root.to_string_lossy();
let disk_facade = facade(
py,
&format!(
"Cache = type('Cache', (), {{'__module__': 'diskcache.core'}})\n\
Disk = type('Disk', (), {{'__module__': 'diskcache.core'}})\n\
store = Cache()\n\
store._disk = Disk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&disk_facade).unwrap()
else {
panic!("disk cache should be supported");
};
let CacheBackendConfig::Disk(disk) = config.backend else {
panic!("expected disk configuration");
};
assert_eq!(disk.directory, root);
let matching = NativeResponseCache::disk(&directory).unwrap();
assert_eq!(
(NativeCacheConfig {
policy: config.policy,
backend: CacheBackendConfig::Disk(disk),
})
.service_mismatch(&matching),
None
);
let other = NativeResponseCache::disk(&root.join("other").to_string_lossy()).unwrap();
let mismatch = NativeCacheConfig {
policy: CachePolicy {
mode: "default-on".into(),
ttl: None,
namespace: None,
supported_call_types: None,
redis_flush_size: None,
semantic_cache_scope: "key".into(),
},
backend: CacheBackendConfig::Disk(DiskCacheConfig {
directory: root.clone(),
}),
};
assert_eq!(
mismatch.service_mismatch(&other),
Some("facade and native backend directories must match")
);
let custom = facade(
py,
&format!(
"CustomCache = type('CustomCache', (), {{'__module__': 'mypkg'}})\n\
CustomDisk = type('CustomDisk', (), {{'__module__': 'mypkg'}})\n\
store = CustomCache()\n\
store._disk = CustomDisk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&custom).unwrap()
else {
panic!("custom disk store must stay on Python");
};
assert_eq!(
reason.message(),
"native disk cache requires the built-in diskcache store"
);
});
}
fn s3_facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> {
let locals = PyDict::new(py);
py.run(
&CString::new(format!(
"from types import SimpleNamespace\n\
S3Client = type('S3', (), {{'__module__': 'botocore.client'}})\n\
client = S3Client()\n\
client.meta = SimpleNamespace(region_name='us-east-1', endpoint_url='https://example.test', config=SimpleNamespace(s3=None, proxies=None, client_cert=None, signature_version='s3v4'))\n\
client._endpoint = SimpleNamespace(http_session=SimpleNamespace(_verify=True))\n\
client._request_signer = SimpleNamespace(_credentials=SimpleNamespace(method='explicit', access_key='key', secret_key='secret', token='token'))\n\
backend = SimpleNamespace(bucket_name='bucket', key_prefix='team/', s3_client=client)\n\
facade = SimpleNamespace(type='s3', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)\n\
{body}"
))
.unwrap(),
None,
Some(&locals),
)
.unwrap();
locals.get_item("facade").unwrap().unwrap()
}
#[test]
fn projects_s3_configuration_with_explicit_credentials_and_custom_endpoint() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(py, "");
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert_eq!(s3.bucket, "bucket");
assert_eq!(s3.key_prefix, "team/");
assert_eq!(s3.region, "us-east-1");
assert_eq!(
s3.endpoint.map(|endpoint| endpoint.url).as_deref(),
Some("https://example.test")
);
assert_eq!(s3.auth.access_key_id.as_deref(), Some("key"));
assert_eq!(s3.auth.secret_access_key.as_deref(), Some("secret"));
assert_eq!(s3.auth.session_token.as_deref(), Some("token"));
assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1"));
});
}
#[test]
fn default_s3_endpoint_projects_no_custom_endpoint() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(
py,
"facade.cache.s3_client.meta.endpoint_url = 'https://s3.us-east-1.amazonaws.com'",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert!(s3.endpoint.is_none());
});
}
#[test]
fn non_sigv4_proxies_and_disabled_verification_stay_on_python() {
Python::initialize();
Python::attach(|py| {
for (body, message) in [
(
"facade.cache.s3_client.meta.config.signature_version = 's3'",
"native S3 configuration requires Python",
),
(
"facade.cache.s3_client.meta.config.proxies = {'https': 'proxy'}",
"native S3 configuration requires Python",
),
(
"facade.cache.s3_client._endpoint.http_session._verify = False",
"native S3 configuration requires Python",
),
(
"del facade.cache.s3_client._endpoint.http_session._verify",
"native S3 configuration requires Python",
),
] {
let facade = s3_facade(py, body);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("{body} must stay on Python");
};
assert_eq!(reason.message(), message);
}
let facade = s3_facade(py, "facade.cache.s3_client = SimpleNamespace()");
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("non-botocore client must stay on Python");
};
assert_eq!(reason.message(), "native S3 client type is not implemented");
let facade = s3_facade(
py,
"facade.cache.s3_client._request_signer._credentials = None",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("missing credentials must stay on Python");
};
assert_eq!(reason.message(), "native S3 credentials require Python");
});
}
#[test]
fn non_explicit_s3_credentials_use_the_default_chain() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(
py,
"facade.cache.s3_client._request_signer._credentials = SimpleNamespace(method='sso', access_key=None, secret_key=None, token=None)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("default-chain credentials should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert_eq!(s3.auth.access_key_id, None);
assert_eq!(s3.auth.secret_access_key, None);
assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1"));
});
}
fn s3_service(py: Python<'_>, region: &str, endpoint: Option<&str>) -> NativeResponseCache {
let config = S3CacheConfig {
bucket: "bucket".to_string(),
key_prefix: "team/".to_string(),
region: region.to_string(),
endpoint: endpoint.map(|url| S3Endpoint {
url: url.to_string(),
}),
auth: AwsAuthConfig::default(),
};
run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) }).unwrap()
}
#[test]
fn s3_binding_rejects_region_and_endpoint_mismatches() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(py, "");
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", Some("https://example.test"))),
None
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-west-2", Some("https://example.test"))),
Some("facade and native backend regions must match")
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", Some("https://other.test"))),
Some("facade and native backend endpoints must match")
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", None)),
Some("facade and native backend endpoints must match")
);
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {
Python::initialize();
Python::attach(|py| {
let facade = cluster_facade(
py,
"[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]",
"client.on_connect",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("cluster startup nodes should project natively");
};
let CacheBackendConfig::Redis(redis) = &config.backend else {
panic!("expected Redis configuration");
};
let expected = RedisTopology::Cluster {
startup_nodes: vec![
RedisNode {
host: "node-a".into(),
port: 7000,
},
RedisNode {
host: "node-b".into(),
port: 7001,
},
],
};
assert_eq!(redis.topology, expected);
assert_eq!(redis.connection.host, "node-a");
assert_eq!(redis.connection.port, 7000);
assert_eq!(redis.connection.password.as_deref(), Some("secret"));
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(
redis
.connection
.tls
.as_ref()
.unwrap()
.certificate_requirement,
CertificateRequirement::None
);
});
}
#[test]
fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() {
Python::initialize();

View file

@ -0,0 +1,63 @@
use std::{future::Future, sync::Arc};
use litellm_cache::Error;
use litellm_cache_valkey_semantic::Embedder;
use litellm_host_python::to_py;
use pyo3::{PyTraverseError, PyVisit, prelude::*};
use serde_json::Value;
#[derive(Clone)]
pub(super) struct PythonEmbedder {
sync_embed: Arc<Py<PyAny>>,
async_embed_callable: Arc<Py<PyAny>>,
}
impl PythonEmbedder {
pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()),
async_embed_callable: Arc::new(backend.getattr("_get_async_embedding")?.unbind()),
})
}
pub(super) fn async_embed_awaitable<'py>(
&self,
py: Python<'py>,
prompt: &str,
metadata: &Option<Value>,
) -> PyResult<Bound<'py, PyAny>> {
let metadata = to_py(py, metadata)?;
self.async_embed_callable.bind(py).call1((prompt, metadata))
}
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&*self.sync_embed)?;
visit.call(&*self.async_embed_callable)
}
}
impl Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
let result = Python::attach(|py| -> PyResult<Vec<f64>> {
let metadata = to_py(py, &metadata)?;
self.sync_embed
.bind(py)
.call1((prompt, metadata))?
.extract()
})
.map_err(|_| Error::Unavailable)?;
Ok(result.into_iter().map(|value| value as f32).collect())
}
#[expect(
clippy::manual_async_fn,
reason = "the shared Embedder trait uses an impl Future return"
)]
fn async_embed(
&self,
_prompt: &str,
_metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
async { Err(Error::Unavailable) }
}
}

View file

@ -31,17 +31,15 @@ struct RedisPoolGuard {
connection_class: Py<PyAny>,
connection_kwargs: Py<PyAny>,
max_connections: Option<usize>,
client_name: &'static str,
attributes: RedisPoolAttributes,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
struct DiskStoreGuard {
reference: Py<PyAny>,
directory: String,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
@ -49,12 +47,18 @@ struct AzureBlobClientGuard {
container_name: String,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
enum ConnectionGuard {
None,
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
S3(S3ClientGuard),
}
#[derive(Clone, Copy)]
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
@ -72,6 +76,9 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
connection_class: "connection_pool_class",
max_connections: None,
};
const VALKEY_POOL: RedisPoolAttributes = STANDALONE_POOL;
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
@ -158,7 +165,9 @@ impl ObjectGuard {
return Ok(false);
}
for (name, value) in &expected.attributes {
if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) {
if (instance.contains(name)? && !self.config_names.contains(&name.as_str()))
|| !attributes.get_item(name)?.is(value.bind(py))
{
return Ok(false);
}
}
@ -179,8 +188,12 @@ impl ObjectGuard {
}
impl RedisPoolGuard {
fn capture(backend: &Bound<'_, PyAny>, attributes: RedisPoolAttributes) -> PyResult<Self> {
let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?;
fn capture(
backend: &Bound<'_, PyAny>,
client_name: &'static str,
attributes: RedisPoolAttributes,
) -> PyResult<Self> {
let pool = backend.getattr(client_name)?.getattr(attributes.pool)?;
Ok(Self {
reference: pool.clone().unbind(),
connection_class: pool.getattr(attributes.connection_class)?.unbind(),
@ -188,31 +201,30 @@ impl RedisPoolGuard {
.getattr("connection_kwargs")?
.call_method0("copy")?
.unbind(),
max_connections: Self::max_connections(&pool, &attributes)?,
max_connections: attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()?,
client_name,
attributes,
})
}
fn max_connections(
pool: &Bound<'_, PyAny>,
attributes: &RedisPoolAttributes,
) -> PyResult<Option<usize>> {
attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let pool = backend
.getattr("redis_client")?
.getattr(self.client_name)?
.getattr(self.attributes.pool)?;
Ok(self.reference.bind(py).is(&pool)
&& self
.connection_class
.bind(py)
.is(&pool.getattr(self.attributes.connection_class)?)
&& self.max_connections == Self::max_connections(&pool, &self.attributes)?
&& self.max_connections
== self
.attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()?
&& self
.connection_kwargs
.bind(py)
@ -226,22 +238,6 @@ impl RedisPoolGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl DiskStoreGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let store = backend.getattr("disk_cache")?;
@ -261,6 +257,7 @@ impl DiskStoreGuard {
visit.call(&self.reference)
}
}
impl AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
@ -289,11 +286,41 @@ impl AzureBlobClientGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl ConnectionGuard {
fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(match (kind, cluster) {
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?),
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"redis_client",
STANDALONE_POOL,
)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"redis_client",
CLUSTER_POOL,
)?),
("valkey-semantic", _) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"sync_client",
VALKEY_POOL,
)?),
("disk", _) => Self::None,
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
("s3", _) => Self::S3(S3ClientGuard::capture(backend)?),
_ => Self::None,
@ -318,6 +345,7 @@ impl ConnectionGuard {
}
}
}
impl FacadeGuard {
pub(super) fn capture(
py: Python<'_>,
@ -340,14 +368,19 @@ impl FacadeGuard {
"RedisClusterCache",
"redis",
),
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
("valkey-semantic", false) => (
"litellm.caching.valkey_semantic_cache",
"ValkeySemanticCache",
"valkey-semantic",
),
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
"azure-blob",
),
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -391,6 +424,11 @@ impl FacadeGuard {
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
"similarity_threshold",
"embedding_model",
"index_name",
"embedding_max_input_tokens",
"embedding_timeout",
"bucket_name",
"key_prefix",
"path_service_account",

View file

@ -1,12 +1,14 @@
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
use super::{
cache_error, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache,
request::duration,
};
#[pyclass(frozen, name = "_CacheTestHandle")]
pub(crate) struct CacheTestHandle {
@ -139,6 +141,29 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[pyo3(signature = (url, similarity_threshold, index_name, embedder))]
fn valkey_semantic(
url: String,
similarity_threshold: f64,
index_name: String,
embedder: &Bound<'_, PyAny>,
) -> PyResult<Self> {
let python_embedder = PythonEmbedder::from_backend(embedder)?;
let service = NativeResponseCache::valkey_semantic(
&url,
similarity_threshold,
index_name,
python_embedder,
)
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (account_url, container))]
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {
@ -153,6 +178,7 @@ impl CacheTestHandle {
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()
@ -161,11 +187,17 @@ impl CacheTestHandle {
fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
let service = self.service()?;
let guard = FacadeGuard::capture(py, facade, &service)?;
let service = service.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
let service = service
.with_scope(
facade
.getattr("semantic_cache_scope")?
.extract::<String>()?,
)
.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
let handle = Py::new(
py,
Self {

View file

@ -1,12 +1,14 @@
mod binding;
mod callback;
mod config;
mod embedder;
mod facade;
mod future;
mod handle;
mod native;
mod request;
mod resolver;
mod semantic_step;
use litellm_cache::Error;
use pyo3::{

View file

@ -1,17 +1,77 @@
use std::{path::Path, sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache::{
CacheCodec, CacheConnectionResult, Error, ExactCacheContext, SemanticCacheContext,
};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_disk::DiskCache;
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, WriteBuffer,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig};
use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig};
use pyo3::prelude::*;
use serde_json::Value;
use super::{
embedder::PythonEmbedder,
request::NativeRequest,
semantic_step::{SemanticEmbedExecution, drive_semantic},
};
fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response::CacheKeyInput {
let mut key = request.key.clone();
if key.preset.is_some() {
return key;
}
key.fields
.retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input"));
const TENANT: [&str; 3] = [
"user_api_key",
"user_api_key_team_id",
"user_api_key_org_id",
];
let end_user = (scope == "end_user").then_some("user_api_key_end_user_id");
for name in TENANT.into_iter().chain(end_user) {
let sources = [
request.metadata.as_ref(),
request.litellm_metadata.as_ref(),
request
.litellm_params
.as_ref()
.and_then(|params| params.get("metadata")),
request
.litellm_params
.as_ref()
.and_then(|params| params.get("litellm_metadata")),
];
let Some(value) = sources.into_iter().flatten().find_map(|source| {
source
.as_object()
.and_then(|values| values.get(name))
.filter(|value| !value.is_null())
}) else {
continue;
};
let value = match value {
Value::Null => continue,
Value::String(text) => text.clone(),
other => other.to_string(),
};
key.fields.push(CacheKeyField {
name: name.to_owned(),
value: Some(value),
api_parameter: true,
internal_parameter: false,
});
}
key
}
#[derive(Clone)]
pub(super) enum NativeResponseCache {
Memory(Arc<ResponseCache<InMemoryCache<CacheEntry>>>),
@ -21,6 +81,11 @@ pub(super) enum NativeResponseCache {
},
S3(Arc<ResponseCache<S3Cache<ResponseCacheCodec>>>),
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
ValkeySemantic {
cache: Arc<ResponseCache<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>>,
embedder: PythonEmbedder,
scope: String,
},
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
@ -53,6 +118,7 @@ impl NativeResponseCache {
buffer: None,
})
}
pub async fn s3(config: S3CacheConfig) -> Self {
let runtime = tokio::runtime::Handle::current();
Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new(
@ -61,6 +127,29 @@ impl NativeResponseCache {
runtime,
)))))
}
pub fn valkey_semantic(
url: &str,
similarity_threshold: f64,
index_name: String,
embedder: PythonEmbedder,
) -> Result<Self, Error> {
let backend = ValkeySemanticCache::new(
url,
embedder.clone(),
ResponseCacheCodec,
ValkeySemanticConfig {
similarity_threshold,
index_name,
},
)?;
Ok(Self::ValkeySemantic {
cache: Arc::new(ResponseCache::new(Arc::new(backend))),
embedder,
scope: String::from("key"),
})
}
pub fn disk(directory: &str) -> Result<Self, Error> {
let cache = DiskCache::open(directory, ResponseCacheCodec)?;
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
@ -97,20 +186,72 @@ impl NativeResponseCache {
cache.backend().account_url(),
cache.backend().container_name(),
)),
Self::Memory(_) | Self::Redis { .. } | Self::S3(_) | Self::Disk(_) | Self::Gcs(_) => {
None
}
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::Gcs(_) => None,
}
}
fn exact(request: &NativeRequest) -> ResponseCacheRequest<ExactCacheContext> {
ResponseCacheRequest {
key: request.key.clone(),
controls: request.controls,
context: ExactCacheContext { ttl: request.ttl },
max_age: request.max_age,
}
}
fn semantic(
request: &NativeRequest,
scope: &str,
) -> ResponseCacheRequest<SemanticCacheContext> {
ResponseCacheRequest {
key: semantic_key(request, scope),
controls: request.controls,
context: SemanticCacheContext {
input: request.input.clone(),
messages: request.messages.clone(),
metadata: request.metadata.clone(),
scope: Some(scope.to_owned()),
ttl: request.ttl,
},
max_age: request.max_age,
}
}
pub fn with_redis_flush_size(self, flush_size: Option<usize>) -> Self {
match self {
Self::Redis { cache, .. } => Self::Redis {
cache,
buffer: flush_size.map(|size| Arc::new(WriteBuffer::new(size))),
},
value => value,
}
}
pub fn with_scope(self, scope: String) -> Self {
match self {
Self::ValkeySemantic {
cache, embedder, ..
} => Self::ValkeySemantic {
cache,
embedder,
scope,
},
value => value,
}
}
}
impl NativeResponseCache {
pub fn kind(&self) -> &'static str {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::S3(_) => "s3",
Self::Gcs(_) => "gcs",
Self::ValkeySemantic { .. } => "valkey-semantic",
Self::Disk(_) => "disk",
Self::AzureBlob(_) => "azure-blob",
}
@ -122,6 +263,7 @@ impl NativeResponseCache {
Self::Redis { cache, .. } => cache.default_ttl(),
Self::S3(cache) => cache.default_ttl(),
Self::Gcs(cache) => cache.default_ttl(),
Self::ValkeySemantic { cache, .. } => cache.default_ttl(),
Self::Disk(cache) => cache.default_ttl(),
Self::AzureBlob(cache) => cache.default_ttl(),
}
@ -157,17 +299,25 @@ impl NativeResponseCache {
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None,
Self::Memory(_)
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
Self::S3(_) | Self::Gcs(_) => None,
}
}
pub fn topology(&self) -> Option<&RedisTopology> {
match self {
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Memory(_)
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
Self::S3(_) => None,
}
}
@ -176,6 +326,7 @@ impl NativeResponseCache {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
@ -187,144 +338,386 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
pub fn with_redis_flush_size(self, flush_size: Option<usize>) -> Self {
match self {
Self::Redis { cache, .. } => Self::Redis {
cache,
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
},
other => other,
}
}
pub fn directory(&self) -> Option<&Path> {
match self {
Self::Disk(cache) => Some(cache.backend().directory()),
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
pub fn lookup(
&self,
request: &ResponseCacheRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
pub fn semantic_config(&self) -> Option<(f64, &str)> {
match self {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::S3(cache) => cache.lookup(request, now),
Self::Gcs(cache) => cache.lookup(request, now),
Self::Disk(cache) => cache.lookup(request, now),
Self::AzureBlob(cache) => cache.lookup(request, now),
Self::ValkeySemantic { cache, .. } => Some((
cache.backend().similarity_threshold(),
cache.backend().index_name(),
)),
_ => None,
}
}
pub fn lookup(&self, request: &NativeRequest, now: Duration) -> Result<Option<Value>, Error> {
match self {
Self::Memory(cache) => cache.lookup(&Self::exact(request), now),
Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now),
Self::S3(cache) => cache.lookup(&Self::exact(request), now),
Self::ValkeySemantic { cache, scope, .. } => {
cache.lookup(&Self::semantic(request, scope), now)
}
Self::Gcs(cache) => cache.lookup(&Self::exact(request), now),
Self::Disk(cache) => cache.lookup(&Self::exact(request), now),
Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now),
}
}
pub fn store(
&self,
request: &ResponseCacheRequest,
request: &NativeRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
Self::S3(cache) => cache.store(request, response, now),
Self::Gcs(cache) => cache.store(request, response, now),
Self::Disk(cache) => cache.store(request, response, now),
Self::AzureBlob(cache) => cache.store(request, response, now),
Self::Memory(cache) => cache.store(&Self::exact(request), response, now),
Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now),
Self::S3(cache) => cache.store(&Self::exact(request), response, now),
Self::ValkeySemantic { cache, scope, .. } => {
cache.store(&Self::semantic(request, scope), response, now)
}
Self::Gcs(cache) => cache.store(&Self::exact(request), response, now),
Self::Disk(cache) => cache.store(&Self::exact(request), response, now),
Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now),
}
}
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[NativeRequest],
now: Duration,
) -> Result<PartialHits, Error> {
match self {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
Self::S3(cache) => cache.lookup_batch(requests, now),
Self::Gcs(cache) => cache.lookup_batch(requests, now),
Self::Disk(cache) => cache.lookup_batch(requests, now),
Self::AzureBlob(cache) => cache.lookup_batch(requests, now),
Self::Memory(cache) => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.lookup_batch(&requests, now)
}
Self::Redis { cache, .. } => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.lookup_batch(&requests, now)
}
Self::S3(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::Gcs(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::Disk(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::AzureBlob(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
}
}
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
request: &NativeRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
match self {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
Self::S3(cache) => cache.async_lookup(request, now).await,
Self::Gcs(cache) => cache.async_lookup(request, now).await,
Self::Disk(cache) => cache.async_lookup(request, now).await,
Self::AzureBlob(cache) => cache.async_lookup(request, now).await,
Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await,
Self::S3(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::ValkeySemantic { cache, scope, .. } => {
cache
.async_lookup(&Self::semantic(request, scope), now)
.await
}
Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::Disk(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::AzureBlob(cache) => cache.async_lookup(&Self::exact(request), now).await,
}
}
pub(super) fn async_lookup_py<'py>(
&self,
py: Python<'py>,
request: NativeRequest,
) -> PyResult<Bound<'py, PyAny>> {
match self {
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => {
let service = self.clone();
litellm_host_python::run_async(
py,
async move { service.async_lookup(&request, super::request::now()).await },
super::cache_error,
)
}
Self::ValkeySemantic {
cache,
embedder,
scope,
} => drive_semantic(
py,
SemanticEmbedExecution::lookup(
Arc::clone(cache.backend_arc()),
embedder.clone(),
Self::semantic(&request, scope),
),
),
}
}
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
request: &NativeRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_store(request, response, now).await,
Self::Memory(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::Redis {
cache,
buffer: None,
} => cache.async_store(request, response, now).await,
} => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::Redis {
cache,
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
Self::S3(cache) => cache.async_store(request, response, now).await,
Self::Gcs(cache) => cache.async_store(request, response, now).await,
Self::Disk(cache) => cache.async_store(request, response, now).await,
Self::AzureBlob(cache) => cache.async_store(request, response, now).await,
} => {
buffer
.async_store(cache, &Self::exact(request), response, now)
.await
}
Self::S3(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::ValkeySemantic { cache, scope, .. } => {
cache
.async_store(&Self::semantic(request, scope), response, now)
.await
}
Self::Gcs(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::Disk(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::AzureBlob(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
}
}
pub(super) fn async_store_py<'py>(
&self,
py: Python<'py>,
request: NativeRequest,
response: Value,
) -> PyResult<Bound<'py, PyAny>> {
match self {
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => {
let service = self.clone();
litellm_host_python::run_async(
py,
async move {
service
.async_store(&request, response, super::request::now())
.await
},
super::cache_error,
)
}
Self::ValkeySemantic {
cache,
embedder,
scope,
} => drive_semantic(
py,
SemanticEmbedExecution::store(
Arc::clone(cache.backend_arc()),
embedder.clone(),
Self::semantic(&request, scope),
response,
),
),
}
}
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[NativeRequest],
now: Duration,
) -> Result<PartialHits, Error> {
match self {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
Self::S3(cache) => cache.async_lookup_batch(requests, now).await,
Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await,
Self::Disk(cache) => cache.async_lookup_batch(requests, now).await,
Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await,
Self::Memory(cache) => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.async_lookup_batch(&requests, now).await
}
Self::Redis { cache, .. } => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.async_lookup_batch(&requests, now).await
}
Self::S3(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::Gcs(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::Disk(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::AzureBlob(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
}
}
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
entries: Vec<(NativeRequest, Value)>,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
Self::S3(cache) => cache.async_store_batch(entries, now).await,
Self::Gcs(cache) => cache.async_store_batch(entries, now).await,
Self::Disk(cache) => cache.async_store_batch(entries, now).await,
Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await,
Self::Memory(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::Redis { cache, .. } => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::S3(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::ValkeySemantic { cache, scope, .. } => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::semantic(&request, scope), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::Gcs(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::Disk(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::AzureBlob(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
}
}
pub(super) fn async_store_batch_py<'py>(
&self,
py: Python<'py>,
entries: Vec<(NativeRequest, Value)>,
) -> PyResult<Bound<'py, PyAny>> {
match self {
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => {
let service = self.clone();
litellm_host_python::run_async(
py,
async move {
service
.async_store_batch(entries, super::request::now())
.await
},
super::cache_error,
)
}
Self::ValkeySemantic {
cache,
embedder,
scope,
} => {
let (requests, responses): (Vec<_>, Vec<_>) = entries
.into_iter()
.map(|(request, response)| (Self::semantic(&request, scope), response))
.unzip();
drive_semantic(
py,
SemanticEmbedExecution::store_batch(
Arc::clone(cache.backend_arc()),
embedder.clone(),
requests,
responses,
),
)
}
}
}
@ -338,6 +731,7 @@ impl NativeResponseCache {
cache.async_flush().await
}
Self::S3(cache) => cache.async_flush().await,
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::Gcs(cache) => cache.async_flush().await,
Self::Disk(cache) => cache.async_flush().await,
Self::AzureBlob(cache) => cache.async_flush().await,
@ -349,6 +743,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::S3(cache) => cache.test_connection().await,
Self::ValkeySemantic { cache, .. } => cache.test_connection().await,
Self::Gcs(cache) => cache.test_connection().await,
Self::Disk(cache) => cache.test_connection().await,
Self::AzureBlob(cache) => cache.test_connection().await,
@ -362,3 +757,79 @@ impl NativeResponseCache {
}
}
}
#[cfg(test)]
mod tests {
use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key};
use serde_json::json;
use sha2::{Digest, Sha256};
use super::*;
fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest {
NativeRequest {
key,
controls: CacheControls::default(),
ttl: None,
max_age: None,
messages: Some(json!([{"role": "user", "content": "prompt"}])),
input: None,
metadata: Some(metadata),
litellm_metadata: None,
litellm_params: None,
}
}
#[test]
fn semantic_key_matches_python_scope_material() {
let key = CacheKeyInput {
fields: vec![
CacheKeyField {
name: "model".to_owned(),
value: Some("gpt-4.1".to_owned()),
api_parameter: true,
internal_parameter: false,
},
CacheKeyField {
name: "messages".to_owned(),
value: Some("prompt".to_owned()),
api_parameter: true,
internal_parameter: false,
},
],
..Default::default()
};
let request = native_request(
key,
json!({"user_api_key": "k1", "user_api_key_team_id": null}),
);
let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1"));
assert_eq!(cache_key(&semantic_key(&request, "key")), expected);
let end_user_request = native_request(
request.key.clone(),
json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}),
);
let expected = format!(
"{:x}",
Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1")
);
assert_eq!(
cache_key(&semantic_key(&end_user_request, "end_user")),
expected
);
let preset_request = native_request(
CacheKeyInput {
preset: Some("preset-key".to_owned()),
..Default::default()
},
json!({"user_api_key": "k1"}),
);
assert_eq!(
semantic_key(&preset_request, "end_user").preset.as_deref(),
Some("preset-key")
);
assert!(semantic_key(&preset_request, "end_user").fields.is_empty());
}
}

View file

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

View file

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

View file

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