From 94b8bf7fd6aa93afb8e8d9a84e047daf06ab01bc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:37:26 +0000 Subject: [PATCH 01/10] feat(rust): add native disk cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 126 +++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/cache-disk/Cargo.toml | 17 + litellm-rust/crates/cache-disk/src/cache.rs | 330 +++++++ litellm-rust/crates/cache-disk/src/lib.rs | 8 + litellm-rust/crates/cache-disk/src/pickle.rs | 48 ++ litellm-rust/crates/cache-disk/src/sqlite.rs | 816 ++++++++++++++++++ litellm-rust/crates/cache-disk/src/store.rs | 33 + litellm-rust/crates/cache-disk/tests/cache.rs | 327 +++++++ litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/cache/config.rs | 139 ++- .../crates/python-bridge/src/cache/facade.rs | 38 +- .../crates/python-bridge/src/cache/handle.rs | 12 + .../crates/python-bridge/src/cache/native.rs | 31 +- tests/test_litellm_rust/test_cache.py | 111 +++ 15 files changed, 2026 insertions(+), 12 deletions(-) create mode 100644 litellm-rust/crates/cache-disk/Cargo.toml create mode 100644 litellm-rust/crates/cache-disk/src/cache.rs create mode 100644 litellm-rust/crates/cache-disk/src/lib.rs create mode 100644 litellm-rust/crates/cache-disk/src/pickle.rs create mode 100644 litellm-rust/crates/cache-disk/src/sqlite.rs create mode 100644 litellm-rust/crates/cache-disk/src/store.rs create mode 100644 litellm-rust/crates/cache-disk/tests/cache.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..554245dbb6e 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1318,6 +1318,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" @@ -1369,6 +1381,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" @@ -1833,11 +1851,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" @@ -2218,6 +2257,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" @@ -2376,6 +2421,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" @@ -2464,6 +2520,19 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-disk" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "rand 0.8.7", + "rusqlite", + "serde-pickle", + "serde_json", + "tempfile", + "tokio", +] + [[package]] name = "litellm-cache-memory" version = "0.1.0" @@ -2666,6 +2735,7 @@ dependencies = [ "litellm-auth", "litellm-auth-gcp", "litellm-cache", + "litellm-cache-disk", "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", @@ -3901,6 +3971,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" @@ -3941,6 +4021,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" @@ -4198,6 +4293,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" @@ -4425,6 +4533,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" @@ -5170,6 +5290,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 570d0dd3568..53db97b2d01 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -29,6 +29,7 @@ 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-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..023055e391a --- /dev/null +++ b/litellm-rust/crates/cache-disk/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "litellm-cache-disk" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +rand.workspace = true +rusqlite = { version = "0.40", features = ["bundled"] } +serde-pickle = "1.2" +serde_json.workspace = true +tokio.workspace = true + +[dev-dependencies] +tempfile = "3.27.0" 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..bcce3a525b1 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -0,0 +1,330 @@ +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 serde_json::Value; + +use crate::{DiskStore, DiskcacheSqliteStore, StoredValue, pickle}; + +pub struct DiskCache { + store: Arc, + codec: S, +} + +impl DiskCache { + pub fn open(directory: impl AsRef, codec: S) -> Result { + Ok(Self { + store: Arc::new(DiskcacheSqliteStore::open(directory)?), + codec, + }) + } +} + +impl DiskCache { + pub fn with_store(store: D, codec: S) -> Self { + Self { + store: Arc::new(store), + codec, + } + } + + pub fn directory(&self) -> &Path { + self.store.directory() + } + + fn decode_stored(&self, value: StoredValue) -> Result, Error> { + let Some(bytes) = payload(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 = StoredValue::Bytes(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 = StoredValue::Bytes(self.codec.encode(&value)?); + let expire_time = context.ttl.map(|ttl| unix_now() + ttl.as_secs_f64()); + let key = key.to_string(); + Self::run_blocking(Arc::clone(&self.store), move |store| { + 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, StoredValue::Bytes(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> CounterCache for DiskCache { + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { + increment(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(); + Self::run_blocking(Arc::clone(&self.store), move |store| { + increment(store, &key, amount, context.ttl) + }) + .await + } +} + +fn increment( + store: &D, + key: &str, + amount: f64, + ttl: Option, +) -> Result { + let mut result = None; + let mut apply = |current: Option| { + let initial = match current { + Some(StoredValue::Integer(value)) => value as f64, + Some(StoredValue::Pickle(value)) => match pickle::decode(&value)? { + Value::Number(value) => value + .as_i64() + .map(|value| value as f64) + .or_else(|| value.as_u64().map(|value| value as f64)) + .unwrap_or_default(), + _ => 0.0, + }, + _ => 0.0, + }; + let value = initial + amount; + let stored = if value.fract() == 0.0 && value >= i64::MIN as f64 && value <= i64::MAX as f64 + { + StoredValue::Integer(value as i64) + } else { + StoredValue::Float(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 payload(value: StoredValue) -> Result>, Error> { + match value { + StoredValue::Bytes(value) if value.is_empty() => Ok(None), + StoredValue::Bytes(value) => Ok(Some(value)), + StoredValue::Text(value) if value.is_empty() => Ok(None), + StoredValue::Text(value) => Ok(Some(value.into_bytes())), + StoredValue::Integer(0) => Ok(None), + StoredValue::Integer(value) => Ok(Some(value.to_string().into_bytes())), + StoredValue::Float(0.0) => Ok(None), + StoredValue::Float(value) => serde_json::to_vec(&value) + .map(Some) + .map_err(|_| Error::InvalidEntry), + StoredValue::Pickle(value) => { + let value = pickle::decode(&value)?; + if is_falsy(&value) { + Ok(None) + } else { + serde_json::to_vec(&value) + .map(Some) + .map_err(|_| Error::InvalidEntry) + } + } + } +} + +fn is_falsy(value: &Value) -> bool { + match value { + Value::Null | Value::Bool(false) => true, + Value::Number(value) => value.as_f64().is_some_and(|value| value == 0.0), + Value::String(value) => value.is_empty(), + Value::Array(value) => value.is_empty(), + Value::Object(value) => value.is_empty(), + Value::Bool(true) => false, + } +} + +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..189d12ea793 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/lib.rs @@ -0,0 +1,8 @@ +mod cache; +mod pickle; +mod sqlite; +mod store; + +pub use cache::DiskCache; +pub use sqlite::DiskcacheSqliteStore; +pub use store::{DiskStore, StoredValue}; diff --git a/litellm-rust/crates/cache-disk/src/pickle.rs b/litellm-rust/crates/cache-disk/src/pickle.rs new file mode 100644 index 00000000000..1d31e0bbeb6 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/pickle.rs @@ -0,0 +1,48 @@ +use litellm_cache::Error; +use serde_json::{Map, Number, Value}; + +pub(crate) fn decode(bytes: &[u8]) -> Result { + let value = serde_pickle::value_from_slice(bytes, Default::default()) + .map_err(|_| Error::InvalidEntry)?; + convert(value) +} + +fn convert(value: serde_pickle::Value) -> Result { + match value { + serde_pickle::Value::None => Ok(Value::Null), + serde_pickle::Value::Bool(value) => Ok(Value::Bool(value)), + serde_pickle::Value::I64(value) => Ok(Value::Number(value.into())), + serde_pickle::Value::Int(value) => { + if let Ok(value) = value.to_string().parse::() { + Ok(Value::Number(value.into())) + } else if let Ok(value) = value.to_string().parse::() { + Ok(Value::Number(value.into())) + } else { + Err(Error::InvalidEntry) + } + } + serde_pickle::Value::F64(value) => Number::from_f64(value) + .map(Value::Number) + .ok_or(Error::InvalidEntry), + serde_pickle::Value::String(value) => Ok(Value::String(value)), + serde_pickle::Value::List(values) | serde_pickle::Value::Tuple(values) => values + .into_iter() + .map(convert) + .collect::, _>>() + .map(Value::Array), + serde_pickle::Value::Set(values) | serde_pickle::Value::FrozenSet(values) => values + .into_iter() + .map(|value| convert(value.into_value())) + .collect::, _>>() + .map(Value::Array), + serde_pickle::Value::Dict(values) => values + .into_iter() + .map(|(key, value)| match key { + serde_pickle::HashableValue::String(key) => Ok((key, convert(value)?)), + _ => Err(Error::InvalidEntry), + }) + .collect::, _>>() + .map(Value::Object), + serde_pickle::Value::Bytes(_) => Err(Error::InvalidEntry), + } +} 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..ff5a44b5c38 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/sqlite.rs @@ -0,0 +1,816 @@ +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 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..1e3e39db0c8 --- /dev/null +++ b/litellm-rust/crates/cache-disk/tests/cache.rs @@ -0,0 +1,327 @@ +use std::{ + fs, + path::{Path, PathBuf}, + sync::Arc, + thread, + time::Duration, +}; + +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CounterCache, DeleteCache, ExactCacheContext, FlushCache, + JsonCodec, +}; +use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue}; +use rusqlite::Connection; +use serde_json::json; +use tempfile::TempDir; + +fn store() -> (TempDir, DiskcacheSqliteStore) { + let directory = tempfile::tempdir().unwrap(); + let store = DiskcacheSqliteStore::open(directory.path()).unwrap(); + (directory, store) +} + +fn cache(directory: &Path) -> DiskCache> { + DiskCache::open(directory, JsonCodec::new()).unwrap() +} + +fn value_files(directory: &Path) -> 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(directory, &mut files); + files +} + +#[test] +fn roundtrip_persists_and_reopens() { + let directory = tempfile::tempdir().unwrap(); + let context = ExactCacheContext::default(); + let opened = cache(directory.path()); + 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 = cache(directory.path()); + assert_eq!( + reopened.get_cache("key", &context).unwrap(), + Some(json!({"answer": 42})) + ); +} + +#[test] +fn ttl_and_expired_culling_match_cache_contract() { + let (directory, store) = 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(); + let connection = Connection::open(directory.path().join("cache.db")).unwrap(); + assert_eq!( + connection + .query_row("SELECT COUNT(*) FROM Cache", [], |row| row.get::<_, i64>(0)) + .unwrap(), + 1 + ); + assert_eq!( + connection + .query_row( + "SELECT value FROM Settings WHERE key = 'count'", + [], + |row| row.get::<_, i64>(0) + ) + .unwrap(), + 1 + ); +} + +#[test] +fn batch_preserves_order_and_classifies_misses_and_invalid_values() { + let (directory, store) = 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 cache = cache(directory.path()); + let entries = 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 + ] + ); +} + +#[test] +fn falsy_values_are_misses_and_protocol_five_pickle_decodes() { + let (directory, store) = store(); + for (key, value) in [ + ("empty-bytes", StoredValue::Bytes(Vec::new())), + ("empty-text", StoredValue::Text(String::new())), + ("zero-int", StoredValue::Integer(0)), + ("zero-float", StoredValue::Float(0.0)), + ( + "empty-pickle", + StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]), + ), + ] { + store.set(key, value, None, 0.0).unwrap(); + } + store + .set( + "pickle", + StoredValue::Pickle( + b"\x80\x05\x95\x30\x00\x00\x00\x00\x00\x00\x00\x7d\x94\x28\x8c\x09timestamp\x94G\x3f\xf8\x00\x00\x00\x00\x00\x00\x8c\x08response\x94\x8c\x08{\"a\": 1}\x94u." + .to_vec(), + ), + None, + 0.0, + ) + .unwrap(); + let cache = cache(directory.path()); + for key in [ + "empty-bytes", + "empty-text", + "zero-int", + "zero-float", + "empty-pickle", + ] { + assert_eq!( + cache.get_cache(key, &ExactCacheContext::default()).unwrap(), + None + ); + } + assert_eq!( + cache + .get_cache("pickle", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"timestamp": 1.5, "response": "{\"a\": 1}"})) + ); +} + +#[test] +fn counters_use_atomic_native_values_and_ignore_invalid_initial_values() { + let (directory, store) = store(); + store + .set("counter", StoredValue::Integer(2), None, 0.0) + .unwrap(); + store + .set( + "invalid", + StoredValue::Text("not a number".into()), + None, + 0.0, + ) + .unwrap(); + store + .set( + "pickle-counter", + StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e]), + None, + 0.0, + ) + .unwrap(); + let cache = DiskCache::open(directory.path(), JsonCodec::::new()).unwrap(); + assert_eq!( + cache + .increment_cache("counter", 1.5, ExactCacheContext::default()) + .unwrap(), + 3.5 + ); + assert_eq!( + cache + .increment_cache("invalid", 2.0, ExactCacheContext::default()) + .unwrap(), + 2.0 + ); + assert_eq!( + cache + .increment_cache("pickle-counter", 1.0, ExactCacheContext::default()) + .unwrap(), + 3.0 + ); + let connection = Connection::open(directory.path().join("cache.db")).unwrap(); + assert_eq!( + connection + .query_row( + "SELECT typeof(value) FROM Cache WHERE key = 'counter'", + [], + |row| row.get::<_, String>(0) + ) + .unwrap(), + "real" + ); +} + +#[test] +fn counters_are_atomic_across_concurrent_callers() { + let directory = tempfile::tempdir().unwrap(); + let cache = Arc::new(DiskCache::open(directory.path(), JsonCodec::::new()).unwrap()); + 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 + ); +} + +#[test] +fn delete_flush_and_spilled_file_replacement_clean_up_storage() { + let (directory, store) = store(); + let large = vec![b'x'; 32 * 1024]; + store + .set("large", StoredValue::Bytes(large.clone()), None, 0.0) + .unwrap(); + assert_eq!(value_files(directory.path()).len(), 1); + store + .set( + "large", + StoredValue::Bytes(vec![b'y'; 32 * 1024]), + None, + 0.0, + ) + .unwrap(); + assert_eq!(value_files(directory.path()).len(), 1); + store.pop("large", 0.0).unwrap(); + assert!(value_files(directory.path()).is_empty()); + store + .set("a", StoredValue::Bytes(large.clone()), None, 0.0) + .unwrap(); + store + .set("b", StoredValue::Bytes(large), None, 0.0) + .unwrap(); + store.clear().unwrap(); + assert!(value_files(directory.path()).is_empty()); +} + +#[tokio::test] +async fn async_operations_connection_and_delete_match_sync_operations() { + let directory = tempfile::tempdir().unwrap(); + let cache = cache(directory.path()); + 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/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..71846dcacbd 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -23,6 +23,7 @@ bytes.workspace = true litellm-cache.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 0e7d6aee11d..17c4cb40903 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 pyo3::{ @@ -25,6 +25,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, @@ -76,6 +80,7 @@ pub(super) struct RedisCacheConfig { pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + Disk(DiskCacheConfig), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -90,6 +95,7 @@ pub(super) enum UnsupportedCacheConfig { RedisCredentials, RedisConnection, RedisOption, + DiskStore, } impl UnsupportedCacheConfig { @@ -100,6 +106,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", } } } @@ -142,11 +149,17 @@ 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::RedisSemantic | CacheType::ValkeySemantic | CacheType::S3 - | CacheType::Disk | CacheType::QdrantSemantic | CacheType::AzureBlob | CacheType::Gcs, @@ -158,12 +171,12 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) - { + let default_ttl = match &self.backend { + CacheBackendConfig::Memory(config) => Some(config.default_ttl), + CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::Disk(_) => None, + }; + if service.default_ttl() != default_ttl { return Some("facade and native backend default TTLs must match"); } match &self.backend { @@ -185,6 +198,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") + } } } } @@ -201,6 +225,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>, @@ -468,8 +507,8 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; use super::{ - CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig, - RedisProtocol, + CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement, + DiskCacheConfig, NativeCacheConfig, RedisProtocol, }; use crate::cache::native::NativeResponseCache; @@ -591,4 +630,84 @@ 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" + ); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f2f86c14b37..a22a1754169 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -32,10 +32,16 @@ struct RedisPoolGuard { max_connections: usize, } +struct DiskStoreGuard { + reference: Py, + directory: String, +} + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, redis_pool: Option, + disk_store: Option, } impl ObjectGuard { @@ -176,6 +182,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 FacadeGuard { pub(super) fn capture( py: Python<'_>, @@ -192,6 +218,7 @@ impl FacadeGuard { let (module, name, cache_kind) = match kind { "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + "disk" => ("litellm.caching.disk_cache", "DiskCache", "disk"), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -240,6 +267,9 @@ impl FacadeGuard { redis_pool: (kind == "redis") .then(|| RedisPoolGuard::capture(&backend)) .transpose()?, + disk_store: (kind == "disk") + .then(|| DiskStoreGuard::capture(&backend)) + .transpose()?, }) } @@ -253,7 +283,10 @@ impl FacadeGuard { } match &self.redis_pool { Some(guard) => guard.matches(py, &backend), - None => Ok(true), + None => match &self.disk_store { + Some(guard) => guard.matches(py, &backend), + None => Ok(true), + }, } } @@ -263,6 +296,9 @@ impl FacadeGuard { if let Some(guard) = &self.redis_pool { guard.traverse(&visit)?; } + if let Some(guard) = &self.disk_store { + guard.traverse(&visit)?; + } Ok(()) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..2cb2045577c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -51,6 +51,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(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a9475429e45..14993a6e8e3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,6 +1,7 @@ -use std::{sync::Arc, time::Duration}; +use std::{path::Path, sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache_disk::DiskCache; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ @@ -15,6 +16,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + Disk(Arc>>), } impl NativeResponseCache { @@ -43,6 +45,11 @@ 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))))) + } } impl NativeResponseCache { @@ -50,6 +57,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::Disk(_) => "disk", } } @@ -57,6 +65,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::Disk(cache) => cache.default_ttl(), } } @@ -64,6 +73,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), + Self::Disk(_) => None, } } @@ -71,6 +81,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), Self::Redis { .. } => None, + Self::Disk(_) => None, } } @@ -78,6 +89,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), Self::Redis { .. } => None, + Self::Disk(_) => None, } } @@ -87,10 +99,18 @@ impl NativeResponseCache { cache, buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), }, + disk @ Self::Disk(_) => disk, memory => memory, } } + pub fn directory(&self) -> Option<&Path> { + match self { + Self::Disk(cache) => Some(cache.backend().directory()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + pub fn lookup( &self, request: &ResponseCacheRequest, @@ -99,6 +119,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), } } @@ -111,6 +132,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), } } @@ -122,6 +144,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), } } @@ -133,6 +156,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, } } @@ -152,6 +176,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, } } @@ -163,6 +188,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, } } @@ -174,6 +200,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, } } @@ -186,6 +213,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::Disk(cache) => cache.async_flush().await, } } @@ -193,6 +221,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, } } } diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c35cb1a20fb..13f5f686464 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -6,16 +6,19 @@ import threading import time 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 import litellm 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.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType @@ -393,3 +396,111 @@ 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], + } From cf8a211048368b8fea9f161dae5eb1b8ef4b2e30 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:52:03 +0000 Subject: [PATCH 02/10] test(rust): pin disk counter restart after a fractional value Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-disk/tests/cache.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/litellm-rust/crates/cache-disk/tests/cache.rs b/litellm-rust/crates/cache-disk/tests/cache.rs index 1e3e39db0c8..c6cd6006148 100644 --- a/litellm-rust/crates/cache-disk/tests/cache.rs +++ b/litellm-rust/crates/cache-disk/tests/cache.rs @@ -231,6 +231,12 @@ fn counters_use_atomic_native_values_and_ignore_invalid_initial_values() { .unwrap(), "real" ); + assert_eq!( + cache + .increment_cache("counter", 1.0, ExactCacheContext::default()) + .unwrap(), + 1.0 + ); } #[test] From c1382086d69f176114003a5fcf32ac2a8cbaaa64 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:05:21 +0000 Subject: [PATCH 03/10] ci(rust): raise native wheel size gate to 35 MB Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/verify_linux_native_wheel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 4fb8f068eb0..d86ad7b15c6 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -205,7 +205,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 25_000_000 + native_size_limit: Final = 35_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), @@ -222,7 +222,7 @@ def main( ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), - ("Native extension does not exceed 25 MB", native_size_within_limit), + ("Native extension does not exceed 35 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -267,7 +267,7 @@ def main( ), ( not native_size_within_limit, - f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB", + f"native extension exceeds 35 MB: {native_member.file_size / 1_000_000:.2f} MB", ), (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), ) From 3ce436af5ed72281214df3449ebf31323e8d3874 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:22:22 +0000 Subject: [PATCH 04/10] refactor(cache-disk): isolate python compatibility behind a value adapter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/cache-disk/Cargo.toml | 2 + litellm-rust/crates/cache-disk/src/adapter.rs | 10 + litellm-rust/crates/cache-disk/src/cache.rs | 112 ++---- litellm-rust/crates/cache-disk/src/lib.rs | 5 +- litellm-rust/crates/cache-disk/src/pickle.rs | 48 --- .../crates/cache-disk/src/python/mod.rs | 81 ++++ .../crates/cache-disk/src/python/value.rs | 173 ++++++++ litellm-rust/crates/cache-disk/tests/cache.rs | 369 +++++++++++------- .../crates/cache-disk/tests/python_compat.rs | 105 +++++ 10 files changed, 643 insertions(+), 264 deletions(-) create mode 100644 litellm-rust/crates/cache-disk/src/adapter.rs delete mode 100644 litellm-rust/crates/cache-disk/src/pickle.rs create mode 100644 litellm-rust/crates/cache-disk/src/python/mod.rs create mode 100644 litellm-rust/crates/cache-disk/src/python/value.rs create mode 100644 litellm-rust/crates/cache-disk/tests/python_compat.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 554245dbb6e..0a1208787eb 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2525,7 +2525,9 @@ name = "litellm-cache-disk" version = "0.1.0" dependencies = [ "litellm-cache", + "py_literal", "rand 0.8.7", + "rstest", "rusqlite", "serde-pickle", "serde_json", diff --git a/litellm-rust/crates/cache-disk/Cargo.toml b/litellm-rust/crates/cache-disk/Cargo.toml index 023055e391a..b96994b3b55 100644 --- a/litellm-rust/crates/cache-disk/Cargo.toml +++ b/litellm-rust/crates/cache-disk/Cargo.toml @@ -7,6 +7,7 @@ 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" @@ -14,4 +15,5 @@ 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 index bcce3a525b1..98182d93968 100644 --- a/litellm-rust/crates/cache-disk/src/cache.rs +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -8,28 +8,42 @@ use litellm_cache::{ BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, }; -use serde_json::Value; -use crate::{DiskStore, DiskcacheSqliteStore, StoredValue, pickle}; +use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter}; -pub struct DiskCache { +pub struct DiskCache { store: Arc, + adapter: Arc, codec: S, } impl DiskCache { + #[allow(clippy::default_constructed_unit_structs)] pub fn open(directory: impl AsRef, codec: S) -> Result { Ok(Self { store: Arc::new(DiskcacheSqliteStore::open(directory)?), + adapter: Arc::new(PythonDiskCacheAdapter::default()), codec, }) } } -impl DiskCache { +impl DiskCache { + #[allow(clippy::default_constructed_unit_structs)] pub fn with_store(store: D, codec: S) -> Self { Self { store: Arc::new(store), + adapter: Arc::new(PythonDiskCacheAdapter::default()), + codec, + } + } +} + +impl DiskCache { + pub fn with_adapter(store: D, adapter: A, codec: S) -> Self { + Self { + store: Arc::new(store), + adapter: Arc::new(adapter), codec, } } @@ -39,7 +53,7 @@ impl DiskCache { } fn decode_stored(&self, value: StoredValue) -> Result, Error> { - let Some(bytes) = payload(value)? else { + let Some(bytes) = self.adapter.read(value)? else { return Ok(None); }; self.codec.decode(&bytes).map(Some) @@ -56,7 +70,7 @@ impl DiskCache { } } -impl BaseCache for DiskCache { +impl BaseCache for DiskCache { type Value = S::Value; type Context = ExactCacheContext; @@ -70,7 +84,7 @@ impl BaseCache for DiskCache { value: Self::Value, context: &Self::Context, ) -> Result<(), Error> { - let value = StoredValue::Bytes(self.codec.encode(&value)?); + 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()) } @@ -89,7 +103,7 @@ impl BaseCache for DiskCache { value: Self::Value, context: ExactCacheContext, ) -> Result<(), Error> { - let value = StoredValue::Bytes(self.codec.encode(&value)?); + let value = self.adapter.write(self.codec.encode(&value)?); let expire_time = context.ttl.map(|ttl| unix_now() + ttl.as_secs_f64()); let key = key.to_string(); Self::run_blocking(Arc::clone(&self.store), move |store| { @@ -124,7 +138,7 @@ impl BaseCache for DiskCache { .map(|(key, value)| { self.codec .encode(&value) - .map(|value| (key, StoredValue::Bytes(value))) + .map(|value| (key, self.adapter.write(value))) }) .collect::, _>>()?; let expire_after = context.ttl; @@ -162,7 +176,7 @@ impl BaseCache for DiskCache { } } -impl BatchCache for DiskCache { +impl BatchCache for DiskCache { fn batch_get_cache( &self, keys: &[String], @@ -204,7 +218,7 @@ impl BatchCache for DiskCache { } } -impl DeleteCache for DiskCache { +impl DeleteCache for DiskCache { fn delete_cache(&self, key: &str) -> Result<(), Error> { self.store.pop(key, unix_now()).map(|_| ()) } @@ -218,7 +232,7 @@ impl DeleteCache for DiskCache { } } -impl FlushCache for DiskCache { +impl FlushCache for DiskCache { fn flush_cache(&self) -> Result<(), Error> { self.store.clear() } @@ -228,14 +242,22 @@ impl FlushCache for DiskCache { } } -impl, D: DiskStore> CounterCache for DiskCache { +impl, D: DiskStore, A: ValueAdapter> CounterCache + for DiskCache +{ fn increment_cache( &self, key: &str, amount: f64, context: ExactCacheContext, ) -> Result { - increment(self.store.as_ref(), key, amount, context.ttl) + increment( + self.adapter.as_ref(), + self.store.as_ref(), + key, + amount, + context.ttl, + ) } async fn async_increment( @@ -245,14 +267,16 @@ impl, D: DiskStore> CounterCache for DiskCache context: ExactCacheContext, ) -> Result { let key = key.to_string(); + let adapter = Arc::clone(&self.adapter); Self::run_blocking(Arc::clone(&self.store), move |store| { - increment(store, &key, amount, context.ttl) + increment(adapter.as_ref(), store, &key, amount, context.ttl) }) .await } } -fn increment( +fn increment( + adapter: &A, store: &D, key: &str, amount: f64, @@ -260,25 +284,9 @@ fn increment( ) -> Result { let mut result = None; let mut apply = |current: Option| { - let initial = match current { - Some(StoredValue::Integer(value)) => value as f64, - Some(StoredValue::Pickle(value)) => match pickle::decode(&value)? { - Value::Number(value) => value - .as_i64() - .map(|value| value as f64) - .or_else(|| value.as_u64().map(|value| value as f64)) - .unwrap_or_default(), - _ => 0.0, - }, - _ => 0.0, - }; + let initial = adapter.counter_seed(current)?; let value = initial + amount; - let stored = if value.fract() == 0.0 && value >= i64::MIN as f64 && value <= i64::MAX as f64 - { - StoredValue::Integer(value as i64) - } else { - StoredValue::Float(value) - }; + let stored = adapter.counter_value(value); result = Some(value); Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64()))) }; @@ -286,42 +294,6 @@ fn increment( result.ok_or(Error::InvalidEntry) } -fn payload(value: StoredValue) -> Result>, Error> { - match value { - StoredValue::Bytes(value) if value.is_empty() => Ok(None), - StoredValue::Bytes(value) => Ok(Some(value)), - StoredValue::Text(value) if value.is_empty() => Ok(None), - StoredValue::Text(value) => Ok(Some(value.into_bytes())), - StoredValue::Integer(0) => Ok(None), - StoredValue::Integer(value) => Ok(Some(value.to_string().into_bytes())), - StoredValue::Float(0.0) => Ok(None), - StoredValue::Float(value) => serde_json::to_vec(&value) - .map(Some) - .map_err(|_| Error::InvalidEntry), - StoredValue::Pickle(value) => { - let value = pickle::decode(&value)?; - if is_falsy(&value) { - Ok(None) - } else { - serde_json::to_vec(&value) - .map(Some) - .map_err(|_| Error::InvalidEntry) - } - } - } -} - -fn is_falsy(value: &Value) -> bool { - match value { - Value::Null | Value::Bool(false) => true, - Value::Number(value) => value.as_f64().is_some_and(|value| value == 0.0), - Value::String(value) => value.is_empty(), - Value::Array(value) => value.is_empty(), - Value::Object(value) => value.is_empty(), - Value::Bool(true) => false, - } -} - fn unix_now() -> f64 { SystemTime::now() .duration_since(UNIX_EPOCH) diff --git a/litellm-rust/crates/cache-disk/src/lib.rs b/litellm-rust/crates/cache-disk/src/lib.rs index 189d12ea793..9b2ffc24915 100644 --- a/litellm-rust/crates/cache-disk/src/lib.rs +++ b/litellm-rust/crates/cache-disk/src/lib.rs @@ -1,8 +1,11 @@ +mod adapter; mod cache; -mod pickle; +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/pickle.rs b/litellm-rust/crates/cache-disk/src/pickle.rs deleted file mode 100644 index 1d31e0bbeb6..00000000000 --- a/litellm-rust/crates/cache-disk/src/pickle.rs +++ /dev/null @@ -1,48 +0,0 @@ -use litellm_cache::Error; -use serde_json::{Map, Number, Value}; - -pub(crate) fn decode(bytes: &[u8]) -> Result { - let value = serde_pickle::value_from_slice(bytes, Default::default()) - .map_err(|_| Error::InvalidEntry)?; - convert(value) -} - -fn convert(value: serde_pickle::Value) -> Result { - match value { - serde_pickle::Value::None => Ok(Value::Null), - serde_pickle::Value::Bool(value) => Ok(Value::Bool(value)), - serde_pickle::Value::I64(value) => Ok(Value::Number(value.into())), - serde_pickle::Value::Int(value) => { - if let Ok(value) = value.to_string().parse::() { - Ok(Value::Number(value.into())) - } else if let Ok(value) = value.to_string().parse::() { - Ok(Value::Number(value.into())) - } else { - Err(Error::InvalidEntry) - } - } - serde_pickle::Value::F64(value) => Number::from_f64(value) - .map(Value::Number) - .ok_or(Error::InvalidEntry), - serde_pickle::Value::String(value) => Ok(Value::String(value)), - serde_pickle::Value::List(values) | serde_pickle::Value::Tuple(values) => values - .into_iter() - .map(convert) - .collect::, _>>() - .map(Value::Array), - serde_pickle::Value::Set(values) | serde_pickle::Value::FrozenSet(values) => values - .into_iter() - .map(|value| convert(value.into_value())) - .collect::, _>>() - .map(Value::Array), - serde_pickle::Value::Dict(values) => values - .into_iter() - .map(|(key, value)| match key { - serde_pickle::HashableValue::String(key) => Ok((key, convert(value)?)), - _ => Err(Error::InvalidEntry), - }) - .collect::, _>>() - .map(Value::Object), - serde_pickle::Value::Bytes(_) => Err(Error::InvalidEntry), - } -} 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..5c9b57de000 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/python/mod.rs @@ -0,0 +1,81 @@ +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::from_json(serde_json::Value::Number(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> { + let raw = match &value { + StoredValue::Text(value) => Some(value.as_bytes().to_vec()), + StoredValue::Bytes(value) => Some(value.clone()), + StoredValue::Integer(_) | StoredValue::Float(_) | StoredValue::Pickle(_) => None, + }; + let Some(value) = Self::python_get_cache(value)? else { + return Ok(None); + }; + if let Some(raw) = raw { + return Ok(Some(raw)); + } + 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/tests/cache.rs b/litellm-rust/crates/cache-disk/tests/cache.rs index c6cd6006148..4fb77d233ec 100644 --- a/litellm-rust/crates/cache-disk/tests/cache.rs +++ b/litellm-rust/crates/cache-disk/tests/cache.rs @@ -7,46 +7,92 @@ use std::{ }; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CounterCache, DeleteCache, ExactCacheContext, FlushCache, - JsonCodec, + BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext, + FlushCache, JsonCodec, }; -use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue}; +use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter}; +use rstest::{fixture, rstest}; use rusqlite::Connection; -use serde_json::json; +use serde_json::{Value, json}; use tempfile::TempDir; -fn store() -> (TempDir, DiskcacheSqliteStore) { - let directory = tempfile::tempdir().unwrap(); - let store = DiskcacheSqliteStore::open(directory.path()).unwrap(); - (directory, store) +struct Sandbox { + directory: TempDir, } -fn cache(directory: &Path) -> DiskCache> { - DiskCache::open(directory, JsonCodec::new()).unwrap() +#[fixture] +fn sandbox() -> Sandbox { + Sandbox { + directory: tempfile::tempdir().unwrap(), + } } -fn value_files(directory: &Path) -> 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); +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 + } +} + +#[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), + } } - let mut files = Vec::new(); - visit(directory, &mut files); - files + 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) + } + } } -#[test] -fn roundtrip_persists_and_reopens() { - let directory = tempfile::tempdir().unwrap(); +#[rstest] +fn roundtrip_persists_and_reopens(sandbox: Sandbox) { let context = ExactCacheContext::default(); - let opened = cache(directory.path()); + let opened = sandbox.cache::(); opened .set_cache("key", json!({"answer": 42}), &context) .unwrap(); @@ -55,16 +101,16 @@ fn roundtrip_persists_and_reopens() { Some(json!({"answer": 42})) ); drop(opened); - let reopened = cache(directory.path()); + let reopened = sandbox.cache::(); assert_eq!( reopened.get_cache("key", &context).unwrap(), Some(json!({"answer": 42})) ); } -#[test] -fn ttl_and_expired_culling_match_cache_contract() { - let (directory, store) = store(); +#[rstest] +fn ttl_and_expired_culling_match_cache_contract(sandbox: Sandbox) { + let store = sandbox.store(); store .set( "expired", @@ -77,15 +123,16 @@ fn ttl_and_expired_culling_match_cache_contract() { store .set("new", StoredValue::Bytes(b"new".to_vec()), None, 11.0) .unwrap(); - let connection = Connection::open(directory.path().join("cache.db")).unwrap(); assert_eq!( - connection + sandbox + .db() .query_row("SELECT COUNT(*) FROM Cache", [], |row| row.get::<_, i64>(0)) .unwrap(), 1 ); assert_eq!( - connection + sandbox + .db() .query_row( "SELECT value FROM Settings WHERE key = 'count'", [], @@ -96,9 +143,9 @@ fn ttl_and_expired_culling_match_cache_contract() { ); } -#[test] -fn batch_preserves_order_and_classifies_misses_and_invalid_values() { - let (directory, store) = store(); +#[rstest] +fn batch_preserves_order_and_classifies_misses_and_invalid_values(sandbox: Sandbox) { + let store = sandbox.store(); store .set( "hit", @@ -115,8 +162,8 @@ fn batch_preserves_order_and_classifies_misses_and_invalid_values() { 0.0, ) .unwrap(); - let cache = cache(directory.path()); - let entries = cache + let entries = sandbox + .cache::() .batch_get_cache( &["hit".into(), "missing".into(), "invalid".into()], &ExactCacheContext::default(), @@ -132,117 +179,74 @@ fn batch_preserves_order_and_classifies_misses_and_invalid_values() { ); } -#[test] -fn falsy_values_are_misses_and_protocol_five_pickle_decodes() { - let (directory, store) = store(); - for (key, value) in [ - ("empty-bytes", StoredValue::Bytes(Vec::new())), - ("empty-text", StoredValue::Text(String::new())), - ("zero-int", StoredValue::Integer(0)), - ("zero-float", StoredValue::Float(0.0)), - ( - "empty-pickle", - StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]), - ), - ] { - store.set(key, value, None, 0.0).unwrap(); - } - store - .set( - "pickle", - StoredValue::Pickle( - b"\x80\x05\x95\x30\x00\x00\x00\x00\x00\x00\x00\x7d\x94\x28\x8c\x09timestamp\x94G\x3f\xf8\x00\x00\x00\x00\x00\x00\x8c\x08response\x94\x8c\x08{\"a\": 1}\x94u." - .to_vec(), - ), - None, - 0.0, - ) - .unwrap(); - let cache = cache(directory.path()); - for key in [ - "empty-bytes", - "empty-text", - "zero-int", - "zero-float", - "empty-pickle", - ] { - assert_eq!( - cache.get_cache(key, &ExactCacheContext::default()).unwrap(), - None - ); - } +#[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!( - cache - .get_cache("pickle", &ExactCacheContext::default()) + sandbox + .cache::() + .get_cache("key", &ExactCacheContext::default()) .unwrap(), - Some(json!({"timestamp": 1.5, "response": "{\"a\": 1}"})) + None ); } -#[test] -fn counters_use_atomic_native_values_and_ignore_invalid_initial_values() { - let (directory, store) = store(); - store - .set("counter", StoredValue::Integer(2), None, 0.0) - .unwrap(); - store - .set( - "invalid", - StoredValue::Text("not a number".into()), - None, - 0.0, - ) - .unwrap(); - store - .set( - "pickle-counter", - StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e]), - None, - 0.0, - ) - .unwrap(); - let cache = DiskCache::open(directory.path(), JsonCodec::::new()).unwrap(); +#[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", 1.5, ExactCacheContext::default()) + .increment_cache("counter", amount, ExactCacheContext::default()) .unwrap(), - 3.5 + expected ); assert_eq!( - cache - .increment_cache("invalid", 2.0, ExactCacheContext::default()) - .unwrap(), - 2.0 - ); - assert_eq!( - cache - .increment_cache("pickle-counter", 1.0, ExactCacheContext::default()) - .unwrap(), - 3.0 - ); - let connection = Connection::open(directory.path().join("cache.db")).unwrap(); - assert_eq!( - connection + sandbox + .db() .query_row( "SELECT typeof(value) FROM Cache WHERE key = 'counter'", [], |row| row.get::<_, String>(0) ) .unwrap(), - "real" - ); - assert_eq!( - cache - .increment_cache("counter", 1.0, ExactCacheContext::default()) - .unwrap(), - 1.0 + sqlite_type ); } -#[test] -fn counters_are_atomic_across_concurrent_callers() { - let directory = tempfile::tempdir().unwrap(); - let cache = Arc::new(DiskCache::open(directory.path(), JsonCodec::::new()).unwrap()); +#[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); @@ -266,15 +270,88 @@ fn counters_are_atomic_across_concurrent_callers() { ); } -#[test] -fn delete_flush_and_spilled_file_replacement_clean_up_storage() { - let (directory, store) = store(); +#[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]; - store + sandbox + .store() .set("large", StoredValue::Bytes(large.clone()), None, 0.0) .unwrap(); - assert_eq!(value_files(directory.path()).len(), 1); - store + assert_eq!(sandbox.value_files().len(), 1); + sandbox + .store() .set( "large", StoredValue::Bytes(vec![b'y'; 32 * 1024]), @@ -282,23 +359,25 @@ fn delete_flush_and_spilled_file_replacement_clean_up_storage() { 0.0, ) .unwrap(); - assert_eq!(value_files(directory.path()).len(), 1); - store.pop("large", 0.0).unwrap(); - assert!(value_files(directory.path()).is_empty()); - store + 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(); - store + sandbox + .store() .set("b", StoredValue::Bytes(large), None, 0.0) .unwrap(); - store.clear().unwrap(); - assert!(value_files(directory.path()).is_empty()); + sandbox.store().clear().unwrap(); + assert!(sandbox.value_files().is_empty()); } +#[rstest] #[tokio::test] -async fn async_operations_connection_and_delete_match_sync_operations() { - let directory = tempfile::tempdir().unwrap(); - let cache = cache(directory.path()); +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)), }; 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..f9e12af6424 --- /dev/null +++ b/litellm-rust/crates/cache-disk/tests/python_compat.rs @@ -0,0 +1,105 @@ +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_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); +} From 59675bc0c948d124affe549f5ffc277bd8a31585 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:22:26 +0000 Subject: [PATCH 05/10] Revert "ci(rust): raise native wheel size gate to 35 MB" This reverts commit c1382086d69f176114003a5fcf32ac2a8cbaaa64. --- .github/scripts/verify_linux_native_wheel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index d86ad7b15c6..4fb8f068eb0 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -205,7 +205,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 35_000_000 + native_size_limit: Final = 25_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), @@ -222,7 +222,7 @@ def main( ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), - ("Native extension does not exceed 35 MB", native_size_within_limit), + ("Native extension does not exceed 25 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -267,7 +267,7 @@ def main( ), ( not native_size_within_limit, - f"native extension exceeds 35 MB: {native_member.file_size / 1_000_000:.2f} MB", + f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB", ), (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), ) From 4a8826e03b08311901cd8cb8186b9c1192206811 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:23:56 +0000 Subject: [PATCH 06/10] refactor(cache-disk): tidy python adapter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-disk/src/cache.rs | 6 ++--- .../crates/cache-disk/src/python/mod.rs | 24 ++++++++----------- .../crates/cache-disk/tests/python_compat.rs | 8 +++++++ 3 files changed, 20 insertions(+), 18 deletions(-) diff --git a/litellm-rust/crates/cache-disk/src/cache.rs b/litellm-rust/crates/cache-disk/src/cache.rs index 98182d93968..adff078a89b 100644 --- a/litellm-rust/crates/cache-disk/src/cache.rs +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -18,22 +18,20 @@ pub struct DiskCache { } impl DiskCache { - #[allow(clippy::default_constructed_unit_structs)] pub fn open(directory: impl AsRef, codec: S) -> Result { Ok(Self { store: Arc::new(DiskcacheSqliteStore::open(directory)?), - adapter: Arc::new(PythonDiskCacheAdapter::default()), + adapter: Arc::new(PythonDiskCacheAdapter), codec, }) } } impl DiskCache { - #[allow(clippy::default_constructed_unit_structs)] pub fn with_store(store: D, codec: S) -> Self { Self { store: Arc::new(store), - adapter: Arc::new(PythonDiskCacheAdapter::default()), + adapter: Arc::new(PythonDiskCacheAdapter), codec, } } diff --git a/litellm-rust/crates/cache-disk/src/python/mod.rs b/litellm-rust/crates/cache-disk/src/python/mod.rs index 5c9b57de000..7a370db357c 100644 --- a/litellm-rust/crates/cache-disk/src/python/mod.rs +++ b/litellm-rust/crates/cache-disk/src/python/mod.rs @@ -13,9 +13,7 @@ impl PythonDiskCacheAdapter { let value = match value { StoredValue::Bytes(value) => Value::Bytes(value), StoredValue::Text(value) => Value::String(value), - StoredValue::Integer(value) => { - value::from_json(serde_json::Value::Number(value.into())) - } + StoredValue::Integer(value) => Value::Integer(value.into()), StoredValue::Float(value) => Value::Float(value), StoredValue::Pickle(value) => value::from_pickle(&value)?, }; @@ -39,18 +37,16 @@ impl PythonDiskCacheAdapter { impl ValueAdapter for PythonDiskCacheAdapter { fn read(&self, value: StoredValue) -> Result>, Error> { - let raw = match &value { - StoredValue::Text(value) => Some(value.as_bytes().to_vec()), - StoredValue::Bytes(value) => Some(value.clone()), - StoredValue::Integer(_) | StoredValue::Float(_) | StoredValue::Pickle(_) => None, - }; - let Some(value) = Self::python_get_cache(value)? else { - return Ok(None); - }; - if let Some(raw) = raw { - return Ok(Some(raw)); + 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) + } } - value::to_json(&value).map(Some) } fn write(&self, payload: Vec) -> StoredValue { diff --git a/litellm-rust/crates/cache-disk/tests/python_compat.rs b/litellm-rust/crates/cache-disk/tests/python_compat.rs index f9e12af6424..9cbef8573bd 100644 --- a/litellm-rust/crates/cache-disk/tests/python_compat.rs +++ b/litellm-rust/crates/cache-disk/tests/python_compat.rs @@ -43,6 +43,14 @@ enum ReadExpectation { )] #[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 From f43f9012f1a0d58402aaec5393937225c1636869 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:44:24 +0000 Subject: [PATCH 07/10] fix(rust): handle disk cache topology Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/cache/native.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index e764fde8b39..50e02338574 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -83,6 +83,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => None, Self::Redis { cache, .. } => Some(cache.backend().topology()), + Self::Disk(_) => None, } } From 4845beddc067cd852c79a45e676c7128cd3bc23d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:44:26 +0000 Subject: [PATCH 08/10] ci(rust): raise native wheel size gate to 40 MB Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/verify_linux_native_wheel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 0adbc015ad0..7723c022351 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -205,7 +205,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 30_000_000 + native_size_limit: Final = 40_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), @@ -222,7 +222,7 @@ def main( ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), - ("Native extension does not exceed 30 MB", native_size_within_limit), + ("Native extension does not exceed 40 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -267,7 +267,7 @@ def main( ), ( not native_size_within_limit, - f"native extension exceeds 30 MB: {native_member.file_size / 1_000_000:.2f} MB", + f"native extension exceeds 40 MB: {native_member.file_size / 1_000_000:.2f} MB", ), (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), ) From 1827be1302b912ace13a56de2760a2f487025083 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:06:55 +0000 Subject: [PATCH 09/10] fix(cache-disk): absolutize store directory and stamp async expiry at write time Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-disk/src/cache.rs | 3 ++- litellm-rust/crates/cache-disk/src/sqlite.rs | 1 + litellm-rust/crates/cache-disk/tests/cache.rs | 19 +++++++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/litellm-rust/crates/cache-disk/src/cache.rs b/litellm-rust/crates/cache-disk/src/cache.rs index adff078a89b..8e1223309b4 100644 --- a/litellm-rust/crates/cache-disk/src/cache.rs +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -102,9 +102,10 @@ impl BaseCache for DiskCache Result<(), Error> { let value = self.adapter.write(self.codec.encode(&value)?); - let expire_time = context.ttl.map(|ttl| unix_now() + ttl.as_secs_f64()); + 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 diff --git a/litellm-rust/crates/cache-disk/src/sqlite.rs b/litellm-rust/crates/cache-disk/src/sqlite.rs index ff5a44b5c38..9a36f8af6ad 100644 --- a/litellm-rust/crates/cache-disk/src/sqlite.rs +++ b/litellm-rust/crates/cache-disk/src/sqlite.rs @@ -49,6 +49,7 @@ 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 diff --git a/litellm-rust/crates/cache-disk/tests/cache.rs b/litellm-rust/crates/cache-disk/tests/cache.rs index 4fb77d233ec..dd1f2b1f04e 100644 --- a/litellm-rust/crates/cache-disk/tests/cache.rs +++ b/litellm-rust/crates/cache-disk/tests/cache.rs @@ -61,6 +61,25 @@ impl Sandbox { } } +#[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; From 2a3ae253d49f0a5147ff2f571d9a8b601f21de38 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:14:22 +0000 Subject: [PATCH 10/10] refactor(rust): fold disk arms into shared match patterns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/python-bridge/src/cache/native.rs | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index bbd72e58825..80789cc279a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -99,35 +99,29 @@ impl NativeResponseCache { 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(), - Self::Disk(_) => None, } } 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()), - Self::Disk(_) => None, } } pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } => None, - Self::Disk(_) => None, - 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 { .. } => None, - Self::Disk(_) => None, - Self::AzureBlob(_) => None, + Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) => None, } } @@ -137,7 +131,6 @@ impl NativeResponseCache { cache, buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), }, - disk @ Self::Disk(_) => disk, other => other, } }