From fc3844e9913829cc99ee39b1270dbc7bcafde163 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:27:04 +0000 Subject: [PATCH] refactor(cache-response): generalize ResponseCache over the backend context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-response/src/response.rs | 57 ++++++++----- .../crates/cache-response/tests/response.rs | 83 ++++++++++++++++++- .../crates/python-bridge/src/cache/request.rs | 13 ++- 3 files changed, 128 insertions(+), 25 deletions(-) diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index e50e68cdabb..eedbf2caf1a 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/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 { pub key: CacheKeyInput, pub controls: CacheControls, - pub context: ExactCacheContext, + pub context: C, pub max_age: Option, } -impl ResponseCacheRequest { +impl ResponseCacheRequest { 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> { +impl ResponseCacheRequest { + pub fn with_context(self, context: D) -> ResponseCacheRequest { + ResponseCacheRequest { + key: self.key, + controls: self.controls, + context, + max_age: self.max_age, + } + } +} + +pub struct ResponseCache> +where + B::Context: Default + PartialEq, +{ backend: Arc, } -impl> ResponseCache { +impl ResponseCache +where + B: BaseCache, + B::Context: Default + PartialEq, +{ pub fn new(backend: Arc) -> Self { Self { backend } } @@ -46,7 +65,7 @@ impl> ResponseCach } pub fn default_ttl(&self) -> Option { - 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> ResponseCach pub fn lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -81,7 +100,7 @@ impl> ResponseCach pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -101,7 +120,7 @@ impl> ResponseCach pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -126,7 +145,7 @@ impl> ResponseCach pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -153,7 +172,7 @@ impl> ResponseCach pub fn store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -172,7 +191,7 @@ impl> ResponseCach pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -193,7 +212,7 @@ impl> ResponseCach pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, ) -> Result<(), Error> { self.async_store_entries( @@ -209,7 +228,7 @@ impl> ResponseCach /// the freshness of its original response. pub async fn async_store_entries( &self, - entries: Vec<(ResponseCacheRequest, Value, Duration)>, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, ) -> Result<(), Error> { let writable = entries .into_iter() @@ -248,9 +267,9 @@ impl> ResponseCach Ok(()) } - fn partial_hits( - requests: &[ResponseCacheRequest], - readable: Vec<(usize, &ResponseCacheRequest)>, + fn partial_hits( + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, entries: Vec>, now: Duration, ) -> Result { diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index e4f78dae8b2..56589063291 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/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>, + contexts: Mutex>, +} + +impl BaseCache for SemanticBackend { + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + 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, 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 { + 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)); diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 0c5343a63d0..0067fc4392b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/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, } -pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { +pub(super) fn request( + value: &Bound<'_, PyAny>, +) -> PyResult> { let input: RequestInput = from_py(value)?; request_input(input) } -fn request_input(input: RequestInput) -> PyResult { - let mut request = ResponseCacheRequest::new(input.key); +fn request_input(input: RequestInput) -> PyResult> { + let mut request: ResponseCacheRequest = ResponseCacheRequest::new(input.key); if let Some(controls) = input.controls { request.controls = controls; } @@ -29,7 +32,9 @@ fn request_input(input: RequestInput) -> PyResult { Ok(request) } -pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { +pub(super) fn requests( + value: &Bound<'_, PyAny>, +) -> PyResult>> { from_py::>(value)? .into_iter() .map(request_input)