mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
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 <yujong@berri.ai>
This commit is contained in:
parent
898fbaf985
commit
33e5d9168e
95 changed files with 3635 additions and 1847 deletions
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -3886,6 +3886,7 @@ dependencies = [
|
|||
"litellm-auth",
|
||||
"litellm-coroutine",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,29 +1,69 @@
|
|||
# Response caching
|
||||
|
||||
Design this crate for shared Rust execution used by the Python SDK and the Rust gateway. The Python SDK will remain, with more core execution moving to Rust and Python callbacks staying in Python. The Rust gateway is still evolving and is intended to replace the Python proxy. Keep response-cache policy independent of Python, HTTP serving, and either proxy's configuration format
|
||||
This crate owns response-cache semantics shared by Rust execution in the Python SDK and in the Rust gateway: what a cache entry is, when it is fresh, how a call's cache policy is expressed, and the contract a cache service offers to core. Keep it independent of Python, HTTP serving and any proxy's configuration format
|
||||
|
||||
Separate what is cached, how a hit is matched, and where entries are stored. Chat Completions, Messages, Responses, and embeddings are API workloads. Exact and semantic matching are lookup behaviors. Memory, Redis, disk, and object stores are storage choices. Embeddings are inference too, so do not use an inference-cache name to imply a category that excludes embeddings. Consult the existing Python cache and caching handler for behavior and compatibility contracts without copying their class structure
|
||||
## Who decides what
|
||||
|
||||
Storage traits, codecs, and backend capabilities belong in `litellm-cache` and the storage crates. Keep storage reusable for value types beyond LLM responses. This crate owns response entries, matching and freshness semantics, the Python-compatible response codec, and deferred-write policy. Core owns route-specific request identity, response encoding and reconstruction, embedding partial-hit orchestration, and stream capture and replay. Boundaries own configuration translation, resource construction, and caller identity
|
||||
Core decides when. It looks up before calling the provider and stores provider output once that output is final, a complete response after the hooks that run once per provider call and a stream once the caller has consumed it. Core never decides whether a call may use the cache, what the key is, or how an entry is laid out
|
||||
|
||||
Construct and inject the response-cache service at the Python bridge or gateway boundary, as with the HTTP client. Reuse it across calls. Core and provider code must not discover cache configuration through Python globals, process configuration, or backend-specific factories
|
||||
The cache service decides what. It decides whether a call may read or write, derives the key, reads and writes storage, and owns the stored layout. When the service delegates to another cache implementation, that implementation owns all four, and Rust only carries values across
|
||||
|
||||
Keep `ResponseCache<B>` generic over its storage backend. Preserve typed backend contexts and capability bounds internally. Inject an object-safe service into core for runtime backend selection, so storage types do not spread through route and host types. Keep API request and response types statically typed. Add a generic parameter only where it preserves a useful type relationship or capability
|
||||
Values crossing that boundary are route-neutral, either a complete response or a stream's bytes. Rust never names a route's stored shape. Each route's stored layout is defined once, in the layer that owns storage
|
||||
|
||||
Keep the core service contract narrow. Lookup and store must not require connection testing, ping, flush, deletion, counters, queues, or scripts. Require batch operations where a consumer needs partial hits, and keep management capabilities on their own interfaces. An exact-only adapter must remain explicit about its matching restriction. Supporting semantic matching requires a defined lookup-context and embedding execution contract, not just a renamed trait
|
||||
## One key for everyone who shares a cache
|
||||
|
||||
Separate reusable resources from per-call policy. Backend configuration, namespace, default expiry, and entry limits belong to the configured service or backend. Read/write controls, expiry and freshness overrides, and authenticated caller scope belong to the call. Passing call options must not replace or mutate the route's configured service
|
||||
Every path that reads or writes the same entries derives the key from the same inputs with the same rule. Two paths that each normalize the call on their own drift apart silently, and nothing fails when they do
|
||||
|
||||
Keep cache misses and storage failures distinguishable in return values. Core owns the decision to continue with provider execution after a cache failure. A read can reject an entry for freshness while the backend still retains it. Preserve the timestamp at which a response was produced when writing it later
|
||||
Key inputs are the call as the caller expressed it, including arguments passed by position. A missing input is a correctness bug, since unrelated requests then share an entry and one caller receives another caller's answer. An extra input, such as an injected default the other paths leave out, splits entries that should be shared. Treat both as key-compatibility breaks
|
||||
|
||||
Define lookup placement explicitly relative to authorization, deployment and credential resolution, and request-transforming callbacks. Cache identity must account for every input that affects reuse, including API surface and caller scope, while preserving intentional Python caching groups. Preserve existing keys and response formats unless changing them is an explicit migration decision
|
||||
Decide whether the cache applies before deriving the key. Disabled caching, opt-in modes, request-size limits and read/write controls must not pay for hashing the request. Derive the key once per call and use that same key for lookup, store and the key reported on a hit
|
||||
|
||||
Cache normalized provider results before caller-specific response transformations. Hits must still run the applicable response processing, success callbacks, and cache-hit accounting. Keep callback execution in the host. Python cache implementations and semantic embedders that require the caller's task must use the existing host-operation mechanism rather than Python calls from a Rust worker. Preserve legacy fallback until that contract is supported
|
||||
Preserve existing keys and stored formats unless a change is an explicit migration decision
|
||||
|
||||
Keep unary caching independent of stream-only methods. Store streams only after successful exhaustion and protocol completion. Errors, incomplete streams, cancellation, and oversized entries must not populate the cache. Embedding batches need ordered partial results and reconstruction around the uncached inputs
|
||||
## Entries
|
||||
|
||||
Test each contract in its owner: storage capabilities in backend tests, envelopes and freshness here, reuse and replay in core, Python callback and fallback behavior at the bridge, and HTTP behavior at the gateway. Run backend contract checks and Python response-codec fixtures before exposing a new backend
|
||||
A hit carries the key that found it. Report that key, never a separately derived one
|
||||
|
||||
`ScopedCache` requires an explicit shared or isolated scope at construction. Per-call `CachePolicy` controls reads, writes, expiry, and freshness without replacing the attached scope or service. `CacheOptions` binds that policy to an explicit scope for storage requests and has no default sharing policy. Versioned native envelopes reject incompatible API surfaces and versions as misses. This envelope is distinct from the legacy Python response codec
|
||||
An entry with no usable output is neither served nor stored. An entry that cannot be decoded, has the wrong surface or version, or does not match what the request asked for is a miss, not an error
|
||||
|
||||
Response storage is not the source of budget or rate-limit coordination dependencies. Keep counters, reservations, and atomic admission operations out of `ResponseCacheService`, including when both services happen to use Redis
|
||||
A stream request is answered with a stream and a complete request with a complete response. Returning the other kind is a miss unless the route defines how to replay one as the other
|
||||
|
||||
Native envelopes are versioned per API surface. They are distinct from any format shared with another cache implementation, which must stay readable and writable by that implementation
|
||||
|
||||
## Streams
|
||||
|
||||
Store a stream only after the caller has consumed it to the end, the protocol reached its terminal event and no error event appeared. Early close, cancellation, errors and incomplete streams store nothing. A replayed stream yields the same bytes the provider sent
|
||||
|
||||
## Failures
|
||||
|
||||
A cache failure is a cache miss, never a request failure. That covers deciding eligibility, deriving the key, lookup and store. Cancellation is not a cache failure and must propagate. A failed store never repeats the provider call
|
||||
|
||||
Keep misses and storage failures distinguishable in return values, and leave the decision to continue with the provider to core. A freshness check may reject an entry the backend still holds. Writes record when the response was produced, not when the write lands. An asynchronous caller gets its response before a background write completes
|
||||
|
||||
## Responses and hooks
|
||||
|
||||
Response hooks sit on two sides of the cache. A hook that runs once per provider call, such as a deployment success hook, sits inside it. It runs before the store, its result is what gets cached, and it does not run again on a hit, since it may execute tools, call providers again or redact what must not be stored. A hook that runs once per request, such as per-caller policy, sits outside it and runs on hits and misses alike
|
||||
|
||||
A hit therefore returns what the first caller received from the inside hooks. Hits still produce success callbacks and cache-hit accounting, and callback execution stays in the host
|
||||
|
||||
## Services, scope and policy
|
||||
|
||||
Construct the cache service at the boundary, from the boundary's configuration, and inject it like the HTTP client. Reuse it across calls. Core and provider code never discover cache configuration from globals or backend factories
|
||||
|
||||
Configured resources and per-call policy stay separate. Backend, namespace, default expiry and entry limits belong to the service. Read and write controls, expiry and freshness overrides and caller scope belong to the call. Call policy never replaces the attached service or its scope, and every scope is explicit, shared or isolated, with no default
|
||||
|
||||
A service's lookup and store need nothing beyond lookup and store. Connection checks, flushing, deletion, counters and scripts live on separate interfaces. Exact and semantic matching are different lookup contracts, and supporting semantic matching needs an embedding execution contract, not a renamed trait
|
||||
|
||||
Implementations that must run in the caller's task, such as a Python cache, run through host operations, never as direct calls from a Rust worker
|
||||
|
||||
Response caching is not budget or rate-limit coordination. Counters, reservations and atomic admission stay out of the cache service, even when both use the same store
|
||||
|
||||
## Layering
|
||||
|
||||
Storage traits, codecs and backend capabilities belong to the storage crates, and stay reusable beyond LLM responses. This crate owns entries, freshness, envelopes, cache policy and the service contract. Core owns route identity, response reconstruction, stream capture and replay, and embedding partial hits. Boundaries own configuration, resource construction and caller identity
|
||||
|
||||
API workloads, lookup behavior and storage are separate axes. Chat Completions, Messages, Responses and embeddings are workloads. Exact and semantic are lookups. Memory, Redis, disk and object stores are storage
|
||||
|
||||
## Tests
|
||||
|
||||
Test each contract where it is owned. Storage capabilities belong in backend tests, entries and freshness here, reuse and replay in core, delegation and callbacks at the bridge, and HTTP behavior at the gateway. Any key or format shared with another implementation needs a test where each side reads what the other wrote
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
use std::{future::Future, pin::Pin, time::Duration};
|
||||
use std::{future::Future, pin::Pin, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{BaseCache, Error, ExactCacheContext};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
|
||||
cache_key,
|
||||
};
|
||||
|
||||
type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
|
|
@ -24,46 +25,78 @@ impl Default for ResponseCacheConfig {
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub struct CacheHit {
|
||||
pub value: Value,
|
||||
pub key: String,
|
||||
}
|
||||
|
||||
pub trait ResponseCacheService: Send + Sync {
|
||||
fn config(&self) -> &ResponseCacheConfig;
|
||||
fn prepare(
|
||||
self: Arc<Self>,
|
||||
options: CacheOptions,
|
||||
surface: &str,
|
||||
input: &dyn Fn() -> Value,
|
||||
) -> PreparedCache;
|
||||
}
|
||||
|
||||
fn lookup<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> CacheFuture<'a, Option<Value>>;
|
||||
pub struct PreparedCache {
|
||||
pub policy: CachePolicy,
|
||||
pub max_entry_bytes: usize,
|
||||
pub call: Arc<dyn ResponseCacheCall>,
|
||||
}
|
||||
|
||||
fn store<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> CacheFuture<'a, ()>;
|
||||
pub trait ResponseCacheCall: Send + Sync {
|
||||
fn lookup<'a>(&'a self, now: Duration) -> CacheFuture<'a, Option<CacheHit>>;
|
||||
|
||||
fn store<'a>(&'a self, response: Value, now: Duration) -> CacheFuture<'a, ()>;
|
||||
}
|
||||
|
||||
impl<B> ResponseCacheService for ResponseCache<B>
|
||||
where
|
||||
B: BaseCache<Value = CacheEntry, Context = ExactCacheContext> + 'static,
|
||||
{
|
||||
fn prepare(
|
||||
self: Arc<Self>,
|
||||
options: CacheOptions,
|
||||
surface: &str,
|
||||
input: &dyn Fn() -> Value,
|
||||
) -> PreparedCache {
|
||||
PreparedCache {
|
||||
policy: options.policy,
|
||||
max_entry_bytes: self.config().max_entry_bytes,
|
||||
call: Arc::new(NativeCacheCall {
|
||||
request: options.request(&self.config().namespace, surface, input()),
|
||||
service: self,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct NativeCacheCall<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
|
||||
service: Arc<ResponseCache<B>>,
|
||||
request: ResponseCacheRequest,
|
||||
}
|
||||
|
||||
impl<B> ResponseCacheCall for NativeCacheCall<B>
|
||||
where
|
||||
B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>,
|
||||
{
|
||||
fn config(&self) -> &ResponseCacheConfig {
|
||||
self.config()
|
||||
fn lookup<'a>(&'a self, now: Duration) -> CacheFuture<'a, Option<CacheHit>> {
|
||||
Box::pin(async move {
|
||||
Ok(self
|
||||
.service
|
||||
.async_lookup(&self.request, now)
|
||||
.await?
|
||||
.map(|value| CacheHit {
|
||||
value,
|
||||
key: cache_key(&self.request.key),
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
fn lookup<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> CacheFuture<'a, Option<Value>> {
|
||||
Box::pin(self.async_lookup(request, now))
|
||||
}
|
||||
|
||||
fn store<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
Box::pin(self.async_store(request, response, now))
|
||||
fn store<'a>(&'a self, response: Value, now: Duration) -> CacheFuture<'a, ()> {
|
||||
Box::pin(self.service.async_store(&self.request, response, now))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -83,8 +116,16 @@ pub struct CachePolicy {
|
|||
}
|
||||
|
||||
impl CachePolicy {
|
||||
pub fn reads(&self) -> bool {
|
||||
self.caching != Some(false) && !self.no_cache
|
||||
}
|
||||
|
||||
pub fn writes(&self) -> bool {
|
||||
self.caching != Some(false) && !self.no_store
|
||||
}
|
||||
|
||||
pub fn enabled(&self) -> bool {
|
||||
self.caching != Some(false) && !(self.no_cache && self.no_store)
|
||||
self.reads() || self.writes()
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -183,3 +224,56 @@ impl ScopedCache {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::cell::Cell;
|
||||
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::shared(CacheScope::Shared)]
|
||||
#[case::isolated(CacheScope::Isolated("tenant".into()))]
|
||||
#[tokio::test]
|
||||
async fn native_preparation_keeps_existing_keys_and_storage(#[case] scope: CacheScope) {
|
||||
let service = Arc::new(
|
||||
ResponseCache::new(Arc::new(InMemoryCache::default())).with_config(
|
||||
ResponseCacheConfig {
|
||||
namespace: "test".into(),
|
||||
max_entry_bytes: 4096,
|
||||
},
|
||||
),
|
||||
);
|
||||
let options = CacheOptions::new(scope);
|
||||
let input = json!({"model": "test", "body": {"z": 1, "a": 2}});
|
||||
let legacy = options.clone().request("test", "test", input.clone());
|
||||
let preparations = Cell::new(0);
|
||||
let prepared = service.clone().prepare(options, "test", &|| {
|
||||
preparations.set(preparations.get() + 1);
|
||||
input.clone()
|
||||
});
|
||||
assert_eq!(prepared.max_entry_bytes, 4096);
|
||||
assert_eq!(prepared.call.lookup(Duration::ZERO).await.unwrap(), None);
|
||||
prepared
|
||||
.call
|
||||
.store(json!({"answer": 42}), Duration::ZERO)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
service.async_lookup(&legacy, Duration::ZERO).await.unwrap(),
|
||||
Some(json!({"answer": 42}))
|
||||
);
|
||||
assert_eq!(
|
||||
prepared.call.lookup(Duration::ZERO).await.unwrap(),
|
||||
Some(CacheHit {
|
||||
value: json!({"answer": 42}),
|
||||
key: cache_key(&legacy.key),
|
||||
})
|
||||
);
|
||||
assert_eq!(preparations.get(), 1);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,11 +6,10 @@ use std::{
|
|||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::ExactCacheContext;
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheConfig, ResponseCacheRequest,
|
||||
ResponseCacheService,
|
||||
CacheEntry, CacheKeyInput, CacheOptions, CachePolicy, CacheScope, ResponseCache,
|
||||
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
|
@ -25,48 +24,61 @@ async fn service_honors_per_call_expiry_and_freshness() {
|
|||
Duration::from_secs(cache_clock.load(Ordering::SeqCst))
|
||||
}),
|
||||
)));
|
||||
let request = ResponseCacheRequest {
|
||||
context: ExactCacheContext {
|
||||
let options = CacheOptions {
|
||||
policy: CachePolicy {
|
||||
ttl: Some(Duration::from_secs(5)),
|
||||
..CachePolicy::default()
|
||||
},
|
||||
..ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some("entry".into()),
|
||||
..Default::default()
|
||||
})
|
||||
scope: CacheScope::Shared,
|
||||
};
|
||||
cache
|
||||
.store(&request, json!({"answer":7}), Duration::ZERO)
|
||||
let key =
|
||||
litellm_cache_response::cache_key(&options.clone().request("", "test", json!("entry")).key);
|
||||
let request = cache
|
||||
.clone()
|
||||
.prepare(options.clone(), "test", &|| json!("entry"));
|
||||
request
|
||||
.call
|
||||
.store(json!({"answer":7}), Duration::ZERO)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::ZERO).await.unwrap(),
|
||||
Some(json!({"answer":7}))
|
||||
request.call.lookup(Duration::ZERO).await.unwrap(),
|
||||
Some(litellm_cache_response::CacheHit {
|
||||
value: json!({"answer":7}),
|
||||
key,
|
||||
})
|
||||
);
|
||||
let stale_request = cache.prepare(
|
||||
CacheOptions {
|
||||
policy: CachePolicy {
|
||||
max_age: Some(Duration::from_secs(1)),
|
||||
..options.policy
|
||||
},
|
||||
..options
|
||||
},
|
||||
"test",
|
||||
&|| json!("entry"),
|
||||
);
|
||||
let stale_request = ResponseCacheRequest {
|
||||
max_age: Some(Duration::from_secs(1)),
|
||||
..request.clone()
|
||||
};
|
||||
clock.store(2, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&stale_request, Duration::from_secs(2))
|
||||
stale_request
|
||||
.call
|
||||
.lookup(Duration::from_secs(2))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert!(
|
||||
cache
|
||||
.lookup(&request, Duration::from_secs(2))
|
||||
request
|
||||
.call
|
||||
.lookup(Duration::from_secs(2))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
clock.store(6, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&request, Duration::from_secs(6))
|
||||
.await
|
||||
.unwrap(),
|
||||
request.call.lookup(Duration::from_secs(6)).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
|
@ -104,20 +116,33 @@ async fn entry_limit_applies_to_sync_async_and_batch_writes() {
|
|||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let service: Arc<dyn ResponseCacheService> = Arc::new(cache);
|
||||
let service = Arc::new(cache).prepare(CacheOptions::new(CacheScope::Shared), "test", &|| {
|
||||
json!("service")
|
||||
});
|
||||
service
|
||||
.store(&request("service"), small.clone(), Duration::ZERO)
|
||||
.call
|
||||
.store(small.clone(), Duration::ZERO)
|
||||
.await
|
||||
.unwrap();
|
||||
for key in ["sync", "async", "batch-large"] {
|
||||
assert!(storage.get_cache(key).unwrap().is_none());
|
||||
}
|
||||
for key in ["batch-small", "service"] {
|
||||
assert_eq!(
|
||||
service.lookup(&request(key), Duration::ZERO).await.unwrap(),
|
||||
Some(small.clone())
|
||||
);
|
||||
}
|
||||
assert!(storage.get_cache("sync").unwrap().is_none());
|
||||
assert!(storage.get_cache("async").unwrap().is_none());
|
||||
assert!(storage.get_cache("batch-large").unwrap().is_none());
|
||||
assert_eq!(
|
||||
storage
|
||||
.get_cache("batch-small")
|
||||
.unwrap()
|
||||
.map(|entry| entry.response),
|
||||
Some(small.clone())
|
||||
);
|
||||
assert_eq!(
|
||||
service
|
||||
.call
|
||||
.lookup(Duration::ZERO)
|
||||
.await
|
||||
.unwrap()
|
||||
.map(|hit| hit.value),
|
||||
Some(small)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
|
|||
|
|
@ -77,6 +77,7 @@
|
|||
"kwargs",
|
||||
"call_type"
|
||||
],
|
||||
"reviews_deployment_success": [],
|
||||
"after_deployment_success": [
|
||||
"kwargs",
|
||||
"response",
|
||||
|
|
|
|||
|
|
@ -311,13 +311,12 @@ impl LegacyLogging {
|
|||
self.prepare(py)
|
||||
}
|
||||
|
||||
fn resume_after_success(
|
||||
fn resume_provider_result(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
_: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<HookStep<Self, Py<PyAny>>> {
|
||||
self.response = Some(result?);
|
||||
self.finalize(py)
|
||||
result.map(HookStep::Ready)
|
||||
}
|
||||
|
||||
fn resume_deployment_failure(
|
||||
|
|
@ -423,20 +422,29 @@ impl LegacyLogging {
|
|||
) -> PyResult<HookStep<Self, Py<PyAny>>> {
|
||||
self.end = Some(datetime(py, timing.end_time)?);
|
||||
self.response = Some(response);
|
||||
if self.runs_deployment_hooks() {
|
||||
return Ok(HookStep::Await(
|
||||
DeploymentHooks::after_success(
|
||||
py,
|
||||
self.call.kwargs(),
|
||||
&self.response,
|
||||
self.call_type(),
|
||||
)?,
|
||||
Self::resume_after_success,
|
||||
));
|
||||
}
|
||||
self.finalize(py)
|
||||
}
|
||||
|
||||
pub(crate) fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool> {
|
||||
Ok(self.runs_deployment_hooks() && DeploymentHooks::reviews_success(py)?)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_result(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
) -> PyResult<HookStep<Self, Py<PyAny>>> {
|
||||
Ok(HookStep::Await(
|
||||
DeploymentHooks::after_success(
|
||||
py,
|
||||
self.call.kwargs(),
|
||||
&Some(response),
|
||||
self.call_type(),
|
||||
)?,
|
||||
Self::resume_provider_result,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) fn result_ready(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
|
|
@ -787,8 +795,24 @@ assert hooked == ([opaque] if asynchronous else []), hooked
|
|||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::synchronous(false, false)]
|
||||
#[case::asynchronous(true, true)]
|
||||
fn only_an_async_call_reviews_provider_results(
|
||||
#[case] asynchronous: bool,
|
||||
#[case] reviews: bool,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = namespace(py, c"kwargs = {'logger': logger}");
|
||||
let (mut logging, _) = begin(py, &locals, asynchronous);
|
||||
|
||||
assert_eq!(logging.reviews_provider_results(py).unwrap(), reviews);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn response_returned_by_the_post_call_hook_is_finalized_and_returned() {
|
||||
fn the_success_hook_reviews_the_provider_result_and_its_answer_is_finalized() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = namespace(
|
||||
|
|
@ -809,18 +833,23 @@ logger.hooks = {'pre': lambda kwargs: kwargs}
|
|||
)
|
||||
.unwrap();
|
||||
let step = logging
|
||||
.transform_response(py, local(&locals, "response").unbind(), TIMING)
|
||||
.provider_result(py, local(&locals, "response").unbind())
|
||||
.unwrap();
|
||||
assert!(awaits_deployment_hook(&step));
|
||||
let step = resume(
|
||||
let HookStep::Ready(reviewed) = resume(
|
||||
&mut logging,
|
||||
step,
|
||||
py,
|
||||
Ok(local(&locals, "replacement").unbind()),
|
||||
)
|
||||
.unwrap();
|
||||
let HookStep::Ready(returned) = step else {
|
||||
panic!("expected the finalized response");
|
||||
.unwrap() else {
|
||||
panic!("expected the hook's answer");
|
||||
};
|
||||
assert!(reviewed.bind(py).is(local(&locals, "replacement")));
|
||||
let HookStep::Ready(returned) =
|
||||
logging.transform_response(py, reviewed, TIMING).unwrap()
|
||||
else {
|
||||
panic!("transforming the public response runs no deployment hook");
|
||||
};
|
||||
assert!(returned.bind(py).is(local(&locals, "replacement")));
|
||||
run(
|
||||
|
|
@ -829,6 +858,7 @@ logger.hooks = {'pre': lambda kwargs: kwargs}
|
|||
c"
|
||||
[finalized] = [value for name, value in logger.calls if name == 'finalize']
|
||||
assert finalized is replacement
|
||||
assert [name for name in logger.names() if name.endswith('_hook')] == ['pre_hook', 'success_hook']
|
||||
",
|
||||
);
|
||||
});
|
||||
|
|
@ -853,7 +883,7 @@ assert finalized is replacement
|
|||
)
|
||||
.unwrap();
|
||||
let step = logging
|
||||
.transform_response(py, local(&locals, "response").unbind(), TIMING)
|
||||
.provider_result(py, local(&locals, "response").unbind())
|
||||
.unwrap();
|
||||
resume(&mut logging, step, py, Err(cancellation))
|
||||
.err()
|
||||
|
|
|
|||
|
|
@ -95,6 +95,12 @@ impl DeploymentHooks {
|
|||
.map(Bound::unbind)
|
||||
}
|
||||
|
||||
pub fn reviews_success(py: Python<'_>) -> PyResult<bool> {
|
||||
python::DeploymentHooks::ReviewsSuccess
|
||||
.call(py, ())?
|
||||
.extract()
|
||||
}
|
||||
|
||||
pub fn after_success(
|
||||
py: Python<'_>,
|
||||
kwargs: &Py<PyDict>,
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ pub enum CallBoundary {
|
|||
PrepareArguments,
|
||||
BeforeProviderRequest,
|
||||
AfterProviderResponse,
|
||||
ProviderResult,
|
||||
TransformResponse,
|
||||
Succeeded,
|
||||
Failed,
|
||||
|
|
@ -53,6 +54,7 @@ type Prepare = fn(&mut LegacyLogging, Python<'_>, Py<PyDict>, f64) -> Step<Py<Py
|
|||
type Before =
|
||||
fn(&mut LegacyLogging, Python<'_>, Box<WireRequest>, &RequestContext) -> Step<Box<WireRequest>>;
|
||||
type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>;
|
||||
type ProviderResult = fn(&mut LegacyLogging, Python<'_>, Py<PyAny>) -> Step<Py<PyAny>>;
|
||||
type Transform = fn(&mut LegacyLogging, Python<'_>, Py<PyAny>, Timing) -> Step<Py<PyAny>>;
|
||||
type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py<PyAny>) -> Step<()>;
|
||||
type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>;
|
||||
|
|
@ -77,10 +79,16 @@ const AFTER: Binding<After> = Binding {
|
|||
callbacks: &["log_post_api_call"],
|
||||
};
|
||||
|
||||
const PROVIDER_RESULT: Binding<ProviderResult> = Binding {
|
||||
boundary: CallBoundary::ProviderResult,
|
||||
invoke: LegacyLogging::provider_result,
|
||||
callbacks: &["async_post_call_success_deployment_hook"],
|
||||
};
|
||||
|
||||
const TRANSFORM: Binding<Transform> = Binding {
|
||||
boundary: CallBoundary::TransformResponse,
|
||||
invoke: LegacyLogging::transform_public_response,
|
||||
callbacks: &["async_post_call_success_deployment_hook"],
|
||||
callbacks: &[],
|
||||
};
|
||||
|
||||
const SUCCESS: Binding<Success> = Binding {
|
||||
|
|
@ -127,6 +135,7 @@ pub fn callback_mappings() -> impl Iterator<Item = CallbackMapping> {
|
|||
.mappings()
|
||||
.chain(BEFORE.mappings())
|
||||
.chain(AFTER.mappings())
|
||||
.chain(PROVIDER_RESULT.mappings())
|
||||
.chain(TRANSFORM.mappings())
|
||||
.chain(SUCCESS.mappings())
|
||||
.chain(FAILURE.mappings())
|
||||
|
|
@ -159,6 +168,14 @@ impl CallHooks<PythonRuntime> for LegacyLogging {
|
|||
(BEFORE.invoke)(self, py, wire, context)
|
||||
}
|
||||
|
||||
fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool> {
|
||||
LegacyLogging::reviews_provider_results(self, py)
|
||||
}
|
||||
|
||||
fn provider_result(&mut self, py: Python<'_>, response: Py<PyAny>) -> Step<Py<PyAny>> {
|
||||
(PROVIDER_RESULT.invoke)(self, py, response)
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
),
|
||||
|
|
|
|||
|
|
@ -47,7 +47,7 @@ async fn all_inference_endpoints_share_native_cache(
|
|||
.expect(2)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let cache: Arc<dyn ResponseCacheService> = Arc::new(
|
||||
let cache = Arc::new(
|
||||
ResponseCache::new(Arc::new(InMemoryCache::new(
|
||||
Some(100),
|
||||
Some(Duration::from_secs(60)),
|
||||
|
|
@ -72,7 +72,7 @@ async fn all_inference_endpoints_share_native_cache(
|
|||
let cache_key = second.headers().get("x-litellm-cache-key").unwrap().clone();
|
||||
assert!(!cache_key.as_bytes().is_empty());
|
||||
let stored = cache
|
||||
.lookup(
|
||||
.async_lookup(
|
||||
&ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some(cache_key.to_str().unwrap().into()),
|
||||
..Default::default()
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ use crate::{InvokeError, PythonOwned};
|
|||
use litellm_host::protocol::Protocol;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyDict;
|
||||
use serde_json::Value;
|
||||
|
||||
/// Converts requests, responses, stream values and errors at the Python boundary.
|
||||
pub trait PythonBinding: PythonOwned {
|
||||
|
|
@ -27,6 +28,16 @@ pub trait PythonBinding: PythonOwned {
|
|||
response: <Self::Protocol as Protocol>::Response,
|
||||
) -> PyResult<Py<PyAny>>;
|
||||
|
||||
/// A provider result as the public object Python hooks review, built from its JSON form.
|
||||
fn public_response(&mut self, py: Python<'_>, response: &Value) -> PyResult<Py<PyAny>> {
|
||||
crate::to_py(py, response)
|
||||
}
|
||||
|
||||
/// The JSON form of the public object a Python hook handed back.
|
||||
fn response_value(&mut self, _py: Python<'_>, public: &Bound<'_, PyAny>) -> PyResult<Value> {
|
||||
crate::from_py(public)
|
||||
}
|
||||
|
||||
/// What the stream carries at hand-off, as the caller's stream receives it.
|
||||
fn encode_stream_head(
|
||||
&mut self,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -71,6 +71,12 @@ pub(super) trait ChainHooks: PythonOwned {
|
|||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<ChainStep<Box<WireRequest>>>;
|
||||
fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool>;
|
||||
fn provider_result(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
) -> PyResult<ChainStep<Py<PyAny>>>;
|
||||
fn transform_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
|
|
@ -138,6 +144,19 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
|
|||
Ok(self.step(step, Continuation::Wire))
|
||||
}
|
||||
|
||||
fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool> {
|
||||
self.hooks.reviews_provider_results(py)
|
||||
}
|
||||
|
||||
fn provider_result(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
) -> PyResult<ChainStep<Py<PyAny>>> {
|
||||
let step = self.hooks.provider_result(py, response)?;
|
||||
Ok(self.step(step, Continuation::Response))
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ pub struct HookChain {
|
|||
enum Pending {
|
||||
Arguments(usize, f64),
|
||||
Wire(usize, RequestContext),
|
||||
ProviderResult(usize),
|
||||
Response(usize, Timing),
|
||||
}
|
||||
|
||||
|
|
@ -124,6 +125,41 @@ impl Transform for Wire {
|
|||
}
|
||||
}
|
||||
|
||||
struct ProviderResult;
|
||||
|
||||
impl Transform for ProviderResult {
|
||||
type Value = Py<PyAny>;
|
||||
type Context = ();
|
||||
|
||||
fn call(
|
||||
hooks: &mut dyn ChainHooks,
|
||||
py: Python<'_>,
|
||||
value: Py<PyAny>,
|
||||
_: &(),
|
||||
) -> PyResult<ChainStep<Py<PyAny>>> {
|
||||
hooks.provider_result(py, value)
|
||||
}
|
||||
|
||||
fn resume(
|
||||
hooks: &mut dyn ChainHooks,
|
||||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<ChainStep<Py<PyAny>>> {
|
||||
hooks.resume_response(py, result)
|
||||
}
|
||||
|
||||
fn suspend(index: usize, (): ()) -> Pending {
|
||||
Pending::ProviderResult(index)
|
||||
}
|
||||
|
||||
fn resumed(pending: Pending) -> Option<(usize, ())> {
|
||||
match pending {
|
||||
Pending::ProviderResult(index) => Some((index, ())),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Response;
|
||||
|
||||
impl Transform for Response {
|
||||
|
|
@ -386,6 +422,23 @@ impl CallHooks<PythonRuntime> for HookChain {
|
|||
self.transform_from::<Wire>(py, 0, wire, context.clone())
|
||||
}
|
||||
|
||||
fn reviews_provider_results(&mut self, py: Python<'_>) -> PyResult<bool> {
|
||||
for hooks in &mut self.hooks {
|
||||
if hooks.reviews_provider_results(py)? {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
fn provider_result(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
) -> PyResult<HookStep<Self, Py<PyAny>>> {
|
||||
self.transform_from::<ProviderResult>(py, 0, response, ())
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ use rstest::rstest;
|
|||
use support::{
|
||||
HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options,
|
||||
failing_machine, handed, hooks, interpreter, lifecycle_module, open_stream, read_all,
|
||||
run_composed, run_scripted, settled, success_machine,
|
||||
reviewed_machine, run_composed, run_scripted, settled, success_machine,
|
||||
};
|
||||
|
||||
fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
|
||||
|
|
@ -325,6 +325,34 @@ fn success_runs_every_step_in_order_and_returns_the_public_response(
|
|||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::reviewed(HookScript::ReviewResult, "reviewed", &["provider_result:provider"])]
|
||||
#[case::not_reviewed(HookScript::Plain, "provider", &[])]
|
||||
fn a_provider_result_keeps_what_the_reviewing_hooks_answer(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
#[case] script: HookScript,
|
||||
#[case] kept: &str,
|
||||
#[case] reviewed: &[&str],
|
||||
) {
|
||||
Python::attach(|py| {
|
||||
let (result, log) = run_scripted(
|
||||
py,
|
||||
reviewed_machine(),
|
||||
OpScript::Answer,
|
||||
script,
|
||||
asynchronous,
|
||||
);
|
||||
let seen: Vec<_> = log
|
||||
.iter()
|
||||
.filter(|entry| entry.starts_with("provider_result:"))
|
||||
.cloned()
|
||||
.collect();
|
||||
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), kept);
|
||||
assert_eq!(seen, reviewed);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
|
||||
#[from(interpreter)] _python: Interpreter,
|
||||
|
|
|
|||
|
|
@ -321,6 +321,7 @@ pub enum HookScript {
|
|||
FailAfterSuccess,
|
||||
CancelTerminal,
|
||||
FailTerminal,
|
||||
ReviewResult,
|
||||
}
|
||||
|
||||
pub struct SyntheticHooks {
|
||||
|
|
@ -336,6 +337,24 @@ pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
|
|||
}
|
||||
|
||||
impl CallHooks<PythonRuntime> for SyntheticHooks {
|
||||
fn reviews_provider_results(&mut self, _: Python<'_>) -> PyResult<bool> {
|
||||
Ok(matches!(self.script, HookScript::ReviewResult))
|
||||
}
|
||||
|
||||
fn provider_result(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
) -> PyResult<HookStep<Self, Py<PyAny>>> {
|
||||
let reviewed = response.bind(py).cast::<PyDict>()?.copy()?;
|
||||
self.log.push(format!(
|
||||
"provider_result:{}",
|
||||
reviewed.get_item("text")?.unwrap()
|
||||
));
|
||||
reviewed.set_item("text", "reviewed")?;
|
||||
Ok(HookStep::Ready(reviewed.into_any().unbind()))
|
||||
}
|
||||
|
||||
fn prepare_arguments(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
|
|
@ -474,6 +493,21 @@ pub fn success_machine() -> MachineFactory {
|
|||
})
|
||||
}
|
||||
|
||||
/// A call that offers its provider result for review and returns the text it keeps.
|
||||
pub fn reviewed_machine() -> MachineFactory {
|
||||
Box::new(move |_| {
|
||||
CallMachine::new(None, move |host| {
|
||||
Box::pin(async move {
|
||||
let kept = host
|
||||
.interceptors
|
||||
.provider_result(serde_json::json!({"text": "provider"}))
|
||||
.await?;
|
||||
Ok(kept["text"].as_str().unwrap_or_default().to_owned())
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn failing_machine() -> MachineFactory {
|
||||
Box::new(move |_| {
|
||||
CallMachine::new(None, |_| {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
||||
|
|
|
|||
|
|
@ -44,6 +44,21 @@ pub trait CallHooks<R: HookRuntime>: Sized {
|
|||
Ok(R::ready(wire))
|
||||
}
|
||||
|
||||
/// Whether `provider_result` would do anything, so a host can skip encoding the response.
|
||||
fn reviews_provider_results(&mut self, _runtime: R::Context<'_>) -> Result<bool, R::Error> {
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
/// A complete provider response before it is cached or delivered. What this returns
|
||||
/// replaces it. Never offered a cache hit.
|
||||
fn provider_result(
|
||||
&mut self,
|
||||
_runtime: R::Context<'_>,
|
||||
response: R::Response,
|
||||
) -> Result<R::Step<Self, R::Response>, R::Error> {
|
||||
Ok(R::ready(response))
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&mut self,
|
||||
_runtime: R::Context<'_>,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
use std::future::Future;
|
||||
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
|
||||
use crate::{
|
||||
machine::MachineFault,
|
||||
protocol::{InterceptRequest, Reply, reply},
|
||||
|
|
@ -145,6 +147,25 @@ where
|
|||
fn result_ready(&self, facts: ExecutionFacts) -> impl Future<Output = Result<(), E>> + Send {
|
||||
self.intercept(|reply| InterceptRequest::ResultReady { facts, reply })
|
||||
}
|
||||
|
||||
/// Offers a complete provider response before it is cached or delivered, and returns the
|
||||
/// response the call keeps. Never raised for a cache hit.
|
||||
fn provider_result<R>(&self, response: R) -> impl Future<Output = Result<R, E>> + Send
|
||||
where
|
||||
R: Serialize + DeserializeOwned + Send,
|
||||
{
|
||||
let value = serde_json::to_value(&response);
|
||||
async move {
|
||||
let value = value.map_err(|_| MachineFault::UnreadableAnswer)?;
|
||||
let answer = self
|
||||
.intercept(|reply| InterceptRequest::ProviderResult {
|
||||
response: value,
|
||||
reply,
|
||||
})
|
||||
.await?;
|
||||
serde_json::from_value(answer).map_err(|_| MachineFault::UnreadableAnswer.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: From<MachineFault>, T: Interceptors<E> + ?Sized> InterceptPoints<E> for T {}
|
||||
|
|
|
|||
|
|
@ -16,6 +16,8 @@ pub enum MachineFault {
|
|||
Abandoned,
|
||||
/// The host resumed the call out of turn.
|
||||
Protocol(ResumeError),
|
||||
/// The host answered with a value the call cannot read back as its own type.
|
||||
UnreadableAnswer,
|
||||
}
|
||||
|
||||
pub type ExecuteFuture<R, C = <R as Protocol>::Response> =
|
||||
|
|
|
|||
|
|
@ -33,6 +33,12 @@ pub enum InterceptRequest {
|
|||
raw: RawResponse,
|
||||
reply: Reply<()>,
|
||||
},
|
||||
/// A complete response the provider just produced, before anything caches or delivers it.
|
||||
/// Never raised for a cache hit. The answer is the response the call keeps.
|
||||
ProviderResult {
|
||||
response: serde_json::Value,
|
||||
reply: Reply<serde_json::Value>,
|
||||
},
|
||||
}
|
||||
|
||||
impl InterceptRequest {
|
||||
|
|
@ -43,6 +49,7 @@ impl InterceptRequest {
|
|||
reply.send(());
|
||||
}
|
||||
Self::BeforeProviderRequest { wire, reply, .. } => reply.send(*wire),
|
||||
Self::ProviderResult { response, reply } => reply.send(response),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -68,6 +75,9 @@ impl InterceptRequest {
|
|||
Self::AfterProviderResponse { raw, reply } => {
|
||||
reply.send(interceptors.after_provider_response(raw).await?);
|
||||
}
|
||||
Self::ProviderResult { response, reply } => {
|
||||
reply.send(interceptors.provider_result(response).await?);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -59,14 +59,11 @@ pub(super) async fn execute(
|
|||
)
|
||||
.await?;
|
||||
let cache = cache.filter(|_| authenticated.signer.is_none());
|
||||
let cache_request = litellm_inference::caching::CacheRequest::from_wire(
|
||||
identity,
|
||||
cache.as_ref().map(|_| &wire),
|
||||
);
|
||||
let cache_request = litellm_inference::caching::CacheRequest::from_wire::<
|
||||
super::route::ChatCompletions,
|
||||
>(identity, &wire, cache.as_ref(), cache_options);
|
||||
litellm_inference::caching::execute_unary::<super::route::ChatCompletions, _, _>(
|
||||
cache_request,
|
||||
cache.as_ref().map(|cache| cache.service.clone()),
|
||||
cache.as_ref().map(|cache| cache.options(cache_options)),
|
||||
interceptors,
|
||||
observers,
|
||||
|| async move {
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ use prepare::{prepare_provider_request, resolve_request};
|
|||
use litellm_auth::AuthServices;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use std::sync::Arc;
|
||||
use types::ChatCompletionsRequest;
|
||||
use types::{ChatCompletionsCall, ChatCompletionsRequest};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ChatCompletionsRoute {
|
||||
|
|
@ -48,6 +48,16 @@ impl ChatCompletionsRoute {
|
|||
request: ChatCompletionsRequest<'_>,
|
||||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
options: impl Into<litellm_inference::CallOptions>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
self.execute_call(request.into(), interceptors, options)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn execute_call(
|
||||
&self,
|
||||
call: ChatCompletionsCall,
|
||||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
options: impl Into<litellm_inference::CallOptions>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let litellm_inference::CallOptions {
|
||||
cache: cache_options,
|
||||
|
|
@ -55,12 +65,7 @@ impl ChatCompletionsRoute {
|
|||
} = options.into();
|
||||
litellm_host::lifecycle::observe_unary(
|
||||
observers.clone(),
|
||||
self.run_call(
|
||||
request.into(),
|
||||
cache_options,
|
||||
interceptors,
|
||||
observers.as_ref(),
|
||||
),
|
||||
self.run_call(call, cache_options, interceptors, observers.as_ref()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -181,6 +181,7 @@ impl Interceptors<RouteError> for ChangingHooks {
|
|||
InterceptRequest::ResultReady { facts, reply } => {
|
||||
reply.send(self.result_ready(facts).await?);
|
||||
}
|
||||
other => other.pass_through(),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -106,6 +106,12 @@ impl MessagesRoute {
|
|||
source: source.clone(),
|
||||
})
|
||||
.await?;
|
||||
let output = litellm_inference::caching::provider_result::<route::Messages>(
|
||||
output,
|
||||
&source,
|
||||
context.interceptors,
|
||||
)
|
||||
.await?;
|
||||
Ok(cache.finish(output, &source).await)
|
||||
})
|
||||
.await
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ use std::{
|
|||
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{
|
||||
CacheOptions, CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, ScopedCache,
|
||||
CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, ScopedCache,
|
||||
};
|
||||
use litellm_host::interceptors::{
|
||||
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext,
|
||||
|
|
@ -58,22 +58,22 @@ async fn messages_cache_identity_includes_provider_native_parameters(
|
|||
for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] {
|
||||
let response =
|
||||
execute_unary::<Messages, _, _>(
|
||||
CacheRequest::from_wire(
|
||||
CacheRequest::from_wire::<Messages>(
|
||||
ProviderIdentity {
|
||||
model: "test".into(),
|
||||
provider: "anthropic".into(),
|
||||
},
|
||||
Some(&WireRequest {
|
||||
&WireRequest {
|
||||
url: "https://example.test/v1/messages".into(),
|
||||
headers: vec![],
|
||||
body: json!({
|
||||
"model":"test", "messages":[{"role":"user","content":"hello"}],
|
||||
"max_tokens":32, (field):value
|
||||
}),
|
||||
}),
|
||||
},
|
||||
Some(&ScopedCache::new(cache.clone(), CacheScope::Shared)),
|
||||
None,
|
||||
),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
&(),
|
||||
None,
|
||||
|| async {
|
||||
|
|
@ -148,6 +148,7 @@ impl Interceptors<RouteError> for ChangingHooks {
|
|||
InterceptRequest::ResultReady { facts, reply } => {
|
||||
reply.send(self.result_ready(facts).await?);
|
||||
}
|
||||
other => other.pass_through(),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ impl Interceptors<Error> for LiveCall {
|
|||
InterceptRequest::AfterProviderResponse { raw, reply } => {
|
||||
reply.send(self.after_provider_response(raw).await?);
|
||||
}
|
||||
other => other.pass_through(),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -95,6 +95,7 @@ impl litellm_host::interceptors::Interceptors<<Messages as litellm_host::protoco
|
|||
InterceptRequest::AfterProviderResponse { raw, reply } => {
|
||||
reply.send(self.after_provider_response(raw).await?);
|
||||
}
|
||||
other => other.pass_through(),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -37,14 +37,11 @@ pub(super) async fn execute(
|
|||
)
|
||||
.await?;
|
||||
let cache = cache.filter(|_| authenticated.signer.is_none());
|
||||
let cache_request = litellm_inference::caching::CacheRequest::from_wire(
|
||||
identity,
|
||||
cache.as_ref().map(|_| &wire),
|
||||
);
|
||||
let cache_request = litellm_inference::caching::CacheRequest::from_wire::<
|
||||
super::route::Responses,
|
||||
>(identity, &wire, cache.as_ref(), cache_options);
|
||||
litellm_inference::caching::execute_streaming::<super::route::Responses, _, _>(
|
||||
cache_request,
|
||||
cache.as_ref().map(|cache| cache.service.clone()),
|
||||
cache.as_ref().map(|cache| cache.options(cache_options)),
|
||||
interceptors,
|
||||
observers,
|
||||
|| async move {
|
||||
|
|
|
|||
|
|
@ -25,14 +25,20 @@ use litellm_inference::{
|
|||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
fn cache_request(input: Value) -> CacheRequest {
|
||||
CacheRequest {
|
||||
identity: ProviderIdentity {
|
||||
fn cache_request<P: litellm_inference::caching::Cachable>(
|
||||
input: Value,
|
||||
cache: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
) -> CacheRequest {
|
||||
CacheRequest::prepare::<P>(
|
||||
ProviderIdentity {
|
||||
model: "test-model".into(),
|
||||
provider: "test-provider".into(),
|
||||
},
|
||||
input,
|
||||
}
|
||||
cache,
|
||||
options,
|
||||
&|_| input.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
|
|
@ -59,31 +65,55 @@ struct InvalidEntryCache(
|
|||
);
|
||||
|
||||
impl ResponseCacheService for InvalidEntryCache {
|
||||
fn config(&self) -> &ResponseCacheConfig {
|
||||
self.0.config()
|
||||
fn prepare(
|
||||
self: Arc<Self>,
|
||||
options: CacheOptions,
|
||||
surface: &str,
|
||||
input: &dyn Fn() -> Value,
|
||||
) -> litellm_cache_response::PreparedCache {
|
||||
litellm_cache_response::PreparedCache {
|
||||
policy: options.policy,
|
||||
max_entry_bytes: self.0.config().max_entry_bytes,
|
||||
call: Arc::new(InvalidEntryCall {
|
||||
request: options.request(&self.0.config().namespace, surface, input()),
|
||||
service: self,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct InvalidEntryCall {
|
||||
service: Arc<InvalidEntryCache>,
|
||||
request: litellm_cache_response::ResponseCacheRequest,
|
||||
}
|
||||
|
||||
impl litellm_cache_response::ResponseCacheCall for InvalidEntryCall {
|
||||
fn lookup<'a>(
|
||||
&'a self,
|
||||
request: &'a litellm_cache_response::ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> futures_util::future::BoxFuture<'a, Result<Option<Value>, litellm_cache::Error>> {
|
||||
) -> futures_util::future::BoxFuture<
|
||||
'a,
|
||||
Result<Option<litellm_cache_response::CacheHit>, litellm_cache::Error>,
|
||||
> {
|
||||
Box::pin(async move {
|
||||
Ok(self
|
||||
.0
|
||||
.async_lookup(request, now)
|
||||
.await?
|
||||
.or_else(|| Some(self.1.clone())))
|
||||
Ok(Some(litellm_cache_response::CacheHit {
|
||||
value: self
|
||||
.service
|
||||
.0
|
||||
.async_lookup(&self.request, now)
|
||||
.await?
|
||||
.unwrap_or_else(|| self.service.1.clone()),
|
||||
key: litellm_cache_response::cache_key(&self.request.key),
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
fn store<'a>(
|
||||
&'a self,
|
||||
request: &'a litellm_cache_response::ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> futures_util::future::BoxFuture<'a, Result<(), litellm_cache::Error>> {
|
||||
Box::pin(self.0.async_store(request, response, now))
|
||||
Box::pin(self.service.0.async_store(&self.request, response, now))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -139,6 +169,7 @@ impl Interceptors<RouteError> for ChangingHooks {
|
|||
InterceptRequest::ResultReady { facts, reply } => {
|
||||
reply.send(self.result_ready(facts).await?);
|
||||
}
|
||||
other => other.pass_through(),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -184,9 +215,11 @@ async fn responses_refetches_instead_of_deserializing_another_api_response(
|
|||
let calls = AtomicUsize::new(0);
|
||||
for _ in 0..2 {
|
||||
let response = execute_unary::<Responses, _, _>(
|
||||
cache_request(json!({"input":"hello"})),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
cache_request::<Responses>(
|
||||
json!({"input":"hello"}),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
),
|
||||
&(),
|
||||
None,
|
||||
|| async {
|
||||
|
|
@ -226,9 +259,11 @@ async fn responses_cache_only_reuses_completed_responses(
|
|||
let calls = AtomicUsize::new(0);
|
||||
for _ in 0..2 {
|
||||
let response = execute_unary::<Responses, _, _>(
|
||||
cache_request(json!({"input":"hello"})),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
cache_request::<Responses>(
|
||||
json!({"input":"hello"}),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
),
|
||||
&(),
|
||||
None,
|
||||
|| async {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -8,8 +8,7 @@ use std::{
|
|||
use bytes::{Bytes, BytesMut};
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use litellm_cache_response::{
|
||||
CacheOptions, CachePolicy, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope,
|
||||
ScopedCache, cache_key,
|
||||
CacheOptions, CachePolicy, PreparedCache, ResponseCacheService, ResponseEnvelope, ScopedCache,
|
||||
};
|
||||
use litellm_host::{
|
||||
call::{CallOutput, OutputOf},
|
||||
|
|
@ -37,13 +36,36 @@ pub trait Cachable: Protocol<Error = RouteError> {
|
|||
|
||||
pub struct CacheRequest {
|
||||
pub identity: ProviderIdentity,
|
||||
pub input: Value,
|
||||
session: Option<CacheSession>,
|
||||
}
|
||||
|
||||
impl CacheRequest {
|
||||
pub fn from_wire(identity: ProviderIdentity, wire: Option<&WireRequest>) -> Self {
|
||||
Self {
|
||||
input: wire.map_or(Value::Null, |wire| {
|
||||
pub fn prepare<P: Cachable>(
|
||||
identity: ProviderIdentity,
|
||||
service: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
input: &dyn Fn(&ProviderIdentity) -> Value,
|
||||
) -> Self {
|
||||
let session = options
|
||||
.filter(|options| options.policy.enabled())
|
||||
.zip(service)
|
||||
.map(|(options, service)| CacheSession {
|
||||
prepared: service.prepare(options, P::SURFACE, &|| input(&identity)),
|
||||
});
|
||||
Self { identity, session }
|
||||
}
|
||||
|
||||
pub fn from_wire<P: Cachable>(
|
||||
identity: ProviderIdentity,
|
||||
wire: &WireRequest,
|
||||
cache: Option<&ScopedCache>,
|
||||
policy: Option<CachePolicy>,
|
||||
) -> Self {
|
||||
Self::prepare::<P>(
|
||||
identity,
|
||||
cache.map(|cache| cache.service.clone()),
|
||||
cache.map(|cache| cache.options(policy)),
|
||||
&|identity| {
|
||||
serde_json::json!({
|
||||
"provider": identity.provider,
|
||||
"model": identity.model,
|
||||
|
|
@ -51,9 +73,8 @@ impl CacheRequest {
|
|||
"headers": wire.headers,
|
||||
"body": wire.body,
|
||||
})
|
||||
}),
|
||||
identity,
|
||||
}
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -72,35 +93,23 @@ pub enum CachedOutput<R> {
|
|||
}
|
||||
|
||||
struct CacheSession {
|
||||
service: Arc<dyn ResponseCacheService>,
|
||||
request: ResponseCacheRequest,
|
||||
prepared: PreparedCache,
|
||||
}
|
||||
|
||||
impl CacheSession {
|
||||
fn prepare<P: Cachable>(
|
||||
service: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
request: &CacheRequest,
|
||||
) -> Option<Self> {
|
||||
let options = options.filter(|options| options.policy.enabled())?;
|
||||
let service = service?;
|
||||
let input = request.input.clone();
|
||||
let request = options.request(&service.config().namespace, P::SURFACE, input);
|
||||
Some(Self { service, request })
|
||||
}
|
||||
|
||||
async fn lookup<P: Cachable>(&self) -> Option<CachedOutput<P::Response>>
|
||||
async fn lookup<P: Cachable>(&self) -> Option<(CachedOutput<P::Response>, String)>
|
||||
where
|
||||
P::Response: DeserializeOwned,
|
||||
{
|
||||
if !self.request.controls.reads() {
|
||||
if !self.prepared.policy.reads() {
|
||||
return None;
|
||||
}
|
||||
match self.service.lookup(&self.request, now()).await {
|
||||
Ok(Some(value)) => {
|
||||
serde_json::from_value::<ResponseEnvelope<CachedOutput<P::Response>>>(value)
|
||||
match self.prepared.call.lookup(now()).await {
|
||||
Ok(Some(hit)) => {
|
||||
serde_json::from_value::<ResponseEnvelope<CachedOutput<P::Response>>>(hit.value)
|
||||
.ok()
|
||||
.and_then(|entry| entry.decode(P::SURFACE))
|
||||
.map(|output| (output, hit.key))
|
||||
}
|
||||
Ok(None) => None,
|
||||
Err(_) => {
|
||||
|
|
@ -111,15 +120,10 @@ impl CacheSession {
|
|||
}
|
||||
|
||||
async fn store(&self, entry: Value) {
|
||||
if !self.request.controls.writes() {
|
||||
if !self.prepared.policy.writes() {
|
||||
return;
|
||||
}
|
||||
if self
|
||||
.service
|
||||
.store(&self.request, entry, now())
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
if self.prepared.call.store(entry, now()).await.is_err() {
|
||||
tracing::warn!("response cache write failed");
|
||||
}
|
||||
}
|
||||
|
|
@ -128,7 +132,7 @@ impl CacheSession {
|
|||
where
|
||||
P::Response: Serialize,
|
||||
{
|
||||
if !self.request.controls.writes() || !P::reusable(response) {
|
||||
if !self.prepared.policy.writes() || !P::reusable(response) {
|
||||
return;
|
||||
}
|
||||
if let Ok(value) = serde_json::to_value(response)
|
||||
|
|
@ -144,8 +148,6 @@ impl CacheSession {
|
|||
|
||||
pub async fn execute_unary<P, F, Fut>(
|
||||
request: CacheRequest,
|
||||
cache: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
interceptors: &impl Interceptors<RouteError>,
|
||||
observers: Option<&ObservationSender>,
|
||||
provider: F,
|
||||
|
|
@ -156,14 +158,16 @@ where
|
|||
F: FnOnce() -> Fut,
|
||||
Fut: Future<Output = Result<P::Response, RouteError>>,
|
||||
{
|
||||
let identity = request.identity.clone();
|
||||
let CacheRequest { identity, session } = request;
|
||||
crate::diagnostic::provider(&identity.model, &identity.provider);
|
||||
let session = CacheSession::prepare::<P>(cache, options, &request);
|
||||
let hit = match &session {
|
||||
Some(session) => session.lookup::<P>().await.and_then(|entry| match entry {
|
||||
CachedOutput::Response(response) => Some((response, cache_key(&session.request.key))),
|
||||
CachedOutput::Stream(_) => None,
|
||||
}),
|
||||
Some(session) => session
|
||||
.lookup::<P>()
|
||||
.await
|
||||
.and_then(|(entry, key)| match entry {
|
||||
CachedOutput::Response(response) => Some((response, key)),
|
||||
CachedOutput::Stream(_) => None,
|
||||
}),
|
||||
None => None,
|
||||
};
|
||||
let (response, source) = match hit {
|
||||
|
|
@ -182,7 +186,11 @@ where
|
|||
observers,
|
||||
)
|
||||
.await?;
|
||||
if from_provider && let Some(session) = session {
|
||||
if !from_provider {
|
||||
return Ok(response);
|
||||
}
|
||||
let response = interceptors.provider_result(response).await?;
|
||||
if let Some(session) = session {
|
||||
session.store_response::<P>(&response).await;
|
||||
}
|
||||
Ok(response)
|
||||
|
|
@ -190,8 +198,6 @@ where
|
|||
|
||||
pub async fn execute_streaming<P, F, Fut>(
|
||||
request: CacheRequest,
|
||||
cache: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
interceptors: &impl Interceptors<RouteError>,
|
||||
observers: Option<&ObservationSender>,
|
||||
provider: F,
|
||||
|
|
@ -202,9 +208,8 @@ where
|
|||
F: FnOnce() -> Fut,
|
||||
Fut: Future<Output = Result<OutputOf<P>, RouteError>>,
|
||||
{
|
||||
let identity = request.identity.clone();
|
||||
let CacheRequest { identity, session } = request;
|
||||
crate::diagnostic::provider(&identity.model, &identity.provider);
|
||||
let session = CacheSession::prepare::<P>(cache, options, &request);
|
||||
let cache = CallCache::<P> {
|
||||
session,
|
||||
protocol: PhantomData,
|
||||
|
|
@ -224,9 +229,29 @@ where
|
|||
observers,
|
||||
)
|
||||
.await?;
|
||||
let output = provider_result::<P>(output, &source, interceptors).await?;
|
||||
Ok(cache.finish(output, &source).await)
|
||||
}
|
||||
|
||||
/// Offers a complete provider response to the host before it is cached or delivered. A cache
|
||||
/// hit and a stream pass through untouched.
|
||||
pub async fn provider_result<P>(
|
||||
output: OutputOf<P>,
|
||||
source: &ResultSource,
|
||||
interceptors: &impl Interceptors<RouteError>,
|
||||
) -> Result<OutputOf<P>, RouteError>
|
||||
where
|
||||
P: Protocol<Error = RouteError>,
|
||||
P::Response: Serialize + DeserializeOwned,
|
||||
{
|
||||
match (source, output) {
|
||||
(ResultSource::Provider, CallOutput::Complete(response)) => Ok(CallOutput::Complete(
|
||||
interceptors.provider_result(response).await?,
|
||||
)),
|
||||
(_, output) => Ok(output),
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CallCache<P> {
|
||||
session: Option<CacheSession>,
|
||||
protocol: PhantomData<P>,
|
||||
|
|
@ -239,23 +264,9 @@ impl<P: StreamCachable> CallCache<P> {
|
|||
identity: &ProviderIdentity,
|
||||
wire: &WireRequest,
|
||||
) -> Self {
|
||||
let session = cache.and_then(|cache| {
|
||||
if !policy.enabled() {
|
||||
return None;
|
||||
}
|
||||
let options = cache.options(Some(policy));
|
||||
let request = CacheRequest::from_wire(identity.clone(), Some(wire));
|
||||
Some(CacheSession {
|
||||
request: options.request(
|
||||
&cache.service.config().namespace,
|
||||
P::SURFACE,
|
||||
request.input,
|
||||
),
|
||||
service: cache.service.clone(),
|
||||
})
|
||||
});
|
||||
let request = CacheRequest::from_wire::<P>(identity.clone(), wire, cache, Some(policy));
|
||||
Self {
|
||||
session,
|
||||
session: request.session,
|
||||
protocol: PhantomData,
|
||||
}
|
||||
}
|
||||
|
|
@ -265,16 +276,12 @@ impl<P: StreamCachable> CallCache<P> {
|
|||
P::Response: DeserializeOwned,
|
||||
{
|
||||
let session = self.session.as_ref()?;
|
||||
let output = match session.lookup::<P>().await? {
|
||||
let (entry, key) = session.lookup::<P>().await?;
|
||||
let output = match entry {
|
||||
CachedOutput::Response(response) => CallOutput::Complete(response),
|
||||
CachedOutput::Stream(data) => P::replay(Bytes::from(data))?,
|
||||
};
|
||||
Some((
|
||||
output,
|
||||
ResultSource::Cache {
|
||||
key: cache_key(&session.request.key),
|
||||
},
|
||||
))
|
||||
Some((output, ResultSource::Cache { key }))
|
||||
}
|
||||
|
||||
pub async fn finish(self, output: OutputOf<P>, source: &ResultSource) -> OutputOf<P>
|
||||
|
|
@ -282,7 +289,7 @@ impl<P: StreamCachable> CallCache<P> {
|
|||
P::Response: Serialize,
|
||||
{
|
||||
let Some(session) = self.session.filter(|session| {
|
||||
*source == ResultSource::Provider && session.request.controls.writes()
|
||||
*source == ResultSource::Provider && session.prepared.policy.writes()
|
||||
}) else {
|
||||
return output;
|
||||
};
|
||||
|
|
@ -310,8 +317,7 @@ fn capture_stream<P: StreamCachable>(
|
|||
Some(chunk) => {
|
||||
let captured = captured.and_then(|mut data| {
|
||||
let bytes = P::bytes(&chunk);
|
||||
if data.len().saturating_add(bytes.len())
|
||||
> session.service.config().max_entry_bytes
|
||||
if data.len().saturating_add(bytes.len()) > session.prepared.max_entry_bytes
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ impl From<litellm_host::machine::MachineFault> for RouteError {
|
|||
Self::InvalidRequest(match fault {
|
||||
MachineFault::Abandoned => "host driver was abandoned".into(),
|
||||
MachineFault::Protocol(message) => format!("host {message}").into(),
|
||||
MachineFault::UnreadableAnswer => "host answered with an unreadable response".into(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -61,14 +61,20 @@ impl StreamCachable for TestRoute {
|
|||
}
|
||||
}
|
||||
|
||||
fn cache_request(input: Value) -> CacheRequest {
|
||||
CacheRequest {
|
||||
identity: ProviderIdentity {
|
||||
fn cache_request<P: litellm_inference::caching::Cachable>(
|
||||
input: Value,
|
||||
cache: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
) -> CacheRequest {
|
||||
CacheRequest::prepare::<P>(
|
||||
ProviderIdentity {
|
||||
model: "test-model".into(),
|
||||
provider: "test-provider".into(),
|
||||
},
|
||||
input,
|
||||
}
|
||||
cache,
|
||||
options,
|
||||
&|_| input.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
|
|
@ -76,6 +82,38 @@ fn cache() -> Arc<dyn ResponseCacheService> {
|
|||
cache_with_limit(4096)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_service(false, true, CachePolicy::default())]
|
||||
#[case::no_options(true, false, CachePolicy::default())]
|
||||
#[case::off(true, true, CachePolicy { caching: Some(false), ..CachePolicy::default() })]
|
||||
#[case::no_io(true, true, CachePolicy { no_cache: true, no_store: true, ..CachePolicy::default() })]
|
||||
#[tokio::test]
|
||||
async fn disabled_caching_does_not_prepare_identity(
|
||||
cache: Arc<dyn ResponseCacheService>,
|
||||
#[case] configured: bool,
|
||||
#[case] options_present: bool,
|
||||
#[case] policy: CachePolicy,
|
||||
) {
|
||||
let request = CacheRequest::prepare::<TestRoute>(
|
||||
ProviderIdentity {
|
||||
model: "test".into(),
|
||||
provider: "test".into(),
|
||||
},
|
||||
configured.then_some(cache),
|
||||
options_present.then_some(CacheOptions {
|
||||
policy,
|
||||
scope: CacheScope::Shared,
|
||||
}),
|
||||
&|_| panic!("disabled cache must not prepare identity"),
|
||||
);
|
||||
let response = execute_unary::<TestRoute, _, _>(request, &(), None, || async {
|
||||
Ok(json!({"source": "provider"}))
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response, json!({"source": "provider"}));
|
||||
}
|
||||
|
||||
fn cache_with_limit(max_entry_bytes: usize) -> Arc<dyn ResponseCacheService> {
|
||||
Arc::new(
|
||||
ResponseCache::new(Arc::new(InMemoryCache::new(
|
||||
|
|
@ -96,9 +134,7 @@ async fn call(
|
|||
request: Value,
|
||||
) -> Value {
|
||||
let output = execute_streaming::<TestRoute, _, _>(
|
||||
cache_request(request),
|
||||
Some(cache.clone()),
|
||||
options,
|
||||
cache_request::<TestRoute>(request, Some(cache.clone()), options),
|
||||
&(),
|
||||
None,
|
||||
|| async {
|
||||
|
|
@ -193,9 +229,11 @@ async fn streamed(
|
|||
fail: bool,
|
||||
) -> OutputOf<TestRoute> {
|
||||
execute_streaming::<TestRoute, _, _>(
|
||||
cache_request(json!({"stream":true})),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
cache_request::<TestRoute>(
|
||||
json!({"stream":true}),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
),
|
||||
&(),
|
||||
None,
|
||||
|| async {
|
||||
|
|
@ -309,9 +347,11 @@ async fn oversized_streams_are_delivered_without_being_stored() {
|
|||
async fn a_provider_failure_never_populates_the_cache(cache: Arc<dyn ResponseCacheService>) {
|
||||
let calls = AtomicUsize::new(0);
|
||||
let first = execute_streaming::<TestRoute, _, _>(
|
||||
cache_request(json!({})),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
cache_request::<TestRoute>(
|
||||
json!({}),
|
||||
Some(cache.clone()),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
),
|
||||
&(),
|
||||
None,
|
||||
|| async {
|
||||
|
|
@ -345,31 +385,55 @@ struct InvalidEntryCache(
|
|||
);
|
||||
|
||||
impl ResponseCacheService for InvalidEntryCache {
|
||||
fn config(&self) -> &ResponseCacheConfig {
|
||||
self.0.config()
|
||||
fn prepare(
|
||||
self: Arc<Self>,
|
||||
options: CacheOptions,
|
||||
surface: &str,
|
||||
input: &dyn Fn() -> Value,
|
||||
) -> litellm_cache_response::PreparedCache {
|
||||
litellm_cache_response::PreparedCache {
|
||||
policy: options.policy,
|
||||
max_entry_bytes: self.0.config().max_entry_bytes,
|
||||
call: Arc::new(InvalidEntryCall {
|
||||
request: options.request(&self.0.config().namespace, surface, input()),
|
||||
service: self,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct InvalidEntryCall {
|
||||
service: Arc<InvalidEntryCache>,
|
||||
request: litellm_cache_response::ResponseCacheRequest,
|
||||
}
|
||||
|
||||
impl litellm_cache_response::ResponseCacheCall for InvalidEntryCall {
|
||||
fn lookup<'a>(
|
||||
&'a self,
|
||||
request: &'a litellm_cache_response::ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> futures_util::future::BoxFuture<'a, Result<Option<Value>, litellm_cache::Error>> {
|
||||
) -> futures_util::future::BoxFuture<
|
||||
'a,
|
||||
Result<Option<litellm_cache_response::CacheHit>, litellm_cache::Error>,
|
||||
> {
|
||||
Box::pin(async move {
|
||||
Ok(self
|
||||
.0
|
||||
.async_lookup(request, now)
|
||||
.await?
|
||||
.or_else(|| Some(self.1.clone())))
|
||||
Ok(Some(litellm_cache_response::CacheHit {
|
||||
value: self
|
||||
.service
|
||||
.0
|
||||
.async_lookup(&self.request, now)
|
||||
.await?
|
||||
.unwrap_or_else(|| self.service.1.clone()),
|
||||
key: litellm_cache_response::cache_key(&self.request.key),
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
fn store<'a>(
|
||||
&'a self,
|
||||
request: &'a litellm_cache_response::ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> futures_util::future::BoxFuture<'a, Result<(), litellm_cache::Error>> {
|
||||
Box::pin(self.0.async_store(request, response, now))
|
||||
Box::pin(self.service.0.async_store(&self.request, response, now))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -482,6 +546,7 @@ impl Interceptors<RouteError> for CacheHitAccounting {
|
|||
InterceptRequest::AfterProviderResponse { raw, reply } => {
|
||||
reply.send(self.after_provider_response(raw).await?);
|
||||
}
|
||||
other => other.pass_through(),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -565,9 +630,11 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure(
|
|||
};
|
||||
let result = if streaming_route {
|
||||
match execute_streaming::<TestRoute, _, _>(
|
||||
cache_request(request),
|
||||
Some(cache),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
cache_request::<TestRoute>(
|
||||
request,
|
||||
Some(cache),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
),
|
||||
&accounting,
|
||||
Some(&observer),
|
||||
|| async { panic!("a cache hit must not call the provider") },
|
||||
|
|
@ -579,9 +646,11 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure(
|
|||
}
|
||||
} else {
|
||||
execute_unary::<UnaryTestRoute, _, _>(
|
||||
cache_request(request),
|
||||
Some(cache),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
cache_request::<UnaryTestRoute>(
|
||||
request,
|
||||
Some(cache),
|
||||
Some(CacheOptions::new(CacheScope::Shared)),
|
||||
),
|
||||
&accounting,
|
||||
Some(&observer),
|
||||
|| async { panic!("a cache hit must not call the provider") },
|
||||
|
|
@ -627,9 +696,7 @@ async fn unary_call(
|
|||
request: Value,
|
||||
) -> Value {
|
||||
execute_unary::<UnaryTestRoute, _, _>(
|
||||
cache_request(request),
|
||||
Some(cache.clone()),
|
||||
options,
|
||||
cache_request::<UnaryTestRoute>(request, Some(cache.clone()), options),
|
||||
&(),
|
||||
None,
|
||||
|| async { Ok(json!({"call":calls.fetch_add(1, Ordering::SeqCst)})) },
|
||||
|
|
|
|||
|
|
@ -118,6 +118,9 @@ impl From<litellm_host::machine::MachineFault> for Error {
|
|||
Self::InvalidRequest(match fault {
|
||||
MachineFault::Abandoned => "OCR host driver was abandoned".into(),
|
||||
MachineFault::Protocol(message) => format!("OCR {message}"),
|
||||
MachineFault::UnreadableAnswer => {
|
||||
"OCR host answered with an unreadable response".into()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# Cache boundary
|
||||
|
||||
This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
|
||||
This folder owns how Rust inference reaches the configured cache. Rust holds no cache rules: `litellm.rust_bridge.host.cache` decides whether a call may read or write, builds the key and stores the entry, using the same Python handler functions the Python route runs. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
|
||||
|
||||
`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters
|
||||
`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `python/service.rs` owns the `Cached<P>` protocol and typed cache operations. `python/binding.rs` owns `CachedBinding<H>`, which wraps a route's Python binding, asks the Python host module for access and serves cache operations. Route bindings never hold cache state
|
||||
|
||||
`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports
|
||||
`python/` delegates operations to the Python cache the host module granted. `native/` owns the experimental V2 native cache handles exposed to Python. Inference does not use them until the Rust cache is ready
|
||||
|
|
|
|||
|
|
@ -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::{
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -12,9 +12,9 @@ use litellm_cache_memory::InMemoryCache;
|
|||
use litellm_cache_redis::{RedisCache, RedisTopology};
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec,
|
||||
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
|
||||
ResponseCacheConfig, ResponseCacheRequest,
|
||||
};
|
||||
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
|
||||
use pyo3::{exceptions::PyValueError, prelude::*};
|
||||
|
||||
#[pyclass(
|
||||
frozen,
|
||||
|
|
@ -22,7 +22,6 @@ use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
|
|||
module = "litellm.rust_bridge._native"
|
||||
)]
|
||||
pub(crate) struct NativeCacheHandle {
|
||||
service: Arc<dyn ResponseCacheService>,
|
||||
backend: Arc<dyn ExactResponseCache>,
|
||||
storage: Storage,
|
||||
pid: u32,
|
||||
|
|
@ -71,7 +70,6 @@ impl NativeCacheHandle {
|
|||
},
|
||||
));
|
||||
Ok(Self {
|
||||
service: backend.clone(),
|
||||
backend,
|
||||
storage: Storage::Memory(storage),
|
||||
pid: std::process::id(),
|
||||
|
|
@ -112,7 +110,6 @@ impl NativeCacheHandle {
|
|||
},
|
||||
));
|
||||
Ok(Self {
|
||||
service: backend.clone(),
|
||||
backend,
|
||||
storage: Storage::Redis(storage),
|
||||
pid: std::process::id(),
|
||||
|
|
@ -259,85 +256,6 @@ fn duration(seconds: f64) -> PyResult<Duration> {
|
|||
.ok_or_else(|| PyValueError::new_err("cache durations must be finite and positive"))
|
||||
}
|
||||
|
||||
pub(in crate::cache) fn native_handle<'py>(
|
||||
configured: &Bound<'py, PyAny>,
|
||||
) -> PyResult<Option<Bound<'py, PyAny>>> {
|
||||
Ok(configured
|
||||
.getattr_opt("cache")?
|
||||
.map(|backend| backend.getattr_opt("native_handle"))
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.filter(|handle| handle.is_instance_of::<NativeCacheHandle>()))
|
||||
}
|
||||
|
||||
pub(in crate::cache) fn configured(
|
||||
configured: &Bound<'_, PyAny>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
) -> PyResult<(
|
||||
Option<Arc<dyn ResponseCacheService>>,
|
||||
litellm_cache_response::CacheOptions,
|
||||
)> {
|
||||
let handle = native_handle(configured)?.ok_or_else(|| {
|
||||
pyo3::exceptions::PyRuntimeError::new_err(
|
||||
"native inference needs a v2 cache, the catalog keeps Python caches on Python",
|
||||
)
|
||||
})?;
|
||||
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;
|
||||
cache.check_process()?;
|
||||
let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none());
|
||||
let controls = controls
|
||||
.as_ref()
|
||||
.map(|value| value.cast::<PyDict>())
|
||||
.transpose()?;
|
||||
if let Some(controls) = controls {
|
||||
for name in controls.keys() {
|
||||
let name = name.extract::<String>()?;
|
||||
if !matches!(
|
||||
name.as_str(),
|
||||
"no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache"
|
||||
) {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"unsupported v2 cache control: {name}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
let boolean = |name: &str| -> PyResult<bool> {
|
||||
controls
|
||||
.map(|values| values.get_item(name))
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.map(|value| value.extract())
|
||||
.transpose()
|
||||
.map(|value| value.unwrap_or(false))
|
||||
};
|
||||
let seconds = |name: &str| -> PyResult<Option<Duration>> {
|
||||
controls
|
||||
.map(|values| values.get_item(name))
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.map(|value| duration(value.extract()?))
|
||||
.transpose()
|
||||
};
|
||||
Ok((
|
||||
Some(cache.service.clone()),
|
||||
litellm_cache_response::CacheOptions {
|
||||
policy: litellm_cache_response::CachePolicy {
|
||||
caching: kwargs
|
||||
.get_item("caching")?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| value.extract())
|
||||
.transpose()?,
|
||||
no_cache: boolean("no-cache")?,
|
||||
no_store: boolean("no-store")?,
|
||||
ttl: seconds("ttl")?,
|
||||
max_age: seconds("s-max-age")?.or(seconds("s-maxage")?),
|
||||
},
|
||||
scope: litellm_cache_response::CacheScope::Shared,
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
fn now() -> Duration {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
338
litellm-rust/crates/python-bridge/src/cache/python/binding.rs
vendored
Normal file
338
litellm-rust/crates/python-bridge/src/cache/python/binding.rs
vendored
Normal file
|
|
@ -0,0 +1,338 @@
|
|||
use super::service::{CacheCall, Cached};
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_response::{CacheHit, CachePolicy};
|
||||
use litellm_host::protocol::{Protocol, Reply};
|
||||
use litellm_host_python::{
|
||||
InvokeError, PythonBinding, PythonHostCalls, PythonOwned, effective_py_args, from_py, to_py,
|
||||
};
|
||||
use pyo3::{
|
||||
exceptions::{PyException, PyRuntimeError},
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
prelude::*,
|
||||
types::PyDict,
|
||||
};
|
||||
|
||||
const MODULE: &str = "litellm.rust_bridge.host.cache";
|
||||
|
||||
pub(crate) struct CachedBinding<H> {
|
||||
inner: H,
|
||||
cache: PythonCache,
|
||||
call_type: &'static str,
|
||||
positional: Option<Py<PyDict>>,
|
||||
}
|
||||
|
||||
impl<H> CachedBinding<H> {
|
||||
pub(crate) fn new(
|
||||
inner: H,
|
||||
call_type: &'static str,
|
||||
asynchronous: bool,
|
||||
positional: Py<PyDict>,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
cache: PythonCache::new(asynchronous),
|
||||
call_type,
|
||||
positional: Some(positional),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type ErrorOf<H> = <<H as PythonBinding>::Protocol as Protocol>::Error;
|
||||
|
||||
impl<H: PythonBinding> PythonBinding for CachedBinding<H> {
|
||||
type Protocol = Cached<H::Protocol>;
|
||||
type Failure = H::Failure;
|
||||
|
||||
fn decode_request(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
arguments: &Bound<'_, PyDict>,
|
||||
) -> Result<<Self::Protocol as Protocol>::Request, InvokeError<ErrorOf<H>>> {
|
||||
let request = self.inner.decode_request(py, arguments)?;
|
||||
let policy = self
|
||||
.positional
|
||||
.as_ref()
|
||||
.ok_or_else(|| PyRuntimeError::new_err("cache binding used after close"))
|
||||
.and_then(|positional| effective_py_args(positional.bind(py), arguments))
|
||||
.and_then(|keyed| self.cache.configure(py, &keyed, self.call_type))
|
||||
.map_err(InvokeError::Python)?;
|
||||
Ok((request, policy))
|
||||
}
|
||||
|
||||
fn encode_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: <Self::Protocol as Protocol>::Response,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.inner.encode_response(py, response)
|
||||
}
|
||||
|
||||
fn public_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: &serde_json::Value,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.inner.public_response(py, response)
|
||||
}
|
||||
|
||||
fn response_value(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
public: &Bound<'_, PyAny>,
|
||||
) -> PyResult<serde_json::Value> {
|
||||
self.inner.response_value(py, public)
|
||||
}
|
||||
|
||||
fn encode_stream_head(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
head: <Self::Protocol as Protocol>::StreamHead,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.inner.encode_stream_head(py, head)
|
||||
}
|
||||
|
||||
fn encode_chunk(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
chunk: <Self::Protocol as Protocol>::Chunk,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.inner.encode_chunk(py, chunk)
|
||||
}
|
||||
|
||||
fn map_error(&self, py: Python<'_>, error: ErrorOf<H>) -> PyResult<H::Failure> {
|
||||
self.inner.map_error(py, error)
|
||||
}
|
||||
|
||||
fn host_error(error: &PyErr) -> ErrorOf<H> {
|
||||
H::host_error(error)
|
||||
}
|
||||
}
|
||||
|
||||
impl<H: PythonBinding> PythonHostCalls<Cached<H::Protocol>> for CachedBinding<H> {
|
||||
fn handle_host_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
call: CacheCall,
|
||||
) -> Result<(), InvokeError<ErrorOf<H>>> {
|
||||
self.begin_host_call(py, call).map(|_| ())
|
||||
}
|
||||
|
||||
fn begin_host_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
call: CacheCall,
|
||||
) -> Result<Option<Py<PyAny>>, InvokeError<ErrorOf<H>>> {
|
||||
self.cache.begin(py, call).map_err(InvokeError::Python)
|
||||
}
|
||||
|
||||
fn resume_host_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> Result<Option<Py<PyAny>>, InvokeError<ErrorOf<H>>> {
|
||||
self.cache.resume(py, result).map_err(InvokeError::Python)
|
||||
}
|
||||
}
|
||||
|
||||
impl<H: PythonOwned> PythonOwned for CachedBinding<H> {
|
||||
fn close(&mut self, py: Python<'_>) {
|
||||
self.inner.close(py);
|
||||
self.cache.close();
|
||||
self.positional = None;
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
self.inner.traverse(visit)?;
|
||||
self.cache.traverse(visit)?;
|
||||
if let Some(positional) = &self.positional {
|
||||
visit.call(positional)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
enum Pending {
|
||||
Lookup(Reply<Result<Option<CacheHit>, Error>>),
|
||||
Store(Reply<Result<(), Error>>),
|
||||
}
|
||||
|
||||
struct Granted {
|
||||
access: Py<PyAny>,
|
||||
}
|
||||
|
||||
fn cache_hit(hit: &Bound<'_, PyAny>) -> PyResult<Option<CacheHit>> {
|
||||
if hit.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(CacheHit {
|
||||
value: from_py(&hit.getattr("response")?)?,
|
||||
key: hit.getattr("key")?.extract()?,
|
||||
}))
|
||||
}
|
||||
|
||||
struct PythonCache {
|
||||
module: &'static str,
|
||||
granted: Option<Granted>,
|
||||
pending: Option<Pending>,
|
||||
asynchronous: bool,
|
||||
}
|
||||
|
||||
impl PythonCache {
|
||||
fn new(asynchronous: bool) -> Self {
|
||||
Self {
|
||||
module: MODULE,
|
||||
granted: None,
|
||||
pending: None,
|
||||
asynchronous,
|
||||
}
|
||||
}
|
||||
|
||||
fn configure(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
arguments: &Bound<'_, PyDict>,
|
||||
call_type: &str,
|
||||
) -> PyResult<Option<CachePolicy>> {
|
||||
let access = py
|
||||
.import(self.module)?
|
||||
.getattr("access")?
|
||||
.call1((arguments, call_type))?;
|
||||
if access.is_none() {
|
||||
self.granted = None;
|
||||
return Ok(None);
|
||||
}
|
||||
let policy = CachePolicy {
|
||||
no_cache: !access.getattr("read")?.extract::<bool>()?,
|
||||
no_store: !access.getattr("write")?.extract::<bool>()?,
|
||||
..CachePolicy::default()
|
||||
};
|
||||
self.granted = Some(Granted {
|
||||
access: access.unbind(),
|
||||
});
|
||||
Ok(Some(policy))
|
||||
}
|
||||
|
||||
fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult<Option<Py<PyAny>>> {
|
||||
let granted = self.granted.as_ref().ok_or_else(|| {
|
||||
PyRuntimeError::new_err("cache operation without granted cache access")
|
||||
})?;
|
||||
let access = granted.access.bind(py);
|
||||
let function = |name: &str| py.import(self.module)?.getattr(name);
|
||||
let (pending, result) = match call {
|
||||
CacheCall::Lookup { reply } => (
|
||||
Pending::Lookup(reply),
|
||||
function(if self.asynchronous {
|
||||
"alookup"
|
||||
} else {
|
||||
"lookup"
|
||||
})
|
||||
.and_then(|lookup| lookup.call1((access,))),
|
||||
),
|
||||
CacheCall::Store { value, reply } => (
|
||||
Pending::Store(reply),
|
||||
function(if self.asynchronous { "astore" } else { "store" })
|
||||
.and_then(|store| store.call1((access, to_py(py, &value)?))),
|
||||
),
|
||||
};
|
||||
let result = result.map(Bound::unbind);
|
||||
self.pending = Some(pending);
|
||||
if self.asynchronous && result.is_ok() {
|
||||
return result.map(Some);
|
||||
}
|
||||
self.resume(py, result)
|
||||
}
|
||||
|
||||
fn resume(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<Option<Py<PyAny>>> {
|
||||
if let Err(error) = &result
|
||||
&& !error.is_instance_of::<PyException>(py)
|
||||
{
|
||||
self.pending = None;
|
||||
return Err(result.err().unwrap());
|
||||
}
|
||||
match self.pending.take() {
|
||||
Some(Pending::Lookup(reply)) => {
|
||||
let hit = result
|
||||
.map_err(|_| Error::Unavailable)
|
||||
.and_then(|hit| cache_hit(hit.bind(py)).map_err(|_| Error::InvalidEntry));
|
||||
reply.send(hit);
|
||||
}
|
||||
Some(Pending::Store(reply)) => {
|
||||
reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable));
|
||||
}
|
||||
None => {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"cache reply without pending operation",
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn close(&mut self) {
|
||||
self.pending = None;
|
||||
self.granted = None;
|
||||
}
|
||||
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
let Some(granted) = &self.granted else {
|
||||
return Ok(());
|
||||
};
|
||||
visit.call(&granted.access)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures_util::FutureExt;
|
||||
use litellm_host::protocol::reply;
|
||||
use rstest::rstest;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
fn a_cache_call_that_cannot_reach_python_answers_unavailable(
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
#[values(false, true)] store: bool,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let mut cache = PythonCache {
|
||||
module: "litellm_missing_cache_host",
|
||||
granted: Some(Granted { access: py.None() }),
|
||||
pending: None,
|
||||
asynchronous,
|
||||
};
|
||||
let answer = if store {
|
||||
let (reply, answer) = reply();
|
||||
assert!(
|
||||
cache
|
||||
.begin(
|
||||
py,
|
||||
CacheCall::Store {
|
||||
value: Value::Null,
|
||||
reply
|
||||
}
|
||||
)
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
answer.now_or_never().map(|answer| answer.unwrap().err())
|
||||
} else {
|
||||
let (reply, answer) = reply();
|
||||
assert!(
|
||||
cache
|
||||
.begin(py, CacheCall::Lookup { reply })
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
answer.now_or_never().map(|answer| answer.unwrap().err())
|
||||
};
|
||||
assert!(matches!(answer, Some(Some(Error::Unavailable))));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -1,141 +0,0 @@
|
|||
use litellm_cache::Error;
|
||||
use litellm_host::protocol::Reply;
|
||||
use litellm_host_python::{from_py, to_py};
|
||||
use pyo3::{
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
prelude::*,
|
||||
types::PyDict,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::service::CacheCall;
|
||||
|
||||
enum Pending {
|
||||
Lookup(Reply<Result<Option<Value>, Error>>),
|
||||
Store(Reply<Result<(), Error>>),
|
||||
}
|
||||
|
||||
pub(crate) struct PythonCache {
|
||||
cache: Option<Py<PyAny>>,
|
||||
arguments: Option<Py<PyDict>>,
|
||||
pending: Option<Pending>,
|
||||
asynchronous: bool,
|
||||
}
|
||||
|
||||
impl PythonCache {
|
||||
pub fn new(asynchronous: bool) -> Self {
|
||||
Self {
|
||||
cache: None,
|
||||
arguments: None,
|
||||
pending: None,
|
||||
asynchronous,
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::cache) fn bind(
|
||||
&mut self,
|
||||
cache: Bound<'_, PyAny>,
|
||||
arguments: &Bound<'_, PyDict>,
|
||||
) {
|
||||
self.cache = Some(cache.unbind());
|
||||
self.arguments = Some(arguments.clone().unbind());
|
||||
}
|
||||
|
||||
pub fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult<Option<Py<PyAny>>> {
|
||||
let Some(cache) = self.cache.as_ref() else {
|
||||
return Err(pyo3::exceptions::PyRuntimeError::new_err(
|
||||
"cache operation without configured cache",
|
||||
));
|
||||
};
|
||||
let arguments = self
|
||||
.arguments
|
||||
.as_ref()
|
||||
.ok_or_else(|| {
|
||||
pyo3::exceptions::PyRuntimeError::new_err("cache arguments unavailable")
|
||||
})?
|
||||
.bind(py)
|
||||
.copy()?;
|
||||
let (method, result) = match call {
|
||||
CacheCall::Lookup { reply } => {
|
||||
self.pending = Some(Pending::Lookup(reply));
|
||||
(
|
||||
if self.asynchronous {
|
||||
"async_get_cache"
|
||||
} else {
|
||||
"get_cache"
|
||||
},
|
||||
None,
|
||||
)
|
||||
}
|
||||
CacheCall::Store { value, reply } => {
|
||||
self.pending = Some(Pending::Store(reply));
|
||||
(
|
||||
if self.asynchronous {
|
||||
"async_add_cache"
|
||||
} else {
|
||||
"add_cache"
|
||||
},
|
||||
Some(to_py(py, &value)?),
|
||||
)
|
||||
}
|
||||
};
|
||||
let result = match result {
|
||||
Some(value) => cache
|
||||
.bind(py)
|
||||
.call_method(method, (value,), Some(&arguments)),
|
||||
None => cache.bind(py).call_method(method, (), Some(&arguments)),
|
||||
}
|
||||
.map(Bound::unbind);
|
||||
if self.asynchronous && result.is_ok() {
|
||||
return result.map(Some);
|
||||
}
|
||||
self.resume(py, result)
|
||||
}
|
||||
|
||||
pub fn resume(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> PyResult<Option<Py<PyAny>>> {
|
||||
if let Err(error) = &result
|
||||
&& !error.is_instance_of::<pyo3::exceptions::PyException>(py)
|
||||
{
|
||||
self.pending = None;
|
||||
return Err(result.err().unwrap());
|
||||
}
|
||||
match self.pending.take() {
|
||||
Some(Pending::Lookup(reply)) => {
|
||||
let value = result.map_err(|_| Error::Unavailable).and_then(|value| {
|
||||
if value.bind(py).is_none() {
|
||||
Ok(None)
|
||||
} else {
|
||||
from_py(value.bind(py))
|
||||
.map(Some)
|
||||
.map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
});
|
||||
reply.send(value);
|
||||
}
|
||||
Some(Pending::Store(reply)) => {
|
||||
reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable));
|
||||
}
|
||||
None => {
|
||||
return Err(pyo3::exceptions::PyRuntimeError::new_err(
|
||||
"cache reply without pending operation",
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub fn close(&mut self) {
|
||||
self.pending = None;
|
||||
self.cache = None;
|
||||
self.arguments = None;
|
||||
}
|
||||
|
||||
pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
visit.call(&self.cache)?;
|
||||
visit.call(&self.arguments)
|
||||
}
|
||||
}
|
||||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -1,44 +1,31 @@
|
|||
use std::{future::Future, pin::Pin, time::Duration};
|
||||
use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_response::{
|
||||
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope,
|
||||
CacheHit, CacheOptions, CachePolicy, CacheScope, PreparedCache, ResponseCacheCall,
|
||||
ResponseCacheService, ResponseEnvelope, ScopedCache,
|
||||
};
|
||||
use litellm_host::{
|
||||
machine::{HostServices, MachineFault},
|
||||
machine::HostServices,
|
||||
protocol::{Protocol, Reply},
|
||||
};
|
||||
use litellm_inference::caching::CachedOutput;
|
||||
use litellm_inference::caching::{Cachable, CachedOutput};
|
||||
use serde_json::Value;
|
||||
|
||||
const STREAM_EVENTS_KEY: &str = "litellm_cached_anthropic_sse_events";
|
||||
pub(crate) struct Cached<P>(PhantomData<P>);
|
||||
|
||||
fn from_python(value: Value) -> Result<Value, Error> {
|
||||
let output = match value.get(STREAM_EVENTS_KEY) {
|
||||
Some(events) => {
|
||||
let events: Vec<String> =
|
||||
serde_json::from_value(events.clone()).map_err(|_| Error::InvalidEntry)?;
|
||||
CachedOutput::Stream(events.concat())
|
||||
}
|
||||
None => CachedOutput::Response(value),
|
||||
};
|
||||
serde_json::to_value(ResponseEnvelope::new("messages", output)).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn to_python(value: Value) -> Result<Value, Error> {
|
||||
let envelope: ResponseEnvelope<CachedOutput<Value>> =
|
||||
serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?;
|
||||
match envelope.decode("messages").ok_or(Error::InvalidEntry)? {
|
||||
CachedOutput::Response(response) => Ok(response),
|
||||
CachedOutput::Stream(text) => Ok(serde_json::json!({
|
||||
STREAM_EVENTS_KEY: text.split_inclusive("\n\n").collect::<Vec<_>>()
|
||||
})),
|
||||
}
|
||||
impl<P: Protocol> Protocol for Cached<P> {
|
||||
type Request = (P::Request, Option<CachePolicy>);
|
||||
type Response = P::Response;
|
||||
type Error = P::Error;
|
||||
type HostCall = CacheCall;
|
||||
type Chunk = P::Chunk;
|
||||
type StreamHead = P::StreamHead;
|
||||
}
|
||||
|
||||
pub(crate) enum CacheCall {
|
||||
Lookup {
|
||||
reply: Reply<Result<Option<Value>, Error>>,
|
||||
reply: Reply<Result<Option<CacheHit>, Error>>,
|
||||
},
|
||||
Store {
|
||||
value: Value,
|
||||
|
|
@ -46,58 +33,76 @@ pub(crate) enum CacheCall {
|
|||
},
|
||||
}
|
||||
|
||||
struct PythonCacheService<P: Protocol> {
|
||||
services: HostServices<P>,
|
||||
config: ResponseCacheConfig,
|
||||
struct PythonCacheService<P: Cachable> {
|
||||
services: HostServices<Cached<P>>,
|
||||
protocol: PhantomData<fn() -> P>,
|
||||
}
|
||||
|
||||
pub(in crate::cache) fn service<P: Protocol<HostCall = CacheCall>>(
|
||||
services: HostServices<P>,
|
||||
namespace: String,
|
||||
) -> std::sync::Arc<dyn ResponseCacheService>
|
||||
where
|
||||
P::Error: From<MachineFault>,
|
||||
{
|
||||
std::sync::Arc::new(PythonCacheService {
|
||||
services,
|
||||
config: ResponseCacheConfig {
|
||||
namespace,
|
||||
..Default::default()
|
||||
},
|
||||
})
|
||||
pub(crate) fn scoped<P: Cachable>(services: HostServices<Cached<P>>) -> ScopedCache {
|
||||
ScopedCache::new(
|
||||
Arc::new(PythonCacheService {
|
||||
services,
|
||||
protocol: PhantomData,
|
||||
}),
|
||||
CacheScope::Shared,
|
||||
)
|
||||
}
|
||||
|
||||
impl<P: Protocol<HostCall = CacheCall>> ResponseCacheService for PythonCacheService<P>
|
||||
where
|
||||
P::Error: From<MachineFault>,
|
||||
{
|
||||
fn config(&self) -> &ResponseCacheConfig {
|
||||
&self.config
|
||||
fn from_python<P: Cachable>(value: Value) -> Result<Value, Error> {
|
||||
let output: CachedOutput<Value> =
|
||||
serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?;
|
||||
serde_json::to_value(ResponseEnvelope::new(P::SURFACE, output)).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn to_python<P: Cachable>(value: Value) -> Result<Value, Error> {
|
||||
let envelope: ResponseEnvelope<CachedOutput<Value>> =
|
||||
serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?;
|
||||
let output = envelope.decode(P::SURFACE).ok_or(Error::InvalidEntry)?;
|
||||
serde_json::to_value(output).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
impl<P: Cachable> ResponseCacheService for PythonCacheService<P> {
|
||||
fn prepare(
|
||||
self: Arc<Self>,
|
||||
options: CacheOptions,
|
||||
_: &str,
|
||||
_: &dyn Fn() -> Value,
|
||||
) -> PreparedCache {
|
||||
PreparedCache {
|
||||
policy: options.policy,
|
||||
max_entry_bytes: usize::MAX,
|
||||
call: self,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<P: Cachable> ResponseCacheCall for PythonCacheService<P> {
|
||||
fn lookup<'a>(
|
||||
&'a self,
|
||||
_: &'a ResponseCacheRequest,
|
||||
_: Duration,
|
||||
) -> Pin<Box<dyn Future<Output = Result<Option<Value>, Error>> + Send + 'a>> {
|
||||
) -> Pin<Box<dyn Future<Output = Result<Option<CacheHit>, Error>> + Send + 'a>> {
|
||||
Box::pin(async move {
|
||||
self.services
|
||||
.call(|reply| CacheCall::Lookup { reply })
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)??
|
||||
.map(from_python)
|
||||
.map(|hit| {
|
||||
Ok(CacheHit {
|
||||
value: from_python::<P>(hit.value)?,
|
||||
key: hit.key,
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
})
|
||||
}
|
||||
|
||||
fn store<'a>(
|
||||
&'a self,
|
||||
_: &'a ResponseCacheRequest,
|
||||
value: Value,
|
||||
_: Duration,
|
||||
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'a>> {
|
||||
Box::pin(async move {
|
||||
let value = to_python(value)?;
|
||||
let value = to_python::<P>(value)?;
|
||||
self.services
|
||||
.call(|reply| CacheCall::Store { value, reply })
|
||||
.await
|
||||
|
|
@ -105,3 +110,122 @@ where
|
|||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::convert::Infallible;
|
||||
|
||||
use litellm_host::{
|
||||
interceptors::ProviderIdentity,
|
||||
machine::{CallMachine, Machine, MachineStep},
|
||||
protocol::{HostRequest, Protocol},
|
||||
};
|
||||
use litellm_inference::{
|
||||
RouteError,
|
||||
caching::{CacheRequest, execute_unary},
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
struct TestRoute;
|
||||
|
||||
impl Protocol for TestRoute {
|
||||
type Request = ();
|
||||
type Response = Value;
|
||||
type Error = RouteError;
|
||||
type HostCall = Infallible;
|
||||
type Chunk = ();
|
||||
type StreamHead = ();
|
||||
}
|
||||
|
||||
impl Cachable for TestRoute {
|
||||
const SURFACE: &'static str = "test";
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::hit(true)]
|
||||
#[case::miss(false)]
|
||||
#[tokio::test]
|
||||
async fn python_cache_calls_never_materialize_native_identity(#[case] hit: bool) {
|
||||
let mut machine = CallMachine::<Cached<TestRoute>>::new(None, |host| {
|
||||
Box::pin(async move {
|
||||
let cache = scoped::<TestRoute>(host.services);
|
||||
let request = CacheRequest::prepare::<TestRoute>(
|
||||
ProviderIdentity {
|
||||
model: "test".into(),
|
||||
provider: "test".into(),
|
||||
},
|
||||
Some(cache.service.clone()),
|
||||
Some(cache.options(None)),
|
||||
&|_| panic!("Python caching must not construct native identity"),
|
||||
);
|
||||
execute_unary::<TestRoute, _, _>(request, &(), None, || async {
|
||||
Ok(json!({"source": "provider"}))
|
||||
})
|
||||
.await
|
||||
})
|
||||
});
|
||||
let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Lookup { reply })) =
|
||||
machine.resume().await.unwrap()
|
||||
else {
|
||||
panic!("expected Python lookup");
|
||||
};
|
||||
reply.send(Ok(hit.then(|| CacheHit {
|
||||
value: json!({"kind": "Response", "value": {"source": "cache"}}),
|
||||
key: "python-owned-key".into(),
|
||||
})));
|
||||
if !hit {
|
||||
let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Store { value, reply })) =
|
||||
machine.resume().await.unwrap()
|
||||
else {
|
||||
panic!("expected Python store on a miss");
|
||||
};
|
||||
assert_eq!(
|
||||
value,
|
||||
json!({"kind": "Response", "value": {"source": "provider"}})
|
||||
);
|
||||
reply.send(Ok(()));
|
||||
}
|
||||
let MachineStep::Complete(response) = machine.resume().await.unwrap() else {
|
||||
panic!("expected completed call");
|
||||
};
|
||||
assert_eq!(
|
||||
response,
|
||||
json!({"source": if hit { "cache" } else { "provider" }})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn python_hit_keeps_the_host_key() {
|
||||
let mut machine = CallMachine::<Cached<TestRoute>, CacheHit>::new(None, |host| {
|
||||
Box::pin(async move {
|
||||
let cache = scoped::<TestRoute>(host.services);
|
||||
let options = cache.options(None);
|
||||
let prepared = cache.service.prepare(options, TestRoute::SURFACE, &|| {
|
||||
panic!("Python caching must not construct native identity")
|
||||
});
|
||||
Ok(prepared.call.lookup(Duration::ZERO).await.unwrap().unwrap())
|
||||
})
|
||||
});
|
||||
let MachineStep::Suspended(HostRequest::HostCall(CacheCall::Lookup { reply })) =
|
||||
machine.resume().await.unwrap()
|
||||
else {
|
||||
panic!("expected Python lookup");
|
||||
};
|
||||
reply.send(Ok(Some(CacheHit {
|
||||
value: json!({"kind": "Response", "value": {"answer": 42}}),
|
||||
key: "python-owned-key".into(),
|
||||
})));
|
||||
let MachineStep::Complete(hit) = machine.resume().await.unwrap() else {
|
||||
panic!("expected completed lookup");
|
||||
};
|
||||
assert_eq!(hit.key, "python-owned-key");
|
||||
assert_eq!(
|
||||
from_python::<TestRoute>(json!({"kind": "Response", "value": {"answer": 42}})).unwrap(),
|
||||
hit.value
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,171 +0,0 @@
|
|||
use super::{native, python};
|
||||
use litellm_cache_response::{
|
||||
CacheOptions, CachePolicy, CacheScope, ResponseCacheService, ScopedCache,
|
||||
};
|
||||
use litellm_host::{
|
||||
machine::{HostServices, MachineFault},
|
||||
protocol::Protocol,
|
||||
};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
use std::sync::Arc;
|
||||
|
||||
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
|
||||
|
||||
pub(crate) struct Cached<P>(std::marker::PhantomData<P>);
|
||||
|
||||
impl<P: Protocol> Protocol for Cached<P> {
|
||||
type Request = (P::Request, Selection);
|
||||
type Response = P::Response;
|
||||
type Error = P::Error;
|
||||
type HostCall = python::CacheCall;
|
||||
type Chunk = P::Chunk;
|
||||
type StreamHead = P::StreamHead;
|
||||
}
|
||||
|
||||
enum Backend {
|
||||
Disabled,
|
||||
Native(Arc<dyn ResponseCacheService>),
|
||||
Python { namespace: String },
|
||||
}
|
||||
|
||||
pub(crate) struct Selection {
|
||||
backend: Backend,
|
||||
options: CacheOptions,
|
||||
}
|
||||
|
||||
impl Selection {
|
||||
pub(crate) fn attach<P: Protocol<HostCall = python::CacheCall>>(
|
||||
self,
|
||||
services: HostServices<P>,
|
||||
) -> (Option<ScopedCache>, CacheOptions)
|
||||
where
|
||||
P::Error: From<MachineFault>,
|
||||
{
|
||||
let service = match self.backend {
|
||||
Backend::Disabled => None,
|
||||
Backend::Native(service) => Some(service),
|
||||
Backend::Python { namespace } => Some(python::service(services, namespace)),
|
||||
};
|
||||
(
|
||||
service.map(|service| ScopedCache::new(service, CacheScope::Shared)),
|
||||
self.options,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn selected_cache<'py>(
|
||||
py: Python<'py>,
|
||||
kwargs: &Bound<'py, PyDict>,
|
||||
call_type: &str,
|
||||
) -> PyResult<Option<Bound<'py, PyAny>>> {
|
||||
let configured = py
|
||||
.import(CACHE_HOST_MODULE)?
|
||||
.getattr("configured_cache")?
|
||||
.call0()?;
|
||||
if configured.is_none()
|
||||
|| kwargs
|
||||
.get_item("caching")?
|
||||
.is_some_and(|value| value.is(pyo3::types::PyBool::new(py, false)))
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let supported = configured.getattr("supported_call_types")?;
|
||||
if supported.is_none() || !supported.contains(call_type)? {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(configured))
|
||||
}
|
||||
|
||||
pub(crate) fn configured_native(
|
||||
py: Python<'_>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
call_type: &str,
|
||||
) -> PyResult<(
|
||||
Option<Arc<dyn ResponseCacheService>>,
|
||||
litellm_cache_response::CacheOptions,
|
||||
)> {
|
||||
let Some(configured) = selected_cache(py, kwargs, call_type)? else {
|
||||
return Ok((
|
||||
None,
|
||||
litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
|
||||
));
|
||||
};
|
||||
native_configuration(&configured, kwargs)
|
||||
}
|
||||
|
||||
fn native_configuration(
|
||||
configured: &Bound<'_, PyAny>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
) -> PyResult<(Option<Arc<dyn ResponseCacheService>>, CacheOptions)> {
|
||||
if !configured
|
||||
.call_method("should_use_cache", (), Some(kwargs))?
|
||||
.extract::<bool>()?
|
||||
{
|
||||
return Ok((
|
||||
None,
|
||||
litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
|
||||
));
|
||||
}
|
||||
native::v2::configured(configured, kwargs)
|
||||
}
|
||||
|
||||
pub(crate) fn configure(
|
||||
python: &mut python::PythonCache,
|
||||
py: Python<'_>,
|
||||
arguments: &Bound<'_, PyDict>,
|
||||
call_type: &str,
|
||||
) -> PyResult<Selection> {
|
||||
let selected = selected_cache(py, arguments, call_type)?;
|
||||
let Some(cache) = selected else {
|
||||
return Ok(Selection {
|
||||
backend: Backend::Disabled,
|
||||
options: CacheOptions::new(CacheScope::Shared),
|
||||
});
|
||||
};
|
||||
if native::v2::native_handle(&cache)?.is_some() {
|
||||
let (native, options) = native_configuration(&cache, arguments)?;
|
||||
return Ok(Selection {
|
||||
backend: native.map_or(Backend::Disabled, Backend::Native),
|
||||
options,
|
||||
});
|
||||
}
|
||||
let enabled = cache
|
||||
.call_method("should_use_cache", (), Some(arguments))?
|
||||
.extract::<bool>()?;
|
||||
let controls = arguments
|
||||
.get_item("cache")?
|
||||
.filter(|value| !value.is_none());
|
||||
let boolean = |name: &str| -> PyResult<bool> {
|
||||
controls
|
||||
.as_ref()
|
||||
.map(|value| value.cast::<PyDict>()?.get_item(name))
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.map(|value| value.extract())
|
||||
.transpose()
|
||||
.map(|value| value.unwrap_or(false))
|
||||
};
|
||||
let options = CacheOptions {
|
||||
policy: CachePolicy {
|
||||
no_cache: boolean("no-cache")?,
|
||||
no_store: boolean("no-store")?,
|
||||
..CachePolicy::default()
|
||||
},
|
||||
..CacheOptions::new(CacheScope::Shared)
|
||||
};
|
||||
let namespace = cache
|
||||
.getattr_opt("namespace")?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| value.extract())
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
python.bind(cache, arguments);
|
||||
Ok(Selection {
|
||||
backend: if enabled {
|
||||
Backend::Python { namespace }
|
||||
} else {
|
||||
Backend::Disabled
|
||||
},
|
||||
options,
|
||||
})
|
||||
}
|
||||
|
|
@ -126,6 +126,17 @@ pub(crate) fn public_response(
|
|||
.map(Bound::unbind)
|
||||
}
|
||||
|
||||
/// The JSON form of a public response a Python hook handed back, models nested in it included.
|
||||
pub(crate) fn response_value(public: &Bound<'_, PyAny>) -> PyResult<serde_json::Value> {
|
||||
from_py(
|
||||
&public
|
||||
.py()
|
||||
.import("pydantic_core")?
|
||||
.getattr("to_jsonable_python")?
|
||||
.call1((public,))?,
|
||||
)
|
||||
}
|
||||
|
||||
struct RequestFieldSources<'py> {
|
||||
body: Option<Bound<'py, PyAny>>,
|
||||
credentials: Option<Bound<'py, PyAny>>,
|
||||
|
|
|
|||
|
|
@ -59,6 +59,22 @@ impl PythonBinding for ChatCompletionsPythonHost {
|
|||
self.0.response(py, &response)
|
||||
}
|
||||
|
||||
fn public_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: &serde_json::Value,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.0.response(py, response)
|
||||
}
|
||||
|
||||
fn response_value(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
public: &Bound<'_, PyAny>,
|
||||
) -> PyResult<serde_json::Value> {
|
||||
crate::marshal::response_value(public)
|
||||
}
|
||||
|
||||
fn encode_stream_head(
|
||||
&mut self,
|
||||
_py: Python<'_>,
|
||||
|
|
|
|||
|
|
@ -4,13 +4,21 @@ use std::sync::Arc;
|
|||
|
||||
use host::ChatCompletionsPythonHost;
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_cache_response::{CachePolicy, ScopedCache};
|
||||
use litellm_cache_response::ScopedCache;
|
||||
use litellm_callbacks_legacy_python::LoggingOperation;
|
||||
use litellm_host::{call::HostedMachine, protocol::Protocol};
|
||||
use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions};
|
||||
use litellm_host::{
|
||||
call::{CallOutput, OutputOf},
|
||||
machine::ChannelInterceptors,
|
||||
};
|
||||
use litellm_inference::{CallOptions, RouteError};
|
||||
use litellm_inference_chat::{
|
||||
ChatCompletionsRoute, route::ChatCompletions, types::ChatCompletionsCall,
|
||||
};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::cache::Cached;
|
||||
|
||||
use super::{
|
||||
NativeCall,
|
||||
inference::{InferenceHost, InferenceRoute, run_inference},
|
||||
|
|
@ -61,11 +69,14 @@ impl InferenceRoute for ChatCompletionsRoute {
|
|||
self.with_cache(cache)
|
||||
}
|
||||
|
||||
fn machine(
|
||||
async fn call(
|
||||
self,
|
||||
call: <ChatCompletions as Protocol>::Request,
|
||||
policy: CachePolicy,
|
||||
) -> HostedMachine<ChatCompletions> {
|
||||
self.machine(call, policy)
|
||||
request: ChatCompletionsCall,
|
||||
interceptors: ChannelInterceptors<Cached<ChatCompletions>>,
|
||||
options: CallOptions,
|
||||
) -> Result<OutputOf<ChatCompletions>, RouteError> {
|
||||
self.execute_call(request, &interceptors, options)
|
||||
.await
|
||||
.map(CallOutput::Complete)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,13 +1,17 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
|
||||
use litellm_cache_response::ScopedCache;
|
||||
use litellm_callbacks_legacy_python::LoggingOperation;
|
||||
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
|
||||
use litellm_host::{call::HostedMachine, protocol::Protocol};
|
||||
use litellm_host_python::{PythonBinding, PythonHostCalls, from_py};
|
||||
use litellm_host::{
|
||||
call::{OutputOf, hosted_call},
|
||||
machine::ChannelInterceptors,
|
||||
protocol::Protocol,
|
||||
};
|
||||
use litellm_host_python::{PythonBinding, from_py};
|
||||
use litellm_http::transport::Error as TransportError;
|
||||
use litellm_inference::RouteError;
|
||||
use litellm_inference::{CallOptions, RouteError, caching::Cachable};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use pyo3::{
|
||||
exceptions::PyValueError,
|
||||
|
|
@ -20,6 +24,7 @@ use serde_json::{Map, Value};
|
|||
|
||||
use super::{NativeCall, RequestView};
|
||||
use crate::{
|
||||
cache::{Cached, CachedBinding},
|
||||
errors::{RustUpstreamError, route_error_to_pyerr},
|
||||
marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds},
|
||||
};
|
||||
|
|
@ -143,8 +148,8 @@ impl InferenceHost {
|
|||
}
|
||||
}
|
||||
|
||||
pub(super) trait InferenceRoute: Sized + 'static {
|
||||
type Protocol: Protocol<Error = RouteError>;
|
||||
pub(super) trait InferenceRoute: Sized + Send + 'static {
|
||||
type Protocol: Cachable;
|
||||
const OPERATION: LoggingOperation;
|
||||
const SYNC_CALL_TYPE: &'static str;
|
||||
const ASYNC_CALL_TYPE: &'static str;
|
||||
|
|
@ -155,11 +160,12 @@ pub(super) trait InferenceRoute: Sized + 'static {
|
|||
secrets: Arc<dyn SecretSource>,
|
||||
) -> Self;
|
||||
fn with_cache(self, cache: ScopedCache) -> Self;
|
||||
fn machine(
|
||||
fn call(
|
||||
self,
|
||||
call: <Self::Protocol as Protocol>::Request,
|
||||
policy: CachePolicy,
|
||||
) -> HostedMachine<Self::Protocol>;
|
||||
request: <Self::Protocol as Protocol>::Request,
|
||||
interceptors: ChannelInterceptors<Cached<Self::Protocol>>,
|
||||
options: CallOptions,
|
||||
) -> impl Future<Output = Result<OutputOf<Self::Protocol>, RouteError>> + Send;
|
||||
}
|
||||
|
||||
pub(super) fn run_inference<R, H>(
|
||||
|
|
@ -170,13 +176,14 @@ pub(super) fn run_inference<R, H>(
|
|||
) -> PyResult<Py<PyAny>>
|
||||
where
|
||||
R: InferenceRoute,
|
||||
H: PythonBinding<Protocol = R::Protocol> + PythonHostCalls<R::Protocol> + 'static,
|
||||
H: PythonBinding<Protocol = R::Protocol> + 'static,
|
||||
{
|
||||
let call_type = if asynchronous {
|
||||
R::ASYNC_CALL_TYPE
|
||||
} else {
|
||||
R::SYNC_CALL_TYPE
|
||||
};
|
||||
let positional = call.positional()?.unbind();
|
||||
let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?;
|
||||
super::run_public_call(
|
||||
py,
|
||||
|
|
@ -188,14 +195,28 @@ where
|
|||
crate::http::resources().auth.clone(),
|
||||
crate::secrets::source(py)?,
|
||||
);
|
||||
let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?;
|
||||
let route = match cache {
|
||||
Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)),
|
||||
None => route,
|
||||
};
|
||||
Ok(route.machine(request, cache_options.policy))
|
||||
Ok(hosted_call(
|
||||
request,
|
||||
None,
|
||||
move |(call, policy), services, interceptors, observers| async move {
|
||||
let route = match policy {
|
||||
Some(_) => route.with_cache(crate::cache::python_cache(services)),
|
||||
None => route,
|
||||
};
|
||||
route
|
||||
.call(
|
||||
call,
|
||||
interceptors,
|
||||
CallOptions {
|
||||
cache: policy,
|
||||
observers,
|
||||
},
|
||||
)
|
||||
.await
|
||||
},
|
||||
))
|
||||
},
|
||||
host,
|
||||
CachedBinding::new(host, call_type, asynchronous, positional),
|
||||
hooks,
|
||||
asynchronous,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
use crate::cache::{CacheCall, Cached, PythonCache, Selection};
|
||||
use litellm_host_python::{PythonHostCalls, PythonOwned};
|
||||
use litellm_host_python::PythonOwned;
|
||||
|
||||
use bytes::Bytes;
|
||||
use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py};
|
||||
|
|
@ -95,15 +94,11 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult<PyErr> {
|
|||
/// public response, chunks and exceptions.
|
||||
pub(super) struct MessagesPythonHost {
|
||||
view: RequestView,
|
||||
cache: PythonCache,
|
||||
}
|
||||
|
||||
impl MessagesPythonHost {
|
||||
pub(super) fn new(view: RequestView, asynchronous: bool) -> Self {
|
||||
Self {
|
||||
view,
|
||||
cache: PythonCache::new(asynchronous),
|
||||
}
|
||||
pub(super) fn new(view: RequestView) -> Self {
|
||||
Self { view }
|
||||
}
|
||||
|
||||
fn projection(
|
||||
|
|
@ -232,24 +227,20 @@ impl MessagesPythonHost {
|
|||
}
|
||||
|
||||
impl PythonBinding for MessagesPythonHost {
|
||||
type Protocol = Cached<Messages>;
|
||||
type Protocol = Messages;
|
||||
type Failure = PyErr;
|
||||
|
||||
fn decode_request(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
arguments: &Bound<'_, PyDict>,
|
||||
) -> Result<(MessagesCall, Selection), InvokeError<Error>> {
|
||||
) -> Result<MessagesCall, InvokeError<Error>> {
|
||||
self.view
|
||||
.resolve(py, arguments)
|
||||
.map_err(InvokeError::Python)?;
|
||||
let selection =
|
||||
crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages")
|
||||
.map_err(InvokeError::Python)?;
|
||||
self.projection(py, arguments)
|
||||
.map_err(|error| InvokeError::Python(self.map_failure(py, error)))?
|
||||
.map_err(InvokeError::Native)
|
||||
.map(|request| (request, selection))
|
||||
}
|
||||
|
||||
fn encode_response(
|
||||
|
|
@ -260,6 +251,22 @@ impl PythonBinding for MessagesPythonHost {
|
|||
public_response(py, ROUTE_HOST_MODULE, response.as_ref())
|
||||
}
|
||||
|
||||
fn public_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: &serde_json::Value,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
public_response(py, ROUTE_HOST_MODULE, response)
|
||||
}
|
||||
|
||||
fn response_value(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
public: &Bound<'_, PyAny>,
|
||||
) -> PyResult<serde_json::Value> {
|
||||
crate::marshal::response_value(public)
|
||||
}
|
||||
|
||||
fn encode_stream_head(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
|
|
@ -289,43 +296,12 @@ impl PythonBinding for MessagesPythonHost {
|
|||
}
|
||||
}
|
||||
|
||||
impl PythonHostCalls<Cached<Messages>> for MessagesPythonHost {
|
||||
fn handle_host_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
op: CacheCall,
|
||||
) -> Result<(), InvokeError<Error>> {
|
||||
self.cache
|
||||
.begin(py, op)
|
||||
.map(|_| ())
|
||||
.map_err(InvokeError::Python)
|
||||
}
|
||||
|
||||
fn begin_host_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
op: CacheCall,
|
||||
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
|
||||
self.cache.begin(py, op).map_err(InvokeError::Python)
|
||||
}
|
||||
|
||||
fn resume_host_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
result: PyResult<Py<PyAny>>,
|
||||
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
|
||||
self.cache.resume(py, result).map_err(InvokeError::Python)
|
||||
}
|
||||
}
|
||||
|
||||
impl PythonOwned for MessagesPythonHost {
|
||||
fn close(&mut self, _: Python<'_>) {
|
||||
self.view.close();
|
||||
self.cache.close();
|
||||
}
|
||||
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
self.view.traverse(visit)?;
|
||||
self.cache.traverse(visit)
|
||||
self.view.traverse(visit)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,53 +1,26 @@
|
|||
mod host;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use host::MessagesPythonHost;
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_cache_response::ScopedCache;
|
||||
use litellm_callbacks_legacy_python::LoggingOperation;
|
||||
use litellm_host::{call::OutputOf, machine::ChannelInterceptors};
|
||||
use litellm_inference::{CallOptions, RouteError};
|
||||
use litellm_inference_messages::{MessagesCall, MessagesRoute, route::Messages};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use super::NativeCall;
|
||||
use super::{
|
||||
NativeCall,
|
||||
inference::{InferenceRoute, run_inference},
|
||||
};
|
||||
use crate::cache::Cached;
|
||||
|
||||
fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
|
||||
let (arguments, hooks) =
|
||||
crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?;
|
||||
crate::routes::run_public_call(
|
||||
py,
|
||||
arguments,
|
||||
move |py, arguments, request| {
|
||||
let route = litellm_inference_messages::MessagesRoute::new(
|
||||
crate::http::provider_client(py, arguments, asynchronous)?
|
||||
.map_err(crate::http::client_error)?,
|
||||
crate::http::resources().auth.clone(),
|
||||
crate::secrets::source(py)?,
|
||||
);
|
||||
Ok(litellm_host::call::hosted_call(
|
||||
request,
|
||||
None,
|
||||
move |(call, selection): (_, crate::cache::Selection),
|
||||
services,
|
||||
interceptors,
|
||||
observers| async move {
|
||||
let (cache, options) = selection.attach(services);
|
||||
let route = match cache {
|
||||
Some(cache) => route.with_cache(cache),
|
||||
None => route,
|
||||
};
|
||||
route
|
||||
.execute(
|
||||
call,
|
||||
&interceptors,
|
||||
litellm_inference::CallOptions {
|
||||
cache: Some(options.policy),
|
||||
observers,
|
||||
},
|
||||
)
|
||||
.await
|
||||
},
|
||||
))
|
||||
},
|
||||
MessagesPythonHost::new(call.view()?, asynchronous),
|
||||
hooks,
|
||||
asynchronous,
|
||||
)
|
||||
let host = MessagesPythonHost::new(call.view()?);
|
||||
run_inference::<MessagesRoute, _>(py, call, asynchronous, host)
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
|
|
@ -59,3 +32,31 @@ pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAn
|
|||
pub(crate) fn amessages(py: Python<'_>, call: NativeCall<'_>) -> PyResult<Py<PyAny>> {
|
||||
run_messages(py, call, true)
|
||||
}
|
||||
|
||||
impl InferenceRoute for MessagesRoute {
|
||||
type Protocol = Messages;
|
||||
const OPERATION: LoggingOperation = LoggingOperation::Messages;
|
||||
const SYNC_CALL_TYPE: &'static str = "anthropic_messages";
|
||||
const ASYNC_CALL_TYPE: &'static str = "anthropic_messages";
|
||||
|
||||
fn new(
|
||||
http: litellm_http::Client,
|
||||
auth: Arc<AuthServices>,
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
) -> Self {
|
||||
Self::new(http, auth, secrets)
|
||||
}
|
||||
|
||||
fn with_cache(self, cache: ScopedCache) -> Self {
|
||||
self.with_cache(cache)
|
||||
}
|
||||
|
||||
async fn call(
|
||||
self,
|
||||
request: MessagesCall,
|
||||
interceptors: ChannelInterceptors<Cached<Messages>>,
|
||||
options: CallOptions,
|
||||
) -> Result<OutputOf<Messages>, RouteError> {
|
||||
self.execute(request, &interceptors, options).await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -41,6 +41,15 @@ impl<'py> NativeCall<'py> {
|
|||
effective_py_args(&self.base, &self.kwargs)
|
||||
}
|
||||
|
||||
/// The positional arguments by name, which the legacy cache folds into its key input.
|
||||
fn positional(&self) -> PyResult<Bound<'py, PyDict>> {
|
||||
let named = PyDict::new(self.base.py());
|
||||
for (name, value) in self.base.iter().take(self.args.len()) {
|
||||
named.set_item(name, value)?;
|
||||
}
|
||||
Ok(named)
|
||||
}
|
||||
|
||||
/// The view a route host reads its arguments from.
|
||||
fn view(&self) -> PyResult<RequestView> {
|
||||
RequestView::new(&self.base, &self.kwargs)
|
||||
|
|
|
|||
|
|
@ -118,6 +118,22 @@ impl PythonBinding for OcrPythonHost {
|
|||
public_response(py, "litellm.rust_bridge.host.routes.ocr", &response)
|
||||
}
|
||||
|
||||
fn public_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: &serde_json::Value,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
public_response(py, "litellm.rust_bridge.host.routes.ocr", response)
|
||||
}
|
||||
|
||||
fn response_value(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
public: &Bound<'_, PyAny>,
|
||||
) -> PyResult<serde_json::Value> {
|
||||
crate::marshal::response_value(public)
|
||||
}
|
||||
|
||||
fn encode_stream_head(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
|
|
|
|||
|
|
@ -80,6 +80,22 @@ impl PythonBinding for ResponsesPythonHost {
|
|||
self.0.response(py, &response)
|
||||
}
|
||||
|
||||
fn public_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: &serde_json::Value,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
self.0.response(py, response)
|
||||
}
|
||||
|
||||
fn response_value(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
public: &Bound<'_, PyAny>,
|
||||
) -> PyResult<serde_json::Value> {
|
||||
crate::marshal::response_value(public)
|
||||
}
|
||||
|
||||
fn encode_stream_head(
|
||||
&mut self,
|
||||
_py: Python<'_>,
|
||||
|
|
|
|||
|
|
@ -5,10 +5,11 @@ use std::sync::Arc;
|
|||
|
||||
use host::ResponsesPythonHost;
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_cache_response::{CachePolicy, ScopedCache};
|
||||
use litellm_cache_response::ScopedCache;
|
||||
use litellm_callbacks_legacy_python::LoggingOperation;
|
||||
use litellm_host::{call::HostedMachine, protocol::Protocol};
|
||||
use litellm_inference_responses::{ResponsesRoute, route::Responses};
|
||||
use litellm_host::{call::OutputOf, machine::ChannelInterceptors};
|
||||
use litellm_inference::{CallOptions, RouteError};
|
||||
use litellm_inference_responses::{ResponsesRoute, route::Responses, types::ResponsesCall};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use pyo3::prelude::*;
|
||||
pub(crate) use websocket::ResponsesWebSocketConnection;
|
||||
|
|
@ -17,6 +18,7 @@ use super::{
|
|||
NativeCall,
|
||||
inference::{InferenceHost, InferenceRoute, run_inference},
|
||||
};
|
||||
use crate::cache::Cached;
|
||||
|
||||
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
|
||||
|
||||
|
|
@ -53,11 +55,12 @@ impl InferenceRoute for ResponsesRoute {
|
|||
self.with_cache(cache)
|
||||
}
|
||||
|
||||
fn machine(
|
||||
async fn call(
|
||||
self,
|
||||
call: <Responses as Protocol>::Request,
|
||||
policy: CachePolicy,
|
||||
) -> HostedMachine<Responses> {
|
||||
self.machine(call, policy)
|
||||
request: ResponsesCall,
|
||||
interceptors: ChannelInterceptors<Cached<Responses>>,
|
||||
options: CallOptions,
|
||||
) -> Result<OutputOf<Responses>, RouteError> {
|
||||
self.execute(request, &interceptors, options).await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2,10 +2,10 @@
|
|||
|
||||
Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name
|
||||
|
||||
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
|
||||
- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (whether a call may read or write the configured cache, and the lookup and store calls), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds)
|
||||
- A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here
|
||||
- Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers
|
||||
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green
|
||||
- The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/cache/python/binding.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green
|
||||
- Inputs and outputs use values and frozen snapshots. Retain caller-owned clients and caches as opaque references when identity is required; Rust never discovers them through module globals
|
||||
- Moving a question into Rust means replacing the function body with a native call, not adding a second resolver. The Python signature stays
|
||||
- `routes/<route>.py` holds the questions one route asks, under the same function names on every route (`response`, `map_failure`, `connection_defaults`), so Rust names only the module. A question two routes ask goes in a top-level module here instead
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
1
tests/test_litellm_rust/cache/__init__.py
vendored
1
tests/test_litellm_rust/cache/__init__.py
vendored
|
|
@ -1 +0,0 @@
|
|||
|
||||
19
tests/test_litellm_rust/cache/conftest.py
vendored
19
tests/test_litellm_rust/cache/conftest.py
vendored
|
|
@ -1,19 +0,0 @@
|
|||
import threading
|
||||
from collections.abc import Generator
|
||||
from typing import Final
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_url() -> Generator[str]:
|
||||
server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis")
|
||||
worker: Final = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
worker.start()
|
||||
try:
|
||||
yield f"redis://127.0.0.1:{server.server_address[1]}"
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
worker.join(timeout=5)
|
||||
121
tests/test_litellm_rust/cache/support.py
vendored
121
tests/test_litellm_rust/cache/support.py
vendored
|
|
@ -1,121 +0,0 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.rust_bridge import runtime
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
|
||||
from litellm.rust_bridge.chat_completions.entrypoints import NATIVE_ACOMPLETION
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES
|
||||
from litellm.types.utils import ModelResponse
|
||||
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
|
||||
from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
|
||||
|
||||
def payload(value: object) -> object:
|
||||
if isinstance(value, ModelResponse):
|
||||
return value.model_dump_json(exclude=MappingProxyType({"id": True, "created": True}))
|
||||
if isinstance(value, dict):
|
||||
fields: Final = TypeAdapter(dict[str, object]).validate_python(value)
|
||||
return {name: field for name, field in fields.items() if name != "_hidden_params"}
|
||||
return value.model_dump_json() if isinstance(value, BaseModel) else value
|
||||
|
||||
|
||||
def cache_key(response: object) -> object:
|
||||
hidden: Final = get_hidden_params_dict(response)
|
||||
headers: Final = TypeAdapter(dict[str, object]).validate_python(hidden.get("additional_headers", {}))
|
||||
return headers.get("x-litellm-cache-key")
|
||||
|
||||
|
||||
async def invoke(
|
||||
route: Literal["chat", "messages", "responses"],
|
||||
server: RecordingServer,
|
||||
options: Mapping[str, object],
|
||||
native: bool = True,
|
||||
) -> object:
|
||||
common: Final = {"api_key": "test-key", "api_base": server.base_url, **options}
|
||||
if route == "responses":
|
||||
server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
|
||||
arguments: Final = {"model": RESPONSES_MODEL, "input": "hello", **common}
|
||||
if not native:
|
||||
return await litellm.aresponses(**arguments)
|
||||
request: Final = NativeCall(
|
||||
args=(),
|
||||
kwargs=arguments,
|
||||
base={
|
||||
"model": RESPONSES_MODEL,
|
||||
"input": "hello",
|
||||
"stream": None,
|
||||
"api_key": "test-key",
|
||||
"api_base": server.base_url,
|
||||
"custom_llm_provider": "openai",
|
||||
"extra_headers": None,
|
||||
**arguments,
|
||||
},
|
||||
)
|
||||
return await runtime.arun(
|
||||
RouteContext(Route.RESPONSES),
|
||||
binding=NATIVE_ARESPONSES,
|
||||
native=lambda hook: hook(request),
|
||||
python=runtime.NO_PYTHON,
|
||||
rules=(RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED),),
|
||||
)
|
||||
server.default_response = (
|
||||
ResponseSpec(body=None, events=MESSAGES_EVENTS)
|
||||
if options.get("stream")
|
||||
else ResponseSpec(body=MESSAGES_RESPONSE)
|
||||
)
|
||||
parameters: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32, **common}
|
||||
if route == "chat":
|
||||
if not native:
|
||||
return await litellm.acompletion(**parameters)
|
||||
chat: Final = NativeCall(args=(), kwargs=parameters, base={})
|
||||
return await runtime.arun(
|
||||
RouteContext(Route.CHAT_COMPLETIONS),
|
||||
binding=NATIVE_ACOMPLETION,
|
||||
native=lambda hook: hook(chat),
|
||||
python=runtime.NO_PYTHON,
|
||||
rules=(RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),),
|
||||
)
|
||||
if not native:
|
||||
return await litellm.anthropic_messages(**parameters)
|
||||
messages: Final = NativeCall(
|
||||
args=(),
|
||||
kwargs=parameters,
|
||||
base={
|
||||
"model": MESSAGES_MODEL,
|
||||
"messages": list(MESSAGES),
|
||||
"max_tokens": 32,
|
||||
"stream": None,
|
||||
"api_key": "test-key",
|
||||
"api_base": server.base_url,
|
||||
"custom_llm_provider": "anthropic",
|
||||
**parameters,
|
||||
},
|
||||
)
|
||||
return await runtime.arun(
|
||||
RouteContext(Route.MESSAGES),
|
||||
binding=NATIVE_AMESSAGES,
|
||||
native=lambda hook: hook(messages),
|
||||
python=runtime.NO_PYTHON,
|
||||
rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),),
|
||||
)
|
||||
|
||||
|
||||
async def collect(stream: object) -> bytes:
|
||||
assert isinstance(stream, AsyncIterator)
|
||||
return b"".join([chunk_bytes(chunk) async for chunk in stream])
|
||||
|
||||
|
||||
def chunk_bytes(value: object) -> bytes:
|
||||
assert isinstance(value, bytes)
|
||||
return value
|
||||
534
tests/test_litellm_rust/cache/test_python_cache.py
vendored
534
tests/test_litellm_rust/cache/test_python_cache.py
vendored
|
|
@ -1,534 +0,0 @@
|
|||
import asyncio
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm import _v2
|
||||
from litellm.caching.caching import Cache, CacheMode
|
||||
from litellm.caching.caching_handler import (
|
||||
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
|
||||
)
|
||||
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
model_budget_spend_cache_key,
|
||||
)
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
|
||||
from litellm.proxy.utils import InternalUsageCache
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.rust_bridge import runtime
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from litellm.types.caching import CachingSupportedCallTypes
|
||||
from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload
|
||||
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
|
||||
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
|
||||
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.routing import native_route
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("route", "stream", "legacy"),
|
||||
(
|
||||
("chat", False, False),
|
||||
("messages", False, False),
|
||||
("responses", False, False),
|
||||
("messages", True, False),
|
||||
("messages", False, True),
|
||||
("messages", True, True),
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage(
|
||||
recording_server: RecordingServer,
|
||||
route: Literal["chat", "messages", "responses"],
|
||||
stream: bool,
|
||||
native: bool,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
legacy: bool,
|
||||
) -> None:
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
litellm.cache = Cache() if legacy else _v2.Cache.memory()
|
||||
counters: Final = litellm.DualCache()
|
||||
budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters)
|
||||
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
|
||||
InternalUsageCache(counters), model_group_resolver=lambda model: model
|
||||
)
|
||||
recorder: Final = RecordingLogger()
|
||||
key_hash: Final = "a" * 64
|
||||
metadata: Final = {
|
||||
"user_api_key": key_hash,
|
||||
"model_group": "cached-model",
|
||||
"user_api_key_model_max_budget": {"cached-model": {"max_budget": 1, "budget_duration": "1h"}},
|
||||
}
|
||||
options: Final = {
|
||||
"callbacks": [budget, limiter, recorder],
|
||||
"metadata": metadata,
|
||||
"stream": stream,
|
||||
}
|
||||
spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h")
|
||||
token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens")
|
||||
first: Final = await invoke(route, recording_server, options, native=native)
|
||||
if stream:
|
||||
await collect(first)
|
||||
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
|
||||
await drain_logging()
|
||||
first_events: Final = await recorder.wait_for_async("async_log_success_event")
|
||||
first_log: Final = TypeAdapter(dict[str, object]).validate_python(first_events[0].kwargs)
|
||||
first_payload: Final = TypeAdapter(dict[str, object]).validate_python(first_log["standard_logging_object"])
|
||||
expected_cost: Final = TypeAdapter(float).validate_python(first_log["response_cost"])
|
||||
usage: Final = RESPONSES_RESPONSE["usage"] if route == "responses" else MESSAGES_RESPONSE["usage"]
|
||||
expected_tokens: Final = usage["input_tokens"] + usage["output_tokens"]
|
||||
assert expected_cost > 0
|
||||
assert counters.get_cache(spend_key) == pytest.approx(expected_cost)
|
||||
assert counters.get_cache(token_key) == first_payload["total_tokens"] == expected_tokens
|
||||
|
||||
second: Final = await invoke(route, recording_server, options, native=native)
|
||||
if stream:
|
||||
await collect(second)
|
||||
await drain_logging()
|
||||
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
|
||||
cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs)
|
||||
cached_payload: Final = TypeAdapter(dict[str, object]).validate_python(cached_log["standard_logging_object"])
|
||||
assert len(recording_server.requests) == 1
|
||||
assert len(successes) == 2
|
||||
assert cached_log["cache_hit"] is True
|
||||
assert cached_log["response_cost"] == cached_payload["response_cost"] == 0
|
||||
assert cached_payload["cache_hit"] is True
|
||||
assert cached_payload["id"] != first_payload["id"]
|
||||
assert cached_payload["custom_llm_provider"] == first_payload["custom_llm_provider"]
|
||||
assert cached_payload["custom_llm_provider"] == ("openai" if route == "responses" else "anthropic"), {
|
||||
"miss_provider": first_log.get("custom_llm_provider"),
|
||||
"hit_provider": cached_log.get("custom_llm_provider"),
|
||||
}
|
||||
assert cached_payload["total_tokens"] == first_payload["total_tokens"]
|
||||
assert counters.get_cache(spend_key) == pytest.approx(expected_cost)
|
||||
assert counters.get_cache(token_key) == 2 * expected_tokens
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("native", (False, True), ids=("python", "rust"))
|
||||
@pytest.mark.parametrize("backend", ("disabled", "memory", "redis"))
|
||||
async def test_response_cache_backend_does_not_control_coordination(
|
||||
recording_server: RecordingServer,
|
||||
native: bool,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
backend: Literal["disabled", "memory", "redis"],
|
||||
redis_url: str,
|
||||
) -> None:
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
litellm.cache = (
|
||||
None
|
||||
if backend == "disabled"
|
||||
else _v2.Cache.memory()
|
||||
if backend == "memory"
|
||||
else _v2.Cache.redis(redis_url, namespace="independent-coordination")
|
||||
)
|
||||
recording_server.expected_requests = 2 if backend == "disabled" else 1
|
||||
counters: Final = litellm.DualCache()
|
||||
budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters)
|
||||
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
|
||||
InternalUsageCache(counters), model_group_resolver=lambda model: model
|
||||
)
|
||||
key_hash: Final = "b" * 64
|
||||
identity: Final = UserAPIKeyAuth(api_key=key_hash, rpm_limit=2, tpm_limit=1000, max_parallel_requests=1)
|
||||
spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h")
|
||||
request_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "requests")
|
||||
token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens")
|
||||
parallel_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "max_parallel_requests")
|
||||
expected_tokens: Final = MESSAGES_RESPONSE["usage"]["input_tokens"] + MESSAGES_RESPONSE["usage"]["output_tokens"]
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
async def request(call_id: str, successes: int) -> object:
|
||||
data: Final = {
|
||||
"model": MESSAGES_MODEL,
|
||||
"messages": list(MESSAGES),
|
||||
"litellm_call_id": call_id,
|
||||
"max_tokens": 32,
|
||||
"metadata": {
|
||||
"user_api_key": key_hash,
|
||||
"model_group": "cached-model",
|
||||
"user_api_key_model_max_budget": {"cached-model": {"max_budget": 1, "budget_duration": "1h"}},
|
||||
},
|
||||
}
|
||||
await limiter.async_pre_call_hook(identity, counters, data, "acompletion")
|
||||
assert len(TypeAdapter(dict[str, float]).validate_python(counters.get_cache(parallel_key))) == 1
|
||||
response: Final = await invoke(
|
||||
"chat", recording_server, {**data, "callbacks": [budget, limiter, recorder]}, native=native
|
||||
)
|
||||
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
|
||||
await recorder.wait_for_async("async_log_success_event", count=successes)
|
||||
return response
|
||||
|
||||
await asyncio.create_task(request("cache-miss", 1))
|
||||
assert counters.get_cache(parallel_key) == {}
|
||||
assert counters.get_cache(request_key) == 1
|
||||
assert counters.get_cache(token_key) == expected_tokens
|
||||
first_events: Final = await recorder.wait_for_async("async_log_success_event")
|
||||
first_cost: Final = TypeAdapter(float).validate_python(first_events[0].kwargs["response_cost"])
|
||||
assert first_cost > 0
|
||||
assert counters.get_cache(spend_key) == pytest.approx(first_cost)
|
||||
await asyncio.create_task(request("cache-hit", 2))
|
||||
expected_spend: Final = first_cost * recording_server.expected_requests
|
||||
assert counters.get_cache(spend_key) == pytest.approx(expected_spend)
|
||||
assert len(recording_server.requests) == recording_server.expected_requests
|
||||
assert counters.get_cache(parallel_key) == {}
|
||||
assert counters.get_cache(request_key) == 2
|
||||
assert counters.get_cache(token_key) == 2 * expected_tokens
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await asyncio.create_task(request("over-rpm-limit", 3))
|
||||
assert len(recording_server.requests) == recording_server.expected_requests
|
||||
assert counters.get_cache(parallel_key) == {}
|
||||
assert counters.get_cache(token_key) == 2 * expected_tokens
|
||||
|
||||
assert counters.get_cache(spend_key) == pytest.approx(expected_spend)
|
||||
if litellm.cache is not None:
|
||||
await litellm.cache.disconnect()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
|
||||
)
|
||||
async def test_cache_controls_and_backend_credential_key_semantics(
|
||||
recording_server: RecordingServer,
|
||||
route: Literal["chat", "messages", "responses"],
|
||||
legacy: bool,
|
||||
) -> None:
|
||||
recording_server.expected_requests = 3 if legacy else 4
|
||||
litellm.cache = Cache() if legacy else _v2.Cache.memory()
|
||||
await invoke(route, recording_server, {"cache": {"no-store": True}})
|
||||
await invoke(route, recording_server, {})
|
||||
await invoke(route, recording_server, {})
|
||||
assert len(recording_server.requests) == 2
|
||||
await invoke(route, recording_server, {"cache": {"no-cache": True}})
|
||||
await invoke(route, recording_server, {"api_key": "another-key"})
|
||||
assert len(recording_server.requests) == recording_server.expected_requests
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("legacy", (False, True))
|
||||
async def test_v2_messages_replays_a_completed_stream(recording_server: RecordingServer, legacy: bool) -> None:
|
||||
recording_server.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS)
|
||||
litellm.cache = Cache() if legacy else _v2.Cache.memory()
|
||||
recorder: Final = RecordingLogger()
|
||||
parameters: Final = {
|
||||
"model": MESSAGES_MODEL,
|
||||
"messages": list(MESSAGES),
|
||||
"max_tokens": 32,
|
||||
"api_key": "test-key",
|
||||
"api_base": recording_server.base_url,
|
||||
"stream": True,
|
||||
"callbacks": [recorder],
|
||||
}
|
||||
with native_route(Route.MESSAGES):
|
||||
first_stream: Final = await litellm.anthropic_messages(**parameters)
|
||||
assert cache_key(first_stream) is None
|
||||
first: Final = await collect(first_stream)
|
||||
await recorder.wait_for_async("async_log_success_event")
|
||||
second_stream: Final = await litellm.anthropic_messages(**parameters)
|
||||
assert isinstance(cache_key(second_stream), str)
|
||||
assert cache_key(second_stream) == get_hidden_params_dict(second_stream)["cache_key"]
|
||||
second: Final = await collect(second_stream)
|
||||
assert payload(first) == payload(second)
|
||||
assert first == b"".join(recording_server.default_response.payloads())
|
||||
assert len(recording_server.requests) == 1
|
||||
await drain_logging()
|
||||
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
|
||||
cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs)
|
||||
assert cached_log["cache_hit"] is True
|
||||
assert cached_log["response_cost"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ("memory", "redis"))
|
||||
async def test_rust_messages_uses_a_legacy_cache_without_python_inference(
|
||||
recording_server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
backend: Literal["memory", "redis"],
|
||||
redis_url: str,
|
||||
) -> None:
|
||||
from litellm.caching.caching import Cache
|
||||
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
litellm.cache = Cache() if backend == "memory" else Cache(type="redis", url=redis_url, namespace="rust-host")
|
||||
logger: Final = RecordingLogger()
|
||||
litellm.callbacks = [logger]
|
||||
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
|
||||
parameters: Final = {
|
||||
"model": MESSAGES_MODEL,
|
||||
"messages": list(MESSAGES),
|
||||
"max_tokens": 32,
|
||||
"api_key": "test-key",
|
||||
"api_base": recording_server.base_url,
|
||||
}
|
||||
first: Final = await invoke("messages", recording_server, parameters)
|
||||
second: Final = await invoke("messages", recording_server, parameters)
|
||||
assert cache_key(second)
|
||||
assert cache_key(first) is None
|
||||
assert payload(first) == payload(second)
|
||||
assert len(recording_server.requests) == 1
|
||||
|
||||
await logger.wait_for_async("async_log_success_event", count=2)
|
||||
assert logger.names.count("async_log_success_event") == 2
|
||||
assert "log_failure_event" not in logger.names
|
||||
assert "async_log_failure_event" not in logger.names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", ("chat", "messages", "responses"))
|
||||
@pytest.mark.parametrize("native", (False, True))
|
||||
@pytest.mark.parametrize("excluded", (None, [], ["embedding"]))
|
||||
async def test_v2_cache_honors_supported_call_types_for_reads_and_writes(
|
||||
recording_server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
route: Literal["chat", "messages", "responses"],
|
||||
native: bool,
|
||||
excluded: list[CachingSupportedCallTypes] | None,
|
||||
) -> None:
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
litellm.cache = _v2.Cache.memory()
|
||||
call_type: Final[CachingSupportedCallTypes] = (
|
||||
"acompletion" if route == "chat" else "anthropic_messages" if route == "messages" else "aresponses"
|
||||
)
|
||||
recording_server.expected_requests = 4
|
||||
litellm.cache.supported_call_types = excluded
|
||||
await invoke(route, recording_server, {}, native=native)
|
||||
await invoke(route, recording_server, {}, native=native)
|
||||
assert len(recording_server.requests) == 2
|
||||
litellm.cache.supported_call_types = [call_type]
|
||||
await invoke(route, recording_server, {}, native=native)
|
||||
assert len(recording_server.requests) == 3
|
||||
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
|
||||
await invoke(route, recording_server, {}, native=native)
|
||||
assert len(recording_server.requests) == 3
|
||||
litellm.cache.supported_call_types = excluded
|
||||
await invoke(route, recording_server, {}, native=native)
|
||||
assert len(recording_server.requests) == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("native", (False, True))
|
||||
@pytest.mark.parametrize(
|
||||
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
|
||||
)
|
||||
async def test_v2_default_off_requires_opt_in_even_for_existing_entries(
|
||||
recording_server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
route: Literal["chat", "messages", "responses"],
|
||||
native: bool,
|
||||
legacy: bool,
|
||||
) -> None:
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if native else "0")
|
||||
litellm.cache = Cache() if legacy else _v2.Cache.memory()
|
||||
litellm.cache.mode = CacheMode.default_off
|
||||
recording_server.expected_requests = 4
|
||||
await invoke(route, recording_server, {}, native=native)
|
||||
await invoke(route, recording_server, {}, native=native)
|
||||
assert len(recording_server.requests) == 2
|
||||
await invoke(route, recording_server, {"cache": {"use-cache": True}}, native=native)
|
||||
assert len(recording_server.requests) == 3
|
||||
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
|
||||
await invoke(route, recording_server, {"cache": {"use-cache": True}}, native=native)
|
||||
assert len(recording_server.requests) == 3
|
||||
await invoke(route, recording_server, {}, native=native)
|
||||
assert len(recording_server.requests) == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True))
|
||||
)
|
||||
async def test_cache_lookup_uses_backend_request_callback_semantics(
|
||||
recording_server: RecordingServer,
|
||||
route: Literal["chat", "messages", "responses"],
|
||||
legacy: bool,
|
||||
) -> None:
|
||||
class Rewrite(RecordingLogger):
|
||||
temperature = 0.1
|
||||
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
|
||||
request_body(kwargs)["temperature"] = self.temperature
|
||||
super().log_pre_api_call(model, messages, kwargs)
|
||||
|
||||
logger: Final = Rewrite()
|
||||
litellm.cache = Cache() if legacy else _v2.Cache.memory()
|
||||
recording_server.expected_requests = 1 if legacy else 2
|
||||
await invoke(route, recording_server, {"callbacks": [logger]})
|
||||
first_hit: Final = await invoke(route, recording_server, {"callbacks": [logger]})
|
||||
logger.temperature = 0.8
|
||||
await invoke(route, recording_server, {"callbacks": [logger]})
|
||||
second_hit: Final = await invoke(route, recording_server, {"callbacks": [logger]})
|
||||
assert logger.names.count("log_pre_api_call") == 4
|
||||
assert len(recording_server.requests) == recording_server.expected_requests
|
||||
assert recording_server.requests[0].body["temperature"] == 0.1
|
||||
if not legacy:
|
||||
assert recording_server.requests[1].body["temperature"] == 0.8
|
||||
assert isinstance(cache_key(first_hit), str)
|
||||
assert isinstance(cache_key(second_hit), str)
|
||||
if not legacy:
|
||||
assert cache_key(first_hit) != cache_key(second_hit)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("cancel_lookup", (False, True))
|
||||
async def test_python_cache_operations_stay_in_the_rust_callers_task(
|
||||
recording_server: RecordingServer,
|
||||
cancel_lookup: bool,
|
||||
) -> None:
|
||||
from litellm.caching.base_cache import BaseCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
|
||||
caller: Final = asyncio.current_task()
|
||||
entered: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
storage: Final = InMemoryCache()
|
||||
|
||||
class CallerCache(BaseCache):
|
||||
async def async_set_cache_pipeline(
|
||||
self, cache_list: list[tuple[str, object]], ttl: float | None = None
|
||||
) -> None:
|
||||
await storage.async_set_cache_pipeline(cache_list, ttl=ttl)
|
||||
|
||||
async def async_get_cache(self, key: str, **kwargs: object) -> object:
|
||||
if cancel_lookup:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
else:
|
||||
assert asyncio.current_task() is caller
|
||||
return storage.get_cache(key, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
assert asyncio.current_task() is caller
|
||||
await asyncio.sleep(0)
|
||||
storage.set_cache(key, value, **kwargs)
|
||||
|
||||
litellm.cache = Cache(_backend=CallerCache())
|
||||
if cancel_lookup:
|
||||
recording_server.expected_requests = 0
|
||||
task: Final = asyncio.create_task(invoke("messages", recording_server, {}))
|
||||
await asyncio.wait_for(entered.wait(), timeout=5)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
release.set()
|
||||
await asyncio.sleep(0)
|
||||
assert len(recording_server.requests) == 0
|
||||
assert storage.cache_dict == {}
|
||||
return
|
||||
first: Final = await invoke("messages", recording_server, {})
|
||||
second: Final = await invoke("messages", recording_server, {})
|
||||
assert payload(first) == payload(second)
|
||||
assert cache_key(second)
|
||||
assert len(recording_server.requests) == 1
|
||||
|
||||
|
||||
def test_sync_rust_messages_calls_python_cache(recording_server: RecordingServer) -> None:
|
||||
from litellm.rust_bridge.messages.entrypoints import NATIVE_MESSAGES
|
||||
|
||||
litellm.cache = Cache()
|
||||
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
|
||||
arguments: Final = {
|
||||
"model": MESSAGES_MODEL,
|
||||
"messages": list(MESSAGES),
|
||||
"max_tokens": 32,
|
||||
"api_key": "test-key",
|
||||
"api_base": recording_server.base_url,
|
||||
}
|
||||
request: Final = NativeCall(
|
||||
args=(),
|
||||
kwargs=arguments,
|
||||
base={
|
||||
"model": MESSAGES_MODEL,
|
||||
"messages": list(MESSAGES),
|
||||
"max_tokens": 32,
|
||||
"stream": None,
|
||||
"api_key": "test-key",
|
||||
"api_base": recording_server.base_url,
|
||||
"custom_llm_provider": "anthropic",
|
||||
**arguments,
|
||||
},
|
||||
)
|
||||
|
||||
def call() -> object:
|
||||
return runtime.run(
|
||||
RouteContext(Route.MESSAGES),
|
||||
binding=NATIVE_MESSAGES,
|
||||
native=lambda hook: hook(request),
|
||||
python=runtime.NO_PYTHON,
|
||||
rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),),
|
||||
)
|
||||
|
||||
first: Final = call()
|
||||
second: Final = call()
|
||||
assert payload(first) == payload(second)
|
||||
assert cache_key(second)
|
||||
assert len(recording_server.requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("namespace_source", ("cache", "metadata"))
|
||||
async def test_rust_messages_legacy_cache_honors_request_namespaces(
|
||||
recording_server: RecordingServer, namespace_source: str
|
||||
) -> None:
|
||||
litellm.cache = Cache()
|
||||
recording_server.expected_requests = 2
|
||||
first_options: Final = (
|
||||
{"cache": {"namespace": "first"}} if namespace_source == "cache" else {"metadata": {"redis_namespace": "first"}}
|
||||
)
|
||||
second_options: Final = (
|
||||
{"cache": {"namespace": "second"}}
|
||||
if namespace_source == "cache"
|
||||
else {"metadata": {"redis_namespace": "second"}}
|
||||
)
|
||||
await invoke("messages", recording_server, first_options)
|
||||
second: Final = await invoke("messages", recording_server, second_options)
|
||||
first_hit: Final = await invoke("messages", recording_server, first_options)
|
||||
second_hit: Final = await invoke("messages", recording_server, second_options)
|
||||
assert cache_key(second) is None
|
||||
assert cache_key(first_hit) is not None
|
||||
assert cache_key(second_hit) is not None
|
||||
assert len(recording_server.requests) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rust_messages_legacy_semantic_cache_preserves_python_scope(
|
||||
recording_server: RecordingServer,
|
||||
) -> None:
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.types.caching import LiteLLMCacheType
|
||||
|
||||
litellm.cache = Cache(type=LiteLLMCacheType.REDIS_SEMANTIC, _backend=InMemoryCache())
|
||||
first_options: Final = {"messages": [{"role": "user", "content": "hello"}]}
|
||||
second_options: Final = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
first: Final = await invoke("messages", recording_server, first_options)
|
||||
second: Final = await invoke("messages", recording_server, second_options)
|
||||
assert cache_key(first) is None
|
||||
assert cache_key(second) is not None
|
||||
assert payload(first) == payload(second)
|
||||
assert len(recording_server.requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("rust_first", (False, True), ids=("python_to_rust", "rust_to_python"))
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
|
||||
async def test_legacy_cache_keeps_public_messages_responses_compatible(
|
||||
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, rust_first: bool, stream: bool
|
||||
) -> None:
|
||||
litellm.cache = Cache()
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
options: Final = {"litellm_params": {"preset_cache_key": "shared-messages"}, "stream": stream}
|
||||
first: Final = await invoke("messages", recording_server, options, native=rust_first)
|
||||
first_payload: Final = await collect(first) if stream else payload(first)
|
||||
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
|
||||
second: Final = await invoke("messages", recording_server, options, native=not rust_first)
|
||||
second_payload: Final = await collect(second) if stream else payload(second)
|
||||
assert second_payload == first_payload
|
||||
assert len(recording_server.requests) == 1
|
||||
0
tests/test_litellm_rust/cache/v2/__init__.py
vendored
0
tests/test_litellm_rust/cache/v2/__init__.py
vendored
207
tests/test_litellm_rust/cache/v2/test_cache.py
vendored
207
tests/test_litellm_rust/cache/v2/test_cache.py
vendored
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
328
tests/test_litellm_rust/chat_completions/test_cache.py
Normal file
328
tests/test_litellm_rust/chat_completions/test_cache.py
Normal file
|
|
@ -0,0 +1,328 @@
|
|||
import asyncio
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.caching.caching import Cache, CacheMode
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.types.caching import CachingSupportedCallTypes
|
||||
from litellm.types.utils import ModelResponse
|
||||
from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE, MESSAGES_TEXT
|
||||
from tests.test_litellm_rust.support.caching import (
|
||||
Call,
|
||||
CallerCache,
|
||||
FailingLookupCache,
|
||||
FakeRedis,
|
||||
HeldWriteCache,
|
||||
assert_cache_hit_keeps_budget_spend_and_counts_usage,
|
||||
settle_cache_writes,
|
||||
use_cache,
|
||||
)
|
||||
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer
|
||||
from tests.test_litellm_rust.support.results import cache_key, stored_cache_key
|
||||
from tests.test_litellm_rust.support.routing import assert_served_natively
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
EMPTY: Final = {"object": "chat.completion", "choices": []}
|
||||
OWN_CALL_TYPES: Final = ("acompletion",)
|
||||
|
||||
|
||||
def content(response: object) -> object:
|
||||
assert isinstance(response, ModelResponse)
|
||||
return response.model_dump()["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_repeated_native_call_is_answered_from_the_python_cache(
|
||||
call: Call, server: RecordingServer, asynchronous: bool
|
||||
) -> None:
|
||||
first: Final = await call(asynchronous=asynchronous)
|
||||
second: Final = await call(asynchronous=asynchronous)
|
||||
|
||||
assert_served_natively(server, first)
|
||||
assert cache_key(first) is None
|
||||
assert isinstance(cache_key(second), str)
|
||||
assert content(second) == content(first) == MESSAGES_TEXT
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("first", "second"),
|
||||
(
|
||||
({"cache": {"no-store": True}}, {}),
|
||||
({}, {"cache": {"no-cache": True}}),
|
||||
({}, {"caching": False}),
|
||||
({}, {"messages": [{"role": "user", "content": "Goodbye"}]}),
|
||||
({}, {"temperature": 0.5}),
|
||||
({"cache": {"namespace": "first"}}, {"cache": {"namespace": "second"}}),
|
||||
({"metadata": {"redis_namespace": "first"}}, {"metadata": {"redis_namespace": "second"}}),
|
||||
),
|
||||
ids=(
|
||||
"no-store",
|
||||
"no-cache",
|
||||
"caching-off-read",
|
||||
"other-messages",
|
||||
"other-param",
|
||||
"cache-namespace",
|
||||
"metadata-namespace",
|
||||
),
|
||||
)
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_cache_controls_and_keys_keep_the_second_call_on_the_provider(
|
||||
call: Call, server: RecordingServer, first: Mapping[str, object], second: Mapping[str, object]
|
||||
) -> None:
|
||||
server.expect(2)
|
||||
|
||||
await call(**first)
|
||||
response: Final = await call(**second)
|
||||
|
||||
assert_served_natively(server, count=2)
|
||||
assert cache_key(response) is None
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_caching_off_skips_the_read_but_still_writes(call: Call, server: RecordingServer) -> None:
|
||||
await call(caching=False)
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert content(hit) == MESSAGES_TEXT
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_each_namespace_keeps_its_own_entry(call: Call, server: RecordingServer) -> None:
|
||||
server.expect(2)
|
||||
|
||||
await call(cache={"namespace": "first"})
|
||||
await call(cache={"namespace": "second"})
|
||||
hit: Final = await call(cache={"namespace": "first"})
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_the_key_is_built_from_the_call_arguments_not_the_wire_body(call: Call, server: RecordingServer) -> None:
|
||||
class Rewrite(RecordingLogger):
|
||||
temperature = 0.1
|
||||
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
|
||||
request_body(kwargs)["temperature"] = self.temperature
|
||||
super().log_pre_api_call(model, messages, kwargs)
|
||||
|
||||
logger: Final = Rewrite()
|
||||
|
||||
await call(callbacks=[logger])
|
||||
logger.temperature = 0.8
|
||||
hit: Final = await call(callbacks=[logger])
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.1
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("supported", "expected_requests"), ((["embedding"], 2), (OWN_CALL_TYPES, 1)), ids=("other-type", "own-type")
|
||||
)
|
||||
async def test_supported_call_types_gate_native_reads_and_writes(
|
||||
call: Call,
|
||||
server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
supported: Sequence[CachingSupportedCallTypes],
|
||||
expected_requests: int,
|
||||
) -> None:
|
||||
use_cache(monkeypatch, Cache(supported_call_types=list(supported)))
|
||||
server.expect(expected_requests)
|
||||
|
||||
await call()
|
||||
second: Final = await call()
|
||||
|
||||
assert len(server.requests) == expected_requests
|
||||
assert isinstance(cache_key(second), str) == (expected_requests == 1)
|
||||
assert_served_natively(server, count=expected_requests)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("options", "expected_requests"), (({}, 2), ({"cache": {"use-cache": True}}, 1)), ids=("plain", "opt-in")
|
||||
)
|
||||
async def test_default_off_mode_caches_only_calls_that_opt_in(
|
||||
call: Call,
|
||||
server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
options: Mapping[str, object],
|
||||
expected_requests: int,
|
||||
) -> None:
|
||||
use_cache(monkeypatch, Cache(mode=CacheMode.default_off))
|
||||
server.expect(expected_requests)
|
||||
|
||||
await call(**options)
|
||||
second: Final = await call(**options)
|
||||
|
||||
assert len(server.requests) == expected_requests
|
||||
assert isinstance(cache_key(second), str) == (expected_requests == 1)
|
||||
assert_served_natively(server, count=expected_requests)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_native_hit_is_logged_as_a_free_cache_hit(call: Call, server: RecordingServer) -> None:
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await call(callbacks=[recorder])
|
||||
await call(callbacks=[recorder])
|
||||
await drain_logging()
|
||||
|
||||
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
|
||||
miss, hit = (_OBJECT.validate_python(event.kwargs) for event in successes)
|
||||
assert len(server.requests) == 1
|
||||
assert miss["cache_hit"] is not True
|
||||
assert TypeAdapter(float).validate_python(miss["response_cost"]) > 0
|
||||
assert hit["cache_hit"] is True
|
||||
assert hit["response_cost"] == 0
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_python_cache_lookups_run_in_the_callers_task(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = CallerCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
await call()
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert backend.lookup_tasks == [asyncio.current_task(), asyncio.current_task()]
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_cancelling_a_pending_lookup_sends_and_stores_nothing(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
entered: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
backend: Final = CallerCache(entered, release)
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
server.expect(0)
|
||||
|
||||
task: Final = asyncio.create_task(call())
|
||||
await asyncio.wait_for(entered.wait(), timeout=5)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
release.set()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert server.requests == []
|
||||
assert backend.storage.cache_dict == {}
|
||||
assert_served_natively(server, count=0)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_cache_hit_keeps_budget_spend_and_counts_usage(call: Call, server: RecordingServer) -> None:
|
||||
usage: Final = TypeAdapter(dict[str, int]).validate_python(MESSAGES_RESPONSE["usage"])
|
||||
|
||||
await assert_cache_hit_keeps_budget_spend_and_counts_usage(
|
||||
call, server, usage["input_tokens"] + usage["output_tokens"]
|
||||
)
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_stored_response_without_output_is_a_miss(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = InMemoryCache()
|
||||
cache: Final = use_cache(monkeypatch, Cache(_backend=backend))
|
||||
await call()
|
||||
(key,) = backend.cache_dict
|
||||
cache.add_cache(EMPTY, cache_key=key)
|
||||
server.expect(2)
|
||||
|
||||
response: Final = await call()
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert content(response) == MESSAGES_TEXT
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_a_failing_cache_lookup_runs_the_call_uncached(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
use_cache(monkeypatch, FailingLookupCache())
|
||||
server.expect(2)
|
||||
|
||||
first: Final = await call()
|
||||
await call()
|
||||
|
||||
assert content(first) == MESSAGES_TEXT
|
||||
assert len(server.requests) == 2
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_the_response_returns_before_the_cache_write_lands(
|
||||
call: Call, native_call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = HeldWriteCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
first: Final = await asyncio.wait_for(native_call(), timeout=5)
|
||||
backend.release.set()
|
||||
await settle_cache_writes()
|
||||
hit: Final = await call()
|
||||
|
||||
assert content(first) == MESSAGES_TEXT
|
||||
assert len(server.requests) == 1
|
||||
assert content(hit) == MESSAGES_TEXT
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_redis_hit_fills_the_worker_memory_tier(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
|
||||
) -> None:
|
||||
use_cache(monkeypatch, fake_redis.cache())
|
||||
|
||||
await call()
|
||||
await call()
|
||||
fake_redis.client.flushall()
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert content(hit) == MESSAGES_TEXT
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_stored_response_lands_in_the_worker_memory_tier(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
|
||||
) -> None:
|
||||
use_cache(monkeypatch, fake_redis.cache())
|
||||
|
||||
first: Final = await call()
|
||||
fake_redis.client.flushall()
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert content(hit) == content(first)
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_hit_reports_the_key_the_entry_is_stored_under(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = InMemoryCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
await call()
|
||||
hit: Final = await call()
|
||||
|
||||
(stored,) = backend.cache_dict
|
||||
assert len(server.requests) == 1
|
||||
assert stored_cache_key(hit) == stored
|
||||
assert_served_natively(server, count=1)
|
||||
|
|
@ -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}),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
497
tests/test_litellm_rust/messages/test_cache.py
Normal file
497
tests/test_litellm_rust/messages/test_cache.py
Normal file
|
|
@ -0,0 +1,497 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache, CacheMode
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.caching import CachingSupportedCallTypes, LiteLLMCacheType
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
|
||||
from litellm.types.utils import LLMResponseTypes
|
||||
from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
|
||||
from tests.test_litellm_rust.support.caching import (
|
||||
Call,
|
||||
CallerCache,
|
||||
FailingLookupCache,
|
||||
FakeRedis,
|
||||
HeldWriteCache,
|
||||
assert_cache_hit_keeps_budget_spend_and_counts_usage,
|
||||
settle_cache_writes,
|
||||
use_cache,
|
||||
)
|
||||
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
|
||||
from tests.test_litellm_rust.support.isolation import rebound
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.results import cache_key, consume, response_id, stored_cache_key
|
||||
from tests.test_litellm_rust.support.routing import assert_served_natively, routed
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
_MESSAGE: Final = TypeAdapter(AnthropicMessagesResponse)
|
||||
_CHUNKS: Final = TypeAdapter(list[bytes])
|
||||
STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS)
|
||||
EMPTY: Final = {**MESSAGES_RESPONSE, "content": []}
|
||||
OWN_CALL_TYPES: Final = ("anthropic_messages",)
|
||||
|
||||
|
||||
def body(response: object) -> Mapping[str, object]:
|
||||
return {name: value for name, value in _OBJECT.validate_python(response).items() if name != "_hidden_params"}
|
||||
|
||||
|
||||
async def collect(stream: object) -> bytes:
|
||||
streamed_async: Final = isinstance(stream, AsyncIterator)
|
||||
collected: Final = _CHUNKS.validate_python(list(await consume(stream)), strict=True)
|
||||
assert streamed_async
|
||||
await settle_cache_writes()
|
||||
return b"".join(collected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_repeated_native_call_is_answered_from_the_python_cache(
|
||||
call: Call, server: RecordingServer, asynchronous: bool
|
||||
) -> None:
|
||||
first: Final = await call(asynchronous=asynchronous)
|
||||
second: Final = await call(asynchronous=asynchronous)
|
||||
|
||||
assert_served_natively(server, first)
|
||||
assert cache_key(first) is None
|
||||
assert isinstance(cache_key(second), str)
|
||||
assert body(second) == body(first) == MESSAGES_RESPONSE
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_completed_native_stream_is_replayed_from_the_python_cache(call: Call, server: RecordingServer) -> None:
|
||||
server.serve(STREAM)
|
||||
|
||||
first_stream: Final = await call(stream=True)
|
||||
assert cache_key(first_stream) is None
|
||||
first: Final = await collect(first_stream)
|
||||
second_stream: Final = await call(stream=True)
|
||||
assert isinstance(cache_key(second_stream), str)
|
||||
second: Final = await collect(second_stream)
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert first == b"".join(STREAM.payloads())
|
||||
assert second == first
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_prompts_passed_positionally_get_their_own_entries(server: RecordingServer) -> None:
|
||||
from litellm.messages import anthropic_messages
|
||||
|
||||
server.expect(2)
|
||||
|
||||
async def ask(content: str) -> object:
|
||||
return await anthropic_messages(
|
||||
64,
|
||||
[{"role": "user", "content": content}],
|
||||
MESSAGES_MODEL,
|
||||
api_key="test-key",
|
||||
api_base=server.base_url,
|
||||
)
|
||||
|
||||
await ask("first prompt")
|
||||
second: Final = await ask("second prompt")
|
||||
await settle_cache_writes()
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert cache_key(second) is None
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("native_writes", (True, False), ids=("native-writes", "python-writes"))
|
||||
@pytest.mark.parametrize(
|
||||
"options",
|
||||
({}, {"metadata": None}, {"cache_key": "shared-entry"}),
|
||||
ids=("plain", "metadata-none", "preset-key"),
|
||||
)
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_native_and_python_replay_each_others_stream_entries(
|
||||
call: Call, server: RecordingServer, native_writes: bool, options: Mapping[str, object]
|
||||
) -> None:
|
||||
server.serve(STREAM)
|
||||
|
||||
async def stream(native: bool) -> bytes:
|
||||
if native:
|
||||
return await collect(await call(stream=True, **options))
|
||||
with routed():
|
||||
return await collect(await call(stream=True, **options))
|
||||
|
||||
first: Final = await stream(native_writes)
|
||||
second: Final = await stream(not native_writes)
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert second == first == b"".join(STREAM.payloads())
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_stream_closed_early_is_not_stored(call: Call, server: RecordingServer) -> None:
|
||||
server.serve(STREAM)
|
||||
server.expect(2)
|
||||
|
||||
stream: Final = await call(stream=True)
|
||||
assert isinstance(stream, AsyncIterator)
|
||||
await anext(stream)
|
||||
await stream.aclose()
|
||||
await collect(await call(stream=True))
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("first", "second"),
|
||||
(
|
||||
({"cache": {"no-store": True}}, {}),
|
||||
({}, {"cache": {"no-cache": True}}),
|
||||
({}, {"caching": False}),
|
||||
({}, {"messages": [{"role": "user", "content": "Goodbye"}]}),
|
||||
({}, {"system": "Be brief"}),
|
||||
({"cache": {"namespace": "first"}}, {"cache": {"namespace": "second"}}),
|
||||
({"metadata": {"redis_namespace": "first"}}, {"metadata": {"redis_namespace": "second"}}),
|
||||
),
|
||||
ids=(
|
||||
"no-store",
|
||||
"no-cache",
|
||||
"caching-off-read",
|
||||
"other-messages",
|
||||
"other-param",
|
||||
"cache-namespace",
|
||||
"metadata-namespace",
|
||||
),
|
||||
)
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_cache_controls_and_keys_keep_the_second_call_on_the_provider(
|
||||
call: Call, server: RecordingServer, first: Mapping[str, object], second: Mapping[str, object]
|
||||
) -> None:
|
||||
server.expect(2)
|
||||
|
||||
await call(**first)
|
||||
response: Final = await call(**second)
|
||||
|
||||
assert_served_natively(server, count=2)
|
||||
assert cache_key(response) is None
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_caching_off_skips_the_read_but_still_writes(call: Call, server: RecordingServer) -> None:
|
||||
await call(caching=False)
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert body(hit) == MESSAGES_RESPONSE
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_each_namespace_keeps_its_own_entry(call: Call, server: RecordingServer) -> None:
|
||||
server.expect(2)
|
||||
|
||||
await call(cache={"namespace": "first"})
|
||||
await call(cache={"namespace": "second"})
|
||||
hit: Final = await call(cache={"namespace": "first"})
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_the_key_is_built_from_the_call_arguments_not_the_wire_body(call: Call, server: RecordingServer) -> None:
|
||||
class Rewrite(RecordingLogger):
|
||||
temperature = 0.1
|
||||
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
|
||||
request_body(kwargs)["temperature"] = self.temperature
|
||||
super().log_pre_api_call(model, messages, kwargs)
|
||||
|
||||
logger: Final = Rewrite()
|
||||
|
||||
await call(callbacks=[logger])
|
||||
logger.temperature = 0.8
|
||||
hit: Final = await call(callbacks=[logger])
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.1
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("supported", "expected_requests"), ((["embedding"], 2), (OWN_CALL_TYPES, 1)), ids=("other-type", "own-type")
|
||||
)
|
||||
async def test_supported_call_types_gate_native_reads_and_writes(
|
||||
call: Call,
|
||||
server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
supported: Sequence[CachingSupportedCallTypes],
|
||||
expected_requests: int,
|
||||
) -> None:
|
||||
use_cache(monkeypatch, Cache(supported_call_types=list(supported)))
|
||||
server.expect(expected_requests)
|
||||
|
||||
await call()
|
||||
second: Final = await call()
|
||||
|
||||
assert len(server.requests) == expected_requests
|
||||
assert isinstance(cache_key(second), str) == (expected_requests == 1)
|
||||
assert_served_natively(server, count=expected_requests)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("options", "expected_requests"), (({}, 2), ({"cache": {"use-cache": True}}, 1)), ids=("plain", "opt-in")
|
||||
)
|
||||
async def test_default_off_mode_caches_only_calls_that_opt_in(
|
||||
call: Call,
|
||||
server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
options: Mapping[str, object],
|
||||
expected_requests: int,
|
||||
) -> None:
|
||||
use_cache(monkeypatch, Cache(mode=CacheMode.default_off))
|
||||
server.expect(expected_requests)
|
||||
|
||||
await call(**options)
|
||||
second: Final = await call(**options)
|
||||
|
||||
assert len(server.requests) == expected_requests
|
||||
assert isinstance(cache_key(second), str) == (expected_requests == 1)
|
||||
assert_served_natively(server, count=expected_requests)
|
||||
|
||||
|
||||
async def test_a_semantic_python_cache_answers_a_similar_prompt(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
use_cache(monkeypatch, Cache(type=LiteLLMCacheType.REDIS_SEMANTIC, _backend=InMemoryCache()))
|
||||
|
||||
first: Final = await call(messages=[{"role": "user", "content": "hello"}])
|
||||
second: Final = await call(messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert cache_key(first) is None
|
||||
assert isinstance(cache_key(second), str)
|
||||
assert body(second) == body(first)
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_native_hit_is_logged_as_a_free_cache_hit(call: Call, server: RecordingServer, stream: bool) -> None:
|
||||
if stream:
|
||||
server.serve(STREAM)
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
async def send() -> None:
|
||||
response: Final = await call(stream=stream, callbacks=[recorder])
|
||||
if stream:
|
||||
await collect(response)
|
||||
|
||||
await send()
|
||||
await send()
|
||||
await drain_logging()
|
||||
|
||||
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
|
||||
miss, hit = (_OBJECT.validate_python(event.kwargs) for event in successes)
|
||||
assert miss["cache_hit"] is not True
|
||||
assert TypeAdapter(float).validate_python(miss["response_cost"]) > 0
|
||||
assert hit["cache_hit"] is True
|
||||
assert hit["response_cost"] == 0
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_python_cache_lookups_run_in_the_callers_task(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = CallerCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
await call()
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert backend.lookup_tasks == [asyncio.current_task(), asyncio.current_task()]
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_cancelling_a_pending_lookup_sends_and_stores_nothing(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
entered: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
backend: Final = CallerCache(entered, release)
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
server.expect(0)
|
||||
|
||||
task: Final = asyncio.create_task(call())
|
||||
await asyncio.wait_for(entered.wait(), timeout=5)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
release.set()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert server.requests == []
|
||||
assert backend.storage.cache_dict == {}
|
||||
assert_served_natively(server, count=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_cache_hit_keeps_budget_spend_and_counts_usage(
|
||||
call: Call, server: RecordingServer, stream: bool
|
||||
) -> None:
|
||||
if stream:
|
||||
server.serve(STREAM)
|
||||
|
||||
async def send(**options: object) -> object:
|
||||
response: Final = await call(stream=stream, **options)
|
||||
if stream:
|
||||
await collect(response)
|
||||
return response
|
||||
|
||||
usage: Final = TypeAdapter(dict[str, int]).validate_python(MESSAGES_RESPONSE["usage"])
|
||||
await assert_cache_hit_keeps_budget_spend_and_counts_usage(
|
||||
send, server, usage["input_tokens"] + usage["output_tokens"]
|
||||
)
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_response_without_output_is_not_stored(call: Call, server: RecordingServer) -> None:
|
||||
server.serve(ResponseSpec(body=EMPTY))
|
||||
server.expect(2)
|
||||
|
||||
await call()
|
||||
await call()
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_a_stored_response_without_output_is_a_miss(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = InMemoryCache()
|
||||
cache: Final = use_cache(monkeypatch, Cache(_backend=backend))
|
||||
await call()
|
||||
(key,) = backend.cache_dict
|
||||
cache.add_cache(EMPTY, cache_key=key)
|
||||
server.expect(2)
|
||||
|
||||
response: Final = await call()
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert body(response) == MESSAGES_RESPONSE
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_a_failing_cache_lookup_runs_the_call_uncached(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
use_cache(monkeypatch, FailingLookupCache())
|
||||
server.expect(2)
|
||||
|
||||
first: Final = await call()
|
||||
await call()
|
||||
|
||||
assert body(first) == MESSAGES_RESPONSE
|
||||
assert len(server.requests) == 2
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_the_response_returns_before_the_cache_write_lands(
|
||||
call: Call, native_call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = HeldWriteCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
first: Final = await asyncio.wait_for(native_call(), timeout=5)
|
||||
backend.release.set()
|
||||
await settle_cache_writes()
|
||||
hit: Final = await call()
|
||||
|
||||
assert body(first) == MESSAGES_RESPONSE
|
||||
assert len(server.requests) == 1
|
||||
assert body(hit) == MESSAGES_RESPONSE
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_redis_hit_fills_the_worker_memory_tier(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
|
||||
) -> None:
|
||||
use_cache(monkeypatch, fake_redis.cache())
|
||||
|
||||
await call()
|
||||
await call()
|
||||
fake_redis.client.flushall()
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert body(hit) == MESSAGES_RESPONSE
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_stored_response_reaches_the_worker_memory_tier_only_when_read(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
|
||||
) -> None:
|
||||
use_cache(monkeypatch, fake_redis.cache())
|
||||
server.expect(2)
|
||||
|
||||
await call()
|
||||
fake_redis.client.flushall()
|
||||
await call()
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_a_native_hit_reports_the_key_the_entry_is_stored_under(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = InMemoryCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
await call()
|
||||
hit: Final = await call()
|
||||
|
||||
(stored,) = backend.cache_dict
|
||||
assert cache_key(hit) == stored_cache_key(hit) == stored
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("served", ("native", "python"))
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_hit_returns_what_the_success_deployment_hook_handed_the_first_caller(
|
||||
call: Call, server: RecordingServer, served: str
|
||||
) -> None:
|
||||
class Replace(CustomLogger):
|
||||
calls = 0
|
||||
|
||||
async def async_post_call_success_deployment_hook(
|
||||
self, request_data: dict[str, object], response: object, call_type: object
|
||||
) -> LLMResponseTypes | None:
|
||||
Replace.calls += 1
|
||||
return _MESSAGE.validate_python({**_OBJECT.validate_python(response), "id": f"hooked-{Replace.calls}"})
|
||||
|
||||
async def ask() -> object:
|
||||
if served == "native":
|
||||
return await call()
|
||||
with routed():
|
||||
return await call()
|
||||
|
||||
with rebound(litellm, "callbacks", [Replace()]):
|
||||
first: Final = await ask()
|
||||
second: Final = await ask()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert response_id(first) == "hooked-1"
|
||||
assert response_id(second) == "hooked-1"
|
||||
assert Replace.calls == 1
|
||||
|
|
@ -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}),
|
||||
(
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -1,16 +1,18 @@
|
|||
from collections.abc import Generator
|
||||
from functools import partial
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.rust_bridge.catalog import Route
|
||||
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
|
||||
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE, call
|
||||
from tests.test_litellm_rust.support.caching import Call
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.routing import native_route
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server(recording_server: RecordingServer) -> RecordingServer:
|
||||
recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
|
||||
recording_server.serve(ResponseSpec(body=RESPONSES_RESPONSE))
|
||||
return recording_server
|
||||
|
||||
|
||||
|
|
@ -18,3 +20,8 @@ def server(recording_server: RecordingServer) -> RecordingServer:
|
|||
def native() -> Generator[None]:
|
||||
with native_route(Route.RESPONSES):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def native_call(server: RecordingServer) -> Call:
|
||||
return partial(call, server)
|
||||
|
|
|
|||
375
tests/test_litellm_rust/responses/test_cache.py
Normal file
375
tests/test_litellm_rust/responses/test_cache.py
Normal file
|
|
@ -0,0 +1,375 @@
|
|||
import asyncio
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache, CacheMode
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.caching import CachingSupportedCallTypes
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.utils import LLMResponseTypes
|
||||
from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE, RESPONSES_TEXT
|
||||
from tests.test_litellm_rust.support.caching import (
|
||||
Call,
|
||||
CallerCache,
|
||||
FailingLookupCache,
|
||||
FakeRedis,
|
||||
HeldWriteCache,
|
||||
assert_cache_hit_keeps_budget_spend_and_counts_usage,
|
||||
settle_cache_writes,
|
||||
use_cache,
|
||||
)
|
||||
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
|
||||
from tests.test_litellm_rust.support.isolation import rebound
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.results import cache_key, response_id, stored_cache_key
|
||||
from tests.test_litellm_rust.support.routing import assert_served_natively, routed
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
EMPTY: Final = {**RESPONSES_RESPONSE, "output": []}
|
||||
OWN_CALL_TYPES: Final = ("responses",)
|
||||
|
||||
|
||||
def output_text(response: object) -> str:
|
||||
assert isinstance(response, ResponsesAPIResponse)
|
||||
return response.output_text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_repeated_native_call_is_answered_from_the_python_cache(
|
||||
call: Call, server: RecordingServer, asynchronous: bool
|
||||
) -> None:
|
||||
first: Final = await call(asynchronous=asynchronous)
|
||||
second: Final = await call(asynchronous=asynchronous)
|
||||
|
||||
assert_served_natively(server, first)
|
||||
assert cache_key(first) is None
|
||||
assert isinstance(cache_key(second), str)
|
||||
assert output_text(second) == output_text(first) == RESPONSES_TEXT
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("first", "second"),
|
||||
(
|
||||
({"cache": {"no-store": True}}, {}),
|
||||
({}, {"cache": {"no-cache": True}}),
|
||||
({}, {"caching": False}),
|
||||
({}, {"input": "Goodbye"}),
|
||||
({}, {"instructions": "Be brief"}),
|
||||
({"cache": {"namespace": "first"}}, {"cache": {"namespace": "second"}}),
|
||||
({"metadata": {"redis_namespace": "first"}}, {"metadata": {"redis_namespace": "second"}}),
|
||||
),
|
||||
ids=(
|
||||
"no-store",
|
||||
"no-cache",
|
||||
"caching-off-read",
|
||||
"other-input",
|
||||
"other-param",
|
||||
"cache-namespace",
|
||||
"metadata-namespace",
|
||||
),
|
||||
)
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_cache_controls_and_keys_keep_the_second_call_on_the_provider(
|
||||
call: Call, server: RecordingServer, first: Mapping[str, object], second: Mapping[str, object]
|
||||
) -> None:
|
||||
server.expect(2)
|
||||
|
||||
await call(**first)
|
||||
response: Final = await call(**second)
|
||||
|
||||
assert_served_natively(server, count=2)
|
||||
assert cache_key(response) is None
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_caching_off_skips_the_read_but_still_writes(call: Call, server: RecordingServer) -> None:
|
||||
await call(caching=False)
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert output_text(hit) == RESPONSES_TEXT
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_each_namespace_keeps_its_own_entry(call: Call, server: RecordingServer) -> None:
|
||||
server.expect(2)
|
||||
|
||||
await call(cache={"namespace": "first"})
|
||||
await call(cache={"namespace": "second"})
|
||||
hit: Final = await call(cache={"namespace": "first"})
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_the_key_is_built_from_the_call_arguments_not_the_wire_body(call: Call, server: RecordingServer) -> None:
|
||||
class Rewrite(RecordingLogger):
|
||||
temperature = 0.1
|
||||
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
|
||||
request_body(kwargs)["temperature"] = self.temperature
|
||||
super().log_pre_api_call(model, messages, kwargs)
|
||||
|
||||
logger: Final = Rewrite()
|
||||
|
||||
await call(callbacks=[logger])
|
||||
logger.temperature = 0.8
|
||||
hit: Final = await call(callbacks=[logger])
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.1
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.parametrize(
|
||||
("supported", "expected_requests"), ((["embedding"], 2), (OWN_CALL_TYPES, 1)), ids=("other-type", "own-type")
|
||||
)
|
||||
async def test_supported_call_types_gate_native_reads_and_writes(
|
||||
call: Call,
|
||||
server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
asynchronous: bool,
|
||||
supported: Sequence[CachingSupportedCallTypes],
|
||||
expected_requests: int,
|
||||
) -> None:
|
||||
use_cache(monkeypatch, Cache(supported_call_types=list(supported)))
|
||||
server.expect(expected_requests)
|
||||
|
||||
await call(asynchronous=asynchronous)
|
||||
second: Final = await call(asynchronous=asynchronous)
|
||||
|
||||
assert len(server.requests) == expected_requests
|
||||
assert isinstance(cache_key(second), str) == (expected_requests == 1)
|
||||
assert_served_natively(server, count=expected_requests)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("options", "expected_requests"), (({}, 2), ({"cache": {"use-cache": True}}, 1)), ids=("plain", "opt-in")
|
||||
)
|
||||
async def test_default_off_mode_caches_only_calls_that_opt_in(
|
||||
call: Call,
|
||||
server: RecordingServer,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
options: Mapping[str, object],
|
||||
expected_requests: int,
|
||||
) -> None:
|
||||
use_cache(monkeypatch, Cache(mode=CacheMode.default_off))
|
||||
server.expect(expected_requests)
|
||||
|
||||
await call(**options)
|
||||
second: Final = await call(**options)
|
||||
|
||||
assert len(server.requests) == expected_requests
|
||||
assert isinstance(cache_key(second), str) == (expected_requests == 1)
|
||||
assert_served_natively(server, count=expected_requests)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_native_hit_is_logged_as_a_free_cache_hit(call: Call, server: RecordingServer) -> None:
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await call(callbacks=[recorder])
|
||||
await call(callbacks=[recorder])
|
||||
await drain_logging()
|
||||
|
||||
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
|
||||
miss, hit = (_OBJECT.validate_python(event.kwargs) for event in successes)
|
||||
assert len(server.requests) == 1
|
||||
assert miss["cache_hit"] is not True
|
||||
assert TypeAdapter(float).validate_python(miss["response_cost"]) > 0
|
||||
assert hit["cache_hit"] is True
|
||||
assert hit["response_cost"] == 0
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_python_cache_lookups_run_in_the_callers_task(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = CallerCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
await call()
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert isinstance(cache_key(hit), str)
|
||||
assert backend.lookup_tasks == [asyncio.current_task(), asyncio.current_task()]
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_cancelling_a_pending_lookup_sends_and_stores_nothing(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
entered: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
backend: Final = CallerCache(entered, release)
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
server.expect(0)
|
||||
|
||||
task: Final = asyncio.create_task(call())
|
||||
await asyncio.wait_for(entered.wait(), timeout=5)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
release.set()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert server.requests == []
|
||||
assert backend.storage.cache_dict == {}
|
||||
assert_served_natively(server, count=0)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_cache_hit_keeps_budget_spend_and_counts_usage(call: Call, server: RecordingServer) -> None:
|
||||
usage: Final = TypeAdapter(dict[str, int]).validate_python(RESPONSES_RESPONSE["usage"])
|
||||
|
||||
await assert_cache_hit_keeps_budget_spend_and_counts_usage(call, server, usage["total_tokens"])
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_response_without_output_is_not_stored(call: Call, server: RecordingServer) -> None:
|
||||
server.serve(ResponseSpec(body=EMPTY))
|
||||
server.expect(2)
|
||||
|
||||
await call()
|
||||
await call()
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_a_stored_response_without_output_is_a_miss(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = InMemoryCache()
|
||||
cache: Final = use_cache(monkeypatch, Cache(_backend=backend))
|
||||
await call()
|
||||
(key,) = backend.cache_dict
|
||||
cache.add_cache(EMPTY, cache_key=key)
|
||||
server.expect(2)
|
||||
|
||||
response: Final = await call()
|
||||
|
||||
assert len(server.requests) == 2
|
||||
assert output_text(response) == RESPONSES_TEXT
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_a_failing_cache_lookup_runs_the_call_uncached(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
use_cache(monkeypatch, FailingLookupCache())
|
||||
server.expect(2)
|
||||
|
||||
first: Final = await call()
|
||||
await call()
|
||||
|
||||
assert output_text(first) == RESPONSES_TEXT
|
||||
assert len(server.requests) == 2
|
||||
assert_served_natively(server, count=2)
|
||||
|
||||
|
||||
async def test_the_response_returns_before_the_cache_write_lands(
|
||||
call: Call, native_call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = HeldWriteCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
first: Final = await asyncio.wait_for(native_call(), timeout=5)
|
||||
backend.release.set()
|
||||
await settle_cache_writes()
|
||||
hit: Final = await call()
|
||||
|
||||
assert output_text(first) == RESPONSES_TEXT
|
||||
assert len(server.requests) == 1
|
||||
assert output_text(hit) == RESPONSES_TEXT
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_redis_hit_fills_the_worker_memory_tier(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
|
||||
) -> None:
|
||||
use_cache(monkeypatch, fake_redis.cache())
|
||||
|
||||
await call()
|
||||
await call()
|
||||
fake_redis.client.flushall()
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert output_text(hit) == RESPONSES_TEXT
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_stored_response_lands_in_the_worker_memory_tier(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch, fake_redis: FakeRedis
|
||||
) -> None:
|
||||
use_cache(monkeypatch, fake_redis.cache())
|
||||
|
||||
first: Final = await call()
|
||||
fake_redis.client.flushall()
|
||||
hit: Final = await call()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert output_text(hit) == output_text(first)
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
async def test_a_hit_reports_the_key_the_entry_is_stored_under(
|
||||
call: Call, server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
backend: Final = InMemoryCache()
|
||||
use_cache(monkeypatch, Cache(_backend=backend))
|
||||
|
||||
await call()
|
||||
hit: Final = await call()
|
||||
|
||||
(stored,) = backend.cache_dict
|
||||
assert len(server.requests) == 1
|
||||
assert stored_cache_key(hit) == stored
|
||||
assert_served_natively(server, count=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("served", ("native", "python"))
|
||||
@pytest.mark.usefixtures("cache")
|
||||
async def test_a_hit_returns_what_the_success_deployment_hook_handed_the_first_caller(
|
||||
call: Call, server: RecordingServer, served: str
|
||||
) -> None:
|
||||
class Replace(CustomLogger):
|
||||
calls = 0
|
||||
|
||||
async def async_post_call_success_deployment_hook(
|
||||
self, request_data: dict[str, object], response: object, call_type: object
|
||||
) -> LLMResponseTypes | None:
|
||||
Replace.calls += 1
|
||||
assert isinstance(response, ResponsesAPIResponse)
|
||||
return response.model_copy(update={"id": f"hooked-{Replace.calls}"})
|
||||
|
||||
async def ask() -> object:
|
||||
if served == "native":
|
||||
return await call()
|
||||
with routed():
|
||||
return await call()
|
||||
|
||||
with rebound(litellm, "callbacks", [Replace()]):
|
||||
first: Final = await ask()
|
||||
second: Final = await ask()
|
||||
|
||||
assert len(server.requests) == 1
|
||||
assert response_id(first) == "hooked-1"
|
||||
assert response_id(second) == "hooked-1"
|
||||
assert Replace.calls == 1
|
||||
|
|
@ -5,7 +5,6 @@ from typing import Final, cast # noqa: TID251 # narrows legacy callable signat
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.responses import dispatch as responses_dispatch
|
||||
from litellm.responses import main as python_responses
|
||||
|
|
@ -287,7 +286,6 @@ ENCODED_ID: Final = ResponsesAPIRequestUtils._build_encrypted_item_id("deploymen
|
|||
(Feature.DEPLOYMENT_RESPONSE_ID, {"litellm_metadata": {"model_info": {"id": "deployment-1"}}}, {}),
|
||||
(Feature.ENCODED_ITEM_IDS, {"input": [{"type": "reasoning", "id": ENCODED_ID, "summary": []}]}, {}),
|
||||
(Feature.MODIFY_PARAMS, {}, {"modify_params": True}),
|
||||
(Feature.RESPONSE_CACHE, {}, {"cache": Cache()}),
|
||||
(
|
||||
Feature.MANAGED_RESOURCE_IDS,
|
||||
{
|
||||
|
|
|
|||
164
tests/test_litellm_rust/support/caching.py
Normal file
164
tests/test_litellm_rust/support/caching.py
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
import asyncio
|
||||
import threading
|
||||
from collections.abc import Callable, Coroutine, Generator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, TypeAlias
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
import redis
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.caching.base_cache import BaseCache
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.caching.caching_handler import (
|
||||
_PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request
|
||||
in_memory_cache_obj,
|
||||
)
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
model_budget_spend_cache_key,
|
||||
)
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
|
||||
from litellm.proxy.utils import InternalUsageCache
|
||||
from litellm.types.caching import LiteLLMCacheType
|
||||
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer
|
||||
|
||||
Call: TypeAlias = Callable[..., Coroutine[object, object, object]]
|
||||
|
||||
|
||||
class CallerCache(BaseCache):
|
||||
def __init__(self, lookup_entered: asyncio.Event | None = None, release: asyncio.Event | None = None) -> None:
|
||||
super().__init__()
|
||||
self.storage: Final = InMemoryCache()
|
||||
self.lookup_tasks: Final[list[asyncio.Task[object] | None]] = [] # mutable-ok: test double records each lookup
|
||||
self.lookup_entered: Final = lookup_entered
|
||||
self.release: Final = release
|
||||
|
||||
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], ttl: float | None = None) -> None:
|
||||
await self.storage.async_set_cache_pipeline(cache_list, ttl=ttl)
|
||||
|
||||
async def async_get_cache(self, key: str, **kwargs: object) -> object:
|
||||
self.lookup_tasks.append(asyncio.current_task())
|
||||
if self.lookup_entered is not None and self.release is not None:
|
||||
self.lookup_entered.set()
|
||||
await self.release.wait()
|
||||
return self.storage.get_cache(key, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
await asyncio.sleep(0)
|
||||
self.storage.set_cache(key, value, **kwargs)
|
||||
|
||||
|
||||
class HeldWriteCache(InMemoryCache):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.release: Final = asyncio.Event()
|
||||
|
||||
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
await self.release.wait()
|
||||
await super().async_set_cache(key, value, **kwargs)
|
||||
|
||||
|
||||
class FailingLookupCache(Cache):
|
||||
def should_use_cache(self, **kwargs: object) -> bool:
|
||||
raise RuntimeError("cache lookup failed")
|
||||
|
||||
|
||||
def use_cache(monkeypatch: pytest.MonkeyPatch, cache: Cache) -> Cache:
|
||||
monkeypatch.setattr(litellm, "cache", cache)
|
||||
return cache
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cache(monkeypatch: pytest.MonkeyPatch) -> Cache:
|
||||
return use_cache(monkeypatch, Cache())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def call(native_call: Call) -> Call:
|
||||
async def send(**kwargs: object) -> object:
|
||||
response: Final = await native_call(**kwargs)
|
||||
await settle_cache_writes()
|
||||
return response
|
||||
|
||||
return send
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FakeRedis:
|
||||
url: str
|
||||
client: redis.Redis
|
||||
|
||||
def cache(self) -> Cache:
|
||||
parsed: Final = urlparse(self.url)
|
||||
in_memory_cache_obj.flush_cache()
|
||||
return Cache(type=LiteLLMCacheType.REDIS, host=parsed.hostname, port=str(parsed.port))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_redis() -> Generator[FakeRedis]:
|
||||
server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis")
|
||||
worker: Final = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
worker.start()
|
||||
url: Final = f"redis://127.0.0.1:{server.server_address[1]}"
|
||||
try:
|
||||
yield FakeRedis(url=url, client=redis.Redis.from_url(url))
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
worker.join(timeout=5)
|
||||
|
||||
|
||||
async def settle_cache_writes() -> None:
|
||||
await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
|
||||
|
||||
|
||||
async def assert_cache_hit_keeps_budget_spend_and_counts_usage(
|
||||
send: Call, server: RecordingServer, tokens: int
|
||||
) -> None:
|
||||
counters: Final = litellm.DualCache()
|
||||
budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters)
|
||||
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
|
||||
InternalUsageCache(counters), model_group_resolver=lambda model: model
|
||||
)
|
||||
recorder: Final = RecordingLogger()
|
||||
key_hash: Final = "a" * 64
|
||||
options: Final = {
|
||||
"callbacks": [budget, limiter, recorder],
|
||||
"metadata": {
|
||||
"user_api_key": key_hash,
|
||||
"model_group": "cached-model",
|
||||
"user_api_key_model_max_budget": {"cached-model": {"max_budget": 1, "budget_duration": "1h"}},
|
||||
},
|
||||
}
|
||||
spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h")
|
||||
token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens")
|
||||
objects: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
await send(**options)
|
||||
await drain_logging()
|
||||
miss: Final = objects.validate_python((await recorder.wait_for_async("async_log_success_event"))[0].kwargs)
|
||||
cost: Final = TypeAdapter(float).validate_python(miss["response_cost"])
|
||||
assert cost > 0
|
||||
assert counters.get_cache(spend_key) == pytest.approx(cost)
|
||||
assert counters.get_cache(token_key) == tokens
|
||||
|
||||
await send(**options)
|
||||
await drain_logging()
|
||||
successes: Final = await recorder.wait_for_async("async_log_success_event", count=2)
|
||||
hit: Final = objects.validate_python(successes[-1].kwargs)
|
||||
miss_payload: Final = objects.validate_python(miss["standard_logging_object"])
|
||||
hit_payload: Final = objects.validate_python(hit["standard_logging_object"])
|
||||
assert len(server.requests) == 1
|
||||
assert hit["cache_hit"] is True
|
||||
assert hit["response_cost"] == hit_payload["response_cost"] == 0
|
||||
assert hit_payload["custom_llm_provider"] == miss_payload["custom_llm_provider"]
|
||||
assert hit_payload["total_tokens"] == miss_payload["total_tokens"] == tokens
|
||||
assert counters.get_cache(spend_key) == pytest.approx(cost)
|
||||
assert counters.get_cache(token_key) == 2 * tokens
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue