scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

feat(rust): add native disk cache backend

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 20:37:26 +00:00
parent d79eadfe8c
commit 94b8bf7fd6
15 changed files with 2026 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

126
litellm-rust/Cargo.lock generated
View file

@ -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"

1
litellm-rust/Cargo.toml
View file

@ -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" }

17
litellm-rust/crates/cache-disk/Cargo.toml Normal file
View file

@ -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"

330
litellm-rust/crates/cache-disk/src/cache.rs Normal file
View file

@ -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<S, D = DiskcacheSqliteStore> {
store: Arc<D>,
codec: S,
}
impl<S: CacheCodec> DiskCache<S> {
pub fn open(directory: impl AsRef<Path>, codec: S) -> Result<Self, Error> {
Ok(Self {
store: Arc::new(DiskcacheSqliteStore::open(directory)?),
codec,
})
}
}
impl<S: CacheCodec, D: DiskStore> DiskCache<S, D> {
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<Option<S::Value>, Error> {
let Some(bytes) = payload(value)? else {
return Ok(None);
};
self.codec.decode(&bytes).map(Some)
}
async fn run_blocking<T, F>(store: Arc<D>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&D) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || operation(&store))
.await
.map_err(|_| Error::Unavailable)?
}
}
impl<S: CacheCodec, D: DiskStore> BaseCache for DiskCache<S, D> {
type Value = S::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
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<Option<Self::Value>, 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<Option<Self::Value>, 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::<Result<Vec<_>, _>>()?;
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<CacheConnectionResult, Error> {
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<S: CacheCodec, D: DiskStore> BatchCache for DiskCache<S, D> {
fn batch_get_cache(
&self,
keys: &[String],
context: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, 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<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, 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::<Result<Vec<_>, _>>()
})
.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<S: CacheCodec, D: DiskStore> DeleteCache for DiskCache<S, D> {
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<S: CacheCodec, D: DiskStore> FlushCache for DiskCache<S, D> {
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<S: CacheCodec<Value = f64>, D: DiskStore> CounterCache for DiskCache<S, D> {
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
increment(self.store.as_ref(), key, amount, context.ttl)
}
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = key.to_string();
Self::run_blocking(Arc::clone(&self.store), move |store| {
increment(store, &key, amount, context.ttl)
})
.await
}
}
fn increment<D: DiskStore>(
store: &D,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> 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 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<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)
.unwrap_or_default()
.as_secs_f64()
}

8
litellm-rust/crates/cache-disk/src/lib.rs Normal file
View file

@ -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};

48
litellm-rust/crates/cache-disk/src/pickle.rs Normal file
View file

@ -0,0 +1,48 @@
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),
}
}

816
litellm-rust/crates/cache-disk/src/sqlite.rs Normal file
View file

@ -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<Connection>,
min_file_size: usize,
eviction_policy: String,
size_limit: i64,
cull_limit: i64,
statistics: bool,
}
struct StoredColumns {
size: i64,
mode: i64,
filename: Option<String>,
value: Option<Value>,
}
struct Row {
rowid: i64,
mode: i64,
filename: Option<String>,
value: Value,
}
impl DiskcacheSqliteStore {
pub fn open(directory: impl AsRef<Path>) -> Result<Self, Error> {
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<f64>,
now: f64,
) -> Result<Vec<String>, 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<String>>(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<Vec<String>, 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<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.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<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.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<i64, Error> {
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<Option<StoredValue>, 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<f64>,
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<Option<StoredValue>, 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<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
if rows.is_empty() {
return Ok(rows);
}
let ids = rows
.iter()
.map(|(rowid, _)| rowid.to_string())
.collect::<Vec<_>>()
.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<StoredValue>) -> Result<(StoredValue, Option<f64>), 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<String, Value> {
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<HashMap<String, Value>, 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::<Result<HashMap<_, _>, _>>()
.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<String, Value>, key: &str) -> Option<i64> {
match settings.get(key) {
Some(Value::Integer(value)) => Some(*value),
_ => None,
}
}
fn setting_string(settings: &HashMap<String, Value>, key: &str) -> Option<String> {
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<Row> {
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<Option<StoredValue>, 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<Option<Vec<u8>>, 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<StoredColumns, Error> {
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<String, Error> {
let mut random = [0_u8; 16];
rand::rngs::OsRng.fill_bytes(&mut random);
let hex = random
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
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<String>) {
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<T>(
connection: &Connection,
operation: impl FnOnce(&Connection) -> Result<T, Error>,
) -> Result<T, Error> {
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)
}
}
}

33
litellm-rust/crates/cache-disk/src/store.rs Normal file
View file

@ -0,0 +1,33 @@
use std::path::Path;
use litellm_cache::Error;
#[derive(Clone, Debug, PartialEq)]
pub enum StoredValue {
Bytes(Vec<u8>),
Text(String),
Integer(i64),
Float(f64),
Pickle(Vec<u8>),
}
pub trait DiskStore: Send + Sync + 'static {
fn directory(&self) -> &Path;
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
fn set(
&self,
key: &str,
value: StoredValue,
expire_time: Option<f64>,
now: f64,
) -> Result<(), Error>;
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
fn clear(&self) -> Result<(), Error>;
fn update(
&self,
key: &str,
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error>;
fn probe(&self) -> Result<(), Error>;
}

327
litellm-rust/crates/cache-disk/tests/cache.rs Normal file
View file

@ -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<JsonCodec<serde_json::Value>> {
DiskCache::open(directory, JsonCodec::new()).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);
}
}
}
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::<f64>::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::<f64>::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::<Vec<_>>();
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
);
}

1
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -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

139
litellm-rust/crates/python-bridge/src/cache/config.rs vendored
View file

@ -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<RedisCacheConfig>),
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<MemoryCacheConfig> {
})
}
#[inline(never)]
fn project_disk(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<DiskCacheConfig, UnsupportedCacheConfig>> {
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::<String>()?),
}))
}
#[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"
);
});
}
}

38
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored
View file

@ -32,10 +32,16 @@ struct RedisPoolGuard {
max_connections: usize,
}
struct DiskStoreGuard {
reference: Py<PyAny>,
directory: String,
}
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
redis_pool: Option<RedisPoolGuard>,
disk_store: Option<DiskStoreGuard>,
}
impl ObjectGuard {
@ -176,6 +182,26 @@ impl RedisPoolGuard {
}
}
impl DiskStoreGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
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<bool> {
let store = backend.getattr("disk_cache")?;
Ok(self.reference.bind(py).is(&store)
&& self.directory == store.getattr("directory")?.extract::<String>()?)
}
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(())
}
}

12
litellm-rust/crates/python-bridge/src/cache/handle.rs vendored
View file

@ -51,6 +51,18 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[pyo3(signature = (directory))]
fn disk(py: Python<'_>, directory: String) -> PyResult<Self> {
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()

31
litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

@ -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<ResponseCache<RedisCache<ResponseCacheCodec>>>,
buffer: Option<Arc<WriteBuffer>>,
},
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
}
impl NativeResponseCache {
@ -43,6 +45,11 @@ impl NativeResponseCache {
buffer: None,
})
}
pub fn disk(directory: &str) -> Result<Self, Error> {
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,
}
}
}

111
tests/test_litellm_rust/test_cache.py
View file

@ -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],
}