mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #42316 from BerriAI/litellm_valkey_semantic_native_cache
This commit is contained in:
commit
0b8faae494
22 changed files with 3243 additions and 602 deletions
19
litellm-rust/Cargo.lock
generated
19
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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> {
|
||||
|
|
|
|||
20
litellm-rust/crates/cache-valkey-semantic/Cargo.toml
Normal file
20
litellm-rust/crates/cache-valkey-semantic/Cargo.toml
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
[package]
|
||||
name = "litellm-cache-valkey-semantic"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
tokio.workspace = true
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
||||
[dev-dependencies]
|
||||
redis-test = "1.0.4"
|
||||
rstest.workspace = true
|
||||
1153
litellm-rust/crates/cache-valkey-semantic/src/lib.rs
Normal file
1153
litellm-rust/crates/cache-valkey-semantic/src/lib.rs
Normal file
File diff suppressed because it is too large
Load diff
50
litellm-rust/crates/cache/src/base_cache.rs
vendored
50
litellm-rust/crates/cache/src/base_cache.rs
vendored
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
2
litellm-rust/crates/cache/src/lib.rs
vendored
2
litellm-rust/crates/cache/src/lib.rs
vendored
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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]]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
63
litellm-rust/crates/python-bridge/src/cache/embedder.rs
vendored
Normal file
63
litellm-rust/crates/python-bridge/src/cache/embedder.rs
vendored
Normal 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) }
|
||||
}
|
||||
}
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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::{
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
249
litellm-rust/crates/python-bridge/src/cache/semantic_step.rs
vendored
Normal file
249
litellm-rust/crates/python-bridge/src/cache/semantic_step.rs
vendored
Normal 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,))
|
||||
}
|
||||
599
tests/test_litellm_rust/test_valkey_semantic_cache_native.py
Normal file
599
tests/test_litellm_rust/test_valkey_semantic_cache_native.py
Normal 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()
|
||||
Loading…
Add table
Reference in a new issue