diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b37202dce1b..1a4eb51af08 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1377,6 +1377,18 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fancy-regex" version = "0.17.0" @@ -1428,6 +1440,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -1892,11 +1910,32 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "foldhash", +] + [[package]] name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashlink" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb" +dependencies = [ + "hashbrown 0.17.1", +] [[package]] name = "heck" @@ -2277,6 +2316,12 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "iter-read" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071ed4cc1afd86650602c7b11aa2e1ce30762a1c27193201cb5cee9c6ebb1294" + [[package]] name = "itertools" version = "0.13.0" @@ -2435,6 +2480,17 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libsqlite3-sys" +version = "0.38.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -2540,6 +2596,21 @@ dependencies = [ "url", ] +[[package]] +name = "litellm-cache-disk" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "py_literal", + "rand 0.8.7", + "rstest", + "rusqlite", + "serde-pickle", + "serde_json", + "tempfile", + "tokio", +] + [[package]] name = "litellm-cache-memory" version = "0.1.0" @@ -2743,6 +2814,7 @@ dependencies = [ "litellm-auth-gcp", "litellm-cache", "litellm-cache-azure-blob", + "litellm-cache-disk", "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", @@ -4033,6 +4105,16 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rsqlite-vfs" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" +dependencies = [ + "hashbrown 0.16.1", + "thiserror 2.0.19", +] + [[package]] name = "rstest" version = "0.26.1" @@ -4073,6 +4155,21 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "rusqlite" +version = "0.40.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3" +dependencies = [ + "bitflags 2.13.1", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", + "sqlite-wasm-rs", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -4330,6 +4427,19 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-pickle" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b641fdc8bcf2781ee78b30c599700d64ad4f412976143e4c5d0b9df906bb4843" +dependencies = [ + "byteorder", + "iter-read", + "num-bigint 0.4.8", + "num-traits", + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -4557,6 +4667,18 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "sqlite-wasm-rs" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75" +dependencies = [ + "cc", + "js-sys", + "rsqlite-vfs", + "wasm-bindgen", +] + [[package]] name = "sse-stream" version = "0.2.6" @@ -5303,6 +5425,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + [[package]] name = "veil" version = "0.3.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 7d94d31db3e..fabec9bcd6c 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -32,6 +32,7 @@ litellm-cache = { path = "crates/cache" } litellm-cache-azure-blob = { path = "crates/cache-azure-blob" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } +litellm-cache-disk = { path = "crates/cache-disk" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } diff --git a/litellm-rust/crates/cache-disk/Cargo.toml b/litellm-rust/crates/cache-disk/Cargo.toml new file mode 100644 index 00000000000..b96994b3b55 --- /dev/null +++ b/litellm-rust/crates/cache-disk/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "litellm-cache-disk" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +py_literal = "0.4.0" +rand.workspace = true +rusqlite = { version = "0.40", features = ["bundled"] } +serde-pickle = "1.2" +serde_json.workspace = true +tokio.workspace = true + +[dev-dependencies] +rstest.workspace = true +tempfile = "3.27.0" diff --git a/litellm-rust/crates/cache-disk/src/adapter.rs b/litellm-rust/crates/cache-disk/src/adapter.rs new file mode 100644 index 00000000000..b6d318d5509 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/adapter.rs @@ -0,0 +1,10 @@ +use litellm_cache::Error; + +use crate::StoredValue; + +pub trait ValueAdapter: Send + Sync + 'static { + fn read(&self, value: StoredValue) -> Result>, Error>; + fn write(&self, payload: Vec) -> StoredValue; + fn counter_seed(&self, value: Option) -> Result; + fn counter_value(&self, value: f64) -> StoredValue; +} diff --git a/litellm-rust/crates/cache-disk/src/cache.rs b/litellm-rust/crates/cache-disk/src/cache.rs new file mode 100644 index 00000000000..8e1223309b4 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -0,0 +1,301 @@ +use std::{ + path::Path, + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, + CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, +}; + +use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter}; + +pub struct DiskCache { + store: Arc, + adapter: Arc, + codec: S, +} + +impl DiskCache { + pub fn open(directory: impl AsRef, codec: S) -> Result { + Ok(Self { + store: Arc::new(DiskcacheSqliteStore::open(directory)?), + adapter: Arc::new(PythonDiskCacheAdapter), + codec, + }) + } +} + +impl DiskCache { + pub fn with_store(store: D, codec: S) -> Self { + Self { + store: Arc::new(store), + adapter: Arc::new(PythonDiskCacheAdapter), + codec, + } + } +} + +impl DiskCache { + pub fn with_adapter(store: D, adapter: A, codec: S) -> Self { + Self { + store: Arc::new(store), + adapter: Arc::new(adapter), + codec, + } + } + + pub fn directory(&self) -> &Path { + self.store.directory() + } + + fn decode_stored(&self, value: StoredValue) -> Result, Error> { + let Some(bytes) = self.adapter.read(value)? else { + return Ok(None); + }; + self.codec.decode(&bytes).map(Some) + } + + async fn run_blocking(store: Arc, operation: F) -> Result + where + T: Send + 'static, + F: FnOnce(&D) -> Result + Send + 'static, + { + tokio::task::spawn_blocking(move || operation(&store)) + .await + .map_err(|_| Error::Unavailable)? + } +} + +impl BaseCache for DiskCache { + type Value = S::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let value = self.adapter.write(self.codec.encode(&value)?); + let expire_time = context.ttl.map(|ttl| unix_now() + ttl.as_secs_f64()); + self.store.set(key, value, expire_time, unix_now()) + } + + fn get_cache(&self, key: &str, _: &Self::Context) -> Result, Error> { + self.store + .get(key, unix_now())? + .map(|value| self.decode_stored(value)) + .transpose() + .map(|value| value.flatten()) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: ExactCacheContext, + ) -> Result<(), Error> { + let value = self.adapter.write(self.codec.encode(&value)?); + let ttl = context.ttl; + let key = key.to_string(); + Self::run_blocking(Arc::clone(&self.store), move |store| { + let expire_time = ttl.map(|ttl| unix_now() + ttl.as_secs_f64()); + store.set(&key, value, expire_time, unix_now()) + }) + .await + } + + async fn async_get_cache( + &self, + key: &str, + _: &ExactCacheContext, + ) -> Result, Error> { + let key = key.to_string(); + let value = Self::run_blocking(Arc::clone(&self.store), move |store| { + store.get(&key, unix_now()) + }) + .await?; + value + .map(|value| self.decode_stored(value)) + .transpose() + .map(|value| value.flatten()) + } + + async fn async_set_cache_pipeline( + &self, + entries: Vec<(String, Self::Value)>, + context: ExactCacheContext, + ) -> Result<(), Error> { + let entries = entries + .into_iter() + .map(|(key, value)| { + self.codec + .encode(&value) + .map(|value| (key, self.adapter.write(value))) + }) + .collect::, _>>()?; + let expire_after = context.ttl; + Self::run_blocking(Arc::clone(&self.store), move |store| { + for (key, value) in entries { + let expire_time = expire_after.map(|ttl| unix_now() + ttl.as_secs_f64()); + store.set(&key, value, expire_time, unix_now())?; + } + Ok(()) + }) + .await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + let result = Self::run_blocking(Arc::clone(&self.store), |store| { + store.probe().map(|_| CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Disk cache connection test successful".into(), + error: None, + }) + }) + .await; + Ok(match result { + Ok(result) => result, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Disk cache connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + } +} + +impl BatchCache for DiskCache { + fn batch_get_cache( + &self, + keys: &[String], + context: &ExactCacheContext, + ) -> Result>, Error> { + keys.iter() + .map(|key| match self.get_cache(key, context) { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + }) + .collect() + } + + async fn async_batch_get_cache( + &self, + keys: Vec, + _: ExactCacheContext, + ) -> Result>, Error> { + let values = Self::run_blocking(Arc::clone(&self.store), move |store| { + keys.into_iter() + .map(|key| store.get(&key, unix_now()).map(|value| (key, value))) + .collect::, _>>() + }) + .await?; + values + .into_iter() + .map(|(_, value)| match value { + None => Ok(BatchEntry::Miss), + Some(value) => match self.decode_stored(value) { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + }, + }) + .collect() + } +} + +impl DeleteCache for DiskCache { + fn delete_cache(&self, key: &str) -> Result<(), Error> { + self.store.pop(key, unix_now()).map(|_| ()) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + let key = key.to_string(); + Self::run_blocking(Arc::clone(&self.store), move |store| { + store.pop(&key, unix_now()).map(|_| ()) + }) + .await + } +} + +impl FlushCache for DiskCache { + fn flush_cache(&self) -> Result<(), Error> { + self.store.clear() + } + + async fn async_flush_cache(&self) -> Result<(), Error> { + Self::run_blocking(Arc::clone(&self.store), |store| store.clear()).await + } +} + +impl, D: DiskStore, A: ValueAdapter> CounterCache + for DiskCache +{ + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { + increment( + self.adapter.as_ref(), + self.store.as_ref(), + key, + amount, + context.ttl, + ) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { + let key = key.to_string(); + let adapter = Arc::clone(&self.adapter); + Self::run_blocking(Arc::clone(&self.store), move |store| { + increment(adapter.as_ref(), store, &key, amount, context.ttl) + }) + .await + } +} + +fn increment( + adapter: &A, + store: &D, + key: &str, + amount: f64, + ttl: Option, +) -> Result { + let mut result = None; + let mut apply = |current: Option| { + let initial = adapter.counter_seed(current)?; + let value = initial + amount; + let stored = adapter.counter_value(value); + result = Some(value); + Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64()))) + }; + store.update(key, unix_now(), &mut apply)?; + result.ok_or(Error::InvalidEntry) +} + +fn unix_now() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64() +} diff --git a/litellm-rust/crates/cache-disk/src/lib.rs b/litellm-rust/crates/cache-disk/src/lib.rs new file mode 100644 index 00000000000..9b2ffc24915 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/lib.rs @@ -0,0 +1,11 @@ +mod adapter; +mod cache; +mod python; +mod sqlite; +mod store; + +pub use adapter::ValueAdapter; +pub use cache::DiskCache; +pub use python::PythonDiskCacheAdapter; +pub use sqlite::DiskcacheSqliteStore; +pub use store::{DiskStore, StoredValue}; diff --git a/litellm-rust/crates/cache-disk/src/python/mod.rs b/litellm-rust/crates/cache-disk/src/python/mod.rs new file mode 100644 index 00000000000..7a370db357c --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/python/mod.rs @@ -0,0 +1,77 @@ +mod value; + +use litellm_cache::Error; +use py_literal::Value; + +use crate::{StoredValue, ValueAdapter}; + +#[derive(Clone, Copy, Debug, Default)] +pub struct PythonDiskCacheAdapter; + +impl PythonDiskCacheAdapter { + fn python_get_cache(value: StoredValue) -> Result, Error> { + let value = match value { + StoredValue::Bytes(value) => Value::Bytes(value), + StoredValue::Text(value) => Value::String(value), + StoredValue::Integer(value) => Value::Integer(value.into()), + StoredValue::Float(value) => Value::Float(value), + StoredValue::Pickle(value) => value::from_pickle(&value)?, + }; + if !value::is_truthy(&value) { + return Ok(None); + } + match value { + Value::String(text) => Ok(Some( + value::from_json_text(&text).unwrap_or(Value::String(text)), + )), + Value::Bytes(bytes) => match std::str::from_utf8(&bytes) { + Ok(text) => Ok(Some( + value::from_json_text(text).unwrap_or(Value::Bytes(bytes)), + )), + Err(_) => Ok(Some(Value::Bytes(bytes))), + }, + value => Ok(Some(value)), + } + } +} + +impl ValueAdapter for PythonDiskCacheAdapter { + fn read(&self, value: StoredValue) -> Result>, Error> { + match value { + StoredValue::Text(value) => Ok((!value.is_empty()).then(|| value.into_bytes())), + StoredValue::Bytes(value) => Ok((!value.is_empty()).then_some(value)), + value => { + let Some(value) = Self::python_get_cache(value)? else { + return Ok(None); + }; + value::to_json(&value).map(Some) + } + } + } + + fn write(&self, payload: Vec) -> StoredValue { + StoredValue::Bytes(payload) + } + + fn counter_seed(&self, value: Option) -> Result { + let Some(value) = value else { + return Ok(0.0); + }; + let Some(value) = Self::python_get_cache(value)? else { + return Ok(0.0); + }; + Ok(if value::is_int(&value) { + value::to_f64(&value).unwrap_or(0.0) + } else { + 0.0 + }) + } + + fn counter_value(&self, value: f64) -> StoredValue { + if value.fract() == 0.0 && value >= i64::MIN as f64 && value <= i64::MAX as f64 { + StoredValue::Integer(value as i64) + } else { + StoredValue::Float(value) + } + } +} diff --git a/litellm-rust/crates/cache-disk/src/python/value.rs b/litellm-rust/crates/cache-disk/src/python/value.rs new file mode 100644 index 00000000000..645eba7b757 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/python/value.rs @@ -0,0 +1,173 @@ +use litellm_cache::Error; +use py_literal::Value; +use serde_json::{Map, Number}; + +pub(crate) fn from_pickle(bytes: &[u8]) -> Result { + let value = serde_pickle::value_from_slice(bytes, Default::default()) + .map_err(|_| Error::InvalidEntry)?; + from_pickle_value(value) +} + +fn from_pickle_value(value: serde_pickle::Value) -> Result { + match value { + serde_pickle::Value::None => Ok(Value::None), + serde_pickle::Value::Bool(value) => Ok(Value::Boolean(value)), + serde_pickle::Value::I64(value) => integer(value.to_string()), + serde_pickle::Value::Int(value) => integer(value.to_string()), + serde_pickle::Value::F64(value) => Ok(Value::Float(value)), + serde_pickle::Value::String(value) => Ok(Value::String(value)), + serde_pickle::Value::Bytes(value) => Ok(Value::Bytes(value)), + serde_pickle::Value::List(values) => values + .into_iter() + .map(from_pickle_value) + .collect::, _>>() + .map(Value::List), + serde_pickle::Value::Tuple(values) => values + .into_iter() + .map(from_pickle_value) + .collect::, _>>() + .map(Value::Tuple), + serde_pickle::Value::Set(values) => values + .into_iter() + .map(from_pickle_hashable) + .collect::, _>>() + .map(Value::Set), + serde_pickle::Value::FrozenSet(values) => values + .into_iter() + .map(from_pickle_hashable) + .collect::, _>>() + .map(Value::Set), + serde_pickle::Value::Dict(values) => values + .into_iter() + .map(|(key, value)| Ok((from_pickle_hashable(key)?, from_pickle_value(value)?))) + .collect::, Error>>() + .map(Value::Dict), + } +} + +fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result { + Ok(match value { + serde_pickle::HashableValue::None => Value::None, + serde_pickle::HashableValue::Bool(value) => Value::Boolean(value), + serde_pickle::HashableValue::I64(value) => integer(value.to_string())?, + serde_pickle::HashableValue::Int(value) => integer(value.to_string())?, + serde_pickle::HashableValue::F64(value) => Value::Float(value), + serde_pickle::HashableValue::Bytes(value) => Value::Bytes(value), + serde_pickle::HashableValue::String(value) => Value::String(value), + serde_pickle::HashableValue::Tuple(values) => Value::Tuple( + values + .into_iter() + .map(from_pickle_hashable) + .collect::, _>>()?, + ), + serde_pickle::HashableValue::FrozenSet(values) => Value::Set( + values + .into_iter() + .map(from_pickle_hashable) + .collect::, _>>()?, + ), + }) +} + +fn integer(value: String) -> Result { + value.parse().map_err(|_| Error::InvalidEntry) +} + +pub(crate) fn from_json(value: serde_json::Value) -> Value { + match value { + serde_json::Value::Null => Value::None, + serde_json::Value::Bool(value) => Value::Boolean(value), + serde_json::Value::Number(value) => { + if value.is_i64() || value.is_u64() { + integer(value.to_string()) + .unwrap_or(Value::Float(value.as_f64().unwrap_or(f64::NAN))) + } else { + Value::Float(value.as_f64().unwrap_or(f64::NAN)) + } + } + serde_json::Value::String(value) => Value::String(value), + serde_json::Value::Array(values) => { + Value::List(values.into_iter().map(from_json).collect()) + } + serde_json::Value::Object(values) => Value::Dict( + values + .into_iter() + .map(|(key, value)| (Value::String(key), from_json(value))) + .collect(), + ), + } +} + +pub(crate) fn from_json_text(value: &str) -> Result { + serde_json::from_str(value) + .map(from_json) + .map_err(|_| Error::InvalidEntry) +} + +pub(crate) fn is_truthy(value: &Value) -> bool { + match value { + Value::None => false, + Value::Boolean(value) => *value, + Value::Integer(value) => value.to_string() != "0", + Value::Float(value) => *value != 0.0, + Value::Complex(value) => value.re != 0.0 || value.im != 0.0, + Value::String(value) => !value.is_empty(), + Value::Bytes(value) => !value.is_empty(), + Value::Tuple(value) | Value::List(value) | Value::Set(value) => !value.is_empty(), + Value::Dict(value) => !value.is_empty(), + } +} + +pub(crate) fn is_int(value: &Value) -> bool { + matches!(value, Value::Integer(_) | Value::Boolean(_)) +} + +pub(crate) fn to_f64(value: &Value) -> Option { + match value { + Value::Integer(value) => value.to_string().parse().ok(), + Value::Boolean(value) => Some(if *value { 1.0 } else { 0.0 }), + _ => None, + } +} + +pub(crate) fn to_json(value: &Value) -> Result, Error> { + serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry) +} + +fn to_json_value(value: &Value) -> Result { + Ok(match value { + Value::None => serde_json::Value::Null, + Value::Boolean(value) => serde_json::Value::Bool(*value), + Value::Integer(value) => serde_json::Value::Number( + value + .to_string() + .parse::() + .map_err(|_| Error::InvalidEntry)?, + ), + Value::Float(value) => { + serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?) + } + Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry), + Value::String(value) => serde_json::Value::String(value.clone()), + Value::Tuple(values) | Value::List(values) | Value::Set(values) => { + serde_json::Value::Array( + values + .iter() + .map(to_json_value) + .collect::, _>>()?, + ) + } + Value::Dict(values) => { + let values = values + .iter() + .map(|(key, value)| { + let Value::String(key) = key else { + return Err(Error::InvalidEntry); + }; + Ok((key.clone(), to_json_value(value)?)) + }) + .collect::, _>>()?; + serde_json::Value::Object(values) + } + }) +} diff --git a/litellm-rust/crates/cache-disk/src/sqlite.rs b/litellm-rust/crates/cache-disk/src/sqlite.rs new file mode 100644 index 00000000000..9a36f8af6ad --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/sqlite.rs @@ -0,0 +1,817 @@ +use std::{ + collections::HashMap, + fs::{self, OpenOptions}, + io::Write, + path::{Path, PathBuf}, + sync::Mutex, +}; + +use litellm_cache::Error; +use rand::RngCore; +use rusqlite::{Connection, OptionalExtension, params, types::Value}; + +use crate::{DiskStore, StoredValue}; + +const MODE_RAW: i64 = 1; +const MODE_BINARY: i64 = 2; +const MODE_TEXT: i64 = 3; +const MODE_PICKLE: i64 = 4; + +const DEFAULT_DISK_MIN_FILE_SIZE: i64 = 2_i64.pow(15); +const DEFAULT_SIZE_LIMIT: i64 = 2_i64.pow(30); +const DEFAULT_CULL_LIMIT: i64 = 10; + +pub struct DiskcacheSqliteStore { + directory: PathBuf, + connection: Mutex, + min_file_size: usize, + eviction_policy: String, + size_limit: i64, + cull_limit: i64, + statistics: bool, +} + +struct StoredColumns { + size: i64, + mode: i64, + filename: Option, + value: Option, +} + +struct Row { + rowid: i64, + mode: i64, + filename: Option, + value: Value, +} + +impl DiskcacheSqliteStore { + pub fn open(directory: impl AsRef) -> Result { + let directory = directory.as_ref().to_path_buf(); + fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?; + let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?; + let database = directory.join("cache.db"); + let connection = Connection::open(database).map_err(|_| Error::Unavailable)?; + connection + .busy_timeout(std::time::Duration::from_secs(60)) + .map_err(|_| Error::Unavailable)?; + + let mut settings = read_settings(&connection)?; + for (key, value) in default_settings() { + settings.entry(key).or_insert(value); + } + for (key, value) in settings + .iter() + .filter(|(key, _)| key.starts_with("sqlite_")) + { + apply_pragma(&connection, key, value)?; + } + + connection + .execute_batch( + "CREATE TABLE IF NOT EXISTS Settings ( + key TEXT NOT NULL UNIQUE, + value + )", + ) + .map_err(|_| Error::Unavailable)?; + for (key, value) in &settings { + if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") { + connection + .execute( + "INSERT OR REPLACE INTO Settings VALUES (?, ?)", + params![key, value], + ) + .map_err(|_| Error::Unavailable)?; + } + } + for (key, value) in [ + ("count", Value::Integer(0)), + ("size", Value::Integer(0)), + ("hits", Value::Integer(0)), + ("misses", Value::Integer(0)), + ] { + connection + .execute( + "INSERT OR IGNORE INTO Settings VALUES (?, ?)", + params![key, value], + ) + .map_err(|_| Error::Unavailable)?; + } + connection + .execute_batch( + "CREATE TABLE IF NOT EXISTS Cache ( + rowid INTEGER PRIMARY KEY, + key BLOB, + raw INTEGER, + store_time REAL, + expire_time REAL, + access_time REAL, + access_count INTEGER DEFAULT 0, + tag BLOB, + size INTEGER DEFAULT 0, + mode INTEGER DEFAULT 0, + filename TEXT, + value BLOB + ); + CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw); + CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);", + ) + .map_err(|_| Error::Unavailable)?; + + let eviction_policy = setting_string(&settings, "eviction_policy") + .unwrap_or_else(|| "least-recently-stored".to_string()); + match eviction_policy.as_str() { + "none" => {} + "least-recently-stored" => { + connection + .execute_batch( + "CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)", + ) + .map_err(|_| Error::Unavailable)?; + } + "least-recently-used" => { + connection + .execute_batch( + "CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)", + ) + .map_err(|_| Error::Unavailable)?; + } + "least-frequently-used" => { + connection + .execute_batch( + "CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)", + ) + .map_err(|_| Error::Unavailable)?; + } + _ => return Err(Error::Unavailable), + } + connection + .execute_batch( + "CREATE TRIGGER IF NOT EXISTS Settings_count_insert + AFTER INSERT ON Cache FOR EACH ROW BEGIN + UPDATE Settings SET value = value + 1 + WHERE key = \"count\"; END; + CREATE TRIGGER IF NOT EXISTS Settings_count_delete + AFTER DELETE ON Cache FOR EACH ROW BEGIN + UPDATE Settings SET value = value - 1 + WHERE key = \"count\"; END; + CREATE TRIGGER IF NOT EXISTS Settings_size_insert + AFTER INSERT ON Cache FOR EACH ROW BEGIN + UPDATE Settings SET value = value + NEW.size + WHERE key = \"size\"; END; + CREATE TRIGGER IF NOT EXISTS Settings_size_update + AFTER UPDATE ON Cache FOR EACH ROW BEGIN + UPDATE Settings + SET value = value + NEW.size - OLD.size + WHERE key = \"size\"; END; + CREATE TRIGGER IF NOT EXISTS Settings_size_delete + AFTER DELETE ON Cache FOR EACH ROW BEGIN + UPDATE Settings SET value = value - OLD.size + WHERE key = \"size\"; END;", + ) + .map_err(|_| Error::Unavailable)?; + + let min_file_size = setting_i64(&settings, "disk_min_file_size") + .unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE) + .try_into() + .map_err(|_| Error::Unavailable)?; + let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT); + let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT); + let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0; + + Ok(Self { + directory, + connection: Mutex::new(connection), + min_file_size, + eviction_policy, + size_limit, + cull_limit, + statistics, + }) + } + + fn set_locked( + &self, + connection: &Connection, + key: &str, + columns: StoredColumns, + expire_time: Option, + now: f64, + ) -> Result, Error> { + let mut cleanup = Vec::new(); + if let Some(old_filename) = connection + .query_row( + "SELECT filename FROM Cache WHERE key = ? AND raw = 1", + params![key], + |row| row.get::<_, Option>(0), + ) + .optional() + .map_err(|_| Error::Unavailable)? + .flatten() + { + cleanup.push(old_filename); + } + let (size, mode, filename, value) = + (columns.size, columns.mode, columns.filename, columns.value); + let rowid = connection + .query_row( + "SELECT rowid FROM Cache WHERE key = ? AND raw = 1", + params![key], + |row| row.get::<_, i64>(0), + ) + .optional() + .map_err(|_| Error::Unavailable)?; + if let Some(rowid) = rowid { + connection + .execute( + "UPDATE Cache SET store_time = ?, expire_time = ?, access_time = ?, + access_count = 0, tag = NULL, size = ?, mode = ?, filename = ?, value = ? + WHERE rowid = ?", + params![now, expire_time, now, size, mode, filename, value, rowid], + ) + .map_err(|_| Error::Unavailable)?; + } else { + connection + .execute( + "INSERT INTO Cache( + key, raw, store_time, expire_time, access_time, access_count, + tag, size, mode, filename, value + ) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)", + params![key, now, expire_time, now, size, mode, filename, value], + ) + .map_err(|_| Error::Unavailable)?; + } + cleanup.extend(self.cull(connection, now)?); + Ok(cleanup) + } + + fn cull(&self, connection: &Connection, now: f64) -> Result, Error> { + if self.cull_limit <= 0 { + return Ok(Vec::new()); + } + let mut cleanup = Vec::new(); + let expired = connection + .prepare( + "SELECT rowid, filename FROM Cache + WHERE expire_time IS NOT NULL AND expire_time < ? + ORDER BY expire_time LIMIT ?", + ) + .map_err(|_| Error::Unavailable)? + .query_map(params![now, self.cull_limit], |row| { + Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) + }) + .map_err(|_| Error::Unavailable)? + .collect::, _>>() + .map_err(|_| Error::Unavailable)?; + for (_, filename) in &expired { + if let Some(filename) = filename { + cleanup.push(filename.clone()); + } + } + for (rowid, _) in &expired { + connection + .execute("DELETE FROM Cache WHERE rowid = ?", params![rowid]) + .map_err(|_| Error::Unavailable)?; + } + let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit); + if remaining <= 0 || self.volume(connection)? < self.size_limit { + return Ok(cleanup); + } + let order = match self.eviction_policy.as_str() { + "none" => return Ok(cleanup), + "least-recently-stored" => "store_time", + "least-recently-used" => "access_time", + "least-frequently-used" => "access_count", + _ => return Err(Error::Unavailable), + }; + let rows = connection + .prepare(&format!( + "SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?" + )) + .map_err(|_| Error::Unavailable)? + .query_map(params![remaining], |row| { + Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) + }) + .map_err(|_| Error::Unavailable)? + .collect::, _>>() + .map_err(|_| Error::Unavailable)?; + for (_, filename) in &rows { + if let Some(filename) = filename { + cleanup.push(filename.clone()); + } + } + for (rowid, _) in rows { + connection + .execute("DELETE FROM Cache WHERE rowid = ?", params![rowid]) + .map_err(|_| Error::Unavailable)?; + } + Ok(cleanup) + } + + fn volume(&self, connection: &Connection) -> Result { + let page_count: i64 = connection + .query_row("PRAGMA page_count", [], |row| row.get(0)) + .map_err(|_| Error::Unavailable)?; + let page_size: i64 = connection + .query_row("PRAGMA page_size", [], |row| row.get(0)) + .map_err(|_| Error::Unavailable)?; + let size: i64 = connection + .query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| { + row.get(0) + }) + .map_err(|_| Error::Unavailable)?; + Ok(page_count.saturating_mul(page_size).saturating_add(size)) + } +} + +impl DiskStore for DiskcacheSqliteStore { + fn directory(&self) -> &Path { + &self.directory + } + + fn get(&self, key: &str, now: f64) -> Result, Error> { + let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache + WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)"; + let row = connection + .query_row(select, params![key, now], row_from_query) + .optional() + .map_err(|_| Error::Unavailable)?; + if !self.statistics && !has_get_update(&self.eviction_policy) { + return row + .map(|row| fetch_row(&self.directory, row)) + .transpose() + .map(|value| value.flatten()); + } + transactional(&connection, |connection| { + let row = connection + .query_row(select, params![key, now], row_from_query) + .optional() + .map_err(|_| Error::Unavailable)?; + let Some(row) = row else { + if self.statistics { + connection + .execute( + "UPDATE Settings SET value = value + 1 WHERE key = 'misses'", + [], + ) + .map_err(|_| Error::Unavailable)?; + } + return Ok(None); + }; + let rowid = row.rowid; + let value = fetch_row(&self.directory, row); + let hit = value.as_ref().is_ok_and(Option::is_some); + if hit && self.statistics { + connection + .execute( + "UPDATE Settings SET value = value + 1 WHERE key = 'hits'", + [], + ) + .map_err(|_| Error::Unavailable)?; + } else if !hit && self.statistics { + connection + .execute( + "UPDATE Settings SET value = value + 1 WHERE key = 'misses'", + [], + ) + .map_err(|_| Error::Unavailable)?; + } + if has_get_update(&self.eviction_policy) && hit { + let update = match self.eviction_policy.as_str() { + "least-recently-used" => "UPDATE Cache SET access_time = ? WHERE rowid = ?", + "least-frequently-used" => { + "UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?" + } + _ => return Err(Error::Unavailable), + }; + if self.eviction_policy == "least-recently-used" { + connection + .execute(update, params![now, rowid]) + .map_err(|_| Error::Unavailable)?; + } else { + connection + .execute(update, params![rowid]) + .map_err(|_| Error::Unavailable)?; + } + } + value + }) + } + + fn set( + &self, + key: &str, + value: StoredValue, + expire_time: Option, + now: f64, + ) -> Result<(), Error> { + let columns = store_value(&self.directory, self.min_file_size, value)?; + let new_filename = columns.filename.clone(); + let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let result = transactional(&connection, |connection| { + self.set_locked(connection, key, columns, expire_time, now) + }); + match result { + Ok(cleanup) => { + cleanup_files(&self.directory, cleanup); + Ok(()) + } + Err(error) => { + if let Some(filename) = new_filename { + remove_file(&self.directory, &filename); + } + Err(error) + } + } + } + + fn pop(&self, key: &str, now: f64) -> Result, Error> { + let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let selected = transactional(&connection, |connection| { + let row = connection + .query_row( + "SELECT rowid, expire_time, mode, filename, value FROM Cache + WHERE key = ? AND raw = 1 + AND (expire_time IS NULL OR expire_time > ?)", + params![key, now], + row_from_query, + ) + .optional() + .map_err(|_| Error::Unavailable)?; + let Some(row) = row else { + return Ok(None); + }; + connection + .execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid]) + .map_err(|_| Error::Unavailable)?; + Ok(Some(row)) + })?; + let Some(row) = selected else { + return Ok(None); + }; + let filename = row.filename.clone(); + let result = fetch_row(&self.directory, row)?; + if let Some(filename) = filename { + remove_file(&self.directory, &filename); + } + Ok(result) + } + + fn clear(&self) -> Result<(), Error> { + let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let mut last_rowid = 0_i64; + loop { + let batch = transactional(&connection, |connection| { + let rows = connection + .prepare( + "SELECT rowid, filename FROM Cache + WHERE rowid > ? ORDER BY rowid LIMIT 100", + ) + .map_err(|_| Error::Unavailable)? + .query_map(params![last_rowid], |row| { + Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) + }) + .map_err(|_| Error::Unavailable)? + .collect::, _>>() + .map_err(|_| Error::Unavailable)?; + if rows.is_empty() { + return Ok(rows); + } + let ids = rows + .iter() + .map(|(rowid, _)| rowid.to_string()) + .collect::>() + .join(","); + connection + .execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), []) + .map_err(|_| Error::Unavailable)?; + Ok(rows) + })?; + if batch.is_empty() { + return Ok(()); + } + last_rowid = batch.last().map(|(rowid, _)| *rowid).unwrap_or(last_rowid); + cleanup_files( + &self.directory, + batch + .into_iter() + .filter_map(|(_, filename)| filename) + .collect(), + ); + } + } + + fn update( + &self, + key: &str, + now: f64, + apply: &mut dyn FnMut(Option) -> Result<(StoredValue, Option), Error>, + ) -> Result<(), Error> { + let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let mut created_filename = None; + let result = transactional(&connection, |connection| { + let current = connection + .query_row( + "SELECT rowid, expire_time, mode, filename, value FROM Cache + WHERE key = ? AND raw = 1 + AND (expire_time IS NULL OR expire_time > ?)", + params![key, now], + row_from_query, + ) + .optional() + .map_err(|_| Error::Unavailable)? + .map(|row| fetch_row(&self.directory, row)) + .transpose()? + .flatten(); + let (value, expire_time) = apply(current)?; + let columns = store_value(&self.directory, self.min_file_size, value)?; + created_filename = columns.filename.clone(); + let cleanup = self.set_locked(connection, key, columns, expire_time, now)?; + Ok(cleanup) + }); + match result { + Ok(cleanup) => { + cleanup_files(&self.directory, cleanup); + Ok(()) + } + Err(error) => { + if let Some(filename) = created_filename { + remove_file(&self.directory, &filename); + } + Err(error) + } + } + } + + fn probe(&self) -> Result<(), Error> { + let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + connection + .query_row( + "SELECT value FROM Settings WHERE key = 'count'", + [], + |row| row.get::<_, i64>(0), + ) + .map(|_| ()) + .map_err(|_| Error::Unavailable) + } +} + +fn default_settings() -> HashMap { + HashMap::from([ + ("statistics".to_string(), Value::Integer(0)), + ("tag_index".to_string(), Value::Integer(0)), + ( + "eviction_policy".to_string(), + Value::Text("least-recently-stored".to_string()), + ), + ("size_limit".to_string(), Value::Integer(DEFAULT_SIZE_LIMIT)), + ("cull_limit".to_string(), Value::Integer(DEFAULT_CULL_LIMIT)), + ("sqlite_auto_vacuum".to_string(), Value::Integer(1)), + ("sqlite_cache_size".to_string(), Value::Integer(8192)), + ( + "sqlite_journal_mode".to_string(), + Value::Text("wal".to_string()), + ), + ( + "sqlite_mmap_size".to_string(), + Value::Integer(2_i64.pow(26)), + ), + ("sqlite_synchronous".to_string(), Value::Integer(1)), + ( + "disk_min_file_size".to_string(), + Value::Integer(DEFAULT_DISK_MIN_FILE_SIZE), + ), + ("disk_pickle_protocol".to_string(), Value::Integer(5)), + ]) +} + +fn read_settings(connection: &Connection) -> Result, Error> { + let mut statement = match connection.prepare("SELECT key, value FROM Settings") { + Ok(statement) => statement, + Err(_) => return Ok(HashMap::new()), + }; + statement + .query_map([], |row| Ok((row.get(0)?, row.get(1)?))) + .map_err(|_| Error::Unavailable)? + .collect::, _>>() + .map_err(|_| Error::Unavailable) +} + +fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> { + let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?; + match value { + Value::Integer(value) => connection + .pragma_update(None, pragma, value) + .map_err(|_| Error::Unavailable), + Value::Text(value) => connection + .pragma_update(None, pragma, value) + .map_err(|_| Error::Unavailable), + _ => Err(Error::Unavailable), + } +} + +fn setting_i64(settings: &HashMap, key: &str) -> Option { + match settings.get(key) { + Some(Value::Integer(value)) => Some(*value), + _ => None, + } +} + +fn setting_string(settings: &HashMap, key: &str) -> Option { + match settings.get(key) { + Some(Value::Text(value)) => Some(value.clone()), + _ => None, + } +} + +fn has_get_update(policy: &str) -> bool { + matches!(policy, "least-recently-used" | "least-frequently-used") +} + +fn row_from_query(row: &rusqlite::Row<'_>) -> rusqlite::Result { + Ok(Row { + rowid: row.get(0)?, + mode: row.get(2)?, + filename: row.get(3)?, + value: row.get(4)?, + }) +} + +fn fetch_row(directory: &Path, row: Row) -> Result, Error> { + match row.mode { + MODE_RAW => match row.value { + Value::Blob(value) => Ok(Some(StoredValue::Bytes(value))), + Value::Text(value) => Ok(Some(StoredValue::Text(value))), + Value::Integer(value) => Ok(Some(StoredValue::Integer(value))), + Value::Real(value) => Ok(Some(StoredValue::Float(value))), + Value::Null => Err(Error::InvalidEntry), + }, + MODE_BINARY | MODE_PICKLE => { + let bytes = match row.value { + Value::Blob(value) => value, + Value::Null => { + let Some(value) = read_file(directory, row.filename.as_deref())? else { + return Ok(None); + }; + value + } + _ => return Err(Error::InvalidEntry), + }; + Ok(Some(if row.mode == MODE_BINARY { + StoredValue::Bytes(bytes) + } else { + StoredValue::Pickle(bytes) + })) + } + MODE_TEXT => { + let bytes = match row.value { + Value::Null => { + let Some(value) = read_file(directory, row.filename.as_deref())? else { + return Ok(None); + }; + value + } + Value::Blob(value) => value, + Value::Text(value) => value.into_bytes(), + _ => return Err(Error::InvalidEntry), + }; + Ok(Some(StoredValue::Text( + String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?, + ))) + } + _ => Err(Error::InvalidEntry), + } +} + +fn read_file(directory: &Path, filename: Option<&str>) -> Result>, Error> { + let Some(filename) = filename else { + return Err(Error::InvalidEntry); + }; + match fs::read(directory.join(filename)) { + Ok(value) => Ok(Some(value)), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(_) => Err(Error::Unavailable), + } +} + +fn store_value( + directory: &Path, + min_file_size: usize, + value: StoredValue, +) -> Result { + match value { + StoredValue::Integer(value) => Ok(StoredColumns { + size: 0, + mode: MODE_RAW, + filename: None, + value: Some(Value::Integer(value)), + }), + StoredValue::Float(value) => Ok(StoredColumns { + size: 0, + mode: MODE_RAW, + filename: None, + value: Some(Value::Real(value)), + }), + StoredValue::Text(value) if value.chars().count() < min_file_size => Ok(StoredColumns { + size: 0, + mode: MODE_RAW, + filename: None, + value: Some(Value::Text(value)), + }), + StoredValue::Text(value) => { + let bytes = value.into_bytes(); + let filename = write_file(directory, &bytes)?; + Ok(StoredColumns { + size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?, + mode: MODE_TEXT, + filename: Some(filename), + value: None, + }) + } + StoredValue::Bytes(value) if value.len() < min_file_size => Ok(StoredColumns { + size: 0, + mode: MODE_RAW, + filename: None, + value: Some(Value::Blob(value)), + }), + StoredValue::Bytes(value) => { + let filename = write_file(directory, &value)?; + Ok(StoredColumns { + size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?, + mode: MODE_BINARY, + filename: Some(filename), + value: None, + }) + } + StoredValue::Pickle(value) if value.len() < min_file_size => Ok(StoredColumns { + size: 0, + mode: MODE_PICKLE, + filename: None, + value: Some(Value::Blob(value)), + }), + StoredValue::Pickle(value) => { + let filename = write_file(directory, &value)?; + Ok(StoredColumns { + size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?, + mode: MODE_PICKLE, + filename: Some(filename), + value: None, + }) + } + } +} + +fn write_file(directory: &Path, bytes: &[u8]) -> Result { + let mut random = [0_u8; 16]; + rand::rngs::OsRng.fill_bytes(&mut random); + let hex = random + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]); + let path = directory.join(&filename); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?; + } + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .open(path) + .map_err(|_| Error::Unavailable)?; + file.write_all(bytes).map_err(|_| Error::Unavailable)?; + Ok(filename) +} + +fn cleanup_files(directory: &Path, filenames: Vec) { + for filename in filenames { + remove_file(directory, &filename); + } +} + +fn remove_file(directory: &Path, filename: &str) { + let path = directory.join(filename); + let _ = fs::remove_file(&path); +} + +fn transactional( + connection: &Connection, + operation: impl FnOnce(&Connection) -> Result, +) -> Result { + connection + .execute_batch("BEGIN IMMEDIATE") + .map_err(|_| Error::Unavailable)?; + match operation(connection) { + Ok(value) => { + connection + .execute_batch("COMMIT") + .map_err(|_| Error::Unavailable)?; + Ok(value) + } + Err(error) => { + let _ = connection.execute_batch("ROLLBACK"); + Err(error) + } + } +} diff --git a/litellm-rust/crates/cache-disk/src/store.rs b/litellm-rust/crates/cache-disk/src/store.rs new file mode 100644 index 00000000000..ed167317cf0 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/store.rs @@ -0,0 +1,33 @@ +use std::path::Path; + +use litellm_cache::Error; + +#[derive(Clone, Debug, PartialEq)] +pub enum StoredValue { + Bytes(Vec), + Text(String), + Integer(i64), + Float(f64), + Pickle(Vec), +} + +pub trait DiskStore: Send + Sync + 'static { + fn directory(&self) -> &Path; + fn get(&self, key: &str, now: f64) -> Result, Error>; + fn set( + &self, + key: &str, + value: StoredValue, + expire_time: Option, + now: f64, + ) -> Result<(), Error>; + fn pop(&self, key: &str, now: f64) -> Result, Error>; + fn clear(&self) -> Result<(), Error>; + fn update( + &self, + key: &str, + now: f64, + apply: &mut dyn FnMut(Option) -> Result<(StoredValue, Option), Error>, + ) -> Result<(), Error>; + fn probe(&self) -> Result<(), Error>; +} diff --git a/litellm-rust/crates/cache-disk/tests/cache.rs b/litellm-rust/crates/cache-disk/tests/cache.rs new file mode 100644 index 00000000000..dd1f2b1f04e --- /dev/null +++ b/litellm-rust/crates/cache-disk/tests/cache.rs @@ -0,0 +1,431 @@ +use std::{ + fs, + path::{Path, PathBuf}, + sync::Arc, + thread, + time::Duration, +}; + +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext, + FlushCache, JsonCodec, +}; +use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter}; +use rstest::{fixture, rstest}; +use rusqlite::Connection; +use serde_json::{Value, json}; +use tempfile::TempDir; + +struct Sandbox { + directory: TempDir, +} + +#[fixture] +fn sandbox() -> Sandbox { + Sandbox { + directory: tempfile::tempdir().unwrap(), + } +} + +impl Sandbox { + fn store(&self) -> DiskcacheSqliteStore { + DiskcacheSqliteStore::open(self.directory.path()).unwrap() + } + + fn cache(&self) -> DiskCache> + where + JsonCodec: CacheCodec, + { + DiskCache::open(self.directory.path(), JsonCodec::new()).unwrap() + } + + fn db(&self) -> Connection { + Connection::open(self.directory.path().join("cache.db")).unwrap() + } + + fn value_files(&self) -> Vec { + fn visit(directory: &Path, files: &mut Vec) { + for entry in fs::read_dir(directory).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + visit(&path, files); + } else if path.extension().is_some_and(|extension| extension == "val") { + files.push(path); + } + } + } + + let mut files = Vec::new(); + visit(self.directory.path(), &mut files); + files + } +} + +#[rstest] +fn relative_store_directory_is_absolutized(sandbox: Sandbox) { + let relative = PathBuf::from(format!( + ".litellm-cache-disk-{}", + sandbox + .directory + .path() + .file_name() + .unwrap() + .to_string_lossy() + )); + let store = DiskcacheSqliteStore::open(&relative).unwrap(); + assert!(store.directory().is_absolute()); + assert!(store.directory().ends_with(&relative)); + let directory = store.directory().to_path_buf(); + drop(store); + fs::remove_dir_all(directory).unwrap(); +} + +#[derive(Clone, Copy, Debug, Default)] +struct TextAdapter; + +impl ValueAdapter for TextAdapter { + fn read(&self, value: StoredValue) -> Result>, litellm_cache::Error> { + match value { + StoredValue::Text(value) => Ok(Some(value.into_bytes())), + _ => Ok(None), + } + } + + fn write(&self, payload: Vec) -> StoredValue { + StoredValue::Text(String::from_utf8(payload).unwrap()) + } + + fn counter_seed(&self, _: Option) -> Result { + Ok(0.0) + } + + fn counter_value(&self, value: f64) -> StoredValue { + if value.fract() == 0.0 { + StoredValue::Integer(value as i64) + } else { + StoredValue::Float(value) + } + } +} + +#[rstest] +fn roundtrip_persists_and_reopens(sandbox: Sandbox) { + let context = ExactCacheContext::default(); + let opened = sandbox.cache::(); + opened + .set_cache("key", json!({"answer": 42}), &context) + .unwrap(); + assert_eq!( + opened.get_cache("key", &context).unwrap(), + Some(json!({"answer": 42})) + ); + drop(opened); + let reopened = sandbox.cache::(); + assert_eq!( + reopened.get_cache("key", &context).unwrap(), + Some(json!({"answer": 42})) + ); +} + +#[rstest] +fn ttl_and_expired_culling_match_cache_contract(sandbox: Sandbox) { + let store = sandbox.store(); + store + .set( + "expired", + StoredValue::Bytes(b"old".to_vec()), + Some(10.0), + 0.0, + ) + .unwrap(); + assert_eq!(store.get("expired", 10.0).unwrap(), None); + store + .set("new", StoredValue::Bytes(b"new".to_vec()), None, 11.0) + .unwrap(); + assert_eq!( + sandbox + .db() + .query_row("SELECT COUNT(*) FROM Cache", [], |row| row.get::<_, i64>(0)) + .unwrap(), + 1 + ); + assert_eq!( + sandbox + .db() + .query_row( + "SELECT value FROM Settings WHERE key = 'count'", + [], + |row| row.get::<_, i64>(0) + ) + .unwrap(), + 1 + ); +} + +#[rstest] +fn batch_preserves_order_and_classifies_misses_and_invalid_values(sandbox: Sandbox) { + let store = sandbox.store(); + store + .set( + "hit", + StoredValue::Bytes(br#"{"ok":true}"#.to_vec()), + None, + 0.0, + ) + .unwrap(); + store + .set( + "invalid", + StoredValue::Pickle(vec![0x80, 0x05, 0x2e]), + None, + 0.0, + ) + .unwrap(); + let entries = sandbox + .cache::() + .batch_get_cache( + &["hit".into(), "missing".into(), "invalid".into()], + &ExactCacheContext::default(), + ) + .unwrap(); + assert_eq!( + entries, + vec![ + BatchEntry::Hit(json!({"ok": true})), + BatchEntry::Miss, + BatchEntry::Invalid + ] + ); +} + +#[rstest] +#[case(StoredValue::Bytes(Vec::new()))] +#[case(StoredValue::Text(String::new()))] +#[case(StoredValue::Integer(0))] +#[case(StoredValue::Float(0.0))] +#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]))] +#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]))] +#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]))] +#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]))] +#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]))] +#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]))] +#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]))] +#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]))] +fn falsy_values_are_misses(sandbox: Sandbox, #[case] value: StoredValue) { + sandbox.store().set("key", value, None, 0.0).unwrap(); + assert_eq!( + sandbox + .cache::() + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), + None + ); +} + +#[rstest] +#[case(Some(StoredValue::Integer(2)), 1.5, 3.5, "real")] +#[case(Some(StoredValue::Integer(2)), 1.0, 3.0, "integer")] +#[case(Some(StoredValue::Float(3.5)), 1.0, 1.0, "integer")] +#[case(Some(StoredValue::Text("not a number".into())), 2.0, 2.0, "integer")] +#[case(Some(StoredValue::Text("5".into())), 2.0, 7.0, "integer")] +#[case(Some(StoredValue::Text("3.5".into())), 2.0, 2.0, "integer")] +#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0, 2.0, "integer")] +#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 1.0, 3.0, "integer")] +#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 1.0, 1.0, "integer")] +#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 4.0, 4.0, "integer")] +fn counters_follow_python_initialization( + sandbox: Sandbox, + #[case] initial: Option, + #[case] amount: f64, + #[case] expected: f64, + #[case] sqlite_type: &str, +) { + if let Some(initial) = initial { + sandbox.store().set("counter", initial, None, 0.0).unwrap(); + } + let cache = sandbox.cache::(); + assert_eq!( + cache + .increment_cache("counter", amount, ExactCacheContext::default()) + .unwrap(), + expected + ); + assert_eq!( + sandbox + .db() + .query_row( + "SELECT typeof(value) FROM Cache WHERE key = 'counter'", + [], + |row| row.get::<_, String>(0) + ) + .unwrap(), + sqlite_type + ); +} + +#[rstest] +fn counters_are_atomic_across_concurrent_callers(sandbox: Sandbox) { + let cache = Arc::new(sandbox.cache::()); + let workers = (0..8) + .map(|_| { + let cache = Arc::clone(&cache); + thread::spawn(move || { + for _ in 0..25 { + cache + .increment_cache("counter", 1.0, ExactCacheContext::default()) + .unwrap(); + } + }) + }) + .collect::>(); + for worker in workers { + worker.join().unwrap(); + } + assert_eq!( + cache + .increment_cache("counter", 0.0, ExactCacheContext::default()) + .unwrap(), + 200.0 + ); +} + +#[rstest] +fn fractional_then_integer_increment_follows_python_behavior(sandbox: Sandbox) { + let cache = sandbox.cache::(); + assert_eq!( + cache + .increment_cache("counter", 3.5, ExactCacheContext::default()) + .unwrap(), + 3.5 + ); + assert_eq!( + cache + .increment_cache("counter", 1.0, ExactCacheContext::default()) + .unwrap(), + 1.0 + ); +} + +#[rstest] +fn increment_ttl_replacement_clears_expiry_without_ttl(sandbox: Sandbox) { + let cache = sandbox.cache::(); + cache + .increment_cache( + "counter", + 1.0, + ExactCacheContext { + ttl: Some(Duration::from_secs(60)), + }, + ) + .unwrap(); + assert!( + sandbox + .db() + .query_row( + "SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'", + [], + |row| row.get::<_, bool>(0) + ) + .unwrap() + ); + cache + .increment_cache("counter", 1.0, ExactCacheContext::default()) + .unwrap(); + assert!( + !sandbox + .db() + .query_row( + "SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'", + [], + |row| row.get::<_, bool>(0) + ) + .unwrap() + ); +} + +#[rstest] +fn custom_adapter_controls_storage_and_reads(sandbox: Sandbox) { + let cache = DiskCache::with_adapter(sandbox.store(), TextAdapter, JsonCodec::::new()); + cache + .set_cache("key", json!({"answer": 42}), &ExactCacheContext::default()) + .unwrap(); + assert!(matches!( + sandbox.store().get("key", 0.0).unwrap(), + Some(StoredValue::Text(_)) + )); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"answer": 42})) + ); +} + +#[rstest] +fn delete_flush_and_spilled_file_replacement_clean_up_storage(sandbox: Sandbox) { + let large = vec![b'x'; 32 * 1024]; + sandbox + .store() + .set("large", StoredValue::Bytes(large.clone()), None, 0.0) + .unwrap(); + assert_eq!(sandbox.value_files().len(), 1); + sandbox + .store() + .set( + "large", + StoredValue::Bytes(vec![b'y'; 32 * 1024]), + None, + 0.0, + ) + .unwrap(); + assert_eq!(sandbox.value_files().len(), 1); + sandbox.store().pop("large", 0.0).unwrap(); + assert!(sandbox.value_files().is_empty()); + sandbox + .store() + .set("a", StoredValue::Bytes(large.clone()), None, 0.0) + .unwrap(); + sandbox + .store() + .set("b", StoredValue::Bytes(large), None, 0.0) + .unwrap(); + sandbox.store().clear().unwrap(); + assert!(sandbox.value_files().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn async_operations_connection_and_delete_match_sync_operations(sandbox: Sandbox) { + let cache = sandbox.cache::(); + let context = ExactCacheContext { + ttl: Some(Duration::from_secs(60)), + }; + cache + .async_set_cache("a", json!(1), context.clone()) + .await + .unwrap(); + cache + .async_set_cache_pipeline( + vec![("b".into(), json!(2)), ("c".into(), json!(3))], + context.clone(), + ) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("a", &context).await.unwrap(), + Some(json!(1)) + ); + assert_eq!( + cache + .async_batch_get_cache(vec!["c".into(), "missing".into()], context.clone()) + .await + .unwrap(), + vec![BatchEntry::Hit(json!(3)), BatchEntry::Miss] + ); + cache.async_delete_cache("a").await.unwrap(); + cache.async_flush_cache().await.unwrap(); + assert_eq!( + cache.test_connection().await.unwrap().status, + litellm_cache::CacheConnectionStatus::Success + ); +} diff --git a/litellm-rust/crates/cache-disk/tests/python_compat.rs b/litellm-rust/crates/cache-disk/tests/python_compat.rs new file mode 100644 index 00000000000..9cbef8573bd --- /dev/null +++ b/litellm-rust/crates/cache-disk/tests/python_compat.rs @@ -0,0 +1,113 @@ +use litellm_cache::Error; +use litellm_cache_disk::{PythonDiskCacheAdapter, StoredValue, ValueAdapter}; +use rstest::rstest; + +enum ReadExpectation { + Bytes(&'static [u8]), + Miss, + Invalid, +} + +#[rstest] +#[case::pickled_dictionary_with_string_keys( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e]), + ReadExpectation::Bytes(br#"{"a":1}"#) +)] +#[case::pickled_list_of_integers( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x5d, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x65, 0x2e]), + ReadExpectation::Bytes(br#"[1,2]"#) +)] +#[case::pickled_tuple_of_integers( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4b, 0x01, 0x4b, 0x02, 0x86, 0x94, 0x2e]), + ReadExpectation::Bytes(br#"[1,2]"#) +)] +#[case::pickled_set_of_integers( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8f, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x90, 0x2e]), + ReadExpectation::Bytes(br#"[1,2]"#) +)] +#[case::pickled_response_envelope( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x28, 0x8c, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x94, 0x47, 0x3f, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x94, 0x8c, 0x08, 0x7b, 0x22, 0x61, 0x22, 0x3a, 0x20, 0x31, 0x7d, 0x94, 0x75, 0x2e]), + ReadExpectation::Bytes(br#"{"response":"{\"a\": 1}","timestamp":1.5}"#) +)] +#[case::non_json_text( + StoredValue::Text("not json".into()), + ReadExpectation::Bytes(b"not json") +)] +#[case::json_text( + StoredValue::Text("{\"a\": 1}".into()), + ReadExpectation::Bytes(br#"{"a": 1}"#) +)] +#[case::non_utf8_bytes( + StoredValue::Bytes(vec![0xff, 0xfe]), + ReadExpectation::Bytes(&[0xff, 0xfe]) +)] +#[case::integer_seven(StoredValue::Integer(7), ReadExpectation::Bytes(b"7"))] +#[case::float_one_point_five(StoredValue::Float(1.5), ReadExpectation::Bytes(b"1.5"))] +#[case::pickled_true( + StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e]), + ReadExpectation::Bytes(b"true") +)] +#[case::pickled_negative_integer( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0xfd, 0xff, 0xff, 0xff, 0x2e]), + ReadExpectation::Bytes(b"-3") +)] +#[case::pickled_bytes( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x43, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x94, 0x2e]), + ReadExpectation::Invalid +)] +#[case::pickled_dictionary_with_integer_key( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x4b, 0x01, 0x8c, 0x01, 0x61, 0x94, 0x73, 0x2e]), + ReadExpectation::Invalid +)] +#[case::pickled_complex( + StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x2e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x62, 0x75, 0x69, 0x6c, 0x74, 0x69, 0x6e, 0x73, 0x94, 0x8c, 0x07, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x94, 0x93, 0x94, 0x47, 0x3f, 0xf0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x86, 0x94, 0x52, 0x94, 0x2e]), + ReadExpectation::Invalid +)] +#[case::truncated_pickle( + StoredValue::Pickle(vec![0x80, 0x05, 0x2e]), + ReadExpectation::Invalid +)] +#[case::empty_bytes(StoredValue::Bytes(Vec::new()), ReadExpectation::Miss)] +#[case::empty_text(StoredValue::Text(String::new()), ReadExpectation::Miss)] +#[case::zero_integer(StoredValue::Integer(0), ReadExpectation::Miss)] +#[case::zero_float(StoredValue::Float(0.0), ReadExpectation::Miss)] +#[case::pickled_none(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]), ReadExpectation::Miss)] +#[case::pickled_false(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]), ReadExpectation::Miss)] +#[case::pickled_zero(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]), ReadExpectation::Miss)] +#[case::pickled_zero_float(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]), ReadExpectation::Miss)] +#[case::pickled_empty_string(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]), ReadExpectation::Miss)] +#[case::pickled_empty_list(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]), ReadExpectation::Miss)] +#[case::pickled_empty_dictionary(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]), ReadExpectation::Miss)] +#[case::pickled_empty_tuple(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]), ReadExpectation::Miss)] +fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation) { + let result = PythonDiskCacheAdapter.read(row); + match expected { + ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected), + ReadExpectation::Miss => assert_eq!(result.unwrap(), None), + ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))), + } +} + +#[rstest] +#[case::integer_two(Some(StoredValue::Integer(2)), 2.0)] +#[case::float_three_point_five(Some(StoredValue::Float(3.5)), 0.0)] +#[case::text_not_a_number(Some(StoredValue::Text("not a number".into())), 0.0)] +#[case::text_five(Some(StoredValue::Text("5".into())), 5.0)] +#[case::text_three_point_five(Some(StoredValue::Text("3.5".into())), 0.0)] +#[case::pickled_true(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0)] +#[case::pickled_two(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 2.0)] +#[case::pickled_dictionary(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 0.0)] +#[case::missing(None, 0.0)] +#[case::pickled_none(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 0.0)] +fn python_counter_seed_cases(#[case] row: Option, #[case] expected: f64) { + assert_eq!(PythonDiskCacheAdapter.counter_seed(row).unwrap(), expected); +} + +#[rstest] +#[case::integer_three(3.0, StoredValue::Integer(3))] +#[case::fractional_three_point_five(3.5, StoredValue::Float(3.5))] +#[case::negative_zero(-0.0, StoredValue::Integer(0))] +#[case::large_float(1e300, StoredValue::Float(1e300))] +fn python_counter_value_cases(#[case] value: f64, #[case] expected: StoredValue) { + assert_eq!(PythonDiskCacheAdapter.counter_value(value), expected); +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 5528e12ee2d..f5d08850b39 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -24,6 +24,7 @@ litellm-cache.workspace = true litellm-cache-azure-blob.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-cache-disk.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 5c962f2bc7a..927017e2e60 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1,4 +1,4 @@ -use std::time::Duration; +use std::{path::PathBuf, time::Duration}; use litellm_cache::CacheType; use litellm_cache_redis::{RedisNode, RedisTopology}; @@ -26,6 +26,10 @@ pub(super) struct MemoryCacheConfig { pub(super) max_entry_bytes: usize, } +pub(super) struct DiskCacheConfig { + pub(super) directory: PathBuf, +} + #[derive(Debug, PartialEq)] pub(super) enum RedisProtocol { Resp2, @@ -94,6 +98,7 @@ const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31; pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + Disk(DiskCacheConfig), AzureBlob(AzureBlobCacheConfig), } @@ -109,6 +114,7 @@ pub(super) enum UnsupportedCacheConfig { RedisCredentials, RedisConnection, RedisOption, + DiskStore, } impl UnsupportedCacheConfig { @@ -119,6 +125,7 @@ impl UnsupportedCacheConfig { Self::RedisCredentials => "native Redis credentials require Python", Self::RedisConnection => "native Redis connection type is not implemented", Self::RedisOption => "native Redis configuration requires Python", + Self::DiskStore => "native disk cache requires the built-in diskcache store", } } } @@ -161,6 +168,13 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::Disk) => match project_disk(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::Disk(backend), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| { CacheConfigProjection::Native(Box::new(Self { policy, @@ -171,7 +185,6 @@ impl NativeCacheConfig { CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::S3 - | CacheType::Disk | CacheType::QdrantSemantic | CacheType::Gcs, ) @@ -185,6 +198,7 @@ impl NativeCacheConfig { let default_ttl = match &self.backend { CacheBackendConfig::Memory(config) => Some(config.default_ttl), CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::Disk(_) => None, CacheBackendConfig::AzureBlob(_) => None, }; if service.default_ttl() != default_ttl { @@ -212,6 +226,17 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::Disk(_) if service.kind() != "disk" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::Disk(config) => { + let Some(directory) = service.directory() else { + return Some("facade and native backend types must match"); + }; + let native = std::fs::canonicalize(directory).ok(); + let facade = std::fs::canonicalize(&config.directory).ok(); + (native != facade).then_some("facade and native backend directories must match") + } CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() { None => Some("facade and native backend types must match"), Some((account_url, container)) @@ -252,6 +277,21 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { }) } +#[inline(never)] +fn project_disk( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let store = backend.getattr("disk_cache")?; + if !instance_class_is(&store, "diskcache.core", "Cache")? + || !instance_class_is(&store.getattr("_disk")?, "diskcache.core", "Disk")? + { + return Ok(Err(UnsupportedCacheConfig::DiskStore)); + } + Ok(Ok(DiskCacheConfig { + directory: PathBuf::from(store.getattr("directory")?.extract::()?), + })) +} + #[inline(never)] fn project_redis( backend: &Bound<'_, PyAny>, @@ -639,8 +679,8 @@ mod tests { use litellm_cache_redis::{RedisNode, RedisTopology}; use super::{ - CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig, - RedisProtocol, + CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement, + DiskCacheConfig, NativeCacheConfig, RedisProtocol, }; use crate::cache::native::NativeResponseCache; @@ -775,6 +815,85 @@ mod tests { assert_eq!(reason.message(), "native Redis credentials require Python"); }); } + #[test] + fn projects_builtin_disk_configuration_and_rejects_custom_stores() { + Python::initialize(); + Python::attach(|py| { + let root = + std::env::temp_dir().join(format!("litellm-disk-config-{}", std::process::id())); + let directory = root.to_string_lossy(); + let disk_facade = facade( + py, + &format!( + "Cache = type('Cache', (), {{'__module__': 'diskcache.core'}})\n\ + Disk = type('Disk', (), {{'__module__': 'diskcache.core'}})\n\ + store = Cache()\n\ + store._disk = Disk()\n\ + store.directory = {directory:?}\n\ + backend = SimpleNamespace(disk_cache=store)\n\ + facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)" + ), + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&disk_facade).unwrap() + else { + panic!("disk cache should be supported"); + }; + let CacheBackendConfig::Disk(disk) = config.backend else { + panic!("expected disk configuration"); + }; + assert_eq!(disk.directory, root); + let matching = NativeResponseCache::disk(&directory).unwrap(); + assert_eq!( + (NativeCacheConfig { + policy: config.policy, + backend: CacheBackendConfig::Disk(disk), + }) + .service_mismatch(&matching), + None + ); + let other = NativeResponseCache::disk(&root.join("other").to_string_lossy()).unwrap(); + let mismatch = NativeCacheConfig { + policy: CachePolicy { + mode: "default-on".into(), + ttl: None, + namespace: None, + supported_call_types: None, + redis_flush_size: None, + semantic_cache_scope: "key".into(), + }, + backend: CacheBackendConfig::Disk(DiskCacheConfig { + directory: root.clone(), + }), + }; + assert_eq!( + mismatch.service_mismatch(&other), + Some("facade and native backend directories must match") + ); + + let custom = facade( + py, + &format!( + "CustomCache = type('CustomCache', (), {{'__module__': 'mypkg'}})\n\ + CustomDisk = type('CustomDisk', (), {{'__module__': 'mypkg'}})\n\ + store = CustomCache()\n\ + store._disk = CustomDisk()\n\ + store.directory = {directory:?}\n\ + backend = SimpleNamespace(disk_cache=store)\n\ + facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)" + ), + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&custom).unwrap() + else { + panic!("custom disk store must stay on Python"); + }; + assert_eq!( + reason.message(), + "native disk cache requires the built-in diskcache store" + ); + }); + } #[test] fn projects_cluster_startup_nodes_as_redis_topology() { diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 0e20676938f..ee6423daf6e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -34,6 +34,11 @@ struct RedisPoolGuard { attributes: RedisPoolAttributes, } +struct DiskStoreGuard { + reference: Py, + directory: String, +} + struct AzureBlobClientGuard { sync_client: Py, async_client: Py, @@ -46,7 +51,6 @@ enum ConnectionGuard { RedisPool(RedisPoolGuard), AzureBlob(AzureBlobClientGuard), } - struct RedisPoolAttributes { pool: &'static str, connection_class: &'static str, @@ -68,6 +72,7 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes { pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, + disk_store: Option, connection: ConnectionGuard, } @@ -218,6 +223,26 @@ impl RedisPoolGuard { } } +impl DiskStoreGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + let store = backend.getattr("disk_cache")?; + Ok(Self { + reference: store.clone().unbind(), + directory: store.getattr("directory")?.extract()?, + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + let store = backend.getattr("disk_cache")?; + Ok(self.reference.bind(py).is(&store) + && self.directory == store.getattr("directory")?.extract::()?) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference) + } +} + impl AzureBlobClientGuard { fn capture(backend: &Bound<'_, PyAny>) -> PyResult { let sync_client = backend.getattr("container_client")?; @@ -295,6 +320,7 @@ impl FacadeGuard { "RedisClusterCache", "redis", ), + ("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"), ("azure-blob", _) => ( "litellm.caching.azure_blob_cache", "AzureBlobCache", @@ -345,6 +371,9 @@ impl FacadeGuard { "redis_flush_size", ], )?, + disk_store: (kind == "disk") + .then(|| DiskStoreGuard::capture(&backend)) + .transpose()?, connection: ConnectionGuard::capture(kind, cluster, &backend)?, }) } @@ -357,12 +386,20 @@ impl FacadeGuard { if !self.backend.matches(py, &backend)? { return Ok(false); } + if let Some(guard) = &self.disk_store + && !guard.matches(py, &backend)? + { + return Ok(false); + } self.connection.matches(py, &backend) } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.outer.traverse(&visit)?; self.backend.traverse(&visit)?; + if let Some(guard) = &self.disk_store { + guard.traverse(&visit)?; + } self.connection.traverse(&visit) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 002cd6b7a33..6e12e12040b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -64,6 +64,18 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (directory))] + fn disk(py: Python<'_>, directory: String) -> PyResult { + let service = + release_gil(py, move || NativeResponseCache::disk(&directory)).map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[staticmethod] #[pyo3(signature = (account_url, container))] fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult { diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 2cd72a7ce14..80789cc279a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,7 +1,8 @@ -use std::{sync::Arc, time::Duration}; +use std::{path::Path, sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_azure_blob::AzureBlobCache; +use litellm_cache_disk::DiskCache; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_response::{ @@ -16,6 +17,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + Disk(Arc>>), AzureBlob(Arc>>), } @@ -47,6 +49,10 @@ impl NativeResponseCache { buffer: None, }) } + pub fn disk(directory: &str) -> Result { + let cache = DiskCache::open(directory, ResponseCacheCodec)?; + Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) + } pub async fn azure_blob(account_url: &str, container: &str) -> Result { let backend = AzureBlobCache::connect( @@ -67,7 +73,7 @@ impl NativeResponseCache { cache.backend().account_url(), cache.backend().container_name(), )), - Self::Memory(_) | Self::Redis { .. } => None, + Self::Memory(_) | Self::Redis { .. } | Self::Disk(_) => None, } } } @@ -77,6 +83,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::Disk(_) => "disk", Self::AzureBlob(_) => "azure-blob", } } @@ -85,20 +92,21 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::Disk(cache) => cache.default_ttl(), Self::AzureBlob(cache) => cache.default_ttl(), } } pub fn namespace(&self) -> Option<&str> { match self { - Self::Memory(_) | Self::AzureBlob(_) => None, + Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), } } pub fn topology(&self) -> Option<&RedisTopology> { match self { - Self::Memory(_) | Self::AzureBlob(_) => None, + Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None, Self::Redis { cache, .. } => Some(cache.backend().topology()), } } @@ -106,14 +114,14 @@ impl NativeResponseCache { pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } | Self::AzureBlob(_) => None, + Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) => None, } } pub fn max_entry_bytes(&self) -> Option { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } | Self::AzureBlob(_) => None, + Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) => None, } } @@ -127,6 +135,13 @@ impl NativeResponseCache { } } + pub fn directory(&self) -> Option<&Path> { + match self { + Self::Disk(cache) => Some(cache.backend().directory()), + Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) => None, + } + } + pub fn lookup( &self, request: &ResponseCacheRequest, @@ -135,6 +150,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup(request, now), Self::Redis { cache, .. } => cache.lookup(request, now), + Self::Disk(cache) => cache.lookup(request, now), Self::AzureBlob(cache) => cache.lookup(request, now), } } @@ -148,6 +164,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.store(request, response, now), Self::Redis { cache, .. } => cache.store(request, response, now), + Self::Disk(cache) => cache.store(request, response, now), Self::AzureBlob(cache) => cache.store(request, response, now), } } @@ -160,6 +177,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup_batch(requests, now), Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Disk(cache) => cache.lookup_batch(requests, now), Self::AzureBlob(cache) => cache.lookup_batch(requests, now), } } @@ -172,6 +190,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::Disk(cache) => cache.async_lookup(request, now).await, Self::AzureBlob(cache) => cache.async_lookup(request, now).await, } } @@ -192,6 +211,7 @@ impl NativeResponseCache { cache, buffer: Some(buffer), } => buffer.async_store(cache, request, response, now).await, + Self::Disk(cache) => cache.async_store(request, response, now).await, Self::AzureBlob(cache) => cache.async_store(request, response, now).await, } } @@ -204,6 +224,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::Disk(cache) => cache.async_lookup_batch(requests, now).await, Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await, } } @@ -216,6 +237,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_store_batch(entries, now).await, Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::Disk(cache) => cache.async_store_batch(entries, now).await, Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await, } } @@ -229,6 +251,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::Disk(cache) => cache.async_flush().await, Self::AzureBlob(cache) => cache.async_flush().await, } } @@ -237,6 +260,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::Disk(cache) => cache.test_connection().await, Self::AzureBlob(cache) => cache.test_connection().await, } } diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index a0f2396a25d..8cddc868073 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -8,10 +8,12 @@ import time import uuid import weakref from collections.abc import Generator +from pathlib import Path from types import SimpleNamespace from typing import Final, Protocol, cast from urllib.parse import urlparse +import diskcache import fakeredis import pytest import redis @@ -20,6 +22,7 @@ from azure.storage.blob import ContainerClient import litellm from litellm.caching.azure_blob_cache import AzureBlobCache from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache +from litellm.caching.disk_cache import DiskCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.rust_bridge import _native @@ -519,6 +522,112 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() +async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None: + disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path)) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} + disk_cache.disk_cache.set( + "sync", + {"timestamp": time.time(), "response": json.dumps(response)}, + ) + disk_cache.disk_cache.set("async", json.dumps({"timestamp": time.time(), "response": response})) + disk_cache.disk_cache.set("raw", json.dumps(response)) + disk_cache.disk_cache.set("invalid", "not a cache entry") + disk_cache.disk_cache.set( + "large", + {"timestamp": time.time(), "response": {"text": "x" * 70_000}}, + ) + binding: Final = _native._CacheTestResolver( + SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) + ).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("large")) == {"text": "x" * 70_000} + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored_response: Final = disk_cache.get_cache("native") + assert isinstance(stored_response, dict) + assert stored_response["response"] == response + stored, expire_time = disk_cache.disk_cache.get("native", expire_time=True) + assert stored is not None + assert time.time() < expire_time <= time.time() + 12.0 + await binding.async_store(request("no-ttl"), response) + _, no_expiry = disk_cache.disk_cache.get("no-ttl", expire_time=True) + assert no_expiry is None + + +async def test_disk_entries_survive_a_fresh_handle_and_expire_on_time(tmp_path: Path) -> None: + first: Final = _native._CacheTestResolver( + SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) + ).resolve() + await first.async_store(request("persistent"), {"value": "persistent"}) + await first.async_store({**request("expiring"), "ttl_seconds": 0.3}, {"value": "expiring"}) + fresh: Final = _native._CacheTestResolver( + SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) + ).resolve() + assert fresh.lookup(request("persistent")) == {"value": "persistent"} + assert fresh.lookup(request("expiring")) == {"value": "expiring"} + await asyncio.sleep(0.4) + assert fresh.lookup(request("expiring")) is None + assert fresh.lookup(request("persistent")) == {"value": "persistent"} + + +def test_disk_facade_registers_and_store_changes_fall_back(tmp_path: Path) -> None: + facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) + with pytest.raises(TypeError, match="directories must match"): + _native._CacheTestHandle.disk(str(tmp_path / "other"))._bind_facade(facade) + handle: Final = _native._CacheTestHandle.disk(str(tmp_path)) + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + binding.store(request("native"), {"value": "native"}) + assert facade.get_cache(cache_key="native") == {"value": "native"} + + with rebound(facade.cache, "disk_cache", diskcache.Cache(str(tmp_path))): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "native" + + class CustomDiskCache(DiskCache): + pass + + with rebound(facade, "cache", CustomDiskCache(disk_cache_dir=str(tmp_path))): + assert resolver.resolve().kind == "python_callback" + + class CustomStore(diskcache.Cache): + pass + + custom_facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) + custom_facade.cache.disk_cache = CustomStore(str(tmp_path)) + with pytest.raises(TypeError, match="built-in diskcache store"): + _native._CacheTestHandle.disk(str(tmp_path))._bind_facade(custom_facade) + + +async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: Path) -> None: + binding: Final = _native._CacheTestResolver( + SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) + ).resolve() + requests: Final = [request("hit"), request("miss"), request("disabled")] + requests[2]["controls"] = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": False, + "no_cache": False, + "no_store": False, + "use_cache": False, + } + await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) + + partial: Final = await binding.async_lookup_batch(requests) + + assert partial == { + "values": [{"value": 1}, {"value": 2}, None], + "missing_indices": [2], + } async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_natively(