refactor(rust-bridge): let cache adapters own request preparation
Some checks are pending
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run

Keep Python cache key construction, lookup and storage in Python. Prepare native identity lazily through the selected adapter, preserve authoritative hit keys and cache-handler parity, and consolidate the cache-local Python binding.

Co-authored-by: Yujong Lee <yujong@berri.ai>
This commit is contained in:
Devin AI 2026-10-17 22:15:21 +00:00
parent 898fbaf985
commit 33e5d9168e
95 changed files with 3635 additions and 1847 deletions

View file

@ -3886,6 +3886,7 @@ dependencies = [
"litellm-auth",
"litellm-coroutine",
"rstest",
"serde",
"serde_json",
"tokio",
]

View file

@ -1,29 +1,69 @@
# Response caching
Design this crate for shared Rust execution used by the Python SDK and the Rust gateway. The Python SDK will remain, with more core execution moving to Rust and Python callbacks staying in Python. The Rust gateway is still evolving and is intended to replace the Python proxy. Keep response-cache policy independent of Python, HTTP serving, and either proxy's configuration format
This crate owns response-cache semantics shared by Rust execution in the Python SDK and in the Rust gateway: what a cache entry is, when it is fresh, how a call's cache policy is expressed, and the contract a cache service offers to core. Keep it independent of Python, HTTP serving and any proxy's configuration format
Separate what is cached, how a hit is matched, and where entries are stored. Chat Completions, Messages, Responses, and embeddings are API workloads. Exact and semantic matching are lookup behaviors. Memory, Redis, disk, and object stores are storage choices. Embeddings are inference too, so do not use an inference-cache name to imply a category that excludes embeddings. Consult the existing Python cache and caching handler for behavior and compatibility contracts without copying their class structure
## Who decides what
Storage traits, codecs, and backend capabilities belong in `litellm-cache` and the storage crates. Keep storage reusable for value types beyond LLM responses. This crate owns response entries, matching and freshness semantics, the Python-compatible response codec, and deferred-write policy. Core owns route-specific request identity, response encoding and reconstruction, embedding partial-hit orchestration, and stream capture and replay. Boundaries own configuration translation, resource construction, and caller identity
Core decides when. It looks up before calling the provider and stores provider output once that output is final, a complete response after the hooks that run once per provider call and a stream once the caller has consumed it. Core never decides whether a call may use the cache, what the key is, or how an entry is laid out
Construct and inject the response-cache service at the Python bridge or gateway boundary, as with the HTTP client. Reuse it across calls. Core and provider code must not discover cache configuration through Python globals, process configuration, or backend-specific factories
The cache service decides what. It decides whether a call may read or write, derives the key, reads and writes storage, and owns the stored layout. When the service delegates to another cache implementation, that implementation owns all four, and Rust only carries values across
Keep `ResponseCache<B>` generic over its storage backend. Preserve typed backend contexts and capability bounds internally. Inject an object-safe service into core for runtime backend selection, so storage types do not spread through route and host types. Keep API request and response types statically typed. Add a generic parameter only where it preserves a useful type relationship or capability
Values crossing that boundary are route-neutral, either a complete response or a stream's bytes. Rust never names a route's stored shape. Each route's stored layout is defined once, in the layer that owns storage
Keep the core service contract narrow. Lookup and store must not require connection testing, ping, flush, deletion, counters, queues, or scripts. Require batch operations where a consumer needs partial hits, and keep management capabilities on their own interfaces. An exact-only adapter must remain explicit about its matching restriction. Supporting semantic matching requires a defined lookup-context and embedding execution contract, not just a renamed trait
## One key for everyone who shares a cache
Separate reusable resources from per-call policy. Backend configuration, namespace, default expiry, and entry limits belong to the configured service or backend. Read/write controls, expiry and freshness overrides, and authenticated caller scope belong to the call. Passing call options must not replace or mutate the route's configured service
Every path that reads or writes the same entries derives the key from the same inputs with the same rule. Two paths that each normalize the call on their own drift apart silently, and nothing fails when they do
Keep cache misses and storage failures distinguishable in return values. Core owns the decision to continue with provider execution after a cache failure. A read can reject an entry for freshness while the backend still retains it. Preserve the timestamp at which a response was produced when writing it later
Key inputs are the call as the caller expressed it, including arguments passed by position. A missing input is a correctness bug, since unrelated requests then share an entry and one caller receives another caller's answer. An extra input, such as an injected default the other paths leave out, splits entries that should be shared. Treat both as key-compatibility breaks
Define lookup placement explicitly relative to authorization, deployment and credential resolution, and request-transforming callbacks. Cache identity must account for every input that affects reuse, including API surface and caller scope, while preserving intentional Python caching groups. Preserve existing keys and response formats unless changing them is an explicit migration decision
Decide whether the cache applies before deriving the key. Disabled caching, opt-in modes, request-size limits and read/write controls must not pay for hashing the request. Derive the key once per call and use that same key for lookup, store and the key reported on a hit
Cache normalized provider results before caller-specific response transformations. Hits must still run the applicable response processing, success callbacks, and cache-hit accounting. Keep callback execution in the host. Python cache implementations and semantic embedders that require the caller's task must use the existing host-operation mechanism rather than Python calls from a Rust worker. Preserve legacy fallback until that contract is supported
Preserve existing keys and stored formats unless a change is an explicit migration decision
Keep unary caching independent of stream-only methods. Store streams only after successful exhaustion and protocol completion. Errors, incomplete streams, cancellation, and oversized entries must not populate the cache. Embedding batches need ordered partial results and reconstruction around the uncached inputs
## Entries
Test each contract in its owner: storage capabilities in backend tests, envelopes and freshness here, reuse and replay in core, Python callback and fallback behavior at the bridge, and HTTP behavior at the gateway. Run backend contract checks and Python response-codec fixtures before exposing a new backend
A hit carries the key that found it. Report that key, never a separately derived one
`ScopedCache` requires an explicit shared or isolated scope at construction. Per-call `CachePolicy` controls reads, writes, expiry, and freshness without replacing the attached scope or service. `CacheOptions` binds that policy to an explicit scope for storage requests and has no default sharing policy. Versioned native envelopes reject incompatible API surfaces and versions as misses. This envelope is distinct from the legacy Python response codec
An entry with no usable output is neither served nor stored. An entry that cannot be decoded, has the wrong surface or version, or does not match what the request asked for is a miss, not an error
Response storage is not the source of budget or rate-limit coordination dependencies. Keep counters, reservations, and atomic admission operations out of `ResponseCacheService`, including when both services happen to use Redis
A stream request is answered with a stream and a complete request with a complete response. Returning the other kind is a miss unless the route defines how to replay one as the other
Native envelopes are versioned per API surface. They are distinct from any format shared with another cache implementation, which must stay readable and writable by that implementation
## Streams
Store a stream only after the caller has consumed it to the end, the protocol reached its terminal event and no error event appeared. Early close, cancellation, errors and incomplete streams store nothing. A replayed stream yields the same bytes the provider sent
## Failures
A cache failure is a cache miss, never a request failure. That covers deciding eligibility, deriving the key, lookup and store. Cancellation is not a cache failure and must propagate. A failed store never repeats the provider call
Keep misses and storage failures distinguishable in return values, and leave the decision to continue with the provider to core. A freshness check may reject an entry the backend still holds. Writes record when the response was produced, not when the write lands. An asynchronous caller gets its response before a background write completes
## Responses and hooks
Response hooks sit on two sides of the cache. A hook that runs once per provider call, such as a deployment success hook, sits inside it. It runs before the store, its result is what gets cached, and it does not run again on a hit, since it may execute tools, call providers again or redact what must not be stored. A hook that runs once per request, such as per-caller policy, sits outside it and runs on hits and misses alike
A hit therefore returns what the first caller received from the inside hooks. Hits still produce success callbacks and cache-hit accounting, and callback execution stays in the host
## Services, scope and policy
Construct the cache service at the boundary, from the boundary's configuration, and inject it like the HTTP client. Reuse it across calls. Core and provider code never discover cache configuration from globals or backend factories
Configured resources and per-call policy stay separate. Backend, namespace, default expiry and entry limits belong to the service. Read and write controls, expiry and freshness overrides and caller scope belong to the call. Call policy never replaces the attached service or its scope, and every scope is explicit, shared or isolated, with no default
A service's lookup and store need nothing beyond lookup and store. Connection checks, flushing, deletion, counters and scripts live on separate interfaces. Exact and semantic matching are different lookup contracts, and supporting semantic matching needs an embedding execution contract, not a renamed trait
Implementations that must run in the caller's task, such as a Python cache, run through host operations, never as direct calls from a Rust worker
Response caching is not budget or rate-limit coordination. Counters, reservations and atomic admission stay out of the cache service, even when both use the same store
## Layering
Storage traits, codecs and backend capabilities belong to the storage crates, and stay reusable beyond LLM responses. This crate owns entries, freshness, envelopes, cache policy and the service contract. Core owns route identity, response reconstruction, stream capture and replay, and embedding partial hits. Boundaries own configuration, resource construction and caller identity
API workloads, lookup behavior and storage are separate axes. Chat Completions, Messages, Responses and embeddings are workloads. Exact and semantic are lookups. Memory, Redis, disk and object stores are storage
## Tests
Test each contract where it is owned. Storage capabilities belong in backend tests, entries and freshness here, reuse and replay in core, delegation and callbacks at the bridge, and HTTP behavior at the gateway. Any key or format shared with another implementation needs a test where each side reads what the other wrote

View file

@ -17,6 +17,6 @@ pub use exact::{ConnectionProbe, ExactResponseCache};
pub use response::{ResponseCache, ResponseCacheRequest};
pub use service::{
CacheOptions, CachePolicy, CacheScope, ResponseCacheConfig, ResponseCacheService,
ResponseEnvelope, ScopedCache,
CacheHit, CacheOptions, CachePolicy, CacheScope, PreparedCache, ResponseCacheCall,
ResponseCacheConfig, ResponseCacheService, ResponseEnvelope, ScopedCache,
};

View file

@ -1,10 +1,11 @@
use std::{future::Future, pin::Pin, time::Duration};
use std::{future::Future, pin::Pin, sync::Arc, time::Duration};
use litellm_cache::{BaseCache, Error, ExactCacheContext};
use serde_json::Value;
use crate::{
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
cache_key,
};
type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
@ -24,46 +25,78 @@ impl Default for ResponseCacheConfig {
}
}
#[derive(Debug, PartialEq)]
pub struct CacheHit {
pub value: Value,
pub key: String,
}
pub trait ResponseCacheService: Send + Sync {
fn config(&self) -> &ResponseCacheConfig;
fn prepare(
self: Arc<Self>,
options: CacheOptions,
surface: &str,
input: &dyn Fn() -> Value,
) -> PreparedCache;
}
fn lookup<'a>(
&'a self,
request: &'a ResponseCacheRequest,
now: Duration,
) -> CacheFuture<'a, Option<Value>>;
pub struct PreparedCache {
pub policy: CachePolicy,
pub max_entry_bytes: usize,
pub call: Arc<dyn ResponseCacheCall>,
}
fn store<'a>(
&'a self,
request: &'a ResponseCacheRequest,
response: Value,
now: Duration,
) -> CacheFuture<'a, ()>;
pub trait ResponseCacheCall: Send + Sync {
fn lookup<'a>(&'a self, now: Duration) -> CacheFuture<'a, Option<CacheHit>>;
fn store<'a>(&'a self, response: Value, now: Duration) -> CacheFuture<'a, ()>;
}
impl<B> ResponseCacheService for ResponseCache<B>
where
B: BaseCache<Value = CacheEntry, Context = ExactCacheContext> + 'static,
{
fn prepare(
self: Arc<Self>,
options: CacheOptions,
surface: &str,
input: &dyn Fn() -> Value,
) -> PreparedCache {
PreparedCache {
policy: options.policy,
max_entry_bytes: self.config().max_entry_bytes,
call: Arc::new(NativeCacheCall {
request: options.request(&self.config().namespace, surface, input()),
service: self,
}),
}
}
}
struct NativeCacheCall<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
service: Arc<ResponseCache<B>>,
request: ResponseCacheRequest,
}
impl<B> ResponseCacheCall for NativeCacheCall<B>
where
B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>,
{
fn config(&self) -> &ResponseCacheConfig {
self.config()
fn lookup<'a>(&'a self, now: Duration) -> CacheFuture<'a, Option<CacheHit>> {
Box::pin(async move {
Ok(self
.service
.async_lookup(&self.request, now)
.await?
.map(|value| CacheHit {
value,
key: cache_key(&self.request.key),
}))
})
}
fn lookup<'a>(
&'a self,
request: &'a ResponseCacheRequest,
now: Duration,
) -> CacheFuture<'a, Option<Value>> {
Box::pin(self.async_lookup(request, now))
}
fn store<'a>(
&'a self,
request: &'a ResponseCacheRequest,
response: Value,
now: Duration,
) -> CacheFuture<'a, ()> {
Box::pin(self.async_store(request, response, now))
fn store<'a>(&'a self, response: Value, now: Duration) -> CacheFuture<'a, ()> {
Box::pin(self.service.async_store(&self.request, response, now))
}
}
@ -83,8 +116,16 @@ pub struct CachePolicy {
}
impl CachePolicy {
pub fn reads(&self) -> bool {
self.caching != Some(false) && !self.no_cache
}
pub fn writes(&self) -> bool {
self.caching != Some(false) && !self.no_store
}
pub fn enabled(&self) -> bool {
self.caching != Some(false) && !(self.no_cache && self.no_store)
self.reads() || self.writes()
}
}
@ -183,3 +224,56 @@ impl ScopedCache {
}
}
}
#[cfg(test)]
mod tests {
use std::cell::Cell;
use litellm_cache_memory::InMemoryCache;
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::shared(CacheScope::Shared)]
#[case::isolated(CacheScope::Isolated("tenant".into()))]
#[tokio::test]
async fn native_preparation_keeps_existing_keys_and_storage(#[case] scope: CacheScope) {
let service = Arc::new(
ResponseCache::new(Arc::new(InMemoryCache::default())).with_config(
ResponseCacheConfig {
namespace: "test".into(),
max_entry_bytes: 4096,
},
),
);
let options = CacheOptions::new(scope);
let input = json!({"model": "test", "body": {"z": 1, "a": 2}});
let legacy = options.clone().request("test", "test", input.clone());
let preparations = Cell::new(0);
let prepared = service.clone().prepare(options, "test", &|| {
preparations.set(preparations.get() + 1);
input.clone()
});
assert_eq!(prepared.max_entry_bytes, 4096);
assert_eq!(prepared.call.lookup(Duration::ZERO).await.unwrap(), None);
prepared
.call
.store(json!({"answer": 42}), Duration::ZERO)
.await
.unwrap();
assert_eq!(
service.async_lookup(&legacy, Duration::ZERO).await.unwrap(),
Some(json!({"answer": 42}))
);
assert_eq!(
prepared.call.lookup(Duration::ZERO).await.unwrap(),
Some(CacheHit {
value: json!({"answer": 42}),
key: cache_key(&legacy.key),
})
);
assert_eq!(preparations.get(), 1);
}
}

View file

@ -6,11 +6,10 @@ use std::{
time::Duration,
};
use litellm_cache::ExactCacheContext;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheConfig, ResponseCacheRequest,
ResponseCacheService,
CacheEntry, CacheKeyInput, CacheOptions, CachePolicy, CacheScope, ResponseCache,
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
};
use rstest::rstest;
use serde_json::json;
@ -25,48 +24,61 @@ async fn service_honors_per_call_expiry_and_freshness() {
Duration::from_secs(cache_clock.load(Ordering::SeqCst))
}),
)));
let request = ResponseCacheRequest {
context: ExactCacheContext {
let options = CacheOptions {
policy: CachePolicy {
ttl: Some(Duration::from_secs(5)),
..CachePolicy::default()
},
..ResponseCacheRequest::new(CacheKeyInput {
preset: Some("entry".into()),
..Default::default()
})
scope: CacheScope::Shared,
};
cache
.store(&request, json!({"answer":7}), Duration::ZERO)
let key =
litellm_cache_response::cache_key(&options.clone().request("", "test", json!("entry")).key);
let request = cache
.clone()
.prepare(options.clone(), "test", &|| json!("entry"));
request
.call
.store(json!({"answer":7}), Duration::ZERO)
.await
.unwrap();
assert_eq!(
cache.lookup(&request, Duration::ZERO).await.unwrap(),
Some(json!({"answer":7}))
request.call.lookup(Duration::ZERO).await.unwrap(),
Some(litellm_cache_response::CacheHit {
value: json!({"answer":7}),
key,
})
);
let stale_request = cache.prepare(
CacheOptions {
policy: CachePolicy {
max_age: Some(Duration::from_secs(1)),
..options.policy
},
..options
},
"test",
&|| json!("entry"),
);
let stale_request = ResponseCacheRequest {
max_age: Some(Duration::from_secs(1)),
..request.clone()
};
clock.store(2, Ordering::SeqCst);
assert_eq!(
cache
.lookup(&stale_request, Duration::from_secs(2))
stale_request
.call
.lookup(Duration::from_secs(2))
.await
.unwrap(),
None
);
assert!(
cache
.lookup(&request, Duration::from_secs(2))
request
.call
.lookup(Duration::from_secs(2))
.await
.unwrap()
.is_some()
);
clock.store(6, Ordering::SeqCst);
assert_eq!(
cache
.lookup(&request, Duration::from_secs(6))
.await
.unwrap(),
request.call.lookup(Duration::from_secs(6)).await.unwrap(),
None
);
}
@ -104,20 +116,33 @@ async fn entry_limit_applies_to_sync_async_and_batch_writes() {
)
.await
.unwrap();
let service: Arc<dyn ResponseCacheService> = Arc::new(cache);
let service = Arc::new(cache).prepare(CacheOptions::new(CacheScope::Shared), "test", &|| {
json!("service")
});
service
.store(&request("service"), small.clone(), Duration::ZERO)
.call
.store(small.clone(), Duration::ZERO)
.await
.unwrap();
for key in ["sync", "async", "batch-large"] {
assert!(storage.get_cache(key).unwrap().is_none());
}
for key in ["batch-small", "service"] {
assert_eq!(
service.lookup(&request(key), Duration::ZERO).await.unwrap(),
Some(small.clone())
);
}
assert!(storage.get_cache("sync").unwrap().is_none());
assert!(storage.get_cache("async").unwrap().is_none());
assert!(storage.get_cache("batch-large").unwrap().is_none());
assert_eq!(
storage
.get_cache("batch-small")
.unwrap()
.map(|entry| entry.response),
Some(small.clone())
);
assert_eq!(
service
.call
.lookup(Duration::ZERO)
.await
.unwrap()
.map(|hit| hit.value),
Some(small)
);
}
#[rstest]

View file

@ -77,6 +77,7 @@
"kwargs",
"call_type"
],
"reviews_deployment_success": [],
"after_deployment_success": [
"kwargs",
"response",

View file

@ -311,13 +311,12 @@ impl LegacyLogging {
self.prepare(py)
}
fn resume_after_success(
fn resume_provider_result(
&mut self,
py: Python<'_>,
_: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
self.response = Some(result?);
self.finalize(py)
result.map(HookStep::Ready)
}
fn resume_deployment_failure(
@ -423,20 +422,29 @@ impl LegacyLogging {
) -> PyResult<HookStep<Self, Py<PyAny>>> {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response);
if self.runs_deployment_hooks() {
return Ok(HookStep::Await(
DeploymentHooks::after_success(
py,
self.call.kwargs(),
&self.response,
self.call_type(),
)?,
Self::resume_after_success,
));
}
self.finalize(py)
}
pub(crate) fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool> {
Ok(self.runs_deployment_hooks() && DeploymentHooks::reviews_success(py)?)
}
pub(crate) fn provider_result(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
Ok(HookStep::Await(
DeploymentHooks::after_success(
py,
self.call.kwargs(),
&Some(response),
self.call_type(),
)?,
Self::resume_provider_result,
))
}
pub(crate) fn result_ready(
&mut self,
py: Python<'_>,
@ -787,8 +795,24 @@ assert hooked == ([opaque] if asynchronous else []), hooked
});
}
#[rstest]
#[case::synchronous(false, false)]
#[case::asynchronous(true, true)]
fn only_an_async_call_reviews_provider_results(
#[case] asynchronous: bool,
#[case] reviews: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"kwargs = {'logger': logger}");
let (mut logging, _) = begin(py, &locals, asynchronous);
assert_eq!(logging.reviews_provider_results(py).unwrap(), reviews);
});
}
#[rstest::rstest]
fn response_returned_by_the_post_call_hook_is_finalized_and_returned() {
fn the_success_hook_reviews_the_provider_result_and_its_answer_is_finalized() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
@ -809,18 +833,23 @@ logger.hooks = {'pre': lambda kwargs: kwargs}
)
.unwrap();
let step = logging
.transform_response(py, local(&locals, "response").unbind(), TIMING)
.provider_result(py, local(&locals, "response").unbind())
.unwrap();
assert!(awaits_deployment_hook(&step));
let step = resume(
let HookStep::Ready(reviewed) = resume(
&mut logging,
step,
py,
Ok(local(&locals, "replacement").unbind()),
)
.unwrap();
let HookStep::Ready(returned) = step else {
panic!("expected the finalized response");
.unwrap() else {
panic!("expected the hook's answer");
};
assert!(reviewed.bind(py).is(local(&locals, "replacement")));
let HookStep::Ready(returned) =
logging.transform_response(py, reviewed, TIMING).unwrap()
else {
panic!("transforming the public response runs no deployment hook");
};
assert!(returned.bind(py).is(local(&locals, "replacement")));
run(
@ -829,6 +858,7 @@ logger.hooks = {'pre': lambda kwargs: kwargs}
c"
[finalized] = [value for name, value in logger.calls if name == 'finalize']
assert finalized is replacement
assert [name for name in logger.names() if name.endswith('_hook')] == ['pre_hook', 'success_hook']
",
);
});
@ -853,7 +883,7 @@ assert finalized is replacement
)
.unwrap();
let step = logging
.transform_response(py, local(&locals, "response").unbind(), TIMING)
.provider_result(py, local(&locals, "response").unbind())
.unwrap();
resume(&mut logging, step, py, Err(cancellation))
.err()

View file

@ -95,6 +95,12 @@ impl DeploymentHooks {
.map(Bound::unbind)
}
pub fn reviews_success(py: Python<'_>) -> PyResult<bool> {
python::DeploymentHooks::ReviewsSuccess
.call(py, ())?
.extract()
}
pub fn after_success(
py: Python<'_>,
kwargs: &Py<PyDict>,

View file

@ -13,6 +13,7 @@ pub enum CallBoundary {
PrepareArguments,
BeforeProviderRequest,
AfterProviderResponse,
ProviderResult,
TransformResponse,
Succeeded,
Failed,
@ -53,6 +54,7 @@ type Prepare = fn(&mut LegacyLogging, Python<'_>, Py<PyDict>, f64) -> Step<Py<Py
type Before =
fn(&mut LegacyLogging, Python<'_>, Box<WireRequest>, &RequestContext) -> Step<Box<WireRequest>>;
type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>;
type ProviderResult = fn(&mut LegacyLogging, Python<'_>, Py<PyAny>) -> Step<Py<PyAny>>;
type Transform = fn(&mut LegacyLogging, Python<'_>, Py<PyAny>, Timing) -> Step<Py<PyAny>>;
type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py<PyAny>) -> Step<()>;
type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>;
@ -77,10 +79,16 @@ const AFTER: Binding<After> = Binding {
callbacks: &["log_post_api_call"],
};
const PROVIDER_RESULT: Binding<ProviderResult> = Binding {
boundary: CallBoundary::ProviderResult,
invoke: LegacyLogging::provider_result,
callbacks: &["async_post_call_success_deployment_hook"],
};
const TRANSFORM: Binding<Transform> = Binding {
boundary: CallBoundary::TransformResponse,
invoke: LegacyLogging::transform_public_response,
callbacks: &["async_post_call_success_deployment_hook"],
callbacks: &[],
};
const SUCCESS: Binding<Success> = Binding {
@ -127,6 +135,7 @@ pub fn callback_mappings() -> impl Iterator<Item = CallbackMapping> {
.mappings()
.chain(BEFORE.mappings())
.chain(AFTER.mappings())
.chain(PROVIDER_RESULT.mappings())
.chain(TRANSFORM.mappings())
.chain(SUCCESS.mappings())
.chain(FAILURE.mappings())
@ -159,6 +168,14 @@ impl CallHooks<PythonRuntime> for LegacyLogging {
(BEFORE.invoke)(self, py, wire, context)
}
fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool> {
LegacyLogging::reviews_provider_results(self, py)
}
fn provider_result(&mut self, py: Python<'_>, response: Py<PyAny>) -> Step<Py<PyAny>> {
(PROVIDER_RESULT.invoke)(self, py, response)
}
fn transform_response(
&mut self,
py: Python<'_>,

View file

@ -65,6 +65,8 @@ pub(crate) enum Logging {
pub(crate) enum DeploymentHooks {
#[strum(serialize = "before_deployment_call")]
BeforeDeploymentCall,
#[strum(serialize = "reviews_deployment_success")]
ReviewsSuccess,
#[strum(serialize = "after_deployment_success")]
AfterDeploymentSuccess,
#[strum(serialize = "after_deployment_failure")]

View file

@ -80,6 +80,7 @@ FAKES = {
'custom_pricing_fields': lambda: ('ocr_cost_per_page',),
'is_internal_call': lambda: legacy.is_internal.get(),
'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type),
'reviews_deployment_success': lambda: True,
'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook(
'success', response, call_type
),

View file

@ -47,7 +47,7 @@ async fn all_inference_endpoints_share_native_cache(
.expect(2)
.mount(&upstream)
.await;
let cache: Arc<dyn ResponseCacheService> = Arc::new(
let cache = Arc::new(
ResponseCache::new(Arc::new(InMemoryCache::new(
Some(100),
Some(Duration::from_secs(60)),
@ -72,7 +72,7 @@ async fn all_inference_endpoints_share_native_cache(
let cache_key = second.headers().get("x-litellm-cache-key").unwrap().clone();
assert!(!cache_key.as_bytes().is_empty());
let stored = cache
.lookup(
.async_lookup(
&ResponseCacheRequest::new(CacheKeyInput {
preset: Some(cache_key.to_str().unwrap().into()),
..Default::default()

View file

@ -13,6 +13,7 @@ bytes.workspace = true
futures-util.workspace = true
parking_lot.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }

View file

@ -2,6 +2,7 @@ use crate::{InvokeError, PythonOwned};
use litellm_host::protocol::Protocol;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use serde_json::Value;
/// Converts requests, responses, stream values and errors at the Python boundary.
pub trait PythonBinding: PythonOwned {
@ -27,6 +28,16 @@ pub trait PythonBinding: PythonOwned {
response: <Self::Protocol as Protocol>::Response,
) -> PyResult<Py<PyAny>>;
/// A provider result as the public object Python hooks review, built from its JSON form.
fn public_response(&mut self, py: Python<'_>, response: &Value) -> PyResult<Py<PyAny>> {
crate::to_py(py, response)
}
/// The JSON form of the public object a Python hook handed back.
fn response_value(&mut self, _py: Python<'_>, public: &Bound<'_, PyAny>) -> PyResult<Value> {
crate::from_py(public)
}
/// What the stream carries at hand-off, as the caller's stream receives it.
fn encode_stream_head(
&mut self,

View file

@ -294,6 +294,30 @@ where
reply.send(*wire);
Ok(())
}
HostRequest::Intercept(InterceptRequest::ProviderResult { response, reply }) => {
let step = {
let mut parts = self.parts();
let Parts { binding, hooks, .. } = &mut *parts;
Python::attach(|py| {
if !hooks.reviews_provider_results(py)? {
return Ok(None);
}
let public = binding.public_response(py, &response)?;
hooks.provider_result(py, public).map(Some)
})?
};
let Some(step) = step else {
reply.send(response);
return Ok(());
};
let public = self.settle(step).await?;
let reviewed = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.response_value(py, public.bind(py)))?
};
reply.send(reviewed);
Ok(())
}
HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => {
self.notify(PythonCallEvent::Execution(ExecutionEvent::ResultReady {
facts,

View file

@ -71,6 +71,12 @@ pub(super) trait ChainHooks: PythonOwned {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Box<WireRequest>>>;
fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool>;
fn provider_result(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
) -> PyResult<ChainStep<Py<PyAny>>>;
fn transform_response(
&mut self,
py: Python<'_>,
@ -138,6 +144,19 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
Ok(self.step(step, Continuation::Wire))
}
fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool> {
self.hooks.reviews_provider_results(py)
}
fn provider_result(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
) -> PyResult<ChainStep<Py<PyAny>>> {
let step = self.hooks.provider_result(py, response)?;
Ok(self.step(step, Continuation::Response))
}
fn transform_response(
&mut self,
py: Python<'_>,

View file

@ -28,6 +28,7 @@ pub struct HookChain {
enum Pending {
Arguments(usize, f64),
Wire(usize, RequestContext),
ProviderResult(usize),
Response(usize, Timing),
}
@ -124,6 +125,41 @@ impl Transform for Wire {
}
}
struct ProviderResult;
impl Transform for ProviderResult {
type Value = Py<PyAny>;
type Context = ();
fn call(
hooks: &mut dyn ChainHooks,
py: Python<'_>,
value: Py<PyAny>,
_: &(),
) -> PyResult<ChainStep<Py<PyAny>>> {
hooks.provider_result(py, value)
}
fn resume(
hooks: &mut dyn ChainHooks,
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Py<PyAny>>> {
hooks.resume_response(py, result)
}
fn suspend(index: usize, (): ()) -> Pending {
Pending::ProviderResult(index)
}
fn resumed(pending: Pending) -> Option<(usize, ())> {
match pending {
Pending::ProviderResult(index) => Some((index, ())),
_ => None,
}
}
}
struct Response;
impl Transform for Response {
@ -386,6 +422,23 @@ impl CallHooks<PythonRuntime> for HookChain {
self.transform_from::<Wire>(py, 0, wire, context.clone())
}
fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool> {
for hooks in &mut self.hooks {
if hooks.reviews_provider_results(py)? {
return Ok(true);
}
}
Ok(false)
}
fn provider_result(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
self.transform_from::<ProviderResult>(py, 0, response, ())
}
fn transform_response(
&mut self,
py: Python<'_>,

View file

@ -24,7 +24,7 @@ use rstest::rstest;
use support::{
HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options,
failing_machine, handed, hooks, interpreter, lifecycle_module, open_stream, read_all,
run_composed, run_scripted, settled, success_machine,
reviewed_machine, run_composed, run_scripted, settled, success_machine,
};
fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
@ -325,6 +325,34 @@ fn success_runs_every_step_in_order_and_returns_the_public_response(
});
}
#[rstest]
#[case::reviewed(HookScript::ReviewResult, "reviewed", &["provider_result:provider"])]
#[case::not_reviewed(HookScript::Plain, "provider", &[])]
fn a_provider_result_keeps_what_the_reviewing_hooks_answer(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
#[case] script: HookScript,
#[case] kept: &str,
#[case] reviewed: &[&str],
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
reviewed_machine(),
OpScript::Answer,
script,
asynchronous,
);
let seen: Vec<_> = log
.iter()
.filter(|entry| entry.starts_with("provider_result:"))
.cloned()
.collect();
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), kept);
assert_eq!(seen, reviewed);
});
}
#[rstest]
fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
#[from(interpreter)] _python: Interpreter,

View file

@ -321,6 +321,7 @@ pub enum HookScript {
FailAfterSuccess,
CancelTerminal,
FailTerminal,
ReviewResult,
}
pub struct SyntheticHooks {
@ -336,6 +337,24 @@ pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
}
impl CallHooks<PythonRuntime> for SyntheticHooks {
fn reviews_provider_results(&mut self, _: Python<'_>) -> PyResult<bool> {
Ok(matches!(self.script, HookScript::ReviewResult))
}
fn provider_result(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
let reviewed = response.bind(py).cast::<PyDict>()?.copy()?;
self.log.push(format!(
"provider_result:{}",
reviewed.get_item("text")?.unwrap()
));
reviewed.set_item("text", "reviewed")?;
Ok(HookStep::Ready(reviewed.into_any().unbind()))
}
fn prepare_arguments(
&mut self,
_: Python<'_>,
@ -474,6 +493,21 @@ pub fn success_machine() -> MachineFactory {
})
}
/// A call that offers its provider result for review and returns the text it keeps.
pub fn reviewed_machine() -> MachineFactory {
Box::new(move |_| {
CallMachine::new(None, move |host| {
Box::pin(async move {
let kept = host
.interceptors
.provider_result(serde_json::json!({"text": "provider"}))
.await?;
Ok(kept["text"].as_str().unwrap_or_default().to_owned())
})
})
})
}
pub fn failing_machine() -> MachineFactory {
Box::new(move |_| {
CallMachine::new(None, |_| {

View file

@ -32,3 +32,5 @@ Keep API policy in gateway-inference and python-bridge, and legacy callback poli
Interceptors run inline and can rewrite values or fail execution. Observers consume owned `CallEvent` snapshots from `observation_channel`. Its bounded `ObservationSender` never waits for delivery and counts events dropped when the queue is full or closed. The host owns receiver processing and draining. Pass the same publisher to machine construction and the driver when one receiver should collect execution and lifecycle events. Legacy Python callbacks retain their existing awaited, fallible behavior through the Python adapter
`ExecutionFacts` and `ResultSource` describe execution without pricing or budget policy. `InterceptPoints::result_ready` delivers these facts through an awaited `InterceptRequest::ResultReady`. Hosts receive them before response transformation or stream delivery. `ExecutionEvent::ResultReady` is the matching lifecycle event and can also be published as a passive snapshot. Accounting must consume the awaited path rather than a lossy observation queue
`InterceptRequest::ProviderResult` offers a complete provider response before a route caches or delivers it, and the route keeps the answer. Routes raise it only for provider-sourced results, never for a cache hit, so hooks that run once per provider call sit inside the cache

View file

@ -9,6 +9,7 @@ repository.workspace = true
futures-util.workspace = true
litellm-auth.workspace = true
litellm-coroutine.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio = { workspace = true, features = ["sync"] }

View file

@ -44,6 +44,21 @@ pub trait CallHooks<R: HookRuntime>: Sized {
Ok(R::ready(wire))
}
/// Whether `provider_result` would do anything, so a host can skip encoding the response.
fn reviews_provider_results(&mut self, _runtime: R::Context<'_>) -> Result<bool, R::Error> {
Ok(false)
}
/// A complete provider response before it is cached or delivered. What this returns
/// replaces it. Never offered a cache hit.
fn provider_result(
&mut self,
_runtime: R::Context<'_>,
response: R::Response,
) -> Result<R::Step<Self, R::Response>, R::Error> {
Ok(R::ready(response))
}
fn transform_response(
&mut self,
_runtime: R::Context<'_>,

View file

@ -1,5 +1,7 @@
use std::future::Future;
use serde::{Serialize, de::DeserializeOwned};
use crate::{
machine::MachineFault,
protocol::{InterceptRequest, Reply, reply},
@ -145,6 +147,25 @@ where
fn result_ready(&self, facts: ExecutionFacts) -> impl Future<Output = Result<(), E>> + Send {
self.intercept(|reply| InterceptRequest::ResultReady { facts, reply })
}
/// Offers a complete provider response before it is cached or delivered, and returns the
/// response the call keeps. Never raised for a cache hit.
fn provider_result<R>(&self, response: R) -> impl Future<Output = Result<R, E>> + Send
where
R: Serialize + DeserializeOwned + Send,
{
let value = serde_json::to_value(&response);
async move {
let value = value.map_err(|_| MachineFault::UnreadableAnswer)?;
let answer = self
.intercept(|reply| InterceptRequest::ProviderResult {
response: value,
reply,
})
.await?;
serde_json::from_value(answer).map_err(|_| MachineFault::UnreadableAnswer.into())
}
}
}
impl<E: From<MachineFault>, T: Interceptors<E> + ?Sized> InterceptPoints<E> for T {}

View file

@ -16,6 +16,8 @@ pub enum MachineFault {
Abandoned,
/// The host resumed the call out of turn.
Protocol(ResumeError),
/// The host answered with a value the call cannot read back as its own type.
UnreadableAnswer,
}
pub type ExecuteFuture<R, C = <R as Protocol>::Response> =

View file

@ -33,6 +33,12 @@ pub enum InterceptRequest {
raw: RawResponse,
reply: Reply<()>,
},
/// A complete response the provider just produced, before anything caches or delivers it.
/// Never raised for a cache hit. The answer is the response the call keeps.
ProviderResult {
response: serde_json::Value,
reply: Reply<serde_json::Value>,
},
}
impl InterceptRequest {
@ -43,6 +49,7 @@ impl InterceptRequest {
reply.send(());
}
Self::BeforeProviderRequest { wire, reply, .. } => reply.send(*wire),
Self::ProviderResult { response, reply } => reply.send(response),
}
}
@ -68,6 +75,9 @@ impl InterceptRequest {
Self::AfterProviderResponse { raw, reply } => {
reply.send(interceptors.after_provider_response(raw).await?);
}
Self::ProviderResult { response, reply } => {
reply.send(interceptors.provider_result(response).await?);
}
}
Ok(())
}

View file

@ -59,14 +59,11 @@ pub(super) async fn execute(
)
.await?;
let cache = cache.filter(|_| authenticated.signer.is_none());
let cache_request = litellm_inference::caching::CacheRequest::from_wire(
identity,
cache.as_ref().map(|_| &wire),
);
let cache_request = litellm_inference::caching::CacheRequest::from_wire::<
super::route::ChatCompletions,
>(identity, &wire, cache.as_ref(), cache_options);
litellm_inference::caching::execute_unary::<super::route::ChatCompletions, _, _>(
cache_request,
cache.as_ref().map(|cache| cache.service.clone()),
cache.as_ref().map(|cache| cache.options(cache_options)),
interceptors,
observers,
|| async move {

View file

@ -12,7 +12,7 @@ use prepare::{prepare_provider_request, resolve_request};
use litellm_auth::AuthServices;
use litellm_secrets::source::SecretSource;
use std::sync::Arc;
use types::ChatCompletionsRequest;
use types::{ChatCompletionsCall, ChatCompletionsRequest};
#[derive(Clone)]
pub struct ChatCompletionsRoute {
@ -48,6 +48,16 @@ impl ChatCompletionsRoute {
request: ChatCompletionsRequest<'_>,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
options: impl Into<litellm_inference::CallOptions>,
) -> Result<ChatCompletionsResponse, Error> {
self.execute_call(request.into(), interceptors, options)
.await
}
pub async fn execute_call(
&self,
call: ChatCompletionsCall,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
options: impl Into<litellm_inference::CallOptions>,
) -> Result<ChatCompletionsResponse, Error> {
let litellm_inference::CallOptions {
cache: cache_options,
@ -55,12 +65,7 @@ impl ChatCompletionsRoute {
} = options.into();
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run_call(
request.into(),
cache_options,
interceptors,
observers.as_ref(),
),
self.run_call(call, cache_options, interceptors, observers.as_ref()),
)
.await
}

View file

@ -181,6 +181,7 @@ impl Interceptors<RouteError> for ChangingHooks {
InterceptRequest::ResultReady { facts, reply } => {
reply.send(self.result_ready(facts).await?);
}
other => other.pass_through(),
}
Ok(())
}

View file

@ -106,6 +106,12 @@ impl MessagesRoute {
source: source.clone(),
})
.await?;
let output = litellm_inference::caching::provider_result::<route::Messages>(
output,
&source,
context.interceptors,
)
.await?;
Ok(cache.finish(output, &source).await)
})
.await

View file

@ -11,7 +11,7 @@ use std::{
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{
CacheOptions, CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, ScopedCache,
CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, ScopedCache,
};
use litellm_host::interceptors::{
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext,
@ -58,22 +58,22 @@ async fn messages_cache_identity_includes_provider_native_parameters(
for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] {
let response =
execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
CacheRequest::from_wire::<Messages>(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
Some(&WireRequest {
&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
},
Some(&ScopedCache::new(cache.clone(), CacheScope::Shared)),
None,
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
@ -148,6 +148,7 @@ impl Interceptors<RouteError> for ChangingHooks {
InterceptRequest::ResultReady { facts, reply } => {
reply.send(self.result_ready(facts).await?);
}
other => other.pass_through(),
}
Ok(())
}

View file

@ -54,6 +54,7 @@ impl Interceptors<Error> for LiveCall {
InterceptRequest::AfterProviderResponse { raw, reply } => {
reply.send(self.after_provider_response(raw).await?);
}
other => other.pass_through(),
}
Ok(())
}

View file

@ -95,6 +95,7 @@ impl litellm_host::interceptors::Interceptors<<Messages as litellm_host::protoco
InterceptRequest::AfterProviderResponse { raw, reply } => {
reply.send(self.after_provider_response(raw).await?);
}
other => other.pass_through(),
}
Ok(())
}

View file

@ -1,3 +1,4 @@
use litellm_host::interceptors::InterceptPoints;
use litellm_host::observation::ObservationSender;
use std::sync::Arc;
@ -58,7 +59,7 @@ impl OcrRoute {
caller_document,
observers,
));
execute.await
interceptors.provider_result(execute.await?).await
})
.await
}

View file

@ -70,6 +70,11 @@ async fn drive_until(
.map(|()| reply.send(()))
.map_err(HostFailure::Error)
}
HostRequest::Intercept(other) => {
ops.push("provider_result");
other.pass_through();
Ok(())
}
};
if let Err(failure) = answer {
break machine.interrupt(failure).await.map(completed);
@ -90,6 +95,7 @@ async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, sto
MachineStep::Suspended(HostRequest::HostCall(op)) => host.handle_host_call(op).await.unwrap(),
MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, reply, .. })) => reply.send(*wire),
MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(()),
MachineStep::Suspended(HostRequest::Intercept(other)) => other.pass_through(),
MachineStep::Suspended(HostRequest::Stream(stream)) => match stream {
litellm_host::protocol::StreamDelivery::Open(head, _) => match head {},
litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {},
@ -117,7 +123,7 @@ async fn a_hand_driven_machine_performs_the_same_call() {
assert_eq!(outcome.unwrap().pages[0].markdown, "native");
assert_eq!(received(&upstream).await.len(), 1);
assert_eq!(ops, ["BeforeSend", "response"]);
assert_eq!(ops, ["BeforeSend", "response", "provider_result"]);
assert!(matches!(
machine.resume().await,
Err(Error::InvalidRequest(_))
@ -146,7 +152,7 @@ async fn a_path_document_is_read_by_core_without_a_host_operation() {
std::fs::remove_dir_all(&dir).unwrap();
assert_eq!(response.unwrap().pages[0].markdown, "path");
assert_eq!(ops, ["BeforeSend", "response"]);
assert_eq!(ops, ["BeforeSend", "response", "provider_result"]);
assert_eq!(
only_request(&upstream).await.json()["document"]["image_url"],
"data:image/png;base64,YWJj"

View file

@ -37,14 +37,11 @@ pub(super) async fn execute(
)
.await?;
let cache = cache.filter(|_| authenticated.signer.is_none());
let cache_request = litellm_inference::caching::CacheRequest::from_wire(
identity,
cache.as_ref().map(|_| &wire),
);
let cache_request = litellm_inference::caching::CacheRequest::from_wire::<
super::route::Responses,
>(identity, &wire, cache.as_ref(), cache_options);
litellm_inference::caching::execute_streaming::<super::route::Responses, _, _>(
cache_request,
cache.as_ref().map(|cache| cache.service.clone()),
cache.as_ref().map(|cache| cache.options(cache_options)),
interceptors,
observers,
|| async move {

View file

@ -25,14 +25,20 @@ use litellm_inference::{
use rstest::{fixture, rstest};
use serde_json::{Value, json};
fn cache_request(input: Value) -> CacheRequest {
CacheRequest {
identity: ProviderIdentity {
fn cache_request<P: litellm_inference::caching::Cachable>(
input: Value,
cache: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
) -> CacheRequest {
CacheRequest::prepare::<P>(
ProviderIdentity {
model: "test-model".into(),
provider: "test-provider".into(),
},
input,
}
cache,
options,
&|_| input.clone(),
)
}
#[fixture]
@ -59,31 +65,55 @@ struct InvalidEntryCache(
);
impl ResponseCacheService for InvalidEntryCache {
fn config(&self) -> &ResponseCacheConfig {
self.0.config()
fn prepare(
self: Arc<Self>,
options: CacheOptions,
surface: &str,
input: &dyn Fn() -> Value,
) -> litellm_cache_response::PreparedCache {
litellm_cache_response::PreparedCache {
policy: options.policy,
max_entry_bytes: self.0.config().max_entry_bytes,
call: Arc::new(InvalidEntryCall {
request: options.request(&self.0.config().namespace, surface, input()),
service: self,
}),
}
}
}
struct InvalidEntryCall {
service: Arc<InvalidEntryCache>,
request: litellm_cache_response::ResponseCacheRequest,
}
impl litellm_cache_response::ResponseCacheCall for InvalidEntryCall {
fn lookup<'a>(
&'a self,
request: &'a litellm_cache_response::ResponseCacheRequest,
now: Duration,
) -> futures_util::future::BoxFuture<'a, Result<Option<Value>, litellm_cache::Error>> {
) -> futures_util::future::BoxFuture<
'a,
Result<Option<litellm_cache_response::CacheHit>, litellm_cache::Error>,
> {
Box::pin(async move {
Ok(self
.0
.async_lookup(request, now)
.await?
.or_else(|| Some(self.1.clone())))
Ok(Some(litellm_cache_response::CacheHit {
value: self
.service
.0
.async_lookup(&self.request, now)
.await?
.unwrap_or_else(|| self.service.1.clone()),
key: litellm_cache_response::cache_key(&self.request.key),
}))
})
}
fn store<'a>(
&'a self,
request: &'a litellm_cache_response::ResponseCacheRequest,
response: Value,
now: Duration,
) -> futures_util::future::BoxFuture<'a, Result<(), litellm_cache::Error>> {
Box::pin(self.0.async_store(request, response, now))
Box::pin(self.service.0.async_store(&self.request, response, now))
}
}
@ -139,6 +169,7 @@ impl Interceptors<RouteError> for ChangingHooks {
InterceptRequest::ResultReady { facts, reply } => {
reply.send(self.result_ready(facts).await?);
}
other => other.pass_through(),
}
Ok(())
}
@ -184,9 +215,11 @@ async fn responses_refetches_instead_of_deserializing_another_api_response(
let calls = AtomicUsize::new(0);
for _ in 0..2 {
let response = execute_unary::<Responses, _, _>(
cache_request(json!({"input":"hello"})),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
cache_request::<Responses>(
json!({"input":"hello"}),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
),
&(),
None,
|| async {
@ -226,9 +259,11 @@ async fn responses_cache_only_reuses_completed_responses(
let calls = AtomicUsize::new(0);
for _ in 0..2 {
let response = execute_unary::<Responses, _, _>(
cache_request(json!({"input":"hello"})),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
cache_request::<Responses>(
json!({"input":"hello"}),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
),
&(),
None,
|| async {

View file

@ -67,7 +67,7 @@
- Attach a `litellm_cache_response::ScopedCache` with `route.with_cache(cache)`. Cached and uncached routes use the same `execute` and `machine` methods
- `CallOptions` carries a scope-free `CachePolicy` and observation. Per-call policy never replaces the attached scope or service
- Provider transport does not own cache orchestration. Stream capture stays in `src/caching.rs`
- This crate and the format crates own request identity, typed response reconstruction, and stream capture and replay
- This crate and the format crates expose lazy effective-request identity and own typed response reconstruction and stream capture and replay. The selected cache service prepares per-call state before lookup. Only the native adapter materializes that identity and constructs native key inputs. Python delegation uses Python kwargs and never invokes the native identity builder
- `cache-response` owns cache policy, namespacing, scope encoding, versioned envelopes, and freshness
- The SDK explicitly chooses shared scope. The gateway derives isolated scope from authenticated identity before attaching its service
- `ExecutionFacts` are delivered through the awaited `ResultReady` host operation for both provider and cached results, before public response processing or stream opening
@ -79,5 +79,5 @@
- native gateway accounting belongs to gateway dependencies, independently of `host-python`
- response-cache services expose no coordination counters or reservation APIs. A shared Redis deployment does not make response storage and accounting coordination the same dependency
- Cache lookup follows provider preparation, credential resolution, and the request interceptor
- keys describe the effective provider URL, authenticated headers, and rewritten body
- native keys describe the effective provider URL, authenticated headers, and rewritten body. Python delegation keeps Python's key semantics and returns the host's hit key unchanged
- signed requests bypass caching until the signing identity has a stable cache representation

View file

@ -8,8 +8,7 @@ use std::{
use bytes::{Bytes, BytesMut};
use futures_util::{StreamExt, TryStreamExt, stream};
use litellm_cache_response::{
CacheOptions, CachePolicy, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope,
ScopedCache, cache_key,
CacheOptions, CachePolicy, PreparedCache, ResponseCacheService, ResponseEnvelope, ScopedCache,
};
use litellm_host::{
call::{CallOutput, OutputOf},
@ -37,13 +36,36 @@ pub trait Cachable: Protocol<Error = RouteError> {
pub struct CacheRequest {
pub identity: ProviderIdentity,
pub input: Value,
session: Option<CacheSession>,
}
impl CacheRequest {
pub fn from_wire(identity: ProviderIdentity, wire: Option<&WireRequest>) -> Self {
Self {
input: wire.map_or(Value::Null, |wire| {
pub fn prepare<P: Cachable>(
identity: ProviderIdentity,
service: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
input: &dyn Fn(&ProviderIdentity) -> Value,
) -> Self {
let session = options
.filter(|options| options.policy.enabled())
.zip(service)
.map(|(options, service)| CacheSession {
prepared: service.prepare(options, P::SURFACE, &|| input(&identity)),
});
Self { identity, session }
}
pub fn from_wire<P: Cachable>(
identity: ProviderIdentity,
wire: &WireRequest,
cache: Option<&ScopedCache>,
policy: Option<CachePolicy>,
) -> Self {
Self::prepare::<P>(
identity,
cache.map(|cache| cache.service.clone()),
cache.map(|cache| cache.options(policy)),
&|identity| {
serde_json::json!({
"provider": identity.provider,
"model": identity.model,
@ -51,9 +73,8 @@ impl CacheRequest {
"headers": wire.headers,
"body": wire.body,
})
}),
identity,
}
},
)
}
}
@ -72,35 +93,23 @@ pub enum CachedOutput<R> {
}
struct CacheSession {
service: Arc<dyn ResponseCacheService>,
request: ResponseCacheRequest,
prepared: PreparedCache,
}
impl CacheSession {
fn prepare<P: Cachable>(
service: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
request: &CacheRequest,
) -> Option<Self> {
let options = options.filter(|options| options.policy.enabled())?;
let service = service?;
let input = request.input.clone();
let request = options.request(&service.config().namespace, P::SURFACE, input);
Some(Self { service, request })
}
async fn lookup<P: Cachable>(&self) -> Option<CachedOutput<P::Response>>
async fn lookup<P: Cachable>(&self) -> Option<(CachedOutput<P::Response>, String)>
where
P::Response: DeserializeOwned,
{
if !self.request.controls.reads() {
if !self.prepared.policy.reads() {
return None;
}
match self.service.lookup(&self.request, now()).await {
Ok(Some(value)) => {
serde_json::from_value::<ResponseEnvelope<CachedOutput<P::Response>>>(value)
match self.prepared.call.lookup(now()).await {
Ok(Some(hit)) => {
serde_json::from_value::<ResponseEnvelope<CachedOutput<P::Response>>>(hit.value)
.ok()
.and_then(|entry| entry.decode(P::SURFACE))
.map(|output| (output, hit.key))
}
Ok(None) => None,
Err(_) => {
@ -111,15 +120,10 @@ impl CacheSession {
}
async fn store(&self, entry: Value) {
if !self.request.controls.writes() {
if !self.prepared.policy.writes() {
return;
}
if self
.service
.store(&self.request, entry, now())
.await
.is_err()
{
if self.prepared.call.store(entry, now()).await.is_err() {
tracing::warn!("response cache write failed");
}
}
@ -128,7 +132,7 @@ impl CacheSession {
where
P::Response: Serialize,
{
if !self.request.controls.writes() || !P::reusable(response) {
if !self.prepared.policy.writes() || !P::reusable(response) {
return;
}
if let Ok(value) = serde_json::to_value(response)
@ -144,8 +148,6 @@ impl CacheSession {
pub async fn execute_unary<P, F, Fut>(
request: CacheRequest,
cache: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
interceptors: &impl Interceptors<RouteError>,
observers: Option<&ObservationSender>,
provider: F,
@ -156,14 +158,16 @@ where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<P::Response, RouteError>>,
{
let identity = request.identity.clone();
let CacheRequest { identity, session } = request;
crate::diagnostic::provider(&identity.model, &identity.provider);
let session = CacheSession::prepare::<P>(cache, options, &request);
let hit = match &session {
Some(session) => session.lookup::<P>().await.and_then(|entry| match entry {
CachedOutput::Response(response) => Some((response, cache_key(&session.request.key))),
CachedOutput::Stream(_) => None,
}),
Some(session) => session
.lookup::<P>()
.await
.and_then(|(entry, key)| match entry {
CachedOutput::Response(response) => Some((response, key)),
CachedOutput::Stream(_) => None,
}),
None => None,
};
let (response, source) = match hit {
@ -182,7 +186,11 @@ where
observers,
)
.await?;
if from_provider && let Some(session) = session {
if !from_provider {
return Ok(response);
}
let response = interceptors.provider_result(response).await?;
if let Some(session) = session {
session.store_response::<P>(&response).await;
}
Ok(response)
@ -190,8 +198,6 @@ where
pub async fn execute_streaming<P, F, Fut>(
request: CacheRequest,
cache: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
interceptors: &impl Interceptors<RouteError>,
observers: Option<&ObservationSender>,
provider: F,
@ -202,9 +208,8 @@ where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<OutputOf<P>, RouteError>>,
{
let identity = request.identity.clone();
let CacheRequest { identity, session } = request;
crate::diagnostic::provider(&identity.model, &identity.provider);
let session = CacheSession::prepare::<P>(cache, options, &request);
let cache = CallCache::<P> {
session,
protocol: PhantomData,
@ -224,9 +229,29 @@ where
observers,
)
.await?;
let output = provider_result::<P>(output, &source, interceptors).await?;
Ok(cache.finish(output, &source).await)
}
/// Offers a complete provider response to the host before it is cached or delivered. A cache
/// hit and a stream pass through untouched.
pub async fn provider_result<P>(
output: OutputOf<P>,
source: &ResultSource,
interceptors: &impl Interceptors<RouteError>,
) -> Result<OutputOf<P>, RouteError>
where
P: Protocol<Error = RouteError>,
P::Response: Serialize + DeserializeOwned,
{
match (source, output) {
(ResultSource::Provider, CallOutput::Complete(response)) => Ok(CallOutput::Complete(
interceptors.provider_result(response).await?,
)),
(_, output) => Ok(output),
}
}
pub struct CallCache<P> {
session: Option<CacheSession>,
protocol: PhantomData<P>,
@ -239,23 +264,9 @@ impl<P: StreamCachable> CallCache<P> {
identity: &ProviderIdentity,
wire: &WireRequest,
) -> Self {
let session = cache.and_then(|cache| {
if !policy.enabled() {
return None;
}
let options = cache.options(Some(policy));
let request = CacheRequest::from_wire(identity.clone(), Some(wire));
Some(CacheSession {
request: options.request(
&cache.service.config().namespace,
P::SURFACE,
request.input,
),
service: cache.service.clone(),
})
});
let request = CacheRequest::from_wire::<P>(identity.clone(), wire, cache, Some(policy));
Self {
session,
session: request.session,
protocol: PhantomData,
}
}
@ -265,16 +276,12 @@ impl<P: StreamCachable> CallCache<P> {
P::Response: DeserializeOwned,
{
let session = self.session.as_ref()?;
let output = match session.lookup::<P>().await? {
let (entry, key) = session.lookup::<P>().await?;
let output = match entry {
CachedOutput::Response(response) => CallOutput::Complete(response),
CachedOutput::Stream(data) => P::replay(Bytes::from(data))?,
};
Some((
output,
ResultSource::Cache {
key: cache_key(&session.request.key),
},
))
Some((output, ResultSource::Cache { key }))
}
pub async fn finish(self, output: OutputOf<P>, source: &ResultSource) -> OutputOf<P>
@ -282,7 +289,7 @@ impl<P: StreamCachable> CallCache<P> {
P::Response: Serialize,
{
let Some(session) = self.session.filter(|session| {
*source == ResultSource::Provider && session.request.controls.writes()
*source == ResultSource::Provider && session.prepared.policy.writes()
}) else {
return output;
};
@ -310,8 +317,7 @@ fn capture_stream<P: StreamCachable>(
Some(chunk) => {
let captured = captured.and_then(|mut data| {
let bytes = P::bytes(&chunk);
if data.len().saturating_add(bytes.len())
> session.service.config().max_entry_bytes
if data.len().saturating_add(bytes.len()) > session.prepared.max_entry_bytes
{
return None;
}

View file

@ -47,6 +47,7 @@ impl From<litellm_host::machine::MachineFault> for RouteError {
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("host {message}").into(),
MachineFault::UnreadableAnswer => "host answered with an unreadable response".into(),
})
}
}

View file

@ -61,14 +61,20 @@ impl StreamCachable for TestRoute {
}
}
fn cache_request(input: Value) -> CacheRequest {
CacheRequest {
identity: ProviderIdentity {
fn cache_request<P: litellm_inference::caching::Cachable>(
input: Value,
cache: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
) -> CacheRequest {
CacheRequest::prepare::<P>(
ProviderIdentity {
model: "test-model".into(),
provider: "test-provider".into(),
},
input,
}
cache,
options,
&|_| input.clone(),
)
}
#[fixture]
@ -76,6 +82,38 @@ fn cache() -> Arc<dyn ResponseCacheService> {
cache_with_limit(4096)
}
#[rstest]
#[case::no_service(false, true, CachePolicy::default())]
#[case::no_options(true, false, CachePolicy::default())]
#[case::off(true, true, CachePolicy { caching: Some(false), ..CachePolicy::default() })]
#[case::no_io(true, true, CachePolicy { no_cache: true, no_store: true, ..CachePolicy::default() })]
#[tokio::test]
async fn disabled_caching_does_not_prepare_identity(
cache: Arc<dyn ResponseCacheService>,
#[case] configured: bool,
#[case] options_present: bool,
#[case] policy: CachePolicy,
) {
let request = CacheRequest::prepare::<TestRoute>(
ProviderIdentity {
model: "test".into(),
provider: "test".into(),
},
configured.then_some(cache),
options_present.then_some(CacheOptions {
policy,
scope: CacheScope::Shared,
}),
&|_| panic!("disabled cache must not prepare identity"),
);
let response = execute_unary::<TestRoute, _, _>(request, &(), None, || async {
Ok(json!({"source": "provider"}))
})
.await
.unwrap();
assert_eq!(response, json!({"source": "provider"}));
}
fn cache_with_limit(max_entry_bytes: usize) -> Arc<dyn ResponseCacheService> {
Arc::new(
ResponseCache::new(Arc::new(InMemoryCache::new(
@ -96,9 +134,7 @@ async fn call(
request: Value,
) -> Value {
let output = execute_streaming::<TestRoute, _, _>(
cache_request(request),
Some(cache.clone()),
options,
cache_request::<TestRoute>(request, Some(cache.clone()), options),
&(),
None,
|| async {
@ -193,9 +229,11 @@ async fn streamed(
fail: bool,
) -> OutputOf<TestRoute> {
execute_streaming::<TestRoute, _, _>(
cache_request(json!({"stream":true})),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
cache_request::<TestRoute>(
json!({"stream":true}),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
),
&(),
None,
|| async {
@ -309,9 +347,11 @@ async fn oversized_streams_are_delivered_without_being_stored() {
async fn a_provider_failure_never_populates_the_cache(cache: Arc<dyn ResponseCacheService>) {
let calls = AtomicUsize::new(0);
let first = execute_streaming::<TestRoute, _, _>(
cache_request(json!({})),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
cache_request::<TestRoute>(
json!({}),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
),
&(),
None,
|| async {
@ -345,31 +385,55 @@ struct InvalidEntryCache(
);
impl ResponseCacheService for InvalidEntryCache {
fn config(&self) -> &ResponseCacheConfig {
self.0.config()
fn prepare(
self: Arc<Self>,
options: CacheOptions,
surface: &str,
input: &dyn Fn() -> Value,
) -> litellm_cache_response::PreparedCache {
litellm_cache_response::PreparedCache {
policy: options.policy,
max_entry_bytes: self.0.config().max_entry_bytes,
call: Arc::new(InvalidEntryCall {
request: options.request(&self.0.config().namespace, surface, input()),
service: self,
}),
}
}
}
struct InvalidEntryCall {
service: Arc<InvalidEntryCache>,
request: litellm_cache_response::ResponseCacheRequest,
}
impl litellm_cache_response::ResponseCacheCall for InvalidEntryCall {
fn lookup<'a>(
&'a self,
request: &'a litellm_cache_response::ResponseCacheRequest,
now: Duration,
) -> futures_util::future::BoxFuture<'a, Result<Option<Value>, litellm_cache::Error>> {
) -> futures_util::future::BoxFuture<
'a,
Result<Option<litellm_cache_response::CacheHit>, litellm_cache::Error>,
> {
Box::pin(async move {
Ok(self
.0
.async_lookup(request, now)
.await?
.or_else(|| Some(self.1.clone())))
Ok(Some(litellm_cache_response::CacheHit {
value: self
.service
.0
.async_lookup(&self.request, now)
.await?
.unwrap_or_else(|| self.service.1.clone()),
key: litellm_cache_response::cache_key(&self.request.key),
}))
})
}
fn store<'a>(
&'a self,
request: &'a litellm_cache_response::ResponseCacheRequest,
response: Value,
now: Duration,
) -> futures_util::future::BoxFuture<'a, Result<(), litellm_cache::Error>> {
Box::pin(self.0.async_store(request, response, now))
Box::pin(self.service.0.async_store(&self.request, response, now))
}
}
@ -482,6 +546,7 @@ impl Interceptors<RouteError> for CacheHitAccounting {
InterceptRequest::AfterProviderResponse { raw, reply } => {
reply.send(self.after_provider_response(raw).await?);
}
other => other.pass_through(),
}
Ok(())
}
@ -565,9 +630,11 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure(
};
let result = if streaming_route {
match execute_streaming::<TestRoute, _, _>(
cache_request(request),
Some(cache),
Some(CacheOptions::new(CacheScope::Shared)),
cache_request::<TestRoute>(
request,
Some(cache),
Some(CacheOptions::new(CacheScope::Shared)),
),
&accounting,
Some(&observer),
|| async { panic!("a cache hit must not call the provider") },
@ -579,9 +646,11 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure(
}
} else {
execute_unary::<UnaryTestRoute, _, _>(
cache_request(request),
Some(cache),
Some(CacheOptions::new(CacheScope::Shared)),
cache_request::<UnaryTestRoute>(
request,
Some(cache),
Some(CacheOptions::new(CacheScope::Shared)),
),
&accounting,
Some(&observer),
|| async { panic!("a cache hit must not call the provider") },
@ -627,9 +696,7 @@ async fn unary_call(
request: Value,
) -> Value {
execute_unary::<UnaryTestRoute, _, _>(
cache_request(request),
Some(cache.clone()),
options,
cache_request::<UnaryTestRoute>(request, Some(cache.clone()), options),
&(),
None,
|| async { Ok(json!({"call":calls.fetch_add(1, Ordering::SeqCst)})) },

View file

@ -118,6 +118,9 @@ impl From<litellm_host::machine::MachineFault> for Error {
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "OCR host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("OCR {message}"),
MachineFault::UnreadableAnswer => {
"OCR host answered with an unreadable response".into()
}
})
}
}

View file

@ -1,7 +1,7 @@
# Cache boundary
This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
This folder owns how Rust inference reaches the configured cache. Rust holds no cache rules: `litellm.rust_bridge.host.cache` decides whether a call may read or write, builds the key and stores the entry, using the same Python handler functions the Python route runs. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters
`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `python/service.rs` owns the `Cached<P>` protocol and typed cache operations. `python/binding.rs` owns `CachedBinding<H>`, which wraps a route's Python binding, asks the Python host module for access and serves cache operations. Route bindings never hold cache state
`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports
`python/` delegates operations to the Python cache the host module granted. `native/` owns the experimental V2 native cache handles exposed to Python. Inference does not use them until the Rust cache is ready

View file

@ -1,10 +1,8 @@
mod native;
mod python;
mod selection;
pub(crate) use native::NativeCacheHandle;
pub(crate) use python::{CacheCall, PythonCache};
pub(crate) use selection::{Cached, Selection, configure, configured_native};
pub(crate) use python::{Cached, CachedBinding, scoped as python_cache};
use litellm_cache::Error;
use pyo3::{

View file

@ -2,8 +2,8 @@
This directory exposes the experimental V2 native cache handles (`NativeCacheHandle`) to Python. Cache algorithms and storage protocols remain in their cache crates
Accept the cache object selected by the parent module. Do not read global `litellm.cache`, decide route admission, or select the Python cache adapter here
Inference does not select these handles. A Python `Cache` backed by one is reached through the Python cache delegation like any other backend
Keep Python-facing cache classes and method signatures stable when reorganizing modules. Native internals stay private to this directory unless the shared cache boundary or Python module registration needs them
Keep Python-facing cache classes and method signatures stable when reorganizing modules. Native internals stay private to this directory
Verify changes with the V2 cache tests using a freshly built extension. Test cache behavior, not module paths or file structure

View file

@ -12,9 +12,9 @@ use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec,
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
ResponseCacheConfig, ResponseCacheRequest,
};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use pyo3::{exceptions::PyValueError, prelude::*};
#[pyclass(
frozen,
@ -22,7 +22,6 @@ use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeCacheHandle {
service: Arc<dyn ResponseCacheService>,
backend: Arc<dyn ExactResponseCache>,
storage: Storage,
pid: u32,
@ -71,7 +70,6 @@ impl NativeCacheHandle {
},
));
Ok(Self {
service: backend.clone(),
backend,
storage: Storage::Memory(storage),
pid: std::process::id(),
@ -112,7 +110,6 @@ impl NativeCacheHandle {
},
));
Ok(Self {
service: backend.clone(),
backend,
storage: Storage::Redis(storage),
pid: std::process::id(),
@ -259,85 +256,6 @@ fn duration(seconds: f64) -> PyResult<Duration> {
.ok_or_else(|| PyValueError::new_err("cache durations must be finite and positive"))
}
pub(in crate::cache) fn native_handle<'py>(
configured: &Bound<'py, PyAny>,
) -> PyResult<Option<Bound<'py, PyAny>>> {
Ok(configured
.getattr_opt("cache")?
.map(|backend| backend.getattr_opt("native_handle"))
.transpose()?
.flatten()
.filter(|handle| handle.is_instance_of::<NativeCacheHandle>()))
}
pub(in crate::cache) fn configured(
configured: &Bound<'_, PyAny>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<(
Option<Arc<dyn ResponseCacheService>>,
litellm_cache_response::CacheOptions,
)> {
let handle = native_handle(configured)?.ok_or_else(|| {
pyo3::exceptions::PyRuntimeError::new_err(
"native inference needs a v2 cache, the catalog keeps Python caches on Python",
)
})?;
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;
cache.check_process()?;
let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none());
let controls = controls
.as_ref()
.map(|value| value.cast::<PyDict>())
.transpose()?;
if let Some(controls) = controls {
for name in controls.keys() {
let name = name.extract::<String>()?;
if !matches!(
name.as_str(),
"no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache"
) {
return Err(PyValueError::new_err(format!(
"unsupported v2 cache control: {name}"
)));
}
}
}
let boolean = |name: &str| -> PyResult<bool> {
controls
.map(|values| values.get_item(name))
.transpose()?
.flatten()
.map(|value| value.extract())
.transpose()
.map(|value| value.unwrap_or(false))
};
let seconds = |name: &str| -> PyResult<Option<Duration>> {
controls
.map(|values| values.get_item(name))
.transpose()?
.flatten()
.map(|value| duration(value.extract()?))
.transpose()
};
Ok((
Some(cache.service.clone()),
litellm_cache_response::CacheOptions {
policy: litellm_cache_response::CachePolicy {
caching: kwargs
.get_item("caching")?
.filter(|value| !value.is_none())
.map(|value| value.extract())
.transpose()?,
no_cache: boolean("no-cache")?,
no_store: boolean("no-store")?,
ttl: seconds("ttl")?,
max_age: seconds("s-max-age")?.or(seconds("s-maxage")?),
},
scope: litellm_cache_response::CacheScope::Shared,
},
))
}
fn now() -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)

View file

@ -1,9 +1,11 @@
# Python cache delegation
This directory lets Rust inference use a selected Python cache. `service.rs` implements the injected Rust response-cache service and yields typed cache operations. `host.rs` calls the Python cache's sync or async API and delivers the result back to Rust.
This directory lets Rust inference use the configured Python cache. `service.rs` owns the typed protocol and injected response-cache service. `binding.rs` wraps the route binding, asks `litellm.rust_bridge.host.cache` for access and delivers lookup and store results back to Rust. Keep PyO3 conversion, suspension and Python references in the binding, and cache rules in the Python host module
Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details
Every call into Python goes through that host module. Do not read `litellm.cache`, call cache methods directly, or re-derive which calls may read or write. Core remains independent of Python objects and cache implementation details
Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
Preparing the Python adapter must not invoke the lazy native identity builder or create a native cache request or key. Keep the host's hit key unchanged. Python owns key construction, lookup, storage and the stored entry format. Rust sends and receives a route-neutral `CachedOutput` (`Response` or `Stream`), and `litellm.rust_bridge.host.cache` maps it to the entry the legacy Python route stores for that call type, so the two replay each other's entries. Never name a route's entry shape in Rust. Rust stores a complete response after the route raised its provider-result point, so what the deployment success hook answered is what gets cached, and a cache hit never raises that point. A stream is stored once the caller has consumed it
Await asynchronous cache lookups through the existing host driver in the caller's task. Do not create another asyncio task or event loop for them. Asynchronous stores are handed to Python's `create_cache_write_task`, so the response returns before the write lands, as on the Python route. Cancellation during a lookup must prevent the provider request and the cache write. A failure to reach Python is a cache failure, never a request failure. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation

View file

@ -0,0 +1,338 @@
use super::service::{CacheCall, Cached};
use litellm_cache::Error;
use litellm_cache_response::{CacheHit, CachePolicy};
use litellm_host::protocol::{Protocol, Reply};
use litellm_host_python::{
InvokeError, PythonBinding, PythonHostCalls, PythonOwned, effective_py_args, from_py, to_py,
};
use pyo3::{
exceptions::{PyException, PyRuntimeError},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
const MODULE: &str = "litellm.rust_bridge.host.cache";
pub(crate) struct CachedBinding<H> {
inner: H,
cache: PythonCache,
call_type: &'static str,
positional: Option<Py<PyDict>>,
}
impl<H> CachedBinding<H> {
pub(crate) fn new(
inner: H,
call_type: &'static str,
asynchronous: bool,
positional: Py<PyDict>,
) -> Self {
Self {
inner,
cache: PythonCache::new(asynchronous),
call_type,
positional: Some(positional),
}
}
}
type ErrorOf<H> = <<H as PythonBinding>::Protocol as Protocol>::Error;
impl<H: PythonBinding> PythonBinding for CachedBinding<H> {
type Protocol = Cached<H::Protocol>;
type Failure = H::Failure;
fn decode_request(
&mut self,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> Result<<Self::Protocol as Protocol>::Request, InvokeError<ErrorOf<H>>> {
let request = self.inner.decode_request(py, arguments)?;
let policy = self
.positional
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("cache binding used after close"))
.and_then(|positional| effective_py_args(positional.bind(py), arguments))
.and_then(|keyed| self.cache.configure(py, &keyed, self.call_type))
.map_err(InvokeError::Python)?;
Ok((request, policy))
}
fn encode_response(
&mut self,
py: Python<'_>,
response: <Self::Protocol as Protocol>::Response,
) -> PyResult<Py<PyAny>> {
self.inner.encode_response(py, response)
}
fn public_response(
&mut self,
py: Python<'_>,
response: &serde_json::Value,
) -> PyResult<Py<PyAny>> {
self.inner.public_response(py, response)
}
fn response_value(
&mut self,
py: Python<'_>,
public: &Bound<'_, PyAny>,
) -> PyResult<serde_json::Value> {
self.inner.response_value(py, public)
}
fn encode_stream_head(
&mut self,
py: Python<'_>,
head: <Self::Protocol as Protocol>::StreamHead,
) -> PyResult<Py<PyAny>> {
self.inner.encode_stream_head(py, head)
}
fn encode_chunk(
&mut self,
py: Python<'_>,
chunk: <Self::Protocol as Protocol>::Chunk,
) -> PyResult<Py<PyAny>> {
self.inner.encode_chunk(py, chunk)
}
fn map_error(&self, py: Python<'_>, error: ErrorOf<H>) -> PyResult<H::Failure> {
self.inner.map_error(py, error)
}
fn host_error(error: &PyErr) -> ErrorOf<H> {
H::host_error(error)
}
}
impl<H: PythonBinding> PythonHostCalls<Cached<H::Protocol>> for CachedBinding<H> {
fn handle_host_call(
&mut self,
py: Python<'_>,
call: CacheCall,
) -> Result<(), InvokeError<ErrorOf<H>>> {
self.begin_host_call(py, call).map(|_| ())
}
fn begin_host_call(
&mut self,
py: Python<'_>,
call: CacheCall,
) -> Result<Option<Py<PyAny>>, InvokeError<ErrorOf<H>>> {
self.cache.begin(py, call).map_err(InvokeError::Python)
}
fn resume_host_call(
&mut self,
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> Result<Option<Py<PyAny>>, InvokeError<ErrorOf<H>>> {
self.cache.resume(py, result).map_err(InvokeError::Python)
}
}
impl<H: PythonOwned> PythonOwned for CachedBinding<H> {
fn close(&mut self, py: Python<'_>) {
self.inner.close(py);
self.cache.close();
self.positional = None;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.inner.traverse(visit)?;
self.cache.traverse(visit)?;
if let Some(positional) = &self.positional {
visit.call(positional)?;
}
Ok(())
}
}
enum Pending {
Lookup(Reply<Result<Option<CacheHit>, Error>>),
Store(Reply<Result<(), Error>>),
}
struct Granted {
access: Py<PyAny>,
}
fn cache_hit(hit: &Bound<'_, PyAny>) -> PyResult<Option<CacheHit>> {
if hit.is_none() {
return Ok(None);
}
Ok(Some(CacheHit {
value: from_py(&hit.getattr("response")?)?,
key: hit.getattr("key")?.extract()?,
}))
}
struct PythonCache {
module: &'static str,
granted: Option<Granted>,
pending: Option<Pending>,
asynchronous: bool,
}
impl PythonCache {
fn new(asynchronous: bool) -> Self {
Self {
module: MODULE,
granted: None,
pending: None,
asynchronous,
}
}
fn configure(
&mut self,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
call_type: &str,
) -> PyResult<Option<CachePolicy>> {
let access = py
.import(self.module)?
.getattr("access")?
.call1((arguments, call_type))?;
if access.is_none() {
self.granted = None;
return Ok(None);
}
let policy = CachePolicy {
no_cache: !access.getattr("read")?.extract::<bool>()?,
no_store: !access.getattr("write")?.extract::<bool>()?,
..CachePolicy::default()
};
self.granted = Some(Granted {
access: access.unbind(),
});
Ok(Some(policy))
}
fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult<Option<Py<PyAny>>> {
let granted = self.granted.as_ref().ok_or_else(|| {
PyRuntimeError::new_err("cache operation without granted cache access")
})?;
let access = granted.access.bind(py);
let function = |name: &str| py.import(self.module)?.getattr(name);
let (pending, result) = match call {
CacheCall::Lookup { reply } => (
Pending::Lookup(reply),
function(if self.asynchronous {
"alookup"
} else {
"lookup"
})
.and_then(|lookup| lookup.call1((access,))),
),
CacheCall::Store { value, reply } => (
Pending::Store(reply),
function(if self.asynchronous { "astore" } else { "store" })
.and_then(|store| store.call1((access, to_py(py, &value)?))),
),
};
let result = result.map(Bound::unbind);
self.pending = Some(pending);
if self.asynchronous && result.is_ok() {
return result.map(Some);
}
self.resume(py, result)
}
fn resume(
&mut self,
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<Option<Py<PyAny>>> {
if let Err(error) = &result
&& !error.is_instance_of::<PyException>(py)
{
self.pending = None;
return Err(result.err().unwrap());
}
match self.pending.take() {
Some(Pending::Lookup(reply)) => {
let hit = result
.map_err(|_| Error::Unavailable)
.and_then(|hit| cache_hit(hit.bind(py)).map_err(|_| Error::InvalidEntry));
reply.send(hit);
}
Some(Pending::Store(reply)) => {
reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable));
}
None => {
return Err(PyRuntimeError::new_err(
"cache reply without pending operation",
));
}
}
Ok(None)
}
fn close(&mut self) {
self.pending = None;
self.granted = None;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
let Some(granted) = &self.granted else {
return Ok(());
};
visit.call(&granted.access)
}
}
#[cfg(test)]
mod tests {
use futures_util::FutureExt;
use litellm_host::protocol::reply;
use rstest::rstest;
use serde_json::Value;
use super::*;
#[rstest]
fn a_cache_call_that_cannot_reach_python_answers_unavailable(
#[values(false, true)] asynchronous: bool,
#[values(false, true)] store: bool,
) {
Python::initialize();
Python::attach(|py| {
let mut cache = PythonCache {
module: "litellm_missing_cache_host",
granted: Some(Granted { access: py.None() }),
pending: None,
asynchronous,
};
let answer = if store {
let (reply, answer) = reply();
assert!(
cache
.begin(
py,
CacheCall::Store {
value: Value::Null,
reply
}
)
.unwrap()
.is_none()
);
answer.now_or_never().map(|answer| answer.unwrap().err())
} else {
let (reply, answer) = reply();
assert!(
cache
.begin(py, CacheCall::Lookup { reply })
.unwrap()
.is_none()
);
answer.now_or_never().map(|answer| answer.unwrap().err())
};
assert!(matches!(answer, Some(Some(Error::Unavailable))));
});
}
}

View file

@ -1,141 +0,0 @@
use litellm_cache::Error;
use litellm_host::protocol::Reply;
use litellm_host_python::{from_py, to_py};
use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use serde_json::Value;
use super::service::CacheCall;
enum Pending {
Lookup(Reply<Result<Option<Value>, Error>>),
Store(Reply<Result<(), Error>>),
}
pub(crate) struct PythonCache {
cache: Option<Py<PyAny>>,
arguments: Option<Py<PyDict>>,
pending: Option<Pending>,
asynchronous: bool,
}
impl PythonCache {
pub fn new(asynchronous: bool) -> Self {
Self {
cache: None,
arguments: None,
pending: None,
asynchronous,
}
}
pub(in crate::cache) fn bind(
&mut self,
cache: Bound<'_, PyAny>,
arguments: &Bound<'_, PyDict>,
) {
self.cache = Some(cache.unbind());
self.arguments = Some(arguments.clone().unbind());
}
pub fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult<Option<Py<PyAny>>> {
let Some(cache) = self.cache.as_ref() else {
return Err(pyo3::exceptions::PyRuntimeError::new_err(
"cache operation without configured cache",
));
};
let arguments = self
.arguments
.as_ref()
.ok_or_else(|| {
pyo3::exceptions::PyRuntimeError::new_err("cache arguments unavailable")
})?
.bind(py)
.copy()?;
let (method, result) = match call {
CacheCall::Lookup { reply } => {
self.pending = Some(Pending::Lookup(reply));
(
if self.asynchronous {
"async_get_cache"
} else {
"get_cache"
},
None,
)
}
CacheCall::Store { value, reply } => {
self.pending = Some(Pending::Store(reply));
(
if self.asynchronous {
"async_add_cache"
} else {
"add_cache"
},
Some(to_py(py, &value)?),
)
}
};
let result = match result {
Some(value) => cache
.bind(py)
.call_method(method, (value,), Some(&arguments)),
None => cache.bind(py).call_method(method, (), Some(&arguments)),
}
.map(Bound::unbind);
if self.asynchronous && result.is_ok() {
return result.map(Some);
}
self.resume(py, result)
}
pub fn resume(
&mut self,
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<Option<Py<PyAny>>> {
if let Err(error) = &result
&& !error.is_instance_of::<pyo3::exceptions::PyException>(py)
{
self.pending = None;
return Err(result.err().unwrap());
}
match self.pending.take() {
Some(Pending::Lookup(reply)) => {
let value = result.map_err(|_| Error::Unavailable).and_then(|value| {
if value.bind(py).is_none() {
Ok(None)
} else {
from_py(value.bind(py))
.map(Some)
.map_err(|_| Error::InvalidEntry)
}
});
reply.send(value);
}
Some(Pending::Store(reply)) => {
reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable));
}
None => {
return Err(pyo3::exceptions::PyRuntimeError::new_err(
"cache reply without pending operation",
));
}
}
Ok(None)
}
pub fn close(&mut self) {
self.pending = None;
self.cache = None;
self.arguments = None;
}
pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.cache)?;
visit.call(&self.arguments)
}
}

View file

@ -1,6 +1,5 @@
mod host;
mod binding;
mod service;
pub(crate) use host::PythonCache;
pub(crate) use service::CacheCall;
pub(super) use service::service;
pub(crate) use binding::CachedBinding;
pub(crate) use service::{Cached, scoped};

View file

@ -1,44 +1,31 @@
use std::{future::Future, pin::Pin, time::Duration};
use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Duration};
use litellm_cache::Error;
use litellm_cache_response::{
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope,
CacheHit, CacheOptions, CachePolicy, CacheScope, PreparedCache, ResponseCacheCall,
ResponseCacheService, ResponseEnvelope, ScopedCache,
};
use litellm_host::{
machine::{HostServices, MachineFault},
machine::HostServices,
protocol::{Protocol, Reply},
};
use litellm_inference::caching::CachedOutput;
use litellm_inference::caching::{Cachable, CachedOutput};
use serde_json::Value;
const STREAM_EVENTS_KEY: &str = "litellm_cached_anthropic_sse_events";
pub(crate) struct Cached<P>(PhantomData<P>);
fn from_python(value: Value) -> Result<Value, Error> {
let output = match value.get(STREAM_EVENTS_KEY) {
Some(events) => {
let events: Vec<String> =
serde_json::from_value(events.clone()).map_err(|_| Error::InvalidEntry)?;
CachedOutput::Stream(events.concat())
}
None => CachedOutput::Response(value),
};
serde_json::to_value(ResponseEnvelope::new("messages", output)).map_err(|_| Error::InvalidEntry)
}
fn to_python(value: Value) -> Result<Value, Error> {
let envelope: ResponseEnvelope<CachedOutput<Value>> =
serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?;
match envelope.decode("messages").ok_or(Error::InvalidEntry)? {
CachedOutput::Response(response) => Ok(response),
CachedOutput::Stream(text) => Ok(serde_json::json!({
STREAM_EVENTS_KEY: text.split_inclusive("\n\n").collect::<Vec<_>>()
})),
}
impl<P: Protocol> Protocol for Cached<P> {
type Request = (P::Request, Option<CachePolicy>);
type Response = P::Response;
type Error = P::Error;
type HostCall = CacheCall;
type Chunk = P::Chunk;
type StreamHead = P::StreamHead;
}
pub(crate) enum CacheCall {
Lookup {
reply: Reply<Result<Option<Value>, Error>>,
reply: Reply<Result<Option<CacheHit>, Error>>,
},
Store {
value: Value,
@ -46,58 +33,76 @@ pub(crate) enum CacheCall {
},
}
struct PythonCacheService<P: Protocol> {
services: HostServices<P>,
config: ResponseCacheConfig,
struct PythonCacheService<P: Cachable> {
services: HostServices<Cached<P>>,
protocol: PhantomData<fn() -> P>,
}
pub(in crate::cache) fn service<P: Protocol<HostCall = CacheCall>>(
services: HostServices<P>,
namespace: String,
) -> std::sync::Arc<dyn ResponseCacheService>
where
P::Error: From<MachineFault>,
{
std::sync::Arc::new(PythonCacheService {
services,
config: ResponseCacheConfig {
namespace,
..Default::default()
},
})
pub(crate) fn scoped<P: Cachable>(services: HostServices<Cached<P>>) -> ScopedCache {
ScopedCache::new(
Arc::new(PythonCacheService {
services,
protocol: PhantomData,
}),
CacheScope::Shared,
)
}
impl<P: Protocol<HostCall = CacheCall>> ResponseCacheService for PythonCacheService<P>
where
P::Error: From<MachineFault>,
{
fn config(&self) -> &ResponseCacheConfig {
&self.config
fn from_python<P: Cachable>(value: Value) -> Result<Value, Error> {
let output: CachedOutput<Value> =
serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?;
serde_json::to_value(ResponseEnvelope::new(P::SURFACE, output)).map_err(|_| Error::InvalidEntry)
}
fn to_python<P: Cachable>(value: Value) -> Result<Value, Error> {
let envelope: ResponseEnvelope<CachedOutput<Value>> =
serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?;
let output = envelope.decode(P::SURFACE).ok_or(Error::InvalidEntry)?;
serde_json::to_value(output).map_err(|_| Error::InvalidEntry)
}
impl<P: Cachable> ResponseCacheService for PythonCacheService<P> {
fn prepare(
self: Arc<Self>,
options: CacheOptions,
_: &str,
_: &dyn Fn() -> Value,
) -> PreparedCache {
PreparedCache {
policy: options.policy,
max_entry_bytes: usize::MAX,
call: self,
}
}
}
impl<P: Cachable> ResponseCacheCall for PythonCacheService<P> {
fn lookup<'a>(
&'a self,
_: &'a ResponseCacheRequest,
_: Duration,
) -> Pin<Box<dyn Future<Output = Result<Option<Value>, Error>> + Send + 'a>> {
) -> Pin<Box<dyn Future<Output = Result<Option<CacheHit>, Error>> + Send + 'a>> {
Box::pin(async move {
self.services
.call(|reply| CacheCall::Lookup { reply })
.await
.map_err(|_| Error::Unavailable)??
.map(from_python)
.map(|hit| {
Ok(CacheHit {
value: from_python::<P>(hit.value)?,
key: hit.key,
})
})
.transpose()
})
}
fn store<'a>(
&'a self,
_: &'a ResponseCacheRequest,
value: Value,
_: Duration,
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'a>> {
Box::pin(async move {
let value = to_python(value)?;
let value = to_python::<P>(value)?;
self.services
.call(|reply| CacheCall::Store { value, reply })
.await
@ -105,3 +110,122 @@ where
})
}
}
#[cfg(test)]
mod tests {
use std::convert::Infallible;
use litellm_host::{
interceptors::ProviderIdentity,
machine::{CallMachine, Machine, MachineStep},
protocol::{HostRequest, Protocol},
};
use litellm_inference::{
RouteError,
caching::{CacheRequest, execute_unary},
};
use rstest::rstest;
use serde_json::json;
use super::*;
struct TestRoute;
impl Protocol for TestRoute {
type Request = ();
type Response = Value;
type Error = RouteError;
type HostCall = Infallible;
type Chunk = ();
type StreamHead = ();
}
impl Cachable for TestRoute {
const SURFACE: &'static str = "test";
}
#[rstest]
#[case::hit(true)]
#[case::miss(false)]
#[tokio::test]
async fn python_cache_calls_never_materialize_native_identity(#[case] hit: bool) {
let mut machine = CallMachine::<Cached<TestRoute>>::new(None, |host| {
Box::pin(async move {
let cache = scoped::<TestRoute>(host.services);
let request = CacheRequest::prepare::<TestRoute>(
ProviderIdentity {
model: "test".into(),
provider: "test".into(),
},
Some(cache.service.clone()),
Some(cache.options(None)),
&|_| panic!("Python caching must not construct native identity"),
);
execute_unary::<TestRoute, _, _>(request, &(), None, || async {
Ok(json!({"source": "provider"}))
})
.await
})
});
let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Lookup { reply })) =
machine.resume().await.unwrap()
else {
panic!("expected Python lookup");
};
reply.send(Ok(hit.then(|| CacheHit {
value: json!({"kind": "Response", "value": {"source": "cache"}}),
key: "python-owned-key".into(),
})));
if !hit {
let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Store { value, reply })) =
machine.resume().await.unwrap()
else {
panic!("expected Python store on a miss");
};
assert_eq!(
value,
json!({"kind": "Response", "value": {"source": "provider"}})
);
reply.send(Ok(()));
}
let MachineStep::Complete(response) = machine.resume().await.unwrap() else {
panic!("expected completed call");
};
assert_eq!(
response,
json!({"source": if hit { "cache" } else { "provider" }})
);
}
#[rstest]
#[tokio::test]
async fn python_hit_keeps_the_host_key() {
let mut machine = CallMachine::<Cached<TestRoute>, CacheHit>::new(None, |host| {
Box::pin(async move {
let cache = scoped::<TestRoute>(host.services);
let options = cache.options(None);
let prepared = cache.service.prepare(options, TestRoute::SURFACE, &|| {
panic!("Python caching must not construct native identity")
});
Ok(prepared.call.lookup(Duration::ZERO).await.unwrap().unwrap())
})
});
let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Lookup { reply })) =
machine.resume().await.unwrap()
else {
panic!("expected Python lookup");
};
reply.send(Ok(Some(CacheHit {
value: json!({"kind": "Response", "value": {"answer": 42}}),
key: "python-owned-key".into(),
})));
let MachineStep::Complete(hit) = machine.resume().await.unwrap() else {
panic!("expected completed lookup");
};
assert_eq!(hit.key, "python-owned-key");
assert_eq!(
from_python::<TestRoute>(json!({"kind": "Response", "value": {"answer": 42}})).unwrap(),
hit.value
);
}
}

View file

@ -1,171 +0,0 @@
use super::{native, python};
use litellm_cache_response::{
CacheOptions, CachePolicy, CacheScope, ResponseCacheService, ScopedCache,
};
use litellm_host::{
machine::{HostServices, MachineFault},
protocol::Protocol,
};
use pyo3::{prelude::*, types::PyDict};
use std::sync::Arc;
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
pub(crate) struct Cached<P>(std::marker::PhantomData<P>);
impl<P: Protocol> Protocol for Cached<P> {
type Request = (P::Request, Selection);
type Response = P::Response;
type Error = P::Error;
type HostCall = python::CacheCall;
type Chunk = P::Chunk;
type StreamHead = P::StreamHead;
}
enum Backend {
Disabled,
Native(Arc<dyn ResponseCacheService>),
Python { namespace: String },
}
pub(crate) struct Selection {
backend: Backend,
options: CacheOptions,
}
impl Selection {
pub(crate) fn attach<P: Protocol<HostCall = python::CacheCall>>(
self,
services: HostServices<P>,
) -> (Option<ScopedCache>, CacheOptions)
where
P::Error: From<MachineFault>,
{
let service = match self.backend {
Backend::Disabled => None,
Backend::Native(service) => Some(service),
Backend::Python { namespace } => Some(python::service(services, namespace)),
};
(
service.map(|service| ScopedCache::new(service, CacheScope::Shared)),
self.options,
)
}
}
fn selected_cache<'py>(
py: Python<'py>,
kwargs: &Bound<'py, PyDict>,
call_type: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
let configured = py
.import(CACHE_HOST_MODULE)?
.getattr("configured_cache")?
.call0()?;
if configured.is_none()
|| kwargs
.get_item("caching")?
.is_some_and(|value| value.is(pyo3::types::PyBool::new(py, false)))
{
return Ok(None);
}
let supported = configured.getattr("supported_call_types")?;
if supported.is_none() || !supported.contains(call_type)? {
return Ok(None);
}
Ok(Some(configured))
}
pub(crate) fn configured_native(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
call_type: &str,
) -> PyResult<(
Option<Arc<dyn ResponseCacheService>>,
litellm_cache_response::CacheOptions,
)> {
let Some(configured) = selected_cache(py, kwargs, call_type)? else {
return Ok((
None,
litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
));
};
native_configuration(&configured, kwargs)
}
fn native_configuration(
configured: &Bound<'_, PyAny>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<(Option<Arc<dyn ResponseCacheService>>, CacheOptions)> {
if !configured
.call_method("should_use_cache", (), Some(kwargs))?
.extract::<bool>()?
{
return Ok((
None,
litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
));
}
native::v2::configured(configured, kwargs)
}
pub(crate) fn configure(
python: &mut python::PythonCache,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
call_type: &str,
) -> PyResult<Selection> {
let selected = selected_cache(py, arguments, call_type)?;
let Some(cache) = selected else {
return Ok(Selection {
backend: Backend::Disabled,
options: CacheOptions::new(CacheScope::Shared),
});
};
if native::v2::native_handle(&cache)?.is_some() {
let (native, options) = native_configuration(&cache, arguments)?;
return Ok(Selection {
backend: native.map_or(Backend::Disabled, Backend::Native),
options,
});
}
let enabled = cache
.call_method("should_use_cache", (), Some(arguments))?
.extract::<bool>()?;
let controls = arguments
.get_item("cache")?
.filter(|value| !value.is_none());
let boolean = |name: &str| -> PyResult<bool> {
controls
.as_ref()
.map(|value| value.cast::<PyDict>()?.get_item(name))
.transpose()?
.flatten()
.map(|value| value.extract())
.transpose()
.map(|value| value.unwrap_or(false))
};
let options = CacheOptions {
policy: CachePolicy {
no_cache: boolean("no-cache")?,
no_store: boolean("no-store")?,
..CachePolicy::default()
},
..CacheOptions::new(CacheScope::Shared)
};
let namespace = cache
.getattr_opt("namespace")?
.filter(|value| !value.is_none())
.map(|value| value.extract())
.transpose()?
.unwrap_or_default();
python.bind(cache, arguments);
Ok(Selection {
backend: if enabled {
Backend::Python { namespace }
} else {
Backend::Disabled
},
options,
})
}

View file

@ -126,6 +126,17 @@ pub(crate) fn public_response(
.map(Bound::unbind)
}
/// The JSON form of a public response a Python hook handed back, models nested in it included.
pub(crate) fn response_value(public: &Bound<'_, PyAny>) -> PyResult<serde_json::Value> {
from_py(
&public
.py()
.import("pydantic_core")?
.getattr("to_jsonable_python")?
.call1((public,))?,
)
}
struct RequestFieldSources<'py> {
body: Option<Bound<'py, PyAny>>,
credentials: Option<Bound<'py, PyAny>>,

View file

@ -59,6 +59,22 @@ impl PythonBinding for ChatCompletionsPythonHost {
self.0.response(py, &response)
}
fn public_response(
&mut self,
py: Python<'_>,
response: &serde_json::Value,
) -> PyResult<Py<PyAny>> {
self.0.response(py, response)
}
fn response_value(
&mut self,
_: Python<'_>,
public: &Bound<'_, PyAny>,
) -> PyResult<serde_json::Value> {
crate::marshal::response_value(public)
}
fn encode_stream_head(
&mut self,
_py: Python<'_>,

View file

@ -4,13 +4,21 @@ use std::sync::Arc;
use host::ChatCompletionsPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_cache_response::ScopedCache;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions};
use litellm_host::{
call::{CallOutput, OutputOf},
machine::ChannelInterceptors,
};
use litellm_inference::{CallOptions, RouteError};
use litellm_inference_chat::{
ChatCompletionsRoute, route::ChatCompletions, types::ChatCompletionsCall,
};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
use crate::cache::Cached;
use super::{
NativeCall,
inference::{InferenceHost, InferenceRoute, run_inference},
@ -61,11 +69,14 @@ impl InferenceRoute for ChatCompletionsRoute {
self.with_cache(cache)
}
fn machine(
async fn call(
self,
call: <ChatCompletions as Protocol>::Request,
policy: CachePolicy,
) -> HostedMachine<ChatCompletions> {
self.machine(call, policy)
request: ChatCompletionsCall,
interceptors: ChannelInterceptors<Cached<ChatCompletions>>,
options: CallOptions,
) -> Result<OutputOf<ChatCompletions>, RouteError> {
self.execute_call(request, &interceptors, options)
.await
.map(CallOutput::Complete)
}
}

View file

@ -1,13 +1,17 @@
use std::sync::Arc;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
use litellm_cache_response::ScopedCache;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_host_python::{PythonBinding, PythonHostCalls, from_py};
use litellm_host::{
call::{OutputOf, hosted_call},
machine::ChannelInterceptors,
protocol::Protocol,
};
use litellm_host_python::{PythonBinding, from_py};
use litellm_http::transport::Error as TransportError;
use litellm_inference::RouteError;
use litellm_inference::{CallOptions, RouteError, caching::Cachable};
use litellm_secrets::source::SecretSource;
use pyo3::{
exceptions::PyValueError,
@ -20,6 +24,7 @@ use serde_json::{Map, Value};
use super::{NativeCall, RequestView};
use crate::{
cache::{Cached, CachedBinding},
errors::{RustUpstreamError, route_error_to_pyerr},
marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds},
};
@ -143,8 +148,8 @@ impl InferenceHost {
}
}
pub(super) trait InferenceRoute: Sized + 'static {
type Protocol: Protocol<Error = RouteError>;
pub(super) trait InferenceRoute: Sized + Send + 'static {
type Protocol: Cachable;
const OPERATION: LoggingOperation;
const SYNC_CALL_TYPE: &'static str;
const ASYNC_CALL_TYPE: &'static str;
@ -155,11 +160,12 @@ pub(super) trait InferenceRoute: Sized + 'static {
secrets: Arc<dyn SecretSource>,
) -> Self;
fn with_cache(self, cache: ScopedCache) -> Self;
fn machine(
fn call(
self,
call: <Self::Protocol as Protocol>::Request,
policy: CachePolicy,
) -> HostedMachine<Self::Protocol>;
request: <Self::Protocol as Protocol>::Request,
interceptors: ChannelInterceptors<Cached<Self::Protocol>>,
options: CallOptions,
) -> impl Future<Output = Result<OutputOf<Self::Protocol>, RouteError>> + Send;
}
pub(super) fn run_inference<R, H>(
@ -170,13 +176,14 @@ pub(super) fn run_inference<R, H>(
) -> PyResult<Py<PyAny>>
where
R: InferenceRoute,
H: PythonBinding<Protocol = R::Protocol> + PythonHostCalls<R::Protocol> + 'static,
H: PythonBinding<Protocol = R::Protocol> + 'static,
{
let call_type = if asynchronous {
R::ASYNC_CALL_TYPE
} else {
R::SYNC_CALL_TYPE
};
let positional = call.positional()?.unbind();
let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?;
super::run_public_call(
py,
@ -188,14 +195,28 @@ where
crate::http::resources().auth.clone(),
crate::secrets::source(py)?,
);
let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?;
let route = match cache {
Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)),
None => route,
};
Ok(route.machine(request, cache_options.policy))
Ok(hosted_call(
request,
None,
move |(call, policy), services, interceptors, observers| async move {
let route = match policy {
Some(_) => route.with_cache(crate::cache::python_cache(services)),
None => route,
};
route
.call(
call,
interceptors,
CallOptions {
cache: policy,
observers,
},
)
.await
},
))
},
host,
CachedBinding::new(host, call_type, asynchronous, positional),
hooks,
asynchronous,
)

View file

@ -1,5 +1,4 @@
use crate::cache::{CacheCall, Cached, PythonCache, Selection};
use litellm_host_python::{PythonHostCalls, PythonOwned};
use litellm_host_python::PythonOwned;
use bytes::Bytes;
use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py};
@ -95,15 +94,11 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult<PyErr> {
/// public response, chunks and exceptions.
pub(super) struct MessagesPythonHost {
view: RequestView,
cache: PythonCache,
}
impl MessagesPythonHost {
pub(super) fn new(view: RequestView, asynchronous: bool) -> Self {
Self {
view,
cache: PythonCache::new(asynchronous),
}
pub(super) fn new(view: RequestView) -> Self {
Self { view }
}
fn projection(
@ -232,24 +227,20 @@ impl MessagesPythonHost {
}
impl PythonBinding for MessagesPythonHost {
type Protocol = Cached<Messages>;
type Protocol = Messages;
type Failure = PyErr;
fn decode_request(
&mut self,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> Result<(MessagesCall, Selection), InvokeError<Error>> {
) -> Result<MessagesCall, InvokeError<Error>> {
self.view
.resolve(py, arguments)
.map_err(InvokeError::Python)?;
let selection =
crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages")
.map_err(InvokeError::Python)?;
self.projection(py, arguments)
.map_err(|error| InvokeError::Python(self.map_failure(py, error)))?
.map_err(InvokeError::Native)
.map(|request| (request, selection))
}
fn encode_response(
@ -260,6 +251,22 @@ impl PythonBinding for MessagesPythonHost {
public_response(py, ROUTE_HOST_MODULE, response.as_ref())
}
fn public_response(
&mut self,
py: Python<'_>,
response: &serde_json::Value,
) -> PyResult<Py<PyAny>> {
public_response(py, ROUTE_HOST_MODULE, response)
}
fn response_value(
&mut self,
_: Python<'_>,
public: &Bound<'_, PyAny>,
) -> PyResult<serde_json::Value> {
crate::marshal::response_value(public)
}
fn encode_stream_head(
&mut self,
py: Python<'_>,
@ -289,43 +296,12 @@ impl PythonBinding for MessagesPythonHost {
}
}
impl PythonHostCalls<Cached<Messages>> for MessagesPythonHost {
fn handle_host_call(
&mut self,
py: Python<'_>,
op: CacheCall,
) -> Result<(), InvokeError<Error>> {
self.cache
.begin(py, op)
.map(|_| ())
.map_err(InvokeError::Python)
}
fn begin_host_call(
&mut self,
py: Python<'_>,
op: CacheCall,
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
self.cache.begin(py, op).map_err(InvokeError::Python)
}
fn resume_host_call(
&mut self,
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
self.cache.resume(py, result).map_err(InvokeError::Python)
}
}
impl PythonOwned for MessagesPythonHost {
fn close(&mut self, _: Python<'_>) {
self.view.close();
self.cache.close();
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.view.traverse(visit)?;
self.cache.traverse(visit)
self.view.traverse(visit)
}
}

View file

@ -1,53 +1,26 @@
mod host;
use std::sync::Arc;
use host::MessagesPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::ScopedCache;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::OutputOf, machine::ChannelInterceptors};
use litellm_inference::{CallOptions, RouteError};
use litellm_inference_messages::{MessagesCall, MessagesRoute, route::Messages};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
use super::NativeCall;
use super::{
NativeCall,
inference::{InferenceRoute, run_inference},
};
use crate::cache::Cached;
fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
let (arguments, hooks) =
crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?;
crate::routes::run_public_call(
py,
arguments,
move |py, arguments, request| {
let route = litellm_inference_messages::MessagesRoute::new(
crate::http::provider_client(py, arguments, asynchronous)?
.map_err(crate::http::client_error)?,
crate::http::resources().auth.clone(),
crate::secrets::source(py)?,
);
Ok(litellm_host::call::hosted_call(
request,
None,
move |(call, selection): (_, crate::cache::Selection),
services,
interceptors,
observers| async move {
let (cache, options) = selection.attach(services);
let route = match cache {
Some(cache) => route.with_cache(cache),
None => route,
};
route
.execute(
call,
&interceptors,
litellm_inference::CallOptions {
cache: Some(options.policy),
observers,
},
)
.await
},
))
},
MessagesPythonHost::new(call.view()?, asynchronous),
hooks,
asynchronous,
)
let host = MessagesPythonHost::new(call.view()?);
run_inference::<MessagesRoute, _>(py, call, asynchronous, host)
}
#[pyfunction]
@ -59,3 +32,31 @@ pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAn
pub(crate) fn amessages(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
run_messages(py, call, true)
}
impl InferenceRoute for MessagesRoute {
type Protocol = Messages;
const OPERATION: LoggingOperation = LoggingOperation::Messages;
const SYNC_CALL_TYPE: &'static str = "anthropic_messages";
const ASYNC_CALL_TYPE: &'static str = "anthropic_messages";
fn new(
http: litellm_http::Client,
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
) -> Self {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
self.with_cache(cache)
}
async fn call(
self,
request: MessagesCall,
interceptors: ChannelInterceptors<Cached<Messages>>,
options: CallOptions,
) -> Result<OutputOf<Messages>, RouteError> {
self.execute(request, &interceptors, options).await
}
}

View file

@ -41,6 +41,15 @@ impl<'py> NativeCall<'py> {
effective_py_args(&self.base, &self.kwargs)
}
/// The positional arguments by name, which the legacy cache folds into its key input.
fn positional(&self) -> PyResult<Bound<'py, PyDict>> {
let named = PyDict::new(self.base.py());
for (name, value) in self.base.iter().take(self.args.len()) {
named.set_item(name, value)?;
}
Ok(named)
}
/// The view a route host reads its arguments from.
fn view(&self) -> PyResult<RequestView> {
RequestView::new(&self.base, &self.kwargs)

View file

@ -118,6 +118,22 @@ impl PythonBinding for OcrPythonHost {
public_response(py, "litellm.rust_bridge.host.routes.ocr", &response)
}
fn public_response(
&mut self,
py: Python<'_>,
response: &serde_json::Value,
) -> PyResult<Py<PyAny>> {
public_response(py, "litellm.rust_bridge.host.routes.ocr", response)
}
fn response_value(
&mut self,
_: Python<'_>,
public: &Bound<'_, PyAny>,
) -> PyResult<serde_json::Value> {
crate::marshal::response_value(public)
}
fn encode_stream_head(
&mut self,
_: Python<'_>,

View file

@ -80,6 +80,22 @@ impl PythonBinding for ResponsesPythonHost {
self.0.response(py, &response)
}
fn public_response(
&mut self,
py: Python<'_>,
response: &serde_json::Value,
) -> PyResult<Py<PyAny>> {
self.0.response(py, response)
}
fn response_value(
&mut self,
_: Python<'_>,
public: &Bound<'_, PyAny>,
) -> PyResult<serde_json::Value> {
crate::marshal::response_value(public)
}
fn encode_stream_head(
&mut self,
_py: Python<'_>,

View file

@ -5,10 +5,11 @@ use std::sync::Arc;
use host::ResponsesPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_cache_response::ScopedCache;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_inference_responses::{ResponsesRoute, route::Responses};
use litellm_host::{call::OutputOf, machine::ChannelInterceptors};
use litellm_inference::{CallOptions, RouteError};
use litellm_inference_responses::{ResponsesRoute, route::Responses, types::ResponsesCall};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
pub(crate) use websocket::ResponsesWebSocketConnection;
@ -17,6 +18,7 @@ use super::{
NativeCall,
inference::{InferenceHost, InferenceRoute, run_inference},
};
use crate::cache::Cached;
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
@ -53,11 +55,12 @@ impl InferenceRoute for ResponsesRoute {
self.with_cache(cache)
}
fn machine(
async fn call(
self,
call: <Responses as Protocol>::Request,
policy: CachePolicy,
) -> HostedMachine<Responses> {
self.machine(call, policy)
request: ResponsesCall,
interceptors: ChannelInterceptors<Cached<Responses>>,
options: CallOptions,
) -> Result<OutputOf<Responses>, RouteError> {
self.execute(request, &interceptors, options).await
}
}

View file

@ -737,7 +737,7 @@ class Cache:
except Exception as e:
raise e
def add_cache(self, result, **kwargs):
def add_cache(self, result: object, **kwargs: object) -> None:
"""
Adds a result to the cache.
@ -764,7 +764,9 @@ class Cache:
return
verbose_logger.error("%s: %s", message, exc)
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
async def async_add_cache(
self, result: object, dynamic_cache_object: BaseCache | None = None, **kwargs: object
) -> None:
"""
Async implementation of add_cache
"""
@ -963,7 +965,7 @@ class Cache:
except Exception as e:
self._log_add_cache_failure(e)
def should_use_cache(self, **kwargs):
def should_use_cache(self, **kwargs: object) -> bool:
"""
Returns true if we should use the cache for LLM API calls

View file

@ -26,8 +26,8 @@ from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError
import litellm
from litellm._internal_context import is_internal_call, post_response_phase
from litellm._logging import print_verbose, verbose_logger
from litellm.caching import InMemoryCache
from litellm.caching.caching import S3Cache, response_cache_phase
from litellm.caching import DualCache, InMemoryCache
from litellm.caching.caching import Cache, RedisCache, S3Cache, response_cache_phase
from litellm.constants import CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS, RESPONSE_CACHE_EXCLUDED_PROVIDERS
from litellm.litellm_core_utils.hidden_params import get_hidden_params
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
@ -118,6 +118,46 @@ def _is_response_cache_excluded(model: str | None, kwargs: Mapping[str, object])
)
_CACHE_CONTROLS: Final = TypeAdapter(Mapping[str, object])
def _cache_controls(kwargs: Mapping[str, object]) -> Mapping[str, object]:
controls: Final = kwargs.get("cache")
return _CACHE_CONTROLS.validate_python(controls) if isinstance(controls, Mapping) else {}
def is_call_type_cached(call_type: str) -> bool:
if litellm.cache is None or litellm.cache.supported_call_types is None:
return False
covering_call_types: Final = ("aresponses", "responses") if call_type == "aresponses" else (call_type,)
return any(name in litellm.cache.supported_call_types for name in covering_call_types)
def _reads_requested(caching: object, controls: Mapping[str, object]) -> bool:
return ((caching is None and litellm.cache is not None) or caching is True) and controls.get("no-cache") is not True
def cache_reads_requested(kwargs: Mapping[str, object]) -> bool:
return _reads_requested(kwargs.get("caching"), _cache_controls(kwargs))
def cache_access(model: str | None, kwargs: Mapping[str, object], call_type: str) -> tuple[bool, bool]:
if _is_response_cache_excluded(model=model, kwargs=kwargs) or not is_call_type_cached(call_type):
return (False, False)
controls: Final = _cache_controls(kwargs)
return (_reads_requested(kwargs.get("caching"), controls), controls.get("no-store") is not True)
def may_store_cache(model: str | None, kwargs: Mapping[str, object], call_type: str) -> bool:
return cache_access(model, kwargs, call_type)[1]
def response_memory_tier(cache: Cache) -> DualCache | None:
if not isinstance(cache.cache, RedisCache):
return None
return DualCache(redis_cache=cache.cache, in_memory_cache=in_memory_cache_obj)
def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
cached_id: Final = cached_result.get("id")
if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"):
@ -222,9 +262,17 @@ def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "as
return task
def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
"""Read the caller-supplied ``cache_key`` off the request kwargs."""
return request_kwargs.get("cache_key", None)
def cache_key_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]:
return {
name: value
for name, value in kwargs.items()
if name != "cache_key" and not (name == "metadata" and value is None)
}
def response_cache_key(cache: Cache, kwargs: Mapping[str, object]) -> object:
preset: Final = kwargs.get("cache_key")
return preset if preset else cache.get_cache_key(**cache_key_kwargs(kwargs))
def _set_cached_hidden_param(response: object, key: str, value: object) -> None:
@ -274,21 +322,13 @@ class LLMCachingHandler:
request_kwargs: dict[str, object],
start_time: datetime.datetime,
):
from litellm.caching import DualCache, RedisCache
self.async_streaming_chunks: list[object] = []
self.sync_streaming_chunks: list[object] = []
self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs)
self.preset_cache_key: str | None = None
self.preset_cache_key: object = None
self.original_function = original_function
self.start_time = start_time
if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):
self.dual_cache: DualCache | None = DualCache(
redis_cache=litellm.cache.cache,
in_memory_cache=in_memory_cache_obj,
)
else:
self.dual_cache = None
self.dual_cache: DualCache | None = None if litellm.cache is None else response_memory_tier(litellm.cache)
async def async_get_cache(
self,
@ -325,11 +365,7 @@ class LLMCachingHandler:
return None
# Check if caching should be performed BEFORE doing expensive operations
if (
(kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) is True
) and (
kwargs.get("cache", {}).get("no-cache", False) is not True
): # allow users to control returning cached responses from the completion function
if cache_reads_requested(kwargs):
args = args or ()
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False
@ -908,12 +944,10 @@ class LLMCachingHandler:
if all(result is None for result in cached_result):
cached_result = None
else:
request_kwargs: Final = new_kwargs.copy()
request_cache_key: Final = _request_cache_key(request_kwargs)
request_kwargs.pop("cache_key", None)
request_kwargs: Final = cache_key_kwargs(new_kwargs)
self.preset_cache_key = response_cache_key(litellm.cache, new_kwargs)
if litellm.cache.supports_async() is True:
## check if dual cache is supported ##
self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
with response_cache_phase("get"):
cached_result = await litellm.cache.async_get_cache(
dynamic_cache_object=self.dual_cache,
@ -921,7 +955,6 @@ class LLMCachingHandler:
**request_kwargs,
)
else: # fallback for caches that don't support async
self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
with response_cache_phase("get"):
cached_result = litellm.cache.get_cache(
dynamic_cache_object=self.dual_cache,
@ -1239,20 +1272,8 @@ class LLMCachingHandler:
Returns:
bool: True if the result should be stored in the cache, False otherwise.
"""
model_value: Final = kwargs.get("model")
model: Final = model_value if isinstance(model_value, str) else None
cache_options_value: Final = kwargs.get("cache")
cache_options: Final[Mapping[str, object]] = (
TypeAdapter(Mapping[str, object]).validate_python(cache_options_value)
if isinstance(cache_options_value, Mapping)
else {}
)
no_store: Final = cache_options.get("no-store", False)
return (
not _is_response_cache_excluded(model=model, kwargs=kwargs)
and self._is_call_type_supported_by_cache(original_function=original_function)
and no_store is not True
)
model: Final = kwargs.get("model")
return may_store_cache(model if isinstance(model, str) else None, kwargs, str(original_function.__name__))
_should_store_result_in_cache = should_store_result_in_cache
@ -1290,11 +1311,7 @@ class LLMCachingHandler:
Returns:
bool: True if the call type is supported by the cache, False otherwise.
"""
if litellm.cache is None or litellm.cache.supported_call_types is None:
return False
call_type: Final = str(original_function.__name__)
covering_call_types: Final = ("aresponses", "responses") if call_type == "aresponses" else (call_type,)
return any(name in litellm.cache.supported_call_types for name in covering_call_types)
return is_call_type_cached(str(original_function.__name__))
async def add_streaming_response_to_cache(self, processed_chunk: ModelResponse) -> None:
"""

View file

@ -70,7 +70,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType(
**detectors.COMMON,
Feature.MOCK_RESPONSE: detectors.chat_mock_response,
Feature.AGENTIC_LOOP_HOOK: detectors.chat_agentic_loop_hook,
Feature.RESPONSE_CACHE: detectors.response_cache,
Feature.STREAMING: detectors.streaming,
Feature.DROP_PARAMS: detectors.drop_params,
Feature.MODIFY_PARAMS: detectors.modify_params,

View file

@ -103,7 +103,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType(
Feature.SAFEGUARDS_CLASSIFIER: lambda request: native_route_classifies(
str(request.resolved["model"]), request.resolved.get("safeguards")
),
Feature.RESPONSE_CACHE: detectors.response_cache,
Feature.MODIFY_PARAMS: detectors.modify_params,
Feature.WORKLOAD_IDENTITY: detectors.anthropic_workload_identity,
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,

View file

@ -77,7 +77,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType(
Feature.INPUT_REWRITE: detectors.input_rewrite,
Feature.EXTRA_BODY: detectors.extra_body,
Feature.PROVIDER_AFFINITY: detectors.provider_affinity,
Feature.RESPONSE_CACHE: detectors.response_cache,
Feature.MANAGED_RESOURCE_IDS: detectors.managed_resource_ids,
Feature.CACHE_CONTROL_MARKS: detectors.cache_control_marks,
Feature.ANTHROPIC_TOOL_SEARCH: detectors.anthropic_tool_search,

View file

@ -74,8 +74,6 @@ class Feature(str, Enum):
"""A registered MessagesInterceptor that takes over the call."""
AGENTIC_LOOP_HOOK = "agentic_loop_hook"
"""A logger overriding async_should_run_agentic_loop, which may run follow-up calls."""
RESPONSE_CACHE = "response_cache"
"""litellm.cache, whose hit and store order differs from native."""
TLS_POLICY = "tls_policy"
"""An OpenSSL security level or ECDH curve that native TLS cannot apply."""
MANAGED_RESOURCE_IDS = "managed_resource_ids"
@ -151,7 +149,6 @@ COMMON_PYTHON_FEATURES: Final = frozenset(
)
CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | {
Feature.RESPONSE_CACHE,
Feature.STREAMING,
Feature.DROP_PARAMS,
Feature.MODIFY_PARAMS,
@ -169,7 +166,6 @@ MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | {
Feature.PRE_REQUEST_HOOK,
Feature.WEBSEARCH_INTERCEPTION,
Feature.MESSAGES_INTERCEPTOR,
Feature.RESPONSE_CACHE,
Feature.MODIFY_PARAMS,
Feature.WORKLOAD_IDENTITY,
Feature.CACHE_CONTROL_INJECTION,
@ -190,7 +186,6 @@ RESPONSES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | {
Feature.INPUT_REWRITE,
Feature.EXTRA_BODY,
Feature.PROVIDER_AFFINITY,
Feature.RESPONSE_CACHE,
Feature.MANAGED_RESOURCE_IDS,
Feature.CACHE_CONTROL_MARKS,
Feature.ANTHROPIC_TOOL_SEARCH,
@ -202,7 +197,6 @@ RESPONSES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | {
}
RULES: Final[Rules] = (
RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY, features=CHAT_PYTHON_FEATURES),
RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY),
RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY),
RouteRule(Route.OCR, Rollout.RUST_REQUIRED),

View file

@ -443,6 +443,10 @@ def _request_success_callbacks(request: NativeCall) -> tuple[object, ...]:
return (*supplied, *(dynamic or ()))
def deployment_success_hook() -> bool:
return _overridden("async_post_call_success_deployment_hook", litellm.callbacks)
def pre_request_hook(_request: NativeCall) -> bool:
return _overridden("async_pre_request_hook", litellm.callbacks)
@ -471,11 +475,6 @@ def websearch_interception(_request: NativeCall) -> bool:
return any(isinstance(callback, WebSearchInterceptionLogger) for callback in litellm.callbacks)
def response_cache(request: NativeCall) -> bool:
cache: Final = optional_mapping(request.resolved.get("cache")) or {}
return litellm.cache is not None and cache.get("no-cache") is not True
_NATIVE_ECDH_CURVES: Final = frozenset({"x25519", "prime256v1", "secp256r1", "p-256", "secp384r1", "p-384"})

View file

@ -2,10 +2,10 @@
Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (whether a call may read or write the configured cache, and the lookup and store calls), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
- A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here
- Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/cache/python/binding.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green
- Inputs and outputs use values and frozen snapshots. Retain caller-owned clients and caches as opaque references when identity is required; Rust never discovers them through module globals
- Moving a question into Rust means replacing the function body with a native call, not adding a second resolver. The Python signature stays
- `routes/<route>.py` holds the questions one route asks, under the same function names on every route (`response`, `map_failure`, `connection_defaults`), so Rust names only the module. A question two routes ask goes in a top-level module here instead

View file

@ -1,12 +1,229 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import re
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from litellm.caching.caching import Cache
from litellm.caching.dual_cache import DualCache
def configured_cache() -> Cache | None:
class ResponseOutput(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["Response"] = "Response"
value: object
class StreamOutput(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["Stream"] = "Stream"
value: str
CachedOutput: TypeAlias = ResponseOutput | StreamOutput
_OUTPUT: Final = TypeAdapter[CachedOutput](Annotated[CachedOutput, Field(discriminator="kind")])
_ENTRY: Final = TypeAdapter(dict[str, object])
_EVENTS: Final = TypeAdapter(list[str])
_SSE_EVENT_BOUNDARY: Final = re.compile(r"(?<=\n\n)")
@dataclass(frozen=True, slots=True)
class EntryFormat:
entry: Callable[[CachedOutput], object | None]
output: Callable[[object, bool], CachedOutput | None]
tiers_responses: bool
tiers_streams: bool
def tiered(self, output: CachedOutput) -> bool:
match output:
case ResponseOutput():
return self.tiers_responses
case StreamOutput():
return self.tiers_streams
def _response_entry(output: CachedOutput) -> object | None:
match output:
case ResponseOutput(value=response):
return response
case StreamOutput():
return None
def _response_output(entry: object, stream: bool) -> CachedOutput | None:
return None if stream else ResponseOutput(value=entry)
def _messages_entry(output: CachedOutput) -> object:
from litellm.types.caching import CACHED_STREAM_EVENTS_KEY
match output:
case ResponseOutput(value=response):
return response
case StreamOutput(value=text):
return {CACHED_STREAM_EVENTS_KEY: [event for event in _SSE_EVENT_BOUNDARY.split(text) if event]}
def _stream_events(entry: object) -> list[str] | None:
from litellm.types.caching import CACHED_STREAM_EVENTS_KEY
if not isinstance(entry, dict):
return None
events: Final = _ENTRY.validate_python(entry).get(CACHED_STREAM_EVENTS_KEY)
return None if events is None else _EVENTS.validate_python(events)
def _messages_output(entry: object, stream: bool) -> CachedOutput:
events: Final = _stream_events(entry)
if events is None:
return ResponseOutput(value=entry)
return StreamOutput(value="".join(events))
RESPONSE_ENTRY: Final = EntryFormat(
entry=_response_entry, output=_response_output, tiers_responses=True, tiers_streams=False
)
MESSAGES_ENTRY: Final = EntryFormat(
entry=_messages_entry, output=_messages_output, tiers_responses=False, tiers_streams=True
)
def entry_format(call_type: CallTypes) -> EntryFormat:
match call_type:
case CallTypes.anthropic_messages:
return MESSAGES_ENTRY
case _:
return RESPONSE_ENTRY
@dataclass(frozen=True, slots=True)
class CacheAccess:
cache: Cache
request: Mapping[str, object]
key: object
read: bool
write: bool
memory_tier: DualCache | None
format: EntryFormat
def _allowed(kwargs: Mapping[str, object], call_type: CallTypes) -> tuple[bool, bool]:
from litellm._logging import verbose_logger
from litellm.caching.caching_handler import cache_access
model: Final = kwargs.get("model")
try:
return cache_access(model if isinstance(model, str) else None, kwargs, call_type)
except Exception: # noqa: BLE001 # a cache failure is a miss, never a request failure
verbose_logger.exception("LiteLLM Cache: deciding cache access failed")
return (False, False)
def _key(cache: Cache, kwargs: Mapping[str, object]) -> object:
from litellm._logging import verbose_logger
from litellm.caching.caching_handler import cache_key_kwargs, response_cache_key
try:
if cache.should_use_cache(**cache_key_kwargs(kwargs)) is not True:
return None
return response_cache_key(cache, kwargs)
except Exception: # noqa: BLE001 # a cache failure is a miss, never a request failure
verbose_logger.exception("LiteLLM Cache: preparing the cache key failed")
return None
def grant(cache: Cache, kwargs: Mapping[str, object], call_type: CallTypes) -> CacheAccess | None:
from litellm.caching.caching_handler import cache_key_kwargs, response_memory_tier
read, write = _allowed(kwargs, call_type)
if not (read or write):
return None
key: Final = _key(cache, kwargs)
if key is None:
return None
return CacheAccess(
cache=cache,
request=cache_key_kwargs(kwargs),
key=key,
read=read,
write=write,
memory_tier=response_memory_tier(cache),
format=entry_format(call_type),
)
def access(kwargs: Mapping[str, object], call_type: str) -> CacheAccess | None:
import litellm
return litellm.cache
cache: Final = litellm.cache
return None if cache is None else grant(cache, kwargs, CallTypes(call_type))
@dataclass(frozen=True, slots=True)
class CacheHit:
response: Mapping[str, object]
key: str
def _hit(access: CacheAccess, entry: object) -> CacheHit | None:
output: Final = access.format.output(entry, access.request.get("stream") is True)
return None if output is None else CacheHit(response=output.model_dump(), key=str(access.key))
def lookup(access: CacheAccess) -> CacheHit | None:
from litellm.caching.caching_handler import is_response_without_output
result: Final = access.cache.get_cache(dynamic_cache_object=None, cache_key=access.key, **access.request) # pyright: ignore[reportUnknownMemberType] # Cache.get_cache kwargs are untyped
if result is None or is_response_without_output(result):
return None
return _hit(access, result)
async def alookup(access: CacheAccess) -> CacheHit | None:
from litellm.caching.caching_handler import is_response_without_output
result: Final = await access.cache.async_get_cache( # pyright: ignore[reportUnknownMemberType] # Cache.async_get_cache kwargs are untyped
dynamic_cache_object=access.memory_tier, cache_key=access.key, **access.request
)
if result is None:
return None
if is_response_without_output(result):
if access.memory_tier is not None and isinstance(access.key, str):
access.memory_tier.in_memory_cache.delete_cache(access.key)
return None
return _hit(access, result)
def _entry(access: CacheAccess, value: object) -> tuple[object, bool] | None:
from litellm.caching.caching_handler import is_response_without_output
output: Final = _OUTPUT.validate_python(value)
entry: Final = access.format.entry(output)
if entry is None or is_response_without_output(entry):
return None
return entry, access.format.tiered(output)
def store(access: CacheAccess, value: object) -> None:
prepared: Final = _entry(access, value)
if prepared is not None:
access.cache.add_cache(prepared[0], cache_key=access.key, **access.request)
async def astore(access: CacheAccess, value: object) -> None:
from litellm.caching.caching_handler import create_cache_write_task
prepared: Final = _entry(access, value)
if prepared is None:
return
entry, tiered = prepared
tier: Final = access.memory_tier if tiered else None
create_cache_write_task(
lambda: access.cache.async_add_cache(entry, dynamic_cache_object=tier, cache_key=access.key, **access.request)
)

View file

@ -322,6 +322,12 @@ def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitab
return hook(kwargs, call_type)
def reviews_deployment_success() -> bool:
from litellm.rust_bridge.detectors import deployment_success_hook
return deployment_success_hook()
def after_deployment_success(kwargs: dict[str, object], response: object, call_type: str) -> Awaitable[object]:
from litellm import utils
from litellm.types.utils import CallTypes

View file

@ -1 +0,0 @@

View file

@ -1,19 +0,0 @@
import threading
from collections.abc import Generator
from typing import Final
import fakeredis
import pytest
@pytest.fixture
def redis_url() -> Generator[str]:
server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis")
worker: Final = threading.Thread(target=server.serve_forever, daemon=True)
worker.start()
try:
yield f"redis://127.0.0.1:{server.server_address[1]}"
finally:
server.shutdown()
server.server_close()
worker.join(timeout=5)

View file

@ -1,121 +0,0 @@
from __future__ import annotations
from collections.abc import AsyncIterator, Mapping
from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel, TypeAdapter
import litellm
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
from litellm.rust_bridge.chat_completions.entrypoints import NATIVE_ACOMPLETION
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES
from litellm.rust_bridge.public_call import NativeCall
from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES
from litellm.types.utils import ModelResponse
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
def payload(value: object) -> object:
if isinstance(value, ModelResponse):
return value.model_dump_json(exclude=MappingProxyType({"id": True, "created": True}))
if isinstance(value, dict):
fields: Final = TypeAdapter(dict[str, object]).validate_python(value)
return {name: field for name, field in fields.items() if name != "_hidden_params"}
return value.model_dump_json() if isinstance(value, BaseModel) else value
def cache_key(response: object) -> object:
hidden: Final = get_hidden_params_dict(response)
headers: Final = TypeAdapter(dict[str, object]).validate_python(hidden.get("additional_headers", {}))
return headers.get("x-litellm-cache-key")
async def invoke(
route: Literal["chat", "messages", "responses"],
server: RecordingServer,
options: Mapping[str, object],
native: bool = True,
) -> object:
common: Final = {"api_key": "test-key", "api_base": server.base_url, **options}
if route == "responses":
server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
arguments: Final = {"model": RESPONSES_MODEL, "input": "hello", **common}
if not native:
return await litellm.aresponses(**arguments)
request: Final = NativeCall(
args=(),
kwargs=arguments,
base={
"model": RESPONSES_MODEL,
"input": "hello",
"stream": None,
"api_key": "test-key",
"api_base": server.base_url,
"custom_llm_provider": "openai",
"extra_headers": None,
**arguments,
},
)
return await runtime.arun(
RouteContext(Route.RESPONSES),
binding=NATIVE_ARESPONSES,
native=lambda hook: hook(request),
python=runtime.NO_PYTHON,
rules=(RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED),),
)
server.default_response = (
ResponseSpec(body=None, events=MESSAGES_EVENTS)
if options.get("stream")
else ResponseSpec(body=MESSAGES_RESPONSE)
)
parameters: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32, **common}
if route == "chat":
if not native:
return await litellm.acompletion(**parameters)
chat: Final = NativeCall(args=(), kwargs=parameters, base={})
return await runtime.arun(
RouteContext(Route.CHAT_COMPLETIONS),
binding=NATIVE_ACOMPLETION,
native=lambda hook: hook(chat),
python=runtime.NO_PYTHON,
rules=(RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),),
)
if not native:
return await litellm.anthropic_messages(**parameters)
messages: Final = NativeCall(
args=(),
kwargs=parameters,
base={
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"max_tokens": 32,
"stream": None,
"api_key": "test-key",
"api_base": server.base_url,
"custom_llm_provider": "anthropic",
**parameters,
},
)
return await runtime.arun(
RouteContext(Route.MESSAGES),
binding=NATIVE_AMESSAGES,
native=lambda hook: hook(messages),
python=runtime.NO_PYTHON,
rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),),
)
async def collect(stream: object) -> bytes:
assert isinstance(stream, AsyncIterator)
return b"".join([chunk_bytes(chunk) async for chunk in stream])
def chunk_bytes(value: object) -> bytes:
assert isinstance(value, bytes)
return value

View file

@ -1,534 +0,0 @@
import asyncio
from typing import Final, Literal
import pytest
from pydantic import TypeAdapter
import litellm
from litellm import _v2
from litellm.caching.caching import Cache, CacheMode
from litellm.caching.caching_handler import (
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
)
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
from litellm.proxy.hooks.model_max_budget_limiter import (
PROXY_VirtualKeyModelMaxBudgetLimiter,
model_budget_spend_cache_key,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.utils import InternalUsageCache
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.public_call import NativeCall
from litellm.types.caching import CachingSupportedCallTypes
from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.routing import native_route
pytestmark = pytest.mark.requires_rust_extension
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "stream", "legacy"),
(
("chat", False, False),
("messages", False, False),
("responses", False, False),
("messages", True, False),
("messages", False, True),
("messages", True, True),
),
)
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage(
recording_server: RecordingServer,
route: Literal["chat", "messages", "responses"],
stream: bool,
native: bool,
monkeypatch: pytest.MonkeyPatch,
legacy: bool,
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
litellm.cache = Cache() if legacy else _v2.Cache.memory()
counters: Final = litellm.DualCache()
budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters)
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
InternalUsageCache(counters), model_group_resolver=lambda model: model
)
recorder: Final = RecordingLogger()
key_hash: Final = "a" * 64
metadata: Final = {
"user_api_key": key_hash,
"model_group": "cached-model",
"user_api_key_model_max_budget": {"cached-model": {"max_budget": 1, "budget_duration": "1h"}},
}
options: Final = {
"callbacks": [budget, limiter, recorder],
"metadata": metadata,
"stream": stream,
}
spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h")
token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens")
first: Final = await invoke(route, recording_server, options, native=native)
if stream:
await collect(first)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await drain_logging()
first_events: Final = await recorder.wait_for_async("async_log_success_event")
first_log: Final = TypeAdapter(dict[str, object]).validate_python(first_events[0].kwargs)
first_payload: Final = TypeAdapter(dict[str, object]).validate_python(first_log["standard_logging_object"])
expected_cost: Final = TypeAdapter(float).validate_python(first_log["response_cost"])
usage: Final = RESPONSES_RESPONSE["usage"] if route == "responses" else MESSAGES_RESPONSE["usage"]
expected_tokens: Final = usage["input_tokens"] + usage["output_tokens"]
assert expected_cost > 0
assert counters.get_cache(spend_key) == pytest.approx(expected_cost)
assert counters.get_cache(token_key) == first_payload["total_tokens"] == expected_tokens
second: Final = await invoke(route, recording_server, options, native=native)
if stream:
await collect(second)
await drain_logging()
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs)
cached_payload: Final = TypeAdapter(dict[str, object]).validate_python(cached_log["standard_logging_object"])
assert len(recording_server.requests) == 1
assert len(successes) == 2
assert cached_log["cache_hit"] is True
assert cached_log["response_cost"] == cached_payload["response_cost"] == 0
assert cached_payload["cache_hit"] is True
assert cached_payload["id"] != first_payload["id"]
assert cached_payload["custom_llm_provider"] == first_payload["custom_llm_provider"]
assert cached_payload["custom_llm_provider"] == ("openai" if route == "responses" else "anthropic"), {
"miss_provider": first_log.get("custom_llm_provider"),
"hit_provider": cached_log.get("custom_llm_provider"),
}
assert cached_payload["total_tokens"] == first_payload["total_tokens"]
assert counters.get_cache(spend_key) == pytest.approx(expected_cost)
assert counters.get_cache(token_key) == 2 * expected_tokens
@pytest.mark.asyncio
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
@pytest.mark.parametrize("backend", ("disabled", "memory", "redis"))
async def test_response_cache_backend_does_not_control_coordination(
recording_server: RecordingServer,
native: bool,
monkeypatch: pytest.MonkeyPatch,
backend: Literal["disabled", "memory", "redis"],
redis_url: str,
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
litellm.cache = (
None
if backend == "disabled"
else _v2.Cache.memory()
if backend == "memory"
else _v2.Cache.redis(redis_url, namespace="independent-coordination")
)
recording_server.expected_requests = 2 if backend == "disabled" else 1
counters: Final = litellm.DualCache()
budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters)
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
InternalUsageCache(counters), model_group_resolver=lambda model: model
)
key_hash: Final = "b" * 64
identity: Final = UserAPIKeyAuth(api_key=key_hash, rpm_limit=2, tpm_limit=1000, max_parallel_requests=1)
spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h")
request_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "requests")
token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens")
parallel_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "max_parallel_requests")
expected_tokens: Final = MESSAGES_RESPONSE["usage"]["input_tokens"] + MESSAGES_RESPONSE["usage"]["output_tokens"]
recorder: Final = RecordingLogger()
async def request(call_id: str, successes: int) -> object:
data: Final = {
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"litellm_call_id": call_id,
"max_tokens": 32,
"metadata": {
"user_api_key": key_hash,
"model_group": "cached-model",
"user_api_key_model_max_budget": {"cached-model": {"max_budget": 1, "budget_duration": "1h"}},
},
}
await limiter.async_pre_call_hook(identity, counters, data, "acompletion")
assert len(TypeAdapter(dict[str, float]).validate_python(counters.get_cache(parallel_key))) == 1
response: Final = await invoke(
"chat", recording_server, {**data, "callbacks": [budget, limiter, recorder]}, native=native
)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await recorder.wait_for_async("async_log_success_event", count=successes)
return response
await asyncio.create_task(request("cache-miss", 1))
assert counters.get_cache(parallel_key) == {}
assert counters.get_cache(request_key) == 1
assert counters.get_cache(token_key) == expected_tokens
first_events: Final = await recorder.wait_for_async("async_log_success_event")
first_cost: Final = TypeAdapter(float).validate_python(first_events[0].kwargs["response_cost"])
assert first_cost > 0
assert counters.get_cache(spend_key) == pytest.approx(first_cost)
await asyncio.create_task(request("cache-hit", 2))
expected_spend: Final = first_cost * recording_server.expected_requests
assert counters.get_cache(spend_key) == pytest.approx(expected_spend)
assert len(recording_server.requests) == recording_server.expected_requests
assert counters.get_cache(parallel_key) == {}
assert counters.get_cache(request_key) == 2
assert counters.get_cache(token_key) == 2 * expected_tokens
with pytest.raises(litellm.RateLimitError):
await asyncio.create_task(request("over-rpm-limit", 3))
assert len(recording_server.requests) == recording_server.expected_requests
assert counters.get_cache(parallel_key) == {}
assert counters.get_cache(token_key) == 2 * expected_tokens
assert counters.get_cache(spend_key) == pytest.approx(expected_spend)
if litellm.cache is not None:
await litellm.cache.disconnect()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
)
async def test_cache_controls_and_backend_credential_key_semantics(
recording_server: RecordingServer,
route: Literal["chat", "messages", "responses"],
legacy: bool,
) -> None:
recording_server.expected_requests = 3 if legacy else 4
litellm.cache = Cache() if legacy else _v2.Cache.memory()
await invoke(route, recording_server, {"cache": {"no-store": True}})
await invoke(route, recording_server, {})
await invoke(route, recording_server, {})
assert len(recording_server.requests) == 2
await invoke(route, recording_server, {"cache": {"no-cache": True}})
await invoke(route, recording_server, {"api_key": "another-key"})
assert len(recording_server.requests) == recording_server.expected_requests
@pytest.mark.asyncio
@pytest.mark.parametrize("legacy", (False, True))
async def test_v2_messages_replays_a_completed_stream(recording_server: RecordingServer, legacy: bool) -> None:
recording_server.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS)
litellm.cache = Cache() if legacy else _v2.Cache.memory()
recorder: Final = RecordingLogger()
parameters: Final = {
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"max_tokens": 32,
"api_key": "test-key",
"api_base": recording_server.base_url,
"stream": True,
"callbacks": [recorder],
}
with native_route(Route.MESSAGES):
first_stream: Final = await litellm.anthropic_messages(**parameters)
assert cache_key(first_stream) is None
first: Final = await collect(first_stream)
await recorder.wait_for_async("async_log_success_event")
second_stream: Final = await litellm.anthropic_messages(**parameters)
assert isinstance(cache_key(second_stream), str)
assert cache_key(second_stream) == get_hidden_params_dict(second_stream)["cache_key"]
second: Final = await collect(second_stream)
assert payload(first) == payload(second)
assert first == b"".join(recording_server.default_response.payloads())
assert len(recording_server.requests) == 1
await drain_logging()
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs)
assert cached_log["cache_hit"] is True
assert cached_log["response_cost"] == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("memory", "redis"))
async def test_rust_messages_uses_a_legacy_cache_without_python_inference(
recording_server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
backend: Literal["memory", "redis"],
redis_url: str,
) -> None:
from litellm.caching.caching import Cache
monkeypatch.setenv("LITELLM_RUST", "1")
litellm.cache = Cache() if backend == "memory" else Cache(type="redis", url=redis_url, namespace="rust-host")
logger: Final = RecordingLogger()
litellm.callbacks = [logger]
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
parameters: Final = {
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"max_tokens": 32,
"api_key": "test-key",
"api_base": recording_server.base_url,
}
first: Final = await invoke("messages", recording_server, parameters)
second: Final = await invoke("messages", recording_server, parameters)
assert cache_key(second)
assert cache_key(first) is None
assert payload(first) == payload(second)
assert len(recording_server.requests) == 1
await logger.wait_for_async("async_log_success_event", count=2)
assert logger.names.count("async_log_success_event") == 2
assert "log_failure_event" not in logger.names
assert "async_log_failure_event" not in logger.names
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ("chat", "messages", "responses"))
@pytest.mark.parametrize("native", (False, True))
@pytest.mark.parametrize("excluded", (None, [], ["embedding"]))
async def test_v2_cache_honors_supported_call_types_for_reads_and_writes(
recording_server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
route: Literal["chat", "messages", "responses"],
native: bool,
excluded: list[CachingSupportedCallTypes] | None,
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
litellm.cache = _v2.Cache.memory()
call_type: Final[CachingSupportedCallTypes] = (
"acompletion" if route == "chat" else "anthropic_messages" if route == "messages" else "aresponses"
)
recording_server.expected_requests = 4
litellm.cache.supported_call_types = excluded
await invoke(route, recording_server, {}, native=native)
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 2
litellm.cache.supported_call_types = [call_type]
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 3
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 3
litellm.cache.supported_call_types = excluded
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 4
@pytest.mark.asyncio
@pytest.mark.parametrize("native", (False, True))
@pytest.mark.parametrize(
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
)
async def test_v2_default_off_requires_opt_in_even_for_existing_entries(
recording_server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
route: Literal["chat", "messages", "responses"],
native: bool,
legacy: bool,
) -> None:
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
litellm.cache = Cache() if legacy else _v2.Cache.memory()
litellm.cache.mode = CacheMode.default_off
recording_server.expected_requests = 4
await invoke(route, recording_server, {}, native=native)
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 2
await invoke(route, recording_server, {"cache": {"use-cache": True}}, native=native)
assert len(recording_server.requests) == 3
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await invoke(route, recording_server, {"cache": {"use-cache": True}}, native=native)
assert len(recording_server.requests) == 3
await invoke(route, recording_server, {}, native=native)
assert len(recording_server.requests) == 4
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
)
async def test_cache_lookup_uses_backend_request_callback_semantics(
recording_server: RecordingServer,
route: Literal["chat", "messages", "responses"],
legacy: bool,
) -> None:
class Rewrite(RecordingLogger):
temperature = 0.1
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
request_body(kwargs)["temperature"] = self.temperature
super().log_pre_api_call(model, messages, kwargs)
logger: Final = Rewrite()
litellm.cache = Cache() if legacy else _v2.Cache.memory()
recording_server.expected_requests = 1 if legacy else 2
await invoke(route, recording_server, {"callbacks": [logger]})
first_hit: Final = await invoke(route, recording_server, {"callbacks": [logger]})
logger.temperature = 0.8
await invoke(route, recording_server, {"callbacks": [logger]})
second_hit: Final = await invoke(route, recording_server, {"callbacks": [logger]})
assert logger.names.count("log_pre_api_call") == 4
assert len(recording_server.requests) == recording_server.expected_requests
assert recording_server.requests[0].body["temperature"] == 0.1
if not legacy:
assert recording_server.requests[1].body["temperature"] == 0.8
assert isinstance(cache_key(first_hit), str)
assert isinstance(cache_key(second_hit), str)
if not legacy:
assert cache_key(first_hit) != cache_key(second_hit)
@pytest.mark.asyncio
@pytest.mark.parametrize("cancel_lookup", (False, True))
async def test_python_cache_operations_stay_in_the_rust_callers_task(
recording_server: RecordingServer,
cancel_lookup: bool,
) -> None:
from litellm.caching.base_cache import BaseCache
from litellm.caching.in_memory_cache import InMemoryCache
caller: Final = asyncio.current_task()
entered: Final = asyncio.Event()
release: Final = asyncio.Event()
storage: Final = InMemoryCache()
class CallerCache(BaseCache):
async def async_set_cache_pipeline(
self, cache_list: list[tuple[str, object]], ttl: float | None = None
) -> None:
await storage.async_set_cache_pipeline(cache_list, ttl=ttl)
async def async_get_cache(self, key: str, **kwargs: object) -> object:
if cancel_lookup:
entered.set()
await release.wait()
else:
assert asyncio.current_task() is caller
return storage.get_cache(key, **kwargs)
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
assert asyncio.current_task() is caller
await asyncio.sleep(0)
storage.set_cache(key, value, **kwargs)
litellm.cache = Cache(_backend=CallerCache())
if cancel_lookup:
recording_server.expected_requests = 0
task: Final = asyncio.create_task(invoke("messages", recording_server, {}))
await asyncio.wait_for(entered.wait(), timeout=5)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
release.set()
await asyncio.sleep(0)
assert len(recording_server.requests) == 0
assert storage.cache_dict == {}
return
first: Final = await invoke("messages", recording_server, {})
second: Final = await invoke("messages", recording_server, {})
assert payload(first) == payload(second)
assert cache_key(second)
assert len(recording_server.requests) == 1
def test_sync_rust_messages_calls_python_cache(recording_server: RecordingServer) -> None:
from litellm.rust_bridge.messages.entrypoints import NATIVE_MESSAGES
litellm.cache = Cache()
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
arguments: Final = {
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"max_tokens": 32,
"api_key": "test-key",
"api_base": recording_server.base_url,
}
request: Final = NativeCall(
args=(),
kwargs=arguments,
base={
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"max_tokens": 32,
"stream": None,
"api_key": "test-key",
"api_base": recording_server.base_url,
"custom_llm_provider": "anthropic",
**arguments,
},
)
def call() -> object:
return runtime.run(
RouteContext(Route.MESSAGES),
binding=NATIVE_MESSAGES,
native=lambda hook: hook(request),
python=runtime.NO_PYTHON,
rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),),
)
first: Final = call()
second: Final = call()
assert payload(first) == payload(second)
assert cache_key(second)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("namespace_source", ("cache", "metadata"))
async def test_rust_messages_legacy_cache_honors_request_namespaces(
recording_server: RecordingServer, namespace_source: str
) -> None:
litellm.cache = Cache()
recording_server.expected_requests = 2
first_options: Final = (
{"cache": {"namespace": "first"}} if namespace_source == "cache" else {"metadata": {"redis_namespace": "first"}}
)
second_options: Final = (
{"cache": {"namespace": "second"}}
if namespace_source == "cache"
else {"metadata": {"redis_namespace": "second"}}
)
await invoke("messages", recording_server, first_options)
second: Final = await invoke("messages", recording_server, second_options)
first_hit: Final = await invoke("messages", recording_server, first_options)
second_hit: Final = await invoke("messages", recording_server, second_options)
assert cache_key(second) is None
assert cache_key(first_hit) is not None
assert cache_key(second_hit) is not None
assert len(recording_server.requests) == 2
@pytest.mark.asyncio
async def test_rust_messages_legacy_semantic_cache_preserves_python_scope(
recording_server: RecordingServer,
) -> None:
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.types.caching import LiteLLMCacheType
litellm.cache = Cache(type=LiteLLMCacheType.REDIS_SEMANTIC, _backend=InMemoryCache())
first_options: Final = {"messages": [{"role": "user", "content": "hello"}]}
second_options: Final = {"messages": [{"role": "user", "content": "hi"}]}
first: Final = await invoke("messages", recording_server, first_options)
second: Final = await invoke("messages", recording_server, second_options)
assert cache_key(first) is None
assert cache_key(second) is not None
assert payload(first) == payload(second)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("rust_first", (False, True), ids=("python_to_rust", "rust_to_python"))
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
async def test_legacy_cache_keeps_public_messages_responses_compatible(
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, rust_first: bool, stream: bool
) -> None:
litellm.cache = Cache()
monkeypatch.setenv("LITELLM_RUST", "0")
options: Final = {"litellm_params": {"preset_cache_key": "shared-messages"}, "stream": stream}
first: Final = await invoke("messages", recording_server, options, native=rust_first)
first_payload: Final = await collect(first) if stream else payload(first)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
second: Final = await invoke("messages", recording_server, options, native=not rust_first)
second_payload: Final = await collect(second) if stream else payload(second)
assert second_payload == first_payload
assert len(recording_server.requests) == 1

View file

@ -1,207 +0,0 @@
import asyncio
from collections.abc import AsyncIterator
from typing import Final, Literal
import pytest
from pydantic import TypeAdapter
import litellm
from litellm import _v2
from litellm._v2.cache import NativeBackend
from litellm.caching.caching_handler import (
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
)
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.types.utils import ModelResponse
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE, RESPONSES_TEXT
from tests.test_litellm_rust.cache.support import cache_key, invoke, payload
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
pytestmark = pytest.mark.requires_rust_extension
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ("chat", "messages", "responses"))
@pytest.mark.parametrize("backend", ("memory", "redis"))
async def test_v2_cache_skips_provider_and_reports_one_success_per_call(
recording_server: RecordingServer,
route: Literal["chat", "messages", "responses"],
backend: Literal["memory", "redis"],
redis_url: str,
) -> None:
recording_server.expected_requests = 2
litellm.cache = _v2.Cache.memory() if backend == "memory" else _v2.Cache.redis(redis_url, namespace="headers")
recorder: Final = RecordingLogger()
first: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
await recorder.wait_for_async("async_log_success_event")
second: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
assert payload(first) == payload(second)
assert cache_key(first) is None
key: Final = cache_key(second)
assert isinstance(key, str)
assert key == get_hidden_params_dict(second)["cache_key"]
assert len(recording_server.requests) == 1
await drain_logging()
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
assert len(successes) == 2
cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs)
assert cached_log["cache_hit"] is True
assert cached_log["response_cost"] == 0
await litellm.cache.delete_cache_keys([key])
refreshed: Final = await invoke(route, recording_server, {"callbacks": [recorder]})
assert cache_key(refreshed) is None
assert len(recording_server.requests) == 2
assert len(await recorder.wait_for_async("async_log_success_event", count=3)) == 3
await litellm.cache.disconnect()
@pytest.mark.asyncio
async def test_v2_global_cache_leaves_legacy_only_calls_usable() -> None:
litellm.cache = _v2.Cache.memory()
response: Final = await litellm.aembedding(
model="openai/cache-test-embedding",
input=["hello"],
api_key="test-key",
mock_response=[0.25, 0.75],
)
assert response.model_dump(include={"data"}) == {
"data": [{"embedding": [0.25, 0.75], "index": 0, "object": "embedding"}]
}
@pytest.mark.parametrize("route", ("chat", "responses"))
def test_v2_cache_works_through_python_inference(
recording_server: RecordingServer, route: Literal["chat", "messages", "responses"], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
litellm.cache = _v2.Cache.memory()
common: Final = {"api_key": "test-key", "api_base": recording_server.base_url}
if route == "responses":
recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
parameters: Final = {"model": RESPONSES_MODEL, "input": "hello", **common}
first: Final = litellm.responses(**parameters)
second: Final = litellm.responses(**parameters)
assert payload(first) == payload(second)
else:
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
arguments: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32, **common}
initial: Final = litellm.completion(**arguments)
cached: Final = litellm.completion(**arguments)
assert isinstance(initial, ModelResponse) and isinstance(cached, ModelResponse)
assert (
initial.choices[0].message.content
== cached.choices[0].message.content
== MESSAGES_RESPONSE["content"][0]["text"]
)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
async def test_v2_facade_and_backend_share_storage_and_management() -> None:
cache: Final = _v2.Cache.memory()
await cache.async_add_cache({"answer": 7}, cache_key="shared")
assert cache.get_cache(cache_key="shared") == {"answer": 7}
assert await cache.ping() is True
await cache.delete_cache_keys(["shared"])
assert await cache.async_get_cache(cache_key="shared") is None
cache.add_cache({"answer": 8}, cache_key="flush")
backend: Final = cache.cache
assert isinstance(backend, NativeBackend)
backend.flush_cache()
assert cache.get_cache(cache_key="flush") is None
await cache.disconnect()
@pytest.mark.asyncio
@pytest.mark.parametrize("control", ("s-maxage", "s-max-age"))
async def test_v2_native_cache_accepts_existing_freshness_aliases(
recording_server: RecordingServer, control: str
) -> None:
litellm.cache = _v2.Cache.memory()
first: Final = await invoke("responses", recording_server, {})
second: Final = await invoke("responses", recording_server, {"cache": {control: 600}})
assert payload(first) == payload(second)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
async def test_v2_cache_does_not_force_native_responses_streaming(
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
litellm.cache = _v2.Cache.memory()
recording_server.default_response = ResponseSpec(
body=None,
events=(
("response.created", {"type": "response.created", "sequence_number": 0, "response": RESPONSES_RESPONSE}),
(
"response.completed",
{"type": "response.completed", "sequence_number": 1, "response": RESPONSES_RESPONSE},
),
),
)
response: Final = await litellm.aresponses(
model=RESPONSES_MODEL,
input="hello",
stream=True,
caching=False,
api_key="test-key",
api_base=recording_server.base_url,
)
assert isinstance(response, AsyncIterator)
chunks: Final = [chunk async for chunk in response]
assert chunks[-1].type == "response.completed"
assert chunks[-1].response.output[0].content[0].text == RESPONSES_TEXT
@pytest.mark.asyncio
async def test_v2_cache_works_through_python_messages(
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_RUST", "0")
litellm.cache = _v2.Cache.memory()
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
parameters: Final = {
"model": MESSAGES_MODEL,
"messages": list(MESSAGES),
"max_tokens": 32,
"api_key": "test-key",
"api_base": recording_server.base_url,
}
first: Final = await litellm.anthropic_messages(**parameters)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
second: Final = await litellm.anthropic_messages(**parameters)
assert payload(first) == payload(second)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True))
async def test_v2_redis_flush_only_removes_its_namespace(
redis_url: str, recording_server: RecordingServer, asynchronous: bool
) -> None:
own: Final = _v2.Cache.redis(redis_url, namespace="flush-own")
other: Final = _v2.Cache.redis(redis_url, namespace="flush-other")
litellm.cache = own
recording_server.expected_requests = 2
await own.async_add_cache({"answer": "own"}, cache_key="shared")
await other.async_add_cache({"answer": "other"}, cache_key="shared")
await invoke("responses", recording_server, {})
hit: Final = await invoke("responses", recording_server, {})
assert isinstance(cache_key(hit), str)
assert await own.async_get_cache(cache_key="shared") == {"answer": "own"}
backend: Final = own.cache
assert isinstance(backend, NativeBackend)
if asynchronous:
await backend.async_flush_cache()
else:
backend.flush_cache()
assert await own.async_get_cache(cache_key="shared") is None
assert await other.async_get_cache(cache_key="shared") == {"answer": "other"}
refreshed: Final = await invoke("responses", recording_server, {})
assert cache_key(refreshed) is None
assert len(recording_server.requests) == 2
await own.disconnect()
await other.disconnect()

View file

@ -1,16 +1,19 @@
from collections.abc import Generator
from functools import partial
import pytest
from litellm.rust_bridge.catalog import Route
from tests.test_litellm_rust.chat_completions.data import call
from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE
from tests.test_litellm_rust.support.caching import Call
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.routing import native_route
@pytest.fixture
def server(recording_server: RecordingServer) -> RecordingServer:
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
recording_server.serve(ResponseSpec(body=MESSAGES_RESPONSE))
return recording_server
@ -18,3 +21,8 @@ def server(recording_server: RecordingServer) -> RecordingServer:
def native() -> Generator[None]:
with native_route(Route.CHAT_COMPLETIONS):
yield
@pytest.fixture
def native_call(server: RecordingServer) -> Call:
return partial(call, server)

View file

@ -0,0 +1,328 @@
import asyncio
from collections.abc import Mapping, Sequence
from typing import Final
import pytest
from pydantic import TypeAdapter
from litellm.caching.caching import Cache, CacheMode
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.types.caching import CachingSupportedCallTypes
from litellm.types.utils import ModelResponse
from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE, MESSAGES_TEXT
from tests.test_litellm_rust.support.caching import (
Call,
CallerCache,
FailingLookupCache,
FakeRedis,
HeldWriteCache,
assert_cache_hit_keeps_budget_spend_and_counts_usage,
settle_cache_writes,
use_cache,
)
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
from tests.test_litellm_rust.support.recording_server import RecordingServer
from tests.test_litellm_rust.support.results import cache_key, stored_cache_key
from tests.test_litellm_rust.support.routing import assert_served_natively
pytestmark = pytest.mark.requires_rust_extension
_OBJECT: Final = TypeAdapter(dict[str, object])
EMPTY: Final = {"object": "chat.completion", "choices": []}
OWN_CALL_TYPES: Final = ("acompletion",)
def content(response: object) -> object:
assert isinstance(response, ModelResponse)
return response.model_dump()["choices"][0]["message"]["content"]
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.usefixtures("cache")
async def test_a_repeated_native_call_is_answered_from_the_python_cache(
call: Call, server: RecordingServer, asynchronous: bool
) -> None:
first: Final = await call(asynchronous=asynchronous)
second: Final = await call(asynchronous=asynchronous)
assert_served_natively(server, first)
assert cache_key(first) is None
assert isinstance(cache_key(second), str)
assert content(second) == content(first) == MESSAGES_TEXT
@pytest.mark.parametrize(
("first", "second"),
(
({"cache": {"no-store": True}}, {}),
({}, {"cache": {"no-cache": True}}),
({}, {"caching": False}),
({}, {"messages": [{"role": "user", "content": "Goodbye"}]}),
({}, {"temperature": 0.5}),
({"cache": {"namespace": "first"}}, {"cache": {"namespace": "second"}}),
({"metadata": {"redis_namespace": "first"}}, {"metadata": {"redis_namespace": "second"}}),
),
ids=(
"no-store",
"no-cache",
"caching-off-read",
"other-messages",
"other-param",
"cache-namespace",
"metadata-namespace",
),
)
@pytest.mark.usefixtures("cache")
async def test_cache_controls_and_keys_keep_the_second_call_on_the_provider(
call: Call, server: RecordingServer, first: Mapping[str, object], second: Mapping[str, object]
) -> None:
server.expect(2)
await call(**first)
response: Final = await call(**second)
assert_served_natively(server, count=2)
assert cache_key(response) is None
@pytest.mark.usefixtures("cache")
async def test_caching_off_skips_the_read_but_still_writes(call: Call, server: RecordingServer) -> None:
await call(caching=False)
hit: Final = await call()
assert len(server.requests) == 1
assert content(hit) == MESSAGES_TEXT
assert_served_natively(server, count=1)
@pytest.mark.usefixtures("cache")
async def test_each_namespace_keeps_its_own_entry(call: Call, server: RecordingServer) -> None:
server.expect(2)
await call(cache={"namespace": "first"})
await call(cache={"namespace": "second"})
hit: Final = await call(cache={"namespace": "first"})
assert len(server.requests) == 2
assert isinstance(cache_key(hit), str)
assert_served_natively(server, count=2)
@pytest.mark.usefixtures("cache")
async def test_the_key_is_built_from_the_call_arguments_not_the_wire_body(call: Call, server: RecordingServer) -> None:
class Rewrite(RecordingLogger):
temperature = 0.1
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
request_body(kwargs)["temperature"] = self.temperature
super().log_pre_api_call(model, messages, kwargs)
logger: Final = Rewrite()
await call(callbacks=[logger])
logger.temperature = 0.8
hit: Final = await call(callbacks=[logger])
assert len(server.requests) == 1
assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.1
assert isinstance(cache_key(hit), str)
assert_served_natively(server, count=1)
@pytest.mark.parametrize(
("supported", "expected_requests"), ((["embedding"], 2), (OWN_CALL_TYPES, 1)), ids=("other-type", "own-type")
)
async def test_supported_call_types_gate_native_reads_and_writes(
call: Call,
server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
supported: Sequence[CachingSupportedCallTypes],
expected_requests: int,
) -> None:
use_cache(monkeypatch, Cache(supported_call_types=list(supported)))
server.expect(expected_requests)
await call()
second: Final = await call()
assert len(server.requests) == expected_requests
assert isinstance(cache_key(second), str) == (expected_requests == 1)
assert_served_natively(server, count=expected_requests)
@pytest.mark.parametrize(
("options", "expected_requests"), (({}, 2), ({"cache": {"use-cache": True}}, 1)), ids=("plain", "opt-in")
)
async def test_default_off_mode_caches_only_calls_that_opt_in(
call: Call,
server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
options: Mapping[str, object],
expected_requests: int,
) -> None:
use_cache(monkeypatch, Cache(mode=CacheMode.default_off))
server.expect(expected_requests)
await call(**options)
second: Final = await call(**options)
assert len(server.requests) == expected_requests
assert isinstance(cache_key(second), str) == (expected_requests == 1)
assert_served_natively(server, count=expected_requests)
@pytest.mark.usefixtures("cache")
async def test_a_native_hit_is_logged_as_a_free_cache_hit(call: Call, server: RecordingServer) -> None:
recorder: Final = RecordingLogger()
await call(callbacks=[recorder])
await call(callbacks=[recorder])
await drain_logging()
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
miss, hit = (_OBJECT.validate_python(event.kwargs) for event in successes)
assert len(server.requests) == 1
assert miss["cache_hit"] is not True
assert TypeAdapter(float).validate_python(miss["response_cost"]) > 0
assert hit["cache_hit"] is True
assert hit["response_cost"] == 0
assert_served_natively(server, count=1)
async def test_python_cache_lookups_run_in_the_callers_task(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = CallerCache()
use_cache(monkeypatch, Cache(_backend=backend))
await call()
hit: Final = await call()
assert len(server.requests) == 1
assert isinstance(cache_key(hit), str)
assert backend.lookup_tasks == [asyncio.current_task(), asyncio.current_task()]
assert_served_natively(server, count=1)
async def test_cancelling_a_pending_lookup_sends_and_stores_nothing(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
entered: Final = asyncio.Event()
release: Final = asyncio.Event()
backend: Final = CallerCache(entered, release)
use_cache(monkeypatch, Cache(_backend=backend))
server.expect(0)
task: Final = asyncio.create_task(call())
await asyncio.wait_for(entered.wait(), timeout=5)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
release.set()
await asyncio.sleep(0)
assert server.requests == []
assert backend.storage.cache_dict == {}
assert_served_natively(server, count=0)
@pytest.mark.usefixtures("cache")
async def test_a_cache_hit_keeps_budget_spend_and_counts_usage(call: Call, server: RecordingServer) -> None:
usage: Final = TypeAdapter(dict[str, int]).validate_python(MESSAGES_RESPONSE["usage"])
await assert_cache_hit_keeps_budget_spend_and_counts_usage(
call, server, usage["input_tokens"] + usage["output_tokens"]
)
assert_served_natively(server, count=1)
async def test_a_stored_response_without_output_is_a_miss(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = InMemoryCache()
cache: Final = use_cache(monkeypatch, Cache(_backend=backend))
await call()
(key,) = backend.cache_dict
cache.add_cache(EMPTY, cache_key=key)
server.expect(2)
response: Final = await call()
assert len(server.requests) == 2
assert content(response) == MESSAGES_TEXT
assert_served_natively(server, count=2)
async def test_a_failing_cache_lookup_runs_the_call_uncached(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
use_cache(monkeypatch, FailingLookupCache())
server.expect(2)
first: Final = await call()
await call()
assert content(first) == MESSAGES_TEXT
assert len(server.requests) == 2
assert_served_natively(server, count=2)
async def test_the_response_returns_before_the_cache_write_lands(
call: Call, native_call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = HeldWriteCache()
use_cache(monkeypatch, Cache(_backend=backend))
first: Final = await asyncio.wait_for(native_call(), timeout=5)
backend.release.set()
await settle_cache_writes()
hit: Final = await call()
assert content(first) == MESSAGES_TEXT
assert len(server.requests) == 1
assert content(hit) == MESSAGES_TEXT
assert_served_natively(server, count=1)
async def test_a_redis_hit_fills_the_worker_memory_tier(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
) -> None:
use_cache(monkeypatch, fake_redis.cache())
await call()
await call()
fake_redis.client.flushall()
hit: Final = await call()
assert len(server.requests) == 1
assert content(hit) == MESSAGES_TEXT
assert_served_natively(server, count=1)
async def test_a_stored_response_lands_in_the_worker_memory_tier(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
) -> None:
use_cache(monkeypatch, fake_redis.cache())
first: Final = await call()
fake_redis.client.flushall()
hit: Final = await call()
assert len(server.requests) == 1
assert content(hit) == content(first)
assert_served_natively(server, count=1)
async def test_a_hit_reports_the_key_the_entry_is_stored_under(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = InMemoryCache()
use_cache(monkeypatch, Cache(_backend=backend))
await call()
hit: Final = await call()
(stored,) = backend.cache_dict
assert len(server.requests) == 1
assert stored_cache_key(hit) == stored
assert_served_natively(server, count=1)

View file

@ -4,7 +4,6 @@ from typing import Final, cast # noqa: TID251 # narrows legacy callable signat
import pytest
import litellm
from litellm.caching.caching import Cache
from litellm.chat_completions import dispatch
from litellm.chat_completions.dispatch import (
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
@ -247,7 +246,6 @@ class _ChatAgentic(CustomLogger):
(Feature.TIMEOUT_VARIANTS, {"request_timeout": 5}, {}),
(Feature.AGENTIC_LOOP_HOOK, {}, {"callbacks": [_ChatAgentic()]}),
(Feature.TLS_POLICY, {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}),
(Feature.RESPONSE_CACHE, {}, {"cache": Cache()}),
(Feature.STREAMING, {"stream": True}, {}),
(Feature.DROP_PARAMS, {}, {"drop_params": True}),
(Feature.MODIFY_PARAMS, {}, {"modify_params": True}),

View file

@ -13,6 +13,15 @@ from litellm import utils
from litellm.litellm_core_utils import litellm_logging, thread_pool_executor
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.rust_bridge.configuration import _CONFIGURATION
from tests.test_litellm_rust.support.caching import (
cache as cache,
)
from tests.test_litellm_rust.support.caching import (
call as call,
)
from tests.test_litellm_rust.support.caching import (
fake_redis as fake_redis,
)
from tests.test_litellm_rust.support.callback_recorder import drain_logging
from tests.test_litellm_rust.support.clickhouse import clickhouse_url as clickhouse_url
from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound

View file

@ -1,16 +1,18 @@
from collections.abc import Generator
from functools import partial
import pytest
from litellm.rust_bridge.catalog import Route
from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE
from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE, call
from tests.test_litellm_rust.support.caching import Call
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.routing import native_route
@pytest.fixture
def server(recording_server: RecordingServer) -> RecordingServer:
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
recording_server.serve(ResponseSpec(body=MESSAGES_RESPONSE))
return recording_server
@ -18,3 +20,8 @@ def server(recording_server: RecordingServer) -> RecordingServer:
def native() -> Generator[None]:
with native_route(Route.MESSAGES):
yield
@pytest.fixture
def native_call(server: RecordingServer) -> Call:
return partial(call, server)

View file

@ -0,0 +1,497 @@
import asyncio
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Final
import pytest
from pydantic import TypeAdapter
import litellm
from litellm.caching.caching import Cache, CacheMode
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.caching import CachingSupportedCallTypes, LiteLLMCacheType
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
from litellm.types.utils import LLMResponseTypes
from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
from tests.test_litellm_rust.support.caching import (
Call,
CallerCache,
FailingLookupCache,
FakeRedis,
HeldWriteCache,
assert_cache_hit_keeps_budget_spend_and_counts_usage,
settle_cache_writes,
use_cache,
)
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
from tests.test_litellm_rust.support.isolation import rebound
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.results import cache_key, consume, response_id, stored_cache_key
from tests.test_litellm_rust.support.routing import assert_served_natively, routed
pytestmark = pytest.mark.requires_rust_extension
_OBJECT: Final = TypeAdapter(dict[str, object])
_MESSAGE: Final = TypeAdapter(AnthropicMessagesResponse)
_CHUNKS: Final = TypeAdapter(list[bytes])
STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS)
EMPTY: Final = {**MESSAGES_RESPONSE, "content": []}
OWN_CALL_TYPES: Final = ("anthropic_messages",)
def body(response: object) -> Mapping[str, object]:
return {name: value for name, value in _OBJECT.validate_python(response).items() if name != "_hidden_params"}
async def collect(stream: object) -> bytes:
streamed_async: Final = isinstance(stream, AsyncIterator)
collected: Final = _CHUNKS.validate_python(list(await consume(stream)), strict=True)
assert streamed_async
await settle_cache_writes()
return b"".join(collected)
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.usefixtures("cache")
async def test_a_repeated_native_call_is_answered_from_the_python_cache(
call: Call, server: RecordingServer, asynchronous: bool
) -> None:
first: Final = await call(asynchronous=asynchronous)
second: Final = await call(asynchronous=asynchronous)
assert_served_natively(server, first)
assert cache_key(first) is None
assert isinstance(cache_key(second), str)
assert body(second) == body(first) == MESSAGES_RESPONSE
@pytest.mark.usefixtures("cache")
async def test_a_completed_native_stream_is_replayed_from_the_python_cache(call: Call, server: RecordingServer) -> None:
server.serve(STREAM)
first_stream: Final = await call(stream=True)
assert cache_key(first_stream) is None
first: Final = await collect(first_stream)
second_stream: Final = await call(stream=True)
assert isinstance(cache_key(second_stream), str)
second: Final = await collect(second_stream)
assert len(server.requests) == 1
assert first == b"".join(STREAM.payloads())
assert second == first
assert_served_natively(server, count=1)
@pytest.mark.usefixtures("cache")
async def test_prompts_passed_positionally_get_their_own_entries(server: RecordingServer) -> None:
from litellm.messages import anthropic_messages
server.expect(2)
async def ask(content: str) -> object:
return await anthropic_messages(
64,
[{"role": "user", "content": content}],
MESSAGES_MODEL,
api_key="test-key",
api_base=server.base_url,
)
await ask("first prompt")
second: Final = await ask("second prompt")
await settle_cache_writes()
assert len(server.requests) == 2
assert cache_key(second) is None
assert_served_natively(server, count=2)
@pytest.mark.parametrize("native_writes", (True, False), ids=("native-writes", "python-writes"))
@pytest.mark.parametrize(
"options",
({}, {"metadata": None}, {"cache_key": "shared-entry"}),
ids=("plain", "metadata-none", "preset-key"),
)
@pytest.mark.usefixtures("cache")
async def test_native_and_python_replay_each_others_stream_entries(
call: Call, server: RecordingServer, native_writes: bool, options: Mapping[str, object]
) -> None:
server.serve(STREAM)
async def stream(native: bool) -> bytes:
if native:
return await collect(await call(stream=True, **options))
with routed():
return await collect(await call(stream=True, **options))
first: Final = await stream(native_writes)
second: Final = await stream(not native_writes)
assert len(server.requests) == 1
assert second == first == b"".join(STREAM.payloads())
@pytest.mark.usefixtures("cache")
async def test_a_stream_closed_early_is_not_stored(call: Call, server: RecordingServer) -> None:
server.serve(STREAM)
server.expect(2)
stream: Final = await call(stream=True)
assert isinstance(stream, AsyncIterator)
await anext(stream)
await stream.aclose()
await collect(await call(stream=True))
assert len(server.requests) == 2
assert_served_natively(server, count=2)
@pytest.mark.parametrize(
("first", "second"),
(
({"cache": {"no-store": True}}, {}),
({}, {"cache": {"no-cache": True}}),
({}, {"caching": False}),
({}, {"messages": [{"role": "user", "content": "Goodbye"}]}),
({}, {"system": "Be brief"}),
({"cache": {"namespace": "first"}}, {"cache": {"namespace": "second"}}),
({"metadata": {"redis_namespace": "first"}}, {"metadata": {"redis_namespace": "second"}}),
),
ids=(
"no-store",
"no-cache",
"caching-off-read",
"other-messages",
"other-param",
"cache-namespace",
"metadata-namespace",
),
)
@pytest.mark.usefixtures("cache")
async def test_cache_controls_and_keys_keep_the_second_call_on_the_provider(
call: Call, server: RecordingServer, first: Mapping[str, object], second: Mapping[str, object]
) -> None:
server.expect(2)
await call(**first)
response: Final = await call(**second)
assert_served_natively(server, count=2)
assert cache_key(response) is None
@pytest.mark.usefixtures("cache")
async def test_caching_off_skips_the_read_but_still_writes(call: Call, server: RecordingServer) -> None:
await call(caching=False)
hit: Final = await call()
assert len(server.requests) == 1
assert body(hit) == MESSAGES_RESPONSE
assert_served_natively(server, count=1)
@pytest.mark.usefixtures("cache")
async def test_each_namespace_keeps_its_own_entry(call: Call, server: RecordingServer) -> None:
server.expect(2)
await call(cache={"namespace": "first"})
await call(cache={"namespace": "second"})
hit: Final = await call(cache={"namespace": "first"})
assert len(server.requests) == 2
assert isinstance(cache_key(hit), str)
assert_served_natively(server, count=2)
@pytest.mark.usefixtures("cache")
async def test_the_key_is_built_from_the_call_arguments_not_the_wire_body(call: Call, server: RecordingServer) -> None:
class Rewrite(RecordingLogger):
temperature = 0.1
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
request_body(kwargs)["temperature"] = self.temperature
super().log_pre_api_call(model, messages, kwargs)
logger: Final = Rewrite()
await call(callbacks=[logger])
logger.temperature = 0.8
hit: Final = await call(callbacks=[logger])
assert len(server.requests) == 1
assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.1
assert isinstance(cache_key(hit), str)
assert_served_natively(server, count=1)
@pytest.mark.parametrize(
("supported", "expected_requests"), ((["embedding"], 2), (OWN_CALL_TYPES, 1)), ids=("other-type", "own-type")
)
async def test_supported_call_types_gate_native_reads_and_writes(
call: Call,
server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
supported: Sequence[CachingSupportedCallTypes],
expected_requests: int,
) -> None:
use_cache(monkeypatch, Cache(supported_call_types=list(supported)))
server.expect(expected_requests)
await call()
second: Final = await call()
assert len(server.requests) == expected_requests
assert isinstance(cache_key(second), str) == (expected_requests == 1)
assert_served_natively(server, count=expected_requests)
@pytest.mark.parametrize(
("options", "expected_requests"), (({}, 2), ({"cache": {"use-cache": True}}, 1)), ids=("plain", "opt-in")
)
async def test_default_off_mode_caches_only_calls_that_opt_in(
call: Call,
server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
options: Mapping[str, object],
expected_requests: int,
) -> None:
use_cache(monkeypatch, Cache(mode=CacheMode.default_off))
server.expect(expected_requests)
await call(**options)
second: Final = await call(**options)
assert len(server.requests) == expected_requests
assert isinstance(cache_key(second), str) == (expected_requests == 1)
assert_served_natively(server, count=expected_requests)
async def test_a_semantic_python_cache_answers_a_similar_prompt(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
use_cache(monkeypatch, Cache(type=LiteLLMCacheType.REDIS_SEMANTIC, _backend=InMemoryCache()))
first: Final = await call(messages=[{"role": "user", "content": "hello"}])
second: Final = await call(messages=[{"role": "user", "content": "hi"}])
assert len(server.requests) == 1
assert cache_key(first) is None
assert isinstance(cache_key(second), str)
assert body(second) == body(first)
assert_served_natively(server, count=1)
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
@pytest.mark.usefixtures("cache")
async def test_a_native_hit_is_logged_as_a_free_cache_hit(call: Call, server: RecordingServer, stream: bool) -> None:
if stream:
server.serve(STREAM)
recorder: Final = RecordingLogger()
async def send() -> None:
response: Final = await call(stream=stream, callbacks=[recorder])
if stream:
await collect(response)
await send()
await send()
await drain_logging()
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
miss, hit = (_OBJECT.validate_python(event.kwargs) for event in successes)
assert miss["cache_hit"] is not True
assert TypeAdapter(float).validate_python(miss["response_cost"]) > 0
assert hit["cache_hit"] is True
assert hit["response_cost"] == 0
assert_served_natively(server, count=1)
async def test_python_cache_lookups_run_in_the_callers_task(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = CallerCache()
use_cache(monkeypatch, Cache(_backend=backend))
await call()
hit: Final = await call()
assert len(server.requests) == 1
assert isinstance(cache_key(hit), str)
assert backend.lookup_tasks == [asyncio.current_task(), asyncio.current_task()]
assert_served_natively(server, count=1)
async def test_cancelling_a_pending_lookup_sends_and_stores_nothing(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
entered: Final = asyncio.Event()
release: Final = asyncio.Event()
backend: Final = CallerCache(entered, release)
use_cache(monkeypatch, Cache(_backend=backend))
server.expect(0)
task: Final = asyncio.create_task(call())
await asyncio.wait_for(entered.wait(), timeout=5)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
release.set()
await asyncio.sleep(0)
assert server.requests == []
assert backend.storage.cache_dict == {}
assert_served_natively(server, count=0)
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
@pytest.mark.usefixtures("cache")
async def test_a_cache_hit_keeps_budget_spend_and_counts_usage(
call: Call, server: RecordingServer, stream: bool
) -> None:
if stream:
server.serve(STREAM)
async def send(**options: object) -> object:
response: Final = await call(stream=stream, **options)
if stream:
await collect(response)
return response
usage: Final = TypeAdapter(dict[str, int]).validate_python(MESSAGES_RESPONSE["usage"])
await assert_cache_hit_keeps_budget_spend_and_counts_usage(
send, server, usage["input_tokens"] + usage["output_tokens"]
)
assert_served_natively(server, count=1)
@pytest.mark.usefixtures("cache")
async def test_a_response_without_output_is_not_stored(call: Call, server: RecordingServer) -> None:
server.serve(ResponseSpec(body=EMPTY))
server.expect(2)
await call()
await call()
assert len(server.requests) == 2
assert_served_natively(server, count=2)
async def test_a_stored_response_without_output_is_a_miss(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = InMemoryCache()
cache: Final = use_cache(monkeypatch, Cache(_backend=backend))
await call()
(key,) = backend.cache_dict
cache.add_cache(EMPTY, cache_key=key)
server.expect(2)
response: Final = await call()
assert len(server.requests) == 2
assert body(response) == MESSAGES_RESPONSE
assert_served_natively(server, count=2)
async def test_a_failing_cache_lookup_runs_the_call_uncached(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
use_cache(monkeypatch, FailingLookupCache())
server.expect(2)
first: Final = await call()
await call()
assert body(first) == MESSAGES_RESPONSE
assert len(server.requests) == 2
assert_served_natively(server, count=2)
async def test_the_response_returns_before_the_cache_write_lands(
call: Call, native_call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = HeldWriteCache()
use_cache(monkeypatch, Cache(_backend=backend))
first: Final = await asyncio.wait_for(native_call(), timeout=5)
backend.release.set()
await settle_cache_writes()
hit: Final = await call()
assert body(first) == MESSAGES_RESPONSE
assert len(server.requests) == 1
assert body(hit) == MESSAGES_RESPONSE
assert_served_natively(server, count=1)
async def test_a_redis_hit_fills_the_worker_memory_tier(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
) -> None:
use_cache(monkeypatch, fake_redis.cache())
await call()
await call()
fake_redis.client.flushall()
hit: Final = await call()
assert len(server.requests) == 1
assert body(hit) == MESSAGES_RESPONSE
assert_served_natively(server, count=1)
async def test_a_stored_response_reaches_the_worker_memory_tier_only_when_read(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
) -> None:
use_cache(monkeypatch, fake_redis.cache())
server.expect(2)
await call()
fake_redis.client.flushall()
await call()
assert len(server.requests) == 2
assert_served_natively(server, count=2)
async def test_a_native_hit_reports_the_key_the_entry_is_stored_under(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = InMemoryCache()
use_cache(monkeypatch, Cache(_backend=backend))
await call()
hit: Final = await call()
(stored,) = backend.cache_dict
assert cache_key(hit) == stored_cache_key(hit) == stored
assert_served_natively(server, count=1)
@pytest.mark.parametrize("served", ("native", "python"))
@pytest.mark.usefixtures("cache")
async def test_a_hit_returns_what_the_success_deployment_hook_handed_the_first_caller(
call: Call, server: RecordingServer, served: str
) -> None:
class Replace(CustomLogger):
calls = 0
async def async_post_call_success_deployment_hook(
self, request_data: dict[str, object], response: object, call_type: object
) -> LLMResponseTypes | None:
Replace.calls += 1
return _MESSAGE.validate_python({**_OBJECT.validate_python(response), "id": f"hooked-{Replace.calls}"})
async def ask() -> object:
if served == "native":
return await call()
with routed():
return await call()
with rebound(litellm, "callbacks", [Replace()]):
first: Final = await ask()
second: Final = await ask()
assert len(server.requests) == 1
assert response_id(first) == "hooked-1"
assert response_id(second) == "hooked-1"
assert Replace.calls == 1

View file

@ -5,7 +5,6 @@ from typing import Final, cast # noqa: TID251 # narrows legacy callable signat
import pytest
import litellm
from litellm.caching.caching import Cache
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
from litellm.llms.anthropic.pass_through.messages import handler as python_messages
@ -457,7 +456,6 @@ ADVISOR_TOOL: Final = {"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "
(Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}),
(Feature.TIMEOUT_VARIANTS, "anthropic/claude-sonnet-4-5", {"request_timeout": 5}, {}),
(Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}),
(Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}),
(Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}),
(Feature.MODIFY_PARAMS, "anthropic/claude-sonnet-4-5", {}, {"modify_params": True}),
(

View file

@ -444,7 +444,7 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl
assert response._hidden_params["response_cost"] is not None
assert response._hidden_params["_response_ms"] > 0
assert trace_id_var.get() == "callback-free-parent"
assert probe.deployments == 2
assert probe.deployments == (2 if failure else 1)
assert probe.submissions == probe.enqueues == 0
assert len(created_loggers) == 1
logger: Final = created_loggers[0]

View file

@ -1,16 +1,18 @@
from collections.abc import Generator
from functools import partial
import pytest
from litellm.rust_bridge.catalog import Route
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE, call
from tests.test_litellm_rust.support.caching import Call
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.routing import native_route
@pytest.fixture
def server(recording_server: RecordingServer) -> RecordingServer:
recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
recording_server.serve(ResponseSpec(body=RESPONSES_RESPONSE))
return recording_server
@ -18,3 +20,8 @@ def server(recording_server: RecordingServer) -> RecordingServer:
def native() -> Generator[None]:
with native_route(Route.RESPONSES):
yield
@pytest.fixture
def native_call(server: RecordingServer) -> Call:
return partial(call, server)

View file

@ -0,0 +1,375 @@
import asyncio
from collections.abc import Mapping, Sequence
from typing import Final
import pytest
from pydantic import TypeAdapter
import litellm
from litellm.caching.caching import Cache, CacheMode
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.caching import CachingSupportedCallTypes
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import LLMResponseTypes
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE, RESPONSES_TEXT
from tests.test_litellm_rust.support.caching import (
Call,
CallerCache,
FailingLookupCache,
FakeRedis,
HeldWriteCache,
assert_cache_hit_keeps_budget_spend_and_counts_usage,
settle_cache_writes,
use_cache,
)
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
from tests.test_litellm_rust.support.isolation import rebound
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.results import cache_key, response_id, stored_cache_key
from tests.test_litellm_rust.support.routing import assert_served_natively, routed
pytestmark = pytest.mark.requires_rust_extension
_OBJECT: Final = TypeAdapter(dict[str, object])
EMPTY: Final = {**RESPONSES_RESPONSE, "output": []}
OWN_CALL_TYPES: Final = ("responses",)
def output_text(response: object) -> str:
assert isinstance(response, ResponsesAPIResponse)
return response.output_text
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.usefixtures("cache")
async def test_a_repeated_native_call_is_answered_from_the_python_cache(
call: Call, server: RecordingServer, asynchronous: bool
) -> None:
first: Final = await call(asynchronous=asynchronous)
second: Final = await call(asynchronous=asynchronous)
assert_served_natively(server, first)
assert cache_key(first) is None
assert isinstance(cache_key(second), str)
assert output_text(second) == output_text(first) == RESPONSES_TEXT
@pytest.mark.parametrize(
("first", "second"),
(
({"cache": {"no-store": True}}, {}),
({}, {"cache": {"no-cache": True}}),
({}, {"caching": False}),
({}, {"input": "Goodbye"}),
({}, {"instructions": "Be brief"}),
({"cache": {"namespace": "first"}}, {"cache": {"namespace": "second"}}),
({"metadata": {"redis_namespace": "first"}}, {"metadata": {"redis_namespace": "second"}}),
),
ids=(
"no-store",
"no-cache",
"caching-off-read",
"other-input",
"other-param",
"cache-namespace",
"metadata-namespace",
),
)
@pytest.mark.usefixtures("cache")
async def test_cache_controls_and_keys_keep_the_second_call_on_the_provider(
call: Call, server: RecordingServer, first: Mapping[str, object], second: Mapping[str, object]
) -> None:
server.expect(2)
await call(**first)
response: Final = await call(**second)
assert_served_natively(server, count=2)
assert cache_key(response) is None
@pytest.mark.usefixtures("cache")
async def test_caching_off_skips_the_read_but_still_writes(call: Call, server: RecordingServer) -> None:
await call(caching=False)
hit: Final = await call()
assert len(server.requests) == 1
assert output_text(hit) == RESPONSES_TEXT
assert_served_natively(server, count=1)
@pytest.mark.usefixtures("cache")
async def test_each_namespace_keeps_its_own_entry(call: Call, server: RecordingServer) -> None:
server.expect(2)
await call(cache={"namespace": "first"})
await call(cache={"namespace": "second"})
hit: Final = await call(cache={"namespace": "first"})
assert len(server.requests) == 2
assert isinstance(cache_key(hit), str)
assert_served_natively(server, count=2)
@pytest.mark.usefixtures("cache")
async def test_the_key_is_built_from_the_call_arguments_not_the_wire_body(call: Call, server: RecordingServer) -> None:
class Rewrite(RecordingLogger):
temperature = 0.1
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
request_body(kwargs)["temperature"] = self.temperature
super().log_pre_api_call(model, messages, kwargs)
logger: Final = Rewrite()
await call(callbacks=[logger])
logger.temperature = 0.8
hit: Final = await call(callbacks=[logger])
assert len(server.requests) == 1
assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.1
assert isinstance(cache_key(hit), str)
assert_served_natively(server, count=1)
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.parametrize(
("supported", "expected_requests"), ((["embedding"], 2), (OWN_CALL_TYPES, 1)), ids=("other-type", "own-type")
)
async def test_supported_call_types_gate_native_reads_and_writes(
call: Call,
server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
asynchronous: bool,
supported: Sequence[CachingSupportedCallTypes],
expected_requests: int,
) -> None:
use_cache(monkeypatch, Cache(supported_call_types=list(supported)))
server.expect(expected_requests)
await call(asynchronous=asynchronous)
second: Final = await call(asynchronous=asynchronous)
assert len(server.requests) == expected_requests
assert isinstance(cache_key(second), str) == (expected_requests == 1)
assert_served_natively(server, count=expected_requests)
@pytest.mark.parametrize(
("options", "expected_requests"), (({}, 2), ({"cache": {"use-cache": True}}, 1)), ids=("plain", "opt-in")
)
async def test_default_off_mode_caches_only_calls_that_opt_in(
call: Call,
server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
options: Mapping[str, object],
expected_requests: int,
) -> None:
use_cache(monkeypatch, Cache(mode=CacheMode.default_off))
server.expect(expected_requests)
await call(**options)
second: Final = await call(**options)
assert len(server.requests) == expected_requests
assert isinstance(cache_key(second), str) == (expected_requests == 1)
assert_served_natively(server, count=expected_requests)
@pytest.mark.usefixtures("cache")
async def test_a_native_hit_is_logged_as_a_free_cache_hit(call: Call, server: RecordingServer) -> None:
recorder: Final = RecordingLogger()
await call(callbacks=[recorder])
await call(callbacks=[recorder])
await drain_logging()
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
miss, hit = (_OBJECT.validate_python(event.kwargs) for event in successes)
assert len(server.requests) == 1
assert miss["cache_hit"] is not True
assert TypeAdapter(float).validate_python(miss["response_cost"]) > 0
assert hit["cache_hit"] is True
assert hit["response_cost"] == 0
assert_served_natively(server, count=1)
async def test_python_cache_lookups_run_in_the_callers_task(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = CallerCache()
use_cache(monkeypatch, Cache(_backend=backend))
await call()
hit: Final = await call()
assert len(server.requests) == 1
assert isinstance(cache_key(hit), str)
assert backend.lookup_tasks == [asyncio.current_task(), asyncio.current_task()]
assert_served_natively(server, count=1)
async def test_cancelling_a_pending_lookup_sends_and_stores_nothing(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
entered: Final = asyncio.Event()
release: Final = asyncio.Event()
backend: Final = CallerCache(entered, release)
use_cache(monkeypatch, Cache(_backend=backend))
server.expect(0)
task: Final = asyncio.create_task(call())
await asyncio.wait_for(entered.wait(), timeout=5)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
release.set()
await asyncio.sleep(0)
assert server.requests == []
assert backend.storage.cache_dict == {}
assert_served_natively(server, count=0)
@pytest.mark.usefixtures("cache")
async def test_a_cache_hit_keeps_budget_spend_and_counts_usage(call: Call, server: RecordingServer) -> None:
usage: Final = TypeAdapter(dict[str, int]).validate_python(RESPONSES_RESPONSE["usage"])
await assert_cache_hit_keeps_budget_spend_and_counts_usage(call, server, usage["total_tokens"])
assert_served_natively(server, count=1)
@pytest.mark.usefixtures("cache")
async def test_a_response_without_output_is_not_stored(call: Call, server: RecordingServer) -> None:
server.serve(ResponseSpec(body=EMPTY))
server.expect(2)
await call()
await call()
assert len(server.requests) == 2
assert_served_natively(server, count=2)
async def test_a_stored_response_without_output_is_a_miss(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = InMemoryCache()
cache: Final = use_cache(monkeypatch, Cache(_backend=backend))
await call()
(key,) = backend.cache_dict
cache.add_cache(EMPTY, cache_key=key)
server.expect(2)
response: Final = await call()
assert len(server.requests) == 2
assert output_text(response) == RESPONSES_TEXT
assert_served_natively(server, count=2)
async def test_a_failing_cache_lookup_runs_the_call_uncached(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
use_cache(monkeypatch, FailingLookupCache())
server.expect(2)
first: Final = await call()
await call()
assert output_text(first) == RESPONSES_TEXT
assert len(server.requests) == 2
assert_served_natively(server, count=2)
async def test_the_response_returns_before_the_cache_write_lands(
call: Call, native_call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = HeldWriteCache()
use_cache(monkeypatch, Cache(_backend=backend))
first: Final = await asyncio.wait_for(native_call(), timeout=5)
backend.release.set()
await settle_cache_writes()
hit: Final = await call()
assert output_text(first) == RESPONSES_TEXT
assert len(server.requests) == 1
assert output_text(hit) == RESPONSES_TEXT
assert_served_natively(server, count=1)
async def test_a_redis_hit_fills_the_worker_memory_tier(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
) -> None:
use_cache(monkeypatch, fake_redis.cache())
await call()
await call()
fake_redis.client.flushall()
hit: Final = await call()
assert len(server.requests) == 1
assert output_text(hit) == RESPONSES_TEXT
assert_served_natively(server, count=1)
async def test_a_stored_response_lands_in_the_worker_memory_tier(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
) -> None:
use_cache(monkeypatch, fake_redis.cache())
first: Final = await call()
fake_redis.client.flushall()
hit: Final = await call()
assert len(server.requests) == 1
assert output_text(hit) == output_text(first)
assert_served_natively(server, count=1)
async def test_a_hit_reports_the_key_the_entry_is_stored_under(
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
) -> None:
backend: Final = InMemoryCache()
use_cache(monkeypatch, Cache(_backend=backend))
await call()
hit: Final = await call()
(stored,) = backend.cache_dict
assert len(server.requests) == 1
assert stored_cache_key(hit) == stored
assert_served_natively(server, count=1)
@pytest.mark.parametrize("served", ("native", "python"))
@pytest.mark.usefixtures("cache")
async def test_a_hit_returns_what_the_success_deployment_hook_handed_the_first_caller(
call: Call, server: RecordingServer, served: str
) -> None:
class Replace(CustomLogger):
calls = 0
async def async_post_call_success_deployment_hook(
self, request_data: dict[str, object], response: object, call_type: object
) -> LLMResponseTypes | None:
Replace.calls += 1
assert isinstance(response, ResponsesAPIResponse)
return response.model_copy(update={"id": f"hooked-{Replace.calls}"})
async def ask() -> object:
if served == "native":
return await call()
with routed():
return await call()
with rebound(litellm, "callbacks", [Replace()]):
first: Final = await ask()
second: Final = await ask()
assert len(server.requests) == 1
assert response_id(first) == "hooked-1"
assert response_id(second) == "hooked-1"
assert Replace.calls == 1

View file

@ -5,7 +5,6 @@ from typing import Final, cast # noqa: TID251 # narrows legacy callable signat
import pytest
import litellm
from litellm.caching.caching import Cache
from litellm.integrations.custom_logger import CustomLogger
from litellm.responses import dispatch as responses_dispatch
from litellm.responses import main as python_responses
@ -287,7 +286,6 @@ ENCODED_ID: Final = ResponsesAPIRequestUtils._build_encrypted_item_id("deploymen
(Feature.DEPLOYMENT_RESPONSE_ID, {"litellm_metadata": {"model_info": {"id": "deployment-1"}}}, {}),
(Feature.ENCODED_ITEM_IDS, {"input": [{"type": "reasoning", "id": ENCODED_ID, "summary": []}]}, {}),
(Feature.MODIFY_PARAMS, {}, {"modify_params": True}),
(Feature.RESPONSE_CACHE, {}, {"cache": Cache()}),
(
Feature.MANAGED_RESOURCE_IDS,
{

View file

@ -0,0 +1,164 @@
import asyncio
import threading
from collections.abc import Callable, Coroutine, Generator
from dataclasses import dataclass
from typing import Final, TypeAlias
from urllib.parse import urlparse
import fakeredis
import pytest
import redis
from pydantic import TypeAdapter
import litellm
from litellm.caching.base_cache import BaseCache
from litellm.caching.caching import Cache
from litellm.caching.caching_handler import (
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
in_memory_cache_obj,
)
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.hooks.model_max_budget_limiter import (
PROXY_VirtualKeyModelMaxBudgetLimiter,
model_budget_spend_cache_key,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.utils import InternalUsageCache
from litellm.types.caching import LiteLLMCacheType
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.recording_server import RecordingServer
Call: TypeAlias = Callable[..., Coroutine[object, object, object]]
class CallerCache(BaseCache):
def __init__(self, lookup_entered: asyncio.Event | None = None, release: asyncio.Event | None = None) -> None:
super().__init__()
self.storage: Final = InMemoryCache()
self.lookup_tasks: Final[list[asyncio.Task[object] | None]] = [] # mutable-ok: test double records each lookup
self.lookup_entered: Final = lookup_entered
self.release: Final = release
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], ttl: float | None = None) -> None:
await self.storage.async_set_cache_pipeline(cache_list, ttl=ttl)
async def async_get_cache(self, key: str, **kwargs: object) -> object:
self.lookup_tasks.append(asyncio.current_task())
if self.lookup_entered is not None and self.release is not None:
self.lookup_entered.set()
await self.release.wait()
return self.storage.get_cache(key, **kwargs)
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
await asyncio.sleep(0)
self.storage.set_cache(key, value, **kwargs)
class HeldWriteCache(InMemoryCache):
def __init__(self) -> None:
super().__init__()
self.release: Final = asyncio.Event()
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
await self.release.wait()
await super().async_set_cache(key, value, **kwargs)
class FailingLookupCache(Cache):
def should_use_cache(self, **kwargs: object) -> bool:
raise RuntimeError("cache lookup failed")
def use_cache(monkeypatch: pytest.MonkeyPatch, cache: Cache) -> Cache:
monkeypatch.setattr(litellm, "cache", cache)
return cache
@pytest.fixture
def cache(monkeypatch: pytest.MonkeyPatch) -> Cache:
return use_cache(monkeypatch, Cache())
@pytest.fixture
def call(native_call: Call) -> Call:
async def send(**kwargs: object) -> object:
response: Final = await native_call(**kwargs)
await settle_cache_writes()
return response
return send
@dataclass(frozen=True, slots=True)
class FakeRedis:
url: str
client: redis.Redis
def cache(self) -> Cache:
parsed: Final = urlparse(self.url)
in_memory_cache_obj.flush_cache()
return Cache(type=LiteLLMCacheType.REDIS, host=parsed.hostname, port=str(parsed.port))
@pytest.fixture
def fake_redis() -> Generator[FakeRedis]:
server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis")
worker: Final = threading.Thread(target=server.serve_forever, daemon=True)
worker.start()
url: Final = f"redis://127.0.0.1:{server.server_address[1]}"
try:
yield FakeRedis(url=url, client=redis.Redis.from_url(url))
finally:
server.shutdown()
server.server_close()
worker.join(timeout=5)
async def settle_cache_writes() -> None:
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
async def assert_cache_hit_keeps_budget_spend_and_counts_usage(
send: Call, server: RecordingServer, tokens: int
) -> None:
counters: Final = litellm.DualCache()
budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters)
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
InternalUsageCache(counters), model_group_resolver=lambda model: model
)
recorder: Final = RecordingLogger()
key_hash: Final = "a" * 64
options: Final = {
"callbacks": [budget, limiter, recorder],
"metadata": {
"user_api_key": key_hash,
"model_group": "cached-model",
"user_api_key_model_max_budget": {"cached-model": {"max_budget": 1, "budget_duration": "1h"}},
},
}
spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h")
token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens")
objects: Final = TypeAdapter(dict[str, object])
await send(**options)
await drain_logging()
miss: Final = objects.validate_python((await recorder.wait_for_async("async_log_success_event"))[0].kwargs)
cost: Final = TypeAdapter(float).validate_python(miss["response_cost"])
assert cost > 0
assert counters.get_cache(spend_key) == pytest.approx(cost)
assert counters.get_cache(token_key) == tokens
await send(**options)
await drain_logging()
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
hit: Final = objects.validate_python(successes[-1].kwargs)
miss_payload: Final = objects.validate_python(miss["standard_logging_object"])
hit_payload: Final = objects.validate_python(hit["standard_logging_object"])
assert len(server.requests) == 1
assert hit["cache_hit"] is True
assert hit["response_cost"] == hit_payload["response_cost"] == 0
assert hit_payload["custom_llm_provider"] == miss_payload["custom_llm_provider"]
assert hit_payload["total_tokens"] == miss_payload["total_tokens"] == tokens
assert counters.get_cache(spend_key) == pytest.approx(cost)
assert counters.get_cache(token_key) == 2 * tokens

View file

@ -51,6 +51,12 @@ class RecordingServer:
def enqueue(self, response: ResponseSpec) -> None:
self.responses.append(response)
def serve(self, response: ResponseSpec) -> None:
self.default_response = response
def expect(self, count: int | None) -> None:
self.expected_requests = count
async def wait_for_requests(self, count: int) -> None:
async with asyncio.timeout(2):
while len(self.requests) < count:

View file

@ -1,8 +1,10 @@
from collections.abc import AsyncIterator, Iterator
from typing import Final
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.types.llms.openai import ResponsesAPIResponse
@ -24,3 +26,14 @@ async def consume(stream: object) -> tuple[object, ...]:
return tuple([chunk async for chunk in stream])
assert isinstance(stream, Iterator)
return tuple(stream)
def cache_key(response: object) -> object:
headers: Final = TypeAdapter(dict[str, object]).validate_python(
get_hidden_params_dict(response).get("additional_headers", {})
)
return headers.get("x-litellm-cache-key")
def stored_cache_key(response: object) -> object:
return get_hidden_params_dict(response).get("cache_key")

View file

@ -3231,3 +3231,20 @@ async def test_a_cache_hit_inside_an_internal_sub_call_is_free_and_fires_no_call
assert async_logging.async_success_handler.called is not internal
assert async_logging.handle_sync_success_callbacks_for_async_calls.called is not internal
assert sync_logging.handle_sync_cache_hit_success_callbacks.called is not internal
def _reply(response: object) -> object:
assert isinstance(response, litellm.ModelResponse)
return response.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # non-streaming choice
@pytest.mark.asyncio
async def test_a_non_string_cache_key_reads_back_the_entry_it_wrote(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
messages: Final = [{"role": "user", "content": "non-string cache key"}]
await litellm.acompletion(model="gpt-4o", messages=messages, mock_response="first", cache_key=123)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
second: Final = await litellm.acompletion(model="gpt-4o", messages=messages, mock_response="second", cache_key=123)
assert _reply(second) == "first"

View file

@ -1,21 +1,337 @@
import asyncio
from typing import Final
import pytest
from pydantic import TypeAdapter
import litellm
from litellm.caching.caching import Cache
from litellm.rust_bridge.host.cache import configured_cache
from litellm.types.caching import LiteLLMCacheType
from litellm.caching.caching import Cache, CacheMode
from litellm.caching.caching_handler import (
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await background cache writes
cache_key_kwargs,
response_cache_key,
)
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import RESPONSE_CACHE_EXCLUDED_PROVIDERS
from litellm.llms.anthropic.pass_through.messages.response_cache import get_cached_stream_events
from litellm.rust_bridge.host.cache import (
MESSAGES_ENTRY,
RESPONSE_ENTRY,
CacheAccess,
EntryFormat,
ResponseOutput,
StreamOutput,
access,
alookup,
astore,
lookup,
store,
)
from litellm.types.caching import CachingSupportedCallTypes, LiteLLMCacheType
REQUEST: Final[dict[str, object]] = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}
STREAM_REQUEST: Final[dict[str, object]] = {**REQUEST, "stream": True}
EMPTY: Final[dict[str, object]] = {"choices": []}
_ENTRY: Final = TypeAdapter(dict[str, object])
MESSAGES_EVENTS: Final = (
'event: message_start\ndata: {"type":"message_start"}\n\n',
'event: content_block_start\ndata: {"type":"content_block_start","index":0}\n\n',
'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n',
'event: message_stop\ndata: {"type":"message_stop"}\n\n',
)
RAW_ENTRY: Final = EntryFormat(
entry=RESPONSE_ENTRY.entry,
output=lambda entry, _: ResponseOutput(value=entry),
tiers_responses=True,
tiers_streams=False,
)
def test_configured_cache_preserves_identity_and_observes_runtime_replacement(monkeypatch: pytest.MonkeyPatch) -> None:
initial: Final = Cache(type=LiteLLMCacheType.LOCAL)
replacement: Final = Cache(type=LiteLLMCacheType.LOCAL)
monkeypatch.setattr(litellm, "cache", initial)
assert configured_cache() is initial
class KeyCountingCache(Cache):
key_builds: int = 0
monkeypatch.setattr(litellm, "cache", replacement)
assert configured_cache() is replacement
def get_cache_key(self, **kwargs: object) -> str:
self.key_builds += 1
return super().get_cache_key(**kwargs)
def response_output(value: object) -> dict[str, object]:
return ResponseOutput(value=value).model_dump()
def configure(monkeypatch: pytest.MonkeyPatch, cache: Cache) -> Cache:
monkeypatch.setattr(litellm, "cache", cache)
return cache
def _cache_access(
cache: Cache,
request: dict[str, object] = REQUEST,
memory_tier: DualCache | None = None,
entry_format: EntryFormat = RESPONSE_ENTRY,
) -> CacheAccess:
return CacheAccess(
cache=cache,
request=cache_key_kwargs(request),
key=response_cache_key(cache, request),
read=True,
write=True,
memory_tier=memory_tier,
format=entry_format,
)
def stored_entry(cache: Cache, request: dict[str, object]) -> object:
hit: Final = lookup(_cache_access(cache, request, entry_format=RAW_ENTRY))
return None if hit is None else hit.response["value"]
async def memory_entry(memory_tier: DualCache, key: object) -> object:
return await memory_tier.in_memory_cache.async_get_cache(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # InMemoryCache.async_get_cache is unannotated
def grants(kwargs: dict[str, object], call_type: str) -> tuple[bool, bool] | None:
granted: Final = access(kwargs, call_type)
return None if granted is None else (granted.read, granted.write)
def test_no_configured_cache_grants_no_access(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "cache", None)
assert configured_cache() is None
assert access(REQUEST, "acompletion") is None
@pytest.mark.parametrize(
("overrides", "expected"),
(
({}, (True, True)),
({"cache": {"no-cache": True}}, (False, True)),
({"cache": {"no-store": True}}, (True, False)),
({"cache": {"no-cache": "true"}}, (True, True)),
({"caching": False}, (False, True)),
({"cache": {"no-cache": True, "no-store": True}}, None),
),
ids=("default", "no-cache", "no-store", "non-bool-control", "caching-off", "both-off"),
)
def test_access_follows_the_python_handler_controls(
monkeypatch: pytest.MonkeyPatch, overrides: dict[str, object], expected: tuple[bool, bool] | None
) -> None:
configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL))
assert grants({**REQUEST, **overrides}, "acompletion") == expected
def test_an_excluded_provider_grants_no_access(monkeypatch: pytest.MonkeyPatch) -> None:
configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL))
provider: Final = next(iter(RESPONSE_CACHE_EXCLUDED_PROVIDERS))
assert access({**REQUEST, "custom_llm_provider": provider}, "acompletion") is None
assert access({**REQUEST, "model": f"{provider}/gpt-4o"}, "acompletion") is None
@pytest.mark.parametrize(
("supported", "call_type", "granted"),
(
(["acompletion"], "acompletion", True),
(["embedding"], "acompletion", False),
(["responses"], "aresponses", True),
(["aresponses"], "responses", False),
),
)
def test_supported_call_types_gate_access(
monkeypatch: pytest.MonkeyPatch, supported: list[CachingSupportedCallTypes], call_type: str, granted: bool
) -> None:
configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL, supported_call_types=supported))
assert (access(REQUEST, call_type) is not None) == granted
@pytest.mark.parametrize(
("overrides", "granted"),
(({}, False), ({"cache": {"use-cache": True}}, True)),
ids=("plain", "opt-in"),
)
def test_default_off_mode_grants_only_calls_that_opt_in_and_hashes_nothing_otherwise(
monkeypatch: pytest.MonkeyPatch, overrides: dict[str, object], granted: bool
) -> None:
cache: Final = KeyCountingCache(type=LiteLLMCacheType.LOCAL, mode=CacheMode.default_off)
configure(monkeypatch, cache)
assert (access({**REQUEST, **overrides}, "acompletion") is not None) is granted
assert cache.key_builds == int(granted)
def test_a_miss_then_store_builds_the_key_once(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = KeyCountingCache(type=LiteLLMCacheType.LOCAL)
configure(monkeypatch, cache)
granted: Final = access(REQUEST, "acompletion")
assert granted is not None
assert lookup(granted) is None
store(granted, response_output({"id": "cached"}))
hit: Final = lookup(granted)
assert hit is not None
assert hit.response == response_output({"id": "cached"})
assert cache.key_builds == 1
def test_a_cache_that_fails_to_decide_grants_no_access(monkeypatch: pytest.MonkeyPatch) -> None:
class FailingCache(Cache):
def should_use_cache(self, **kwargs: object) -> bool:
raise RuntimeError("cache unavailable")
configure(monkeypatch, FailingCache(type=LiteLLMCacheType.LOCAL))
assert access(REQUEST, "acompletion") is None
@pytest.mark.parametrize("cache_key", (None, "preset-key"), ids=("generated-key", "preset-key"))
def test_lookup_returns_cached_response_and_key(cache_key: str | None) -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
request: Final = REQUEST if cache_key is None else {**REQUEST, "cache_key": cache_key}
cache_access: Final = _cache_access(cache, request)
expected_key: Final = cache.get_cache_key(**REQUEST) if cache_key is None else cache_key
store(cache_access, response_output({"id": "cached"}))
hit: Final = lookup(cache_access)
assert hit is not None
assert hit.response == response_output({"id": "cached"})
assert hit.key == expected_key
def test_lookup_skips_responses_without_output() -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
cache_access: Final = _cache_access(cache)
store(cache_access, response_output(EMPTY))
assert lookup(cache_access) is None
cache.add_cache(EMPTY, **REQUEST)
assert lookup(cache_access) is None
@pytest.mark.asyncio
async def test_async_lookup_uses_memory_tier_and_discards_empty_responses() -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
memory_tier: Final = DualCache(in_memory_cache=InMemoryCache())
cache_access: Final = _cache_access(cache, memory_tier=memory_tier)
assert await alookup(cache_access) is None
await cache.async_add_cache({"id": "cached"}, dynamic_cache_object=memory_tier, **REQUEST)
hit: Final = await alookup(cache_access)
assert hit is not None
assert hit.response == response_output({"id": "cached"})
assert hit.key == cache.get_cache_key(**REQUEST)
empty_request: Final = {**REQUEST, "cache_key": "empty-response"}
await cache.async_add_cache(EMPTY, dynamic_cache_object=memory_tier, **empty_request)
assert await alookup(_cache_access(cache, empty_request, memory_tier)) is None
assert await memory_entry(memory_tier, "empty-response") is None
@pytest.mark.asyncio
async def test_async_store_routes_writes_and_skips_responses_without_output() -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
memory_tier: Final = DualCache(in_memory_cache=InMemoryCache())
cache_access: Final = _cache_access(cache, memory_tier=memory_tier)
await astore(cache_access, response_output({"id": "cached"}))
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
hit: Final = await alookup(cache_access)
assert hit is not None
assert hit.response == response_output({"id": "cached"})
assert hit.key == cache.get_cache_key(**REQUEST)
empty_request: Final = {**REQUEST, "cache_key": "empty-response"}
empty_access: Final = _cache_access(cache, empty_request, memory_tier)
await astore(empty_access, response_output(EMPTY))
assert await alookup(empty_access) is None
assert await memory_entry(memory_tier, "empty-response") is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
("entry_format", "request_kwargs", "output", "tiered"),
(
(RESPONSE_ENTRY, REQUEST, ResponseOutput(value={"id": "cached"}), True),
(MESSAGES_ENTRY, REQUEST, ResponseOutput(value={"id": "cached"}), False),
(MESSAGES_ENTRY, STREAM_REQUEST, StreamOutput(value="".join(MESSAGES_EVENTS)), True),
),
ids=("response", "messages-response", "messages-stream"),
)
async def test_an_async_store_reaches_the_worker_memory_tier_as_the_python_route_does(
entry_format: EntryFormat,
request_kwargs: dict[str, object],
output: ResponseOutput | StreamOutput,
tiered: bool,
) -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
memory_tier: Final = DualCache(in_memory_cache=InMemoryCache())
cache_access: Final = _cache_access(cache, request_kwargs, memory_tier, entry_format)
await astore(cache_access, output.model_dump())
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
assert (await memory_entry(memory_tier, cache_access.key) is not None) is tiered
assert (stored_entry(cache, request_kwargs) is not None) is not tiered
@pytest.mark.asyncio
async def test_only_a_redis_backed_cache_answers_from_the_worker_memory_tier(monkeypatch: pytest.MonkeyPatch) -> None:
request: Final = {**REQUEST, "messages": [{"role": "user", "content": "worker tier"}]}
configure(monkeypatch, Cache(type=LiteLLMCacheType.REDIS, host="127.0.0.1", port="1"))
redis_access: Final = access(request, "acompletion")
configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL))
local_access: Final = access(request, "acompletion")
assert redis_access is not None and local_access is not None
await astore(redis_access, response_output({"id": "from-the-worker"}))
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
hit: Final = await alookup(redis_access)
assert hit is not None
assert hit.response == response_output({"id": "from-the-worker"})
assert local_access.memory_tier is None
def test_cache_controls_that_cannot_be_read_grant_no_access(monkeypatch: pytest.MonkeyPatch) -> None:
configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL))
assert access({**REQUEST, "cache": {1: True}}, "acompletion") is None
def test_a_messages_stream_is_stored_as_the_legacy_event_list_and_replayed() -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
cache_access: Final = _cache_access(cache, STREAM_REQUEST, entry_format=MESSAGES_ENTRY)
stream: Final = StreamOutput(value="".join(MESSAGES_EVENTS)).model_dump()
store(cache_access, stream)
hit: Final = lookup(cache_access)
assert get_cached_stream_events(_ENTRY.validate_python(stored_entry(cache, STREAM_REQUEST))) == MESSAGES_EVENTS
assert hit is not None
assert hit.response == stream
def test_a_stream_on_a_route_without_a_stream_entry_is_not_stored() -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
store(_cache_access(cache, STREAM_REQUEST), StreamOutput(value="".join(MESSAGES_EVENTS)).model_dump())
assert stored_entry(cache, STREAM_REQUEST) is None
@pytest.mark.parametrize(("request_kwargs", "hit"), ((REQUEST, True), (STREAM_REQUEST, False)), ids=("unary", "stream"))
def test_a_stored_response_answers_only_a_unary_request(request_kwargs: dict[str, object], hit: bool) -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
cache.add_cache({"id": "cached"}, **request_kwargs)
assert (lookup(_cache_access(cache, request_kwargs)) is not None) is hit

View file

@ -8,7 +8,6 @@ import httpx
import pytest
import litellm
from litellm.caching.caching import Cache
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.responses.utils import ResponsesAPIRequestUtils
@ -408,22 +407,6 @@ def test_websearch_interception_needs_the_interception_logger(monkeypatch: pytes
assert detectors.websearch_interception(call())
@pytest.mark.parametrize(
("cache", "kwargs", "expected"),
(
(None, {}, False),
(Cache(), {}, True),
(Cache(), {"cache": {"ttl": 60}}, True),
(Cache(), {"cache": {"no-cache": True}}, False),
),
)
def test_response_cache_is_skipped_only_by_no_cache(
monkeypatch: pytest.MonkeyPatch, cache: Cache | None, kwargs: Mapping[str, object], expected: bool
) -> None:
monkeypatch.setattr(litellm, "cache", cache)
assert detectors.response_cache(call(**kwargs)) is expected
@pytest.mark.parametrize(
("environment", "setting", "expected"),
((None, None, False), ("DEFAULT@SECLEVEL=1", None, True), (None, "DEFAULT@SECLEVEL=1", True)),

View file

@ -32,6 +32,7 @@ _CACHE_PARAMETER_HELPERS: Final = frozenset(
"litellm/llms/anthropic/prompt_cache_prediction.py",
"litellm/proxy/common_utils/cache_coordinator.py",
"litellm/proxy/common_utils/user_api_key_cache.py",
"litellm/rust_bridge/host/cache.py",
"litellm/utils.py",
}
)