mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
feat(cache): add native backend operation primitives
This commit is contained in:
parent
ef14fdaf33
commit
dc5f0c58a4
11 changed files with 910 additions and 9 deletions
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -3712,6 +3712,8 @@ dependencies = [
|
|||
"itoa",
|
||||
"num-bigint 0.5.1",
|
||||
"percent-encoding",
|
||||
"rustls 0.23.42",
|
||||
"rustls-native-certs",
|
||||
"ryu",
|
||||
"sha1_smol",
|
||||
"socket2 0.6.5",
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ use std::{
|
|||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache,
|
||||
Error,
|
||||
Error, IncrementOperation,
|
||||
};
|
||||
|
||||
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
|
||||
|
|
@ -134,6 +134,25 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
.copied())
|
||||
}
|
||||
|
||||
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
self.expires_at(key)
|
||||
}
|
||||
|
||||
pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result<Vec<String>, Error> {
|
||||
let state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
let mut expirations = state
|
||||
.expirations
|
||||
.iter()
|
||||
.map(|(key, expiration)| (key.clone(), *expiration))
|
||||
.collect::<Vec<_>>();
|
||||
expirations.sort_unstable_by_key(|(_, expiration)| *expiration);
|
||||
Ok(expirations
|
||||
.into_iter()
|
||||
.take(count)
|
||||
.map(|(key, _)| key)
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
|
||||
Self::remove(&mut state, key);
|
||||
|
|
@ -240,6 +259,27 @@ impl CounterCache for InMemoryCache<f64> {
|
|||
}
|
||||
}
|
||||
|
||||
impl InMemoryCache<f64> {
|
||||
pub async fn async_increment_pipeline(
|
||||
&self,
|
||||
operations: Vec<IncrementOperation>,
|
||||
) -> Result<Vec<f64>, Error> {
|
||||
operations
|
||||
.into_iter()
|
||||
.map(|operation| {
|
||||
self.increment_cache(
|
||||
&operation.key,
|
||||
operation.amount,
|
||||
CacheKwargs {
|
||||
ttl: operation.ttl,
|
||||
..CacheKwargs::default()
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
|
||||
type Value = V;
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ use std::{
|
|||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error,
|
||||
get_cache, set_cache,
|
||||
IncrementOperation, get_cache, set_cache,
|
||||
};
|
||||
use litellm_cache_memory::{CacheWrite, InMemoryCache};
|
||||
use rstest::{fixture, rstest};
|
||||
|
|
@ -305,3 +305,46 @@ fn disabled_cache_does_not_retain_claims_or_counters() {
|
|||
);
|
||||
assert_eq!(counters.get_cache("key").unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ttl_and_oldest_key_operations_use_the_stored_expirations() {
|
||||
let clock = Arc::new(AtomicU64::new(100));
|
||||
let cache = cache(clock, 3);
|
||||
cache
|
||||
.set_cache("later", "2".into(), Some(Duration::from_secs(20)))
|
||||
.unwrap();
|
||||
cache
|
||||
.set_cache("first", "1".into(), Some(Duration::from_secs(10)))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cache.async_get_ttl("first").await.unwrap(),
|
||||
Some(Duration::from_secs(110))
|
||||
);
|
||||
assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]);
|
||||
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn increment_pipeline_preserves_operation_order() {
|
||||
let cache = InMemoryCache::<f64>::new(Some(3), None);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_increment_pipeline(vec![
|
||||
IncrementOperation {
|
||||
key: "a".into(),
|
||||
amount: 1.0,
|
||||
ttl: Some(Duration::from_secs(10)),
|
||||
},
|
||||
IncrementOperation {
|
||||
key: "a".into(),
|
||||
amount: 2.0,
|
||||
ttl: Some(Duration::from_secs(20)),
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1.0, 3.0]
|
||||
);
|
||||
assert_eq!(cache.get_cache("a").unwrap(), Some(3.0));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
redis = "1.7.0"
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
r2d2 = "0.8.10"
|
||||
tokio.workspace = true
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,10 @@ use litellm_cache::{
|
|||
};
|
||||
use redis::Commands;
|
||||
|
||||
mod operations;
|
||||
|
||||
pub use operations::{RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
|
||||
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const REDIS_POOL_SIZE: u32 = 16;
|
||||
|
|
|
|||
500
litellm-rust/crates/cache-redis/src/cache/operations.rs
vendored
Normal file
500
litellm-rust/crates/cache-redis/src/cache/operations.rs
vendored
Normal file
|
|
@ -0,0 +1,500 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{CacheCodec, Error, IncrementOperation};
|
||||
use redis::Commands;
|
||||
|
||||
use super::{ConnectionRef, RedisCache};
|
||||
|
||||
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
|
||||
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
|
||||
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
|
||||
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return count"
|
||||
);
|
||||
const SET_MAX_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
|
||||
"redis.call('SET', KEYS[1], ARGV[1]); ",
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return ARGV[1]; end; return current"
|
||||
);
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum RedisArg {
|
||||
Bytes(Vec<u8>),
|
||||
Integer(i64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
impl From<&str> for RedisArg {
|
||||
fn from(value: &str) -> Self {
|
||||
Self::Bytes(value.as_bytes().to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for RedisArg {
|
||||
fn from(value: String) -> Self {
|
||||
Self::Bytes(value.into_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Vec<u8>> for RedisArg {
|
||||
fn from(value: Vec<u8>) -> Self {
|
||||
Self::Bytes(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i64> for RedisArg {
|
||||
fn from(value: i64) -> Self {
|
||||
Self::Integer(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<f64> for RedisArg {
|
||||
fn from(value: f64) -> Self {
|
||||
Self::Float(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl redis::ToRedisArgs for RedisArg {
|
||||
fn write_redis_args<W>(&self, out: &mut W)
|
||||
where
|
||||
W: ?Sized + redis::RedisWrite,
|
||||
{
|
||||
match self {
|
||||
Self::Bytes(value) => value.write_redis_args(out),
|
||||
Self::Integer(value) => value.write_redis_args(out),
|
||||
Self::Float(value) => value.write_redis_args(out),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct RedisRpushOperation {
|
||||
pub key: String,
|
||||
pub values: Vec<RedisArg>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RedisLpopOperation {
|
||||
pub key: String,
|
||||
pub count: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum RedisLpopResult {
|
||||
Missing,
|
||||
Value(Vec<u8>),
|
||||
Values(Vec<Vec<u8>>),
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
|
||||
if keys.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let keys = keys
|
||||
.into_iter()
|
||||
.map(|key| self.namespaced_key(&key))
|
||||
.collect::<Vec<_>>();
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
connection.del(keys).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
|
||||
let keys = keys
|
||||
.iter()
|
||||
.map(|key| self.namespaced_key(key))
|
||||
.collect::<Vec<_>>();
|
||||
let values = self.connections.execute(|connection| {
|
||||
redis::cmd("MGET")
|
||||
.arg(keys)
|
||||
.query::<Vec<redis::Value>>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})?;
|
||||
values.into_iter().map(count).collect()
|
||||
}
|
||||
|
||||
pub async fn async_batch_get_counts(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
) -> Result<Vec<Option<i64>>, Error> {
|
||||
let keys = keys
|
||||
.iter()
|
||||
.map(|key| self.namespaced_key(key))
|
||||
.collect::<Vec<_>>();
|
||||
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("MGET")
|
||||
.arg(keys)
|
||||
.query::<Vec<redis::Value>>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
values.into_iter().map(count).collect()
|
||||
}
|
||||
|
||||
pub fn sync_ping(&self) -> Result<bool, Error> {
|
||||
self.connections.execute(|connection| {
|
||||
redis::cmd("PING")
|
||||
.query::<String>(connection)
|
||||
.map(|response| response == "PONG")
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn ping(&self) -> Result<bool, Error> {
|
||||
Self::run_blocking(Arc::clone(&self.connections), |connection| {
|
||||
redis::cmd("PING")
|
||||
.query::<String>(connection)
|
||||
.map(|response| response == "PONG")
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
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| {
|
||||
redis::cmd("TTL")
|
||||
.arg(key)
|
||||
.query::<i64>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
Ok((ttl >= 0).then_some(ttl))
|
||||
}
|
||||
|
||||
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| {
|
||||
let mut cursor = 0u64;
|
||||
let mut matches = Vec::new();
|
||||
loop {
|
||||
let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
|
||||
.cursor_arg(cursor)
|
||||
.arg("MATCH")
|
||||
.arg(&pattern)
|
||||
.arg("COUNT")
|
||||
.arg(count)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
matches.extend(keys);
|
||||
if matches.len() >= count || next_cursor == 0 {
|
||||
matches.truncate(count);
|
||||
return Ok(matches);
|
||||
}
|
||||
cursor = next_cursor;
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_set_cache_sadd(
|
||||
&self,
|
||||
key: &str,
|
||||
values: Vec<RedisArg>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<usize, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
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| {
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline.cmd("SADD").arg(&key).arg(values);
|
||||
pipeline.cmd("EXPIRE").arg(&key).arg(ttl).ignore();
|
||||
pipeline
|
||||
.query::<(usize,)>(connection)
|
||||
.map(|(added,)| added)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_rpush(&self, key: &str, values: Vec<RedisArg>) -> Result<usize, Error> {
|
||||
if values.is_empty() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
let key = self.namespaced_key(key);
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("RPUSH")
|
||||
.arg(key)
|
||||
.arg(values)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_rpush_pipeline(
|
||||
&self,
|
||||
operations: Vec<RedisRpushOperation>,
|
||||
) -> Result<Vec<usize>, Error> {
|
||||
let operations = operations
|
||||
.into_iter()
|
||||
.map(|operation| {
|
||||
if operation.values.is_empty() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
Ok((self.namespaced_key(&operation.key), operation.values))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, values) in operations {
|
||||
pipeline.cmd("RPUSH").arg(key).arg(values);
|
||||
}
|
||||
pipeline.query(connection).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_lpop(
|
||||
&self,
|
||||
key: &str,
|
||||
count: Option<usize>,
|
||||
) -> 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 mut command = redis::cmd("LPOP");
|
||||
command.arg(key);
|
||||
if let Some(count) = count {
|
||||
command.arg(count);
|
||||
}
|
||||
command
|
||||
.query::<redis::Value>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
lpop_result(value, multiple)
|
||||
}
|
||||
|
||||
pub async fn async_lpop_pipeline(
|
||||
&self,
|
||||
operations: Vec<RedisLpopOperation>,
|
||||
) -> Result<Vec<RedisLpopResult>, Error> {
|
||||
let operations = operations
|
||||
.into_iter()
|
||||
.map(|operation| (self.namespaced_key(&operation.key), operation.count))
|
||||
.collect::<Vec<_>>();
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let multiple = operations
|
||||
.iter()
|
||||
.map(|(_, count)| count.is_some())
|
||||
.collect::<Vec<_>>();
|
||||
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, count) in operations {
|
||||
let command = pipeline.cmd("LPOP").arg(key);
|
||||
if let Some(count) = count {
|
||||
command.arg(count);
|
||||
}
|
||||
}
|
||||
pipeline
|
||||
.query::<Vec<redis::Value>>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
values
|
||||
.into_iter()
|
||||
.zip(multiple)
|
||||
.map(|(value, multiple)| lpop_result(value, multiple))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn async_eval(
|
||||
&self,
|
||||
script: String,
|
||||
keys: Vec<String>,
|
||||
arguments: Vec<RedisArg>,
|
||||
) -> Result<redis::Value, Error> {
|
||||
let keys = keys
|
||||
.into_iter()
|
||||
.map(|key| self.namespaced_key(&key))
|
||||
.collect::<Vec<_>>();
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("EVAL")
|
||||
.arg(script)
|
||||
.arg(keys.len())
|
||||
.arg(keys)
|
||||
.arg(arguments)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn client_list(&self) -> Result<String, Error> {
|
||||
self.connections.execute(|connection| {
|
||||
redis::cmd("CLIENT")
|
||||
.arg("LIST")
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn info(&self) -> Result<String, Error> {
|
||||
self.connections.execute(|connection| {
|
||||
redis::cmd("INFO")
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn flushall(&self) -> Result<(), Error> {
|
||||
self.connections.execute(|connection| {
|
||||
redis::cmd("FLUSHALL")
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, C> RedisCache<S, C>
|
||||
where
|
||||
S: CacheCodec<Value = f64>,
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
pub fn increment_with_floor(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: i64,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(ttl);
|
||||
self.connections
|
||||
.execute(|connection| increment_with_floor(connection, key, amount, ttl))
|
||||
}
|
||||
|
||||
pub async fn async_increment_pipeline(
|
||||
&self,
|
||||
operations: Vec<IncrementOperation>,
|
||||
) -> Result<Vec<f64>, Error> {
|
||||
let operations = operations
|
||||
.into_iter()
|
||||
.map(|operation| {
|
||||
(
|
||||
self.namespaced_key(&operation.key),
|
||||
operation.amount,
|
||||
operation.ttl.map(Self::ttl_seconds),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if operations.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
let mut pipeline = redis::pipe();
|
||||
for (key, amount, ttl) in operations {
|
||||
pipeline.cmd("INCRBYFLOAT").arg(&key).arg(amount);
|
||||
if let Some(ttl) = ttl {
|
||||
pipeline.cmd("EXPIRE").arg(key).arg(ttl).ignore();
|
||||
}
|
||||
}
|
||||
pipeline.query(connection).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_increment_with_floor(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: i64,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(ttl);
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
increment_with_floor(connection, key, amount, ttl)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn async_set_max(
|
||||
&self,
|
||||
key: &str,
|
||||
value: f64,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<f64, Error> {
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
|
||||
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
|
||||
redis::cmd("EVAL")
|
||||
.arg(SET_MAX_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(value)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
|
||||
match value {
|
||||
redis::Value::BulkString(bytes) => Ok(bytes),
|
||||
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(RedisLpopResult::Missing),
|
||||
redis::Value::Array(values) if multiple => values
|
||||
.into_iter()
|
||||
.map(redis_bytes)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(RedisLpopResult::Values),
|
||||
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::Int(value) => Ok(Some(value)),
|
||||
redis::Value::BulkString(value) => std::str::from_utf8(&value)
|
||||
.ok()
|
||||
.and_then(|value| value.parse().ok())
|
||||
.map(Some)
|
||||
.ok_or(Error::InvalidEntry),
|
||||
redis::Value::SimpleString(value) => {
|
||||
value.parse().map(Some).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn increment_with_floor(
|
||||
connection: &mut ConnectionRef<'_>,
|
||||
key: String,
|
||||
amount: i64,
|
||||
ttl: u64,
|
||||
) -> Result<i64, Error> {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(key)
|
||||
.arg(amount)
|
||||
.arg(ttl)
|
||||
.query(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
|
@ -1,3 +1,3 @@
|
|||
mod cache;
|
||||
|
||||
pub use cache::RedisCache;
|
||||
pub use cache::{RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@ use std::time::Duration;
|
|||
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, ClaimCache,
|
||||
CounterCache, Error, JsonCodec, get_cache, set_cache,
|
||||
CounterCache, Error, IncrementOperation, JsonCodec, get_cache, set_cache,
|
||||
};
|
||||
use litellm_cache_redis::{
|
||||
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation,
|
||||
};
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
|
||||
struct TaggedByteCodec(u8);
|
||||
|
|
@ -260,6 +262,307 @@ async fn async_flush_deletes_each_scan_page_separately() {
|
|||
cache.async_flush_cache().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() {
|
||||
let mut sadd_pipeline = redis::pipe();
|
||||
sadd_pipeline
|
||||
.cmd("SADD")
|
||||
.arg("team:members")
|
||||
.arg("a")
|
||||
.arg("b")
|
||||
.cmd("EXPIRE")
|
||||
.arg("team:members")
|
||||
.arg(600u64)
|
||||
.ignore();
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("MGET").arg("team:count").arg("team:missing"),
|
||||
Ok(redis_test::redis_value!(["7", nil])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("MGET").arg("team:count").arg("team:missing"),
|
||||
Ok(redis_test::redis_value!(["7", nil])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
|
||||
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
|
||||
MockCmd::new(redis::cmd("TTL").arg("team:missing"), Ok(-2i64)),
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("team:job-*")
|
||||
.arg("COUNT")
|
||||
.arg(25),
|
||||
Ok(redis_test::redis_value!(["4", ["team:job-a"]])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(4)
|
||||
.arg("MATCH")
|
||||
.arg("team:job-*")
|
||||
.arg("COUNT")
|
||||
.arg(25),
|
||||
Ok(redis_test::redis_value!(["0", ["team:job-b"]])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("DEL").arg("team:job-a").arg("team:job-b"),
|
||||
Ok(2u32),
|
||||
),
|
||||
MockCmd::with_values(
|
||||
sadd_pipeline,
|
||||
Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"),
|
||||
Ok(2u32),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("LPOP").arg("team:queue").arg(2usize),
|
||||
Ok(redis_test::redis_value!(["a", "b"])),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("EVAL")
|
||||
.arg("return KEYS[1]")
|
||||
.arg(1usize)
|
||||
.arg("team:key"),
|
||||
Ok("team:key"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")),
|
||||
MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")),
|
||||
MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK")),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.batch_get_counts(&["count".into(), "missing".into()])
|
||||
.unwrap(),
|
||||
[Some(7), None]
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_batch_get_counts(vec!["count".into(), "missing".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
[Some(7), None]
|
||||
);
|
||||
assert!(cache.sync_ping().unwrap());
|
||||
assert!(cache.ping().await.unwrap());
|
||||
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
|
||||
assert_eq!(
|
||||
cache.async_scan_iter("job-", 25).await.unwrap(),
|
||||
["team:job-a", "team:job-b"]
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.delete_cache_keys(vec!["job-a".into(), "job-b".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache_sadd("members", vec!["a".into(), "b".into()], None)
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_rpush("queue", vec!["a".into(), "b".into()])
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_lpop("queue", Some(2)).await.unwrap(),
|
||||
RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()])
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_eval("return KEYS[1]".into(), vec!["key".into()], Vec::new())
|
||||
.await
|
||||
.unwrap(),
|
||||
redis::Value::BulkString(b"team:key".to_vec())
|
||||
);
|
||||
assert_eq!(cache.client_list().unwrap(), "id=1");
|
||||
assert_eq!(cache.info().unwrap(), "redis_version:7");
|
||||
cache.flushall().unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_redis_pipelines_preserve_operation_order() {
|
||||
let mut rpush_pipeline = redis::pipe();
|
||||
rpush_pipeline
|
||||
.cmd("RPUSH")
|
||||
.arg("team:a")
|
||||
.arg("one")
|
||||
.cmd("RPUSH")
|
||||
.arg("team:b")
|
||||
.arg("two");
|
||||
let mut lpop_pipeline = redis::pipe();
|
||||
lpop_pipeline
|
||||
.cmd("LPOP")
|
||||
.arg("team:a")
|
||||
.arg(2usize)
|
||||
.cmd("LPOP")
|
||||
.arg("team:b");
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::with_values(
|
||||
rpush_pipeline,
|
||||
Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]),
|
||||
),
|
||||
MockCmd::with_values(
|
||||
lpop_pipeline,
|
||||
Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let queue = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
queue
|
||||
.async_rpush_pipeline(vec![
|
||||
RedisRpushOperation {
|
||||
key: "a".into(),
|
||||
values: vec![RedisArg::from("one")],
|
||||
},
|
||||
RedisRpushOperation {
|
||||
key: "b".into(),
|
||||
values: vec![RedisArg::from("two")],
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1, 2]
|
||||
);
|
||||
assert_eq!(
|
||||
queue
|
||||
.async_lpop_pipeline(vec![
|
||||
RedisLpopOperation {
|
||||
key: "a".into(),
|
||||
count: Some(2),
|
||||
},
|
||||
RedisLpopOperation {
|
||||
key: "b".into(),
|
||||
count: None,
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[
|
||||
RedisLpopResult::Values(vec![b"one".to_vec()]),
|
||||
RedisLpopResult::Missing,
|
||||
]
|
||||
);
|
||||
|
||||
let mut increment_pipeline = redis::pipe();
|
||||
increment_pipeline
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg("team:counter")
|
||||
.arg(1.5f64)
|
||||
.cmd("EXPIRE")
|
||||
.arg("team:counter")
|
||||
.arg(10u64)
|
||||
.ignore()
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg("team:counter")
|
||||
.arg(2.0f64);
|
||||
let connection = MockRedisConnection::new([MockCmd::with_values(
|
||||
increment_pipeline,
|
||||
Ok(vec![
|
||||
redis::Value::BulkString(b"1.5".to_vec()),
|
||||
redis::Value::Int(1),
|
||||
redis::Value::BulkString(b"3.5".to_vec()),
|
||||
]),
|
||||
)])
|
||||
.assert_all_commands_consumed();
|
||||
let counters = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
assert_eq!(
|
||||
counters
|
||||
.async_increment_pipeline(vec![
|
||||
IncrementOperation {
|
||||
key: "counter".into(),
|
||||
amount: 1.5,
|
||||
ttl: Some(Duration::from_secs(10)),
|
||||
},
|
||||
IncrementOperation {
|
||||
key: "counter".into(),
|
||||
amount: 2.0,
|
||||
ttl: None,
|
||||
},
|
||||
])
|
||||
.await
|
||||
.unwrap(),
|
||||
[1.5, 3.5]
|
||||
);
|
||||
}
|
||||
|
||||
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
|
||||
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
|
||||
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
|
||||
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return count"
|
||||
);
|
||||
const SET_MAX_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
|
||||
"redis.call('SET', KEYS[1], ARGV[1]); ",
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
|
||||
"return ARGV[1]; end; return current"
|
||||
);
|
||||
|
||||
#[tokio::test]
|
||||
async fn counter_repairs_are_atomic_and_use_default_ttl() {
|
||||
let floor = || {
|
||||
redis::cmd("EVAL")
|
||||
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg("team:counter")
|
||||
.arg(-2i64)
|
||||
.arg(30u64)
|
||||
.clone()
|
||||
};
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(floor(), Ok(0i64)),
|
||||
MockCmd::new(floor(), Ok(0i64)),
|
||||
MockCmd::new(
|
||||
redis::cmd("EVAL")
|
||||
.arg(SET_MAX_SCRIPT)
|
||||
.arg(1)
|
||||
.arg("team:counter")
|
||||
.arg(4.5f64)
|
||||
.arg(600u64),
|
||||
Ok("4.5"),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
|
||||
assert_eq!(
|
||||
cache
|
||||
.increment_with_floor("counter", -2, Duration::from_secs(30))
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_increment_with_floor("counter", -2, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_set_max("counter", 4.5, None).await.unwrap(),
|
||||
4.5
|
||||
);
|
||||
}
|
||||
|
||||
const CLAIM_SCRIPT: &str = concat!(
|
||||
"local current = redis.call('GET', KEYS[1]); ",
|
||||
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
|
||||
|
|
|
|||
|
|
@ -42,10 +42,12 @@ The resolver reads the namespace's `cache` attribute each time it resolves. A ca
|
|||
|
||||
Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend
|
||||
|
||||
The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and configuration changes before selecting native execution. It does not compare Redis connection settings. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python
|
||||
The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python
|
||||
|
||||
Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy
|
||||
|
||||
The Redis backend also provides the primitives needed to preserve its direct Python surface later: TLS URLs, ping, bulk delete, counter batches, TTL, scan, set membership, raw queue push and pop, queue and counter pipelines, counter floor and maximum operations, script evaluation, client information, namespaced flush, and full flush. These are backend operations only and are not exported to Python by this PR. Memory provides TTL, oldest-key, and counter-pipeline operations
|
||||
|
||||
## Adding another backend
|
||||
|
||||
Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python
|
||||
|
|
@ -56,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, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees
|
||||
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
|
||||
|
|
|
|||
|
|
@ -2,6 +2,13 @@ use std::future::Future;
|
|||
|
||||
use crate::{BaseCache, CacheKwargs, Error};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct IncrementOperation {
|
||||
pub key: String,
|
||||
pub amount: f64,
|
||||
pub ttl: Option<std::time::Duration>,
|
||||
}
|
||||
|
||||
pub trait CounterCache: BaseCache<Value = f64> {
|
||||
fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result<f64, Error>;
|
||||
|
||||
|
|
|
|||
2
litellm-rust/crates/cache/src/lib.rs
vendored
2
litellm-rust/crates/cache/src/lib.rs
vendored
|
|
@ -9,6 +9,6 @@ pub use base_cache::{
|
|||
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheKwargs,
|
||||
};
|
||||
pub use caching::{Cache, CacheBackend, get_cache, set_cache};
|
||||
pub use capabilities::{ClaimCache, CounterCache};
|
||||
pub use capabilities::{ClaimCache, CounterCache, IncrementOperation};
|
||||
pub use codec::{CacheCodec, JsonCodec};
|
||||
pub use error::Error;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue