scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

refactor(rust-bridge): serve the Python cache around native calls
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details

Browse source 
Python owns the response cache for native calls: the host driver looks the call up after argument preparation and before the native machine starts, delivers a hit without touching the provider, and stores the public response the deployment success hook answered or the stream the caller read to its end. Hits and hook-reviewed responses no longer pass through native response types, so fields Rust does not model survive. The ProviderResult intercept and the Python-backed Rust cache service are removed.

Co-authored-by: Yujong Lee <yujong@berri.ai>
This commit is contained in:
Devin AI 2026-10-17 22:15:21 +00:00
parent c682c05396
commit 64b8d8a8d3
82 changed files with 3782 additions and 1949 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
litellm-rust/crates/cache-response/AGENTS.md
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

7
litellm-rust/crates/cache-response/src/caching.rs
View file

@ -73,10 +73,17 @@ pub fn get_cache_key(input: &CacheKeyInput) -> String {
cache_key(input)
}
#[cfg(test)]
thread_local! {
pub(crate) static KEY_DERIVATIONS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}
pub fn cache_key(input: &CacheKeyInput) -> String {
if let Some(preset) = &input.preset {
return preset.clone();
}
#[cfg(test)]
KEY_DERIVATIONS.with(|count| count.set(count.get() + 1));
let mut digest = Sha256::new();
for field in &input.fields {
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))

4
litellm-rust/crates/cache-response/src/lib.rs
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,
};

29
litellm-rust/crates/cache-response/src/response.rs
View file

@ -124,15 +124,21 @@ where
&self,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
self.async_lookup_with_key(request, || cache_key(&request.key), now)
.await
}
pub(crate) async fn async_lookup_with_key(
&self,
request: &ResponseCacheRequest<B::Context>,
key: impl FnOnce() -> String + Send,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
return Ok(None);
}
let entry = match self
.backend
.async_get_cache(&cache_key(&request.key), &request.context)
.await
{
let entry = match self.backend.async_get_cache(&key(), &request.context).await {
Ok(entry) => entry,
Err(Error::InvalidEntry) => None,
Err(error) => return Err(error),
@ -253,13 +259,24 @@ where
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
self.async_store_with_key(request, || cache_key(&request.key), response, now)
.await
}
pub(crate) async fn async_store_with_key(
&self,
request: &ResponseCacheRequest<B::Context>,
key: impl FnOnce() -> String + Send,
response: Value,
now: Duration,
) -> Result<(), Error> {
if !request.controls.writes() || !self.fits(&response) {
return Ok(());
}
self.backend
.async_set_cache(
&cache_key(&request.key),
&key(),
CacheEntry {
timestamp: Some(now.as_secs_f64()),
response,

240
litellm-rust/crates/cache-response/src/service.rs
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,87 @@ 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 {
let policy = options.policy;
let request = options.request(&self.config().namespace, surface, input());
PreparedCache {
policy,
max_entry_bytes: self.config().max_entry_bytes,
call: Arc::new(NativeCacheCall {
key: cache_key(&request.key),
request,
service: self,
}),
}
}
}
struct NativeCacheCall<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
service: Arc<ResponseCache<B>>,
request: ResponseCacheRequest,
key: String,
}
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_with_key(&self.request, || self.key.clone(), now)
.await?
.map(|value| CacheHit {
value,
key: self.key.clone(),
}))
})
}
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_with_key(
&self.request,
|| self.key.clone(),
response,
now,
))
}
}
@ -83,8 +125,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 +233,131 @@ 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);
}
#[rstest]
#[tokio::test]
async fn native_lookup_store_and_hit_reporting_hash_once() {
let service = Arc::new(ResponseCache::new(Arc::new(InMemoryCache::default())));
let before = crate::caching::KEY_DERIVATIONS.with(std::cell::Cell::get);
let call = service
.prepare(
CacheOptions::new(CacheScope::Shared),
"test",
&|| json!({"prompt": "hello"}),
)
.call;
assert_eq!(call.lookup(Duration::ZERO).await.unwrap(), None);
call.store(json!({"answer": 42}), Duration::ZERO)
.await
.unwrap();
let hit = call.lookup(Duration::ZERO).await.unwrap().unwrap();
assert_eq!(hit.value, json!({"answer": 42}));
assert!(!hit.key.is_empty());
assert_eq!(
crate::caching::KEY_DERIVATIONS.with(std::cell::Cell::get) - before,
1
);
}
#[rstest]
#[tokio::test]
async fn prepared_native_keys_keep_freshness_and_entry_limits() {
let service = Arc::new(
ResponseCache::new(Arc::new(InMemoryCache::default())).with_config(
ResponseCacheConfig {
namespace: "test".into(),
max_entry_bytes: 16,
},
),
);
let options = CacheOptions {
policy: CachePolicy {
max_age: Some(Duration::from_secs(1)),
..Default::default()
},
scope: CacheScope::Shared,
};
let input = json!({"body": "same request"});
let request = options.clone().request("test", "test", input.clone());
let call = service
.clone()
.prepare(options, "test", &|| input.clone())
.call;
call.store(json!({"answer": 42}), Duration::ZERO)
.await
.unwrap();
let hit = call.lookup(Duration::from_secs(1)).await.unwrap().unwrap();
assert_eq!(hit.key, cache_key(&request.key));
assert_eq!(hit.value, json!({"answer": 42}));
assert_eq!(call.lookup(Duration::from_secs(2)).await.unwrap(), None);
call.store(
json!({"answer": "too large for this cache"}),
Duration::from_secs(2),
)
.await
.unwrap();
assert_eq!(call.lookup(Duration::from_secs(2)).await.unwrap(), None);
assert_eq!(
service
.async_lookup(&request, Duration::ZERO)
.await
.unwrap(),
Some(json!({"answer": 42}))
);
}
}

97
litellm-rust/crates/cache-response/tests/service.rs
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]

76
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
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,28 @@ 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 provider_result(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
if !self.runs_deployment_hooks() {
return Ok(HookStep::Ready(response));
}
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 +794,27 @@ 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}\nresponse = object()");
let (mut logging, _) = begin(py, &locals, asynchronous);
let step = logging
.provider_result(py, local(&locals, "response").unbind())
.unwrap();
assert_eq!(awaits_deployment_hook(&step), 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 +835,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 +860,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 +885,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()

15
litellm-rust/crates/callbacks-legacy-python/src/mapping.rs
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,10 @@ impl CallHooks<PythonRuntime> for LegacyLogging {
(BEFORE.invoke)(self, py, wire, context)
}
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<'_>,

4
litellm-rust/crates/gateway-inference/tests/caching.rs
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()

2
litellm-rust/crates/host-python/AGENTS.md
View file

@ -10,7 +10,7 @@ This crate owns CPython execution mechanics for generic `litellm-host` machines
Creating a resolved asyncio Future from an already constructed Python value belongs here, alongside runtime waiting, interpreter detachment and panic containment. Choosing which callable exceptions become a public `RuntimeError` belongs to the consumer. `python-bridge::callable::wrap_failure` owns that policy
The driver owns ordering: start, argument preparation, prepared-argument hooks, binding decode and machine start. Fallible per-call resource setup supplied by the consumer runs after all argument hooks, using the prepared argument view, and before provider work. Setup failure follows the existing terminal failure path. Creating or discarding an unstarted coroutine must not initialize clients, acquire credentials or capture execution context
The driver owns ordering: start, argument preparation, prepared-argument hooks, binding decode, the binding's lookup and machine start. A lookup hit is delivered through the same result-ready, transform and terminal steps as a provider result and never starts the machine. A provider's complete response is offered to the `provider_result` hooks as the public object, the binding stores what they answered, and that object is what the caller receives without converting it back into a native type. Fallible per-call resource setup supplied by the consumer runs after all argument hooks, using the prepared argument view, and before provider work. Setup failure follows the existing terminal failure path. Creating or discarding an unstarted coroutine must not initialize clients, acquire credentials or capture execution context
Boundary tests exercise behavior with a supplied lifecycle binding without importing the LiteLLM Python package. Pin inline awaiting, awaitable final values, exception identity, cancellation, re-entry and release of retained objects, rather than module names or source layout

1
litellm-rust/crates/host-python/Cargo.toml
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"] }

55
litellm-rust/crates/host-python/src/binding.rs
View file

@ -1,7 +1,29 @@
use crate::{InvokeError, PythonOwned};
use litellm_host::protocol::Protocol;
use crate::{Error, InvokeError, PythonOwned};
use litellm_host::{interceptors::ExecutionFacts, protocol::Protocol};
use pyo3::prelude::*;
use pyo3::types::PyDict;
use serde_json::Value;
/// A host lookup's progress: no answer, Python work to await first, or the call's answer.
pub enum Lookup<P: Protocol> {
Miss,
Await(Py<PyAny>),
Hit(Hit<P>),
}
/// An answer the host found itself, delivered without starting the native machine.
pub struct Hit<P: Protocol> {
pub facts: ExecutionFacts,
pub output: HitOutput<P>,
}
pub enum HitOutput<P: Protocol> {
Response(Py<PyAny>),
Stream {
head: P::StreamHead,
chunks: Vec<P::Chunk>,
},
}
/// Converts requests, responses, stream values and errors at the Python boundary.
pub trait PythonBinding: PythonOwned {
@ -27,6 +49,35 @@ 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)
}
/// Looks the decoded call up in the host's own store before the native machine starts.
fn lookup(&mut self, _py: Python<'_>) -> PyResult<Lookup<Self::Protocol>> {
Ok(Lookup::Miss)
}
/// Continues a lookup with the result of the Python work it awaited.
fn resume_lookup(
&mut self,
_py: Python<'_>,
_answer: PyResult<Py<PyAny>>,
) -> PyResult<Lookup<Self::Protocol>> {
Err(Error::MissingState.into())
}
/// Keeps the response a provider call handed the caller, after its hooks reviewed it.
fn store_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> {
Ok(())
}
/// Keeps the provider stream the caller read to its end.
fn store_stream(&mut self, _py: Python<'_>) -> PyResult<()> {
Ok(())
}
/// What the stream carries at hand-off, as the caller's stream receives it.
fn encode_stream_head(
&mut self,

164
litellm-rust/crates/host-python/src/driver.rs
View file

@ -1,3 +1,4 @@
use std::ops::ControlFlow;
use std::pin::pin;
use std::sync::Arc;
use std::task::{Context, Poll, Waker};
@ -11,6 +12,7 @@ use pyo3::types::PyDict;
use litellm_coroutine::{Coroutine, CoroutineState};
use litellm_host::{
call::HostedCompletion,
interceptors::ResultSource,
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin, Timing, epoch_seconds},
machine::{HostFailure, Machine, MachineStep},
observation::ObservationSender,
@ -22,7 +24,7 @@ use crate::caller::{Answer, PythonCaller, Yielded};
use crate::handle::{Execution, ExecutionBody, ExecutionStep, StepTypes};
use crate::hooks::{HookStep, PythonCallEvent, PythonCallHooks};
use crate::native::{NativeMachine, NativePoll};
use crate::{Error, InvokeError, PythonBinding};
use crate::{Error, Hit, HitOutput, InvokeError, Lookup, PythonBinding};
type ProtocolOf<H> = <H as PythonBinding>::Protocol;
type ErrorOf<H> = <ProtocolOf<H> as Protocol>::Error;
@ -133,6 +135,7 @@ where
native: NativeMachine<M>,
start: Option<StartMachine<ProtocolOf<H>, M>>,
streaming: bool,
cached: bool,
}
impl<H, M, L> CallBody<H, M, L>
@ -148,7 +151,8 @@ where
async fn run(mut self) -> PyResult<Py<PyAny>> {
match self.begin().await {
Ok(()) => self.execute().await,
Ok(None) => self.execute().await,
Ok(Some(hit)) => self.replay(hit).await,
Err(Setup::Host(error)) => self.failure(error, FailureOrigin::Host).await,
Err(Setup::Decode(InvokeError::Python(error))) => {
self.failure(error, FailureOrigin::Call).await
@ -157,7 +161,7 @@ where
}
}
async fn begin(&mut self) -> Result<(), Setup<ErrorOf<H>>> {
async fn begin(&mut self) -> Result<Option<Hit<ProtocolOf<H>>>, Setup<ErrorOf<H>>> {
let started_at = epoch_seconds();
let started = {
let mut parts = self.parts();
@ -198,6 +202,9 @@ where
decoded.map_err(Setup::Decode)
})?
};
if let Some(hit) = self.look_up().await.map_err(Setup::Host)? {
return Ok(Some(hit));
}
let start = self
.start
.take()
@ -211,9 +218,100 @@ where
Python::attach(|py| start(py, arguments.bind(py), request)).map_err(Setup::Host)?
};
self.native.start(machine);
Ok(None)
}
async fn look_up(&self) -> PyResult<Option<Hit<ProtocolOf<H>>>> {
let mut lookup = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.lookup(py))?
};
loop {
let awaitable = match lookup {
Lookup::Miss => return Ok(None),
Lookup::Hit(hit) => return Ok(Some(hit)),
Lookup::Await(awaitable) => awaitable,
};
let answer = self.caller.await_python(awaitable).await;
lookup = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.resume_lookup(py, answer))?
};
}
}
/// Delivers a host hit as the call's result, through the same events and hooks a provider
/// result takes, minus the ones that review a provider's answer.
async fn replay(&mut self, hit: Hit<ProtocolOf<H>>) -> PyResult<Py<PyAny>> {
self.cached = true;
let ready = PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts: hit.facts });
if let Err(error) = self.notify(ready).await {
return self.failure(error, FailureOrigin::Call).await;
}
let (head, chunks) = match hit.output {
HitOutput::Response(public) => {
self.parts().observation.ended_at = Some(epoch_seconds());
return self.deliver(public).await;
}
HitOutput::Stream { head, chunks } => (head, chunks),
};
if let Err(error) = self.replay_stream(head, chunks).await {
return self.failure(error, FailureOrigin::Call).await;
}
self.parts().observation.ended_at = Some(epoch_seconds());
let none = Python::attach(|py| py.None());
self.succeeded(none).await
}
async fn replay_stream(
&mut self,
head: <ProtocolOf<H> as Protocol>::StreamHead,
chunks: Vec<<ProtocolOf<H> as Protocol>::Chunk>,
) -> PyResult<()> {
if self.open_stream(head).await?.is_break() {
return Ok(());
}
for chunk in chunks {
if self.deliver_chunk(chunk).await?.is_break() {
break;
}
}
Ok(())
}
async fn open_stream(
&mut self,
head: <ProtocolOf<H> as Protocol>::StreamHead,
) -> PyResult<ControlFlow<()>> {
self.streaming = true;
let head = {
let mut parts = self.parts();
let Parts { binding, hooks, .. } = &mut *parts;
Python::attach(|py| {
let head = binding.encode_stream_head(py, head)?;
hooks.on_stream_open(py, &head)?;
Ok::<_, PyErr>(head)
})?
};
Ok(self.caller.open_stream(head).await)
}
async fn deliver_chunk(
&mut self,
chunk: <ProtocolOf<H> as Protocol>::Chunk,
) -> PyResult<ControlFlow<()>> {
let chunk = {
let mut parts = self.parts();
let Parts { binding, hooks, .. } = &mut *parts;
Python::attach(|py| {
let chunk = binding.encode_chunk(py, chunk)?;
hooks.on_stream_chunk(py, &chunk)?;
Ok::<_, PyErr>(chunk)
})?
};
Ok(self.caller.yield_chunk(chunk).await)
}
/// Runs one hook step to its value, handing each awaitable it produces to the caller.
async fn settle<T>(&self, step: HookStep<L, T>) -> PyResult<T> {
let mut step = step;
@ -295,6 +393,7 @@ where
Ok(())
}
HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => {
self.cached = matches!(facts.source, ResultSource::Cache { .. });
self.notify(PythonCallEvent::Execution(ExecutionEvent::ResultReady {
facts,
}))
@ -311,30 +410,11 @@ where
Ok(())
}
HostRequest::Stream(StreamDelivery::Open(head, reply)) => {
self.streaming = true;
let head = {
let mut parts = self.parts();
let Parts { binding, hooks, .. } = &mut *parts;
Python::attach(|py| {
let head = binding.encode_stream_head(py, head)?;
hooks.on_stream_open(py, &head)?;
Ok::<_, PyErr>(head)
})?
};
reply.send(self.caller.open_stream(head).await);
reply.send(self.open_stream(head).await?);
Ok(())
}
HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => {
let chunk = {
let mut parts = self.parts();
let Parts { binding, hooks, .. } = &mut *parts;
Python::attach(|py| {
let chunk = binding.encode_chunk(py, chunk)?;
hooks.on_stream_chunk(py, &chunk)?;
Ok::<_, PyErr>(chunk)
})?
};
reply.send(self.caller.yield_chunk(chunk).await);
reply.send(self.deliver_chunk(chunk).await?);
Ok(())
}
}
@ -374,6 +454,17 @@ where
self.parts().observation.ended_at = Some(epoch_seconds());
let response = match completion {
HostedCompletion::Complete(response) => response,
HostedCompletion::StreamEnded if !self.cached => {
let stored = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.store_stream(py))
};
if let Err(error) = stored {
return self.failure(error, FailureOrigin::Host).await;
}
let none = Python::attach(|py| py.None());
return self.succeeded(none).await;
}
HostedCompletion::StreamEnded | HostedCompletion::Detached => {
let none = Python::attach(|py| py.None());
return self.succeeded(none).await;
@ -390,6 +481,30 @@ where
if self.streaming {
return self.succeeded(public).await;
}
if self.cached {
return self.deliver(public).await;
}
let review = {
let mut parts = self.parts();
Python::attach(|py| parts.hooks.provider_result(py, public))
};
let reviewed = match review {
Ok(step) => self.settle(step).await,
Err(error) => Err(error),
};
let stored = reviewed.and_then(|public| {
let mut parts = self.parts();
Python::attach(|py| parts.binding.store_response(py, public.bind(py)))?;
Ok(public)
});
match stored {
Ok(public) => self.deliver(public).await,
Err(error) => self.failure(error, FailureOrigin::Host).await,
}
}
/// Hands the caller's public response to the transform hooks and settles the call.
async fn deliver(&mut self, public: Py<PyAny>) -> PyResult<Py<PyAny>> {
let transform = {
let mut parts = self.parts();
let timing = parts.observation.timing();
@ -629,6 +744,7 @@ where
native: NativeMachine::new(asynchronous),
start: Some(Box::new(start)),
streaming: false,
cached: false,
}
.run()
}

14
litellm-rust/crates/host-python/src/hooks/chain/adapter.rs
View file

@ -71,6 +71,11 @@ pub(super) trait ChainHooks: PythonOwned {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Box<WireRequest>>>;
fn provider_result(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
) -> PyResult<ChainStep<Py<PyAny>>>;
fn transform_response(
&mut self,
py: Python<'_>,
@ -138,6 +143,15 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
Ok(self.step(step, Continuation::Wire))
}
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<'_>,

44
litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs
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,14 @@ impl CallHooks<PythonRuntime> for HookChain {
self.transform_from::<Wire>(py, 0, wire, context.clone())
}
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<'_>,

2
litellm-rust/crates/host-python/src/lib.rs
View file

@ -21,7 +21,7 @@ mod runtime;
mod services;
pub use argument::{effective_py_args, lookup, present};
pub use binding::PythonBinding;
pub use binding::{Hit, HitOutput, Lookup, PythonBinding};
pub use conversion_cache::{FromPythonCache, ToPythonCache};
pub use driver::{CallOptions, run_call};
pub use error::{Error, InvokeError};

68
litellm-rust/crates/host-python/tests/driver.rs
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,72 @@ 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]
#[case::hit(Some("cached"), "cached", &["lookup"])]
#[case::miss(
None,
"reviewed",
&["lookup", "complete", "provider_result:provider", "store:reviewed"]
)]
fn a_host_hit_skips_the_provider_and_a_miss_stores_what_the_hooks_answered(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
#[case] cached: Option<&'static str>,
#[case] returned: &str,
#[case] steps: &[&str],
) {
Python::attach(|py| {
let log = Log::default();
let (result, log) = run_composed(
py,
reviewed_machine(),
binding(&log, OpScript::Answer).with_cache(cached),
hooks(&log, HookScript::ReviewResult),
std::convert::identity,
call_options(asynchronous),
);
let seen: Vec<_> = log
.iter()
.filter(|entry| {
["lookup", "complete", "provider_result:", "store:"]
.iter()
.any(|step| entry.starts_with(step))
})
.cloned()
.collect();
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), returned);
assert_eq!(seen, steps);
});
}
#[rstest]
fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
#[from(interpreter)] _python: Interpreter,

75
litellm-rust/crates/host-python/tests/support/mod.rs
View file

@ -6,14 +6,17 @@ use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use litellm_host::{
call::{CallOutput, HostedMachine, hosted_call},
hooks::CallHooks,
interceptors::{RawResponse, RequestContext, WireRequest},
interceptors::{
Cost, ExecutionFacts, ProviderIdentity, RawResponse, RequestContext, ResultSource,
WireRequest,
},
lifecycle::{ExecutionEvent, Timing},
machine::{CallMachine, MachineFault},
protocol::{Protocol, Reply},
};
use litellm_host_python::{
CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
PythonHostCalls, PythonOwned, PythonRuntime, run_call,
CallOptions, Hit, HitOutput, HookStep, InvokeError, Lookup, PythonBinding, PythonCallEvent,
PythonCallHooks, PythonHostCalls, PythonOwned, PythonRuntime, run_call,
};
use pyo3::{
exceptions::{PyTypeError, PyValueError, asyncio::CancelledError},
@ -169,6 +172,7 @@ pub struct SyntheticBinding {
op: OpScript,
classifier_fails: bool,
pending_reply: Option<Reply<String>>,
cache: Option<Option<&'static str>>,
}
pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
@ -177,10 +181,19 @@ pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
op,
classifier_fails: false,
pending_reply: None,
cache: None,
}
}
impl SyntheticBinding {
/// Looks every call up in a host cache holding `cached`, or nothing.
pub fn with_cache(self, cached: Option<&'static str>) -> Self {
Self {
cache: Some(cached),
..self
}
}
pub fn with_failing_classifier(self) -> Self {
Self {
classifier_fails: true,
@ -241,6 +254,34 @@ impl PythonBinding for SyntheticBinding {
Ok(PyString::new(py, &response).into_any().unbind())
}
fn lookup(&mut self, py: Python<'_>) -> PyResult<Lookup<Synthetic>> {
let Some(cache) = self.cache else {
return Ok(Lookup::Miss);
};
self.log.push("lookup");
let Some(cached) = cache else {
return Ok(Lookup::Miss);
};
Ok(Lookup::Hit(Hit {
facts: ExecutionFacts {
provider: ProviderIdentity {
model: "m".into(),
provider: "p".into(),
},
source: ResultSource::Cache { key: "k".into() },
cost: Cost::CacheHit,
},
output: HitOutput::Response(PyString::new(py, cached).into_any().unbind()),
}))
}
fn store_response(&mut self, _: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
if self.cache.is_some() {
self.log.push(format!("store:{}", response.str()?));
}
Ok(())
}
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
self.log.push(format!("classify:{error}"));
if self.classifier_fails {
@ -321,6 +362,7 @@ pub enum HookScript {
FailAfterSuccess,
CancelTerminal,
FailTerminal,
ReviewResult,
}
pub struct SyntheticHooks {
@ -336,6 +378,21 @@ pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
}
impl CallHooks<PythonRuntime> for SyntheticHooks {
fn provider_result(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
if !matches!(self.script, HookScript::ReviewResult) {
return Ok(HookStep::Ready(response));
}
self.log
.push(format!("provider_result:{}", response.bind(py).str()?));
Ok(HookStep::Ready(
PyString::new(py, "reviewed").into_any().unbind(),
))
}
fn prepare_arguments(
&mut self,
_: Python<'_>,
@ -474,6 +531,18 @@ 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 _ = host;
Ok("provider".to_owned())
})
})
})
}
pub fn failing_machine() -> MachineFactory {
Box::new(move |_| {
CallMachine::new(None, |_| {

10
litellm-rust/crates/host/src/hooks.rs
View file

@ -44,6 +44,16 @@ pub trait CallHooks<R: HookRuntime>: Sized {
Ok(R::ready(wire))
}
/// A complete provider response before the host caches or delivers it. 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<'_>,

9
litellm-rust/crates/inference-chat/src/handler.rs
View file

@ -81,14 +81,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 {

17
litellm-rust/crates/inference-chat/src/lib.rs
View file

@ -47,6 +47,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,
@ -54,12 +64,7 @@ impl ChatCompletionsRoute {
} = options.into();
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run(
request.into(),
cache_options,
interceptors,
observers.as_ref(),
),
self.run(call, cache_options, interceptors, observers.as_ref()),
)
.await
}

12
litellm-rust/crates/inference-messages/tests/caching.rs
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 {

9
litellm-rust/crates/inference-responses/src/handler.rs
View file

@ -59,14 +59,11 @@ pub(super) async fn execute(
.cache
.clone()
.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 {

78
litellm-rust/crates/inference-responses/tests/caching.rs
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))
}
}
@ -184,9 +214,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 +258,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 {

4
litellm-rust/crates/inference/AGENTS.md
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 and materializes that identity only when it needs native key inputs. A Python-hosted call caches through the host around the machine and attaches no service here
- `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
- signed requests bypass caching until the signing identity has a stable cache representation

135
litellm-rust/crates/inference/src/caching.rs
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,7 @@ where
observers,
)
.await?;
if from_provider && let Some(session) = session {
if let Some(session) = session.filter(|_| from_provider) {
session.store_response::<P>(&response).await;
}
Ok(response)
@ -190,8 +194,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 +204,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,
@ -239,23 +240,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 +252,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 +265,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 +293,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;
}
@ -345,7 +327,8 @@ fn now() -> Duration {
.unwrap_or_default()
}
fn successful_stream(text: &str, terminal: &str) -> bool {
/// Whether captured SSE text is a whole stream that ended in `terminal` without an error event.
pub fn successful_stream(text: &str, terminal: &str) -> bool {
let mut pending = BytesMut::from(text.as_bytes());
let mut codec = litellm_framer::sse::SseCodec::default();
let mut complete = false;

134
litellm-rust/crates/inference/tests/caching.rs
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))
}
}
@ -565,9 +629,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 +645,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 +695,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)})) },

34
litellm-rust/crates/llms-types/src/formats/chat_completions/response.rs
View file

@ -75,14 +75,40 @@ pub struct ChatCompletionsChoice {
}
/// The normalized response handed back to the host.
///
/// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the
/// `ModelResponse` it already created, and echoing the provider's own id here
/// would change it. Pinned by `response_carries_no_id` in the Anthropic chat transformation tests.
#[macro_rules_attribute::apply(crate::wire_type)]
pub struct ChatCompletionsResponse {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
pub created: u64,
pub model: String,
pub choices: Vec<ChatCompletionsChoice>,
pub usage: ChatCompletionsUsage,
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::provider(None)]
#[case::reviewed(Some("reviewed-1".to_owned()))]
fn response_round_trip_keeps_an_optional_host_id(#[case] id: Option<String>) {
let response = ChatCompletionsResponse {
id: id.clone(),
created: 0,
model: "test".into(),
choices: Vec::new(),
usage: ChatCompletionsUsage::default(),
};
let value = serde_json::to_value(&response).unwrap();
assert_eq!(value.get("id"), id.as_ref().map(|id| json!(id)).as_ref());
assert_eq!(
serde_json::from_value::<ChatCompletionsResponse>(value).unwrap(),
response
);
}
}

1
litellm-rust/crates/llms/src/anthropic/chat/transformation.rs
View file

@ -150,6 +150,7 @@ impl BaseConfig for AnthropicConfig {
.collect();
Ok(ChatCompletionsResponse {
id: None,
created: unix_now(),
model: body.model,
choices: vec![ChatCompletionsChoice {

1
litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs
View file

@ -193,6 +193,7 @@ impl BaseConfig for AmazonConverseConfig {
};
Ok(ChatCompletionsResponse {
id: None,
created: unix_now(),
// Converse echoes no model id, so Python reports the requested one.
model: model.to_string(),

1
litellm-rust/crates/llms/src/openai_like/chat/transformation.rs
View file

@ -146,6 +146,7 @@ impl BaseConfig for OpenAILikeChatConfig {
let details = usage.and_then(|usage| usage.get("prompt_tokens_details"));
Ok(ChatCompletionsResponse {
id: None,
created: body
.get("created")
.and_then(Value::as_u64)

6
litellm-rust/crates/python-bridge/src/cache/AGENTS.md vendored
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 a Python-hosted native call 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, picks the entry format 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/binding.rs` owns `CachedBinding<H>`, which wraps a route's Python binding and answers the driver's lookup and store seams. 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/` serves the Python cache around the native machine, never from inside it. `native/` owns the experimental V2 native cache handles exposed to Python. Python-hosted inference does not attach them until the Rust cache is ready

4
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored
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::{CachedBinding, CachedStream};
use litellm_cache::Error;
use pyo3::{

4
litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md vendored
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

86
litellm-rust/crates/python-bridge/src/cache/native/v2.rs vendored
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)

10
litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md vendored
View file

@ -1,9 +1,9 @@
# Python cache delegation
# Python cache around native calls
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.
`CachedBinding` lets the configured Python cache answer a native call. The `host-python` driver asks the binding to look the call up after argument preparation, before the native machine starts, the way `@client` reads its cache before the provider call. A hit is delivered as the call's result without starting the machine, so no wire-request hook, credential resolution or provider-success hook runs. A miss starts the unchanged native route with no cache attached, and the driver hands the binding the public response the provider-success hook answered, or the provider stream the caller read to its end
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 `litellm.rust_bridge.host.cache`. Do not read `litellm.cache`, call cache methods directly, or re-derive which calls may read or write. Python owns key construction, read and write gates, lookup, storage and the stored entry format. A hit's complete response becomes the public response straight from the stored entry and a hook-reviewed response is stored as the Python object the caller receives, so neither passes through a native response type. A stream is kept as its SSE text and replayed through `CachedStream`, and it is stored only once it ended with the route's terminal event
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
Await asynchronous lookups through the existing 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 lookup that raises an ordinary exception is a miss, never a request failure, and a cancellation still ends the call
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
Traverse retained Python references for GC and release them 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

280
litellm-rust/crates/python-bridge/src/cache/python/binding.rs vendored Normal file
View file

@ -0,0 +1,280 @@
use std::convert::Infallible;
use litellm_host::{
interceptors::{Cost, ExecutionFacts, ProviderIdentity, ResultSource},
protocol::Protocol,
};
use litellm_host_python::{
Hit, HitOutput, InvokeError, Lookup, PythonBinding, PythonHostCalls, PythonOwned,
effective_py_args, from_py,
};
use litellm_inference::caching::successful_stream;
use pyo3::{
exceptions::{PyException, PyRuntimeError},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{IntoPyDict, PyDict},
};
const MODULE: &str = "litellm.rust_bridge.host.cache";
/// How a route's stream is kept in and replayed from the Python cache.
pub(crate) trait CachedStream: Protocol {
/// The event a whole stream ends with, or `None` when the route never caches a stream.
const TERMINAL_EVENT: Option<&'static str>;
fn replay(text: String) -> Option<(Self::StreamHead, Vec<Self::Chunk>)>;
fn bytes(chunk: &Self::Chunk) -> &[u8];
}
/// Wraps a route binding so the Python cache answers the call before the native machine
/// starts and keeps what the provider answered. Python owns the key, policy, entry format and
/// storage; this binding only carries values between the driver and
/// `litellm.rust_bridge.host.cache`.
pub(crate) struct CachedBinding<H> {
inner: H,
call_type: &'static str,
asynchronous: bool,
positional: Option<Py<PyDict>>,
access: Option<Py<PyAny>>,
captured: Vec<u8>,
}
impl<H> CachedBinding<H> {
pub(crate) fn new(
inner: H,
call_type: &'static str,
asynchronous: bool,
positional: Py<PyDict>,
) -> Self {
Self {
inner,
call_type,
asynchronous,
positional: Some(positional),
access: None,
captured: Vec::new(),
}
}
fn call(&self, py: Python<'_>, name: &str, extra: Option<Bound<'_, PyAny>>) -> PyResult<()> {
let Some(access) = &self.access else {
return Ok(());
};
let function = py.import(MODULE)?.getattr(name)?;
match extra {
Some(value) => function.call1((access, value, self.asynchronous))?,
None => function.call1((access, self.asynchronous))?,
};
Ok(())
}
}
type ErrorOf<H> = <<H as PythonBinding>::Protocol as Protocol>::Error;
impl<H> CachedBinding<H>
where
H: PythonBinding,
H::Protocol: CachedStream,
{
/// A lookup that failed in Python is a miss; a cancellation still ends the call.
fn found(
&mut self,
py: Python<'_>,
answer: PyResult<Py<PyAny>>,
) -> PyResult<Lookup<H::Protocol>> {
let hit = match answer {
Ok(hit) => hit,
Err(error) if error.is_instance_of::<PyException>(py) => return Ok(Lookup::Miss),
Err(error) => return Err(error),
};
let hit = hit.bind(py);
if hit.is_none() {
return Ok(Lookup::Miss);
}
let output = hit.getattr("output")?;
let output = match output.getattr("kind")?.extract::<String>()?.as_str() {
"Stream" => {
match <H::Protocol as CachedStream>::replay(output.getattr("value")?.extract()?) {
Some((head, chunks)) => HitOutput::Stream { head, chunks },
None => return Ok(Lookup::Miss),
}
}
_ => {
let dumped = output.call_method(
"model_dump",
(),
Some(&[("mode", "json")].into_py_dict(py)?),
)?;
HitOutput::Response(
self.inner
.public_response(py, &from_py(&dumped.get_item("value")?)?)?,
)
}
};
Ok(Lookup::Hit(Hit {
facts: ExecutionFacts {
provider: ProviderIdentity {
model: hit.getattr("model")?.extract()?,
provider: hit.getattr("provider")?.extract()?,
},
source: ResultSource::Cache {
key: hit.getattr("key")?.extract()?,
},
cost: Cost::CacheHit,
},
output,
}))
}
}
impl<H> PythonBinding for CachedBinding<H>
where
H: PythonBinding,
H::Protocol: CachedStream,
{
type Protocol = 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 positional = self
.positional
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("cache binding used after close"))?;
let keyed = effective_py_args(positional.bind(py), arguments)?;
let access = py
.import(MODULE)?
.getattr("access")?
.call1((keyed, self.call_type))?;
self.access = (!access.is_none()).then(|| access.unbind());
Ok(request)
}
fn lookup(&mut self, py: Python<'_>) -> PyResult<Lookup<Self::Protocol>> {
let Some(access) = &self.access else {
return Ok(Lookup::Miss);
};
let lookup = py
.import(MODULE)?
.getattr(if self.asynchronous {
"alookup"
} else {
"lookup"
})?
.call1((access,))
.map(Bound::unbind);
match lookup {
Ok(awaitable) if self.asynchronous => Ok(Lookup::Await(awaitable)),
answer => self.found(py, answer),
}
}
fn resume_lookup(
&mut self,
py: Python<'_>,
answer: PyResult<Py<PyAny>>,
) -> PyResult<Lookup<Self::Protocol>> {
self.found(py, answer)
}
fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
self.call(py, "store_response", Some(response.clone()))
}
fn store_stream(&mut self, py: Python<'_>) -> PyResult<()> {
let captured = std::mem::take(&mut self.captured);
let Some(terminal) = <H::Protocol as CachedStream>::TERMINAL_EVENT else {
return Ok(());
};
let Ok(text) = String::from_utf8(captured) else {
return Ok(());
};
if !successful_stream(&text, terminal) {
return Ok(());
}
self.call(py, "store_stream", Some(text.into_pyobject(py)?.into_any()))
}
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 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>> {
if self.access.is_some() {
self.captured
.extend_from_slice(<H::Protocol as CachedStream>::bytes(&chunk));
}
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> PythonHostCalls<H::Protocol> for CachedBinding<H>
where
H: PythonBinding,
H::Protocol: CachedStream + Protocol<HostCall = Infallible>,
{
fn handle_host_call(
&mut self,
_: Python<'_>,
call: Infallible,
) -> Result<(), InvokeError<ErrorOf<H>>> {
match call {}
}
}
impl<H: PythonOwned> PythonOwned for CachedBinding<H> {
fn close(&mut self, py: Python<'_>) {
self.inner.close(py);
self.positional = None;
self.access = None;
self.captured = Vec::new();
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.inner.traverse(visit)?;
if let Some(positional) = &self.positional {
visit.call(positional)?;
}
if let Some(access) = &self.access {
visit.call(access)?;
}
Ok(())
}
}

141
litellm-rust/crates/python-bridge/src/cache/python/host.rs vendored
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)
}
}

7
litellm-rust/crates/python-bridge/src/cache/python/mod.rs vendored
View file

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

107
litellm-rust/crates/python-bridge/src/cache/python/service.rs vendored
View file

@ -1,107 +0,0 @@
use std::{future::Future, pin::Pin, time::Duration};
use litellm_cache::Error;
use litellm_cache_response::{
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope,
};
use litellm_host::{
machine::{HostServices, MachineFault},
protocol::{Protocol, Reply},
};
use litellm_inference::caching::CachedOutput;
use serde_json::Value;
const STREAM_EVENTS_KEY: &str = "litellm_cached_anthropic_sse_events";
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<_>>()
})),
}
}
pub(crate) enum CacheCall {
Lookup {
reply: Reply<Result<Option<Value>, Error>>,
},
Store {
value: Value,
reply: Reply<Result<(), Error>>,
},
}
struct PythonCacheService<P: Protocol> {
services: HostServices<P>,
config: ResponseCacheConfig,
}
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()
},
})
}
impl<P: Protocol<HostCall = CacheCall>> ResponseCacheService for PythonCacheService<P>
where
P::Error: From<MachineFault>,
{
fn config(&self) -> &ResponseCacheConfig {
&self.config
}
fn lookup<'a>(
&'a self,
_: &'a ResponseCacheRequest,
_: Duration,
) -> Pin<Box<dyn Future<Output = Result<Option<Value>, Error>> + Send + 'a>> {
Box::pin(async move {
self.services
.call(|reply| CacheCall::Lookup { reply })
.await
.map_err(|_| Error::Unavailable)??
.map(from_python)
.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)?;
self.services
.call(|reply| CacheCall::Store { value, reply })
.await
.map_err(|_| Error::Unavailable)?
})
}
}

171
litellm-rust/crates/python-bridge/src/cache/selection.rs vendored
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,
})
}

8
litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs
View file

@ -59,6 +59,14 @@ 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 encode_stream_head(
&mut self,
_py: Python<'_>,

40
litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs
View file

@ -4,10 +4,15 @@ use std::sync::Arc;
use host::ChatCompletionsPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, 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::*;
@ -57,15 +62,26 @@ impl InferenceRoute for ChatCompletionsRoute {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
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<ChatCompletions>,
options: CallOptions,
) -> Result<OutputOf<ChatCompletions>, RouteError> {
self.execute_call(request, &interceptors, options)
.await
.map(CallOutput::Complete)
}
}
impl crate::cache::CachedStream for ChatCompletions {
const TERMINAL_EVENT: Option<&'static str> = None;
fn replay(_: String) -> Option<(Self::StreamHead, Vec<Self::Chunk>)> {
None
}
fn bytes(chunk: &Self::Chunk) -> &[u8] {
match *chunk {}
}
}

53
litellm-rust/crates/python-bridge/src/routes/inference.rs
View file

@ -1,13 +1,16 @@
use std::sync::Arc;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, CacheScope, 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};
use litellm_secrets::source::SecretSource;
use pyo3::{
exceptions::PyValueError,
@ -20,6 +23,7 @@ use serde_json::{Map, Value};
use super::{NativeCall, RequestView};
use crate::{
cache::{CachedBinding, CachedStream},
errors::{RustUpstreamError, route_error_to_pyerr},
marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds},
};
@ -143,8 +147,8 @@ impl InferenceHost {
}
}
pub(super) trait InferenceRoute: Sized + 'static {
type Protocol: Protocol<Error = RouteError>;
pub(super) trait InferenceRoute: Sized + Send + 'static {
type Protocol: CachedStream + Protocol<Error = RouteError, HostCall = std::convert::Infallible>;
const OPERATION: LoggingOperation;
const SYNC_CALL_TYPE: &'static str;
const ASYNC_CALL_TYPE: &'static str;
@ -154,12 +158,12 @@ pub(super) trait InferenceRoute: Sized + 'static {
auth: Arc<AuthServices>,
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<Self::Protocol>,
options: CallOptions,
) -> impl Future<Output = Result<OutputOf<Self::Protocol>, RouteError>> + Send;
}
pub(super) fn run_inference<R, H>(
@ -170,13 +174,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 +193,24 @@ 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, _, interceptors, observers| async move {
route
.call(
call,
interceptors,
CallOptions {
cache: None,
observers,
},
)
.await
},
))
},
host,
CachedBinding::new(host, call_type, asynchronous, positional),
hooks,
asynchronous,
)

60
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
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,14 @@ 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 encode_stream_head(
&mut self,
py: Python<'_>,
@ -289,43 +288,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)
}
}

96
litellm-rust/crates/python-bridge/src/routes/messages/mod.rs
View file

@ -1,53 +1,24 @@
mod host;
use std::sync::Arc;
use host::MessagesPythonHost;
use litellm_auth::AuthServices;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::OutputOf, machine::ChannelInterceptors};
use litellm_inference::{CallOptions, RouteError, caching::StreamCachable};
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},
};
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 +30,44 @@ 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)
}
async fn call(
self,
request: MessagesCall,
interceptors: ChannelInterceptors<Messages>,
options: CallOptions,
) -> Result<OutputOf<Messages>, RouteError> {
self.execute(request, &interceptors, options).await
}
}
impl crate::cache::CachedStream for Messages {
const TERMINAL_EVENT: Option<&'static str> = Some(<Self as StreamCachable>::TERMINAL_EVENT);
fn replay(text: String) -> Option<(Self::StreamHead, Vec<Self::Chunk>)> {
Some((
litellm_inference_messages::route::MessagesStreamHead {
headers: Vec::new(),
},
vec![text.into()],
))
}
fn bytes(chunk: &Self::Chunk) -> &[u8] {
<Self as StreamCachable>::bytes(chunk)
}
}

9
litellm-rust/crates/python-bridge/src/routes/mod.rs
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)

8
litellm-rust/crates/python-bridge/src/routes/ocr/host.rs
View file

@ -118,6 +118,14 @@ 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 encode_stream_head(
&mut self,
_: Python<'_>,

8
litellm-rust/crates/python-bridge/src/routes/responses/host.rs
View file

@ -135,6 +135,14 @@ impl PythonBinding for ResponsesPythonHost {
self.host.response(py, &response)
}
fn public_response(
&mut self,
py: Python<'_>,
response: &serde_json::Value,
) -> PyResult<Py<PyAny>> {
self.host.response(py, response)
}
fn encode_stream_head(
&mut self,
_py: Python<'_>,

38
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
View file

@ -5,10 +5,10 @@ use std::sync::Arc;
use host::ResponsesPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, 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, caching::StreamCachable};
use litellm_inference_responses::{ResponsesRoute, route::Responses, types::ResponsesCall};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
pub(crate) use websocket::ResponsesWebSocketConnection;
@ -49,15 +49,29 @@ impl InferenceRoute for ResponsesRoute {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
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<Responses>,
options: CallOptions,
) -> Result<OutputOf<Responses>, RouteError> {
self.execute(request, &interceptors, options).await
}
}
impl crate::cache::CachedStream for Responses {
const TERMINAL_EVENT: Option<&'static str> = Some(<Self as StreamCachable>::TERMINAL_EVENT);
fn replay(text: String) -> Option<(Self::StreamHead, Vec<Self::Chunk>)> {
Some((
litellm_inference_responses::types::ResponsesStreamHead {
headers: Vec::new(),
},
vec![text.into()],
))
}
fn bytes(chunk: &Self::Chunk) -> &[u8] {
<Self as StreamCachable>::bytes(chunk)
}
}

8
litellm/caching/caching.py
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

107
litellm/caching/caching_handler.py
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]) -> Mapping[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,
@ -937,7 +970,7 @@ class LLMCachingHandler:
def _forget_worker_copy(self) -> None:
"""Drop this worker's memory copy of a stored response without output, so the next read reaches Redis,
where the refill lands: the text completion and messages writers skip the memory tier."""
if self.dual_cache is None or self.preset_cache_key is None:
if self.dual_cache is None or not isinstance(self.preset_cache_key, str):
return
self.dual_cache.in_memory_cache.delete_cache(self.preset_cache_key)
@ -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:
"""

1
litellm/chat_completions/dispatch.py
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,

1
litellm/messages/dispatch.py
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,

1
litellm/responses/dispatch.py
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,

6
litellm/rust_bridge/catalog.py
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),

5
litellm/rust_bridge/detectors.py
View file

@ -471,11 +471,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"})

4
litellm/rust_bridge/host/AGENTS.md
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

252
litellm/rust_bridge/host/cache.py
View file

@ -1,12 +1,258 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import re
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, 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
_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) -> Sequence[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 | None:
events: Final = _stream_events(entry)
if events is None:
return None if stream else ResponseOutput(value=entry)
return StreamOutput(value="".join(events)) if stream else None
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:
output: CachedOutput
key: str
model: str
provider: str
def _identity(request: Mapping[str, object]) -> tuple[str, str]:
from litellm import get_llm_provider
model: Final = request.get("model")
custom_llm_provider: Final = request.get("custom_llm_provider")
name: Final = model if isinstance(model, str) else ""
fallback: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else ""
try:
resolved, provider, _, _ = get_llm_provider(model=name, custom_llm_provider=fallback or None)
except Exception: # noqa: BLE001 # an unresolvable model still reports the hit under its own name
return name, fallback
return resolved, provider
def _hit(access: CacheAccess, entry: object) -> CacheHit | None:
output: Final = access.format.output(entry, access.request.get("stream") is True)
if output is None:
return None
model, provider = _identity(access.request)
return CacheHit(
output=output,
key=str(access.key),
model=model,
provider=provider,
)
def lookup(access: CacheAccess) -> CacheHit | None:
from litellm.caching.caching_handler import is_response_without_output
if not access.read:
return None
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
if not access.read:
return None
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 _store(access: CacheAccess, output: CachedOutput, asynchronous: bool) -> None:
from litellm._logging import verbose_logger
from litellm.caching.caching_handler import create_cache_write_task, is_response_without_output
if not access.write:
return
entry: Final = access.format.entry(output)
if entry is None or is_response_without_output(entry):
return
try:
if not asynchronous:
access.cache.add_cache(entry, cache_key=access.key, **access.request)
return
tier: Final = access.memory_tier if access.format.tiered(output) else None
create_cache_write_task(
lambda: access.cache.async_add_cache(
entry, dynamic_cache_object=tier, cache_key=access.key, **access.request
)
)
except Exception: # noqa: BLE001 # a cache failure never fails the request
verbose_logger.exception("LiteLLM Cache: storing the response failed")
def store_response(access: CacheAccess, response: object, asynchronous: bool) -> None:
_store(access, ResponseOutput(value=response), asynchronous)
def store_stream(access: CacheAccess, text: str, asynchronous: bool) -> None:
_store(access, StreamOutput(value=text), asynchronous)

1
tests/test_litellm_rust/cache/__init__.py vendored
View file

@ -1 +0,0 @@

19
tests/test_litellm_rust/cache/conftest.py vendored
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)

121
tests/test_litellm_rust/cache/support.py vendored
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

534
tests/test_litellm_rust/cache/test_python_cache.py vendored
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

0
tests/test_litellm_rust/cache/v2/__init__.py vendored
View file

207
tests/test_litellm_rust/cache/v2/test_cache.py vendored
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()

10
tests/test_litellm_rust/chat_completions/conftest.py
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)

369
tests/test_litellm_rust/chat_completions/test_cache.py Normal file
View file

@ -0,0 +1,369 @@
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.utils import LLMResponseTypes, 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.isolation import rebound
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"]
class _Review(CustomLogger):
calls = 0
async def async_post_call_success_deployment_hook(
self, request_data: dict[str, object], response: object, call_type: object
) -> LLMResponseTypes | None:
self.calls += 1
assert isinstance(response, ModelResponse)
return response.model_copy(update={"id": "reviewed-1", "system_fingerprint": "reviewed-fingerprint"})
def _reviewed(response: object) -> tuple[object, object]:
assert isinstance(response, ModelResponse)
return response.id, response.system_fingerprint
async def test_the_deployment_hooks_response_reaches_the_caller_whole(call: Call, server: RecordingServer) -> None:
logger: Final = _Review()
with rebound(litellm, "callbacks", [logger]):
response: Final = await call()
assert _reviewed(response) == ("reviewed-1", "reviewed-fingerprint")
assert logger.calls == 1
assert_served_natively(server, response, count=1)
@pytest.mark.usefixtures("cache")
async def test_a_cached_chat_replays_the_deployment_hooks_response_whole(call: Call, server: RecordingServer) -> None:
logger: Final = _Review()
with rebound(litellm, "callbacks", [logger]):
first: Final = await call()
second: Final = await call()
assert _reviewed(first) == _reviewed(second) == ("reviewed-1", "reviewed-fingerprint")
assert logger.calls == 1
assert_served_natively(server, first, count=1)
@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)

2
tests/test_litellm_rust/chat_completions/test_dispatch.py
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}),

9
tests/test_litellm_rust/conftest.py
View file

@ -6,6 +6,15 @@ import pytest
import pytest_asyncio
import litellm
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.clickhouse import clickhouse_url as clickhouse_url
from tests.test_litellm_rust.support.isolation import isolated_test_state
from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service

11
tests/test_litellm_rust/messages/conftest.py
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)

549
tests/test_litellm_rust/messages/test_cache.py Normal file
View file

@ -0,0 +1,549 @@
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("stored_stream", (False, True), ids=("response-to-stream", "stream-to-response"))
@pytest.mark.usefixtures("cache")
async def test_a_shared_key_with_the_other_response_kind_is_a_miss(
call: Call, server: RecordingServer, stored_stream: bool
) -> None:
server.expect(2)
server.serve(STREAM if stored_stream else ResponseSpec(body=MESSAGES_RESPONSE))
first: Final = await call(stream=stored_stream, cache_key="shared-entry")
if stored_stream:
await collect(first)
server.serve(ResponseSpec(body=MESSAGES_RESPONSE) if stored_stream else STREAM)
second: Final = await call(stream=not stored_stream, cache_key="shared-entry")
assert cache_key(second) is None
if stored_stream:
assert body(second) == MESSAGES_RESPONSE
else:
assert await collect(second) == b"".join(STREAM.payloads())
third: Final = await call(stream=not stored_stream, cache_key="shared-entry")
assert cache_key(third) == "shared-entry"
if stored_stream:
assert body(third) == MESSAGES_RESPONSE
else:
assert await collect(third) == b"".join(STREAM.payloads())
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)
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
@pytest.mark.usefixtures("cache")
async def test_a_hit_is_answered_before_the_provider_request_is_built(
call: Call, server: RecordingServer, stream: bool
) -> None:
if stream:
server.serve(STREAM)
async def send() -> RecordingLogger:
recorder: Final = RecordingLogger()
response: Final = await call(stream=stream, callbacks=[recorder])
if stream:
await collect(response)
return recorder
miss: Final = await send()
hit: Final = await send()
await drain_logging()
assert "log_pre_api_call" in miss.names
assert "log_pre_api_call" not in hit.names
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

2
tests/test_litellm_rust/messages/test_dispatch.py
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}),
(

11
tests/test_litellm_rust/responses/conftest.py
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)

375
tests/test_litellm_rust/responses/test_cache.py Normal file
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

2
tests/test_litellm_rust/responses/test_dispatch.py
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,
{

164
tests/test_litellm_rust/support/caching.py Normal file
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

6
tests/test_litellm_rust/support/recording_server.py
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:

15
tests/test_litellm_rust/support/results.py
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")

17
tests/unit/caching/test_caching_handler.py
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"

358
tests/unit/rust_bridge/host/test_cache.py
View file

@ -1,21 +1,357 @@
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,
lookup,
store_response,
store_stream,
)
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 keep(cache_access: CacheAccess, output: ResponseOutput | StreamOutput, asynchronous: bool = False) -> None:
match output:
case ResponseOutput(value=response):
store_response(cache_access, response, asynchronous)
case StreamOutput(value=text):
store_stream(cache_access, text, asynchronous)
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.output.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_response(granted, {"id": "cached"}, False)
hit: Final = lookup(granted)
assert hit is not None
assert hit.output == ResponseOutput(value={"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_response(cache_access, {"id": "cached"}, False)
hit: Final = lookup(cache_access)
assert hit is not None
assert hit.output == ResponseOutput(value={"id": "cached"})
assert hit.key == expected_key
assert (hit.model, hit.provider) == ("gpt-4o", "openai")
def test_lookup_skips_responses_without_output() -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
cache_access: Final = _cache_access(cache)
store_response(cache_access, EMPTY, False)
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.output == ResponseOutput(value={"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)
store_response(cache_access, {"id": "cached"}, True)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
hit: Final = await alookup(cache_access)
assert hit is not None
assert hit.output == ResponseOutput(value={"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)
store_response(empty_access, EMPTY, True)
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)
keep(cache_access, output, asynchronous=True)
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
store_response(redis_access, {"id": "from-the-worker"}, True)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
hit: Final = await alookup(redis_access)
assert hit is not None
assert hit.output == ResponseOutput(value={"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)
store_stream(cache_access, "".join(MESSAGES_EVENTS), False)
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.output == StreamOutput(value="".join(MESSAGES_EVENTS))
def test_a_stream_on_a_route_without_a_stream_entry_is_not_stored() -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
store_stream(_cache_access(cache, STREAM_REQUEST), "".join(MESSAGES_EVENTS), False)
assert stored_entry(cache, STREAM_REQUEST) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
@pytest.mark.parametrize("stored_stream", (False, True), ids=("response-to-stream", "stream-to-response"))
async def test_messages_rejects_a_cached_output_of_the_other_kind(asynchronous: bool, stored_stream: bool) -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
request: Final = {**REQUEST, "cache_key": "shared-entry", "stream": stored_stream}
writer: Final = _cache_access(cache, request, entry_format=MESSAGES_ENTRY)
reader: Final = _cache_access(cache, {**request, "stream": not stored_stream}, entry_format=MESSAGES_ENTRY)
output: Final = (
StreamOutput(value="".join(MESSAGES_EVENTS)) if stored_stream else ResponseOutput(value={"id": "cached"})
)
keep(writer, output)
assert (await alookup(writer) if asynchronous else lookup(writer)) is not None
assert (await alookup(reader) if asynchronous else lookup(reader)) 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

17
tests/unit/rust_bridge/test_detectors.py
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)),

1
tests/unit/test_internal_context.py
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",
}
)