scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

Merge pull request #42313 from BerriAI/litellm_rust_cache_s3

Browse source 
feat(rust): add native S3 cache backend
This commit is contained in:
yujonglee 2026-09-21 16:12:54 -07:00 • committed by GitHub
commit a5431244ef
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1664 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

176
litellm-rust/Cargo.lock generated
View file

@ -40,6 +40,12 @@ dependencies = [
"cc",
]
[[package]]
name = "allocator-api2"
version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
[[package]]
name = "android_system_properties"
version = "0.1.6"
@ -230,6 +236,7 @@ dependencies = [
"aws-credential-types",
"aws-sigv4",
"aws-smithy-async",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-runtime",
"aws-smithy-runtime-api",
@ -238,7 +245,9 @@ dependencies = [
"bytes",
"bytes-utils",
"fastrand",
"http 0.2.12",
"http 1.4.2",
"http-body 0.4.6",
"http-body 1.1.0",
"percent-encoding",
"pin-project-lite",
@ -272,6 +281,43 @@ dependencies = [
"tracing",
]
[[package]]
name = "aws-sdk-s3"
version = "1.146.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2cd651b4400d4011b8927b83a9552bf90ff11e6e5da0b9f0a7583247aceec971"
dependencies = [
"arc-swap",
"aws-credential-types",
"aws-runtime",
"aws-sigv4",
"aws-smithy-async",
"aws-smithy-checksums",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-json",
"aws-smithy-observability",
"aws-smithy-runtime",
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"aws-smithy-xml 0.62.1",
"aws-types",
"bytes",
"fastrand",
"hex",
"hmac",
"http 0.2.12",
"http 1.4.2",
"http-body 1.1.0",
"lru",
"percent-encoding",
"regex-lite",
"sha2 0.11.0",
"tracing",
"url",
]
[[package]]
name = "aws-sdk-secretsmanager"
version = "1.117.0"
@ -316,7 +362,7 @@ dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"aws-smithy-xml",
"aws-smithy-xml 0.61.1",
"aws-types",
"fastrand",
"http 0.2.12",
@ -332,6 +378,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70"
dependencies = [
"aws-credential-types",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-runtime-api",
"aws-smithy-types",
@ -359,10 +406,31 @@ dependencies = [
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.1"
name = "aws-smithy-checksums"
version = "0.65.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
checksum = "b67ecd999972b58e67cab052f5129906c08c25883bd0788ceefc55ef97d61307"
dependencies = [
"aws-smithy-http",
"aws-smithy-types",
"bytes",
"crc-fast",
"hex",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"md-5",
"pin-project-lite",
"sha1 0.11.0",
"sha2 0.11.0",
"tracing",
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "80c2051c2f1016fb8e6548dd07b8bc2ac9c3fe583721444b92f515e856d31609"
dependencies = [
"aws-smithy-types",
"bytes",
@ -375,6 +443,7 @@ version = "0.64.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-runtime-api",
"aws-smithy-types",
"bytes",
@ -554,6 +623,18 @@ dependencies = [
"xmlparser",
]
[[package]]
name = "aws-smithy-xml"
version = "0.62.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b932c8d6dc127fc980eecd78f8694ae9b9551b69a93a7def2a199c1c0033daf"
dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"xmlparser",
]
[[package]]
name = "aws-types"
version = "1.6.0"
@ -980,6 +1061,16 @@ dependencies = [
"libc",
]
[[package]]
name = "crc-fast"
version = "1.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5"
dependencies = [
"digest 0.10.7",
"spin",
]
[[package]]
name = "crc16"
version = "0.4.0"
@ -1348,7 +1439,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -1925,6 +2016,8 @@ version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
dependencies = [
"allocator-api2",
"equivalent",
"foldhash",
]
@ -2148,7 +2241,7 @@ dependencies = [
"libc",
"percent-encoding",
"pin-project-lite",
"socket2 0.6.5",
"socket2 0.5.10",
"tokio",
"tower-service",
"tracing",
@ -2664,6 +2757,21 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-s3"
version = "0.1.0"
dependencies = [
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-types",
"aws-types",
"litellm-auth-aws",
"litellm-cache",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -2826,6 +2934,7 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
@ -2834,6 +2943,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-cache-s3",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -3075,6 +3185,15 @@ version = "0.4.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
[[package]]
name = "lru"
version = "0.18.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ff9840bcc50b71349309900da0ce7279aa336ae71d73250b07998932c7d97c25"
dependencies = [
"hashbrown 0.17.1",
]
[[package]]
name = "lru-slab"
version = "0.1.2"
@ -3097,6 +3216,16 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
[[package]]
name = "md-5"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98"
dependencies = [
"cfg-if",
"digest 0.11.3",
]
[[package]]
name = "memchr"
version = "2.8.3"
@ -3739,7 +3868,7 @@ dependencies = [
"quinn-udp",
"rustc-hash",
"rustls 0.23.42",
"socket2 0.6.5",
"socket2 0.5.10",
"thiserror 2.0.19",
"tokio",
"tracing",
@ -3778,9 +3907,9 @@ dependencies = [
"cfg_aliases",
"libc",
"once_cell",
"socket2 0.6.5",
"socket2 0.5.10",
"tracing",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4266,7 +4395,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4337,7 +4466,7 @@ dependencies = [
"security-framework",
"security-framework-sys",
"webpki-root-certs",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4610,6 +4739,17 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "sha1"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214"
dependencies = [
"cfg-if",
"cpufeatures 0.3.0",
"digest 0.11.3",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
@ -4726,6 +4866,12 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "spin"
version = "0.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3"
[[package]]
name = "spm_precompiled"
version = "0.1.4"
@ -4892,10 +5038,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"getrandom 0.3.4",
"once_cell",
"rustix",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -5359,7 +5505,7 @@ dependencies = [
"rand 0.8.7",
"rustls 0.23.42",
"rustls-pki-types",
"sha1",
"sha1 0.10.7",
"thiserror 1.0.69",
"utf-8",
]
@ -5768,7 +5914,7 @@ version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]

1
litellm-rust/Cargo.toml
View file

@ -33,6 +33,7 @@ litellm-cache = { path = "crates/cache" }
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-s3 = { path = "crates/cache-s3" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-response = { path = "crates/cache-response" }

20
litellm-rust/crates/cache-s3/Cargo.toml Normal file
View file

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-s3"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-auth-aws.workspace = true
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
aws-credential-types = "1.3.0"
aws-smithy-types = "1.6.0"
aws-types = "1.6.0"
tokio.workspace = true
[dev-dependencies]
wiremock = "0.6.5"
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

101
litellm-rust/crates/cache-s3/src/auth.rs Normal file
View file

@ -0,0 +1,101 @@
use aws_credential_types::{
Credentials as AwsCredentials,
provider::{ProvideCredentials, error::CredentialsError, future},
};
use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
#[derive(Clone)]
pub(crate) struct Credentials {
config: AwsAuthConfig,
env: fn(&str) -> Option<String>,
}
impl Credentials {
pub(crate) fn new(config: AwsAuthConfig) -> Self {
Self::with_env(config, |name| std::env::var(name).ok())
}
pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
Self { config, env }
}
}
impl ProvideCredentials for Credentials {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::new(async {
if let (Some(access_key_id), Some(secret_access_key)) = (
self.config.access_key_id.clone(),
self.config.secret_access_key.clone(),
) {
return Ok(AwsCredentials::new(
access_key_id,
secret_access_key,
self.config.session_token.clone(),
None,
"litellm-s3-cache",
));
}
resolve_credentials(self.config.clone(), &self.env)
.await
.map_err(|_| CredentialsError::provider_error("S3 cache authentication failed"))
})
}
}
impl std::fmt::Debug for Credentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Credentials").finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn explicit_keys_ignore_an_ambient_session_token() {
let provider = Credentials::with_env(
AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
|name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()),
);
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "key");
assert_eq!(credentials.secret_access_key(), "secret");
assert_eq!(credentials.session_token(), None);
}
#[tokio::test]
async fn explicit_keys_keep_their_session_token() {
let provider = Credentials::new(AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
session_token: Some("t".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.session_token(), Some("t"));
}
#[tokio::test]
async fn environment_keys_resolve_with_their_session_token() {
let provider = Credentials::with_env(AwsAuthConfig::default(), |name| match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "env-key");
assert_eq!(credentials.secret_access_key(), "env-secret");
assert_eq!(credentials.session_token(), Some("env-token"));
}
}

220
litellm-rust/crates/cache-s3/src/cache.rs Normal file
View file

@ -0,0 +1,220 @@
use std::{
future::Future,
sync::Arc,
time::{Duration, SystemTime},
};
use aws_sdk_s3::{
config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
error::SdkError,
primitives::ByteStream,
};
use aws_smithy_types::{DateTime, date_time::Format};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use crate::auth::Credentials;
pub struct S3Endpoint {
pub url: String,
}
pub struct S3CacheConfig {
pub bucket: String,
pub key_prefix: String,
pub region: String,
pub endpoint: Option<S3Endpoint>,
pub auth: AwsAuthConfig,
}
pub struct S3Cache<C: CacheCodec> {
client: aws_sdk_s3::Client,
codec: C,
runtime: Handle,
bucket: Arc<str>,
key_prefix: Arc<str>,
region: Arc<str>,
endpoint: Option<Arc<str>>,
}
impl<C: CacheCodec> S3Cache<C> {
pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self {
let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
let base = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(config.region.clone()))
.credentials_provider(Credentials::new(config.auth))
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
let builder = match &endpoint_url {
Some(url) => base.endpoint_url(url).force_path_style(true),
None => base,
};
Self {
client: aws_sdk_s3::Client::from_conf(builder.build()),
codec,
runtime,
bucket: config.bucket.into(),
key_prefix: config.key_prefix.into(),
region: config.region.into(),
endpoint: endpoint_url.map(Into::into),
}
}
pub fn bucket(&self) -> &str {
&self.bucket
}
pub fn key_prefix(&self) -> &str {
&self.key_prefix
}
pub fn region(&self) -> &str {
&self.region
}
pub fn endpoint(&self) -> Option<&str> {
self.endpoint.as_deref()
}
pub fn to_s3_key(&self, key: &str) -> String {
format!("{}{}", self.key_prefix, key.replace(':', "/"))
}
fn block_on<F: Future>(&self, future: F) -> F::Output {
if Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
async fn put(
&self,
key: &str,
value: C::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let s3_key = self.to_s3_key(key);
let body = self.codec.encode(&value)?;
let request = self
.client
.put_object()
.bucket(self.bucket.as_ref())
.key(&s3_key)
.body(ByteStream::from(body))
.content_type("application/json")
.content_language("en")
.content_disposition(format!("inline; filename=\"{s3_key}.json\""));
let request = match context.ttl {
Some(ttl) => {
let seconds = ttl.as_secs_f64();
request
.cache_control(format!("immutable, max-age={seconds}, s-maxage={seconds}"))
.expires(DateTime::from(SystemTime::now() + ttl))
}
None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"),
};
request.send().await.map_err(|_| Error::Unavailable)?;
Ok(())
}
async fn get(&self, key: &str) -> Result<Option<C::Value>, Error> {
let output = match self
.client
.get_object()
.bucket(self.bucket.as_ref())
.key(self.to_s3_key(key))
.send()
.await
{
Ok(output) => output,
Err(error) => {
if let SdkError::ServiceError(service) = &error {
let status = error
.raw_response()
.map(|response| response.status().as_u16());
let not_found = service.err().is_no_such_key()
|| service.err().meta().code() == Some("AccessDenied")
|| status == Some(404)
|| status == Some(403);
if not_found {
return Ok(None);
}
}
return Err(Error::Unavailable);
}
};
if let Some(expires) = output.expires_string()
&& let Ok(expires) = DateTime::from_str(expires, Format::HttpDate)
&& expires < DateTime::from(SystemTime::now())
{
return Ok(None);
}
let bytes = output
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?
.into_bytes();
self.codec.decode(&bytes).map(Some)
}
}
impl<C: CacheCodec> BaseCache for S3Cache<C> {
type Value = C::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
self.block_on(self.put(key, value, context))
}
fn get_cache(&self, key: &str, _context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.block_on(self.get(key))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
self.put(key, value, &context).await
}
async fn async_get_cache(
&self,
key: &str,
_context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.get(key).await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<C: CacheCodec> BatchCache for S3Cache<C> {}
impl<C: CacheCodec> FlushCache for S3Cache<C> {
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
}

4
litellm-rust/crates/cache-s3/src/lib.rs Normal file
View file

@ -0,0 +1,4 @@
mod auth;
mod cache;
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

278
litellm-rust/crates/cache-s3/tests/cache.rs Normal file
View file

@ -0,0 +1,278 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
use serde_json::{Value, json};
use tokio::runtime::Handle;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{method, path},
};
fn config(endpoint: String) -> S3CacheConfig {
S3CacheConfig {
bucket: "cache-bucket".to_string(),
key_prefix: "team/".to_string(),
region: "us-east-1".to_string(),
endpoint: Some(S3Endpoint { url: endpoint }),
auth: AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
}
}
fn cache(endpoint: &str) -> S3Cache<JsonCodec<Value>> {
S3Cache::new(
config(endpoint.to_string()),
JsonCodec::<Value>::new(),
Handle::current(),
)
}
async fn mock_server() -> MockServer {
let server = MockServer::start().await;
Mock::given(method("PUT"))
.respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\""))
.mount(&server)
.await;
server
}
fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option<SystemTime> {
use aws_smithy_types::{DateTime, date_time::Format};
headers
.get(name)
.and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok())
.map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos()))
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn set_writes_python_metadata_with_and_without_ttl() {
let server = mock_server().await;
let cache = cache(&server.uri());
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(90)),
};
cache
.set_cache("alpha:beta", json!({"answer": 1}), &context)
.unwrap();
cache
.set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default())
.unwrap();
let requests = server.received_requests().await.unwrap();
let ttl_request = requests
.iter()
.find(|request| request.url.path() == "/cache-bucket/team/alpha/beta")
.expect("ttl write should hit the converted S3 key");
assert_eq!(
ttl_request.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=90, s-maxage=90"
);
assert_eq!(
ttl_request.headers["content-type"].to_str().unwrap(),
"application/json"
);
assert_eq!(
ttl_request.headers["content-language"].to_str().unwrap(),
"en"
);
assert_eq!(
ttl_request.headers["content-disposition"].to_str().unwrap(),
"inline; filename=\"team/alpha/beta.json\""
);
let expires = http_date_from(&ttl_request.headers, "expires").expect("ttl write sets Expires");
let remaining = expires.duration_since(SystemTime::now()).unwrap();
assert!(remaining > Duration::from_secs(60) && remaining <= Duration::from_secs(91));
assert_eq!(
serde_json::from_slice::<Value>(&ttl_request.body).unwrap(),
json!({"answer": 1})
);
let plain = requests
.iter()
.find(|request| request.url.path() == "/cache-bucket/team/plain")
.expect("no-ttl write should hit the converted S3 key");
assert_eq!(
plain.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=31536000, s-maxage=31536000"
);
assert!(plain.headers.get("expires").is_none());
assert_eq!(
plain.headers["content-disposition"].to_str().unwrap(),
"inline; filename=\"team/plain.json\""
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_hit_miss_expired_and_invalid_entries() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/missing"))
.respond_with(
ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/denied"))
.respond_with(
ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/expired"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
.set_body_json(json!({"answer": 4})),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/malformed"))
.respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry"))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
assert_eq!(
cache.get_cache("hit", &context).unwrap(),
Some(json!({"answer": 3}))
);
assert_eq!(cache.get_cache("missing", &context).unwrap(), None);
assert_eq!(cache.get_cache("denied", &context).unwrap(), None);
assert_eq!(cache.get_cache("expired", &context).unwrap(), None);
assert_eq!(
cache.get_cache("malformed", &context),
Err(Error::InvalidEntry)
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
let server = mock_server().await;
for (key, status, body) in [
("first", 200, "{\"answer\": 1}"),
("invalid", 200, "garbage"),
] {
Mock::given(method("GET"))
.and(path(format!("/cache-bucket/team/{key}")))
.respond_with(ResponseTemplate::new(status).set_body_string(body))
.mount(&server)
.await;
}
Mock::given(method("GET"))
.and(path("/cache-bucket/team/miss"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
let keys = vec![
"first".to_string(),
"miss".to_string(),
"invalid".to_string(),
];
let entries = cache.batch_get_cache(&keys, &context).unwrap();
assert_eq!(
entries,
vec![
BatchEntry::Hit(json!({"answer": 1})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unsupported_and_noop_capabilities_match_python() {
let server = mock_server().await;
let cache = cache(&server.uri());
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
cache.flush_cache().unwrap();
cache.disconnect().await.unwrap();
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext {
ttl: Some(Duration::from_secs(45)),
}),
Some(Duration::from_secs(45))
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[test]
fn key_conversion_prefixes_and_splits_colons() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.unwrap();
let _guard = runtime.enter();
let cache = S3Cache::new(
S3CacheConfig {
key_prefix: "team/".to_string(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(cache.bucket(), "cache-bucket");
assert_eq!(cache.key_prefix(), "team/");
assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c");
assert_eq!(cache.to_s3_key("plain"), "team/plain");
let unprefixed = S3Cache::new(
S3CacheConfig {
key_prefix: String::new(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(unprefixed.to_s3_key("a:b"), "a/b");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sync_methods_block_inside_and_outside_the_runtime() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9})))
.mount(&server)
.await;
let uri = server.uri();
let cache = tokio::task::spawn_blocking(move || {
let cache = cache(&uri);
let context = ExactCacheContext::default();
cache
.set_cache("key", json!({"answer": 9}), &context)
.unwrap();
cache.get_cache("key", &context).unwrap()
})
.await
.unwrap();
assert_eq!(cache, Some(json!({"answer": 9})));
}

2
litellm-rust/crates/framer/Cargo.toml
View file

@ -11,7 +11,7 @@ aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
sse = ["dep:sse-stream"]
[dependencies]
aws-smithy-eventstream = { version = "=0.61.1", optional = true }
aws-smithy-eventstream = { version = "=0.61.4", optional = true }
aws-smithy-types = { version = "1.6.1", optional = true }
bytes = "1"
futures-util.workspace = true

2
litellm-rust/crates/llms/Cargo.toml
View file

@ -34,7 +34,7 @@ tokio = { workspace = true, features = ["sync"] }
url.workspace = true
[dev-dependencies]
aws-smithy-eventstream = "=0.61.1"
aws-smithy-eventstream = "=0.61.4"
aws-smithy-types = "1.6.1"
rstest.workspace = true
tokio.workspace = true

2
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -24,11 +24,13 @@ litellm-cache.workspace = true
litellm-cache-azure-blob.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-s3.workspace = true
litellm-cache-gcs.workspace = true
litellm-cache-disk.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-callbacks-legacy-python.workspace = true
litellm-core.workspace = true
litellm-core-utils.workspace = true

334
litellm-rust/crates/python-bridge/src/cache/config.rs vendored
View file

@ -1,11 +1,13 @@
use std::{path::PathBuf, time::Duration};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::CacheType;
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
exceptions::{PyAttributeError, PyTypeError, PyValueError},
prelude::*,
types::{PyAny, PyDict, PyList, PyString},
types::{PyAny, PyBool, PyDict, PyList, PyString},
};
use super::{native::NativeResponseCache, request::duration};
@ -105,6 +107,7 @@ const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
S3(Box<S3CacheConfig>),
Gcs(GcsCacheConfig),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
@ -122,6 +125,9 @@ pub(super) enum UnsupportedCacheConfig {
RedisCredentials,
RedisConnection,
RedisOption,
S3Client,
S3Credentials,
S3Option,
GcsBucket,
DiskStore,
}
@ -134,6 +140,9 @@ impl UnsupportedCacheConfig {
Self::RedisCredentials => "native Redis credentials require Python",
Self::RedisConnection => "native Redis connection type is not implemented",
Self::RedisOption => "native Redis configuration requires Python",
Self::S3Client => "native S3 client type is not implemented",
Self::S3Credentials => "native S3 credentials require Python",
Self::S3Option => "native S3 configuration requires Python",
Self::GcsBucket => "native GCS cache requires a configured bucket name",
Self::DiskStore => "native disk cache requires the built-in diskcache store",
}
@ -178,6 +187,13 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::S3) => match project_s3(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::S3(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Gcs) => match project_gcs(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
@ -199,10 +215,7 @@ impl NativeCacheConfig {
}))
}),
Some(
CacheType::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::QdrantSemantic,
CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::QdrantSemantic,
)
| None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
@ -214,6 +227,7 @@ impl NativeCacheConfig {
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::Disk(_)
| CacheBackendConfig::AzureBlob(_)
| CacheBackendConfig::Gcs(_) => None,
@ -243,6 +257,30 @@ impl NativeCacheConfig {
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
CacheBackendConfig::S3(_) if service.kind() != "s3" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => {
Some("facade and native backend buckets must match")
}
CacheBackendConfig::S3(config)
if service.key_prefix() != Some(config.key_prefix.as_str()) =>
{
Some("facade and native backend key prefixes must match")
}
CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => {
Some("facade and native backend regions must match")
}
CacheBackendConfig::S3(config)
if service.endpoint()
!= config
.endpoint
.as_ref()
.map(|endpoint| endpoint.url.as_str()) =>
{
Some("facade and native backend endpoints must match")
}
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => {
Some("facade and native backend types must match")
}
@ -445,6 +483,77 @@ fn project_redis(
}))
}
#[inline(never)]
fn project_s3(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<S3CacheConfig, UnsupportedCacheConfig>> {
let client = backend.getattr("s3_client")?;
if !instance_class_is(&client, "botocore.client", "S3")? {
return Ok(Err(UnsupportedCacheConfig::S3Client));
}
let meta = client.getattr("meta")?;
let Some(region) = optional_string(meta.getattr("region_name")?)? else {
return Ok(Err(UnsupportedCacheConfig::S3Option));
};
let Some(endpoint_url) = optional_string(meta.getattr("endpoint_url")?)? else {
return Ok(Err(UnsupportedCacheConfig::S3Option));
};
let client_config = meta.getattr("config")?;
for name in ["s3", "proxies", "client_cert"] {
if optional_attribute(&client_config, name)?.is_some_and(|value| !value.is_none()) {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
}
let signature = match optional_attribute(&client_config, "signature_version")? {
Some(value) => value.extract::<Option<String>>()?,
None => None,
};
if signature.as_deref() != Some("s3v4") {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
let insecure = endpoint_url.starts_with("http://");
let verify = optional_attribute_chain(&client, &["_endpoint", "http_session", "_verify"])?;
let verified = verify
.and_then(|value| value.cast::<PyBool>().ok().map(|value| value.is_true()))
.unwrap_or(false);
if !verified && !insecure {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
let credentials = optional_attribute_chain(&client, &["_request_signer", "_credentials"])?
.ok_or(UnsupportedCacheConfig::S3Credentials);
let credentials = match credentials {
Ok(credentials) if !credentials.is_none() => credentials,
_ => return Ok(Err(UnsupportedCacheConfig::S3Credentials)),
};
let auth = if credentials.getattr("method")?.extract::<String>()?.as_str() == "explicit" {
AwsAuthConfig {
access_key_id: credentials
.getattr("access_key")?
.extract::<Option<String>>()?,
secret_access_key: credentials
.getattr("secret_key")?
.extract::<Option<String>>()?,
session_token: credentials.getattr("token")?.extract::<Option<String>>()?,
region_name: Some(region.clone()),
..Default::default()
}
} else {
AwsAuthConfig {
region_name: Some(region.clone()),
..Default::default()
}
};
let default_endpoint = endpoint_url == format!("https://s3.{region}.amazonaws.com")
|| (region == "us-east-1" && endpoint_url == "https://s3.amazonaws.com");
Ok(Ok(S3CacheConfig {
bucket: backend.getattr("bucket_name")?.extract::<String>()?,
key_prefix: backend.getattr("key_prefix")?.extract::<String>()?,
region,
endpoint: (!default_endpoint).then_some(S3Endpoint { url: endpoint_url }),
auth,
}))
}
#[inline(never)]
fn project_standalone_client<'py>(
client: &Bound<'py, PyAny>,
@ -648,6 +757,31 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult<O
}
}
#[inline(never)]
fn optional_attribute<'py>(
value: &Bound<'py, PyAny>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match value.getattr(name) {
Ok(value) => Ok(Some(value)),
Err(error) if error.is_instance_of::<PyAttributeError>(value.py()) => Ok(None),
Err(error) => Err(error),
}
}
#[inline(never)]
fn optional_attribute_chain<'py>(
value: &Bound<'py, PyAny>,
names: &[&str],
) -> PyResult<Option<Bound<'py, PyAny>>> {
names
.iter()
.try_fold(Some(value.clone()), |current, name| match current {
Some(current) => optional_attribute(&current, name),
None => Ok(None),
})
}
#[inline(never)]
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
Ok(value
@ -735,13 +869,16 @@ mod tests {
use pyo3::{prelude::*, types::PyDict};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::run_sync_value;
use super::{
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
};
use crate::cache::native::NativeResponseCache;
use litellm_cache_redis::{RedisNode, RedisTopology};
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
facade(
@ -1019,6 +1156,189 @@ mod tests {
});
}
fn s3_facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> {
let locals = PyDict::new(py);
py.run(
&CString::new(format!(
"from types import SimpleNamespace\n\
S3Client = type('S3', (), {{'__module__': 'botocore.client'}})\n\
client = S3Client()\n\
client.meta = SimpleNamespace(region_name='us-east-1', endpoint_url='https://example.test', config=SimpleNamespace(s3=None, proxies=None, client_cert=None, signature_version='s3v4'))\n\
client._endpoint = SimpleNamespace(http_session=SimpleNamespace(_verify=True))\n\
client._request_signer = SimpleNamespace(_credentials=SimpleNamespace(method='explicit', access_key='key', secret_key='secret', token='token'))\n\
backend = SimpleNamespace(bucket_name='bucket', key_prefix='team/', s3_client=client)\n\
facade = SimpleNamespace(type='s3', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)\n\
{body}"
))
.unwrap(),
None,
Some(&locals),
)
.unwrap();
locals.get_item("facade").unwrap().unwrap()
}
#[test]
fn projects_s3_configuration_with_explicit_credentials_and_custom_endpoint() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(py, "");
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert_eq!(s3.bucket, "bucket");
assert_eq!(s3.key_prefix, "team/");
assert_eq!(s3.region, "us-east-1");
assert_eq!(
s3.endpoint.map(|endpoint| endpoint.url).as_deref(),
Some("https://example.test")
);
assert_eq!(s3.auth.access_key_id.as_deref(), Some("key"));
assert_eq!(s3.auth.secret_access_key.as_deref(), Some("secret"));
assert_eq!(s3.auth.session_token.as_deref(), Some("token"));
assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1"));
});
}
#[test]
fn default_s3_endpoint_projects_no_custom_endpoint() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(
py,
"facade.cache.s3_client.meta.endpoint_url = 'https://s3.us-east-1.amazonaws.com'",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert!(s3.endpoint.is_none());
});
}
#[test]
fn non_sigv4_proxies_and_disabled_verification_stay_on_python() {
Python::initialize();
Python::attach(|py| {
for (body, message) in [
(
"facade.cache.s3_client.meta.config.signature_version = 's3'",
"native S3 configuration requires Python",
),
(
"facade.cache.s3_client.meta.config.proxies = {'https': 'proxy'}",
"native S3 configuration requires Python",
),
(
"facade.cache.s3_client._endpoint.http_session._verify = False",
"native S3 configuration requires Python",
),
(
"del facade.cache.s3_client._endpoint.http_session._verify",
"native S3 configuration requires Python",
),
] {
let facade = s3_facade(py, body);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("{body} must stay on Python");
};
assert_eq!(reason.message(), message);
}
let facade = s3_facade(py, "facade.cache.s3_client = SimpleNamespace()");
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("non-botocore client must stay on Python");
};
assert_eq!(reason.message(), "native S3 client type is not implemented");
let facade = s3_facade(
py,
"facade.cache.s3_client._request_signer._credentials = None",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("missing credentials must stay on Python");
};
assert_eq!(reason.message(), "native S3 credentials require Python");
});
}
#[test]
fn non_explicit_s3_credentials_use_the_default_chain() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(
py,
"facade.cache.s3_client._request_signer._credentials = SimpleNamespace(method='sso', access_key=None, secret_key=None, token=None)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("default-chain credentials should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert_eq!(s3.auth.access_key_id, None);
assert_eq!(s3.auth.secret_access_key, None);
assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1"));
});
}
fn s3_service(py: Python<'_>, region: &str, endpoint: Option<&str>) -> NativeResponseCache {
let config = S3CacheConfig {
bucket: "bucket".to_string(),
key_prefix: "team/".to_string(),
region: region.to_string(),
endpoint: endpoint.map(|url| S3Endpoint {
url: url.to_string(),
}),
auth: AwsAuthConfig::default(),
};
run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) }).unwrap()
}
#[test]
fn s3_binding_rejects_region_and_endpoint_mismatches() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(py, "");
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", Some("https://example.test"))),
None
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-west-2", Some("https://example.test"))),
Some("facade and native backend regions must match")
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", Some("https://other.test"))),
Some("facade and native backend endpoints must match")
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", None)),
Some("facade and native backend endpoints must match")
);
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {
Python::initialize();

29
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored
View file

@ -34,11 +34,14 @@ struct RedisPoolGuard {
attributes: RedisPoolAttributes,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
struct DiskStoreGuard {
reference: Py<PyAny>,
directory: String,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
@ -50,6 +53,7 @@ enum ConnectionGuard {
None,
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
S3(S3ClientGuard),
}
struct RedisPoolAttributes {
pool: &'static str,
@ -68,7 +72,6 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
connection_class: "connection_pool_class",
max_connections: None,
};
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
@ -223,6 +226,22 @@ impl RedisPoolGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl DiskStoreGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let store = backend.getattr("disk_cache")?;
@ -242,7 +261,6 @@ impl DiskStoreGuard {
visit.call(&self.reference)
}
}
impl AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
@ -277,6 +295,7 @@ impl ConnectionGuard {
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?),
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
("s3", _) => Self::S3(S3ClientGuard::capture(backend)?),
_ => Self::None,
})
}
@ -286,6 +305,7 @@ impl ConnectionGuard {
Self::None => Ok(true),
Self::RedisPool(guard) => guard.matches(py, backend),
Self::AzureBlob(guard) => guard.matches(py, backend),
Self::S3(guard) => guard.matches(py, backend),
}
}
@ -294,10 +314,10 @@ impl ConnectionGuard {
Self::None => Ok(()),
Self::RedisPool(guard) => guard.traverse(visit),
Self::AzureBlob(guard) => guard.traverse(visit),
Self::S3(guard) => guard.traverse(visit),
}
}
}
impl FacadeGuard {
pub(super) fn capture(
py: Python<'_>,
@ -320,6 +340,7 @@ impl FacadeGuard {
"RedisClusterCache",
"redis",
),
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
("azure-blob", _) => (

37
litellm-rust/crates/python-bridge/src/cache/handle.rs vendored
View file

@ -1,4 +1,6 @@
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
@ -66,6 +68,40 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (bucket, *, region, endpoint_url=None, key_prefix="", access_key_id=None, secret_access_key=None, session_token=None))]
fn s3(
py: Python<'_>,
bucket: String,
region: String,
endpoint_url: Option<String>,
key_prefix: &str,
access_key_id: Option<String>,
secret_access_key: Option<String>,
session_token: Option<String>,
) -> PyResult<Self> {
let config = S3CacheConfig {
bucket,
key_prefix: key_prefix.to_string(),
region: region.clone(),
endpoint: endpoint_url.map(|url| S3Endpoint { url }),
auth: AwsAuthConfig {
access_key_id,
secret_access_key,
session_token,
region_name: Some(region),
..Default::default()
},
};
let service = run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) })?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (bucket_name, *, gcs_path=None, path_service_account=None, endpoint=None, token=None))]
fn gcs(
@ -117,7 +153,6 @@ impl CacheTestHandle {
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()

74
litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

@ -9,6 +9,7 @@ use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig};
use serde_json::Value;
#[derive(Clone)]
@ -18,6 +19,7 @@ pub(super) enum NativeResponseCache {
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
buffer: Option<Arc<WriteBuffer>>,
},
S3(Arc<ResponseCache<S3Cache<ResponseCacheCodec>>>),
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
@ -51,6 +53,14 @@ impl NativeResponseCache {
buffer: None,
})
}
pub async fn s3(config: S3CacheConfig) -> Self {
let runtime = tokio::runtime::Handle::current();
Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new(
config,
ResponseCacheCodec,
runtime,
)))))
}
pub fn disk(directory: &str) -> Result<Self, Error> {
let cache = DiskCache::open(directory, ResponseCacheCodec)?;
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
@ -87,7 +97,9 @@ impl NativeResponseCache {
cache.backend().account_url(),
cache.backend().container_name(),
)),
Self::Memory(_) | Self::Redis { .. } | Self::Disk(_) | Self::Gcs(_) => None,
Self::Memory(_) | Self::Redis { .. } | Self::S3(_) | Self::Disk(_) | Self::Gcs(_) => {
None
}
}
}
}
@ -97,6 +109,7 @@ impl NativeResponseCache {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::S3(_) => "s3",
Self::Gcs(_) => "gcs",
Self::Disk(_) => "disk",
Self::AzureBlob(_) => "azure-blob",
@ -107,17 +120,46 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
Self::S3(cache) => cache.default_ttl(),
Self::Gcs(cache) => cache.default_ttl(),
Self::Disk(cache) => cache.default_ttl(),
Self::AzureBlob(cache) => cache.default_ttl(),
}
}
pub fn bucket(&self) -> Option<&str> {
match self {
Self::S3(cache) => Some(cache.backend().bucket()),
_ => None,
}
}
pub fn key_prefix(&self) -> Option<&str> {
match self {
Self::S3(cache) => Some(cache.backend().key_prefix()),
_ => None,
}
}
pub fn region(&self) -> Option<&str> {
match self {
Self::S3(cache) => Some(cache.backend().region()),
_ => None,
}
}
pub fn endpoint(&self) -> Option<&str> {
match self {
Self::S3(cache) => cache.backend().endpoint(),
_ => None,
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
Self::Gcs(_) => None,
Self::S3(_) | Self::Gcs(_) => None,
}
}
@ -125,20 +167,29 @@ impl NativeResponseCache {
match self {
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
Self::S3(_) => None,
}
}
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
@ -155,7 +206,11 @@ impl NativeResponseCache {
pub fn directory(&self) -> Option<&Path> {
match self {
Self::Disk(cache) => Some(cache.backend().directory()),
Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
@ -167,6 +222,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::S3(cache) => cache.lookup(request, now),
Self::Gcs(cache) => cache.lookup(request, now),
Self::Disk(cache) => cache.lookup(request, now),
Self::AzureBlob(cache) => cache.lookup(request, now),
@ -182,6 +238,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
Self::S3(cache) => cache.store(request, response, now),
Self::Gcs(cache) => cache.store(request, response, now),
Self::Disk(cache) => cache.store(request, response, now),
Self::AzureBlob(cache) => cache.store(request, response, now),
@ -196,6 +253,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
Self::S3(cache) => cache.lookup_batch(requests, now),
Self::Gcs(cache) => cache.lookup_batch(requests, now),
Self::Disk(cache) => cache.lookup_batch(requests, now),
Self::AzureBlob(cache) => cache.lookup_batch(requests, now),
@ -210,6 +268,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
Self::S3(cache) => cache.async_lookup(request, now).await,
Self::Gcs(cache) => cache.async_lookup(request, now).await,
Self::Disk(cache) => cache.async_lookup(request, now).await,
Self::AzureBlob(cache) => cache.async_lookup(request, now).await,
@ -232,6 +291,7 @@ impl NativeResponseCache {
cache,
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
Self::S3(cache) => cache.async_store(request, response, now).await,
Self::Gcs(cache) => cache.async_store(request, response, now).await,
Self::Disk(cache) => cache.async_store(request, response, now).await,
Self::AzureBlob(cache) => cache.async_store(request, response, now).await,
@ -246,6 +306,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
Self::S3(cache) => cache.async_lookup_batch(requests, now).await,
Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await,
Self::Disk(cache) => cache.async_lookup_batch(requests, now).await,
Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await,
@ -260,6 +321,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
Self::S3(cache) => cache.async_store_batch(entries, now).await,
Self::Gcs(cache) => cache.async_store_batch(entries, now).await,
Self::Disk(cache) => cache.async_store_batch(entries, now).await,
Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await,
@ -275,6 +337,7 @@ impl NativeResponseCache {
}
cache.async_flush().await
}
Self::S3(cache) => cache.async_flush().await,
Self::Gcs(cache) => cache.async_flush().await,
Self::Disk(cache) => cache.async_flush().await,
Self::AzureBlob(cache) => cache.async_flush().await,
@ -285,6 +348,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::S3(cache) => cache.test_connection().await,
Self::Gcs(cache) => cache.test_connection().await,
Self::Disk(cache) => cache.test_connection().await,
Self::AzureBlob(cache) => cache.test_connection().await,

104
litellm/rust_bridge/_native.pyi
View file

@ -93,6 +93,110 @@ class ResponsesWebSocketConnection:
def recv_text(self) -> Future[str | None]: ...
def close(self) -> Future[None]: ...
@final
class _CacheTestBinding:
@property
def kind(self) -> str: ...
def lookup(
self,
request: object,
*,
callback_kwargs: Mapping[str, object] | Sequence[object] | None = None,
) -> object: ...
def store(
self,
request: object,
response: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> None: ...
def lookup_batch(
self,
requests: Sequence[object],
*,
callback_kwargs: Sequence[object] | None = None,
) -> object: ...
def async_lookup(
self,
request: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[object]: ...
def async_store(
self,
request: object,
response: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[None]: ...
def async_lookup_batch(
self,
requests: Sequence[object],
*,
callback_kwargs: Sequence[object] | None = None,
) -> Future[object]: ...
def async_store_batch(
self,
requests: Sequence[object],
responses: Sequence[object],
*,
callback_result: object = None,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[object]: ...
def async_flush(self) -> Future[None]: ...
def ping(self) -> Future[object]: ...
@final
class _CacheTestHandle:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
@staticmethod
def memory(
*,
capacity: int = 200,
ttl_seconds: float = 600.0,
max_entry_bytes: int = 1048576,
) -> _CacheTestHandle: ...
@staticmethod
def redis(
url: str,
*,
ttl_seconds: float = 60.0,
namespace: str | None = None,
startup_nodes: Sequence[tuple[str, int]] | None = None,
) -> _CacheTestHandle: ...
@staticmethod
def disk(directory: str) -> _CacheTestHandle: ...
@staticmethod
def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
@staticmethod
def gcs(
bucket_name: str,
*,
gcs_path: str | None = None,
path_service_account: str | None = None,
endpoint: str | None = None,
token: str | None = None,
) -> _CacheTestHandle: ...
@staticmethod
def s3(
bucket: str,
*,
region: str,
endpoint_url: str | None = None,
key_prefix: str = "",
access_key_id: str | None = None,
secret_access_key: str | None = None,
session_token: str | None = None,
) -> _CacheTestHandle: ...
@property
def backend(self) -> str: ...
def _bind_facade(self, facade: object) -> None: ...
@final
class _CacheTestResolver:
def __new__(cls, namespace: object) -> _CacheTestResolver: ...
def resolve(self) -> _CacheTestBinding: ...
@final
class TokenCounter:
def __new__(cls, tokenizer_json: str) -> TokenCounter: ...

112
tests/test_litellm_rust/support/s3_stub.py Normal file
View file

@ -0,0 +1,112 @@
"""In-process path-style S3 stub for native cache parity tests."""
import threading
from dataclasses import dataclass, field
from email.utils import parsedate_to_datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
from urllib.parse import unquote, urlsplit
_STORED_HEADERS: Final = (
"cache-control",
"content-type",
"content-language",
"content-disposition",
"expires",
)
@dataclass
class S3Object:
body: bytes
headers: dict[str, str] = field(default_factory=dict)
class S3Stub:
"""Minimal path-style S3 endpoint serving PUT and GET object operations."""
def __init__(self) -> None:
self._objects: dict[str, S3Object] = {}
stub: Final = self
class Handler(BaseHTTPRequestHandler):
def _key(self) -> str:
parts: Final = urlsplit(self.path).path.lstrip("/").split("/", 1)
return unquote(parts[1]) if len(parts) == 2 else ""
def _read_body(self) -> bytes:
transfer: Final = self.headers.get("transfer-encoding", "")
if "chunked" not in transfer:
return self.rfile.read(int(self.headers.get("content-length", 0)))
chunks: Final = bytearray()
while True:
size = int(self.rfile.readline().split(b";")[0].strip(), 16)
if size == 0:
while self.rfile.readline().strip():
pass
return bytes(chunks)
chunks.extend(self.rfile.read(size))
self.rfile.readline()
def do_PUT(self) -> None:
body: Final = self._read_body()
headers: Final = {name: self.headers[name] for name in _STORED_HEADERS if name in self.headers}
stub._objects = {**stub._objects, self._key(): S3Object(body=body, headers=headers)}
self.send_response(200)
self.send_header("ETag", '"stub"')
self.send_header("Content-Length", "0")
self.end_headers()
def do_HEAD(self) -> None:
self._object(send_body=False)
def do_GET(self) -> None:
self._object(send_body=True)
def _object(self, send_body: bool) -> None:
entry: Final = stub._objects.get(self._key())
if entry is None:
self.send_response(404)
self.send_header("Content-Type", "application/xml")
body: Final = b'<?xml version="1.0" encoding="UTF-8"?><Error><Code>NoSuchKey</Code></Error>'
self.send_header("Content-Length", str(len(body)))
self.end_headers()
if send_body:
self.wfile.write(body)
return
self.send_response(200)
for name, value in entry.headers.items():
self.send_header(name, value)
self.send_header("ETag", '"stub"')
self.send_header("Content-Length", str(len(entry.body)))
self.end_headers()
if send_body:
self.wfile.write(entry.body)
def log_message(self, format: str, *args: object) -> None:
pass
self._server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
self._worker: Final = threading.Thread(target=self._server.serve_forever, daemon=True)
self._worker.start()
@property
def url(self) -> str:
host, port = self._server.server_address[:2]
return f"http://{host}:{port}"
@property
def objects(self) -> dict[str, S3Object]:
return self._objects
def put_object(self, key: str, body: bytes, headers: dict[str, str] | None = None) -> None:
self._objects = {**self._objects, key: S3Object(body=body, headers=headers or {})}
def expires(self, key: str) -> object:
header: Final = self._objects[key].headers.get("expires")
return parsedate_to_datetime(header) if header else None
def close(self) -> None:
self._server.shutdown()
self._server.server_close()
self._worker.join(timeout=5)

209
tests/test_litellm_rust/test_cache.py
View file

@ -8,11 +8,15 @@ import time
import uuid
import weakref
from collections.abc import Generator
from datetime import datetime
from pathlib import Path
from types import SimpleNamespace
from typing import Final, Protocol, cast
from unittest.mock import Mock
from urllib.parse import urlparse
import boto3
import botocore.config
import diskcache
import fakeredis
import pytest
@ -22,14 +26,16 @@ from azure.storage.blob import ContainerClient
import litellm
from litellm.caching.azure_blob_cache import AzureBlobCache
from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache
from litellm.caching.gcs_cache import GCSCache
from litellm.caching.disk_cache import DiskCache
from litellm.caching.gcs_cache import GCSCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.caching.s3_cache import S3Cache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
from tests.test_litellm_rust.support.fake_gcs import FakeGcs
from tests.test_litellm_rust.support.isolation import rebound
from tests.test_litellm_rust.support.s3_stub import S3Stub
pytestmark: Final = pytest.mark.requires_rust_extension
@ -426,15 +432,20 @@ def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure
assert handle.backend == "azure-blob"
account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}")
with pytest.raises(TypeError, match="containers must match"):
_native._CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade(
azure_blob_facade
)
_native._CacheTestHandle.azure_blob(
account_url, f"{backend.container_client.container_name}-other"
)._bind_facade(azure_blob_facade)
handle._bind_facade(azure_blob_facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade))
native: Final = resolver.resolve()
assert native.kind == "native"
response: Final = {"choices": [{"text": "caf\u00e9 \u2603"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
response: Final = {
"choices": [{"text": "caf\u00e9 \u2603"}],
"usage": {"total_tokens": 3},
"flag": True,
"empty": None,
}
native.store({**request("sync"), "ttl_seconds": 0.001}, response)
native.store(request("sync"), {"choices": [{"text": "second"}]})
time.sleep(0.01)
@ -489,7 +500,9 @@ async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_pyt
await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2})
time.sleep(0.01)
assert await binding.async_lookup(request("async")) == {"value": 2}
assert await backend.async_get_cache("async") == json.loads(backend.container_client.download_blob("async").readall())
assert await backend.async_get_cache("async") == json.loads(
backend.container_client.download_blob("async").readall()
)
assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2}
await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}])
@ -534,6 +547,8 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None:
assert client.get("second") is not None
await facade.cache.disconnect()
client.close()
async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None:
disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path))
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}}
@ -642,6 +657,183 @@ async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path:
}
@pytest.fixture
def s3_stub() -> Generator[S3Stub]:
stub: Final = S3Stub()
try:
yield stub
finally:
stub.close()
def python_s3(url: str) -> S3Cache:
return S3Cache(
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None:
python_cache: Final = python_s3(s3_stub.url)
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}}
python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90)
python_cache.set_cache("plain", {"timestamp": time.time(), "response": response})
s3_stub.put_object("team/malformed", b"not a cache entry")
s3_stub.put_object(
"team/expired",
json.dumps({"timestamp": time.time(), "response": response}).encode(),
{"expires": "Thu, 01 Jan 1970 00:00:00 GMT"},
)
binding: Final = _native._CacheTestResolver(
SimpleNamespace(
cache=_native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
)
).resolve()
assert binding.lookup(request("sync:key")) == response
assert await binding.async_lookup(request("plain")) == response
assert binding.lookup(request("malformed")) is None
assert binding.lookup(request("expired")) is None
assert binding.lookup(request("absent")) is None
binding.store({**request("native:key"), "ttl_seconds": 90.0}, response)
await binding.async_store(request("no_ttl"), response)
stored: Final = s3_stub.objects["team/native/key"]
assert stored.headers["content-type"] == "application/json"
assert stored.headers["content-language"] == "en"
assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"'
assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90"
expires: Final = cast(datetime, s3_stub.expires("team/native/key"))
remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds()
assert 60 < remaining <= 91
no_ttl: Final = s3_stub.objects["team/no_ttl"]
assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000"
assert "expires" not in no_ttl.headers
assert python_cache.get_cache("native:key")["response"] == response
partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")])
assert partial == {"values": [response, None, None], "missing_indices": [1, 2]}
def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None:
facade: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
handle: Final = _native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
with pytest.raises(TypeError, match="buckets must match"):
_native._CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade)
with pytest.raises(TypeError, match="key prefixes must match"):
_native._CacheTestHandle.s3(
"cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/"
)._bind_facade(facade)
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
binding: Final = resolver.resolve()
assert binding.kind == "native"
handler: Final = Mock()
facade.cache.s3_client.meta.events.register("before-call.s3.*", handler)
binding.store(request("native"), {"answer": 1})
assert binding.lookup(request("native")) == {"answer": 1}
assert handler.call_count == 0
assert "team/native" in s3_stub.objects
with rebound(facade.cache, "bucket_name", "other"):
assert resolver.resolve().kind == "python_callback"
other_client: Final = boto3.client(
"s3",
region_name="us-east-1",
endpoint_url=s3_stub.url,
aws_access_key_id="key",
aws_secret_access_key="secret",
)
with rebound(facade.cache, "s3_client", other_client):
assert resolver.resolve().kind == "python_callback"
class CustomS3Cache(S3Cache):
pass
subclassed: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
subclassed.cache = CustomS3Cache(
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
with pytest.raises(TypeError):
handle._bind_facade(subclassed)
assert _native._CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback"
def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None:
handle: Final = _native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
unverified: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url="https://s3.example.test",
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
s3_verify=False,
)
with pytest.raises(TypeError, match="requires Python"):
handle._bind_facade(unverified)
proxied: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}),
)
with pytest.raises(TypeError, match="requires Python"):
handle._bind_facade(proxied)
async def test_gcs_reads_python_entries_and_writes_python_compatible_objects(
fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch
) -> None:
@ -777,6 +969,7 @@ async def test_gcs_facade_binds_only_exact_matching_configuration(
assert resolver.resolve().kind == "python_callback"
with rebound(facade.cache, "path_service_account", "sa.json"):
assert resolver.resolve().kind == "python_callback"
def no_get_cache(*args: object, **kwargs: object) -> None:
return None
@ -911,7 +1104,9 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n
await binding.async_flush()
remaining: Final = tuple(sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)))
remaining: Final = tuple(
sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node))
)
assert remaining == (), remaining
assert client.get("unscoped") == b"stays"
client.delete("unscoped")