refactor(rust): organize cache errors and tests

litellm-rust/crates/cache-memory/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "litellm-cache-memory"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+litellm-cache.workspace = true
+serde_json.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true

litellm-rust/crates/cache-memory/src/cache.rs

@@ -0,0 +1,226 @@
+use std::cmp::Reverse;
+use std::collections::{BinaryHeap, HashMap};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use litellm_cache::{BaseCache, CacheEntry, CacheKwargs, Error};
+
+const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
+const DEFAULT_TTL: Duration = Duration::from_secs(600);
+
+type ValueMeasure<V> = Arc<dyn Fn(&V) -> Result<usize, Error> + Send + Sync>;
+type ValueValidator<V> = Arc<dyn Fn(&V) -> Result<(), Error> + Send + Sync>;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum CacheWrite {
+    Stored,
+    Disabled,
+    TooLarge,
+}
+
+struct CacheState<V> {
+    values: HashMap<String, V>,
+    expirations: HashMap<String, Duration>,
+    expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
+}
+
+pub struct InMemoryCache<V: Clone> {
+    state: Mutex<CacheState<V>>,
+    max_size_in_memory: usize,
+    default_ttl: Duration,
+    max_entry_bytes: Option<usize>,
+    measure_value: Option<ValueMeasure<V>>,
+    validate_value: Option<ValueValidator<V>>,
+    now: Arc<dyn Fn() -> Duration + Send + Sync>,
+}
+
+impl<V: Clone> Default for InMemoryCache<V> {
+    fn default() -> Self {
+        Self::new(None, None)
+    }
+}
+
+impl<V: Clone> InMemoryCache<V> {
+    pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
+        Self::with_clock(max_size_in_memory, default_ttl, || {
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap_or_default()
+        })
+    }
+
+    pub fn with_clock(
+        max_size_in_memory: Option<usize>,
+        default_ttl: Option<Duration>,
+        now: impl Fn() -> Duration + Send + Sync + 'static,
+    ) -> Self {
+        Self::with_clock_and_size_measurement(max_size_in_memory, default_ttl, None, None, now)
+    }
+
+    pub fn with_clock_and_size_measurement(
+        max_size_in_memory: Option<usize>,
+        default_ttl: Option<Duration>,
+        max_entry_bytes: Option<usize>,
+        measure_value: Option<ValueMeasure<V>>,
+        now: impl Fn() -> Duration + Send + Sync + 'static,
+    ) -> Self {
+        Self {
+            state: Mutex::new(CacheState {
+                values: HashMap::new(),
+                expirations: HashMap::new(),
+                expiration_heap: BinaryHeap::new(),
+            }),
+            max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
+            default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
+            max_entry_bytes,
+            measure_value,
+            validate_value: None,
+            now: Arc::new(now),
+        }
+    }
+
+    pub fn set_cache(
+        &self,
+        key: impl Into<String>,
+        value: V,
+        ttl: Option<Duration>,
+    ) -> Result<CacheWrite, Error> {
+        if self.max_size_in_memory == 0 {
+            return Ok(CacheWrite::Disabled);
+        }
+        if let Some(validate) = &self.validate_value {
+            validate(&value)?;
+        }
+        if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
+            && measure(&value)? > limit
+        {
+            return Ok(CacheWrite::TooLarge);
+        }
+        let now = (self.now)();
+        let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
+        Self::evict(&mut state, self.max_size_in_memory, now);
+        let key = key.into();
+        state.values.insert(key.clone(), value);
+        let expiration = state.expirations.get(&key).copied();
+        if expiration.is_none_or(|expiration| expiration < now) {
+            let expiration = now + ttl.unwrap_or(self.default_ttl);
+            state.expirations.insert(key.clone(), expiration);
+            state.expiration_heap.push(Reverse((expiration, key)));
+        }
+        Ok(CacheWrite::Stored)
+    }
+
+    pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
+        let now = (self.now)();
+        let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
+        if state
+            .expirations
+            .get(key)
+            .is_some_and(|expiration| *expiration < now)
+        {
+            Self::remove(&mut state, key);
+        }
+        Ok(state.values.get(key).cloned())
+    }
+
+    pub fn get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
+        Ok(self
+            .state
+            .lock()
+            .map_err(|_| Error::Unavailable)?
+            .expirations
+            .get(key)
+            .copied())
+    }
+    pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
+        let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
+        Self::remove(&mut state, key);
+        Ok(())
+    }
+    pub fn flush_cache(&self) -> Result<(), Error> {
+        let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
+        state.values.clear();
+        state.expirations.clear();
+        state.expiration_heap.clear();
+        Ok(())
+    }
+
+    fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration) {
+        while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
+            if state.expirations.get(&key).copied() != Some(expiration) {
+                state.expiration_heap.pop();
+            } else if expiration < now {
+                state.expiration_heap.pop();
+                Self::remove(state, &key);
+            } else {
+                break;
+            }
+        }
+        while state.values.len() >= capacity {
+            let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else {
+                break;
+            };
+            if state.expirations.get(&key).copied() == Some(expiration) {
+                Self::remove(state, &key);
+            }
+        }
+    }
+
+    fn remove(state: &mut CacheState<V>, key: &str) {
+        state.values.remove(key);
+        state.expirations.remove(key);
+    }
+}
+
+impl InMemoryCache<CacheEntry> {
+    pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
+        Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || {
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap_or_default()
+        })
+    }
+
+    pub fn response_cache_with_clock(
+        capacity: usize,
+        ttl: Duration,
+        max_entry_bytes: usize,
+        now: impl Fn() -> Duration + Send + Sync + 'static,
+    ) -> Self {
+        let mut cache = Self::with_clock_and_size_measurement(
+            Some(capacity),
+            Some(ttl),
+            Some(max_entry_bytes),
+            Some(Arc::new(|entry: &CacheEntry| {
+                serde_json::to_vec(entry)
+                    .map(|bytes| bytes.len())
+                    .map_err(|_| Error::InvalidEntry)
+            })),
+            now,
+        );
+        cache.validate_value = Some(Arc::new(|entry: &CacheEntry| {
+            entry
+                .timestamp
+                .is_finite()
+                .then_some(())
+                .ok_or(Error::InvalidEntry)
+        }));
+        cache
+    }
+}
+
+impl BaseCache for InMemoryCache<CacheEntry> {
+    type Value = CacheEntry;
+    fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
+        self.set_cache(key, value, kwargs.ttl).map(|_| ())
+    }
+    fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
+        self.get_cache(key)
+    }
+    fn delete_cache(&self, key: &str) -> Result<(), Error> {
+        self.delete_cache(key)
+    }
+    fn flush_cache(&self) -> Result<(), Error> {
+        self.flush_cache()
+    }
+}

litellm-rust/crates/cache-memory/src/lib.rs

@@ -0,0 +1,2 @@
+pub mod cache;
+pub use cache::*;

litellm-rust/crates/cache-memory/tests/cache.rs

@@ -0,0 +1,124 @@
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::time::Duration;
+
+use litellm_cache::{CacheEntry, Error};
+use litellm_cache_memory::{CacheWrite, InMemoryCache};
+use rstest::{fixture, rstest};
+
+#[fixture]
+fn clock() -> Arc<AtomicU64> {
+    Arc::new(AtomicU64::new(100))
+}
+
+fn cache(clock: Arc<AtomicU64>, capacity: usize) -> InMemoryCache<String> {
+    InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || {
+        Duration::from_secs(clock.load(Ordering::SeqCst))
+    })
+}
+
+#[rstest]
+fn default_explicit_and_override_ttls_follow_python_rules(clock: Arc<AtomicU64>) {
+    let cache = cache(clock.clone(), 4);
+    cache.set_cache("key", "first".into(), None).unwrap();
+    assert_eq!(
+        cache.get_ttl("key").unwrap(),
+        Some(Duration::from_secs(160))
+    );
+    cache
+        .set_cache("key", "second".into(), Some(Duration::from_secs(10)))
+        .unwrap();
+    assert_eq!(
+        cache.get_ttl("key").unwrap(),
+        Some(Duration::from_secs(160))
+    );
+    clock.store(160, Ordering::SeqCst);
+    assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
+    clock.store(161, Ordering::SeqCst);
+    assert_eq!(cache.get_cache("key").unwrap(), None);
+    cache
+        .set_cache("key", "third".into(), Some(Duration::from_secs(10)))
+        .unwrap();
+    assert_eq!(
+        cache.get_ttl("key").unwrap(),
+        Some(Duration::from_secs(171))
+    );
+}
+
+#[rstest]
+fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>) {
+    let cache = cache(clock, 2);
+    cache
+        .set_cache("early", "a".into(), Some(Duration::from_secs(10)))
+        .unwrap();
+    cache
+        .set_cache("late", "b".into(), Some(Duration::from_secs(20)))
+        .unwrap();
+    cache.delete_cache("early").unwrap();
+    cache
+        .set_cache("new", "c".into(), Some(Duration::from_secs(30)))
+        .unwrap();
+    assert_eq!(cache.get_cache("late").unwrap(), Some("b".into()));
+    cache
+        .set_cache("last", "d".into(), Some(Duration::from_secs(40)))
+        .unwrap();
+    assert_eq!(cache.get_cache("late").unwrap(), None);
+}
+
+#[test]
+fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
+    let disabled = InMemoryCache::<CacheEntry>::response_cache(0, Duration::from_secs(60), 80);
+    assert_eq!(
+        disabled
+            .set_cache(
+                "a",
+                CacheEntry {
+                    timestamp: 1.0,
+                    response: serde_json::json!("x")
+                },
+                None
+            )
+            .unwrap(),
+        CacheWrite::Disabled
+    );
+    let cache = InMemoryCache::<CacheEntry>::response_cache(2, Duration::from_secs(60), 80);
+    assert_eq!(
+        cache
+            .set_cache(
+                "large",
+                CacheEntry {
+                    timestamp: 1.0,
+                    response: serde_json::json!("x".repeat(100))
+                },
+                None
+            )
+            .unwrap(),
+        CacheWrite::TooLarge
+    );
+    cache
+        .set_cache(
+            "small",
+            CacheEntry {
+                timestamp: 1.0,
+                response: serde_json::json!("ok"),
+            },
+            None,
+        )
+        .unwrap();
+    assert!(cache.get_cache("small").unwrap().is_some());
+    assert_eq!(
+        cache
+            .set_cache(
+                "invalid",
+                CacheEntry {
+                    timestamp: f64::NAN,
+                    response: serde_json::json!("bad"),
+                },
+                None,
+            )
+            .unwrap_err(),
+        Error::InvalidEntry
+    );
+    cache.delete_cache("small").unwrap();
+    cache.flush_cache().unwrap();
+}

litellm-rust/crates/cache/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "litellm-cache"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+serde.workspace = true
+serde_json.workspace = true
+sha2.workspace = true
+thiserror.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true

litellm-rust/crates/cache/src/base_cache.rs

@@ -0,0 +1,78 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::time::Duration;
+
+use serde_json::{Map, Value};
+
+use crate::Error;
+
+pub type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
+
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct CacheKwargs {
+    pub ttl: Option<Duration>,
+    pub extras: Map<String, Value>,
+}
+
+pub trait BaseCache: Send + Sync {
+    type Value: Clone + Send + Sync + 'static;
+
+    fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error>;
+
+    fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<Self::Value>, Error>;
+
+    fn async_set_cache<'a>(
+        &'a self,
+        key: &'a str,
+        value: Self::Value,
+        kwargs: CacheKwargs,
+    ) -> CacheFuture<'a, ()> {
+        Box::pin(async move { self.set_cache(key, value, kwargs) })
+    }
+
+    fn async_get_cache<'a>(
+        &'a self,
+        key: &'a str,
+        kwargs: &'a CacheKwargs,
+    ) -> CacheFuture<'a, Option<Self::Value>> {
+        Box::pin(async move { self.get_cache(key, kwargs) })
+    }
+
+    fn async_set_cache_pipeline<'a>(
+        &'a self,
+        cache_list: Vec<(String, Self::Value)>,
+        kwargs: CacheKwargs,
+    ) -> CacheFuture<'a, ()> {
+        Box::pin(async move {
+            for (key, value) in cache_list {
+                self.set_cache(&key, value, kwargs.clone())?;
+            }
+            Ok(())
+        })
+    }
+
+    fn batch_cache_write<'a>(
+        &'a self,
+        key: &'a str,
+        value: Self::Value,
+        kwargs: CacheKwargs,
+    ) -> CacheFuture<'a, ()> {
+        self.async_set_cache(key, value, kwargs)
+    }
+
+    fn delete_cache(&self, key: &str) -> Result<(), Error>;
+
+    fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
+        Box::pin(async move { self.delete_cache(key) })
+    }
+
+    fn flush_cache(&self) -> Result<(), Error>;
+
+    fn disconnect(&self) -> CacheFuture<'_, ()> {
+        Box::pin(async { Ok(()) })
+    }
+
+    fn test_connection(&self) -> CacheFuture<'_, ()> {
+        Box::pin(async { Ok(()) })
+    }
+}

litellm-rust/crates/cache/src/caching.rs

@@ -0,0 +1,164 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use sha2::{Digest, Sha256};
+
+use crate::{BaseCache, CacheKwargs, Error};
+
+pub use crate::BaseCache as Cache;
+
+#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
+pub enum CacheMode {
+    #[default]
+    #[serde(rename = "default_on")]
+    DefaultOn,
+    #[serde(rename = "default_off")]
+    DefaultOff,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct CacheKeyField {
+    pub name: String,
+    pub value: Option<String>,
+    pub api_parameter: bool,
+    pub internal_parameter: bool,
+}
+
+#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+pub struct CacheKeyInput {
+    pub fields: Vec<CacheKeyField>,
+    pub preset: Option<String>,
+    pub namespace: Option<String>,
+    pub include_provider_parameters: bool,
+}
+
+#[derive(Default)]
+pub struct CacheKeyContext {
+    pub model_group: Option<String>,
+    pub caching_groups: Vec<(Vec<String>, String)>,
+    pub file_checksum: Option<String>,
+    pub file_object_name: Option<String>,
+    pub metadata_file_name: Option<String>,
+    pub parameters_file_name: Option<String>,
+}
+
+impl CacheKeyContext {
+    pub fn apply(self, input: &mut CacheKeyInput) {
+        let group = self.model_group.as_ref().and_then(|model| {
+            self.caching_groups
+                .iter()
+                .find(|(models, _)| models.contains(model))
+        });
+        for field in &mut input.fields {
+            match field.name.as_str() {
+                "model" => {
+                    field.value = group
+                        .map(|(_, formatted)| formatted.clone())
+                        .or_else(|| self.model_group.clone())
+                        .or_else(|| field.value.take())
+                }
+                "file" => {
+                    field.value = self
+                        .file_checksum
+                        .clone()
+                        .or_else(|| self.file_object_name.clone())
+                        .or_else(|| self.metadata_file_name.clone())
+                        .or_else(|| self.parameters_file_name.clone())
+                }
+                _ => {}
+            }
+        }
+    }
+}
+
+pub fn get_cache_key(input: &CacheKeyInput) -> String {
+    cache_key(input)
+}
+
+pub fn cache_key(input: &CacheKeyInput) -> String {
+    if let Some(preset) = &input.preset {
+        return preset.clone();
+    }
+    let mut digest = Sha256::new();
+    for field in &input.fields {
+        if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
+            && let Some(value) = &field.value
+        {
+            digest.update(field.name.as_bytes());
+            digest.update(b": ");
+            digest.update(value.as_bytes());
+        }
+    }
+    let hash = format!("{:x}", digest.finalize());
+    input
+        .namespace
+        .as_deref()
+        .filter(|namespace| !namespace.is_empty())
+        .map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
+}
+
+#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
+pub struct CacheControls {
+    pub supported_call_type: bool,
+    pub configured: bool,
+    pub native_backend: bool,
+    pub default_on: bool,
+    pub caching: Option<bool>,
+    pub no_cache: bool,
+    pub no_store: bool,
+    #[serde(default)]
+    pub use_cache: bool,
+}
+
+impl CacheControls {
+    pub fn reads(self) -> bool {
+        self.supported_call_type
+            && self.configured
+            && self.caching.unwrap_or(true)
+            && !self.no_cache
+            && (self.default_on || self.use_cache)
+    }
+    pub fn writes(self) -> bool {
+        self.supported_call_type
+            && self.configured
+            && !self.no_store
+            && (self.default_on || self.use_cache)
+    }
+}
+
+pub fn should_use_cache(controls: CacheControls) -> bool {
+    controls.reads() || controls.writes()
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct CacheEntry {
+    pub timestamp: f64,
+    pub response: Value,
+}
+
+impl CacheEntry {
+    pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
+        self.timestamp.is_finite()
+            && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64())
+    }
+}
+
+pub fn get_cache(
+    cache: &dyn BaseCache<Value = CacheEntry>,
+    key: &str,
+    kwargs: &CacheKwargs,
+) -> Result<Option<CacheEntry>, Error> {
+    cache.get_cache(key, kwargs)
+}
+pub fn set_cache(
+    cache: &dyn BaseCache<Value = CacheEntry>,
+    key: &str,
+    entry: CacheEntry,
+    kwargs: CacheKwargs,
+) -> Result<(), Error> {
+    cache.set_cache(key, entry, kwargs)
+}
+
+pub type CacheBackend = Arc<dyn BaseCache<Value = CacheEntry>>;

litellm-rust/crates/cache/src/error.rs

@@ -0,0 +1,7 @@
+#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
+pub enum Error {
+    #[error("cache is unavailable")]
+    Unavailable,
+    #[error("invalid cache entry")]
+    InvalidEntry,
+}

litellm-rust/crates/cache/src/lib.rs

@@ -0,0 +1,10 @@
+pub mod base_cache;
+pub mod caching;
+pub mod error;
+
+pub use base_cache::{BaseCache, CacheFuture, CacheKwargs};
+pub use caching::{
+    Cache, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode,
+    cache_key, get_cache, get_cache_key, set_cache, should_use_cache,
+};
+pub use error::Error;

litellm-rust/crates/cache/tests/caching.rs

@@ -0,0 +1,83 @@
+use litellm_cache::{
+    CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key,
+};
+use sha2::{Digest, Sha256};
+
+#[test]
+fn keys_match_python_order_groups_files_presets_and_namespaces() {
+    let mut input = CacheKeyInput {
+        fields: vec![
+            CacheKeyField {
+                name: "model".into(),
+                value: Some("deployment".into()),
+                api_parameter: true,
+                internal_parameter: false,
+            },
+            CacheKeyField {
+                name: "file".into(),
+                value: None,
+                api_parameter: true,
+                internal_parameter: false,
+            },
+        ],
+        namespace: Some("team".into()),
+        ..Default::default()
+    };
+    CacheKeyContext {
+        model_group: Some("group".into()),
+        caching_groups: vec![(vec!["group".into()], "['group']".into())],
+        file_checksum: Some("checksum".into()),
+        ..Default::default()
+    }
+    .apply(&mut input);
+    assert_eq!(
+        cache_key(&input),
+        format!(
+            "team:{:x}",
+            Sha256::digest(b"model: ['group']file: checksum")
+        )
+    );
+    input.preset = Some("preset".into());
+    assert_eq!(get_cache_key(&input), "preset");
+}
+
+#[test]
+fn cache_controls_honor_default_modes_and_directives() {
+    let enabled = CacheControls {
+        supported_call_type: true,
+        configured: true,
+        default_on: true,
+        ..Default::default()
+    };
+    assert!(enabled.reads());
+    assert!(enabled.writes());
+    assert!(
+        !CacheControls {
+            default_on: false,
+            ..enabled
+        }
+        .reads()
+    );
+    assert!(
+        CacheControls {
+            default_on: false,
+            use_cache: true,
+            ..enabled
+        }
+        .reads()
+    );
+    assert!(
+        !CacheControls {
+            no_cache: true,
+            ..enabled
+        }
+        .reads()
+    );
+    assert!(
+        !CacheControls {
+            no_store: true,
+            ..enabled
+        }
+        .writes()
+    );
+}