refactor(cache-response): generalize ResponseCache over the backend context

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

@@ -1,21 +1,22 @@
 use std::{sync::Arc, time::Duration};
 
 use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
+    BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error,
+    ExactCacheContext, FlushCache,
 };
 use serde_json::Value;
 
 use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
 
 #[derive(Clone)]
-pub struct ResponseCacheRequest {
+pub struct ResponseCacheRequest<C: CacheContext = ExactCacheContext> {
     pub key: CacheKeyInput,
     pub controls: CacheControls,
-    pub context: ExactCacheContext,
+    pub context: C,
     pub max_age: Option<Duration>,
 }
 
-impl ResponseCacheRequest {
+impl<C: CacheContext + Default> ResponseCacheRequest<C> {
     pub fn new(key: CacheKeyInput) -> Self {
         Self {
             key,
@@ -26,17 +27,35 @@ impl ResponseCacheRequest {
                 default_on: true,
                 ..Default::default()
             },
-            context: ExactCacheContext::default(),
+            context: C::default(),
             max_age: None,
         }
     }
 }
 
-pub struct ResponseCache<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
+impl<C: CacheContext> ResponseCacheRequest<C> {
+    pub fn with_context<D: CacheContext>(self, context: D) -> ResponseCacheRequest<D> {
+        ResponseCacheRequest {
+            key: self.key,
+            controls: self.controls,
+            context,
+            max_age: self.max_age,
+        }
+    }
+}
+
+pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
+where
+    B::Context: Default + PartialEq,
+{
     backend: Arc<B>,
 }
 
-impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
+impl<B> ResponseCache<B>
+where
+    B: BaseCache<Value = CacheEntry>,
+    B::Context: Default + PartialEq,
+{
     pub fn new(backend: Arc<B>) -> Self {
         Self { backend }
     }
@@ -46,7 +65,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
     }
 
     pub fn default_ttl(&self) -> Option<Duration> {
-        self.backend.get_ttl(&ExactCacheContext::default())
+        self.backend.get_ttl(&B::Context::default())
     }
 
     pub async fn async_flush(&self) -> Result<(), Error>
@@ -62,7 +81,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
 
     pub fn lookup(
         &self,
-        request: &ResponseCacheRequest,
+        request: &ResponseCacheRequest<B::Context>,
         now: Duration,
     ) -> Result<Option<Value>, Error> {
         if !request.controls.reads() {
@@ -81,7 +100,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
 
     pub async fn async_lookup(
         &self,
-        request: &ResponseCacheRequest,
+        request: &ResponseCacheRequest<B::Context>,
         now: Duration,
     ) -> Result<Option<Value>, Error> {
         if !request.controls.reads() {
@@ -101,7 +120,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
 
     pub fn lookup_batch(
         &self,
-        requests: &[ResponseCacheRequest],
+        requests: &[ResponseCacheRequest<B::Context>],
         now: Duration,
     ) -> Result<PartialHits, Error>
     where
@@ -126,7 +145,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
 
     pub async fn async_lookup_batch(
         &self,
-        requests: &[ResponseCacheRequest],
+        requests: &[ResponseCacheRequest<B::Context>],
         now: Duration,
     ) -> Result<PartialHits, Error>
     where
@@ -153,7 +172,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
 
     pub fn store(
         &self,
-        request: &ResponseCacheRequest,
+        request: &ResponseCacheRequest<B::Context>,
         response: Value,
         now: Duration,
     ) -> Result<(), Error> {
@@ -172,7 +191,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
 
     pub async fn async_store(
         &self,
-        request: &ResponseCacheRequest,
+        request: &ResponseCacheRequest<B::Context>,
         response: Value,
         now: Duration,
     ) -> Result<(), Error> {
@@ -193,7 +212,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
 
     pub async fn async_store_batch(
         &self,
-        entries: Vec<(ResponseCacheRequest, Value)>,
+        entries: Vec<(ResponseCacheRequest<B::Context>, Value)>,
         now: Duration,
     ) -> Result<(), Error> {
         self.async_store_entries(
@@ -209,7 +228,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
     /// the freshness of its original response.
     pub async fn async_store_entries(
         &self,
-        entries: Vec<(ResponseCacheRequest, Value, Duration)>,
+        entries: Vec<(ResponseCacheRequest<B::Context>, Value, Duration)>,
     ) -> Result<(), Error> {
         let writable = entries
             .into_iter()
@@ -248,9 +267,9 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
         Ok(())
     }
 
-    fn partial_hits(
-        requests: &[ResponseCacheRequest],
-        readable: Vec<(usize, &ResponseCacheRequest)>,
+    fn partial_hits<C: CacheContext>(
+        requests: &[ResponseCacheRequest<C>],
+        readable: Vec<(usize, &ResponseCacheRequest<C>)>,
         entries: Vec<BatchEntry<CacheEntry>>,
         now: Duration,
     ) -> Result<PartialHits, Error> {

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

@@ -1,12 +1,15 @@
 use std::{
     sync::{
-        Arc,
+        Arc, Mutex,
         atomic::{AtomicU64, Ordering},
     },
     time::Duration,
 };
 
-use litellm_cache::{BaseCache, CacheCodec, Error};
+use litellm_cache::{
+    BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
+    SemanticCacheContext,
+};
 use litellm_cache_memory::InMemoryCache;
 use litellm_cache_redis::RedisCache;
 use litellm_cache_response::{
@@ -30,6 +33,82 @@ fn request() -> ResponseCacheRequest {
     })
 }
 
+struct SemanticBackend {
+    entries: Mutex<Vec<(String, CacheEntry)>>,
+    contexts: Mutex<Vec<SemanticCacheContext>>,
+}
+
+impl BaseCache for SemanticBackend {
+    type Value = CacheEntry;
+    type Context = SemanticCacheContext;
+
+    fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
+        None
+    }
+
+    fn set_cache(
+        &self,
+        key: &str,
+        value: Self::Value,
+        context: &Self::Context,
+    ) -> Result<(), Error> {
+        self.contexts.lock().unwrap().push(context.clone());
+        self.entries.lock().unwrap().push((key.to_owned(), value));
+        Ok(())
+    }
+
+    fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
+        self.contexts.lock().unwrap().push(context.clone());
+        Ok(self
+            .entries
+            .lock()
+            .unwrap()
+            .iter()
+            .find(|(entry_key, _)| entry_key == key)
+            .map(|(_, entry)| entry.clone()))
+    }
+
+    async fn disconnect(&self) -> Result<(), Error> {
+        Ok(())
+    }
+
+    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
+        Ok(CacheConnectionResult {
+            status: CacheConnectionStatus::Success,
+            message: "ok".into(),
+            error: None,
+        })
+    }
+}
+
+#[test]
+fn semantic_context_reaches_backend_for_store_and_lookup() {
+    let backend = Arc::new(SemanticBackend {
+        entries: Mutex::new(Vec::new()),
+        contexts: Mutex::new(Vec::new()),
+    });
+    let cache = ResponseCache::new(backend.clone());
+    let context = SemanticCacheContext {
+        messages: vec![json!({"role": "user", "content": "hello"})],
+        ..Default::default()
+    };
+    let request = request().with_context(context.clone());
+    let response = json!({"answer": 42});
+
+    cache
+        .store(&request, response.clone(), Duration::from_secs(100))
+        .unwrap();
+
+    assert_eq!(
+        cache.lookup(&request, Duration::from_secs(100)).unwrap(),
+        Some(response)
+    );
+    assert_eq!(
+        backend.contexts.lock().unwrap().as_slice(),
+        &[context.clone(), context]
+    );
+}
+
 #[tokio::test]
 async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
     let clock = Arc::new(AtomicU64::new(100));

litellm-rust/crates/python-bridge/src/cache/request.rs

@@ -1,5 +1,6 @@
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
+use litellm_cache::ExactCacheContext;
 use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest};
 use litellm_host_python::from_py;
 use pyo3::{exceptions::PyValueError, prelude::*};
@@ -14,13 +15,15 @@ struct RequestInput {
     max_age_seconds: Option<f64>,
 }
 
-pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<ResponseCacheRequest> {
+pub(super) fn request(
+    value: &Bound<'_, PyAny>,
+) -> PyResult<ResponseCacheRequest<ExactCacheContext>> {
     let input: RequestInput = from_py(value)?;
     request_input(input)
 }
 
-fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest> {
-    let mut request = ResponseCacheRequest::new(input.key);
+fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest<ExactCacheContext>> {
+    let mut request: ResponseCacheRequest<ExactCacheContext> = ResponseCacheRequest::new(input.key);
     if let Some(controls) = input.controls {
         request.controls = controls;
     }
@@ -29,7 +32,9 @@ fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest> {
     Ok(request)
 }
 
-pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<ResponseCacheRequest>> {
+pub(super) fn requests(
+    value: &Bound<'_, PyAny>,
+) -> PyResult<Vec<ResponseCacheRequest<ExactCacheContext>>> {
     from_py::<Vec<RequestInput>>(value)?
         .into_iter()
         .map(request_input)