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

feat(python-bridge): serve ValkeySemanticCache natively

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 20:33:32 +00:00
parent 1f86bb8e46
commit 4db3481244
10 changed files with 598 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/Cargo.lock generated
View file

@ -2685,6 +2685,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-cache-valkey-semantic",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -2695,6 +2696,7 @@ dependencies = [
"litellm-types",
"pyo3",
"pyo3-async-runtimes",
"redis",
"rstest",
"serde",
"serde_json",

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

@ -24,6 +24,7 @@ litellm-cache.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
serde.workspace = true
litellm-auth.workspace = true
litellm-callbacks-legacy-python.workspace = true
@ -37,6 +38,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"] }

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

@ -73,9 +73,18 @@ pub(super) struct RedisCacheConfig {
pub(super) connection: RedisConnectionConfig,
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct ValkeySemanticCacheConfig {
pub(super) similarity_threshold: f64,
pub(super) index_name: String,
pub(super) embedding_model: String,
pub(super) connection: RedisConnectionConfig,
}
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
@ -142,9 +151,15 @@ 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::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic
@ -158,11 +173,13 @@ impl NativeCacheConfig {
}
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
if service.default_ttl()
!= Some(match &self.backend {
CacheBackendConfig::Memory(config) => config.default_ttl,
CacheBackendConfig::Redis(config) => config.default_ttl,
})
if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_))
&& service.default_ttl()
!= Some(match &self.backend {
CacheBackendConfig::Memory(config) => config.default_ttl,
CacheBackendConfig::Redis(config) => config.default_ttl,
CacheBackendConfig::ValkeySemantic(_) => Duration::ZERO,
})
{
return Some("facade and native backend default TTLs must match");
}
@ -185,6 +202,16 @@ impl NativeCacheConfig {
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
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")
}
}
}
}
@ -299,6 +326,51 @@ fn project_redis(
}))
}
#[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")?;
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")?);
if !class_is(&connection_class, "redis.connection", "Connection")?
&& !class_is(&connection_class, "redis.connection", "SSLConnection")?
{
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_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
Ok(RedisTlsConfig {

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

@ -0,0 +1,58 @@
use std::{future::Future, sync::Arc};
use litellm_cache::Error;
use litellm_cache_valkey_semantic::Embedder;
use litellm_host_python::to_py;
use pyo3::prelude::*;
use serde_json::Value;
#[derive(Clone)]
pub(super) struct PythonEmbedder {
sync_embed: Arc<Py<PyAny>>,
async_embed: 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: Arc::new(backend.getattr("_get_async_embedding")?.unbind()),
})
}
}
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())
}
fn async_embed(
&self,
prompt: &str,
metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
let callable = Arc::clone(&self.async_embed);
let prompt = prompt.to_owned();
let metadata = metadata.cloned();
async move {
let future = Python::attach(|py| -> PyResult<_> {
let metadata = to_py(py, &metadata)?;
let awaitable = callable.bind(py).call1((prompt, metadata))?;
pyo3_async_runtimes::tokio::into_future(awaitable)
})
.map_err(|_| Error::Unavailable)?;
let result = future.await.map_err(|_| Error::Unavailable)?;
let result = Python::attach(|py| result.bind(py).extract::<Vec<f64>>())
.map_err(|_| Error::Unavailable)?;
Ok(result.into_iter().map(|value| value as f32).collect())
}
}
}

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

@ -36,6 +36,7 @@ pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
redis_pool: Option<RedisPoolGuard>,
redis_client_name: Option<&'static str>,
}
impl ObjectGuard {
@ -117,7 +118,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);
}
}
@ -138,10 +141,8 @@ impl ObjectGuard {
}
impl RedisPoolGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let pool = backend
.getattr("redis_client")?
.getattr("connection_pool")?;
fn capture(backend: &Bound<'_, PyAny>, client_name: &str) -> PyResult<Self> {
let pool = backend.getattr(client_name)?.getattr("connection_pool")?;
Ok(Self {
reference: pool.clone().unbind(),
connection_class: pool.getattr("connection_class")?.unbind(),
@ -153,10 +154,13 @@ impl RedisPoolGuard {
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let pool = backend
.getattr("redis_client")?
.getattr("connection_pool")?;
fn matches(
&self,
py: Python<'_>,
backend: &Bound<'_, PyAny>,
client_name: &str,
) -> PyResult<bool> {
let pool = backend.getattr(client_name)?.getattr("connection_pool")?;
Ok(self.reference.bind(py).is(&pool)
&& self
.connection_class
@ -192,6 +196,11 @@ impl FacadeGuard {
let (module, name, cache_kind) = match kind {
"memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
"redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"),
"valkey-semantic" => (
"litellm.caching.valkey_semantic_cache",
"ValkeySemanticCache",
"valkey-semantic",
),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -235,11 +244,32 @@ impl FacadeGuard {
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
"similarity_threshold",
"embedding_model",
"index_name",
"embedding_max_input_tokens",
"embedding_timeout",
],
)?,
redis_pool: (kind == "redis")
.then(|| RedisPoolGuard::capture(&backend))
redis_pool: (kind == "redis" || kind == "valkey-semantic")
.then(|| {
RedisPoolGuard::capture(
&backend,
if kind == "redis" {
"redis_client"
} else {
"sync_client"
},
)
})
.transpose()?,
redis_client_name: (kind == "redis" || kind == "valkey-semantic").then_some(
if kind == "redis" {
"redis_client"
} else {
"sync_client"
},
),
})
}
@ -252,7 +282,11 @@ impl FacadeGuard {
return Ok(false);
}
match &self.redis_pool {
Some(guard) => guard.matches(py, &backend),
Some(guard) => guard.matches(
py,
&backend,
self.redis_client_name.unwrap_or("redis_client"),
),
None => Ok(true),
}
}

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

@ -1,7 +1,10 @@
use litellm_host_python::release_gil;
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
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 {
@ -51,6 +54,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(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()
@ -59,11 +85,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 {

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

@ -1,6 +1,7 @@
mod binding;
mod callback;
mod config;
mod embedder;
mod facade;
mod future;
mod handle;
@ -10,7 +11,7 @@ mod resolver;
use litellm_cache::Error;
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
prelude::*,
};
@ -21,6 +22,7 @@ pub(crate) use self::{
fn cache_error(error: Error) -> PyErr {
match error {
Error::InvalidEntry => PyValueError::new_err(error.to_string()),
Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()),
_ => PyRuntimeError::new_err(error.to_string()),
}
}

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

@ -1,13 +1,18 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache::{
CacheCodec, CacheConnectionResult, Error, ExactCacheContext, SemanticCacheContext,
};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::RedisCache;
use litellm_cache_response::{
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
};
use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig};
use serde_json::Value;
use super::{embedder::PythonEmbedder, request::NativeRequest};
#[derive(Clone)]
pub(super) enum NativeResponseCache {
Memory(Arc<ResponseCache<InMemoryCache<CacheEntry>>>),
@ -15,6 +20,10 @@ pub(super) enum NativeResponseCache {
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
buffer: Option<Arc<WriteBuffer>>,
},
ValkeySemantic {
cache: Arc<ResponseCache<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>>,
scope: String,
},
}
impl NativeResponseCache {
@ -43,41 +52,52 @@ impl NativeResponseCache {
buffer: None,
})
}
}
impl NativeResponseCache {
pub fn kind(&self) -> &'static str {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
pub fn valkey_semantic(
url: &str,
similarity_threshold: f64,
index_name: String,
embedder: PythonEmbedder,
) -> Result<Self, Error> {
let backend = ValkeySemanticCache::new(
url,
embedder,
ResponseCacheCodec,
ValkeySemanticConfig {
similarity_threshold,
index_name,
},
)?;
Ok(Self::ValkeySemantic {
cache: Arc::new(ResponseCache::new(Arc::new(backend))),
scope: String::from("key"),
})
}
fn exact(request: &NativeRequest) -> ResponseCacheRequest<ExactCacheContext> {
ResponseCacheRequest {
key: request.key.clone(),
controls: request.controls,
context: ExactCacheContext { ttl: request.ttl },
max_age: request.max_age,
}
}
pub fn default_ttl(&self) -> Option<Duration> {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
}
}
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } => None,
fn semantic(
request: &NativeRequest,
scope: &str,
) -> ResponseCacheRequest<SemanticCacheContext> {
ResponseCacheRequest {
key: request.key.clone(),
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,
}
}
@ -85,95 +105,206 @@ impl NativeResponseCache {
match self {
Self::Redis { cache, .. } => Self::Redis {
cache,
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
buffer: flush_size.map(|size| Arc::new(WriteBuffer::new(size))),
},
memory => memory,
value => value,
}
}
pub fn lookup(
&self,
request: &ResponseCacheRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
pub fn with_scope(self, scope: String) -> Self {
match self {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::ValkeySemantic { cache, .. } => Self::ValkeySemantic { cache, scope },
value => value,
}
}
pub fn kind(&self) -> &'static str {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::ValkeySemantic { .. } => "valkey-semantic",
}
}
pub fn default_ttl(&self) -> Option<Duration> {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
Self::ValkeySemantic { cache, .. } => cache.default_ttl(),
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) | Self::ValkeySemantic { .. } => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
}
}
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } | Self::ValkeySemantic { .. } => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } | Self::ValkeySemantic { .. } => None,
}
}
pub fn semantic_config(&self) -> Option<(f64, &str)> {
match self {
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::ValkeySemantic { cache, scope } => {
cache.lookup(&Self::semantic(request, scope), 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::Memory(cache) => cache.store(&Self::exact(request), response, now),
Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now),
Self::ValkeySemantic { cache, scope } => {
cache.store(&Self::semantic(request, scope), 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::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::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
}
}
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::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await,
Self::ValkeySemantic { cache, scope } => {
cache
.async_lookup(&Self::semantic(request, scope), now)
.await
}
}
}
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,
} => {
buffer
.async_store(cache, &Self::exact(request), response, now)
.await
}
Self::ValkeySemantic { cache, scope } => {
cache
.async_store(&Self::semantic(request, scope), response, now)
.await
}
}
}
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::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::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
}
}
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::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::ValkeySemantic { cache, scope } => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::semantic(&request, scope), value))
.collect();
cache.async_store_batch(entries, now).await
}
}
}
@ -186,6 +317,7 @@ impl NativeResponseCache {
}
cache.async_flush().await
}
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
}
}
@ -193,6 +325,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::ValkeySemantic { cache, .. } => cache.test_connection().await,
}
}
}

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

@ -5,6 +5,7 @@ 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)]
@ -13,24 +14,42 @@ struct RequestInput {
controls: Option<CacheControls>,
ttl_seconds: Option<f64>,
max_age_seconds: Option<f64>,
messages: Option<Value>,
input: Option<Value>,
metadata: 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) 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::<ExactCacheContext>::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,
})
}
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)

149
tests/test_litellm_rust/test_valkey_semantic_cache.py Normal file
View file

@ -0,0 +1,149 @@
import os
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
@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() -> dict[str, object]:
return {
"key": {"preset": "key"},
"messages": [{"role": "user", "content": "semantic cache prompt"}],
}
def _backend(url: str, index_name: str) -> ValkeySemanticCache:
backend: Final = ValkeySemanticCache(
redis_url=url,
similarity_threshold=0.8,
index_name=index_name,
)
backend._get_embedding = lambda prompt, metadata=None: [1.0, 0.0]
async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
return [1.0, 0.0]
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"}
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()])