scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

feat(messages): serve response cache reads on rust

Browse source
This commit is contained in:
narmstrong-berry 2026-11-10 20:28:06 +00:00
parent 7728b426ce
commit 3262029c99
No known key found for this signature in database
31 changed files with 1865 additions and 1353 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -4158,6 +4158,7 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",

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

@ -27,6 +27,10 @@ impl Default for ResponseCacheConfig {
pub trait ResponseCacheService: Send + Sync {
fn config(&self) -> &ResponseCacheConfig;
fn keyed_by_caller(&self) -> bool {
false
}
fn lookup<'a>(
&'a self,
request: &'a ResponseCacheRequest,

12
litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

@ -11,7 +11,8 @@
"logger",
"kwargs",
"start_time",
"end_time"
"end_time",
"cache_hit"
],
"update_logging": [
"logger",
@ -95,6 +96,12 @@
"start",
"end"
],
"cache_hit_success": [
"logger",
"response",
"start",
"end"
],
"enqueue_logging": [
"coroutine"
],
@ -126,7 +133,8 @@
"logger",
"kwargs",
"start_time",
"end_time"
"end_time",
"cache_hit"
],
"stream_success": [
"logger",

46
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -273,6 +273,7 @@ impl LegacyLogging {
self.call.kwargs(),
&self.start,
&self.end,
self.messages_cache_hit(),
)?;
self.response
.as_ref()
@ -280,6 +281,25 @@ impl LegacyLogging {
.ok_or_else(|| HostError::MissingState.into())
}
fn messages_cache_hit(&self) -> bool {
self.operation == LoggingOperation::Messages && self.cache_key.is_some()
}
fn dispatch_cache_hit(&self, py: Python<'_>) -> PyResult<()> {
let logger = self.logger()?;
let logged = Logging::CacheHitSuccess.call(
py,
(logger.object(py), &self.response, &self.start, &self.end),
);
match logged {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(logger.object(py)));
Ok(())
}
result => result.map(|_| ()),
}
}
fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
match self.try_dispatch_success(py) {
Err(error) if error.is_instance_of::<PyException>(py) => {
@ -576,7 +596,7 @@ impl LegacyLogging {
) -> PyResult<HookStep<Self, Py<PyAny>>> {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response);
if self.runs_deployment_hooks() {
if self.runs_deployment_hooks() && !self.messages_cache_hit() {
return Ok(HookStep::Await(
DeploymentHooks::after_success(
py,
@ -597,7 +617,17 @@ impl LegacyLogging {
) -> PyResult<HookStep<Self, ()>> {
use litellm_host::interceptors::{Cost, ResultSource};
self.cache_key = match &facts.source {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
};
let logger = self.logger()?.object(py);
if self.messages_cache_hit() {
logger
.getattr("model_call_details")?
.set_item("cache_hit", true)?;
return Ok(HookStep::Ready(()));
}
let params = logger
.getattr("litellm_params")?
.cast_into::<PyDict>()?
@ -625,10 +655,6 @@ impl LegacyLogging {
};
logger.setattr("_provider_reported_cost", reported_cost)?;
details.call_method1("pop", ("response_cost", py.None()))?;
self.cache_key = match &facts.source {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
};
if self.cache_key.is_some() {
details.set_item("cache_hit", true)?;
} else {
@ -717,6 +743,7 @@ impl LegacyLogging {
self.response = Some(response.clone_ref(py));
match &self.stream {
Some(stream) => self.stream_success(py, stream, false)?,
None if self.messages_cache_hit() => self.dispatch_cache_hit(py)?,
None => self.dispatch_success(py)?,
}
Ok(HookStep::Ready(()))
@ -751,7 +778,11 @@ impl LegacyLogging {
if self.stream_billing().is_none() {
return Err(HostError::MissingState.into());
}
if let Some(key) = &self.cache_key {
if let Some(key) = self
.cache_key
.as_ref()
.filter(|_| !self.messages_cache_hit())
{
head.bind(py).set_item("cache_key", key)?;
head.bind(py).set_item("cache_hit", true)?;
}
@ -766,6 +797,7 @@ impl LegacyLogging {
self.call.kwargs(),
&self.start,
datetime(py, epoch_seconds())?,
self.messages_cache_hit(),
),
)?;
}
@ -790,7 +822,7 @@ impl LegacyLogging {
// A caller that goes away mid-stream cancels the call before it ends: Python's Messages
// stream iterator still bills the chunks it read, straight away rather than deferred.
fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> {
if !self.asynchronous || self.end.is_some() {
if !self.asynchronous || self.end.is_some() || self.messages_cache_hit() {
return Ok(());
}
let Some(stream) = self.stream.take() else {

6
litellm-rust/crates/callbacks-legacy-python/src/logger.rs
View file

@ -77,8 +77,12 @@ pub fn finalize(
kwargs: &Py<PyDict>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
cache_hit: bool,
) -> PyResult<()> {
Wrapper::Finalize.call(py, (response, logger.object(py), kwargs, start, end))?;
Wrapper::Finalize.call(
py,
(response, logger.object(py), kwargs, start, end, cache_hit),
)?;
Ok(())
}

2
litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -64,6 +64,8 @@ pub(crate) enum Logging {
SubmitSuccess,
#[strum(serialize = "async_success_handler")]
AsyncSuccessHandler,
#[strum(serialize = "cache_hit_success")]
CacheHitSuccess,
#[strum(serialize = "mark_async_success_enqueued")]
MarkAsyncSuccessEnqueued,
#[strum(serialize = "enqueue_logging")]

5
litellm-rust/crates/callbacks-legacy-python/src/test_support.rs
View file

@ -53,7 +53,7 @@ def setup(call_type, args, kwargs, start, asynchronous):
FAKES = {
'setup': setup,
'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response),
'finalize': lambda response, logger, kwargs, start, end, cache_hit: logger.record('finalize', response),
'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs(
kwargs=kwargs,
model=model,
@ -90,6 +90,7 @@ FAKES = {
)(error, ''.join(traceback.format_exception(error)), start, end),
'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)),
'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end),
'cache_hit_success': lambda logger, response, start, end: logger.record('cache_hit_success', response),
'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None),
'enqueue_logging': lambda coroutine: coroutine.enqueue(),
'restore_context': lambda logger: logger.record('restore', None),
@ -101,7 +102,7 @@ FAKES = {
),
'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
'stream_opened': lambda logger, head: logger.record('stream_opened', head),
'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head),
'stamp_stream_head': lambda head, logger, kwargs, start, end, cache_hit: logger.record('stamp_stream_head', head),
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record(
'stream_success', list(chunks)
),

1
litellm-rust/crates/inference-messages/Cargo.toml
View file

@ -30,6 +30,7 @@ tracing.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-framer = { workspace = true, features = ["sse"] }
litellm-cache.workspace = true
litellm-cache-memory.workspace = true
litellm-host-native.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }

81
litellm-rust/crates/inference-messages/src/handler.rs
View file

@ -35,7 +35,11 @@ use super::{
route::Messages,
types::invalid_request,
};
use litellm_inference::{caching::CallCache, context::CallContext, outbound::outbound_request};
use litellm_inference::{
caching::{CacheRequest, CallCache},
context::CallContext,
outbound::outbound_request,
};
struct ProviderCall {
identity: ProviderIdentity,
@ -54,6 +58,42 @@ impl ProviderCall {
}
}
pub(super) async fn cached(
route: &MessagesRoute,
provider: MessagesProvider,
model: &str,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<Option<MessagesCallResponse>, Error> {
let Some(cache) = caller_keyed(route) else {
return Ok(None);
};
let identity = ProviderIdentity {
model: model.to_owned(),
provider: provider.as_str().to_string(),
};
let Some((output, source)) = CallCache::<Messages>::new(
Some(cache),
context.cache,
CacheRequest::from_wire(identity.clone(), None),
)
.lookup()
.await
else {
return Ok(None);
};
context
.result_ready(execution_facts(provider, identity, &output, &source))
.await?;
Ok(Some(output))
}
fn caller_keyed(route: &MessagesRoute) -> Option<&litellm_cache_response::ScopedCache> {
route
.cache
.as_ref()
.filter(|cache| cache.service.keyed_by_caller())
}
pub(super) async fn execute(
route: &MessagesRoute,
provider: MessagesProvider,
@ -62,6 +102,7 @@ pub(super) async fn execute(
secrets: Secrets,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
let caller_keyed = caller_keyed(route).is_some();
let request = provider_call(
&route.auth,
provider,
@ -73,34 +114,50 @@ pub(super) async fn execute(
.boxed()
.await?;
let cache = CallCache::<Messages>::from_wire(
route.cache.as_ref().filter(|_| request.cacheable()),
route
.cache
.as_ref()
.filter(|_| caller_keyed || request.cacheable()),
context.cache,
&request.identity,
&request.wire,
);
let identity = request.identity.clone();
let (output, source) = match cache.lookup().await {
let hit = match caller_keyed {
true => None,
false => cache.lookup().await,
};
let (output, source) = match hit {
Some(hit) => hit,
None => (
call_provider(&route.http, request, context).await?,
ResultSource::Provider,
),
};
let reported = match &output {
context
.result_ready(execution_facts(provider, identity, &output, &source))
.await?;
Ok(cache.finish(output, &source).await)
}
fn execution_facts(
provider: MessagesProvider,
identity: ProviderIdentity,
output: &MessagesCallResponse,
source: &ResultSource,
) -> ExecutionFacts {
let reported = match output {
MessagesCallResponse::Complete(response) => match response {
MessagesBody::Message(message) => provider.config().reported_cost(message),
MessagesBody::Unchecked(_) => None,
},
MessagesCallResponse::Stream { .. } => None,
};
context
.result_ready(ExecutionFacts {
provider: identity,
cost: Cost::settle(&source, reported),
source: source.clone(),
})
.await?;
Ok(cache.finish(output, &source).await)
ExecutionFacts {
provider: identity,
cost: Cost::settle(source, reported),
source: source.clone(),
}
}
async fn provider_call(

3
litellm-rust/crates/inference-messages/src/lib.rs
View file

@ -94,6 +94,9 @@ impl MessagesRoute {
<&'static str>::from(provider.provider),
);
let model = provider.model.to_owned();
if let Some(output) = handler::cached(self, config, &model, &context).await? {
return Ok(output);
}
let secrets = self.secrets.resolve(config.config().secret_names()).await?;
handler::execute(self, config, model, call, secrets, &context)
.boxed()

86
litellm-rust/crates/inference-messages/tests/caching.rs
View file

@ -301,3 +301,89 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
assert!(matches!(cached.source, ResultSource::Cache { .. }));
assert_eq!(cached.cost, Cost::CacheHit);
}
#[derive(Default)]
struct CallerKeyed {
config: ResponseCacheConfig,
entry: std::sync::Mutex<Option<Value>>,
}
impl ResponseCacheService for CallerKeyed {
fn config(&self) -> &ResponseCacheConfig {
&self.config
}
fn keyed_by_caller(&self) -> bool {
true
}
fn lookup<'a>(
&'a self,
_: &'a litellm_cache_response::ResponseCacheRequest,
_: Duration,
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<Option<Value>, litellm_cache::Error>>
+ Send
+ 'a,
>,
> {
Box::pin(async move { Ok(self.entry.lock().unwrap().clone()) })
}
fn store<'a>(
&'a self,
_: &'a litellm_cache_response::ResponseCacheRequest,
value: Value,
_: Duration,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<(), litellm_cache::Error>> + Send + 'a>,
> {
Box::pin(async move {
*self.entry.lock().unwrap() = Some(value);
Ok(())
})
}
}
#[tokio::test]
async fn a_caller_keyed_hit_is_served_before_the_request_is_built_or_intercepted() {
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
let response = json!({"id": "msg_caller", "type": "message", "role": "assistant", "model": "native-test",
"content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null,
"usage": {"input_tokens": 3, "output_tokens": 2}});
let upstream = support::upstream([support::json_response(response)]).await;
let hooks = ChangingHooks::default();
let route = support::messages_route(litellm_inference_testing::no_secrets()).with_cache(
ScopedCache::new(Arc::new(CallerKeyed::default()), CacheScope::Shared),
);
let request = |api_key: Option<&str>| {
MessagesCall {
body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})).unwrap(),
api_key: api_key.map(Into::into), api_base: Some(upstream.uri()), custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(), extra_headers: None, provider_specific_header: None,
timeout: None, shaping: Default::default(),
}
};
let first = route
.execute(request(Some("key")), &hooks, None)
.await
.unwrap();
let second = route.execute(request(None), &hooks, None).await.unwrap();
let (MessagesCallResponse::Complete(first), MessagesCallResponse::Complete(second)) =
(first, second)
else {
panic!("expected Messages responses");
};
assert_eq!(first, second);
assert_eq!(support::received(&upstream).await.len(), 1);
assert_eq!(hooks.calls.load(Ordering::SeqCst), 1);
let facts = hooks.facts.lock().unwrap();
let [provider, cached] = facts.as_slice() else {
panic!("expected provider and cache results");
};
assert_eq!(provider.source, ResultSource::Provider);
assert!(matches!(cached.source, ResultSource::Cache { .. }));
assert_eq!(cached.cost, Cost::CacheHit);
}

1
litellm-rust/crates/inference/AGENTS.md
View file

@ -81,3 +81,4 @@
- Cache lookup follows provider preparation, credential resolution, and the request interceptor
- keys describe the effective provider URL, authenticated headers, and rewritten body
- signed requests bypass caching until the signing identity has a stable cache representation
- A service that keys entries by the caller's arguments (`keyed_by_caller`, the Python host's legacy cache) is looked up before provider preparation, as Python's `@client` wrapper does, and a hit skips request building and the request interceptor

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

@ -238,12 +238,19 @@ impl<P: StreamCachable> CallCache<P> {
identity: &ProviderIdentity,
wire: &WireRequest,
) -> Self {
Self::new(
cache,
policy,
CacheRequest::from_wire(identity.clone(), Some(wire)),
)
}
pub fn new(cache: Option<&ScopedCache>, policy: CachePolicy, request: CacheRequest) -> 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,

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

@ -1,6 +1,6 @@
# Python cache delegation
This directory lets Rust inference use a selected Python cache. `service.rs` implements the injected Rust response-cache service and yields typed cache operations. `host.rs` calls the Python cache's sync or async API and delivers the result back to Rust.
This directory lets Rust inference use a selected Python cache. `service.rs` implements the injected Rust response-cache service and yields typed cache operations. `host.rs` delivers each operation's result back to Rust. Asynchronous lookups and stream stores go through `litellm.rust_bridge.host.cache`, which reads with the `LLMCachingHandler` the `@client` wrapper uses, for the legacy function the route names, and writes with its dual cache. `configure` disables the cache for synchronous calls.
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

25
litellm-rust/crates/python-bridge/src/cache/python/host.rs vendored
View file

@ -18,6 +18,7 @@ enum Pending {
pub(crate) struct PythonCache {
cache: Option<Py<PyAny>>,
function: Option<Py<PyAny>>,
arguments: Option<Py<PyDict>>,
pending: Option<Pending>,
asynchronous: bool,
@ -28,6 +29,7 @@ impl PythonCache {
pub fn new(asynchronous: bool) -> Self {
Self {
cache: None,
function: None,
arguments: None,
pending: None,
asynchronous,
@ -35,12 +37,18 @@ impl PythonCache {
}
}
pub(in crate::cache) fn asynchronous(&self) -> bool {
self.asynchronous
}
pub(in crate::cache) fn bind(
&mut self,
cache: Bound<'_, PyAny>,
function: &Bound<'_, PyAny>,
arguments: &Bound<'_, PyDict>,
) {
self.cache = Some(cache.unbind());
self.function = Some(function.clone().unbind());
self.arguments = Some(arguments.clone().unbind());
}
@ -59,7 +67,7 @@ impl PythonCache {
self.pending = Some(Pending::Lookup(reply));
(
if self.asynchronous {
"async_get_cache"
"lookup"
} else {
"get_cache"
},
@ -77,7 +85,7 @@ impl PythonCache {
self.pending = Some(Pending::Store(reply));
(
if self.asynchronous {
"async_add_cache"
"store_stream"
} else {
"add_cache"
},
@ -86,9 +94,20 @@ impl PythonCache {
}
};
let result = match result {
Some(value) if self.asynchronous => py
.import(CACHE_HOST_MODULE)?
.getattr(method)?
.call((cache, value), Some(&arguments)),
Some(value) => cache
.bind(py)
.call_method(method, (value,), Some(&arguments)),
None if self.asynchronous => py.import(CACHE_HOST_MODULE)?.getattr(method)?.call(
(
cache,
self.function.as_ref().ok_or(HandleError::Unconfigured)?,
),
Some(&arguments),
),
None => cache.bind(py).call_method(method, (), Some(&arguments)),
}
.map(Bound::unbind);
@ -148,11 +167,13 @@ impl PythonCache {
self.pending = None;
self.deferred = false;
self.cache = None;
self.function = None;
self.arguments = None;
}
pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.cache)?;
visit.call(&self.function)?;
visit.call(&self.arguments)
}
}

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

@ -79,6 +79,10 @@ where
&self.config
}
fn keyed_by_caller(&self) -> bool {
true
}
fn lookup<'a>(
&'a self,
_: &'a ResponseCacheRequest,

32
litellm-rust/crates/python-bridge/src/cache/selection.rs vendored
View file

@ -51,15 +51,21 @@ fn selected_cache<'py>(
kwargs: &Bound<'py, PyDict>,
call_type: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
if kwargs
.get_item("caching")?
.is_some_and(|value| value.is(pyo3::types::PyBool::new(py, false)))
{
return Ok(None);
}
configured_cache(py, call_type)
}
fn configured_cache<'py>(py: Python<'py>, 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)))
{
if configured.is_none() {
return Ok(None);
}
let supported = configured.getattr("supported_call_types")?;
@ -90,9 +96,14 @@ pub(crate) fn configure(
python: &mut python::PythonCache,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
call_type: &str,
function: &Bound<'_, PyAny>,
) -> PyResult<Selection> {
let selected = selected_cache(py, arguments, call_type)?;
let call_type: String = function.getattr("__name__")?.extract()?;
let selected = if python.asynchronous() {
configured_cache(py, &call_type)?
} else {
None
};
let Some(cache) = selected else {
return Ok(Selection {
backend: Backend::Disabled,
@ -102,6 +113,9 @@ pub(crate) fn configure(
let enabled = cache
.call_method("should_use_cache", (), Some(arguments))?
.extract::<bool>()?;
let skips_read = arguments
.get_item("caching")?
.is_some_and(|value| !value.is_none() && !value.is(pyo3::types::PyBool::new(py, true)));
let controls = arguments
.get_item("cache")?
.filter(|value| !value.is_none());
@ -117,7 +131,7 @@ pub(crate) fn configure(
};
let options = CacheOptions {
policy: CachePolicy {
no_cache: boolean("no-cache")?,
no_cache: boolean("no-cache")? || skips_read,
no_store: boolean("no-store")?,
..CachePolicy::default()
},
@ -129,7 +143,7 @@ pub(crate) fn configure(
.map(|value| value.extract())
.transpose()?
.unwrap_or_default();
python.bind(cache, arguments);
python.bind(cache, function, arguments);
Ok(Selection {
backend: if enabled {
Backend::Python { namespace }

9
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -252,9 +252,12 @@ impl PythonBinding for MessagesPythonHost {
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)?;
let function = py
.import(ROUTE_HOST_MODULE)
.and_then(|module| module.call_method0("cached_function"))
.map_err(InvokeError::Python)?;
let selection = crate::cache::configure(&mut self.cache, py, arguments, &function)
.map_err(InvokeError::Python)?;
let settings = self
.host_settings(py, arguments)
.map_err(InvokeError::Python)?;

282
litellm/caching/caching_handler.py
View file

@ -19,6 +19,7 @@ import datetime
import inspect
import time
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError
@ -78,6 +79,19 @@ from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
EmbeddingCacheInputElement = str | list[int] | ChatCompletionFileObject
@dataclass(frozen=True, slots=True)
class CacheRead:
kwargs: dict[str, object] # mutable-ok: the @client call kwargs the embedding path keeps reading and rewriting
checked: bool = False
cached_result: object | None = None
model: str = ""
end_time: datetime.datetime | None = None
@property
def hit(self) -> bool:
return self.end_time is not None
class CachingHandlerResponse(LiteLLMBaseModel):
"""
This is the response object for the caching handler. We need to separate embedding cached responses and (completion / text_completion / transcription) cached responses
@ -114,14 +128,7 @@ def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str
return {k: v for k, v in request_kwargs.items() if k != "litellm_logging_obj"}
def is_call_type_supported_by_cache(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 is_response_cache_excluded(model: str | None, kwargs: Mapping[str, object]) -> bool:
def _is_response_cache_excluded(model: str | None, kwargs: Mapping[str, object]) -> bool:
from litellm.llms.github_copilot.per_user_auth import is_github_copilot_per_user_request
custom_llm_provider: Final = kwargs.get("custom_llm_provider")
@ -402,121 +409,146 @@ class LLMCachingHandler:
Raises:
None
"""
if is_response_cache_excluded(model=model, kwargs=kwargs):
self.request_kwargs = _drop_logging_obj_from_kwargs(kwargs)
if not self.reads_cache(model=model, kwargs=kwargs):
return None
args = args or ()
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False
read: Final = await self.async_read_cache(
model=model,
original_function=original_function,
logging_obj=logging_obj,
call_type=call_type,
kwargs=kwargs,
args=args,
)
kwargs = read.kwargs
cached_result: object | None = read.cached_result
if litellm.cache is not None and read.checked:
if read.end_time is not None:
cached_result = self._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=original_function.__name__,
kwargs=kwargs,
logging_obj=logging_obj,
model=read.model,
custom_llm_provider=kwargs.get("custom_llm_provider", None),
args=args,
)
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
# LOG SUCCESS
self._async_log_cache_hit_on_callbacks(
logging_obj=logging_obj,
cached_result=cached_result,
start_time=start_time,
end_time=read.end_time,
cache_hit=True,
)
_set_cached_hidden_param(cached_result, "cache_key", self.hit_cache_key())
return CachingHandlerResponse(cached_result=cached_result)
elif (
call_type == CallTypes.aembedding.value
and cached_result is not None
and isinstance(cached_result, list)
and litellm.cache is not None
and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
):
(
final_embedding_cached_response,
embedding_all_elements_cache_hit,
) = self._process_async_embedding_cached_response(
final_embedding_cached_response=final_embedding_cached_response,
cached_result=cached_result,
kwargs=kwargs,
logging_obj=logging_obj,
start_time=start_time,
model=model,
)
return CachingHandlerResponse(
final_embedding_cached_response=final_embedding_cached_response,
embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
embedding_uncached_input=self.handle_kwargs_input_list_or_str(kwargs),
)
# Check if caching should be performed BEFORE doing expensive operations
if (
verbose_logger.debug("CACHE RESULT: %s", cached_result)
return CachingHandlerResponse(
cached_result=cached_result,
final_embedding_cached_response=final_embedding_cached_response,
)
_async_get_cache = async_get_cache
def reads_cache(
self,
model: str,
kwargs: dict[str, Any], # pyright: ignore[reportExplicitAny] # mutable-ok: the @client call kwargs, read as the wrapper reads them
) -> bool:
if _is_response_cache_excluded(model=model, kwargs=kwargs):
self.request_kwargs = _drop_logging_obj_from_kwargs(kwargs)
return False
return (
(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
args = args or ()
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False
cached_result: object | None = None
kwargs = kwargs.copy()
#########################################################
# Init cache timing metrics
#########################################################
cache_check_start_time: Final = time.perf_counter()
cache_check_end_time: float | None = None
#########################################################
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
kwargs["parent_otel_span"] = parent_otel_span
) # allow users to control returning cached responses from the completion function
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
verbose_logger.debug("Checking Async Cache")
logging_obj.reset_caching_details()
cached_result = await self._retrieve_from_cache(
call_type=call_type,
kwargs=kwargs,
args=args,
)
cache_check_end_time = time.perf_counter()
async def async_read_cache(
self,
model: str,
original_function: Callable[..., object],
logging_obj: LiteLLMLoggingObj,
call_type: str,
kwargs: Mapping[str, Any], # pyright: ignore[reportExplicitAny] # the @client call kwargs, read as the wrapper reads them
args: tuple[object, ...],
) -> CacheRead:
cache_check_start_time: Final = time.perf_counter()
read_kwargs: Final = dict(kwargs)
read_kwargs["parent_otel_span"] = get_parent_otel_span_from_kwargs(read_kwargs)
if litellm.cache is None or not self._is_call_type_supported_by_cache(original_function=original_function):
return CacheRead(kwargs=read_kwargs)
verbose_logger.debug("Checking Async Cache")
logging_obj.reset_caching_details()
cached_result: Final = await self._retrieve_from_cache(
call_type=call_type,
kwargs=read_kwargs,
args=args,
)
cache_check_end_time: Final = time.perf_counter()
if cached_result is None or isinstance(cached_result, list):
return CacheRead(kwargs=read_kwargs, checked=True, cached_result=cached_result)
verbose_logger.debug("Cache Hit!")
end_time: Final = datetime.datetime.now()
provider_model, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=model,
custom_llm_provider=read_kwargs.get("custom_llm_provider", None),
api_base=read_kwargs.get("api_base", None),
api_key=read_kwargs.get("api_key", None),
)
self._update_litellm_logging_obj_environment(
logging_obj=logging_obj,
model=provider_model,
kwargs=read_kwargs,
cached_result=cached_result,
is_async=True,
custom_llm_provider=custom_llm_provider,
cache_duration_ms=(cache_check_end_time - cache_check_start_time) * 1000,
)
return CacheRead(
kwargs=read_kwargs,
checked=True,
cached_result=cached_result,
model=provider_model,
end_time=end_time,
)
if cached_result is not None and not isinstance(cached_result, list):
verbose_logger.debug("Cache Hit!")
cache_hit: Final = True
end_time: Final = datetime.datetime.now()
provider_model, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=model,
custom_llm_provider=kwargs.get("custom_llm_provider", None),
api_base=kwargs.get("api_base", None),
api_key=kwargs.get("api_key", None),
)
cache_duration_ms: Final = (cache_check_end_time - cache_check_start_time) * 1000
self._update_litellm_logging_obj_environment(
logging_obj=logging_obj,
model=provider_model,
kwargs=kwargs,
cached_result=cached_result,
is_async=True,
custom_llm_provider=custom_llm_provider,
cache_duration_ms=cache_duration_ms,
)
call_type = original_function.__name__
cached_result = self._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=call_type,
kwargs=kwargs,
logging_obj=logging_obj,
model=provider_model,
custom_llm_provider=kwargs.get("custom_llm_provider", None),
args=args,
)
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
# LOG SUCCESS
self._async_log_cache_hit_on_callbacks(
logging_obj=logging_obj,
cached_result=cached_result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
)
cache_key: Final = (
self.preset_cache_key
or self.request_kwargs.get("cache_key")
or litellm.cache.get_cache_key(**self.request_kwargs)
)
_set_cached_hidden_param(cached_result, "cache_key", cache_key)
return CachingHandlerResponse(cached_result=cached_result)
elif (
call_type == CallTypes.aembedding.value
and cached_result is not None
and isinstance(cached_result, list)
and litellm.cache is not None
and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
):
(
final_embedding_cached_response,
embedding_all_elements_cache_hit,
) = self._process_async_embedding_cached_response(
final_embedding_cached_response=final_embedding_cached_response,
cached_result=cached_result,
kwargs=kwargs,
logging_obj=logging_obj,
start_time=start_time,
model=model,
)
return CachingHandlerResponse(
final_embedding_cached_response=final_embedding_cached_response,
embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
embedding_uncached_input=self.handle_kwargs_input_list_or_str(kwargs),
)
verbose_logger.debug("CACHE RESULT: %s", cached_result)
return CachingHandlerResponse(
cached_result=cached_result,
final_embedding_cached_response=final_embedding_cached_response,
)
# Caching disabled - return None to indicate no caching attempted
return None
_async_get_cache = async_get_cache
def hit_cache_key(self) -> object:
if litellm.cache is None:
return None
return (
self.preset_cache_key
or self.request_kwargs.get("cache_key")
or litellm.cache.get_cache_key(**self.request_kwargs)
)
def sync_get_cache(
self,
@ -531,7 +563,7 @@ class LLMCachingHandler:
cached_result: Any | None = None
# Check if caching should be performed BEFORE doing expensive kwargs copy
if is_response_cache_excluded(model=model, kwargs=kwargs):
if _is_response_cache_excluded(model=model, kwargs=kwargs):
self.request_kwargs = _drop_logging_obj_from_kwargs(kwargs)
return CachingHandlerResponse(cached_result=None)
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
@ -889,14 +921,14 @@ class LLMCachingHandler:
_combine_cached_embedding_response_with_api_result = combine_cached_embedding_response_with_api_result
def _async_log_cache_hit_on_callbacks(
def async_log_cache_hit_on_callbacks(
self,
logging_obj: LiteLLMLoggingObj,
cached_result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
cache_hit: bool,
):
) -> None:
"""
Helper function to log the success of a cached result on callbacks
@ -937,6 +969,8 @@ class LLMCachingHandler:
cache_hit=cache_hit,
)
_async_log_cache_hit_on_callbacks = async_log_cache_hit_on_callbacks
async def _retrieve_from_cache(
self, call_type: str, kwargs: dict[str, object], args: tuple[object, ...]
) -> object | None:
@ -972,7 +1006,7 @@ class LLMCachingHandler:
model_value: Final = new_kwargs.get("model")
model: Final = model_value if isinstance(model_value, str) else None
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
if is_response_cache_excluded(model=model, kwargs=new_kwargs):
if _is_response_cache_excluded(model=model, kwargs=new_kwargs):
return None
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
@ -1330,7 +1364,7 @@ class LLMCachingHandler:
)
no_store: Final = cache_options.get("no-store", False)
return (
not is_response_cache_excluded(model=model, kwargs=kwargs)
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
)
@ -1371,7 +1405,11 @@ class LLMCachingHandler:
Returns:
bool: True if the call type is supported by the cache, False otherwise.
"""
return is_call_type_supported_by_cache(str(original_function.__name__))
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)
async def add_streaming_response_to_cache(self, processed_chunk: ModelResponse) -> None:
"""

2
litellm/messages/dispatch.py
View file

@ -139,7 +139,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**{feature: detect for feature, detect in detectors.COMMON.items() if feature not in _NATIVE_COMMON_FEATURES},
Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier,
Feature.RESPONSE_CACHE: route_detectors.sync_response_cache,
Feature.MODIFY_PARAMS: detectors.modify_params,
Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity,
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
@ -150,7 +149,6 @@ _ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**_FEATURES,
Feature.MOCK_RESPONSE: route_detectors.streamed_mock_response,
Feature.RESPONSE_CACHE: route_detectors.response_cache,
Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook,
Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception,
Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted,

1
litellm/rust_bridge/catalog.py
View file

@ -169,7 +169,6 @@ MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VAR
Feature.PRE_REQUEST_HOOK,
Feature.WEBSEARCH_INTERCEPTION,
Feature.MESSAGES_INTERCEPTOR,
Feature.RESPONSE_CACHE,
Feature.MODIFY_PARAMS,
Feature.WORKLOAD_IDENTITY,
Feature.CACHE_CONTROL_INJECTION,

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

@ -1,7 +1,8 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
import weakref
from collections.abc import Callable, Mapping
from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # bounded cast of the served response
from pydantic import TypeAdapter
@ -9,9 +10,12 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_par
if TYPE_CHECKING:
from litellm.caching.caching import Cache
from litellm.litellm_core_utils.litellm_logging import Logging
_RUST_RESPONSE_HEADER: Final = "x-litellm-rust"
_HEADERS: Final = TypeAdapter(dict[str, object])
_EVENTS: Final = TypeAdapter(list[str])
_HITS: Final[weakref.WeakKeyDictionary[Logging, Mapping[str, object]]] = weakref.WeakKeyDictionary()
def configured_cache() -> Cache | None:
@ -25,17 +29,14 @@ def _python_entry(response: Mapping[str, object]) -> Mapping[str, object]:
headers: Final = hidden.get("additional_headers")
if not isinstance(headers, Mapping) or _RUST_RESPONSE_HEADER not in headers:
return response
return {
**response,
"_hidden_params": {
**hidden,
"additional_headers": {
name: value
for name, value in _HEADERS.validate_python(headers).items()
if name != _RUST_RESPONSE_HEADER
},
},
kept: Final = {
name: value for name, value in _HEADERS.validate_python(headers).items() if name != _RUST_RESPONSE_HEADER
}
stored: Final = {name: value for name, value in hidden.items() if name != "additional_headers" or kept} | (
{"additional_headers": kept} if kept else {}
)
entry: Final = {name: value for name, value in response.items() if name != "_hidden_params"}
return {**entry, "_hidden_params": stored} if stored else entry
def store_response(
@ -52,3 +53,87 @@ def store_response(
**kwargs, # pyright: ignore[reportArgumentType] # the cache call's own keyword arguments, as Python passes them
)
)
async def store_stream(
cache: Cache,
entry: Mapping[str, object],
/,
**kwargs: object, # kwargs-ok: the cache call's own arguments
) -> None:
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.types.caching import CACHED_STREAM_EVENTS_KEY
events: Final = entry.get(CACHED_STREAM_EVENTS_KEY)
handler: Final = getattr(kwargs.get("litellm_logging_obj"), "llm_caching_handler", None)
await cache.async_add_cache( # pyright: ignore[reportUnknownMemberType] # Cache.async_add_cache leaves its result untyped
{**entry, CACHED_STREAM_EVENTS_KEY: tuple(_EVENTS.validate_python(events))}
if isinstance(events, list)
else entry,
dynamic_cache_object=handler.dual_cache if isinstance(handler, LLMCachingHandler) else None,
**kwargs, # pyright: ignore[reportArgumentType] # the cache call's own keyword arguments, as Python passes them
)
async def lookup(
_cache: Cache,
original_function: Callable[..., object],
/,
**kwargs: object, # kwargs-ok: the cache call's own arguments
) -> object | None:
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.litellm_core_utils.litellm_logging import Logging
logging_obj: Final = kwargs.get("litellm_logging_obj")
if not isinstance(logging_obj, Logging):
return None
handler: Final = LLMCachingHandler(
original_function=original_function, request_kwargs=dict(kwargs), start_time=logging_obj.start_time
)
logging_obj.llm_caching_handler = handler
model: Final = kwargs.get("model")
model_name: Final = model if isinstance(model, str) else ""
if not handler.reads_cache(model=model_name, kwargs=kwargs):
return None
read: Final = await handler.async_read_cache(
model=model_name,
original_function=original_function,
logging_obj=logging_obj,
call_type=original_function.__name__,
kwargs=kwargs,
args=(),
)
if not read.hit:
return None
_remember(logging_obj, read.cached_result)
return read.cached_result
def _remember(logging_obj: Logging, cached: object) -> None:
if isinstance(cached, Mapping):
_HITS[logging_obj] = _HEADERS.validate_python(cached)
def _replace(response: object, cached: Mapping[str, object] | None) -> None:
if isinstance(response, dict) and cached is not None:
served: Final = cast( # cast-ok: a Messages response is the str-keyed dict the route serves
"dict[str, object]", response
)
served.clear()
served.update(cached)
def restore_hit(response: object, logger: Logging) -> None:
from litellm.caching.caching_handler import (
_set_cached_hidden_param, # pyright: ignore[reportPrivateUsage] # the @client cache hit's own hidden params
)
_replace(response, _HITS.pop(logger, None))
_set_cached_hidden_param(response, "cache_hit", True)
def hit_cache_key(logger: Logging) -> object:
from litellm.caching.caching_handler import LLMCachingHandler
handler: Final = logger.llm_caching_handler
return handler.hit_cache_key() if isinstance(handler, LLMCachingHandler) else None

25
litellm/rust_bridge/host/callbacks_legacy_python.py
View file

@ -42,6 +42,7 @@ class MetadataUpdater(Protocol):
kwargs: dict[str, object],
start_time: datetime.datetime,
end_time: datetime.datetime,
cache_hit: bool | None = None,
) -> None: ...
@ -98,6 +99,7 @@ def finalize(
kwargs: dict[str, object],
start_time: datetime.datetime,
end_time: datetime.datetime,
cache_hit: bool,
) -> None:
from litellm.litellm_core_utils.llm_response_utils import response_metadata
@ -105,6 +107,16 @@ def finalize(
update: Final = cast( # cast-ok: legacy metadata function accepts concrete kwargs
MetadataUpdater, response_metadata.update_response_metadata
)
if cache_hit:
from litellm.caching.caching_handler import (
_set_cached_hidden_param, # pyright: ignore[reportPrivateUsage] # the @client cache hit's own hidden params
)
from litellm.rust_bridge.host.cache import hit_cache_key, restore_hit
restore_hit(response, logger)
update(response, logger, logger.model, kwargs, start_time, end_time, cache_hit=True)
_set_cached_hidden_param(response, "cache_key", hit_cache_key(logger))
return
update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time)
cache_key: Final = logger.model_call_details.get("cache_key")
if logger.model_call_details.get("cache_hit") is True and isinstance(cache_key, str):
@ -401,6 +413,16 @@ def async_success_handler(
return logger.async_success_handler(response, start, end, cache_hit=_cache_hit(logger))
def cache_hit_success(logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime) -> None:
from litellm.caching.caching_handler import LLMCachingHandler
handler: Final = logger.llm_caching_handler
if isinstance(handler, LLMCachingHandler):
handler.async_log_cache_hit_on_callbacks(
logging_obj=logger, cached_result=response, start_time=start, end_time=end, cache_hit=True
)
def enqueue_logging(coroutine: Coroutine[object, object, None]) -> None:
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@ -474,9 +496,10 @@ def stamp_stream_head(
kwargs: dict[str, object], # mutable-ok: `finalize` takes the call kwargs as a dict
start_time: datetime.datetime,
end_time: datetime.datetime,
cache_hit: bool,
) -> None:
carrier: Final = _HiddenParamsCarrier(head)
finalize(carrier, logger, kwargs, start_time, end_time)
finalize(carrier, logger, kwargs, start_time, end_time, cache_hit)
stamped: Final = carrier._hidden_params # pyright: ignore[reportPrivateUsage] # the carrier's own field
if stamped is not head and isinstance(stamped, dict):
head.update(cast("dict[str, object]", stamped)) # cast-ok: update_response_metadata writes a str-keyed dict

7
litellm/rust_bridge/host/routes/messages.py
View file

@ -139,6 +139,13 @@ def _with_deployment_effort(params: Mapping[str, object], request: Mapping[str,
)
def cached_function() -> Callable[..., object]:
return cast( # cast-ok: the @client cache read names and checks the legacy handler it wraps
"Callable[..., object]",
handler.anthropic_messages, # noqa: TID251 # the @client cache read names and checks the legacy handler it wraps
)
def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]:
from litellm.llms.anthropic.messages.native.streaming_iterator import (
anthropic_messages_stream_hidden_params,

27
litellm/rust_bridge/messages/detectors.py
View file

@ -10,33 +10,6 @@ from litellm.rust_bridge.detectors import overridden, present
from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str
def _response_cache_consulted(request: NativeCall, call_type: str, *, honours_cache_controls: bool) -> bool:
from litellm.caching.caching_handler import is_call_type_supported_by_cache, is_response_cache_excluded
cache: Final = litellm.cache
kwargs: Final = request.resolved
if (
cache is None
or is_response_cache_excluded(optional_str(kwargs.get("model")), kwargs)
or not is_call_type_supported_by_cache(call_type)
):
return False
if honours_cache_controls and (
kwargs.get("caching") not in (None, True)
or (optional_mapping(kwargs.get("cache")) or {}).get("no-cache") is True
):
return False
return cache.should_use_cache(**kwargs) is True # pyright: ignore[reportUnknownMemberType] # Cache.should_use_cache takes untyped kwargs
def response_cache(request: NativeCall) -> bool:
return _response_cache_consulted(request, "anthropic_messages", honours_cache_controls=True)
def sync_response_cache(request: NativeCall) -> bool:
return _response_cache_consulted(request, "anthropic_messages_handler", honours_cache_controls=False)
def streamed_mock_response(request: NativeCall) -> bool:
return present(request, "mock_response") and request.resolved.get("stream") is True

81
tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py
View file

@ -35,6 +35,9 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.caching.caching import Cache, CacheMode
from litellm.caching.caching_handler import (
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # a primed entry must land before the measured call reads it
)
from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.messages import dispatch
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@ -322,6 +325,7 @@ class Case:
env: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
settings: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
private_keys: bool = False
cache: Callable[[], Cache] | None = None
gaps: Mapping[str, str | None] = field(default_factory=lambda: MappingProxyType({}))
def params(self) -> Iterator[object]:
@ -862,6 +866,36 @@ THINKING_RETRY_CASES: Final = (
),
)
def _default_off_cache() -> Cache:
return Cache(mode=CacheMode.default_off)
def _completion_only_cache() -> Cache:
return Cache(supported_call_types=["acompletion"])
CACHE_CASES: Final = (
Case("cache_hit", upstream=(_ok(), _ok()), cache=Cache),
Case("cache_hit_stream", _kw(stream=True), upstream=(_stream(TEXT_STREAM), _stream(TEXT_STREAM)), cache=Cache),
Case("cache_no_cache", _kw(cache={"no-cache": True}), upstream=(_ok(), _ok()), cache=Cache),
Case("cache_no_store", _kw(cache={"no-store": True}), upstream=(_ok(), _ok()), cache=Cache),
Case("cache_caching_false", _kw(caching=False), upstream=(_ok(), _ok()), cache=Cache),
Case(
"cache_ttl_and_namespace", _kw(cache={"ttl": 60, "namespace": "parity"}), upstream=(_ok(), _ok()), cache=Cache
),
Case("cache_default_off", upstream=(_ok(), _ok()), cache=_default_off_cache),
Case(
"cache_default_off_use_cache",
_kw(cache={"use-cache": True}),
upstream=(_ok(), _ok()),
cache=_default_off_cache,
),
Case("cache_unsupported_call_type", upstream=(_ok(), _ok()), cache=_completion_only_cache),
Case("cache_hit_with_metadata", _kw(metadata={"user_id": "user-123"}), upstream=(_ok(), _ok()), cache=Cache),
Case("cache_sync", sync=True, upstream=(_ok(), _ok()), cache=Cache),
)
CASES: Final = (
*REQUEST_CASES,
*RESPONSE_CASES,
@ -871,6 +905,7 @@ CASES: Final = (
*SERVED_NATIVELY_CASES,
*THINKING_RETRY_CASES,
*FALLBACK_CASES,
*CACHE_CASES,
)
@ -945,6 +980,15 @@ def _settings(settings: Mapping[str, object]) -> Generator[None]:
setattr(litellm, name, value)
@contextmanager
def _cache(cache: Callable[[], Cache] | None) -> Generator[None]:
if cache is None:
yield
return
with rebound(litellm, "cache", cache()):
yield
def _kwargs(case: Case, server: RecordingServer) -> dict[str, object]:
# Deep: Python writes into the caller's own `litellm_metadata`, which must not leak into the
# other leg.
@ -1062,21 +1106,24 @@ async def _invoke(
}, hidden
async def _terminal_events(recorder: RecordingLogger) -> tuple[str, ...]:
async def _terminal_events(recorder: RecordingLogger, since: int = 0) -> tuple[str, ...]:
for _ in range(100):
await drain_logging()
if any(name in _TERMINAL_EVENTS for name in recorder.names):
if any(name in _TERMINAL_EVENTS for name in recorder.names[since:]):
break
await asyncio.sleep(0.05)
await asyncio.sleep(0.2)
await drain_logging()
return tuple(name for name in recorder.names if name in _TERMINAL_EVENTS)
return tuple(name for name in recorder.names[since:] if name in _TERMINAL_EVENTS)
def _logged(recorder: RecordingLogger, events: tuple[str, ...]) -> tuple[Mapping[str, object], Mapping[str, object]]:
terminal: Final = [event for event in recorder.events if event.name in _TERMINAL_EVENTS]
def _logged(
recorder: RecordingLogger, events: tuple[str, ...], since: int = 0
) -> tuple[Mapping[str, object], Mapping[str, object]]:
terminal: Final = [event for event in recorder.events[since:] if event.name in _TERMINAL_EVENTS]
hooks: Final = tuple(sorted(set(recorder.names[since:])))
if not terminal:
return {"events": events}, {}
return {"events": events, "hooks": hooks}, {}
details: Final = _mapping(terminal[-1].kwargs)
standard: Final = _mapping(details.get("standard_logging_object"))
error_class: Final = _mapping(standard.get("error_information")).get("error_class")
@ -1085,6 +1132,7 @@ def _logged(recorder: RecordingLogger, events: tuple[str, ...]) -> tuple[Mapping
metadata_keys: Final = _mapping(metadata).keys()
logged: Final[Mapping[str, object]] = {
"events": events,
"hooks": hooks,
**{name: standard.get(name) for name in _COMPARED_LOG_FIELDS},
"error_class": error_class,
}
@ -1107,18 +1155,29 @@ def _logged(recorder: RecordingLogger, events: tuple[str, ...]) -> tuple[Mapping
return logged, callback_kwargs
async def _prime(case: Case, server: RecordingServer, recorder: RecordingLogger, native: bool) -> int:
if case.cache is None:
return 0
outcome: Final = await _invoke(case, _kwargs(case, server), native=native)
assert case.sync or outcome[0]["kind"] != "error", outcome
await _terminal_events(recorder)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
return len(recorder.events)
async def _leg(case: Case, server: RecordingServer, probe: FallbackProbe | None) -> Leg:
with isolated_callback_registries(), _environment(case.env), _settings(case.settings):
recorder: Final = RecordingLogger()
litellm.callbacks.append(recorder)
with isolated_callback_registries(), _environment(case.env), _settings(case.settings), _cache(case.cache):
for spec in case.upstream:
server.enqueue(replace(spec))
request_count: Final = len(server.requests)
recorder: Final = RecordingLogger()
litellm.callbacks.append(recorder)
fallbacks_before: Final = 0 if probe is None else probe.calls
primed: Final = await _prime(case, server, recorder, native=probe is not None)
outcome, hidden = await _invoke(case, _kwargs(case, server), native=probe is not None)
events: Final = await _terminal_events(recorder)
events: Final = await _terminal_events(recorder, primed)
server.responses.clear()
logged, callback_kwargs = _logged(recorder, events)
logged, callback_kwargs = _logged(recorder, events, primed)
return Leg(
fell_back=probe is not None and probe.calls > fallbacks_before,
requests=tuple(

122
tests/test_litellm_rust/cache/test_python_cache.py vendored
View file

@ -6,6 +6,7 @@ from pydantic import TypeAdapter
import litellm
from litellm.caching.caching import Cache, CacheMode
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.caching_handler import (
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
)
@ -33,6 +34,14 @@ from tests.test_litellm_rust.support.routing import native_route
pytestmark = pytest.mark.requires_rust_extension
_STORED: Final = TypeAdapter(dict[str, object])
def _entries(storage: InMemoryCache) -> dict[str, object]:
return _STORED.validate_python(
storage.cache_dict # pyright: ignore[reportUnknownMemberType] # InMemoryCache leaves its dict untyped
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
@ -110,7 +119,9 @@ async def test_legacy_cache_controls_and_backend_credential_key_semantics(record
recording_server.expected_requests = 3
litellm.cache = Cache()
await invoke("messages", recording_server, {"cache": {"no-store": True}})
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await invoke("messages", recording_server, {})
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
await invoke("messages", recording_server, {})
assert len(recording_server.requests) == 2
await invoke("messages", recording_server, {"cache": {"no-cache": True}})
@ -174,13 +185,18 @@ async def test_rust_messages_uses_a_legacy_cache_without_python_inference(
"api_base": recording_server.base_url,
}
first: Final = await invoke("messages", recording_server, parameters)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
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
assert cache_key(first) is None
assert cache_key(second) is None
await logger.wait_for_async("async_log_success_event", count=2)
successes: Final = await logger.wait_for_async("async_log_success_event", count=2)
assert [TypeAdapter(dict[str, object]).validate_python(event.kwargs).get("cache_hit") for event in successes] == [
None,
True,
]
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
@ -210,7 +226,9 @@ async def test_legacy_cache_default_off_requires_opt_in_even_for_existing_entrie
@pytest.mark.asyncio
async def test_legacy_cache_lookup_uses_backend_request_callback_semantics(recording_server: RecordingServer) -> None:
async def test_a_legacy_cache_hit_skips_the_pre_call_callback_that_would_rewrite_the_request(
recording_server: RecordingServer,
) -> None:
class Rewrite(RecordingLogger):
temperature = 0.1
@ -222,15 +240,16 @@ async def test_legacy_cache_lookup_uses_backend_request_callback_semantics(recor
litellm.cache = Cache()
recording_server.expected_requests = 1
await invoke("messages", recording_server, {"callbacks": [logger]})
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
first_hit: Final = await invoke("messages", recording_server, {"callbacks": [logger]})
logger.temperature = 0.8
await invoke("messages", recording_server, {"callbacks": [logger]})
second_hit: Final = await invoke("messages", recording_server, {"callbacks": [logger]})
assert logger.names.count("log_pre_api_call") == 4
assert logger.names.count("log_pre_api_call") == 1
assert cache_key(first_hit) is None
assert cache_key(second_hit) is None
assert len(recording_server.requests) == recording_server.expected_requests
assert recording_server.requests[0].body["temperature"] == 0.1
assert isinstance(cache_key(first_hit), str)
assert isinstance(cache_key(second_hit), str)
@pytest.mark.asyncio
@ -314,7 +333,11 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task_except_the_
second: Final = await invoke("messages", recording_server, {"stream": stream})
second_payload: Final = await collect(second) if stream else payload(second)
assert first_payload == second_payload
assert cache_key(second)
assert cache_key(first) is None
if stream:
assert cache_key(second) in _entries(storage)
else:
assert cache_key(second) is None
assert len(recording_server.requests) == 1
@ -389,6 +412,8 @@ async def test_messages_cache_stores_the_response_a_yielding_deployment_hook_ret
async def test_rust_messages_legacy_cache_honors_request_namespaces(
recording_server: RecordingServer, namespace_source: str
) -> None:
from litellm.caching.in_memory_cache import InMemoryCache
litellm.cache = Cache()
recording_server.expected_requests = 2
first_options: Final = (
@ -399,13 +424,20 @@ async def test_rust_messages_legacy_cache_honors_request_namespaces(
if namespace_source == "cache"
else {"metadata": {"redis_namespace": "second"}}
)
storage: Final = litellm.cache.cache
assert isinstance(storage, InMemoryCache)
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
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
first_keys: Final = set(_entries(storage))
await invoke("messages", recording_server, second_options)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
second_keys: Final = set(_entries(storage)) - first_keys
assert len(recording_server.requests) == 2
assert [key.split(":", 1)[0] for key in first_keys] == ["first"]
assert [key.split(":", 1)[0] for key in second_keys] == ["second"]
assert {key.split(":", 1)[1] for key in first_keys} == {key.split(":", 1)[1] for key in second_keys}
await invoke("messages", recording_server, first_options)
await invoke("messages", recording_server, second_options)
assert len(recording_server.requests) == 2
@ -420,9 +452,8 @@ async def test_rust_messages_legacy_semantic_cache_preserves_python_scope(
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)
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
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
@ -443,3 +474,62 @@ async def test_legacy_cache_keeps_public_messages_responses_compatible(
second_payload: Final = await collect(second) if stream else payload(second)
assert second_payload == first_payload
assert len(recording_server.requests) == 1
_PER_CALL_HIDDEN_PARAMS: Final = frozenset(
{"_response_ms", "litellm_overhead_time_ms", "callback_duration_ms", "litellm_call_id"}
)
async def _write_then_read(
recording_server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
writer: bool,
reader: bool,
stream: bool,
) -> tuple[object, object, object, int]:
from litellm.caching.in_memory_cache import InMemoryCache
storage: Final = InMemoryCache()
monkeypatch.setattr(litellm, "cache", Cache(_backend=storage))
requests_before: Final = len(recording_server.requests)
monkeypatch.setenv("LITELLM_RUST", "1" if writer else "0")
first: Final = await invoke("messages", recording_server, {"stream": stream}, native=writer)
if stream:
await collect(first)
else:
add_retry_headers_to_response(
_add_headers_to_response(first, {"x-litellm-model-group": "group"}), attempted_retries=0
)
get_hidden_params_dict(first)["model_id"] = "deployment"
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
(entry,) = _entries(storage).values()
stored: Final = TypeAdapter(dict[str, object]).validate_python(entry)["response"]
monkeypatch.setenv("LITELLM_RUST", "1" if reader else "0")
hit: Final = await invoke("messages", recording_server, {"stream": stream}, native=reader)
hidden: Final = {
name: value for name, value in get_hidden_params_dict(hit).items() if name not in _PER_CALL_HIDDEN_PARAMS
}
headers: Final = TypeAdapter(dict[str, object]).validate_python(hidden.pop("additional_headers", {}))
kept: Final = {
name: value for name, value in headers.items() if name not in {"x-litellm-rust", "x-litellm-cache-key"}
}
served: Final = await collect(hit) if stream else payload(hit)
return stored, served, {**hidden, "additional_headers": kept}, len(recording_server.requests) - requests_before
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
@pytest.mark.parametrize(("writer", "reader"), ((True, False), (False, True)), ids=("rust_to_python", "python_to_rust"))
async def test_a_cache_entry_written_by_one_leg_is_stored_and_served_like_the_readers_own(
recording_server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
writer: bool,
reader: bool,
stream: bool,
) -> None:
recording_server.expected_requests = 2
crossed: Final = await _write_then_read(recording_server, monkeypatch, writer, reader, stream)
own: Final = await _write_then_read(recording_server, monkeypatch, reader, reader, stream)
assert crossed == own
assert crossed[3] == 1

51
tests/test_litellm_rust/messages/test_dispatch.py
View file

@ -18,7 +18,7 @@ from litellm.messages.dispatch import (
)
from litellm.proxy import proxy_server
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import RULES, Feature, Route, RouteRule, Rules, Target
from litellm.rust_bridge.catalog import MESSAGES_PYTHON_FEATURES, RULES, Feature, Route, RouteRule, Rules, Target
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES, NativeAmessages, NativeMessages
from litellm.rust_bridge.public_call import NativeCall
@ -464,10 +464,6 @@ ADVISOR_TOOL: Final = {"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "
STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": True}
def _handler_cache() -> Cache:
return Cache(supported_call_types=["anthropic_messages_handler"]) # pyright: ignore[reportArgumentType] # the sync handler name the cache gate keys on
@pytest.mark.parametrize(
("feature", "model", "kwargs", "settings"),
(
@ -480,12 +476,6 @@ def _handler_cache() -> Cache:
),
(Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}),
(Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}),
(
Feature.RESPONSE_CACHE,
"anthropic/claude-sonnet-4-5",
{},
{"cache": _handler_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}),
(
@ -648,6 +638,45 @@ async def test_strict_tools_stay_native_unless_python_could_drop_strict(
assert served is native_result
@pytest.mark.asyncio
@pytest.mark.parametrize(
"kwargs",
(
pytest.param({}, id="reads-and-writes"),
pytest.param({"cache": {"no-store": True}}, id="no-store"),
pytest.param({"cache": {"ttl": 60, "namespace": "team"}}, id="ttl-and-namespace"),
pytest.param({"stream": True}, id="stream"),
),
)
async def test_a_configured_response_cache_stays_native(
monkeypatch: pytest.MonkeyPatch, kwargs: Mapping[str, object]
) -> None:
rules: Final[Rules] = (
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=MESSAGES_PYTHON_FEATURES),
RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),
)
native_result: Final = response()
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "cache", Cache())
async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub
pytest.fail("Rust reads and writes the configured response cache")
async def native(request: NativeCall) -> AnthropicMessagesResponse:
return native_result
served: Final = await _ADISPATCH.arun(
(16, MESSAGES, "anthropic/claude-sonnet-4-5"),
kwargs,
python=python,
binding=amessages_binding(native),
native=lambda hook, request, call_args, native_kwargs: hook(request),
rules=rules,
)
assert served is native_result
ASYNC_ONLY_FEATURES: Final = (
pytest.param(Feature.PRE_REQUEST_HOOK, {}, {"callbacks": [_Rewrite()]}, id="pre-request-hook"),
pytest.param(

3
tests/unit/rust_bridge/host/test_callbacks_legacy_python.py
View file

@ -73,7 +73,7 @@ def test_cache_hit_finalization_preserves_execution_provider_attribution(explici
)
prepared.logger.model_call_details.update({"cache_hit": True, "cache_key": "cached-response"})
response: Final = ModelResponse(model="cache-test-model")
legacy.finalize(response, prepared.logger, prepared.kwargs, now, now)
legacy.finalize(response, prepared.logger, prepared.kwargs, now, now, False)
assert prepared.logger.model_call_details["custom_llm_provider"] == "azure"
assert prepared.logger.model_call_details["model"] == "resolved-cache-model"
assert prepared.logger.litellm_params["metadata"]["user_api_key"] == "key-hash"
@ -404,6 +404,7 @@ def test_stream_head_is_stamped_in_place_with_the_call_hidden_params() -> None:
{"model": "anthropic/claude-sonnet-4-5", "litellm_call_id": "messages"},
started,
started + datetime.timedelta(milliseconds=5),
False,
)
assert head["additional_headers"] == {"request-id": "req_native", "llm_provider-request-id": "req_native"}

49
tests/unit/rust_bridge/test_detectors.py
View file

@ -9,7 +9,7 @@ import httpx
import pytest
import litellm
from litellm.caching.caching import Cache, CacheMode
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
@ -505,53 +505,6 @@ def test_response_cache_is_skipped_only_by_no_cache(
assert detectors.response_cache(call(**kwargs)) is expected
@pytest.mark.parametrize(
("cache", "kwargs", "expected"),
(
pytest.param(None, {}, False, id="no-cache-configured"),
pytest.param(Cache(), {}, True, id="default-cache"),
pytest.param(Cache(), {"cache": {"no-cache": True}}, False, id="no-cache-control"),
pytest.param(Cache(mode=CacheMode.default_off), {}, False, id="default-off"),
pytest.param(
Cache(mode=CacheMode.default_off), {"cache": {"use-cache": True}}, True, id="default-off-opted-in"
),
pytest.param(Cache(supported_call_types=["acompletion"]), {}, False, id="messages-not-cached"),
pytest.param(Cache(supported_call_types=["anthropic_messages"]), {}, True, id="messages-cached"),
pytest.param(Cache(), {"caching": False}, False, id="caching-off"),
pytest.param(Cache(), {"model": "databricks_agent/agent"}, False, id="excluded-provider"),
),
)
def test_messages_response_cache_fires_only_when_python_touches_the_cache(
monkeypatch: pytest.MonkeyPatch, cache: Cache | None, kwargs: Mapping[str, object], expected: bool
) -> None:
monkeypatch.setattr(litellm, "cache", cache)
assert messages_detectors.response_cache(call(**kwargs)) is expected
def _handler_cache() -> Cache:
return Cache(supported_call_types=["anthropic_messages_handler"]) # pyright: ignore[reportArgumentType] # the sync handler name the cache gate keys on
@pytest.mark.parametrize(
("cache", "kwargs", "expected"),
(
pytest.param(_handler_cache(), {}, True, id="handler-cached"),
pytest.param(Cache(supported_call_types=["anthropic_messages"]), {}, False, id="only-async-cached"),
pytest.param(
_handler_cache(),
{"cache": {"no-cache": True}},
True,
id="sync-lookup-ignores-no-cache",
),
),
)
def test_sync_messages_response_cache_keys_on_the_handler_name(
monkeypatch: pytest.MonkeyPatch, cache: Cache, kwargs: Mapping[str, object], expected: bool
) -> None:
monkeypatch.setattr(litellm, "cache", cache)
assert messages_detectors.sync_response_cache(call(**kwargs)) is expected
@pytest.mark.parametrize(
("kwargs", "expected"),
(

2130
tests/unit/test_private_usage_aliases.py
View file

File diff suppressed because it is too large Load diff