From 3ce436af5ed72281214df3449ebf31323e8d3874 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:22:22 +0000 Subject: [PATCH] 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); +}