mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
feat(cache): add native response cache and Python binding foundations
This commit is contained in:
Yujong Lee
parent
95cf7066d1
commit
081c93908f
18 changed files with 1702 additions and 50 deletions
100
litellm-rust/Cargo.lock
generated
100
litellm-rust/Cargo.lock
generated
|
|
@ -2486,6 +2486,21 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-response"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"py_literal",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-callbacks-legacy-python"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2649,6 +2664,8 @@ dependencies = [
|
|||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-cache",
|
||||
"litellm-cache-response",
|
||||
"litellm-callbacks-legacy-python",
|
||||
"litellm-core",
|
||||
"litellm-core-utils",
|
||||
|
|
@ -2660,6 +2677,7 @@ dependencies = [
|
|||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
|
|
@ -2948,6 +2966,16 @@ dependencies = [
|
|||
"minimal-lexical",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-bigint"
|
||||
version = "0.4.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367"
|
||||
dependencies = [
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-bigint"
|
||||
version = "0.5.1"
|
||||
|
|
@ -2958,6 +2986,15 @@ dependencies = [
|
|||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-complex"
|
||||
version = "0.4.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-conv"
|
||||
version = "0.2.2"
|
||||
|
|
@ -3131,6 +3168,48 @@ version = "2.3.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
|
||||
|
||||
[[package]]
|
||||
name = "pest"
|
||||
version = "2.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6d45aeb61b4bf818e12d4205f2466f8c4748f85f4fce0146d1c03d69d753f0ad"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
"ucd-trie",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pest_derive"
|
||||
version = "2.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "89cc5a242e25ed4e7704d0be240f2cfbe20a8c27e7e252d94835be93d92dc39f"
|
||||
dependencies = [
|
||||
"pest",
|
||||
"pest_generator",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pest_generator"
|
||||
version = "2.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7abf21475cc3820fe4b2ca2dc2142902f67a02189f3b5b3a229f4febc01a43e5"
|
||||
dependencies = [
|
||||
"pest",
|
||||
"pest_meta",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pest_meta"
|
||||
version = "2.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "adba4db388f687393c18c51348d44a41d870ca9df71a2c98172ea3035dc6936e"
|
||||
dependencies = [
|
||||
"pest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project"
|
||||
version = "1.1.13"
|
||||
|
|
@ -3305,6 +3384,19 @@ dependencies = [
|
|||
"prost",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "py_literal"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1"
|
||||
dependencies = [
|
||||
"num-bigint 0.4.8",
|
||||
"num-complex",
|
||||
"num-traits",
|
||||
"pest",
|
||||
"pest_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3"
|
||||
version = "0.29.2"
|
||||
|
|
@ -3604,7 +3696,7 @@ dependencies = [
|
|||
"arcstr",
|
||||
"combine",
|
||||
"itoa",
|
||||
"num-bigint",
|
||||
"num-bigint 0.5.1",
|
||||
"percent-encoding",
|
||||
"ryu",
|
||||
"sha1_smol",
|
||||
|
|
@ -4955,6 +5047,12 @@ dependencies = [
|
|||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ucd-trie"
|
||||
version = "0.1.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
|
||||
|
||||
[[package]]
|
||||
name = "unarray"
|
||||
version = "0.1.4"
|
||||
|
|
|
|||
|
|
@ -28,6 +28,8 @@ litellm-types = { path = "crates/types" }
|
|||
litellm-core-utils = { path = "crates/core-utils" }
|
||||
litellm-cache = { path = "crates/cache" }
|
||||
litellm-cache-memory = { path = "crates/cache-memory" }
|
||||
litellm-cache-redis = { path = "crates/cache-redis" }
|
||||
litellm-cache-response = { path = "crates/cache-response" }
|
||||
litellm-token-counter = { path = "crates/token-counter" }
|
||||
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
|
||||
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }
|
||||
|
|
|
|||
|
|
@ -7,12 +7,12 @@ use litellm_cache::{
|
|||
use redis::Commands;
|
||||
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
const KEY_PREFIX: &str = "litellm-cache:";
|
||||
|
||||
pub struct RedisCache<S, C = redis::Connection> {
|
||||
connection: Arc<Mutex<C>>,
|
||||
default_ttl: Duration,
|
||||
codec: S,
|
||||
namespace: Option<String>,
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> RedisCache<S> {
|
||||
|
|
@ -33,6 +33,7 @@ where
|
|||
connection: Arc::new(Mutex::new(connection)),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
codec,
|
||||
namespace: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -40,13 +41,44 @@ where
|
|||
self.connection.lock().map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn namespaced_key(key: &str) -> String {
|
||||
format!("{KEY_PREFIX}{key}")
|
||||
pub fn with_namespace(self, namespace: Option<String>) -> Self {
|
||||
Self {
|
||||
namespace: namespace.filter(|value| !value.is_empty()),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
fn namespaced_pattern() -> &'static str {
|
||||
const PATTERN: &str = "litellm-cache:*";
|
||||
PATTERN
|
||||
fn namespaced_key(&self, key: &str) -> String {
|
||||
match &self.namespace {
|
||||
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
|
||||
format!("{namespace}:{key}")
|
||||
}
|
||||
_ => key.into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn namespaced_pattern(&self) -> Result<String, Error> {
|
||||
let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
|
||||
let escaped: String = namespace
|
||||
.chars()
|
||||
.flat_map(|ch| {
|
||||
if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
|
||||
vec!['\\', ch]
|
||||
} else {
|
||||
vec![ch]
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Ok(format!("{escaped}:*"))
|
||||
}
|
||||
|
||||
fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
|
||||
match value {
|
||||
redis::Value::Nil => Ok(None),
|
||||
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
|
||||
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn ttl_seconds(ttl: Duration) -> u64 {
|
||||
|
|
@ -84,28 +116,29 @@ where
|
|||
let payload = self.codec.encode(&value)?;
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
self.connection()?
|
||||
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
|
||||
.set_ex::<_, _, ()>(self.namespaced_key(key), payload, ttl)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
|
||||
let bytes = self
|
||||
let value = self
|
||||
.connection()?
|
||||
.get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
|
||||
.get::<_, redis::Value>(self.namespaced_key(key))
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
bytes.map(|bytes| self.codec.decode(&bytes)).transpose()
|
||||
self.decode_response(value)
|
||||
}
|
||||
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.connection()?
|
||||
.del::<_, ()>(Self::namespaced_key(key))
|
||||
.del::<_, ()>(self.namespaced_key(key))
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
let pattern = self.namespaced_pattern()?;
|
||||
let mut connection = self.connection()?;
|
||||
let keys = connection
|
||||
.scan_match(Self::namespaced_pattern())
|
||||
.scan_match(pattern)
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.collect::<redis::RedisResult<Vec<String>>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
|
|
@ -125,7 +158,7 @@ where
|
|||
kwargs: CacheKwargs,
|
||||
) -> Result<(), Error> {
|
||||
let payload = self.codec.encode(&value)?;
|
||||
let key = Self::namespaced_key(key);
|
||||
let key = self.namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
connection
|
||||
|
|
@ -140,15 +173,14 @@ where
|
|||
key: &str,
|
||||
_: &CacheKwargs,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
let key = Self::namespaced_key(key);
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
let key = self.namespaced_key(key);
|
||||
let value = Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
connection
|
||||
.get::<_, Option<Vec<u8>>>(key)
|
||||
.get::<_, redis::Value>(key)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?
|
||||
.map(|bytes| self.codec.decode(&bytes))
|
||||
.transpose()
|
||||
.await?;
|
||||
self.decode_response(value)
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
|
|
@ -161,7 +193,7 @@ where
|
|||
.map(|(key, value)| {
|
||||
self.codec
|
||||
.encode(&value)
|
||||
.map(|payload| (Self::namespaced_key(&key), payload))
|
||||
.map(|payload| (self.namespaced_key(&key), payload))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
|
|
@ -177,7 +209,7 @@ where
|
|||
}
|
||||
|
||||
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let key = Self::namespaced_key(key);
|
||||
let key = self.namespaced_key(key);
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
|
|
@ -250,7 +282,8 @@ mod tests {
|
|||
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<CacheEntry>::new());
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<CacheEntry>::new())
|
||||
.with_namespace(Some("litellm-cache".into()));
|
||||
|
||||
cache
|
||||
.set_cache("key", value.clone(), CacheKwargs::default())
|
||||
|
|
@ -275,7 +308,8 @@ mod tests {
|
|||
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<CacheEntry>::new());
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<CacheEntry>::new())
|
||||
.with_namespace(Some("litellm-cache".into()));
|
||||
|
||||
cache.flush_cache().unwrap();
|
||||
}
|
||||
|
|
@ -284,7 +318,8 @@ mod tests {
|
|||
async fn test_connection_runs_ping_off_executor() {
|
||||
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<CacheEntry>::new());
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<CacheEntry>::new())
|
||||
.with_namespace(Some("litellm-cache".into()));
|
||||
|
||||
assert_eq!(
|
||||
cache.test_connection().await.unwrap().status,
|
||||
|
|
|
|||
|
|
@ -34,15 +34,12 @@ fn generic_helpers_use_the_injected_codec_and_ttl() {
|
|||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("litellm-cache:counter")
|
||||
.arg("counter")
|
||||
.arg(2)
|
||||
.arg([42u8, 7].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("litellm-cache:counter"),
|
||||
Ok(vec![42u8, 7]),
|
||||
),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
|
||||
|
|
@ -59,27 +56,21 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() {
|
|||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("litellm-cache:counter")
|
||||
.arg("counter")
|
||||
.arg(9)
|
||||
.arg([42u8, 7].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("litellm-cache:counter"),
|
||||
Ok(vec![42u8, 7]),
|
||||
),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("litellm-cache:batch")
|
||||
.arg("batch")
|
||||
.arg(2)
|
||||
.arg([42u8, 8].as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:counter"), Ok(1u32)),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("litellm-cache:counter"),
|
||||
Ok(redis::Value::Nil),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)),
|
||||
MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(
|
||||
|
|
@ -116,14 +107,8 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() {
|
|||
#[tokio::test]
|
||||
async fn codec_errors_propagate_without_writing_partial_batches() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("litellm-cache:invalid"),
|
||||
Ok(vec![99u8, 7]),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("litellm-cache:invalid"),
|
||||
Ok(vec![99u8, 7]),
|
||||
),
|
||||
MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
|
||||
MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
|
||||
|
|
@ -154,3 +139,48 @@ async fn codec_errors_propagate_without_writing_partial_batches() {
|
|||
Err(Error::InvalidEntry)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
|
||||
MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team".into()));
|
||||
assert_eq!(
|
||||
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.get_cache("team:key", &CacheKwargs::default())
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_requires_a_namespace_and_escapes_glob_metacharacters() {
|
||||
let unscoped = RedisCache::with_connection(
|
||||
MockRedisConnection::new([]).assert_all_commands_consumed(),
|
||||
None,
|
||||
JsonCodec::<String>::new(),
|
||||
);
|
||||
assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush));
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("team\\*:*"),
|
||||
Ok(redis_test::redis_value!(["0", ["team*:key"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let scoped = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
|
||||
.with_namespace(Some("team*".into()));
|
||||
scoped.flush_cache().unwrap();
|
||||
}
|
||||
|
|
|
|||
19
litellm-rust/crates/cache-response/Cargo.toml
Normal file
19
litellm-rust/crates/cache-response/Cargo.toml
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
[package]
|
||||
name = "litellm-cache-response"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
py_literal = "0.4.0"
|
||||
redis = "1.7.0"
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
redis-test = "1.0.4"
|
||||
tokio.workspace = true
|
||||
100
litellm-rust/crates/cache-response/src/codec.rs
Normal file
100
litellm-rust/crates/cache-response/src/codec.rs
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
use litellm_cache::{CacheCodec, CacheEntry, Error};
|
||||
use serde_json::Value;
|
||||
|
||||
pub struct ResponseCacheCodec;
|
||||
|
||||
impl CacheCodec for ResponseCacheCodec {
|
||||
type Value = CacheEntry;
|
||||
|
||||
fn encode(&self, value: &CacheEntry) -> Result<Vec<u8>, Error> {
|
||||
if !value.timestamp.is_finite() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn decode(&self, bytes: &[u8]) -> Result<CacheEntry, Error> {
|
||||
let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?;
|
||||
let entry: CacheEntry =
|
||||
serde_json::from_value(decode_value(text)?).map_err(|_| Error::InvalidEntry)?;
|
||||
if !entry.timestamp.is_finite() {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
Ok(entry)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn decode_value(text: &str) -> Result<Value, Error> {
|
||||
if let Ok(value) = serde_json::from_str(text) {
|
||||
return Ok(value);
|
||||
}
|
||||
check_literal_depth(text)?;
|
||||
let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?;
|
||||
literal_value(literal, 0)
|
||||
}
|
||||
|
||||
fn literal_value(value: py_literal::Value, depth: usize) -> Result<Value, Error> {
|
||||
use py_literal::Value as Literal;
|
||||
if depth > 128 {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
match value {
|
||||
Literal::String(text) => Ok(Value::String(text)),
|
||||
Literal::Boolean(value) => Ok(Value::Bool(value)),
|
||||
Literal::None => Ok(Value::Null),
|
||||
Literal::Integer(value) => {
|
||||
serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
Literal::Float(value) => serde_json::Number::from_f64(value)
|
||||
.map(Value::Number)
|
||||
.ok_or(Error::InvalidEntry),
|
||||
Literal::List(values) | Literal::Tuple(values) => values
|
||||
.into_iter()
|
||||
.map(|value| literal_value(value, depth + 1))
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Array),
|
||||
Literal::Dict(entries) => entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
let Literal::String(key) = key else {
|
||||
return Err(Error::InvalidEntry);
|
||||
};
|
||||
Ok((key, literal_value(value, depth + 1)?))
|
||||
})
|
||||
.collect::<Result<serde_json::Map<_, _>, _>>()
|
||||
.map(Value::Object),
|
||||
_ => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_literal_depth(text: &str) -> Result<(), Error> {
|
||||
let mut quote = None;
|
||||
let mut escaped = false;
|
||||
let mut depth = 0usize;
|
||||
for ch in text.chars() {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if let Some(delimiter) = quote {
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
} else if ch == delimiter {
|
||||
quote = None;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' | '"' => quote = Some(ch),
|
||||
'[' | '{' | '(' => {
|
||||
depth += 1;
|
||||
if depth > 128 {
|
||||
return Err(Error::InvalidEntry);
|
||||
}
|
||||
}
|
||||
']' | '}' | ')' => depth = depth.saturating_sub(1),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
7
litellm-rust/crates/cache-response/src/lib.rs
Normal file
7
litellm-rust/crates/cache-response/src/lib.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
mod codec;
|
||||
mod native;
|
||||
mod response;
|
||||
|
||||
pub use codec::ResponseCacheCodec;
|
||||
pub use native::NativeResponseCache;
|
||||
pub use response::{ResponseCache, ResponseCacheRequest};
|
||||
97
litellm-rust/crates/cache-response/src/native.rs
Normal file
97
litellm-rust/crates/cache-response/src/native.rs
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{CacheEntry, Error};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest};
|
||||
|
||||
pub enum NativeResponseCache<C = redis::Connection>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
Memory(Arc<ResponseCache<InMemoryCache<CacheEntry>>>),
|
||||
Redis(Arc<ResponseCache<RedisCache<ResponseCacheCodec, C>>>),
|
||||
}
|
||||
|
||||
impl<C: redis::ConnectionLike + Send + 'static> Clone for NativeResponseCache<C> {
|
||||
fn clone(&self) -> Self {
|
||||
match self {
|
||||
Self::Memory(cache) => Self::Memory(Arc::clone(cache)),
|
||||
Self::Redis(cache) => Self::Redis(Arc::clone(cache)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NativeResponseCache {
|
||||
pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
|
||||
Self::Memory(Arc::new(ResponseCache::new(Arc::new(
|
||||
InMemoryCache::response_cache(capacity, ttl, max_entry_bytes),
|
||||
))))
|
||||
}
|
||||
|
||||
pub fn redis(
|
||||
url: &str,
|
||||
ttl: Option<Duration>,
|
||||
namespace: Option<String>,
|
||||
) -> Result<Self, Error> {
|
||||
let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace);
|
||||
Ok(Self::Redis(Arc::new(ResponseCache::new(Arc::new(backend)))))
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: redis::ConnectionLike + Send + 'static> NativeResponseCache<C> {
|
||||
pub fn kind(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Memory(_) => "memory",
|
||||
Self::Redis(_) => "redis",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.lookup(request, now),
|
||||
Self::Redis(cache) => cache.lookup(request, now),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.store(request, response, now),
|
||||
Self::Redis(cache) => cache.store(request, response, now),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.async_store(request, response, now).await,
|
||||
Self::Redis(cache) => cache.async_store(request, response, now).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
124
litellm-rust/crates/cache-response/src/response.rs
Normal file
124
litellm-rust/crates/cache-response/src/response.rs
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheControls, CacheEntry, CacheKeyInput, CacheKwargs, Error, cache_key,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponseCacheRequest {
|
||||
pub key: CacheKeyInput,
|
||||
pub controls: CacheControls,
|
||||
pub kwargs: CacheKwargs,
|
||||
pub max_age: Option<Duration>,
|
||||
}
|
||||
|
||||
impl ResponseCacheRequest {
|
||||
pub fn new(key: CacheKeyInput) -> Self {
|
||||
Self {
|
||||
key,
|
||||
controls: CacheControls {
|
||||
configured: true,
|
||||
supported_call_type: true,
|
||||
native_backend: true,
|
||||
default_on: true,
|
||||
..Default::default()
|
||||
},
|
||||
kwargs: CacheKwargs::default(),
|
||||
max_age: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>> {
|
||||
backend: Arc<B>,
|
||||
}
|
||||
|
||||
impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
|
||||
pub fn new(backend: Arc<B>) -> Self {
|
||||
Self { backend }
|
||||
}
|
||||
|
||||
pub fn lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
if !request.controls.reads() {
|
||||
return Ok(None);
|
||||
}
|
||||
let entry = self
|
||||
.backend
|
||||
.get_cache(&cache_key(&request.key), &request.kwargs)?;
|
||||
Self::fresh_response(entry, now, request.max_age)
|
||||
}
|
||||
|
||||
pub async fn async_lookup(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
if !request.controls.reads() {
|
||||
return Ok(None);
|
||||
}
|
||||
let entry = self
|
||||
.backend
|
||||
.async_get_cache(&cache_key(&request.key), &request.kwargs)
|
||||
.await?;
|
||||
Self::fresh_response(entry, now, request.max_age)
|
||||
}
|
||||
|
||||
pub fn store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
if !request.controls.writes() {
|
||||
return Ok(());
|
||||
}
|
||||
self.backend.set_cache(
|
||||
&cache_key(&request.key),
|
||||
CacheEntry {
|
||||
timestamp: now.as_secs_f64(),
|
||||
response,
|
||||
},
|
||||
request.kwargs.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn async_store(
|
||||
&self,
|
||||
request: &ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
if !request.controls.writes() {
|
||||
return Ok(());
|
||||
}
|
||||
self.backend
|
||||
.async_set_cache(
|
||||
&cache_key(&request.key),
|
||||
CacheEntry {
|
||||
timestamp: now.as_secs_f64(),
|
||||
response,
|
||||
},
|
||||
request.kwargs.clone(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn fresh_response(
|
||||
entry: Option<CacheEntry>,
|
||||
now: Duration,
|
||||
max_age: Option<Duration>,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
entry
|
||||
.filter(|entry| entry.fresh(now, max_age))
|
||||
.map(|entry| match entry.response {
|
||||
Value::String(text) => crate::codec::decode_value(&text),
|
||||
value => Ok(value),
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
}
|
||||
290
litellm-rust/crates/cache-response/tests/response.rs
Normal file
290
litellm-rust/crates/cache-response/tests/response.rs
Normal file
|
|
@ -0,0 +1,290 @@
|
|||
use std::{
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKeyField, CacheKeyInput, Error};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use litellm_cache_response::{
|
||||
NativeResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
|
||||
};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use serde_json::json;
|
||||
|
||||
fn request() -> ResponseCacheRequest {
|
||||
ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some("tenant:key".into()),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
|
||||
let clock = Arc::new(AtomicU64::new(100));
|
||||
let backend = Arc::new(InMemoryCache::with_clock(
|
||||
Some(8),
|
||||
Some(Duration::from_secs(600)),
|
||||
{
|
||||
let clock = clock.clone();
|
||||
move || Duration::from_secs(clock.load(Ordering::SeqCst))
|
||||
},
|
||||
));
|
||||
let cache = ResponseCache::new(backend.clone());
|
||||
let mut request = request();
|
||||
request.kwargs.ttl = Some(Duration::from_secs(10));
|
||||
request.max_age = Some(Duration::from_secs(5));
|
||||
cache
|
||||
.store(
|
||||
&request,
|
||||
json!({"choices": [1], "usage": {"total_tokens": 7}}),
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
backend.expires_at("tenant:key").unwrap(),
|
||||
Some(Duration::from_secs(110))
|
||||
);
|
||||
assert!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(105))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(106)).unwrap(),
|
||||
None
|
||||
);
|
||||
request.max_age = None;
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&request, Duration::from_secs(106))
|
||||
.unwrap()
|
||||
.unwrap()["usage"]["total_tokens"],
|
||||
7
|
||||
);
|
||||
clock.store(111, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(111))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
cache
|
||||
.async_store(&request, json!({"choices": [2]}), Duration::from_secs(111))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
|
||||
Some(json!({"choices": [2]}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn directives_skip_io_and_keep_reads_and_writes_independent() {
|
||||
let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024);
|
||||
let mut request = request();
|
||||
let now = Duration::from_secs(100);
|
||||
request.controls.no_store = true;
|
||||
cache
|
||||
.async_store(&request, json!({"v": 1}), now)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
request.controls.no_store = false;
|
||||
request.controls.no_cache = true;
|
||||
cache.store(&request, json!({"v": 2}), now).unwrap();
|
||||
assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None);
|
||||
request.controls.no_cache = false;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
|
||||
request.controls.default_on = false;
|
||||
cache.store(&request, json!({"v": 3}), now).unwrap();
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
request.controls.use_cache = true;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
|
||||
request.controls.supported_call_type = false;
|
||||
assert_eq!(cache.lookup(&request, now).unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_json() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()),
|
||||
),
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("tenant:key")
|
||||
.arg(600)
|
||||
.arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()),
|
||||
Ok("OK"),
|
||||
),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec)
|
||||
.with_namespace(Some("tenant".into()));
|
||||
let cache = NativeResponseCache::Redis(Arc::new(ResponseCache::new(Arc::new(backend))));
|
||||
let request = request();
|
||||
let expected = json!({"ok": true, "text": "cached"});
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(101)).unwrap(),
|
||||
Some(expected.clone())
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::from_secs(101))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(expected.clone())
|
||||
);
|
||||
cache
|
||||
.async_store(&request, expected, Duration::from_secs(100))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn captured_enum_keeps_the_selected_backend_for_background_writes() {
|
||||
let original = NativeResponseCache::memory(8, Duration::from_secs(600), 1024);
|
||||
let captured = original.clone();
|
||||
let replacement = NativeResponseCache::memory(8, Duration::from_secs(600), 1024);
|
||||
let request = request();
|
||||
let writer = tokio::spawn({
|
||||
let request = request.clone();
|
||||
async move {
|
||||
captured
|
||||
.async_store(
|
||||
&request,
|
||||
json!({"selected": "original"}),
|
||||
Duration::from_secs(100),
|
||||
)
|
||||
.await
|
||||
}
|
||||
});
|
||||
writer.await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
original.lookup(&request, Duration::from_secs(100)).unwrap(),
|
||||
Some(json!({"selected":"original"}))
|
||||
);
|
||||
assert_eq!(
|
||||
replacement
|
||||
.lookup(&request, Duration::from_secs(100))
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generated_keys_preserve_namespace_and_explicit_keys() {
|
||||
let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024);
|
||||
let key = CacheKeyInput {
|
||||
fields: vec![CacheKeyField {
|
||||
name: "model".into(),
|
||||
value: Some("a".into()),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
}],
|
||||
namespace: Some("tenant".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let generated = ResponseCacheRequest::new(key.clone());
|
||||
let explicit = ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some(litellm_cache::cache_key(&key)),
|
||||
..Default::default()
|
||||
});
|
||||
cache
|
||||
.store(&generated, json!({"value": 7}), Duration::from_secs(100))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&explicit, Duration::from_secs(100)).unwrap(),
|
||||
Some(json!({"value":7}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_codec_accepts_python_literals_without_executing_code() {
|
||||
let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#;
|
||||
let entry = ResponseCacheCodec.decode(bytes).unwrap();
|
||||
assert_eq!(
|
||||
entry.response,
|
||||
json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]})
|
||||
);
|
||||
for bytes in [
|
||||
b"__import__('os').system('false')".as_slice(),
|
||||
b"{'timestamp': 'invalid', 'response': {}}",
|
||||
b"{'timestamp': 1e9999, 'response': {}}",
|
||||
] {
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(bytes).unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000));
|
||||
assert_eq!(
|
||||
ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
assert_eq!(
|
||||
ResponseCacheCodec
|
||||
.encode(&CacheEntry {
|
||||
timestamp: f64::NAN,
|
||||
response: json!({})
|
||||
})
|
||||
.unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redis() {
|
||||
let connection = MockRedisConnection::new([MockCmd::new(
|
||||
redis::cmd("GET").arg("tenant:key"),
|
||||
Ok(b"invalid".to_vec()),
|
||||
)])
|
||||
.assert_all_commands_consumed();
|
||||
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec);
|
||||
let cache = ResponseCache::new(Arc::new(backend));
|
||||
let mut request = request();
|
||||
request.controls.no_cache = true;
|
||||
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
|
||||
request.controls.no_cache = false;
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_lookup(&request, Duration::ZERO)
|
||||
.await
|
||||
.unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_memory_entries_are_rejected_by_the_response_consumer() {
|
||||
let backend = Arc::new(InMemoryCache::default());
|
||||
BaseCache::set_cache(
|
||||
backend.as_ref(),
|
||||
"tenant:key",
|
||||
CacheEntry {
|
||||
timestamp: 100.0,
|
||||
response: json!("not a serialized response"),
|
||||
},
|
||||
Default::default(),
|
||||
)
|
||||
.unwrap();
|
||||
let cache = ResponseCache::new(backend);
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&request(), Duration::from_secs(100))
|
||||
.unwrap_err(),
|
||||
Error::InvalidEntry
|
||||
);
|
||||
}
|
||||
1
litellm-rust/crates/cache/src/caching.rs
vendored
1
litellm-rust/crates/cache/src/caching.rs
vendored
|
|
@ -27,6 +27,7 @@ pub struct CacheKeyField {
|
|||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct CacheKeyInput {
|
||||
pub fields: Vec<CacheKeyField>,
|
||||
pub preset: Option<String>,
|
||||
|
|
|
|||
2
litellm-rust/crates/cache/src/error.rs
vendored
2
litellm-rust/crates/cache/src/error.rs
vendored
|
|
@ -4,4 +4,6 @@ pub enum Error {
|
|||
Unavailable,
|
||||
#[error("invalid cache entry")]
|
||||
InvalidEntry,
|
||||
#[error("flushing Redis requires an explicit namespace")]
|
||||
UnscopedFlush,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,9 @@ tiktoken = ["litellm-token-counter/tiktoken"]
|
|||
|
||||
[dependencies]
|
||||
bytes.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
serde.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
litellm-callbacks-legacy-python.workspace = true
|
||||
litellm-core.workspace = true
|
||||
|
|
|
|||
210
litellm-rust/crates/python-bridge/src/cache/facade.rs
vendored
Normal file
210
litellm-rust/crates/python-bridge/src/cache/facade.rs
vendored
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
use litellm_cache_response::NativeResponseCache;
|
||||
use litellm_host_python::from_py;
|
||||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::PyTypeError,
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple, PyType},
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::NativeCacheHandle;
|
||||
|
||||
struct ClassGuard {
|
||||
class: Py<PyType>,
|
||||
attributes: Vec<(String, Py<PyAny>)>,
|
||||
}
|
||||
|
||||
struct ObjectGuard {
|
||||
reference: Py<PyAny>,
|
||||
classes: Vec<ClassGuard>,
|
||||
config_names: &'static [&'static str],
|
||||
config: Vec<Value>,
|
||||
}
|
||||
|
||||
pub(super) struct FacadeGuard {
|
||||
outer: ObjectGuard,
|
||||
backend: ObjectGuard,
|
||||
}
|
||||
|
||||
impl ObjectGuard {
|
||||
fn capture(
|
||||
py: Python<'_>,
|
||||
object: &Bound<'_, PyAny>,
|
||||
config_names: &'static [&'static str],
|
||||
) -> PyResult<Self> {
|
||||
let classes = object
|
||||
.get_type()
|
||||
.getattr("__mro__")?
|
||||
.cast_into::<PyTuple>()?
|
||||
.iter()
|
||||
.map(|class| {
|
||||
let class = class.cast_into::<PyType>()?;
|
||||
let attributes = class
|
||||
.getattr("__dict__")?
|
||||
.call_method0("items")?
|
||||
.try_iter()?
|
||||
.map(|item| item?.extract::<(String, Py<PyAny>)>())
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
Ok(ClassGuard {
|
||||
class: class.unbind(),
|
||||
attributes,
|
||||
})
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
let guard = Self {
|
||||
reference: py
|
||||
.import("weakref")?
|
||||
.getattr("ref")?
|
||||
.call1((object,))?
|
||||
.unbind(),
|
||||
classes,
|
||||
config_names,
|
||||
config: Self::config(object, config_names)?,
|
||||
};
|
||||
if !guard.matches(py, object)? {
|
||||
return Err(PyTypeError::new_err(
|
||||
"native facade registration requires unmodified built-in methods",
|
||||
));
|
||||
}
|
||||
Ok(guard)
|
||||
}
|
||||
|
||||
fn config(object: &Bound<'_, PyAny>, names: &[&str]) -> PyResult<Vec<Value>> {
|
||||
names
|
||||
.iter()
|
||||
.map(|name| match object.getattr(*name) {
|
||||
Ok(value) => from_py(&value),
|
||||
Err(error)
|
||||
if error.is_instance_of::<pyo3::exceptions::PyAttributeError>(object.py()) =>
|
||||
{
|
||||
Ok(Value::Null)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn matches(&self, py: Python<'_>, object: &Bound<'_, PyAny>) -> PyResult<bool> {
|
||||
if !self.reference.bind(py).call0()?.is(object) {
|
||||
return Ok(false);
|
||||
}
|
||||
let mro = object
|
||||
.get_type()
|
||||
.getattr("__mro__")?
|
||||
.cast_into::<PyTuple>()?;
|
||||
if mro.len() != self.classes.len() {
|
||||
return Ok(false);
|
||||
}
|
||||
let instance = object.getattr("__dict__")?.cast_into::<PyDict>()?;
|
||||
for (class, expected) in mro.iter().zip(&self.classes) {
|
||||
if !class.is(expected.class.bind(py)) {
|
||||
return Ok(false);
|
||||
}
|
||||
let attributes = class.getattr("__dict__")?;
|
||||
if attributes.len()? != expected.attributes.len() {
|
||||
return Ok(false);
|
||||
}
|
||||
for (name, value) in &expected.attributes {
|
||||
if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Self::config(object, self.config_names)? == self.config)
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.reference)?;
|
||||
for class in &self.classes {
|
||||
visit.call(&class.class)?;
|
||||
for (_, value) in &class.attributes {
|
||||
visit.call(value)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl FacadeGuard {
|
||||
pub(super) fn capture(py: Python<'_>, facade: &Bound<'_, PyAny>, kind: &str) -> PyResult<Self> {
|
||||
let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?;
|
||||
if !facade.get_type().is(&cache_type) {
|
||||
return Err(PyTypeError::new_err(
|
||||
"only exact built-in Cache facades can be registered",
|
||||
));
|
||||
}
|
||||
let (module, name, cache_kind) = match kind {
|
||||
"memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
|
||||
"redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
let backend = facade.getattr("cache")?;
|
||||
if facade.getattr("type")?.extract::<String>()? != cache_kind
|
||||
|| !backend.get_type().is(&py.import(module)?.getattr(name)?)
|
||||
{
|
||||
return Err(PyTypeError::new_err(
|
||||
"facade and native backend types must match",
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
outer: ObjectGuard::capture(
|
||||
py,
|
||||
facade,
|
||||
&[
|
||||
"type",
|
||||
"mode",
|
||||
"ttl",
|
||||
"namespace",
|
||||
"supported_call_types",
|
||||
"redis_flush_size",
|
||||
],
|
||||
)?,
|
||||
backend: ObjectGuard::capture(
|
||||
py,
|
||||
&backend,
|
||||
&[
|
||||
"namespace",
|
||||
"default_ttl",
|
||||
"max_size_in_memory",
|
||||
"max_size_per_item",
|
||||
],
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<bool> {
|
||||
Ok(self.outer.matches(py, facade)?
|
||||
&& self.backend.matches(py, &facade.getattr("cache")?)?)
|
||||
}
|
||||
|
||||
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
self.outer.traverse(&visit)?;
|
||||
self.backend.traverse(&visit)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn resolve(
|
||||
py: Python<'_>,
|
||||
facade: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Option<NativeResponseCache>> {
|
||||
let Ok(dict) = facade
|
||||
.getattr("__dict__")
|
||||
.and_then(|dict| dict.cast_into::<PyDict>().map_err(Into::into))
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(handle) = dict.get_item("_native_cache_handle")? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Ok(handle) = handle.extract::<PyRef<'_, NativeCacheHandle>>() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(guard) = &handle.guard else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !guard.matches(py, facade).unwrap_or(false) {
|
||||
return Ok(None);
|
||||
}
|
||||
handle.service().map(Some)
|
||||
}
|
||||
350
litellm-rust/crates/python-bridge/src/cache/mod.rs
vendored
Normal file
350
litellm-rust/crates/python-bridge/src/cache/mod.rs
vendored
Normal file
|
|
@ -0,0 +1,350 @@
|
|||
mod facade;
|
||||
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_cache::{CacheControls, CacheKeyInput, Error};
|
||||
use litellm_cache_response::{NativeResponseCache, ResponseCacheRequest};
|
||||
use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py};
|
||||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::{PyRuntimeError, PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
types::PyDict,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use facade::FacadeGuard;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct RequestInput {
|
||||
key: CacheKeyInput,
|
||||
controls: Option<CacheControls>,
|
||||
ttl_seconds: Option<f64>,
|
||||
max_age_seconds: Option<f64>,
|
||||
}
|
||||
|
||||
fn request(value: &Bound<'_, PyAny>) -> PyResult<ResponseCacheRequest> {
|
||||
let input: RequestInput = from_py(value)?;
|
||||
let mut request = ResponseCacheRequest::new(input.key);
|
||||
if let Some(controls) = input.controls {
|
||||
request.controls = controls;
|
||||
}
|
||||
request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?;
|
||||
request.max_age = input.max_age_seconds.map(duration).transpose()?;
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
fn duration(seconds: f64) -> PyResult<Duration> {
|
||||
Duration::try_from_secs_f64(seconds)
|
||||
.map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative"))
|
||||
}
|
||||
|
||||
fn now() -> Duration {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn cache_error(error: Error) -> PyErr {
|
||||
match error {
|
||||
Error::InvalidEntry => PyValueError::new_err(error.to_string()),
|
||||
_ => PyRuntimeError::new_err(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(frozen)]
|
||||
pub(crate) struct NativeCacheHandle {
|
||||
service: NativeResponseCache,
|
||||
guard: Option<FacadeGuard>,
|
||||
pid: u32,
|
||||
}
|
||||
|
||||
impl NativeCacheHandle {
|
||||
fn service(&self) -> PyResult<NativeResponseCache> {
|
||||
if self.pid != std::process::id() {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"native cache handles must be recreated after fork",
|
||||
));
|
||||
}
|
||||
Ok(self.service.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeCacheHandle {
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))]
|
||||
fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes),
|
||||
guard: None,
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))]
|
||||
fn redis(
|
||||
py: Python<'_>,
|
||||
url: String,
|
||||
ttl_seconds: Option<f64>,
|
||||
namespace: Option<String>,
|
||||
) -> PyResult<Self> {
|
||||
let ttl = ttl_seconds.map(duration).transpose()?;
|
||||
let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace))
|
||||
.map_err(cache_error)?;
|
||||
Ok(Self {
|
||||
service,
|
||||
guard: None,
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn backend(&self) -> &'static str {
|
||||
self.service.kind()
|
||||
}
|
||||
|
||||
fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
|
||||
let service = self.service()?;
|
||||
let guard = FacadeGuard::capture(py, facade, self.backend())?;
|
||||
let handle = Py::new(
|
||||
py,
|
||||
Self {
|
||||
service,
|
||||
guard: Some(guard),
|
||||
pid: self.pid,
|
||||
},
|
||||
)?;
|
||||
facade.setattr("_native_cache_handle", handle)
|
||||
}
|
||||
|
||||
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
if let Some(guard) = &self.guard {
|
||||
guard.traverse(visit)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
enum CacheBinding {
|
||||
Disabled,
|
||||
Native(NativeResponseCache),
|
||||
PythonCallback(Py<PyAny>),
|
||||
}
|
||||
|
||||
#[pyclass(frozen, name = "CacheBinding")]
|
||||
pub(crate) struct ResolvedCache {
|
||||
binding: CacheBinding,
|
||||
pid: u32,
|
||||
}
|
||||
|
||||
impl ResolvedCache {
|
||||
fn check_process(&self) -> PyResult<()> {
|
||||
if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"native cache bindings must be resolved again after fork",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn lookup_step(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
input: &Bound<'_, PyAny>,
|
||||
kwargs: Option<&Bound<'_, PyDict>>,
|
||||
) -> PyResult<ExecutionStep> {
|
||||
self.check_process()?;
|
||||
let awaitable = match &self.binding {
|
||||
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,
|
||||
)?
|
||||
}
|
||||
CacheBinding::PythonCallback(object) => object.bind(py).call_method(
|
||||
"async_get_cache",
|
||||
(),
|
||||
Some(callback_kwargs(kwargs)?),
|
||||
)?,
|
||||
};
|
||||
Ok(ExecutionStep::Await(awaitable.unbind()))
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl ResolvedCache {
|
||||
#[getter]
|
||||
fn kind(&self) -> &'static str {
|
||||
match self.binding {
|
||||
CacheBinding::Disabled => "disabled",
|
||||
CacheBinding::Native(_) => "native",
|
||||
CacheBinding::PythonCallback(_) => "python_callback",
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request, *, callback_kwargs=None))]
|
||||
fn lookup(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
request: &Bound<'_, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'_, PyDict>>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => Ok(py.None()),
|
||||
CacheBinding::Native(service) => {
|
||||
let request = self::request(request)?;
|
||||
let service = service.clone();
|
||||
let response = release_gil(py, move || service.lookup(&request, now()))
|
||||
.map_err(cache_error)?;
|
||||
to_py(py, &response)
|
||||
}
|
||||
CacheBinding::PythonCallback(object) => object
|
||||
.bind(py)
|
||||
.call_method(
|
||||
"get_cache",
|
||||
(),
|
||||
Some(self::callback_kwargs(callback_kwargs)?),
|
||||
)
|
||||
.map(Bound::unbind),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request, response, *, callback_kwargs=None))]
|
||||
fn store(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
request: &Bound<'_, PyAny>,
|
||||
response: &Bound<'_, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'_, PyDict>>,
|
||||
) -> PyResult<()> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => Ok(()),
|
||||
CacheBinding::Native(service) => {
|
||||
let request = self::request(request)?;
|
||||
let response: Value = from_py(response)?;
|
||||
let service = service.clone();
|
||||
release_gil(py, move || service.store(&request, response, now()))
|
||||
.map_err(cache_error)
|
||||
}
|
||||
CacheBinding::PythonCallback(object) => object
|
||||
.bind(py)
|
||||
.call_method(
|
||||
"add_cache",
|
||||
(response,),
|
||||
Some(self::callback_kwargs(callback_kwargs)?),
|
||||
)
|
||||
.map(|_| ()),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request, *, callback_kwargs=None))]
|
||||
fn async_lookup<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
request: &Bound<'py, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)?
|
||||
else {
|
||||
unreachable!()
|
||||
};
|
||||
Ok(awaitable.into_bound(py))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request, response, *, callback_kwargs=None))]
|
||||
fn async_store<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
request: &Bound<'py, PyAny>,
|
||||
response: &Bound<'py, PyAny>,
|
||||
callback_kwargs: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.check_process()?;
|
||||
match &self.binding {
|
||||
CacheBinding::Disabled => ready_none(py),
|
||||
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,
|
||||
)
|
||||
}
|
||||
CacheBinding::PythonCallback(object) => object.bind(py).call_method(
|
||||
"async_add_cache",
|
||||
(response,),
|
||||
Some(self::callback_kwargs(callback_kwargs)?),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
if let CacheBinding::PythonCallback(object) = &self.binding {
|
||||
visit.call(object)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn callback_kwargs<'a, 'py>(
|
||||
kwargs: Option<&'a Bound<'py, PyDict>>,
|
||||
) -> PyResult<&'a Bound<'py, PyDict>> {
|
||||
kwargs.ok_or_else(|| {
|
||||
PyTypeError::new_err("Python cache callbacks require their original callback_kwargs")
|
||||
})
|
||||
}
|
||||
|
||||
fn ready_none(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let future = py
|
||||
.import("asyncio")?
|
||||
.call_method0("get_running_loop")?
|
||||
.call_method0("create_future")?;
|
||||
future.call_method1("set_result", (py.None(),))?;
|
||||
Ok(future)
|
||||
}
|
||||
|
||||
#[pyclass(frozen)]
|
||||
pub(crate) struct CacheResolver {
|
||||
namespace: Py<PyAny>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl CacheResolver {
|
||||
#[new]
|
||||
fn new(namespace: Py<PyAny>) -> Self {
|
||||
Self { namespace }
|
||||
}
|
||||
|
||||
pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult<ResolvedCache> {
|
||||
let object = self.namespace.bind(py).getattr("cache")?;
|
||||
let binding = if object.is_none() {
|
||||
CacheBinding::Disabled
|
||||
} else if let Ok(handle) = object.extract::<PyRef<'_, NativeCacheHandle>>() {
|
||||
CacheBinding::Native(handle.service()?)
|
||||
} else if let Some(service) = facade::resolve(py, &object)? {
|
||||
CacheBinding::Native(service)
|
||||
} else {
|
||||
CacheBinding::PythonCallback(object.unbind())
|
||||
};
|
||||
Ok(ResolvedCache {
|
||||
binding,
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.namespace)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
mod cache;
|
||||
mod credentials;
|
||||
mod diagnostics;
|
||||
mod errors;
|
||||
|
|
@ -9,6 +10,8 @@ mod token_counter;
|
|||
|
||||
#[pymodule(gil_used = true)]
|
||||
mod _native {
|
||||
#[pymodule_export]
|
||||
use crate::cache::{CacheResolver, NativeCacheHandle, ResolvedCache};
|
||||
#[cfg(feature = "panic-test")]
|
||||
#[pymodule_export]
|
||||
use crate::diagnostics::_panic_for_test;
|
||||
|
|
@ -65,6 +68,9 @@ mod tests {
|
|||
"achat_completions",
|
||||
"ResponsesWebSocketConnection",
|
||||
"TokenCounter",
|
||||
"CacheResolver",
|
||||
"NativeCacheHandle",
|
||||
"CacheBinding",
|
||||
"gil_stats",
|
||||
"process_state_started",
|
||||
"reserve_process_for_forking",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from asyncio import Future
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
|
||||
from collections.abc import AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence
|
||||
from typing import Never, final
|
||||
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
|
|
@ -93,6 +93,50 @@ class ResponsesWebSocketConnection:
|
|||
def recv_text(self) -> Future[str | None]: ...
|
||||
def close(self) -> Future[None]: ...
|
||||
|
||||
@final
|
||||
class NativeCacheHandle:
|
||||
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
|
||||
@staticmethod
|
||||
def memory(
|
||||
*, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576
|
||||
) -> NativeCacheHandle: ...
|
||||
@staticmethod
|
||||
def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> NativeCacheHandle: ...
|
||||
@property
|
||||
def backend(self) -> str: ...
|
||||
def bind_facade(self, facade: object) -> None: ...
|
||||
|
||||
@final
|
||||
class CacheResolver:
|
||||
def __new__(cls, namespace: object) -> CacheResolver: ...
|
||||
def resolve(self) -> CacheBinding: ...
|
||||
|
||||
@final
|
||||
class CacheBinding:
|
||||
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
|
||||
@property
|
||||
def kind(self) -> str: ...
|
||||
def lookup(
|
||||
self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None
|
||||
) -> object: ...
|
||||
def store(
|
||||
self,
|
||||
request: Mapping[str, object] | None,
|
||||
response: object,
|
||||
*,
|
||||
callback_kwargs: dict[str, object] | None = None,
|
||||
) -> None: ...
|
||||
def async_lookup(
|
||||
self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None
|
||||
) -> Awaitable[object]: ...
|
||||
def async_store(
|
||||
self,
|
||||
request: Mapping[str, object] | None,
|
||||
response: object,
|
||||
*,
|
||||
callback_kwargs: dict[str, object] | None = None,
|
||||
) -> Awaitable[object]: ...
|
||||
|
||||
@final
|
||||
class TokenCounter:
|
||||
def __new__(cls, tokenizer_json: str) -> TokenCounter: ...
|
||||
|
|
@ -109,7 +153,10 @@ def process_state_started() -> bool: ...
|
|||
def reserve_process_for_forking() -> None: ...
|
||||
|
||||
__all__ = [
|
||||
"CacheBinding",
|
||||
"CacheResolver",
|
||||
"ForkedAfterNativeRuntimeStarted",
|
||||
"NativeCacheHandle",
|
||||
"ProcessReservedForForking",
|
||||
"ResponsesWebSocketConnection",
|
||||
"RustBridgeDeclined",
|
||||
|
|
|
|||
231
tests/test_litellm_rust/test_cache.py
Normal file
231
tests/test_litellm_rust/test_cache.py
Normal file
|
|
@ -0,0 +1,231 @@
|
|||
import asyncio
|
||||
import contextvars
|
||||
import gc
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import Generator
|
||||
from types import SimpleNamespace
|
||||
from typing import Final, Protocol, cast
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
import redis
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.rust_bridge import _native
|
||||
from litellm.types.caching import LiteLLMCacheType
|
||||
from tests.test_litellm_rust.support.isolation import rebound
|
||||
|
||||
pytestmark: Final = pytest.mark.requires_rust_extension
|
||||
|
||||
|
||||
class CacheLookup(Protocol):
|
||||
def get_cache(self, **kwargs: object) -> object: ...
|
||||
|
||||
|
||||
def request(key: str = "key") -> dict[str, object]:
|
||||
return {"key": {"preset": key}}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_url() -> Generator[str]:
|
||||
server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis")
|
||||
worker: Final = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
worker.start()
|
||||
try:
|
||||
yield f"redis://127.0.0.1:{server.server_address[1]}"
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
worker.join(timeout=5)
|
||||
|
||||
|
||||
def test_existing_constructor_and_global_are_unchanged() -> None:
|
||||
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
|
||||
assert type(facade.cache) is InMemoryCache
|
||||
assert "_native_cache_handle" not in vars(facade)
|
||||
with rebound(litellm, "cache", facade):
|
||||
resolver: Final = _native.CacheResolver(litellm)
|
||||
assert resolver.resolve().kind == "python_callback"
|
||||
resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"})
|
||||
assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7}
|
||||
|
||||
|
||||
async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None:
|
||||
namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.memory())
|
||||
resolver: Final = _native.CacheResolver(namespace)
|
||||
selected: Final = resolver.resolve()
|
||||
assert selected.kind == "native"
|
||||
selected.store(request(), {"answer": 1})
|
||||
assert await selected.async_lookup(request()) == {"answer": 1}
|
||||
with rebound(namespace, "cache", _native.NativeCacheHandle.memory()):
|
||||
replacement: Final = resolver.resolve()
|
||||
await selected.async_store(request(), {"answer": 2})
|
||||
assert replacement.lookup(request()) is None
|
||||
assert selected.lookup(request()) == {"answer": 2}
|
||||
with rebound(namespace, "cache", None):
|
||||
disabled: Final = resolver.resolve()
|
||||
assert disabled.kind == "disabled"
|
||||
assert disabled.lookup(None) is None
|
||||
await disabled.async_store(None, object())
|
||||
assert await disabled.async_lookup(None) is None
|
||||
assert selected.lookup(request()) == {"answer": 2}
|
||||
|
||||
|
||||
async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None:
|
||||
context: Final = contextvars.ContextVar("cache_context", default="caller")
|
||||
caller: Final = asyncio.current_task()
|
||||
sentinel: Final = object()
|
||||
failure: Final = RuntimeError("callback failed")
|
||||
|
||||
class CustomCache:
|
||||
async def async_get_cache(self, *, marker: object) -> object:
|
||||
assert marker is sentinel
|
||||
assert asyncio.current_task() is caller
|
||||
context.set("callback")
|
||||
return marker
|
||||
|
||||
async def async_add_cache(self, response: object, *, marker: object) -> None:
|
||||
assert response is sentinel
|
||||
assert marker is sentinel
|
||||
raise failure
|
||||
|
||||
namespace: Final = SimpleNamespace(cache=CustomCache())
|
||||
binding: Final = _native.CacheResolver(namespace).resolve()
|
||||
assert binding.kind == "python_callback"
|
||||
assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel
|
||||
assert context.get() == "callback"
|
||||
with pytest.raises(RuntimeError) as caught:
|
||||
await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel})
|
||||
assert caught.value is failure
|
||||
|
||||
|
||||
async def test_callback_cancellation_stays_in_the_callers_task() -> None:
|
||||
entered: Final = asyncio.Event()
|
||||
finished: Final = asyncio.Event()
|
||||
|
||||
class CustomCache:
|
||||
async def async_get_cache(self) -> None:
|
||||
entered.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
finally:
|
||||
finished.set()
|
||||
|
||||
binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve()
|
||||
|
||||
async def lookup() -> object:
|
||||
return await binding.async_lookup(None, callback_kwargs={})
|
||||
|
||||
task: Final = asyncio.create_task(lookup())
|
||||
await entered.wait()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert finished.is_set()
|
||||
|
||||
|
||||
def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None:
|
||||
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
|
||||
handle: Final = _native.NativeCacheHandle.memory()
|
||||
handle.bind_facade(facade)
|
||||
resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade))
|
||||
native: Final = resolver.resolve()
|
||||
assert native.kind == "native"
|
||||
native.store(request(), {"source": "native"})
|
||||
assert native.lookup(request()) == {"source": "native"}
|
||||
assert cast(CacheLookup, facade).get_cache(cache_key="key") is None
|
||||
sentinel: Final = object()
|
||||
|
||||
def outer_override(**_kwargs: object) -> object:
|
||||
return sentinel
|
||||
|
||||
def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]:
|
||||
return {"source": "override"}
|
||||
|
||||
with rebound(facade, "get_cache", outer_override):
|
||||
fallback: Final = resolver.resolve()
|
||||
assert fallback.kind == "python_callback"
|
||||
assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel
|
||||
assert resolver.resolve().kind == "python_callback"
|
||||
delattr(facade, "get_cache")
|
||||
assert resolver.resolve().kind == "native"
|
||||
with rebound(facade.cache, "get_cache", backend_override):
|
||||
backend_fallback: Final = resolver.resolve()
|
||||
assert backend_fallback.kind == "python_callback"
|
||||
assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"}
|
||||
|
||||
|
||||
def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None:
|
||||
class CustomCache(Cache):
|
||||
pass
|
||||
|
||||
handle: Final = _native.NativeCacheHandle.memory()
|
||||
with pytest.raises(TypeError):
|
||||
handle.bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL))
|
||||
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
|
||||
handle.bind_facade(facade)
|
||||
resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade))
|
||||
with rebound(facade, "cache", InMemoryCache()):
|
||||
assert resolver.resolve().kind == "python_callback"
|
||||
with rebound(facade, "ttl", 12):
|
||||
assert resolver.resolve().kind == "python_callback"
|
||||
|
||||
def custom_key(**_kwargs: object) -> str:
|
||||
return "custom"
|
||||
|
||||
with rebound(facade, "get_cache_key", custom_key):
|
||||
assert resolver.resolve().kind == "python_callback"
|
||||
assert resolver.resolve().kind == "python_callback"
|
||||
delattr(facade, "get_cache_key")
|
||||
assert resolver.resolve().kind == "native"
|
||||
|
||||
|
||||
def test_resolver_and_callback_cycles_can_be_collected() -> None:
|
||||
class CustomCache:
|
||||
pass
|
||||
|
||||
def cyclic_reference() -> weakref.ReferenceType[CustomCache]:
|
||||
callback: Final = CustomCache()
|
||||
namespace: Final = SimpleNamespace(cache=callback)
|
||||
binding: Final = _native.CacheResolver(namespace).resolve()
|
||||
setattr(callback, "binding", binding)
|
||||
return weakref.ref(callback)
|
||||
|
||||
reference: Final = cyclic_reference()
|
||||
gc.collect()
|
||||
assert reference() is None
|
||||
|
||||
|
||||
async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None:
|
||||
client: Final = redis.Redis.from_url(redis_url)
|
||||
namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url, namespace="team"))
|
||||
binding: Final = _native.CacheResolver(namespace).resolve()
|
||||
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
|
||||
envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)}
|
||||
client.set("team:sync", str(envelope))
|
||||
client.set("team:async", json.dumps({"timestamp": time.time(), "response": response}))
|
||||
assert binding.lookup(request("sync")) == response
|
||||
assert await binding.async_lookup(request("team:async")) == response
|
||||
await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response)
|
||||
stored: Final = client.get("team:native")
|
||||
assert isinstance(stored, bytes)
|
||||
assert json.loads(stored)["response"] == response
|
||||
assert 0 < client.ttl("team:native") <= 12
|
||||
assert client.get("litellm-cache:team:native") is None
|
||||
assert client.get("team:team:async") is None
|
||||
client.close()
|
||||
|
||||
|
||||
def test_invalid_duration_and_request_shape_fail_before_storage() -> None:
|
||||
binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve()
|
||||
for seconds in (-1.0, float("nan"), float("inf")):
|
||||
with pytest.raises(ValueError):
|
||||
binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1})
|
||||
assert binding.lookup(request()) is None
|
||||
with pytest.raises(ValueError):
|
||||
_native.NativeCacheHandle.memory(ttl_seconds=-1)
|
||||
Loading…
Add table
Reference in a new issue