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

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:23:29 +00:00
parent 8dc960c928
commit 1320eeeb41

View file

@ -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 ResponseCacheRequest<ExactCacheContext> {
pub fn new(key: CacheKeyInput) -> Self {
Self {
key,
@ -32,11 +33,11 @@ impl ResponseCacheRequest {
}
}
pub struct ResponseCache<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>> {
backend: Arc<B>,
}
impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
}
@ -45,8 +46,11 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration> {
self.backend.get_ttl(&ExactCacheContext::default())
pub fn default_ttl(&self) -> Option<Duration>
where
B::Context: Default,
{
self.backend.get_ttl(&B::Context::default())
}
pub async fn async_flush(&self) -> Result<(), Error>
@ -62,7 +66,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 +85,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 +105,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 +130,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 +157,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 +176,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,9 +197,12 @@ 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> {
) -> Result<(), Error>
where
B::Context: PartialEq,
{
self.async_store_entries(
entries
.into_iter()
@ -209,8 +216,11 @@ 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)>,
) -> Result<(), Error> {
entries: Vec<(ResponseCacheRequest<B::Context>, Value, Duration)>,
) -> Result<(), Error>
where
B::Context: PartialEq,
{
let writable = entries
.into_iter()
.filter(|(request, _, _)| request.controls.writes())
@ -249,8 +259,8 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
}
fn partial_hits(
requests: &[ResponseCacheRequest],
readable: Vec<(usize, &ResponseCacheRequest)>,
requests: &[ResponseCacheRequest<B::Context>],
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,
entries: Vec<BatchEntry<CacheEntry>>,
now: Duration,
) -> Result<PartialHits, Error> {