From 33e5d9168efbce6d03ca4c60aba36ca87517436e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 22:15:21 +0000 Subject: [PATCH] refactor(rust-bridge): let cache adapters own request preparation Keep Python cache key construction, lookup and storage in Python. Prepare native identity lazily through the selected adapter, preserve authoritative hit keys and cache-handler parity, and consolidate the cache-local Python binding. Co-authored-by: Yujong Lee --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/cache-response/AGENTS.md | 68 ++- litellm-rust/crates/cache-response/src/lib.rs | 4 +- .../crates/cache-response/src/service.rs | 156 ++++- .../crates/cache-response/tests/service.rs | 97 ++-- .../python_contract.json | 1 + .../callbacks-legacy-python/src/adapter.rs | 74 ++- .../callbacks-legacy-python/src/logger.rs | 6 + .../callbacks-legacy-python/src/mapping.rs | 19 +- .../callbacks-legacy-python/src/python.rs | 2 + .../src/test_support.rs | 1 + .../crates/gateway-inference/tests/caching.rs | 4 +- litellm-rust/crates/host-python/Cargo.toml | 1 + .../crates/host-python/src/binding.rs | 11 + litellm-rust/crates/host-python/src/driver.rs | 24 + .../host-python/src/hooks/chain/adapter.rs | 19 + .../host-python/src/hooks/chain/dispatch.rs | 53 ++ .../crates/host-python/tests/driver.rs | 30 +- .../crates/host-python/tests/support/mod.rs | 34 ++ litellm-rust/crates/host/AGENTS.md | 2 + litellm-rust/crates/host/Cargo.toml | 1 + litellm-rust/crates/host/src/hooks.rs | 15 + litellm-rust/crates/host/src/interceptors.rs | 21 + .../crates/host/src/machine/coroutine.rs | 2 + litellm-rust/crates/host/src/protocol.rs | 10 + .../crates/inference-chat/src/handler.rs | 9 +- litellm-rust/crates/inference-chat/src/lib.rs | 19 +- .../crates/inference-chat/tests/caching.rs | 1 + .../crates/inference-messages/src/lib.rs | 6 + .../inference-messages/tests/caching.rs | 13 +- .../inference-messages/tests/live/support.rs | 1 + .../inference-messages/tests/messages/host.rs | 1 + .../crates/inference-ocr/src/client.rs | 3 +- .../crates/inference-ocr/tests/ocr/machine.rs | 10 +- .../crates/inference-responses/src/handler.rs | 9 +- .../inference-responses/tests/caching.rs | 79 ++- litellm-rust/crates/inference/AGENTS.md | 4 +- litellm-rust/crates/inference/src/caching.rs | 156 ++--- litellm-rust/crates/inference/src/error.rs | 1 + .../crates/inference/tests/caching.rs | 135 +++-- .../crates/llms/src/base_llm/ocr/error.rs | 3 + .../crates/python-bridge/src/cache/AGENTS.md | 6 +- .../crates/python-bridge/src/cache/mod.rs | 4 +- .../python-bridge/src/cache/native/AGENTS.md | 4 +- .../python-bridge/src/cache/native/v2.rs | 86 +-- .../python-bridge/src/cache/python/AGENTS.md | 8 +- .../python-bridge/src/cache/python/binding.rs | 338 +++++++++++ .../python-bridge/src/cache/python/host.rs | 141 ----- .../python-bridge/src/cache/python/mod.rs | 7 +- .../python-bridge/src/cache/python/service.rs | 234 ++++++-- .../python-bridge/src/cache/selection.rs | 171 ------ .../crates/python-bridge/src/marshal.rs | 11 + .../src/routes/chat_completions/host.rs | 16 + .../src/routes/chat_completions/mod.rs | 27 +- .../python-bridge/src/routes/inference.rs | 57 +- .../python-bridge/src/routes/messages/host.rs | 68 +-- .../python-bridge/src/routes/messages/mod.rs | 85 +-- .../crates/python-bridge/src/routes/mod.rs | 9 + .../python-bridge/src/routes/ocr/host.rs | 16 + .../src/routes/responses/host.rs | 16 + .../python-bridge/src/routes/responses/mod.rs | 19 +- litellm/caching/caching.py | 8 +- litellm/caching/caching_handler.py | 105 ++-- litellm/chat_completions/dispatch.py | 1 - litellm/messages/dispatch.py | 1 - litellm/responses/dispatch.py | 1 - litellm/rust_bridge/catalog.py | 6 - litellm/rust_bridge/detectors.py | 9 +- litellm/rust_bridge/host/AGENTS.md | 4 +- litellm/rust_bridge/host/cache.py | 223 +++++++- .../host/callbacks_legacy_python.py | 6 + tests/test_litellm_rust/cache/__init__.py | 1 - tests/test_litellm_rust/cache/conftest.py | 19 - tests/test_litellm_rust/cache/support.py | 121 ---- .../cache/test_python_cache.py | 534 ------------------ tests/test_litellm_rust/cache/v2/__init__.py | 0 .../test_litellm_rust/cache/v2/test_cache.py | 207 ------- .../chat_completions/conftest.py | 10 +- .../chat_completions/test_cache.py | 328 +++++++++++ .../chat_completions/test_dispatch.py | 2 - tests/test_litellm_rust/conftest.py | 9 + tests/test_litellm_rust/messages/conftest.py | 11 +- .../test_litellm_rust/messages/test_cache.py | 497 ++++++++++++++++ .../messages/test_dispatch.py | 2 - tests/test_litellm_rust/ocr/test_lifecycle.py | 2 +- tests/test_litellm_rust/responses/conftest.py | 11 +- .../test_litellm_rust/responses/test_cache.py | 375 ++++++++++++ .../responses/test_dispatch.py | 2 - tests/test_litellm_rust/support/caching.py | 164 ++++++ .../support/recording_server.py | 6 + tests/test_litellm_rust/support/results.py | 15 +- tests/unit/caching/test_caching_handler.py | 17 + tests/unit/rust_bridge/host/test_cache.py | 338 ++++++++++- tests/unit/rust_bridge/test_detectors.py | 17 - tests/unit/test_internal_context.py | 1 + 95 files changed, 3635 insertions(+), 1847 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/python/binding.rs delete mode 100644 litellm-rust/crates/python-bridge/src/cache/python/host.rs delete mode 100644 litellm-rust/crates/python-bridge/src/cache/selection.rs delete mode 100644 tests/test_litellm_rust/cache/__init__.py delete mode 100644 tests/test_litellm_rust/cache/conftest.py delete mode 100644 tests/test_litellm_rust/cache/support.py delete mode 100644 tests/test_litellm_rust/cache/test_python_cache.py delete mode 100644 tests/test_litellm_rust/cache/v2/__init__.py delete mode 100644 tests/test_litellm_rust/cache/v2/test_cache.py create mode 100644 tests/test_litellm_rust/chat_completions/test_cache.py create mode 100644 tests/test_litellm_rust/messages/test_cache.py create mode 100644 tests/test_litellm_rust/responses/test_cache.py create mode 100644 tests/test_litellm_rust/support/caching.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 02b48a3be3b..4ddce5edc78 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3886,6 +3886,7 @@ dependencies = [ "litellm-auth", "litellm-coroutine", "rstest", + "serde", "serde_json", "tokio", ] diff --git a/litellm-rust/crates/cache-response/AGENTS.md b/litellm-rust/crates/cache-response/AGENTS.md index edd6a2f4ea9..fcae5eada20 100644 --- a/litellm-rust/crates/cache-response/AGENTS.md +++ b/litellm-rust/crates/cache-response/AGENTS.md @@ -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` 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 diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 78de27f2d9f..da0c960e3ba 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -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, }; diff --git a/litellm-rust/crates/cache-response/src/service.rs b/litellm-rust/crates/cache-response/src/service.rs index 0bdf948ec48..4c13b53ab87 100644 --- a/litellm-rust/crates/cache-response/src/service.rs +++ b/litellm-rust/crates/cache-response/src/service.rs @@ -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> + Send + 'a>>; @@ -24,46 +25,78 @@ impl Default for ResponseCacheConfig { } } +#[derive(Debug, PartialEq)] +pub struct CacheHit { + pub value: Value, + pub key: String, +} + pub trait ResponseCacheService: Send + Sync { - fn config(&self) -> &ResponseCacheConfig; + fn prepare( + self: Arc, + options: CacheOptions, + surface: &str, + input: &dyn Fn() -> Value, + ) -> PreparedCache; +} - fn lookup<'a>( - &'a self, - request: &'a ResponseCacheRequest, - now: Duration, - ) -> CacheFuture<'a, Option>; +pub struct PreparedCache { + pub policy: CachePolicy, + pub max_entry_bytes: usize, + pub call: Arc, +} - 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>; + + fn store<'a>(&'a self, response: Value, now: Duration) -> CacheFuture<'a, ()>; } impl ResponseCacheService for ResponseCache +where + B: BaseCache + 'static, +{ + fn prepare( + self: Arc, + options: CacheOptions, + surface: &str, + input: &dyn Fn() -> Value, + ) -> PreparedCache { + PreparedCache { + policy: options.policy, + max_entry_bytes: self.config().max_entry_bytes, + call: Arc::new(NativeCacheCall { + request: options.request(&self.config().namespace, surface, input()), + service: self, + }), + } + } +} + +struct NativeCacheCall> { + service: Arc>, + request: ResponseCacheRequest, +} + +impl ResponseCacheCall for NativeCacheCall where B: BaseCache, { - fn config(&self) -> &ResponseCacheConfig { - self.config() + fn lookup<'a>(&'a self, now: Duration) -> CacheFuture<'a, Option> { + Box::pin(async move { + Ok(self + .service + .async_lookup(&self.request, now) + .await? + .map(|value| CacheHit { + value, + key: cache_key(&self.request.key), + })) + }) } - fn lookup<'a>( - &'a self, - request: &'a ResponseCacheRequest, - now: Duration, - ) -> CacheFuture<'a, Option> { - Box::pin(self.async_lookup(request, now)) - } - - fn store<'a>( - &'a self, - request: &'a ResponseCacheRequest, - response: Value, - now: Duration, - ) -> CacheFuture<'a, ()> { - Box::pin(self.async_store(request, response, now)) + fn store<'a>(&'a self, response: Value, now: Duration) -> CacheFuture<'a, ()> { + Box::pin(self.service.async_store(&self.request, response, now)) } } @@ -83,8 +116,16 @@ pub struct CachePolicy { } impl CachePolicy { + pub fn reads(&self) -> bool { + self.caching != Some(false) && !self.no_cache + } + + pub fn writes(&self) -> bool { + self.caching != Some(false) && !self.no_store + } + pub fn enabled(&self) -> bool { - self.caching != Some(false) && !(self.no_cache && self.no_store) + self.reads() || self.writes() } } @@ -183,3 +224,56 @@ impl ScopedCache { } } } + +#[cfg(test)] +mod tests { + use std::cell::Cell; + + use litellm_cache_memory::InMemoryCache; + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::shared(CacheScope::Shared)] + #[case::isolated(CacheScope::Isolated("tenant".into()))] + #[tokio::test] + async fn native_preparation_keeps_existing_keys_and_storage(#[case] scope: CacheScope) { + let service = Arc::new( + ResponseCache::new(Arc::new(InMemoryCache::default())).with_config( + ResponseCacheConfig { + namespace: "test".into(), + max_entry_bytes: 4096, + }, + ), + ); + let options = CacheOptions::new(scope); + let input = json!({"model": "test", "body": {"z": 1, "a": 2}}); + let legacy = options.clone().request("test", "test", input.clone()); + let preparations = Cell::new(0); + let prepared = service.clone().prepare(options, "test", &|| { + preparations.set(preparations.get() + 1); + input.clone() + }); + assert_eq!(prepared.max_entry_bytes, 4096); + assert_eq!(prepared.call.lookup(Duration::ZERO).await.unwrap(), None); + prepared + .call + .store(json!({"answer": 42}), Duration::ZERO) + .await + .unwrap(); + assert_eq!( + service.async_lookup(&legacy, Duration::ZERO).await.unwrap(), + Some(json!({"answer": 42})) + ); + assert_eq!( + prepared.call.lookup(Duration::ZERO).await.unwrap(), + Some(CacheHit { + value: json!({"answer": 42}), + key: cache_key(&legacy.key), + }) + ); + assert_eq!(preparations.get(), 1); + } +} diff --git a/litellm-rust/crates/cache-response/tests/service.rs b/litellm-rust/crates/cache-response/tests/service.rs index be1cf1f8ea7..f25f086bcb0 100644 --- a/litellm-rust/crates/cache-response/tests/service.rs +++ b/litellm-rust/crates/cache-response/tests/service.rs @@ -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 = 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] diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index f65516265c5..4cb49da396a 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -77,6 +77,7 @@ "kwargs", "call_type" ], + "reviews_deployment_success": [], "after_deployment_success": [ "kwargs", "response", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index d891b5a1f0c..7cc683b50bc 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -311,13 +311,12 @@ impl LegacyLogging { self.prepare(py) } - fn resume_after_success( + fn resume_provider_result( &mut self, - py: Python<'_>, + _: Python<'_>, result: PyResult>, ) -> PyResult>> { - self.response = Some(result?); - self.finalize(py) + result.map(HookStep::Ready) } fn resume_deployment_failure( @@ -423,20 +422,29 @@ impl LegacyLogging { ) -> PyResult>> { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response); - if self.runs_deployment_hooks() { - return Ok(HookStep::Await( - DeploymentHooks::after_success( - py, - self.call.kwargs(), - &self.response, - self.call_type(), - )?, - Self::resume_after_success, - )); - } self.finalize(py) } + pub(crate) fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult { + Ok(self.runs_deployment_hooks() && DeploymentHooks::reviews_success(py)?) + } + + pub(crate) fn provider_result( + &mut self, + py: Python<'_>, + response: Py, + ) -> PyResult>> { + Ok(HookStep::Await( + DeploymentHooks::after_success( + py, + self.call.kwargs(), + &Some(response), + self.call_type(), + )?, + Self::resume_provider_result, + )) + } + pub(crate) fn result_ready( &mut self, py: Python<'_>, @@ -787,8 +795,24 @@ assert hooked == ([opaque] if asynchronous else []), hooked }); } + #[rstest] + #[case::synchronous(false, false)] + #[case::asynchronous(true, true)] + fn only_an_async_call_reviews_provider_results( + #[case] asynchronous: bool, + #[case] reviews: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"kwargs = {'logger': logger}"); + let (mut logging, _) = begin(py, &locals, asynchronous); + + assert_eq!(logging.reviews_provider_results(py).unwrap(), reviews); + }); + } + #[rstest::rstest] - fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { + fn the_success_hook_reviews_the_provider_result_and_its_answer_is_finalized() { Python::initialize(); Python::attach(|py| { let locals = namespace( @@ -809,18 +833,23 @@ logger.hooks = {'pre': lambda kwargs: kwargs} ) .unwrap(); let step = logging - .transform_response(py, local(&locals, "response").unbind(), TIMING) + .provider_result(py, local(&locals, "response").unbind()) .unwrap(); assert!(awaits_deployment_hook(&step)); - let step = resume( + let HookStep::Ready(reviewed) = resume( &mut logging, step, py, Ok(local(&locals, "replacement").unbind()), ) - .unwrap(); - let HookStep::Ready(returned) = step else { - panic!("expected the finalized response"); + .unwrap() else { + panic!("expected the hook's answer"); + }; + assert!(reviewed.bind(py).is(local(&locals, "replacement"))); + let HookStep::Ready(returned) = + logging.transform_response(py, reviewed, TIMING).unwrap() + else { + panic!("transforming the public response runs no deployment hook"); }; assert!(returned.bind(py).is(local(&locals, "replacement"))); run( @@ -829,6 +858,7 @@ logger.hooks = {'pre': lambda kwargs: kwargs} c" [finalized] = [value for name, value in logger.calls if name == 'finalize'] assert finalized is replacement +assert [name for name in logger.names() if name.endswith('_hook')] == ['pre_hook', 'success_hook'] ", ); }); @@ -853,7 +883,7 @@ assert finalized is replacement ) .unwrap(); let step = logging - .transform_response(py, local(&locals, "response").unbind(), TIMING) + .provider_result(py, local(&locals, "response").unbind()) .unwrap(); resume(&mut logging, step, py, Err(cancellation)) .err() diff --git a/litellm-rust/crates/callbacks-legacy-python/src/logger.rs b/litellm-rust/crates/callbacks-legacy-python/src/logger.rs index 38f3bf29828..0298689125b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/logger.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/logger.rs @@ -95,6 +95,12 @@ impl DeploymentHooks { .map(Bound::unbind) } + pub fn reviews_success(py: Python<'_>) -> PyResult { + python::DeploymentHooks::ReviewsSuccess + .call(py, ())? + .extract() + } + pub fn after_success( py: Python<'_>, kwargs: &Py, diff --git a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs index a4aee4eb4da..22749d4141c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs @@ -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, f64) -> Step, Box, &RequestContext) -> Step>; type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>; +type ProviderResult = fn(&mut LegacyLogging, Python<'_>, Py) -> Step>; type Transform = fn(&mut LegacyLogging, Python<'_>, Py, Timing) -> Step>; type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py) -> Step<()>; type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>; @@ -77,10 +79,16 @@ const AFTER: Binding = Binding { callbacks: &["log_post_api_call"], }; +const PROVIDER_RESULT: Binding = Binding { + boundary: CallBoundary::ProviderResult, + invoke: LegacyLogging::provider_result, + callbacks: &["async_post_call_success_deployment_hook"], +}; + const TRANSFORM: Binding = Binding { boundary: CallBoundary::TransformResponse, invoke: LegacyLogging::transform_public_response, - callbacks: &["async_post_call_success_deployment_hook"], + callbacks: &[], }; const SUCCESS: Binding = Binding { @@ -127,6 +135,7 @@ pub fn callback_mappings() -> impl Iterator { .mappings() .chain(BEFORE.mappings()) .chain(AFTER.mappings()) + .chain(PROVIDER_RESULT.mappings()) .chain(TRANSFORM.mappings()) .chain(SUCCESS.mappings()) .chain(FAILURE.mappings()) @@ -159,6 +168,14 @@ impl CallHooks for LegacyLogging { (BEFORE.invoke)(self, py, wire, context) } + fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult { + LegacyLogging::reviews_provider_results(self, py) + } + + fn provider_result(&mut self, py: Python<'_>, response: Py) -> Step> { + (PROVIDER_RESULT.invoke)(self, py, response) + } + fn transform_response( &mut self, py: Python<'_>, diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6f1ee9dfd78..578ba39b12e 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -65,6 +65,8 @@ pub(crate) enum Logging { pub(crate) enum DeploymentHooks { #[strum(serialize = "before_deployment_call")] BeforeDeploymentCall, + #[strum(serialize = "reviews_deployment_success")] + ReviewsSuccess, #[strum(serialize = "after_deployment_success")] AfterDeploymentSuccess, #[strum(serialize = "after_deployment_failure")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 093ce60e510..b85c78d9022 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -80,6 +80,7 @@ FAKES = { 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), 'is_internal_call': lambda: legacy.is_internal.get(), 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'reviews_deployment_success': lambda: True, 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( 'success', response, call_type ), diff --git a/litellm-rust/crates/gateway-inference/tests/caching.rs b/litellm-rust/crates/gateway-inference/tests/caching.rs index 143bc1ba6db..9e75365ac40 100644 --- a/litellm-rust/crates/gateway-inference/tests/caching.rs +++ b/litellm-rust/crates/gateway-inference/tests/caching.rs @@ -47,7 +47,7 @@ async fn all_inference_endpoints_share_native_cache( .expect(2) .mount(&upstream) .await; - let cache: Arc = 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() diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index 8d220b19def..b097d720ef0 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index 91fbc8c0e52..7f88efc1a8d 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -2,6 +2,7 @@ use crate::{InvokeError, PythonOwned}; use litellm_host::protocol::Protocol; use pyo3::prelude::*; use pyo3::types::PyDict; +use serde_json::Value; /// Converts requests, responses, stream values and errors at the Python boundary. pub trait PythonBinding: PythonOwned { @@ -27,6 +28,16 @@ pub trait PythonBinding: PythonOwned { response: ::Response, ) -> PyResult>; + /// 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> { + crate::to_py(py, response) + } + + /// The JSON form of the public object a Python hook handed back. + fn response_value(&mut self, _py: Python<'_>, public: &Bound<'_, PyAny>) -> PyResult { + crate::from_py(public) + } + /// What the stream carries at hand-off, as the caller's stream receives it. fn encode_stream_head( &mut self, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index f8e35aa09ba..479e470259f 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -294,6 +294,30 @@ where reply.send(*wire); Ok(()) } + HostRequest::Intercept(InterceptRequest::ProviderResult { response, reply }) => { + let step = { + let mut parts = self.parts(); + let Parts { binding, hooks, .. } = &mut *parts; + Python::attach(|py| { + if !hooks.reviews_provider_results(py)? { + return Ok(None); + } + let public = binding.public_response(py, &response)?; + hooks.provider_result(py, public).map(Some) + })? + }; + let Some(step) = step else { + reply.send(response); + return Ok(()); + }; + let public = self.settle(step).await?; + let reviewed = { + let mut parts = self.parts(); + Python::attach(|py| parts.binding.response_value(py, public.bind(py)))? + }; + reply.send(reviewed); + Ok(()) + } HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => { self.notify(PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts, diff --git a/litellm-rust/crates/host-python/src/hooks/chain/adapter.rs b/litellm-rust/crates/host-python/src/hooks/chain/adapter.rs index c7a5dad3858..89181ea7389 100644 --- a/litellm-rust/crates/host-python/src/hooks/chain/adapter.rs +++ b/litellm-rust/crates/host-python/src/hooks/chain/adapter.rs @@ -71,6 +71,12 @@ pub(super) trait ChainHooks: PythonOwned { py: Python<'_>, result: PyResult>, ) -> PyResult>>; + fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult; + fn provider_result( + &mut self, + py: Python<'_>, + response: Py, + ) -> PyResult>>; fn transform_response( &mut self, py: Python<'_>, @@ -138,6 +144,19 @@ impl ChainHooks for HookAdapter { Ok(self.step(step, Continuation::Wire)) } + fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult { + self.hooks.reviews_provider_results(py) + } + + fn provider_result( + &mut self, + py: Python<'_>, + response: Py, + ) -> PyResult>> { + let step = self.hooks.provider_result(py, response)?; + Ok(self.step(step, Continuation::Response)) + } + fn transform_response( &mut self, py: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs index eec4d9034d0..545a58546db 100644 --- a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs +++ b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs @@ -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; + type Context = (); + + fn call( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + value: Py, + _: &(), + ) -> PyResult>> { + hooks.provider_result(py, value) + } + + fn resume( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + result: PyResult>, + ) -> PyResult>> { + hooks.resume_response(py, result) + } + + fn suspend(index: usize, (): ()) -> Pending { + Pending::ProviderResult(index) + } + + fn resumed(pending: Pending) -> Option<(usize, ())> { + match pending { + Pending::ProviderResult(index) => Some((index, ())), + _ => None, + } + } +} + struct Response; impl Transform for Response { @@ -386,6 +422,23 @@ impl CallHooks for HookChain { self.transform_from::(py, 0, wire, context.clone()) } + fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult { + for hooks in &mut self.hooks { + if hooks.reviews_provider_results(py)? { + return Ok(true); + } + } + Ok(false) + } + + fn provider_result( + &mut self, + py: Python<'_>, + response: Py, + ) -> PyResult>> { + self.transform_from::(py, 0, response, ()) + } + fn transform_response( &mut self, py: Python<'_>, diff --git a/litellm-rust/crates/host-python/tests/driver.rs b/litellm-rust/crates/host-python/tests/driver.rs index 157cfc92148..c1c5f8496a3 100644 --- a/litellm-rust/crates/host-python/tests/driver.rs +++ b/litellm-rust/crates/host-python/tests/driver.rs @@ -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) { @@ -325,6 +325,34 @@ fn success_runs_every_step_in_order_and_returns_the_public_response( }); } +#[rstest] +#[case::reviewed(HookScript::ReviewResult, "reviewed", &["provider_result:provider"])] +#[case::not_reviewed(HookScript::Plain, "provider", &[])] +fn a_provider_result_keeps_what_the_reviewing_hooks_answer( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, + #[case] script: HookScript, + #[case] kept: &str, + #[case] reviewed: &[&str], +) { + Python::attach(|py| { + let (result, log) = run_scripted( + py, + reviewed_machine(), + OpScript::Answer, + script, + asynchronous, + ); + let seen: Vec<_> = log + .iter() + .filter(|entry| entry.starts_with("provider_result:")) + .cloned() + .collect(); + assert_eq!(result.unwrap().extract::(py).unwrap(), kept); + assert_eq!(seen, reviewed); + }); +} + #[rstest] fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks( #[from(interpreter)] _python: Interpreter, diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index 84778e0fdde..3e82bcf394f 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -321,6 +321,7 @@ pub enum HookScript { FailAfterSuccess, CancelTerminal, FailTerminal, + ReviewResult, } pub struct SyntheticHooks { @@ -336,6 +337,24 @@ pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks { } impl CallHooks for SyntheticHooks { + fn reviews_provider_results(&mut self, _: Python<'_>) -> PyResult { + Ok(matches!(self.script, HookScript::ReviewResult)) + } + + fn provider_result( + &mut self, + py: Python<'_>, + response: Py, + ) -> PyResult>> { + let reviewed = response.bind(py).cast::()?.copy()?; + self.log.push(format!( + "provider_result:{}", + reviewed.get_item("text")?.unwrap() + )); + reviewed.set_item("text", "reviewed")?; + Ok(HookStep::Ready(reviewed.into_any().unbind())) + } + fn prepare_arguments( &mut self, _: Python<'_>, @@ -474,6 +493,21 @@ pub fn success_machine() -> MachineFactory { }) } +/// A call that offers its provider result for review and returns the text it keeps. +pub fn reviewed_machine() -> MachineFactory { + Box::new(move |_| { + CallMachine::new(None, move |host| { + Box::pin(async move { + let kept = host + .interceptors + .provider_result(serde_json::json!({"text": "provider"})) + .await?; + Ok(kept["text"].as_str().unwrap_or_default().to_owned()) + }) + }) + }) +} + pub fn failing_machine() -> MachineFactory { Box::new(move |_| { CallMachine::new(None, |_| { diff --git a/litellm-rust/crates/host/AGENTS.md b/litellm-rust/crates/host/AGENTS.md index 5349ed14e40..f1c438ea5ee 100644 --- a/litellm-rust/crates/host/AGENTS.md +++ b/litellm-rust/crates/host/AGENTS.md @@ -32,3 +32,5 @@ Keep API policy in gateway-inference and python-bridge, and legacy callback poli Interceptors run inline and can rewrite values or fail execution. Observers consume owned `CallEvent` snapshots from `observation_channel`. Its bounded `ObservationSender` never waits for delivery and counts events dropped when the queue is full or closed. The host owns receiver processing and draining. Pass the same publisher to machine construction and the driver when one receiver should collect execution and lifecycle events. Legacy Python callbacks retain their existing awaited, fallible behavior through the Python adapter `ExecutionFacts` and `ResultSource` describe execution without pricing or budget policy. `InterceptPoints::result_ready` delivers these facts through an awaited `InterceptRequest::ResultReady`. Hosts receive them before response transformation or stream delivery. `ExecutionEvent::ResultReady` is the matching lifecycle event and can also be published as a passive snapshot. Accounting must consume the awaited path rather than a lossy observation queue + +`InterceptRequest::ProviderResult` offers a complete provider response before a route caches or delivers it, and the route keeps the answer. Routes raise it only for provider-sourced results, never for a cache hit, so hooks that run once per provider call sit inside the cache diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index 51015e99906..6fb216f9d6f 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true futures-util.workspace = true litellm-auth.workspace = true litellm-coroutine.workspace = true +serde.workspace = true serde_json.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/host/src/hooks.rs b/litellm-rust/crates/host/src/hooks.rs index da16bef258c..868828e3078 100644 --- a/litellm-rust/crates/host/src/hooks.rs +++ b/litellm-rust/crates/host/src/hooks.rs @@ -44,6 +44,21 @@ pub trait CallHooks: Sized { Ok(R::ready(wire)) } + /// Whether `provider_result` would do anything, so a host can skip encoding the response. + fn reviews_provider_results(&mut self, _runtime: R::Context<'_>) -> Result { + Ok(false) + } + + /// A complete provider response before it is cached or delivered. What this returns + /// replaces it. Never offered a cache hit. + fn provider_result( + &mut self, + _runtime: R::Context<'_>, + response: R::Response, + ) -> Result, R::Error> { + Ok(R::ready(response)) + } + fn transform_response( &mut self, _runtime: R::Context<'_>, diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index c66551a8a93..dc0577ce110 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -1,5 +1,7 @@ use std::future::Future; +use serde::{Serialize, de::DeserializeOwned}; + use crate::{ machine::MachineFault, protocol::{InterceptRequest, Reply, reply}, @@ -145,6 +147,25 @@ where fn result_ready(&self, facts: ExecutionFacts) -> impl Future> + Send { self.intercept(|reply| InterceptRequest::ResultReady { facts, reply }) } + + /// Offers a complete provider response before it is cached or delivered, and returns the + /// response the call keeps. Never raised for a cache hit. + fn provider_result(&self, response: R) -> impl Future> + Send + where + R: Serialize + DeserializeOwned + Send, + { + let value = serde_json::to_value(&response); + async move { + let value = value.map_err(|_| MachineFault::UnreadableAnswer)?; + let answer = self + .intercept(|reply| InterceptRequest::ProviderResult { + response: value, + reply, + }) + .await?; + serde_json::from_value(answer).map_err(|_| MachineFault::UnreadableAnswer.into()) + } + } } impl, T: Interceptors + ?Sized> InterceptPoints for T {} diff --git a/litellm-rust/crates/host/src/machine/coroutine.rs b/litellm-rust/crates/host/src/machine/coroutine.rs index 825536e1d9b..7ea678fef19 100644 --- a/litellm-rust/crates/host/src/machine/coroutine.rs +++ b/litellm-rust/crates/host/src/machine/coroutine.rs @@ -16,6 +16,8 @@ pub enum MachineFault { Abandoned, /// The host resumed the call out of turn. Protocol(ResumeError), + /// The host answered with a value the call cannot read back as its own type. + UnreadableAnswer, } pub type ExecuteFuture::Response> = diff --git a/litellm-rust/crates/host/src/protocol.rs b/litellm-rust/crates/host/src/protocol.rs index 91d15891cf9..23f9fdc1c5d 100644 --- a/litellm-rust/crates/host/src/protocol.rs +++ b/litellm-rust/crates/host/src/protocol.rs @@ -33,6 +33,12 @@ pub enum InterceptRequest { raw: RawResponse, reply: Reply<()>, }, + /// A complete response the provider just produced, before anything caches or delivers it. + /// Never raised for a cache hit. The answer is the response the call keeps. + ProviderResult { + response: serde_json::Value, + reply: Reply, + }, } impl InterceptRequest { @@ -43,6 +49,7 @@ impl InterceptRequest { reply.send(()); } Self::BeforeProviderRequest { wire, reply, .. } => reply.send(*wire), + Self::ProviderResult { response, reply } => reply.send(response), } } @@ -68,6 +75,9 @@ impl InterceptRequest { Self::AfterProviderResponse { raw, reply } => { reply.send(interceptors.after_provider_response(raw).await?); } + Self::ProviderResult { response, reply } => { + reply.send(interceptors.provider_result(response).await?); + } } Ok(()) } diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index f3fc17f1bc6..a479f856351 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -59,14 +59,11 @@ pub(super) async fn execute( ) .await?; let cache = cache.filter(|_| authenticated.signer.is_none()); - let cache_request = litellm_inference::caching::CacheRequest::from_wire( - identity, - cache.as_ref().map(|_| &wire), - ); + let cache_request = litellm_inference::caching::CacheRequest::from_wire::< + super::route::ChatCompletions, + >(identity, &wire, cache.as_ref(), cache_options); litellm_inference::caching::execute_unary::( cache_request, - cache.as_ref().map(|cache| cache.service.clone()), - cache.as_ref().map(|cache| cache.options(cache_options)), interceptors, observers, || async move { diff --git a/litellm-rust/crates/inference-chat/src/lib.rs b/litellm-rust/crates/inference-chat/src/lib.rs index 65e82af6a8d..9f6c5431eb7 100644 --- a/litellm-rust/crates/inference-chat/src/lib.rs +++ b/litellm-rust/crates/inference-chat/src/lib.rs @@ -12,7 +12,7 @@ use prepare::{prepare_provider_request, resolve_request}; use litellm_auth::AuthServices; use litellm_secrets::source::SecretSource; use std::sync::Arc; -use types::ChatCompletionsRequest; +use types::{ChatCompletionsCall, ChatCompletionsRequest}; #[derive(Clone)] pub struct ChatCompletionsRoute { @@ -48,6 +48,16 @@ impl ChatCompletionsRoute { request: ChatCompletionsRequest<'_>, interceptors: &impl litellm_host::interceptors::Interceptors, options: impl Into, + ) -> Result { + self.execute_call(request.into(), interceptors, options) + .await + } + + pub async fn execute_call( + &self, + call: ChatCompletionsCall, + interceptors: &impl litellm_host::interceptors::Interceptors, + options: impl Into, ) -> Result { let litellm_inference::CallOptions { cache: cache_options, @@ -55,12 +65,7 @@ impl ChatCompletionsRoute { } = options.into(); litellm_host::lifecycle::observe_unary( observers.clone(), - self.run_call( - request.into(), - cache_options, - interceptors, - observers.as_ref(), - ), + self.run_call(call, cache_options, interceptors, observers.as_ref()), ) .await } diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index a08670a9b92..f6f361c88f4 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -181,6 +181,7 @@ impl Interceptors for ChangingHooks { InterceptRequest::ResultReady { facts, reply } => { reply.send(self.result_ready(facts).await?); } + other => other.pass_through(), } Ok(()) } diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index c690c7a2ef4..985d50e1d23 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -106,6 +106,12 @@ impl MessagesRoute { source: source.clone(), }) .await?; + let output = litellm_inference::caching::provider_result::( + output, + &source, + context.interceptors, + ) + .await?; Ok(cache.finish(output, &source).await) }) .await diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 3725ddaac7c..634760a616e 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -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::( - CacheRequest::from_wire( + CacheRequest::from_wire::( ProviderIdentity { model: "test".into(), provider: "anthropic".into(), }, - Some(&WireRequest { + &WireRequest { url: "https://example.test/v1/messages".into(), headers: vec![], body: json!({ "model":"test", "messages":[{"role":"user","content":"hello"}], "max_tokens":32, (field):value }), - }), + }, + Some(&ScopedCache::new(cache.clone(), CacheScope::Shared)), + None, ), - Some(cache.clone()), - Some(CacheOptions::new(CacheScope::Shared)), &(), None, || async { @@ -148,6 +148,7 @@ impl Interceptors for ChangingHooks { InterceptRequest::ResultReady { facts, reply } => { reply.send(self.result_ready(facts).await?); } + other => other.pass_through(), } Ok(()) } diff --git a/litellm-rust/crates/inference-messages/tests/live/support.rs b/litellm-rust/crates/inference-messages/tests/live/support.rs index 5493ee47f59..abe057ecacf 100644 --- a/litellm-rust/crates/inference-messages/tests/live/support.rs +++ b/litellm-rust/crates/inference-messages/tests/live/support.rs @@ -54,6 +54,7 @@ impl Interceptors for LiveCall { InterceptRequest::AfterProviderResponse { raw, reply } => { reply.send(self.after_provider_response(raw).await?); } + other => other.pass_through(), } Ok(()) } diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index c4637f75056..a6802ed743c 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -95,6 +95,7 @@ impl litellm_host::interceptors::Interceptors< { reply.send(self.after_provider_response(raw).await?); } + other => other.pass_through(), } Ok(()) } diff --git a/litellm-rust/crates/inference-ocr/src/client.rs b/litellm-rust/crates/inference-ocr/src/client.rs index 21a4490048d..7584ed6f1f9 100644 --- a/litellm-rust/crates/inference-ocr/src/client.rs +++ b/litellm-rust/crates/inference-ocr/src/client.rs @@ -1,3 +1,4 @@ +use litellm_host::interceptors::InterceptPoints; use litellm_host::observation::ObservationSender; use std::sync::Arc; @@ -58,7 +59,7 @@ impl OcrRoute { caller_document, observers, )); - execute.await + interceptors.provider_result(execute.await?).await }) .await } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs index a3002db409f..86b86724a4c 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs @@ -70,6 +70,11 @@ async fn drive_until( .map(|()| reply.send(())) .map_err(HostFailure::Error) } + HostRequest::Intercept(other) => { + ops.push("provider_result"); + other.pass_through(); + Ok(()) + } }; if let Err(failure) = answer { break machine.interrupt(failure).await.map(completed); @@ -90,6 +95,7 @@ async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, sto MachineStep::Suspended(HostRequest::HostCall(op)) => host.handle_host_call(op).await.unwrap(), MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, reply, .. })) => reply.send(*wire), MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(()), + MachineStep::Suspended(HostRequest::Intercept(other)) => other.pass_through(), MachineStep::Suspended(HostRequest::Stream(stream)) => match stream { litellm_host::protocol::StreamDelivery::Open(head, _) => match head {}, litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {}, @@ -117,7 +123,7 @@ async fn a_hand_driven_machine_performs_the_same_call() { assert_eq!(outcome.unwrap().pages[0].markdown, "native"); assert_eq!(received(&upstream).await.len(), 1); - assert_eq!(ops, ["BeforeSend", "response"]); + assert_eq!(ops, ["BeforeSend", "response", "provider_result"]); assert!(matches!( machine.resume().await, Err(Error::InvalidRequest(_)) @@ -146,7 +152,7 @@ async fn a_path_document_is_read_by_core_without_a_host_operation() { std::fs::remove_dir_all(&dir).unwrap(); assert_eq!(response.unwrap().pages[0].markdown, "path"); - assert_eq!(ops, ["BeforeSend", "response"]); + assert_eq!(ops, ["BeforeSend", "response", "provider_result"]); assert_eq!( only_request(&upstream).await.json()["document"]["image_url"], "data:image/png;base64,YWJj" diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index 0bf8aa974e6..7f4beba887d 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -37,14 +37,11 @@ pub(super) async fn execute( ) .await?; let cache = cache.filter(|_| authenticated.signer.is_none()); - let cache_request = litellm_inference::caching::CacheRequest::from_wire( - identity, - cache.as_ref().map(|_| &wire), - ); + let cache_request = litellm_inference::caching::CacheRequest::from_wire::< + super::route::Responses, + >(identity, &wire, cache.as_ref(), cache_options); litellm_inference::caching::execute_streaming::( cache_request, - cache.as_ref().map(|cache| cache.service.clone()), - cache.as_ref().map(|cache| cache.options(cache_options)), interceptors, observers, || async move { diff --git a/litellm-rust/crates/inference-responses/tests/caching.rs b/litellm-rust/crates/inference-responses/tests/caching.rs index b9206ffae96..0ac328b03db 100644 --- a/litellm-rust/crates/inference-responses/tests/caching.rs +++ b/litellm-rust/crates/inference-responses/tests/caching.rs @@ -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( + input: Value, + cache: Option>, + options: Option, +) -> CacheRequest { + CacheRequest::prepare::

( + 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, + 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, + 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, litellm_cache::Error>> { + ) -> futures_util::future::BoxFuture< + 'a, + Result, litellm_cache::Error>, + > { Box::pin(async move { - Ok(self - .0 - .async_lookup(request, now) - .await? - .or_else(|| Some(self.1.clone()))) + Ok(Some(litellm_cache_response::CacheHit { + value: self + .service + .0 + .async_lookup(&self.request, now) + .await? + .unwrap_or_else(|| self.service.1.clone()), + key: litellm_cache_response::cache_key(&self.request.key), + })) }) } fn store<'a>( &'a self, - request: &'a litellm_cache_response::ResponseCacheRequest, response: Value, now: Duration, ) -> futures_util::future::BoxFuture<'a, Result<(), litellm_cache::Error>> { - Box::pin(self.0.async_store(request, response, now)) + Box::pin(self.service.0.async_store(&self.request, response, now)) } } @@ -139,6 +169,7 @@ impl Interceptors for ChangingHooks { InterceptRequest::ResultReady { facts, reply } => { reply.send(self.result_ready(facts).await?); } + other => other.pass_through(), } Ok(()) } @@ -184,9 +215,11 @@ async fn responses_refetches_instead_of_deserializing_another_api_response( let calls = AtomicUsize::new(0); for _ in 0..2 { let response = execute_unary::( - cache_request(json!({"input":"hello"})), - Some(cache.clone()), - Some(CacheOptions::new(CacheScope::Shared)), + cache_request::( + json!({"input":"hello"}), + Some(cache.clone()), + Some(CacheOptions::new(CacheScope::Shared)), + ), &(), None, || async { @@ -226,9 +259,11 @@ async fn responses_cache_only_reuses_completed_responses( let calls = AtomicUsize::new(0); for _ in 0..2 { let response = execute_unary::( - cache_request(json!({"input":"hello"})), - Some(cache.clone()), - Some(CacheOptions::new(CacheScope::Shared)), + cache_request::( + json!({"input":"hello"}), + Some(cache.clone()), + Some(CacheOptions::new(CacheScope::Shared)), + ), &(), None, || async { diff --git a/litellm-rust/crates/inference/AGENTS.md b/litellm-rust/crates/inference/AGENTS.md index 2c4bd00108b..ae71177050d 100644 --- a/litellm-rust/crates/inference/AGENTS.md +++ b/litellm-rust/crates/inference/AGENTS.md @@ -67,7 +67,7 @@ - Attach a `litellm_cache_response::ScopedCache` with `route.with_cache(cache)`. Cached and uncached routes use the same `execute` and `machine` methods - `CallOptions` carries a scope-free `CachePolicy` and observation. Per-call policy never replaces the attached scope or service - Provider transport does not own cache orchestration. Stream capture stays in `src/caching.rs` -- This crate and the format crates own request identity, typed response reconstruction, and stream capture and replay +- This crate and the format crates expose lazy effective-request identity and own typed response reconstruction and stream capture and replay. The selected cache service prepares per-call state before lookup. Only the native adapter materializes that identity and constructs native key inputs. Python delegation uses Python kwargs and never invokes the native identity builder - `cache-response` owns cache policy, namespacing, scope encoding, versioned envelopes, and freshness - The SDK explicitly chooses shared scope. The gateway derives isolated scope from authenticated identity before attaching its service - `ExecutionFacts` are delivered through the awaited `ResultReady` host operation for both provider and cached results, before public response processing or stream opening @@ -79,5 +79,5 @@ - native gateway accounting belongs to gateway dependencies, independently of `host-python` - response-cache services expose no coordination counters or reservation APIs. A shared Redis deployment does not make response storage and accounting coordination the same dependency - Cache lookup follows provider preparation, credential resolution, and the request interceptor - - keys describe the effective provider URL, authenticated headers, and rewritten body + - native keys describe the effective provider URL, authenticated headers, and rewritten body. Python delegation keeps Python's key semantics and returns the host's hit key unchanged - signed requests bypass caching until the signing identity has a stable cache representation diff --git a/litellm-rust/crates/inference/src/caching.rs b/litellm-rust/crates/inference/src/caching.rs index d6350f654b5..6f8a75e9dbe 100644 --- a/litellm-rust/crates/inference/src/caching.rs +++ b/litellm-rust/crates/inference/src/caching.rs @@ -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 { pub struct CacheRequest { pub identity: ProviderIdentity, - pub input: Value, + session: Option, } impl CacheRequest { - pub fn from_wire(identity: ProviderIdentity, wire: Option<&WireRequest>) -> Self { - Self { - input: wire.map_or(Value::Null, |wire| { + pub fn prepare( + identity: ProviderIdentity, + service: Option>, + options: Option, + 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( + identity: ProviderIdentity, + wire: &WireRequest, + cache: Option<&ScopedCache>, + policy: Option, + ) -> Self { + Self::prepare::

( + 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 { } struct CacheSession { - service: Arc, - request: ResponseCacheRequest, + prepared: PreparedCache, } impl CacheSession { - fn prepare( - service: Option>, - options: Option, - request: &CacheRequest, - ) -> Option { - 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(&self) -> Option> + async fn lookup(&self) -> Option<(CachedOutput, 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::>>(value) + match self.prepared.call.lookup(now()).await { + Ok(Some(hit)) => { + serde_json::from_value::>>(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( request: CacheRequest, - cache: Option>, - options: Option, interceptors: &impl Interceptors, observers: Option<&ObservationSender>, provider: F, @@ -156,14 +158,16 @@ where F: FnOnce() -> Fut, Fut: Future>, { - let identity = request.identity.clone(); + let CacheRequest { identity, session } = request; crate::diagnostic::provider(&identity.model, &identity.provider); - let session = CacheSession::prepare::

(cache, options, &request); let hit = match &session { - Some(session) => session.lookup::

().await.and_then(|entry| match entry { - CachedOutput::Response(response) => Some((response, cache_key(&session.request.key))), - CachedOutput::Stream(_) => None, - }), + Some(session) => session + .lookup::

() + .await + .and_then(|(entry, key)| match entry { + CachedOutput::Response(response) => Some((response, key)), + CachedOutput::Stream(_) => None, + }), None => None, }; let (response, source) = match hit { @@ -182,7 +186,11 @@ where observers, ) .await?; - if from_provider && let Some(session) = session { + if !from_provider { + return Ok(response); + } + let response = interceptors.provider_result(response).await?; + if let Some(session) = session { session.store_response::

(&response).await; } Ok(response) @@ -190,8 +198,6 @@ where pub async fn execute_streaming( request: CacheRequest, - cache: Option>, - options: Option, interceptors: &impl Interceptors, observers: Option<&ObservationSender>, provider: F, @@ -202,9 +208,8 @@ where F: FnOnce() -> Fut, Fut: Future, RouteError>>, { - let identity = request.identity.clone(); + let CacheRequest { identity, session } = request; crate::diagnostic::provider(&identity.model, &identity.provider); - let session = CacheSession::prepare::

(cache, options, &request); let cache = CallCache::

{ session, protocol: PhantomData, @@ -224,9 +229,29 @@ where observers, ) .await?; + let output = provider_result::

(output, &source, interceptors).await?; Ok(cache.finish(output, &source).await) } +/// Offers a complete provider response to the host before it is cached or delivered. A cache +/// hit and a stream pass through untouched. +pub async fn provider_result

( + output: OutputOf

, + source: &ResultSource, + interceptors: &impl Interceptors, +) -> Result, RouteError> +where + P: Protocol, + P::Response: Serialize + DeserializeOwned, +{ + match (source, output) { + (ResultSource::Provider, CallOutput::Complete(response)) => Ok(CallOutput::Complete( + interceptors.provider_result(response).await?, + )), + (_, output) => Ok(output), + } +} + pub struct CallCache

{ session: Option, protocol: PhantomData

, @@ -239,23 +264,9 @@ impl CallCache

{ 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::

(identity.clone(), wire, cache, Some(policy)); Self { - session, + session: request.session, protocol: PhantomData, } } @@ -265,16 +276,12 @@ impl CallCache

{ P::Response: DeserializeOwned, { let session = self.session.as_ref()?; - let output = match session.lookup::

().await? { + let (entry, key) = session.lookup::

().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

, source: &ResultSource) -> OutputOf

@@ -282,7 +289,7 @@ impl CallCache

{ P::Response: Serialize, { let Some(session) = self.session.filter(|session| { - *source == ResultSource::Provider && session.request.controls.writes() + *source == ResultSource::Provider && session.prepared.policy.writes() }) else { return output; }; @@ -310,8 +317,7 @@ fn capture_stream( 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; } diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index e5f02754f3d..3942d9a96f9 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -47,6 +47,7 @@ impl From for RouteError { Self::InvalidRequest(match fault { MachineFault::Abandoned => "host driver was abandoned".into(), MachineFault::Protocol(message) => format!("host {message}").into(), + MachineFault::UnreadableAnswer => "host answered with an unreadable response".into(), }) } } diff --git a/litellm-rust/crates/inference/tests/caching.rs b/litellm-rust/crates/inference/tests/caching.rs index 41180d4a53f..a0d913d7335 100644 --- a/litellm-rust/crates/inference/tests/caching.rs +++ b/litellm-rust/crates/inference/tests/caching.rs @@ -61,14 +61,20 @@ impl StreamCachable for TestRoute { } } -fn cache_request(input: Value) -> CacheRequest { - CacheRequest { - identity: ProviderIdentity { +fn cache_request( + input: Value, + cache: Option>, + options: Option, +) -> CacheRequest { + CacheRequest::prepare::

( + ProviderIdentity { model: "test-model".into(), provider: "test-provider".into(), }, - input, - } + cache, + options, + &|_| input.clone(), + ) } #[fixture] @@ -76,6 +82,38 @@ fn cache() -> Arc { 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, + #[case] configured: bool, + #[case] options_present: bool, + #[case] policy: CachePolicy, +) { + let request = CacheRequest::prepare::( + 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::(request, &(), None, || async { + Ok(json!({"source": "provider"})) + }) + .await + .unwrap(); + assert_eq!(response, json!({"source": "provider"})); +} + fn cache_with_limit(max_entry_bytes: usize) -> Arc { Arc::new( ResponseCache::new(Arc::new(InMemoryCache::new( @@ -96,9 +134,7 @@ async fn call( request: Value, ) -> Value { let output = execute_streaming::( - cache_request(request), - Some(cache.clone()), - options, + cache_request::(request, Some(cache.clone()), options), &(), None, || async { @@ -193,9 +229,11 @@ async fn streamed( fail: bool, ) -> OutputOf { execute_streaming::( - cache_request(json!({"stream":true})), - Some(cache.clone()), - Some(CacheOptions::new(CacheScope::Shared)), + cache_request::( + 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) { let calls = AtomicUsize::new(0); let first = execute_streaming::( - cache_request(json!({})), - Some(cache.clone()), - Some(CacheOptions::new(CacheScope::Shared)), + cache_request::( + 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, + 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, + 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, litellm_cache::Error>> { + ) -> futures_util::future::BoxFuture< + 'a, + Result, litellm_cache::Error>, + > { Box::pin(async move { - Ok(self - .0 - .async_lookup(request, now) - .await? - .or_else(|| Some(self.1.clone()))) + Ok(Some(litellm_cache_response::CacheHit { + value: self + .service + .0 + .async_lookup(&self.request, now) + .await? + .unwrap_or_else(|| self.service.1.clone()), + key: litellm_cache_response::cache_key(&self.request.key), + })) }) } fn store<'a>( &'a self, - request: &'a litellm_cache_response::ResponseCacheRequest, response: Value, now: Duration, ) -> futures_util::future::BoxFuture<'a, Result<(), litellm_cache::Error>> { - Box::pin(self.0.async_store(request, response, now)) + Box::pin(self.service.0.async_store(&self.request, response, now)) } } @@ -482,6 +546,7 @@ impl Interceptors for CacheHitAccounting { InterceptRequest::AfterProviderResponse { raw, reply } => { reply.send(self.after_provider_response(raw).await?); } + other => other.pass_through(), } Ok(()) } @@ -565,9 +630,11 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure( }; let result = if streaming_route { match execute_streaming::( - cache_request(request), - Some(cache), - Some(CacheOptions::new(CacheScope::Shared)), + cache_request::( + request, + Some(cache), + Some(CacheOptions::new(CacheScope::Shared)), + ), &accounting, Some(&observer), || async { panic!("a cache hit must not call the provider") }, @@ -579,9 +646,11 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure( } } else { execute_unary::( - cache_request(request), - Some(cache), - Some(CacheOptions::new(CacheScope::Shared)), + cache_request::( + request, + Some(cache), + Some(CacheOptions::new(CacheScope::Shared)), + ), &accounting, Some(&observer), || async { panic!("a cache hit must not call the provider") }, @@ -627,9 +696,7 @@ async fn unary_call( request: Value, ) -> Value { execute_unary::( - cache_request(request), - Some(cache.clone()), - options, + cache_request::(request, Some(cache.clone()), options), &(), None, || async { Ok(json!({"call":calls.fetch_add(1, Ordering::SeqCst)})) }, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 392dd56412c..cd4443303dd 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -118,6 +118,9 @@ impl From for Error { Self::InvalidRequest(match fault { MachineFault::Abandoned => "OCR host driver was abandoned".into(), MachineFault::Protocol(message) => format!("OCR {message}"), + MachineFault::UnreadableAnswer => { + "OCR host answered with an unreadable response".into() + } }) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md index 41c82187373..30dd7ec3a4a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md @@ -1,7 +1,7 @@ # Cache boundary -This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates +This folder owns how Rust inference reaches the configured cache. Rust holds no cache rules: `litellm.rust_bridge.host.cache` decides whether a call may read or write, builds the key and stores the entry, using the same Python handler functions the Python route runs. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates -`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters +`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `python/service.rs` owns the `Cached

` protocol and typed cache operations. `python/binding.rs` owns `CachedBinding`, which wraps a route's Python binding, asks the Python host module for access and serves cache operations. Route bindings never hold cache state -`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports +`python/` delegates operations to the Python cache the host module granted. `native/` owns the experimental V2 native cache handles exposed to Python. Inference does not use them until the Rust cache is ready diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index e190b12e665..83289151a30 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,10 +1,8 @@ mod native; mod python; -mod selection; pub(crate) use native::NativeCacheHandle; -pub(crate) use python::{CacheCall, PythonCache}; -pub(crate) use selection::{Cached, Selection, configure, configured_native}; +pub(crate) use python::{Cached, CachedBinding, scoped as python_cache}; use litellm_cache::Error; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md index 09a34d188cd..0753d8f5228 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md @@ -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 diff --git a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs index eb09b8131c3..6d7ad043629 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs @@ -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, backend: Arc, 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 { .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>> { - Ok(configured - .getattr_opt("cache")? - .map(|backend| backend.getattr_opt("native_handle")) - .transpose()? - .flatten() - .filter(|handle| handle.is_instance_of::())) -} - -pub(in crate::cache) fn configured( - configured: &Bound<'_, PyAny>, - kwargs: &Bound<'_, PyDict>, -) -> PyResult<( - Option>, - 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::>()?; - cache.check_process()?; - let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none()); - let controls = controls - .as_ref() - .map(|value| value.cast::()) - .transpose()?; - if let Some(controls) = controls { - for name in controls.keys() { - let name = name.extract::()?; - 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 { - 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> { - 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) diff --git a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md index 81a0997f02a..1fb8deaa13b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md @@ -1,9 +1,11 @@ # Python cache delegation -This directory lets Rust inference use a selected Python cache. `service.rs` implements the injected Rust response-cache service and yields typed cache operations. `host.rs` calls the Python cache's sync or async API and delivers the result back to Rust. +This directory lets Rust inference use the configured Python cache. `service.rs` owns the typed protocol and injected response-cache service. `binding.rs` wraps the route binding, asks `litellm.rust_bridge.host.cache` for access and delivers lookup and store results back to Rust. Keep PyO3 conversion, suspension and Python references in the binding, and cache rules in the Python host module -Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details +Every call into Python goes through that host module. Do not read `litellm.cache`, call cache methods directly, or re-derive which calls may read or write. Core remains independent of Python objects and cache implementation details -Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions +Preparing the Python adapter must not invoke the lazy native identity builder or create a native cache request or key. Keep the host's hit key unchanged. Python owns key construction, lookup, storage and the stored entry format. Rust sends and receives a route-neutral `CachedOutput` (`Response` or `Stream`), and `litellm.rust_bridge.host.cache` maps it to the entry the legacy Python route stores for that call type, so the two replay each other's entries. Never name a route's entry shape in Rust. Rust stores a complete response after the route raised its provider-result point, so what the deployment success hook answered is what gets cached, and a cache hit never raises that point. A stream is stored once the caller has consumed it + +Await asynchronous cache lookups through the existing host driver in the caller's task. Do not create another asyncio task or event loop for them. Asynchronous stores are handed to Python's `create_cache_write_task`, so the response returns before the write lands, as on the Python route. Cancellation during a lookup must prevent the provider request and the cache write. A failure to reach Python is a cache failure, never a request failure. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation diff --git a/litellm-rust/crates/python-bridge/src/cache/python/binding.rs b/litellm-rust/crates/python-bridge/src/cache/python/binding.rs new file mode 100644 index 00000000000..ef0e06e16eb --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/python/binding.rs @@ -0,0 +1,338 @@ +use super::service::{CacheCall, Cached}; +use litellm_cache::Error; +use litellm_cache_response::{CacheHit, CachePolicy}; +use litellm_host::protocol::{Protocol, Reply}; +use litellm_host_python::{ + InvokeError, PythonBinding, PythonHostCalls, PythonOwned, effective_py_args, from_py, to_py, +}; +use pyo3::{ + exceptions::{PyException, PyRuntimeError}, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::PyDict, +}; + +const MODULE: &str = "litellm.rust_bridge.host.cache"; + +pub(crate) struct CachedBinding { + inner: H, + cache: PythonCache, + call_type: &'static str, + positional: Option>, +} + +impl CachedBinding { + pub(crate) fn new( + inner: H, + call_type: &'static str, + asynchronous: bool, + positional: Py, + ) -> Self { + Self { + inner, + cache: PythonCache::new(asynchronous), + call_type, + positional: Some(positional), + } + } +} + +type ErrorOf = <::Protocol as Protocol>::Error; + +impl PythonBinding for CachedBinding { + type Protocol = Cached; + type Failure = H::Failure; + + fn decode_request( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> Result<::Request, InvokeError>> { + let request = self.inner.decode_request(py, arguments)?; + let policy = self + .positional + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("cache binding used after close")) + .and_then(|positional| effective_py_args(positional.bind(py), arguments)) + .and_then(|keyed| self.cache.configure(py, &keyed, self.call_type)) + .map_err(InvokeError::Python)?; + Ok((request, policy)) + } + + fn encode_response( + &mut self, + py: Python<'_>, + response: ::Response, + ) -> PyResult> { + self.inner.encode_response(py, response) + } + + fn public_response( + &mut self, + py: Python<'_>, + response: &serde_json::Value, + ) -> PyResult> { + self.inner.public_response(py, response) + } + + fn response_value( + &mut self, + py: Python<'_>, + public: &Bound<'_, PyAny>, + ) -> PyResult { + self.inner.response_value(py, public) + } + + fn encode_stream_head( + &mut self, + py: Python<'_>, + head: ::StreamHead, + ) -> PyResult> { + self.inner.encode_stream_head(py, head) + } + + fn encode_chunk( + &mut self, + py: Python<'_>, + chunk: ::Chunk, + ) -> PyResult> { + self.inner.encode_chunk(py, chunk) + } + + fn map_error(&self, py: Python<'_>, error: ErrorOf) -> PyResult { + self.inner.map_error(py, error) + } + + fn host_error(error: &PyErr) -> ErrorOf { + H::host_error(error) + } +} + +impl PythonHostCalls> for CachedBinding { + fn handle_host_call( + &mut self, + py: Python<'_>, + call: CacheCall, + ) -> Result<(), InvokeError>> { + self.begin_host_call(py, call).map(|_| ()) + } + + fn begin_host_call( + &mut self, + py: Python<'_>, + call: CacheCall, + ) -> Result>, InvokeError>> { + self.cache.begin(py, call).map_err(InvokeError::Python) + } + + fn resume_host_call( + &mut self, + py: Python<'_>, + result: PyResult>, + ) -> Result>, InvokeError>> { + self.cache.resume(py, result).map_err(InvokeError::Python) + } +} + +impl PythonOwned for CachedBinding { + fn close(&mut self, py: Python<'_>) { + self.inner.close(py); + self.cache.close(); + self.positional = None; + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.inner.traverse(visit)?; + self.cache.traverse(visit)?; + if let Some(positional) = &self.positional { + visit.call(positional)?; + } + Ok(()) + } +} + +enum Pending { + Lookup(Reply, Error>>), + Store(Reply>), +} + +struct Granted { + access: Py, +} + +fn cache_hit(hit: &Bound<'_, PyAny>) -> PyResult> { + if hit.is_none() { + return Ok(None); + } + Ok(Some(CacheHit { + value: from_py(&hit.getattr("response")?)?, + key: hit.getattr("key")?.extract()?, + })) +} + +struct PythonCache { + module: &'static str, + granted: Option, + pending: Option, + asynchronous: bool, +} + +impl PythonCache { + fn new(asynchronous: bool) -> Self { + Self { + module: MODULE, + granted: None, + pending: None, + asynchronous, + } + } + + fn configure( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + call_type: &str, + ) -> PyResult> { + let access = py + .import(self.module)? + .getattr("access")? + .call1((arguments, call_type))?; + if access.is_none() { + self.granted = None; + return Ok(None); + } + let policy = CachePolicy { + no_cache: !access.getattr("read")?.extract::()?, + no_store: !access.getattr("write")?.extract::()?, + ..CachePolicy::default() + }; + self.granted = Some(Granted { + access: access.unbind(), + }); + Ok(Some(policy)) + } + + fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult>> { + let granted = self.granted.as_ref().ok_or_else(|| { + PyRuntimeError::new_err("cache operation without granted cache access") + })?; + let access = granted.access.bind(py); + let function = |name: &str| py.import(self.module)?.getattr(name); + let (pending, result) = match call { + CacheCall::Lookup { reply } => ( + Pending::Lookup(reply), + function(if self.asynchronous { + "alookup" + } else { + "lookup" + }) + .and_then(|lookup| lookup.call1((access,))), + ), + CacheCall::Store { value, reply } => ( + Pending::Store(reply), + function(if self.asynchronous { "astore" } else { "store" }) + .and_then(|store| store.call1((access, to_py(py, &value)?))), + ), + }; + let result = result.map(Bound::unbind); + self.pending = Some(pending); + if self.asynchronous && result.is_ok() { + return result.map(Some); + } + self.resume(py, result) + } + + fn resume( + &mut self, + py: Python<'_>, + result: PyResult>, + ) -> PyResult>> { + if let Err(error) = &result + && !error.is_instance_of::(py) + { + self.pending = None; + return Err(result.err().unwrap()); + } + match self.pending.take() { + Some(Pending::Lookup(reply)) => { + let hit = result + .map_err(|_| Error::Unavailable) + .and_then(|hit| cache_hit(hit.bind(py)).map_err(|_| Error::InvalidEntry)); + reply.send(hit); + } + Some(Pending::Store(reply)) => { + reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable)); + } + None => { + return Err(PyRuntimeError::new_err( + "cache reply without pending operation", + )); + } + } + Ok(None) + } + + fn close(&mut self) { + self.pending = None; + self.granted = None; + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + let Some(granted) = &self.granted else { + return Ok(()); + }; + visit.call(&granted.access) + } +} + +#[cfg(test)] +mod tests { + use futures_util::FutureExt; + use litellm_host::protocol::reply; + use rstest::rstest; + use serde_json::Value; + + use super::*; + + #[rstest] + fn a_cache_call_that_cannot_reach_python_answers_unavailable( + #[values(false, true)] asynchronous: bool, + #[values(false, true)] store: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let mut cache = PythonCache { + module: "litellm_missing_cache_host", + granted: Some(Granted { access: py.None() }), + pending: None, + asynchronous, + }; + let answer = if store { + let (reply, answer) = reply(); + assert!( + cache + .begin( + py, + CacheCall::Store { + value: Value::Null, + reply + } + ) + .unwrap() + .is_none() + ); + answer.now_or_never().map(|answer| answer.unwrap().err()) + } else { + let (reply, answer) = reply(); + assert!( + cache + .begin(py, CacheCall::Lookup { reply }) + .unwrap() + .is_none() + ); + answer.now_or_never().map(|answer| answer.unwrap().err()) + }; + assert!(matches!(answer, Some(Some(Error::Unavailable)))); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs deleted file mode 100644 index d8f8ac7c181..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ /dev/null @@ -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, Error>>), - Store(Reply>), -} - -pub(crate) struct PythonCache { - cache: Option>, - arguments: Option>, - pending: Option, - 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>> { - 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>, - ) -> PyResult>> { - if let Err(error) = &result - && !error.is_instance_of::(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) - } -} diff --git a/litellm-rust/crates/python-bridge/src/cache/python/mod.rs b/litellm-rust/crates/python-bridge/src/cache/python/mod.rs index 973dfe89f1f..5af7b6f542b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/mod.rs @@ -1,6 +1,5 @@ -mod host; +mod binding; mod service; -pub(crate) use host::PythonCache; -pub(crate) use service::CacheCall; -pub(super) use service::service; +pub(crate) use binding::CachedBinding; +pub(crate) use service::{Cached, scoped}; diff --git a/litellm-rust/crates/python-bridge/src/cache/python/service.rs b/litellm-rust/crates/python-bridge/src/cache/python/service.rs index f215e5e97f5..82e22e103e3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/service.rs @@ -1,44 +1,31 @@ -use std::{future::Future, pin::Pin, time::Duration}; +use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Duration}; use litellm_cache::Error; use litellm_cache_response::{ - ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope, + CacheHit, CacheOptions, CachePolicy, CacheScope, PreparedCache, ResponseCacheCall, + ResponseCacheService, ResponseEnvelope, ScopedCache, }; use litellm_host::{ - machine::{HostServices, MachineFault}, + machine::HostServices, protocol::{Protocol, Reply}, }; -use litellm_inference::caching::CachedOutput; +use litellm_inference::caching::{Cachable, CachedOutput}; use serde_json::Value; -const STREAM_EVENTS_KEY: &str = "litellm_cached_anthropic_sse_events"; +pub(crate) struct Cached

(PhantomData

); -fn from_python(value: Value) -> Result { - let output = match value.get(STREAM_EVENTS_KEY) { - Some(events) => { - let events: Vec = - 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 { - let envelope: ResponseEnvelope> = - 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::>() - })), - } +impl Protocol for Cached

{ + type Request = (P::Request, Option); + type Response = P::Response; + type Error = P::Error; + type HostCall = CacheCall; + type Chunk = P::Chunk; + type StreamHead = P::StreamHead; } pub(crate) enum CacheCall { Lookup { - reply: Reply, Error>>, + reply: Reply, Error>>, }, Store { value: Value, @@ -46,58 +33,76 @@ pub(crate) enum CacheCall { }, } -struct PythonCacheService { - services: HostServices

, - config: ResponseCacheConfig, +struct PythonCacheService { + services: HostServices>, + protocol: PhantomData P>, } -pub(in crate::cache) fn service>( - services: HostServices

, - namespace: String, -) -> std::sync::Arc -where - P::Error: From, -{ - std::sync::Arc::new(PythonCacheService { - services, - config: ResponseCacheConfig { - namespace, - ..Default::default() - }, - }) +pub(crate) fn scoped(services: HostServices>) -> ScopedCache { + ScopedCache::new( + Arc::new(PythonCacheService { + services, + protocol: PhantomData, + }), + CacheScope::Shared, + ) } -impl> ResponseCacheService for PythonCacheService

-where - P::Error: From, -{ - fn config(&self) -> &ResponseCacheConfig { - &self.config +fn from_python(value: Value) -> Result { + let output: CachedOutput = + serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?; + serde_json::to_value(ResponseEnvelope::new(P::SURFACE, output)).map_err(|_| Error::InvalidEntry) +} + +fn to_python(value: Value) -> Result { + let envelope: ResponseEnvelope> = + serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?; + let output = envelope.decode(P::SURFACE).ok_or(Error::InvalidEntry)?; + serde_json::to_value(output).map_err(|_| Error::InvalidEntry) +} + +impl ResponseCacheService for PythonCacheService

{ + fn prepare( + self: Arc, + options: CacheOptions, + _: &str, + _: &dyn Fn() -> Value, + ) -> PreparedCache { + PreparedCache { + policy: options.policy, + max_entry_bytes: usize::MAX, + call: self, + } } +} +impl ResponseCacheCall for PythonCacheService

{ fn lookup<'a>( &'a self, - _: &'a ResponseCacheRequest, _: Duration, - ) -> Pin, Error>> + Send + 'a>> { + ) -> Pin, Error>> + Send + 'a>> { Box::pin(async move { self.services .call(|reply| CacheCall::Lookup { reply }) .await .map_err(|_| Error::Unavailable)?? - .map(from_python) + .map(|hit| { + Ok(CacheHit { + value: from_python::

(hit.value)?, + key: hit.key, + }) + }) .transpose() }) } fn store<'a>( &'a self, - _: &'a ResponseCacheRequest, value: Value, _: Duration, ) -> Pin> + Send + 'a>> { Box::pin(async move { - let value = to_python(value)?; + let value = to_python::

(value)?; self.services .call(|reply| CacheCall::Store { value, reply }) .await @@ -105,3 +110,122 @@ where }) } } + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + + use litellm_host::{ + interceptors::ProviderIdentity, + machine::{CallMachine, Machine, MachineStep}, + protocol::{HostRequest, Protocol}, + }; + use litellm_inference::{ + RouteError, + caching::{CacheRequest, execute_unary}, + }; + use rstest::rstest; + use serde_json::json; + + use super::*; + + struct TestRoute; + + impl Protocol for TestRoute { + type Request = (); + type Response = Value; + type Error = RouteError; + type HostCall = Infallible; + type Chunk = (); + type StreamHead = (); + } + + impl Cachable for TestRoute { + const SURFACE: &'static str = "test"; + } + + #[rstest] + #[case::hit(true)] + #[case::miss(false)] + #[tokio::test] + async fn python_cache_calls_never_materialize_native_identity(#[case] hit: bool) { + let mut machine = CallMachine::>::new(None, |host| { + Box::pin(async move { + let cache = scoped::(host.services); + let request = CacheRequest::prepare::( + ProviderIdentity { + model: "test".into(), + provider: "test".into(), + }, + Some(cache.service.clone()), + Some(cache.options(None)), + &|_| panic!("Python caching must not construct native identity"), + ); + execute_unary::(request, &(), None, || async { + Ok(json!({"source": "provider"})) + }) + .await + }) + }); + let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Lookup { reply })) = + machine.resume().await.unwrap() + else { + panic!("expected Python lookup"); + }; + reply.send(Ok(hit.then(|| CacheHit { + value: json!({"kind": "Response", "value": {"source": "cache"}}), + key: "python-owned-key".into(), + }))); + if !hit { + let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Store { value, reply })) = + machine.resume().await.unwrap() + else { + panic!("expected Python store on a miss"); + }; + assert_eq!( + value, + json!({"kind": "Response", "value": {"source": "provider"}}) + ); + reply.send(Ok(())); + } + let MachineStep::Complete(response) = machine.resume().await.unwrap() else { + panic!("expected completed call"); + }; + assert_eq!( + response, + json!({"source": if hit { "cache" } else { "provider" }}) + ); + } + + #[rstest] + #[tokio::test] + async fn python_hit_keeps_the_host_key() { + let mut machine = CallMachine::, CacheHit>::new(None, |host| { + Box::pin(async move { + let cache = scoped::(host.services); + let options = cache.options(None); + let prepared = cache.service.prepare(options, TestRoute::SURFACE, &|| { + panic!("Python caching must not construct native identity") + }); + Ok(prepared.call.lookup(Duration::ZERO).await.unwrap().unwrap()) + }) + }); + let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Lookup { reply })) = + machine.resume().await.unwrap() + else { + panic!("expected Python lookup"); + }; + reply.send(Ok(Some(CacheHit { + value: json!({"kind": "Response", "value": {"answer": 42}}), + key: "python-owned-key".into(), + }))); + let MachineStep::Complete(hit) = machine.resume().await.unwrap() else { + panic!("expected completed lookup"); + }; + assert_eq!(hit.key, "python-owned-key"); + assert_eq!( + from_python::(json!({"kind": "Response", "value": {"answer": 42}})).unwrap(), + hit.value + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs deleted file mode 100644 index cd19533fc73..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ /dev/null @@ -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

(std::marker::PhantomData

); - -impl Protocol for Cached

{ - 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), - Python { namespace: String }, -} - -pub(crate) struct Selection { - backend: Backend, - options: CacheOptions, -} - -impl Selection { - pub(crate) fn attach>( - self, - services: HostServices

, - ) -> (Option, CacheOptions) - where - P::Error: From, - { - 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>> { - 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>, - 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>, CacheOptions)> { - if !configured - .call_method("should_use_cache", (), Some(kwargs))? - .extract::()? - { - 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 { - 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::()?; - let controls = arguments - .get_item("cache")? - .filter(|value| !value.is_none()); - let boolean = |name: &str| -> PyResult { - controls - .as_ref() - .map(|value| value.cast::()?.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, - }) -} diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 038fa8c9630..ffe78e5b11b 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -126,6 +126,17 @@ pub(crate) fn public_response( .map(Bound::unbind) } +/// The JSON form of a public response a Python hook handed back, models nested in it included. +pub(crate) fn response_value(public: &Bound<'_, PyAny>) -> PyResult { + from_py( + &public + .py() + .import("pydantic_core")? + .getattr("to_jsonable_python")? + .call1((public,))?, + ) +} + struct RequestFieldSources<'py> { body: Option>, credentials: Option>, diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 5ad9f8df8b8..89b901e5bd8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -59,6 +59,22 @@ impl PythonBinding for ChatCompletionsPythonHost { self.0.response(py, &response) } + fn public_response( + &mut self, + py: Python<'_>, + response: &serde_json::Value, + ) -> PyResult> { + self.0.response(py, response) + } + + fn response_value( + &mut self, + _: Python<'_>, + public: &Bound<'_, PyAny>, + ) -> PyResult { + crate::marshal::response_value(public) + } + fn encode_stream_head( &mut self, _py: Python<'_>, diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs index d78d141f0f2..1d38fde046e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -4,13 +4,21 @@ use std::sync::Arc; use host::ChatCompletionsPythonHost; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, ScopedCache}; +use litellm_cache_response::ScopedCache; use litellm_callbacks_legacy_python::LoggingOperation; -use litellm_host::{call::HostedMachine, protocol::Protocol}; -use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions}; +use litellm_host::{ + call::{CallOutput, OutputOf}, + machine::ChannelInterceptors, +}; +use litellm_inference::{CallOptions, RouteError}; +use litellm_inference_chat::{ + ChatCompletionsRoute, route::ChatCompletions, types::ChatCompletionsCall, +}; use litellm_secrets::source::SecretSource; use pyo3::prelude::*; +use crate::cache::Cached; + use super::{ NativeCall, inference::{InferenceHost, InferenceRoute, run_inference}, @@ -61,11 +69,14 @@ impl InferenceRoute for ChatCompletionsRoute { self.with_cache(cache) } - fn machine( + async fn call( self, - call: ::Request, - policy: CachePolicy, - ) -> HostedMachine { - self.machine(call, policy) + request: ChatCompletionsCall, + interceptors: ChannelInterceptors>, + options: CallOptions, + ) -> Result, RouteError> { + self.execute_call(request, &interceptors, options) + .await + .map(CallOutput::Complete) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index 3b590a522fe..15feb0a856d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -1,13 +1,17 @@ use std::sync::Arc; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache}; +use litellm_cache_response::ScopedCache; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; -use litellm_host::{call::HostedMachine, protocol::Protocol}; -use litellm_host_python::{PythonBinding, PythonHostCalls, from_py}; +use litellm_host::{ + call::{OutputOf, hosted_call}, + machine::ChannelInterceptors, + protocol::Protocol, +}; +use litellm_host_python::{PythonBinding, from_py}; use litellm_http::transport::Error as TransportError; -use litellm_inference::RouteError; +use litellm_inference::{CallOptions, RouteError, caching::Cachable}; use litellm_secrets::source::SecretSource; use pyo3::{ exceptions::PyValueError, @@ -20,6 +24,7 @@ use serde_json::{Map, Value}; use super::{NativeCall, RequestView}; use crate::{ + cache::{Cached, CachedBinding}, errors::{RustUpstreamError, route_error_to_pyerr}, marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds}, }; @@ -143,8 +148,8 @@ impl InferenceHost { } } -pub(super) trait InferenceRoute: Sized + 'static { - type Protocol: Protocol; +pub(super) trait InferenceRoute: Sized + Send + 'static { + type Protocol: Cachable; const OPERATION: LoggingOperation; const SYNC_CALL_TYPE: &'static str; const ASYNC_CALL_TYPE: &'static str; @@ -155,11 +160,12 @@ pub(super) trait InferenceRoute: Sized + 'static { secrets: Arc, ) -> Self; fn with_cache(self, cache: ScopedCache) -> Self; - fn machine( + fn call( self, - call: ::Request, - policy: CachePolicy, - ) -> HostedMachine; + request: ::Request, + interceptors: ChannelInterceptors>, + options: CallOptions, + ) -> impl Future, RouteError>> + Send; } pub(super) fn run_inference( @@ -170,13 +176,14 @@ pub(super) fn run_inference( ) -> PyResult> where R: InferenceRoute, - H: PythonBinding + PythonHostCalls + 'static, + H: PythonBinding + 'static, { let call_type = if asynchronous { R::ASYNC_CALL_TYPE } else { R::SYNC_CALL_TYPE }; + let positional = call.positional()?.unbind(); let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?; super::run_public_call( py, @@ -188,14 +195,28 @@ where crate::http::resources().auth.clone(), crate::secrets::source(py)?, ); - let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?; - let route = match cache { - Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)), - None => route, - }; - Ok(route.machine(request, cache_options.policy)) + Ok(hosted_call( + request, + None, + move |(call, policy), services, interceptors, observers| async move { + let route = match policy { + Some(_) => route.with_cache(crate::cache::python_cache(services)), + None => route, + }; + route + .call( + call, + interceptors, + CallOptions { + cache: policy, + observers, + }, + ) + .await + }, + )) }, - host, + CachedBinding::new(host, call_type, asynchronous, positional), hooks, asynchronous, ) diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 5e445fcaa7d..aff91349fc4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -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 { /// 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; + type Protocol = Messages; type Failure = PyErr; fn decode_request( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - ) -> Result<(MessagesCall, Selection), InvokeError> { + ) -> Result> { self.view .resolve(py, arguments) .map_err(InvokeError::Python)?; - let selection = - crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages") - .map_err(InvokeError::Python)?; self.projection(py, arguments) .map_err(|error| InvokeError::Python(self.map_failure(py, error)))? .map_err(InvokeError::Native) - .map(|request| (request, selection)) } fn encode_response( @@ -260,6 +251,22 @@ impl PythonBinding for MessagesPythonHost { public_response(py, ROUTE_HOST_MODULE, response.as_ref()) } + fn public_response( + &mut self, + py: Python<'_>, + response: &serde_json::Value, + ) -> PyResult> { + public_response(py, ROUTE_HOST_MODULE, response) + } + + fn response_value( + &mut self, + _: Python<'_>, + public: &Bound<'_, PyAny>, + ) -> PyResult { + crate::marshal::response_value(public) + } + fn encode_stream_head( &mut self, py: Python<'_>, @@ -289,43 +296,12 @@ impl PythonBinding for MessagesPythonHost { } } -impl PythonHostCalls> for MessagesPythonHost { - fn handle_host_call( - &mut self, - py: Python<'_>, - op: CacheCall, - ) -> Result<(), InvokeError> { - self.cache - .begin(py, op) - .map(|_| ()) - .map_err(InvokeError::Python) - } - - fn begin_host_call( - &mut self, - py: Python<'_>, - op: CacheCall, - ) -> Result>, InvokeError> { - self.cache.begin(py, op).map_err(InvokeError::Python) - } - - fn resume_host_call( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> Result>, InvokeError> { - 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) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index c788760b268..4ff61b4fac4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,53 +1,26 @@ mod host; +use std::sync::Arc; + use host::MessagesPythonHost; +use litellm_auth::AuthServices; +use litellm_cache_response::ScopedCache; use litellm_callbacks_legacy_python::LoggingOperation; +use litellm_host::{call::OutputOf, machine::ChannelInterceptors}; +use litellm_inference::{CallOptions, RouteError}; +use litellm_inference_messages::{MessagesCall, MessagesRoute, route::Messages}; +use litellm_secrets::source::SecretSource; use pyo3::prelude::*; -use super::NativeCall; +use super::{ + NativeCall, + inference::{InferenceRoute, run_inference}, +}; +use crate::cache::Cached; fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { - 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::(py, call, asynchronous, host) } #[pyfunction] @@ -59,3 +32,31 @@ pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult, call: NativeCall<'_>) -> PyResult> { 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, + secrets: Arc, + ) -> Self { + Self::new(http, auth, secrets) + } + + fn with_cache(self, cache: ScopedCache) -> Self { + self.with_cache(cache) + } + + async fn call( + self, + request: MessagesCall, + interceptors: ChannelInterceptors>, + options: CallOptions, + ) -> Result, RouteError> { + self.execute(request, &interceptors, options).await + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index f51c755c33e..f1bff752bfb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -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> { + 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::new(&self.base, &self.kwargs) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index fc1b70392e2..47076e74ba5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -118,6 +118,22 @@ impl PythonBinding for OcrPythonHost { public_response(py, "litellm.rust_bridge.host.routes.ocr", &response) } + fn public_response( + &mut self, + py: Python<'_>, + response: &serde_json::Value, + ) -> PyResult> { + public_response(py, "litellm.rust_bridge.host.routes.ocr", response) + } + + fn response_value( + &mut self, + _: Python<'_>, + public: &Bound<'_, PyAny>, + ) -> PyResult { + crate::marshal::response_value(public) + } + fn encode_stream_head( &mut self, _: Python<'_>, diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index 8261e66406b..cd03243c0e2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -80,6 +80,22 @@ impl PythonBinding for ResponsesPythonHost { self.0.response(py, &response) } + fn public_response( + &mut self, + py: Python<'_>, + response: &serde_json::Value, + ) -> PyResult> { + self.0.response(py, response) + } + + fn response_value( + &mut self, + _: Python<'_>, + public: &Bound<'_, PyAny>, + ) -> PyResult { + crate::marshal::response_value(public) + } + fn encode_stream_head( &mut self, _py: Python<'_>, diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index 14012152e60..fc7dd42e75b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -5,10 +5,11 @@ use std::sync::Arc; use host::ResponsesPythonHost; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, ScopedCache}; +use litellm_cache_response::ScopedCache; use litellm_callbacks_legacy_python::LoggingOperation; -use litellm_host::{call::HostedMachine, protocol::Protocol}; -use litellm_inference_responses::{ResponsesRoute, route::Responses}; +use litellm_host::{call::OutputOf, machine::ChannelInterceptors}; +use litellm_inference::{CallOptions, RouteError}; +use litellm_inference_responses::{ResponsesRoute, route::Responses, types::ResponsesCall}; use litellm_secrets::source::SecretSource; use pyo3::prelude::*; pub(crate) use websocket::ResponsesWebSocketConnection; @@ -17,6 +18,7 @@ use super::{ NativeCall, inference::{InferenceHost, InferenceRoute, run_inference}, }; +use crate::cache::Cached; const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses"; @@ -53,11 +55,12 @@ impl InferenceRoute for ResponsesRoute { self.with_cache(cache) } - fn machine( + async fn call( self, - call: ::Request, - policy: CachePolicy, - ) -> HostedMachine { - self.machine(call, policy) + request: ResponsesCall, + interceptors: ChannelInterceptors>, + options: CallOptions, + ) -> Result, RouteError> { + self.execute(request, &interceptors, options).await } } diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5f5b8210096..4281b150fec 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -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 diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 7089ff78835..5c9eba05d13 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -26,8 +26,8 @@ from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError import litellm from litellm._internal_context import is_internal_call, post_response_phase from litellm._logging import print_verbose, verbose_logger -from litellm.caching import InMemoryCache -from litellm.caching.caching import S3Cache, response_cache_phase +from litellm.caching import DualCache, InMemoryCache +from litellm.caching.caching import Cache, RedisCache, S3Cache, response_cache_phase from litellm.constants import CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS, RESPONSE_CACHE_EXCLUDED_PROVIDERS from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( @@ -118,6 +118,46 @@ def _is_response_cache_excluded(model: str | None, kwargs: Mapping[str, object]) ) +_CACHE_CONTROLS: Final = TypeAdapter(Mapping[str, object]) + + +def _cache_controls(kwargs: Mapping[str, object]) -> Mapping[str, object]: + controls: Final = kwargs.get("cache") + return _CACHE_CONTROLS.validate_python(controls) if isinstance(controls, Mapping) else {} + + +def is_call_type_cached(call_type: str) -> bool: + if litellm.cache is None or litellm.cache.supported_call_types is None: + return False + covering_call_types: Final = ("aresponses", "responses") if call_type == "aresponses" else (call_type,) + return any(name in litellm.cache.supported_call_types for name in covering_call_types) + + +def _reads_requested(caching: object, controls: Mapping[str, object]) -> bool: + return ((caching is None and litellm.cache is not None) or caching is True) and controls.get("no-cache") is not True + + +def cache_reads_requested(kwargs: Mapping[str, object]) -> bool: + return _reads_requested(kwargs.get("caching"), _cache_controls(kwargs)) + + +def cache_access(model: str | None, kwargs: Mapping[str, object], call_type: str) -> tuple[bool, bool]: + if _is_response_cache_excluded(model=model, kwargs=kwargs) or not is_call_type_cached(call_type): + return (False, False) + controls: Final = _cache_controls(kwargs) + return (_reads_requested(kwargs.get("caching"), controls), controls.get("no-store") is not True) + + +def may_store_cache(model: str | None, kwargs: Mapping[str, object], call_type: str) -> bool: + return cache_access(model, kwargs, call_type)[1] + + +def response_memory_tier(cache: Cache) -> DualCache | None: + if not isinstance(cache.cache, RedisCache): + return None + return DualCache(redis_cache=cache.cache, in_memory_cache=in_memory_cache_obj) + + def _is_chat_completion_cached_dict(cached_result: dict) -> bool: cached_id: Final = cached_result.get("id") if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"): @@ -222,9 +262,17 @@ def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "as return task -def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None: - """Read the caller-supplied ``cache_key`` off the request kwargs.""" - return request_kwargs.get("cache_key", None) +def cache_key_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]: + return { + name: value + for name, value in kwargs.items() + if name != "cache_key" and not (name == "metadata" and value is None) + } + + +def response_cache_key(cache: Cache, kwargs: Mapping[str, object]) -> object: + preset: Final = kwargs.get("cache_key") + return preset if preset else cache.get_cache_key(**cache_key_kwargs(kwargs)) def _set_cached_hidden_param(response: object, key: str, value: object) -> None: @@ -274,21 +322,13 @@ class LLMCachingHandler: request_kwargs: dict[str, object], start_time: datetime.datetime, ): - from litellm.caching import DualCache, RedisCache - self.async_streaming_chunks: list[object] = [] self.sync_streaming_chunks: list[object] = [] self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs) - self.preset_cache_key: str | None = None + self.preset_cache_key: object = None self.original_function = original_function self.start_time = start_time - if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache): - self.dual_cache: DualCache | None = DualCache( - redis_cache=litellm.cache.cache, - in_memory_cache=in_memory_cache_obj, - ) - else: - self.dual_cache = None + self.dual_cache: DualCache | None = None if litellm.cache is None else response_memory_tier(litellm.cache) async def async_get_cache( self, @@ -325,11 +365,7 @@ class LLMCachingHandler: return None # Check if caching should be performed BEFORE doing expensive operations - if ( - (kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) is True - ) and ( - kwargs.get("cache", {}).get("no-cache", False) is not True - ): # allow users to control returning cached responses from the completion function + if cache_reads_requested(kwargs): args = args or () final_embedding_cached_response: EmbeddingResponse | None = None embedding_all_elements_cache_hit: bool = False @@ -908,12 +944,10 @@ class LLMCachingHandler: if all(result is None for result in cached_result): cached_result = None else: - request_kwargs: Final = new_kwargs.copy() - request_cache_key: Final = _request_cache_key(request_kwargs) - request_kwargs.pop("cache_key", None) + request_kwargs: Final = cache_key_kwargs(new_kwargs) + self.preset_cache_key = response_cache_key(litellm.cache, new_kwargs) if litellm.cache.supports_async() is True: ## check if dual cache is supported ## - self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs) with response_cache_phase("get"): cached_result = await litellm.cache.async_get_cache( dynamic_cache_object=self.dual_cache, @@ -921,7 +955,6 @@ class LLMCachingHandler: **request_kwargs, ) else: # fallback for caches that don't support async - self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs) with response_cache_phase("get"): cached_result = litellm.cache.get_cache( dynamic_cache_object=self.dual_cache, @@ -1239,20 +1272,8 @@ class LLMCachingHandler: Returns: bool: True if the result should be stored in the cache, False otherwise. """ - model_value: Final = kwargs.get("model") - model: Final = model_value if isinstance(model_value, str) else None - cache_options_value: Final = kwargs.get("cache") - cache_options: Final[Mapping[str, object]] = ( - TypeAdapter(Mapping[str, object]).validate_python(cache_options_value) - if isinstance(cache_options_value, Mapping) - else {} - ) - no_store: Final = cache_options.get("no-store", False) - return ( - not _is_response_cache_excluded(model=model, kwargs=kwargs) - and self._is_call_type_supported_by_cache(original_function=original_function) - and no_store is not True - ) + model: Final = kwargs.get("model") + return may_store_cache(model if isinstance(model, str) else None, kwargs, str(original_function.__name__)) _should_store_result_in_cache = should_store_result_in_cache @@ -1290,11 +1311,7 @@ class LLMCachingHandler: Returns: bool: True if the call type is supported by the cache, False otherwise. """ - if litellm.cache is None or litellm.cache.supported_call_types is None: - return False - call_type: Final = str(original_function.__name__) - covering_call_types: Final = ("aresponses", "responses") if call_type == "aresponses" else (call_type,) - return any(name in litellm.cache.supported_call_types for name in covering_call_types) + return is_call_type_cached(str(original_function.__name__)) async def add_streaming_response_to_cache(self, processed_chunk: ModelResponse) -> None: """ diff --git a/litellm/chat_completions/dispatch.py b/litellm/chat_completions/dispatch.py index 0197e3bb1ee..268c5394524 100644 --- a/litellm/chat_completions/dispatch.py +++ b/litellm/chat_completions/dispatch.py @@ -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, diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index f0fdfaecddc..bf8c93571d1 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -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, diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py index f857eb5010f..7145e2a1b83 100644 --- a/litellm/responses/dispatch.py +++ b/litellm/responses/dispatch.py @@ -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, diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index b84af456dba..7f6d560d1db 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -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), diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index b3ef8530871..4404bc39230 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -443,6 +443,10 @@ def _request_success_callbacks(request: NativeCall) -> tuple[object, ...]: return (*supplied, *(dynamic or ())) +def deployment_success_hook() -> bool: + return _overridden("async_post_call_success_deployment_hook", litellm.callbacks) + + def pre_request_hook(_request: NativeCall) -> bool: return _overridden("async_pre_request_hook", litellm.callbacks) @@ -471,11 +475,6 @@ def websearch_interception(_request: NativeCall) -> bool: return any(isinstance(callback, WebSearchInterceptionLogger) for callback in litellm.callbacks) -def response_cache(request: NativeCall) -> bool: - cache: Final = optional_mapping(request.resolved.get("cache")) or {} - return litellm.cache is not None and cache.get("no-cache") is not True - - _NATIVE_ECDH_CURVES: Final = frozenset({"x25519", "prime256v1", "secp256r1", "p-256", "secp384r1", "p-384"}) diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 9b751a1ada2..08ee62fc7dc 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -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/.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 diff --git a/litellm/rust_bridge/host/cache.py b/litellm/rust_bridge/host/cache.py index cac4b75677f..a3b964335c6 100644 --- a/litellm/rust_bridge/host/cache.py +++ b/litellm/rust_bridge/host/cache.py @@ -1,12 +1,229 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import re +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter + +from litellm.types.utils import CallTypes if TYPE_CHECKING: from litellm.caching.caching import Cache + from litellm.caching.dual_cache import DualCache -def configured_cache() -> Cache | None: +class ResponseOutput(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["Response"] = "Response" + value: object + + +class StreamOutput(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["Stream"] = "Stream" + value: str + + +CachedOutput: TypeAlias = ResponseOutput | StreamOutput +_OUTPUT: Final = TypeAdapter[CachedOutput](Annotated[CachedOutput, Field(discriminator="kind")]) +_ENTRY: Final = TypeAdapter(dict[str, object]) +_EVENTS: Final = TypeAdapter(list[str]) +_SSE_EVENT_BOUNDARY: Final = re.compile(r"(?<=\n\n)") + + +@dataclass(frozen=True, slots=True) +class EntryFormat: + entry: Callable[[CachedOutput], object | None] + output: Callable[[object, bool], CachedOutput | None] + tiers_responses: bool + tiers_streams: bool + + def tiered(self, output: CachedOutput) -> bool: + match output: + case ResponseOutput(): + return self.tiers_responses + case StreamOutput(): + return self.tiers_streams + + +def _response_entry(output: CachedOutput) -> object | None: + match output: + case ResponseOutput(value=response): + return response + case StreamOutput(): + return None + + +def _response_output(entry: object, stream: bool) -> CachedOutput | None: + return None if stream else ResponseOutput(value=entry) + + +def _messages_entry(output: CachedOutput) -> object: + from litellm.types.caching import CACHED_STREAM_EVENTS_KEY + + match output: + case ResponseOutput(value=response): + return response + case StreamOutput(value=text): + return {CACHED_STREAM_EVENTS_KEY: [event for event in _SSE_EVENT_BOUNDARY.split(text) if event]} + + +def _stream_events(entry: object) -> list[str] | None: + from litellm.types.caching import CACHED_STREAM_EVENTS_KEY + + if not isinstance(entry, dict): + return None + events: Final = _ENTRY.validate_python(entry).get(CACHED_STREAM_EVENTS_KEY) + return None if events is None else _EVENTS.validate_python(events) + + +def _messages_output(entry: object, stream: bool) -> CachedOutput: + events: Final = _stream_events(entry) + if events is None: + return ResponseOutput(value=entry) + return StreamOutput(value="".join(events)) + + +RESPONSE_ENTRY: Final = EntryFormat( + entry=_response_entry, output=_response_output, tiers_responses=True, tiers_streams=False +) +MESSAGES_ENTRY: Final = EntryFormat( + entry=_messages_entry, output=_messages_output, tiers_responses=False, tiers_streams=True +) + + +def entry_format(call_type: CallTypes) -> EntryFormat: + match call_type: + case CallTypes.anthropic_messages: + return MESSAGES_ENTRY + case _: + return RESPONSE_ENTRY + + +@dataclass(frozen=True, slots=True) +class CacheAccess: + cache: Cache + request: Mapping[str, object] + key: object + read: bool + write: bool + memory_tier: DualCache | None + format: EntryFormat + + +def _allowed(kwargs: Mapping[str, object], call_type: CallTypes) -> tuple[bool, bool]: + from litellm._logging import verbose_logger + from litellm.caching.caching_handler import cache_access + + model: Final = kwargs.get("model") + try: + return cache_access(model if isinstance(model, str) else None, kwargs, call_type) + except Exception: # noqa: BLE001 # a cache failure is a miss, never a request failure + verbose_logger.exception("LiteLLM Cache: deciding cache access failed") + return (False, False) + + +def _key(cache: Cache, kwargs: Mapping[str, object]) -> object: + from litellm._logging import verbose_logger + from litellm.caching.caching_handler import cache_key_kwargs, response_cache_key + + try: + if cache.should_use_cache(**cache_key_kwargs(kwargs)) is not True: + return None + return response_cache_key(cache, kwargs) + except Exception: # noqa: BLE001 # a cache failure is a miss, never a request failure + verbose_logger.exception("LiteLLM Cache: preparing the cache key failed") + return None + + +def grant(cache: Cache, kwargs: Mapping[str, object], call_type: CallTypes) -> CacheAccess | None: + from litellm.caching.caching_handler import cache_key_kwargs, response_memory_tier + + read, write = _allowed(kwargs, call_type) + if not (read or write): + return None + key: Final = _key(cache, kwargs) + if key is None: + return None + return CacheAccess( + cache=cache, + request=cache_key_kwargs(kwargs), + key=key, + read=read, + write=write, + memory_tier=response_memory_tier(cache), + format=entry_format(call_type), + ) + + +def access(kwargs: Mapping[str, object], call_type: str) -> CacheAccess | None: import litellm - return litellm.cache + cache: Final = litellm.cache + return None if cache is None else grant(cache, kwargs, CallTypes(call_type)) + + +@dataclass(frozen=True, slots=True) +class CacheHit: + response: Mapping[str, object] + key: str + + +def _hit(access: CacheAccess, entry: object) -> CacheHit | None: + output: Final = access.format.output(entry, access.request.get("stream") is True) + return None if output is None else CacheHit(response=output.model_dump(), key=str(access.key)) + + +def lookup(access: CacheAccess) -> CacheHit | None: + from litellm.caching.caching_handler import is_response_without_output + + result: Final = access.cache.get_cache(dynamic_cache_object=None, cache_key=access.key, **access.request) # pyright: ignore[reportUnknownMemberType] # Cache.get_cache kwargs are untyped + if result is None or is_response_without_output(result): + return None + return _hit(access, result) + + +async def alookup(access: CacheAccess) -> CacheHit | None: + from litellm.caching.caching_handler import is_response_without_output + + result: Final = await access.cache.async_get_cache( # pyright: ignore[reportUnknownMemberType] # Cache.async_get_cache kwargs are untyped + dynamic_cache_object=access.memory_tier, cache_key=access.key, **access.request + ) + if result is None: + return None + if is_response_without_output(result): + if access.memory_tier is not None and isinstance(access.key, str): + access.memory_tier.in_memory_cache.delete_cache(access.key) + return None + return _hit(access, result) + + +def _entry(access: CacheAccess, value: object) -> tuple[object, bool] | None: + from litellm.caching.caching_handler import is_response_without_output + + output: Final = _OUTPUT.validate_python(value) + entry: Final = access.format.entry(output) + if entry is None or is_response_without_output(entry): + return None + return entry, access.format.tiered(output) + + +def store(access: CacheAccess, value: object) -> None: + prepared: Final = _entry(access, value) + if prepared is not None: + access.cache.add_cache(prepared[0], cache_key=access.key, **access.request) + + +async def astore(access: CacheAccess, value: object) -> None: + from litellm.caching.caching_handler import create_cache_write_task + + prepared: Final = _entry(access, value) + if prepared is None: + return + entry, tiered = prepared + tier: Final = access.memory_tier if tiered else None + create_cache_write_task( + lambda: access.cache.async_add_cache(entry, dynamic_cache_object=tier, cache_key=access.key, **access.request) + ) diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index a82e734b47d..54bf6ae5430 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -322,6 +322,12 @@ def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitab return hook(kwargs, call_type) +def reviews_deployment_success() -> bool: + from litellm.rust_bridge.detectors import deployment_success_hook + + return deployment_success_hook() + + def after_deployment_success(kwargs: dict[str, object], response: object, call_type: str) -> Awaitable[object]: from litellm import utils from litellm.types.utils import CallTypes diff --git a/tests/test_litellm_rust/cache/__init__.py b/tests/test_litellm_rust/cache/__init__.py deleted file mode 100644 index 8b137891791..00000000000 --- a/tests/test_litellm_rust/cache/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/tests/test_litellm_rust/cache/conftest.py b/tests/test_litellm_rust/cache/conftest.py deleted file mode 100644 index 35f69683986..00000000000 --- a/tests/test_litellm_rust/cache/conftest.py +++ /dev/null @@ -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) diff --git a/tests/test_litellm_rust/cache/support.py b/tests/test_litellm_rust/cache/support.py deleted file mode 100644 index a30f6b68cc6..00000000000 --- a/tests/test_litellm_rust/cache/support.py +++ /dev/null @@ -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 diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py deleted file mode 100644 index d855d212fec..00000000000 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ /dev/null @@ -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 diff --git a/tests/test_litellm_rust/cache/v2/__init__.py b/tests/test_litellm_rust/cache/v2/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm_rust/cache/v2/test_cache.py b/tests/test_litellm_rust/cache/v2/test_cache.py deleted file mode 100644 index 0ba0a17007f..00000000000 --- a/tests/test_litellm_rust/cache/v2/test_cache.py +++ /dev/null @@ -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() diff --git a/tests/test_litellm_rust/chat_completions/conftest.py b/tests/test_litellm_rust/chat_completions/conftest.py index 04f3f64a91e..9256992411d 100644 --- a/tests/test_litellm_rust/chat_completions/conftest.py +++ b/tests/test_litellm_rust/chat_completions/conftest.py @@ -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) diff --git a/tests/test_litellm_rust/chat_completions/test_cache.py b/tests/test_litellm_rust/chat_completions/test_cache.py new file mode 100644 index 00000000000..2a1d3f87b5a --- /dev/null +++ b/tests/test_litellm_rust/chat_completions/test_cache.py @@ -0,0 +1,328 @@ +import asyncio +from collections.abc import Mapping, Sequence +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.caching.caching import Cache, CacheMode +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.types.caching import CachingSupportedCallTypes +from litellm.types.utils import ModelResponse +from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE, MESSAGES_TEXT +from tests.test_litellm_rust.support.caching import ( + Call, + CallerCache, + FailingLookupCache, + FakeRedis, + HeldWriteCache, + assert_cache_hit_keeps_budget_spend_and_counts_usage, + settle_cache_writes, + use_cache, +) +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.results import cache_key, stored_cache_key +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) +EMPTY: Final = {"object": "chat.completion", "choices": []} +OWN_CALL_TYPES: Final = ("acompletion",) + + +def content(response: object) -> object: + assert isinstance(response, ModelResponse) + return response.model_dump()["choices"][0]["message"]["content"] + + +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.usefixtures("cache") +async def test_a_repeated_native_call_is_answered_from_the_python_cache( + call: Call, server: RecordingServer, asynchronous: bool +) -> None: + first: Final = await call(asynchronous=asynchronous) + second: Final = await call(asynchronous=asynchronous) + + assert_served_natively(server, first) + assert cache_key(first) is None + assert isinstance(cache_key(second), str) + assert content(second) == content(first) == MESSAGES_TEXT + + +@pytest.mark.parametrize( + ("first", "second"), + ( + ({"cache": {"no-store": True}}, {}), + ({}, {"cache": {"no-cache": True}}), + ({}, {"caching": False}), + ({}, {"messages": [{"role": "user", "content": "Goodbye"}]}), + ({}, {"temperature": 0.5}), + ({"cache": {"namespace": "first"}}, {"cache": {"namespace": "second"}}), + ({"metadata": {"redis_namespace": "first"}}, {"metadata": {"redis_namespace": "second"}}), + ), + ids=( + "no-store", + "no-cache", + "caching-off-read", + "other-messages", + "other-param", + "cache-namespace", + "metadata-namespace", + ), +) +@pytest.mark.usefixtures("cache") +async def test_cache_controls_and_keys_keep_the_second_call_on_the_provider( + call: Call, server: RecordingServer, first: Mapping[str, object], second: Mapping[str, object] +) -> None: + server.expect(2) + + await call(**first) + response: Final = await call(**second) + + assert_served_natively(server, count=2) + assert cache_key(response) is None + + +@pytest.mark.usefixtures("cache") +async def test_caching_off_skips_the_read_but_still_writes(call: Call, server: RecordingServer) -> None: + await call(caching=False) + hit: Final = await call() + + assert len(server.requests) == 1 + assert content(hit) == MESSAGES_TEXT + assert_served_natively(server, count=1) + + +@pytest.mark.usefixtures("cache") +async def test_each_namespace_keeps_its_own_entry(call: Call, server: RecordingServer) -> None: + server.expect(2) + + await call(cache={"namespace": "first"}) + await call(cache={"namespace": "second"}) + hit: Final = await call(cache={"namespace": "first"}) + + assert len(server.requests) == 2 + assert isinstance(cache_key(hit), str) + assert_served_natively(server, count=2) + + +@pytest.mark.usefixtures("cache") +async def test_the_key_is_built_from_the_call_arguments_not_the_wire_body(call: Call, server: RecordingServer) -> None: + class Rewrite(RecordingLogger): + temperature = 0.1 + + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + request_body(kwargs)["temperature"] = self.temperature + super().log_pre_api_call(model, messages, kwargs) + + logger: Final = Rewrite() + + await call(callbacks=[logger]) + logger.temperature = 0.8 + hit: Final = await call(callbacks=[logger]) + + assert len(server.requests) == 1 + assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.1 + assert isinstance(cache_key(hit), str) + assert_served_natively(server, count=1) + + +@pytest.mark.parametrize( + ("supported", "expected_requests"), ((["embedding"], 2), (OWN_CALL_TYPES, 1)), ids=("other-type", "own-type") +) +async def test_supported_call_types_gate_native_reads_and_writes( + call: Call, + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + supported: Sequence[CachingSupportedCallTypes], + expected_requests: int, +) -> None: + use_cache(monkeypatch, Cache(supported_call_types=list(supported))) + server.expect(expected_requests) + + await call() + second: Final = await call() + + assert len(server.requests) == expected_requests + assert isinstance(cache_key(second), str) == (expected_requests == 1) + assert_served_natively(server, count=expected_requests) + + +@pytest.mark.parametrize( + ("options", "expected_requests"), (({}, 2), ({"cache": {"use-cache": True}}, 1)), ids=("plain", "opt-in") +) +async def test_default_off_mode_caches_only_calls_that_opt_in( + call: Call, + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + options: Mapping[str, object], + expected_requests: int, +) -> None: + use_cache(monkeypatch, Cache(mode=CacheMode.default_off)) + server.expect(expected_requests) + + await call(**options) + second: Final = await call(**options) + + assert len(server.requests) == expected_requests + assert isinstance(cache_key(second), str) == (expected_requests == 1) + assert_served_natively(server, count=expected_requests) + + +@pytest.mark.usefixtures("cache") +async def test_a_native_hit_is_logged_as_a_free_cache_hit(call: Call, server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + await call(callbacks=[recorder]) + await call(callbacks=[recorder]) + await drain_logging() + + successes: Final = await recorder.wait_for_async("async_log_success_event", count=2) + miss, hit = (_OBJECT.validate_python(event.kwargs) for event in successes) + assert len(server.requests) == 1 + assert miss["cache_hit"] is not True + assert TypeAdapter(float).validate_python(miss["response_cost"]) > 0 + assert hit["cache_hit"] is True + assert hit["response_cost"] == 0 + assert_served_natively(server, count=1) + + +async def test_python_cache_lookups_run_in_the_callers_task( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + backend: Final = CallerCache() + use_cache(monkeypatch, Cache(_backend=backend)) + + await call() + hit: Final = await call() + + assert len(server.requests) == 1 + assert isinstance(cache_key(hit), str) + assert backend.lookup_tasks == [asyncio.current_task(), asyncio.current_task()] + assert_served_natively(server, count=1) + + +async def test_cancelling_a_pending_lookup_sends_and_stores_nothing( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + backend: Final = CallerCache(entered, release) + use_cache(monkeypatch, Cache(_backend=backend)) + server.expect(0) + + task: Final = asyncio.create_task(call()) + await asyncio.wait_for(entered.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + release.set() + await asyncio.sleep(0) + + assert server.requests == [] + assert backend.storage.cache_dict == {} + assert_served_natively(server, count=0) + + +@pytest.mark.usefixtures("cache") +async def test_a_cache_hit_keeps_budget_spend_and_counts_usage(call: Call, server: RecordingServer) -> None: + usage: Final = TypeAdapter(dict[str, int]).validate_python(MESSAGES_RESPONSE["usage"]) + + await assert_cache_hit_keeps_budget_spend_and_counts_usage( + call, server, usage["input_tokens"] + usage["output_tokens"] + ) + assert_served_natively(server, count=1) + + +async def test_a_stored_response_without_output_is_a_miss( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + backend: Final = InMemoryCache() + cache: Final = use_cache(monkeypatch, Cache(_backend=backend)) + await call() + (key,) = backend.cache_dict + cache.add_cache(EMPTY, cache_key=key) + server.expect(2) + + response: Final = await call() + + assert len(server.requests) == 2 + assert content(response) == MESSAGES_TEXT + assert_served_natively(server, count=2) + + +async def test_a_failing_cache_lookup_runs_the_call_uncached( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + use_cache(monkeypatch, FailingLookupCache()) + server.expect(2) + + first: Final = await call() + await call() + + assert content(first) == MESSAGES_TEXT + assert len(server.requests) == 2 + assert_served_natively(server, count=2) + + +async def test_the_response_returns_before_the_cache_write_lands( + call: Call, native_call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + backend: Final = HeldWriteCache() + use_cache(monkeypatch, Cache(_backend=backend)) + + first: Final = await asyncio.wait_for(native_call(), timeout=5) + backend.release.set() + await settle_cache_writes() + hit: Final = await call() + + assert content(first) == MESSAGES_TEXT + assert len(server.requests) == 1 + assert content(hit) == MESSAGES_TEXT + assert_served_natively(server, count=1) + + +async def test_a_redis_hit_fills_the_worker_memory_tier( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis +) -> None: + use_cache(monkeypatch, fake_redis.cache()) + + await call() + await call() + fake_redis.client.flushall() + hit: Final = await call() + + assert len(server.requests) == 1 + assert content(hit) == MESSAGES_TEXT + assert_served_natively(server, count=1) + + +async def test_a_stored_response_lands_in_the_worker_memory_tier( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis +) -> None: + use_cache(monkeypatch, fake_redis.cache()) + + first: Final = await call() + fake_redis.client.flushall() + hit: Final = await call() + + assert len(server.requests) == 1 + assert content(hit) == content(first) + assert_served_natively(server, count=1) + + +async def test_a_hit_reports_the_key_the_entry_is_stored_under( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + backend: Final = InMemoryCache() + use_cache(monkeypatch, Cache(_backend=backend)) + + await call() + hit: Final = await call() + + (stored,) = backend.cache_dict + assert len(server.requests) == 1 + assert stored_cache_key(hit) == stored + assert_served_natively(server, count=1) diff --git a/tests/test_litellm_rust/chat_completions/test_dispatch.py b/tests/test_litellm_rust/chat_completions/test_dispatch.py index 69d9e8e12df..07fad33fd5f 100644 --- a/tests/test_litellm_rust/chat_completions/test_dispatch.py +++ b/tests/test_litellm_rust/chat_completions/test_dispatch.py @@ -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}), diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index c1d452e838e..91de0600fa4 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -13,6 +13,15 @@ from litellm import utils from litellm.litellm_core_utils import litellm_logging, thread_pool_executor from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.rust_bridge.configuration import _CONFIGURATION +from tests.test_litellm_rust.support.caching import ( + cache as cache, +) +from tests.test_litellm_rust.support.caching import ( + call as call, +) +from tests.test_litellm_rust.support.caching import ( + fake_redis as fake_redis, +) from tests.test_litellm_rust.support.callback_recorder import drain_logging from tests.test_litellm_rust.support.clickhouse import clickhouse_url as clickhouse_url from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound diff --git a/tests/test_litellm_rust/messages/conftest.py b/tests/test_litellm_rust/messages/conftest.py index 18a303dcd69..23ed7abbe78 100644 --- a/tests/test_litellm_rust/messages/conftest.py +++ b/tests/test_litellm_rust/messages/conftest.py @@ -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) diff --git a/tests/test_litellm_rust/messages/test_cache.py b/tests/test_litellm_rust/messages/test_cache.py new file mode 100644 index 00000000000..1784581444b --- /dev/null +++ b/tests/test_litellm_rust/messages/test_cache.py @@ -0,0 +1,497 @@ +import asyncio +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.caching.caching import Cache, CacheMode +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.caching import CachingSupportedCallTypes, LiteLLMCacheType +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.utils import LLMResponseTypes +from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE +from tests.test_litellm_rust.support.caching import ( + Call, + CallerCache, + FailingLookupCache, + FakeRedis, + HeldWriteCache, + assert_cache_hit_keeps_budget_spend_and_counts_usage, + settle_cache_writes, + use_cache, +) +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.results import cache_key, consume, response_id, stored_cache_key +from tests.test_litellm_rust.support.routing import assert_served_natively, routed + +pytestmark = pytest.mark.requires_rust_extension + +_OBJECT: Final = TypeAdapter(dict[str, object]) +_MESSAGE: Final = TypeAdapter(AnthropicMessagesResponse) +_CHUNKS: Final = TypeAdapter(list[bytes]) +STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS) +EMPTY: Final = {**MESSAGES_RESPONSE, "content": []} +OWN_CALL_TYPES: Final = ("anthropic_messages",) + + +def body(response: object) -> Mapping[str, object]: + return {name: value for name, value in _OBJECT.validate_python(response).items() if name != "_hidden_params"} + + +async def collect(stream: object) -> bytes: + streamed_async: Final = isinstance(stream, AsyncIterator) + collected: Final = _CHUNKS.validate_python(list(await consume(stream)), strict=True) + assert streamed_async + await settle_cache_writes() + return b"".join(collected) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.usefixtures("cache") +async def test_a_repeated_native_call_is_answered_from_the_python_cache( + call: Call, server: RecordingServer, asynchronous: bool +) -> None: + first: Final = await call(asynchronous=asynchronous) + second: Final = await call(asynchronous=asynchronous) + + assert_served_natively(server, first) + assert cache_key(first) is None + assert isinstance(cache_key(second), str) + assert body(second) == body(first) == MESSAGES_RESPONSE + + +@pytest.mark.usefixtures("cache") +async def test_a_completed_native_stream_is_replayed_from_the_python_cache(call: Call, server: RecordingServer) -> None: + server.serve(STREAM) + + first_stream: Final = await call(stream=True) + assert cache_key(first_stream) is None + first: Final = await collect(first_stream) + second_stream: Final = await call(stream=True) + assert isinstance(cache_key(second_stream), str) + second: Final = await collect(second_stream) + + assert len(server.requests) == 1 + assert first == b"".join(STREAM.payloads()) + assert second == first + assert_served_natively(server, count=1) + + +@pytest.mark.usefixtures("cache") +async def test_prompts_passed_positionally_get_their_own_entries(server: RecordingServer) -> None: + from litellm.messages import anthropic_messages + + server.expect(2) + + async def ask(content: str) -> object: + return await anthropic_messages( + 64, + [{"role": "user", "content": content}], + MESSAGES_MODEL, + api_key="test-key", + api_base=server.base_url, + ) + + await ask("first prompt") + second: Final = await ask("second prompt") + await settle_cache_writes() + + assert len(server.requests) == 2 + assert cache_key(second) is None + assert_served_natively(server, count=2) + + +@pytest.mark.parametrize("native_writes", (True, False), ids=("native-writes", "python-writes")) +@pytest.mark.parametrize( + "options", + ({}, {"metadata": None}, {"cache_key": "shared-entry"}), + ids=("plain", "metadata-none", "preset-key"), +) +@pytest.mark.usefixtures("cache") +async def test_native_and_python_replay_each_others_stream_entries( + call: Call, server: RecordingServer, native_writes: bool, options: Mapping[str, object] +) -> None: + server.serve(STREAM) + + async def stream(native: bool) -> bytes: + if native: + return await collect(await call(stream=True, **options)) + with routed(): + return await collect(await call(stream=True, **options)) + + first: Final = await stream(native_writes) + second: Final = await stream(not native_writes) + + assert len(server.requests) == 1 + assert second == first == b"".join(STREAM.payloads()) + + +@pytest.mark.usefixtures("cache") +async def test_a_stream_closed_early_is_not_stored(call: Call, server: RecordingServer) -> None: + server.serve(STREAM) + server.expect(2) + + stream: Final = await call(stream=True) + assert isinstance(stream, AsyncIterator) + await anext(stream) + await stream.aclose() + await collect(await call(stream=True)) + + assert len(server.requests) == 2 + assert_served_natively(server, count=2) + + +@pytest.mark.parametrize( + ("first", "second"), + ( + ({"cache": {"no-store": True}}, {}), + ({}, {"cache": {"no-cache": True}}), + ({}, {"caching": False}), + ({}, {"messages": [{"role": "user", "content": "Goodbye"}]}), + ({}, {"system": "Be brief"}), + ({"cache": {"namespace": "first"}}, {"cache": {"namespace": "second"}}), + ({"metadata": {"redis_namespace": "first"}}, {"metadata": {"redis_namespace": "second"}}), + ), + ids=( + "no-store", + "no-cache", + "caching-off-read", + "other-messages", + "other-param", + "cache-namespace", + "metadata-namespace", + ), +) +@pytest.mark.usefixtures("cache") +async def test_cache_controls_and_keys_keep_the_second_call_on_the_provider( + call: Call, server: RecordingServer, first: Mapping[str, object], second: Mapping[str, object] +) -> None: + server.expect(2) + + await call(**first) + response: Final = await call(**second) + + assert_served_natively(server, count=2) + assert cache_key(response) is None + + +@pytest.mark.usefixtures("cache") +async def test_caching_off_skips_the_read_but_still_writes(call: Call, server: RecordingServer) -> None: + await call(caching=False) + hit: Final = await call() + + assert len(server.requests) == 1 + assert body(hit) == MESSAGES_RESPONSE + assert_served_natively(server, count=1) + + +@pytest.mark.usefixtures("cache") +async def test_each_namespace_keeps_its_own_entry(call: Call, server: RecordingServer) -> None: + server.expect(2) + + await call(cache={"namespace": "first"}) + await call(cache={"namespace": "second"}) + hit: Final = await call(cache={"namespace": "first"}) + + assert len(server.requests) == 2 + assert isinstance(cache_key(hit), str) + assert_served_natively(server, count=2) + + +@pytest.mark.usefixtures("cache") +async def test_the_key_is_built_from_the_call_arguments_not_the_wire_body(call: Call, server: RecordingServer) -> None: + class Rewrite(RecordingLogger): + temperature = 0.1 + + def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None: + request_body(kwargs)["temperature"] = self.temperature + super().log_pre_api_call(model, messages, kwargs) + + logger: Final = Rewrite() + + await call(callbacks=[logger]) + logger.temperature = 0.8 + hit: Final = await call(callbacks=[logger]) + + assert len(server.requests) == 1 + assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.1 + assert isinstance(cache_key(hit), str) + assert_served_natively(server, count=1) + + +@pytest.mark.parametrize( + ("supported", "expected_requests"), ((["embedding"], 2), (OWN_CALL_TYPES, 1)), ids=("other-type", "own-type") +) +async def test_supported_call_types_gate_native_reads_and_writes( + call: Call, + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + supported: Sequence[CachingSupportedCallTypes], + expected_requests: int, +) -> None: + use_cache(monkeypatch, Cache(supported_call_types=list(supported))) + server.expect(expected_requests) + + await call() + second: Final = await call() + + assert len(server.requests) == expected_requests + assert isinstance(cache_key(second), str) == (expected_requests == 1) + assert_served_natively(server, count=expected_requests) + + +@pytest.mark.parametrize( + ("options", "expected_requests"), (({}, 2), ({"cache": {"use-cache": True}}, 1)), ids=("plain", "opt-in") +) +async def test_default_off_mode_caches_only_calls_that_opt_in( + call: Call, + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + options: Mapping[str, object], + expected_requests: int, +) -> None: + use_cache(monkeypatch, Cache(mode=CacheMode.default_off)) + server.expect(expected_requests) + + await call(**options) + second: Final = await call(**options) + + assert len(server.requests) == expected_requests + assert isinstance(cache_key(second), str) == (expected_requests == 1) + assert_served_natively(server, count=expected_requests) + + +async def test_a_semantic_python_cache_answers_a_similar_prompt( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + use_cache(monkeypatch, Cache(type=LiteLLMCacheType.REDIS_SEMANTIC, _backend=InMemoryCache())) + + first: Final = await call(messages=[{"role": "user", "content": "hello"}]) + second: Final = await call(messages=[{"role": "user", "content": "hi"}]) + + assert len(server.requests) == 1 + assert cache_key(first) is None + assert isinstance(cache_key(second), str) + assert body(second) == body(first) + assert_served_natively(server, count=1) + + +@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream")) +@pytest.mark.usefixtures("cache") +async def test_a_native_hit_is_logged_as_a_free_cache_hit(call: Call, server: RecordingServer, stream: bool) -> None: + if stream: + server.serve(STREAM) + recorder: Final = RecordingLogger() + + async def send() -> None: + response: Final = await call(stream=stream, callbacks=[recorder]) + if stream: + await collect(response) + + await send() + await send() + await drain_logging() + + successes: Final = await recorder.wait_for_async("async_log_success_event", count=2) + miss, hit = (_OBJECT.validate_python(event.kwargs) for event in successes) + assert miss["cache_hit"] is not True + assert TypeAdapter(float).validate_python(miss["response_cost"]) > 0 + assert hit["cache_hit"] is True + assert hit["response_cost"] == 0 + assert_served_natively(server, count=1) + + +async def test_python_cache_lookups_run_in_the_callers_task( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + backend: Final = CallerCache() + use_cache(monkeypatch, Cache(_backend=backend)) + + await call() + hit: Final = await call() + + assert len(server.requests) == 1 + assert isinstance(cache_key(hit), str) + assert backend.lookup_tasks == [asyncio.current_task(), asyncio.current_task()] + assert_served_natively(server, count=1) + + +async def test_cancelling_a_pending_lookup_sends_and_stores_nothing( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + backend: Final = CallerCache(entered, release) + use_cache(monkeypatch, Cache(_backend=backend)) + server.expect(0) + + task: Final = asyncio.create_task(call()) + await asyncio.wait_for(entered.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + release.set() + await asyncio.sleep(0) + + assert server.requests == [] + assert backend.storage.cache_dict == {} + assert_served_natively(server, count=0) + + +@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream")) +@pytest.mark.usefixtures("cache") +async def test_a_cache_hit_keeps_budget_spend_and_counts_usage( + call: Call, server: RecordingServer, stream: bool +) -> None: + if stream: + server.serve(STREAM) + + async def send(**options: object) -> object: + response: Final = await call(stream=stream, **options) + if stream: + await collect(response) + return response + + usage: Final = TypeAdapter(dict[str, int]).validate_python(MESSAGES_RESPONSE["usage"]) + await assert_cache_hit_keeps_budget_spend_and_counts_usage( + send, server, usage["input_tokens"] + usage["output_tokens"] + ) + assert_served_natively(server, count=1) + + +@pytest.mark.usefixtures("cache") +async def test_a_response_without_output_is_not_stored(call: Call, server: RecordingServer) -> None: + server.serve(ResponseSpec(body=EMPTY)) + server.expect(2) + + await call() + await call() + + assert len(server.requests) == 2 + assert_served_natively(server, count=2) + + +async def test_a_stored_response_without_output_is_a_miss( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + backend: Final = InMemoryCache() + cache: Final = use_cache(monkeypatch, Cache(_backend=backend)) + await call() + (key,) = backend.cache_dict + cache.add_cache(EMPTY, cache_key=key) + server.expect(2) + + response: Final = await call() + + assert len(server.requests) == 2 + assert body(response) == MESSAGES_RESPONSE + assert_served_natively(server, count=2) + + +async def test_a_failing_cache_lookup_runs_the_call_uncached( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + use_cache(monkeypatch, FailingLookupCache()) + server.expect(2) + + first: Final = await call() + await call() + + assert body(first) == MESSAGES_RESPONSE + assert len(server.requests) == 2 + assert_served_natively(server, count=2) + + +async def test_the_response_returns_before_the_cache_write_lands( + call: Call, native_call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + backend: Final = HeldWriteCache() + use_cache(monkeypatch, Cache(_backend=backend)) + + first: Final = await asyncio.wait_for(native_call(), timeout=5) + backend.release.set() + await settle_cache_writes() + hit: Final = await call() + + assert body(first) == MESSAGES_RESPONSE + assert len(server.requests) == 1 + assert body(hit) == MESSAGES_RESPONSE + assert_served_natively(server, count=1) + + +async def test_a_redis_hit_fills_the_worker_memory_tier( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis +) -> None: + use_cache(monkeypatch, fake_redis.cache()) + + await call() + await call() + fake_redis.client.flushall() + hit: Final = await call() + + assert len(server.requests) == 1 + assert body(hit) == MESSAGES_RESPONSE + assert_served_natively(server, count=1) + + +async def test_a_stored_response_reaches_the_worker_memory_tier_only_when_read( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis +) -> None: + use_cache(monkeypatch, fake_redis.cache()) + server.expect(2) + + await call() + fake_redis.client.flushall() + await call() + + assert len(server.requests) == 2 + assert_served_natively(server, count=2) + + +async def test_a_native_hit_reports_the_key_the_entry_is_stored_under( + call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + backend: Final = InMemoryCache() + use_cache(monkeypatch, Cache(_backend=backend)) + + await call() + hit: Final = await call() + + (stored,) = backend.cache_dict + assert cache_key(hit) == stored_cache_key(hit) == stored + assert_served_natively(server, count=1) + + +@pytest.mark.parametrize("served", ("native", "python")) +@pytest.mark.usefixtures("cache") +async def test_a_hit_returns_what_the_success_deployment_hook_handed_the_first_caller( + call: Call, server: RecordingServer, served: str +) -> None: + class Replace(CustomLogger): + calls = 0 + + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: object, call_type: object + ) -> LLMResponseTypes | None: + Replace.calls += 1 + return _MESSAGE.validate_python({**_OBJECT.validate_python(response), "id": f"hooked-{Replace.calls}"}) + + async def ask() -> object: + if served == "native": + return await call() + with routed(): + return await call() + + with rebound(litellm, "callbacks", [Replace()]): + first: Final = await ask() + second: Final = await ask() + + assert len(server.requests) == 1 + assert response_id(first) == "hooked-1" + assert response_id(second) == "hooked-1" + assert Replace.calls == 1 diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 309a0e3776c..b003d356693 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -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}), ( diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index c35f6836144..3e2c5fbbeec 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -444,7 +444,7 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl assert response._hidden_params["response_cost"] is not None assert response._hidden_params["_response_ms"] > 0 assert trace_id_var.get() == "callback-free-parent" - assert probe.deployments == 2 + assert probe.deployments == (2 if failure else 1) assert probe.submissions == probe.enqueues == 0 assert len(created_loggers) == 1 logger: Final = created_loggers[0] diff --git a/tests/test_litellm_rust/responses/conftest.py b/tests/test_litellm_rust/responses/conftest.py index 73b398a4329..f8580d82ca1 100644 --- a/tests/test_litellm_rust/responses/conftest.py +++ b/tests/test_litellm_rust/responses/conftest.py @@ -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) diff --git a/tests/test_litellm_rust/responses/test_cache.py b/tests/test_litellm_rust/responses/test_cache.py new file mode 100644 index 00000000000..ee524b04a28 --- /dev/null +++ b/tests/test_litellm_rust/responses/test_cache.py @@ -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 diff --git a/tests/test_litellm_rust/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py index 226e64b2742..1b60d3a09c7 100644 --- a/tests/test_litellm_rust/responses/test_dispatch.py +++ b/tests/test_litellm_rust/responses/test_dispatch.py @@ -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, { diff --git a/tests/test_litellm_rust/support/caching.py b/tests/test_litellm_rust/support/caching.py new file mode 100644 index 00000000000..7098f3cf085 --- /dev/null +++ b/tests/test_litellm_rust/support/caching.py @@ -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 diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index d7b82767b97..08b2adf4e47 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -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: diff --git a/tests/test_litellm_rust/support/results.py b/tests/test_litellm_rust/support/results.py index 814cebe3efc..97ece1e2a5c 100644 --- a/tests/test_litellm_rust/support/results.py +++ b/tests/test_litellm_rust/support/results.py @@ -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") diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 51754c80dd6..f24e4882178 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -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" diff --git a/tests/unit/rust_bridge/host/test_cache.py b/tests/unit/rust_bridge/host/test_cache.py index 9d532631ef2..e560d3fd669 100644 --- a/tests/unit/rust_bridge/host/test_cache.py +++ b/tests/unit/rust_bridge/host/test_cache.py @@ -1,21 +1,337 @@ +import asyncio from typing import Final import pytest +from pydantic import TypeAdapter import litellm -from litellm.caching.caching import Cache -from litellm.rust_bridge.host.cache import configured_cache -from litellm.types.caching import LiteLLMCacheType +from litellm.caching.caching import Cache, CacheMode +from litellm.caching.caching_handler import ( + _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await background cache writes + cache_key_kwargs, + response_cache_key, +) +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import RESPONSE_CACHE_EXCLUDED_PROVIDERS +from litellm.llms.anthropic.pass_through.messages.response_cache import get_cached_stream_events +from litellm.rust_bridge.host.cache import ( + MESSAGES_ENTRY, + RESPONSE_ENTRY, + CacheAccess, + EntryFormat, + ResponseOutput, + StreamOutput, + access, + alookup, + astore, + lookup, + store, +) +from litellm.types.caching import CachingSupportedCallTypes, LiteLLMCacheType + +REQUEST: Final[dict[str, object]] = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} +STREAM_REQUEST: Final[dict[str, object]] = {**REQUEST, "stream": True} +EMPTY: Final[dict[str, object]] = {"choices": []} +_ENTRY: Final = TypeAdapter(dict[str, object]) +MESSAGES_EVENTS: Final = ( + 'event: message_start\ndata: {"type":"message_start"}\n\n', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0}\n\n', + 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n', + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', +) +RAW_ENTRY: Final = EntryFormat( + entry=RESPONSE_ENTRY.entry, + output=lambda entry, _: ResponseOutput(value=entry), + tiers_responses=True, + tiers_streams=False, +) -def test_configured_cache_preserves_identity_and_observes_runtime_replacement(monkeypatch: pytest.MonkeyPatch) -> None: - initial: Final = Cache(type=LiteLLMCacheType.LOCAL) - replacement: Final = Cache(type=LiteLLMCacheType.LOCAL) - monkeypatch.setattr(litellm, "cache", initial) - assert configured_cache() is initial +class KeyCountingCache(Cache): + key_builds: int = 0 - monkeypatch.setattr(litellm, "cache", replacement) - assert configured_cache() is replacement + def get_cache_key(self, **kwargs: object) -> str: + self.key_builds += 1 + return super().get_cache_key(**kwargs) + +def response_output(value: object) -> dict[str, object]: + return ResponseOutput(value=value).model_dump() + + +def configure(monkeypatch: pytest.MonkeyPatch, cache: Cache) -> Cache: + monkeypatch.setattr(litellm, "cache", cache) + return cache + + +def _cache_access( + cache: Cache, + request: dict[str, object] = REQUEST, + memory_tier: DualCache | None = None, + entry_format: EntryFormat = RESPONSE_ENTRY, +) -> CacheAccess: + return CacheAccess( + cache=cache, + request=cache_key_kwargs(request), + key=response_cache_key(cache, request), + read=True, + write=True, + memory_tier=memory_tier, + format=entry_format, + ) + + +def stored_entry(cache: Cache, request: dict[str, object]) -> object: + hit: Final = lookup(_cache_access(cache, request, entry_format=RAW_ENTRY)) + return None if hit is None else hit.response["value"] + + +async def memory_entry(memory_tier: DualCache, key: object) -> object: + return await memory_tier.in_memory_cache.async_get_cache(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # InMemoryCache.async_get_cache is unannotated + + +def grants(kwargs: dict[str, object], call_type: str) -> tuple[bool, bool] | None: + granted: Final = access(kwargs, call_type) + return None if granted is None else (granted.read, granted.write) + + +def test_no_configured_cache_grants_no_access(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "cache", None) - assert configured_cache() is None + + assert access(REQUEST, "acompletion") is None + + +@pytest.mark.parametrize( + ("overrides", "expected"), + ( + ({}, (True, True)), + ({"cache": {"no-cache": True}}, (False, True)), + ({"cache": {"no-store": True}}, (True, False)), + ({"cache": {"no-cache": "true"}}, (True, True)), + ({"caching": False}, (False, True)), + ({"cache": {"no-cache": True, "no-store": True}}, None), + ), + ids=("default", "no-cache", "no-store", "non-bool-control", "caching-off", "both-off"), +) +def test_access_follows_the_python_handler_controls( + monkeypatch: pytest.MonkeyPatch, overrides: dict[str, object], expected: tuple[bool, bool] | None +) -> None: + configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL)) + + assert grants({**REQUEST, **overrides}, "acompletion") == expected + + +def test_an_excluded_provider_grants_no_access(monkeypatch: pytest.MonkeyPatch) -> None: + configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL)) + provider: Final = next(iter(RESPONSE_CACHE_EXCLUDED_PROVIDERS)) + + assert access({**REQUEST, "custom_llm_provider": provider}, "acompletion") is None + assert access({**REQUEST, "model": f"{provider}/gpt-4o"}, "acompletion") is None + + +@pytest.mark.parametrize( + ("supported", "call_type", "granted"), + ( + (["acompletion"], "acompletion", True), + (["embedding"], "acompletion", False), + (["responses"], "aresponses", True), + (["aresponses"], "responses", False), + ), +) +def test_supported_call_types_gate_access( + monkeypatch: pytest.MonkeyPatch, supported: list[CachingSupportedCallTypes], call_type: str, granted: bool +) -> None: + configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL, supported_call_types=supported)) + + assert (access(REQUEST, call_type) is not None) == granted + + +@pytest.mark.parametrize( + ("overrides", "granted"), + (({}, False), ({"cache": {"use-cache": True}}, True)), + ids=("plain", "opt-in"), +) +def test_default_off_mode_grants_only_calls_that_opt_in_and_hashes_nothing_otherwise( + monkeypatch: pytest.MonkeyPatch, overrides: dict[str, object], granted: bool +) -> None: + cache: Final = KeyCountingCache(type=LiteLLMCacheType.LOCAL, mode=CacheMode.default_off) + configure(monkeypatch, cache) + + assert (access({**REQUEST, **overrides}, "acompletion") is not None) is granted + assert cache.key_builds == int(granted) + + +def test_a_miss_then_store_builds_the_key_once(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = KeyCountingCache(type=LiteLLMCacheType.LOCAL) + configure(monkeypatch, cache) + granted: Final = access(REQUEST, "acompletion") + assert granted is not None + + assert lookup(granted) is None + store(granted, response_output({"id": "cached"})) + hit: Final = lookup(granted) + + assert hit is not None + assert hit.response == response_output({"id": "cached"}) + assert cache.key_builds == 1 + + +def test_a_cache_that_fails_to_decide_grants_no_access(monkeypatch: pytest.MonkeyPatch) -> None: + class FailingCache(Cache): + def should_use_cache(self, **kwargs: object) -> bool: + raise RuntimeError("cache unavailable") + + configure(monkeypatch, FailingCache(type=LiteLLMCacheType.LOCAL)) + + assert access(REQUEST, "acompletion") is None + + +@pytest.mark.parametrize("cache_key", (None, "preset-key"), ids=("generated-key", "preset-key")) +def test_lookup_returns_cached_response_and_key(cache_key: str | None) -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = REQUEST if cache_key is None else {**REQUEST, "cache_key": cache_key} + cache_access: Final = _cache_access(cache, request) + expected_key: Final = cache.get_cache_key(**REQUEST) if cache_key is None else cache_key + + store(cache_access, response_output({"id": "cached"})) + hit: Final = lookup(cache_access) + + assert hit is not None + assert hit.response == response_output({"id": "cached"}) + assert hit.key == expected_key + + +def test_lookup_skips_responses_without_output() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + cache_access: Final = _cache_access(cache) + + store(cache_access, response_output(EMPTY)) + assert lookup(cache_access) is None + + cache.add_cache(EMPTY, **REQUEST) + assert lookup(cache_access) is None + + +@pytest.mark.asyncio +async def test_async_lookup_uses_memory_tier_and_discards_empty_responses() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + memory_tier: Final = DualCache(in_memory_cache=InMemoryCache()) + cache_access: Final = _cache_access(cache, memory_tier=memory_tier) + + assert await alookup(cache_access) is None + + await cache.async_add_cache({"id": "cached"}, dynamic_cache_object=memory_tier, **REQUEST) + hit: Final = await alookup(cache_access) + + assert hit is not None + assert hit.response == response_output({"id": "cached"}) + assert hit.key == cache.get_cache_key(**REQUEST) + + empty_request: Final = {**REQUEST, "cache_key": "empty-response"} + await cache.async_add_cache(EMPTY, dynamic_cache_object=memory_tier, **empty_request) + + assert await alookup(_cache_access(cache, empty_request, memory_tier)) is None + assert await memory_entry(memory_tier, "empty-response") is None + + +@pytest.mark.asyncio +async def test_async_store_routes_writes_and_skips_responses_without_output() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + memory_tier: Final = DualCache(in_memory_cache=InMemoryCache()) + cache_access: Final = _cache_access(cache, memory_tier=memory_tier) + + await astore(cache_access, response_output({"id": "cached"})) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + hit: Final = await alookup(cache_access) + + assert hit is not None + assert hit.response == response_output({"id": "cached"}) + assert hit.key == cache.get_cache_key(**REQUEST) + + empty_request: Final = {**REQUEST, "cache_key": "empty-response"} + empty_access: Final = _cache_access(cache, empty_request, memory_tier) + await astore(empty_access, response_output(EMPTY)) + + assert await alookup(empty_access) is None + assert await memory_entry(memory_tier, "empty-response") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("entry_format", "request_kwargs", "output", "tiered"), + ( + (RESPONSE_ENTRY, REQUEST, ResponseOutput(value={"id": "cached"}), True), + (MESSAGES_ENTRY, REQUEST, ResponseOutput(value={"id": "cached"}), False), + (MESSAGES_ENTRY, STREAM_REQUEST, StreamOutput(value="".join(MESSAGES_EVENTS)), True), + ), + ids=("response", "messages-response", "messages-stream"), +) +async def test_an_async_store_reaches_the_worker_memory_tier_as_the_python_route_does( + entry_format: EntryFormat, + request_kwargs: dict[str, object], + output: ResponseOutput | StreamOutput, + tiered: bool, +) -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + memory_tier: Final = DualCache(in_memory_cache=InMemoryCache()) + cache_access: Final = _cache_access(cache, request_kwargs, memory_tier, entry_format) + + await astore(cache_access, output.model_dump()) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + + assert (await memory_entry(memory_tier, cache_access.key) is not None) is tiered + assert (stored_entry(cache, request_kwargs) is not None) is not tiered + + +@pytest.mark.asyncio +async def test_only_a_redis_backed_cache_answers_from_the_worker_memory_tier(monkeypatch: pytest.MonkeyPatch) -> None: + request: Final = {**REQUEST, "messages": [{"role": "user", "content": "worker tier"}]} + configure(monkeypatch, Cache(type=LiteLLMCacheType.REDIS, host="127.0.0.1", port="1")) + redis_access: Final = access(request, "acompletion") + configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL)) + local_access: Final = access(request, "acompletion") + assert redis_access is not None and local_access is not None + + await astore(redis_access, response_output({"id": "from-the-worker"})) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + hit: Final = await alookup(redis_access) + + assert hit is not None + assert hit.response == response_output({"id": "from-the-worker"}) + assert local_access.memory_tier is None + + +def test_cache_controls_that_cannot_be_read_grant_no_access(monkeypatch: pytest.MonkeyPatch) -> None: + configure(monkeypatch, Cache(type=LiteLLMCacheType.LOCAL)) + + assert access({**REQUEST, "cache": {1: True}}, "acompletion") is None + + +def test_a_messages_stream_is_stored_as_the_legacy_event_list_and_replayed() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + cache_access: Final = _cache_access(cache, STREAM_REQUEST, entry_format=MESSAGES_ENTRY) + stream: Final = StreamOutput(value="".join(MESSAGES_EVENTS)).model_dump() + + store(cache_access, stream) + hit: Final = lookup(cache_access) + + assert get_cached_stream_events(_ENTRY.validate_python(stored_entry(cache, STREAM_REQUEST))) == MESSAGES_EVENTS + assert hit is not None + assert hit.response == stream + + +def test_a_stream_on_a_route_without_a_stream_entry_is_not_stored() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + + store(_cache_access(cache, STREAM_REQUEST), StreamOutput(value="".join(MESSAGES_EVENTS)).model_dump()) + + assert stored_entry(cache, STREAM_REQUEST) is None + + +@pytest.mark.parametrize(("request_kwargs", "hit"), ((REQUEST, True), (STREAM_REQUEST, False)), ids=("unary", "stream")) +def test_a_stored_response_answers_only_a_unary_request(request_kwargs: dict[str, object], hit: bool) -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + cache.add_cache({"id": "cached"}, **request_kwargs) + + assert (lookup(_cache_access(cache, request_kwargs)) is not None) is hit diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index b78d4c85bbf..593a7f4f5f9 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -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)), diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 8d3e341096c..a089a2c3cdd 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -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", } )