mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
refactor(cache-disk): isolate python compatibility behind a value adapter
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
parent
c1382086d6
commit
3ce436af5e
10 changed files with 643 additions and 264 deletions
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
10
litellm-rust/crates/cache-disk/src/adapter.rs
Normal file
10
litellm-rust/crates/cache-disk/src/adapter.rs
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
use litellm_cache::Error;
|
||||
|
||||
use crate::StoredValue;
|
||||
|
||||
pub trait ValueAdapter: Send + Sync + 'static {
|
||||
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error>;
|
||||
fn write(&self, payload: Vec<u8>) -> StoredValue;
|
||||
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error>;
|
||||
fn counter_value(&self, value: f64) -> StoredValue;
|
||||
}
|
||||
|
|
@ -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<S, D = DiskcacheSqliteStore> {
|
||||
pub struct DiskCache<S, D = DiskcacheSqliteStore, A = PythonDiskCacheAdapter> {
|
||||
store: Arc<D>,
|
||||
adapter: Arc<A>,
|
||||
codec: S,
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> DiskCache<S> {
|
||||
#[allow(clippy::default_constructed_unit_structs)]
|
||||
pub fn open(directory: impl AsRef<Path>, codec: S) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
store: Arc::new(DiskcacheSqliteStore::open(directory)?),
|
||||
adapter: Arc::new(PythonDiskCacheAdapter::default()),
|
||||
codec,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore> DiskCache<S, D> {
|
||||
impl<S: CacheCodec, D: DiskStore> DiskCache<S, D, PythonDiskCacheAdapter> {
|
||||
#[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<S: CacheCodec, D: DiskStore, A: ValueAdapter> DiskCache<S, D, A> {
|
||||
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<S: CacheCodec, D: DiskStore> DiskCache<S, D> {
|
|||
}
|
||||
|
||||
fn decode_stored(&self, value: StoredValue) -> Result<Option<S::Value>, 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<S: CacheCodec, D: DiskStore> DiskCache<S, D> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore> BaseCache for DiskCache<S, D> {
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BaseCache for DiskCache<S, D, A> {
|
||||
type Value = S::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
|
|
@ -70,7 +84,7 @@ impl<S: CacheCodec, D: DiskStore> BaseCache for DiskCache<S, D> {
|
|||
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<S: CacheCodec, D: DiskStore> BaseCache for DiskCache<S, D> {
|
|||
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<S: CacheCodec, D: DiskStore> BaseCache for DiskCache<S, D> {
|
|||
.map(|(key, value)| {
|
||||
self.codec
|
||||
.encode(&value)
|
||||
.map(|value| (key, StoredValue::Bytes(value)))
|
||||
.map(|value| (key, self.adapter.write(value)))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let expire_after = context.ttl;
|
||||
|
|
@ -162,7 +176,7 @@ impl<S: CacheCodec, D: DiskStore> BaseCache for DiskCache<S, D> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore> BatchCache for DiskCache<S, D> {
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D, A> {
|
||||
fn batch_get_cache(
|
||||
&self,
|
||||
keys: &[String],
|
||||
|
|
@ -204,7 +218,7 @@ impl<S: CacheCodec, D: DiskStore> BatchCache for DiskCache<S, D> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore> DeleteCache for DiskCache<S, D> {
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DeleteCache for DiskCache<S, D, A> {
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.store.pop(key, unix_now()).map(|_| ())
|
||||
}
|
||||
|
|
@ -218,7 +232,7 @@ impl<S: CacheCodec, D: DiskStore> DeleteCache for DiskCache<S, D> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec, D: DiskStore> FlushCache for DiskCache<S, D> {
|
||||
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> FlushCache for DiskCache<S, D, A> {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
self.store.clear()
|
||||
}
|
||||
|
|
@ -228,14 +242,22 @@ impl<S: CacheCodec, D: DiskStore> FlushCache for DiskCache<S, D> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec<Value = f64>, D: DiskStore> CounterCache for DiskCache<S, D> {
|
||||
impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
|
||||
for DiskCache<S, D, A>
|
||||
{
|
||||
fn increment_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
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<S: CacheCodec<Value = f64>, D: DiskStore> CounterCache for DiskCache<S, D>
|
|||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
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<D: DiskStore>(
|
||||
fn increment<A: ValueAdapter, D: DiskStore>(
|
||||
adapter: &A,
|
||||
store: &D,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
|
|
@ -260,25 +284,9 @@ fn increment<D: DiskStore>(
|
|||
) -> Result<f64, Error> {
|
||||
let mut result = None;
|
||||
let mut apply = |current: Option<StoredValue>| {
|
||||
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<D: DiskStore>(
|
|||
result.ok_or(Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn payload(value: StoredValue) -> Result<Option<Vec<u8>>, 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)
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -1,48 +0,0 @@
|
|||
use litellm_cache::Error;
|
||||
use serde_json::{Map, Number, Value};
|
||||
|
||||
pub(crate) fn decode(bytes: &[u8]) -> Result<Value, Error> {
|
||||
let value = serde_pickle::value_from_slice(bytes, Default::default())
|
||||
.map_err(|_| Error::InvalidEntry)?;
|
||||
convert(value)
|
||||
}
|
||||
|
||||
fn convert(value: serde_pickle::Value) -> Result<Value, Error> {
|
||||
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::<i64>() {
|
||||
Ok(Value::Number(value.into()))
|
||||
} else if let Ok(value) = value.to_string().parse::<u64>() {
|
||||
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::<Result<Vec<_>, _>>()
|
||||
.map(Value::Array),
|
||||
serde_pickle::Value::Set(values) | serde_pickle::Value::FrozenSet(values) => values
|
||||
.into_iter()
|
||||
.map(|value| convert(value.into_value()))
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.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::<Result<Map<_, _>, _>>()
|
||||
.map(Value::Object),
|
||||
serde_pickle::Value::Bytes(_) => Err(Error::InvalidEntry),
|
||||
}
|
||||
}
|
||||
81
litellm-rust/crates/cache-disk/src/python/mod.rs
Normal file
81
litellm-rust/crates/cache-disk/src/python/mod.rs
Normal file
|
|
@ -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<Option<Value>, 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<Option<Vec<u8>>, 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<u8>) -> StoredValue {
|
||||
StoredValue::Bytes(payload)
|
||||
}
|
||||
|
||||
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
173
litellm-rust/crates/cache-disk/src/python/value.rs
Normal file
173
litellm-rust/crates/cache-disk/src/python/value.rs
Normal file
|
|
@ -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<Value, Error> {
|
||||
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<Value, Error> {
|
||||
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::<Result<Vec<_>, _>>()
|
||||
.map(Value::List),
|
||||
serde_pickle::Value::Tuple(values) => values
|
||||
.into_iter()
|
||||
.map(from_pickle_value)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Tuple),
|
||||
serde_pickle::Value::Set(values) => values
|
||||
.into_iter()
|
||||
.map(from_pickle_hashable)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Set),
|
||||
serde_pickle::Value::FrozenSet(values) => values
|
||||
.into_iter()
|
||||
.map(from_pickle_hashable)
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(Value::Set),
|
||||
serde_pickle::Value::Dict(values) => values
|
||||
.into_iter()
|
||||
.map(|(key, value)| Ok((from_pickle_hashable(key)?, from_pickle_value(value)?)))
|
||||
.collect::<Result<Vec<_>, Error>>()
|
||||
.map(Value::Dict),
|
||||
}
|
||||
}
|
||||
|
||||
fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result<Value, Error> {
|
||||
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::<Result<Vec<_>, _>>()?,
|
||||
),
|
||||
serde_pickle::HashableValue::FrozenSet(values) => Value::Set(
|
||||
values
|
||||
.into_iter()
|
||||
.map(from_pickle_hashable)
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
fn integer(value: String) -> Result<Value, Error> {
|
||||
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<Value, Error> {
|
||||
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<f64> {
|
||||
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<Vec<u8>, Error> {
|
||||
serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
|
||||
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::<Number>()
|
||||
.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::<Result<Vec<_>, _>>()?,
|
||||
)
|
||||
}
|
||||
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::<Result<Map<String, serde_json::Value>, _>>()?;
|
||||
serde_json::Value::Object(values)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -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<JsonCodec<serde_json::Value>> {
|
||||
DiskCache::open(directory, JsonCodec::new()).unwrap()
|
||||
#[fixture]
|
||||
fn sandbox() -> Sandbox {
|
||||
Sandbox {
|
||||
directory: tempfile::tempdir().unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
fn value_files(directory: &Path) -> Vec<PathBuf> {
|
||||
fn visit(directory: &Path, files: &mut Vec<PathBuf>) {
|
||||
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<V>(&self) -> DiskCache<JsonCodec<V>>
|
||||
where
|
||||
JsonCodec<V>: 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<PathBuf> {
|
||||
fn visit(directory: &Path, files: &mut Vec<PathBuf>) {
|
||||
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<Option<Vec<u8>>, 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<u8>) -> StoredValue {
|
||||
StoredValue::Text(String::from_utf8(payload).unwrap())
|
||||
}
|
||||
|
||||
fn counter_seed(&self, _: Option<StoredValue>) -> Result<f64, litellm_cache::Error> {
|
||||
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::<Value>();
|
||||
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::<Value>();
|
||||
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::<Value>()
|
||||
.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::<Value>()
|
||||
.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::<f64>::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<StoredValue>,
|
||||
#[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::<f64>();
|
||||
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::<f64>::new()).unwrap());
|
||||
#[rstest]
|
||||
fn counters_are_atomic_across_concurrent_callers(sandbox: Sandbox) {
|
||||
let cache = Arc::new(sandbox.cache::<f64>());
|
||||
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::<f64>();
|
||||
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::<f64>();
|
||||
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::<Value>::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::<Value>();
|
||||
let context = ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(60)),
|
||||
};
|
||||
|
|
|
|||
105
litellm-rust/crates/cache-disk/tests/python_compat.rs
Normal file
105
litellm-rust/crates/cache-disk/tests/python_compat.rs
Normal file
|
|
@ -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<StoredValue>, #[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);
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue