mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
feat(rust): add native S3 cache backend
Mirror the Redis vertical slice for S3Cache: a litellm-cache-s3 crate built on aws-sdk-s3 with path-style custom endpoints, python-identical put_object metadata (cache-control, expires, content headers), expires-aware get_object, and no-op flush/unsupported test_connection. Wire it through python-bridge config projection, facade guards, test handle, and binding dispatch, plus an in-process S3 stub and parity tests. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
parent
662e5b6e32
commit
788e24655a
19 changed files with 1465 additions and 21 deletions
165
litellm-rust/Cargo.lock
generated
165
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
@ -208,6 +214,7 @@ dependencies = [
|
|||
"aws-credential-types",
|
||||
"aws-sigv4",
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-eventstream",
|
||||
"aws-smithy-http",
|
||||
"aws-smithy-runtime",
|
||||
"aws-smithy-runtime-api",
|
||||
|
|
@ -216,7 +223,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",
|
||||
|
|
@ -250,6 +259,43 @@ dependencies = [
|
|||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-sdk-s3"
|
||||
version = "1.148.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8a61c1db3987ab6c8740fb87f248864687d50549401115244a84f02c3047b2d5"
|
||||
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"
|
||||
|
|
@ -294,7 +340,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",
|
||||
|
|
@ -310,6 +356,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",
|
||||
|
|
@ -337,10 +384,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",
|
||||
|
|
@ -353,6 +421,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",
|
||||
|
|
@ -532,6 +601,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"
|
||||
|
|
@ -927,6 +1008,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 = "crc32fast"
|
||||
version = "1.5.1"
|
||||
|
|
@ -1369,6 +1460,12 @@ version = "1.0.7"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
||||
|
||||
[[package]]
|
||||
name = "foldhash"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
|
||||
|
||||
[[package]]
|
||||
name = "form_urlencoded"
|
||||
version = "1.2.2"
|
||||
|
|
@ -1838,6 +1935,11 @@ name = "hashbrown"
|
|||
version = "0.17.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
|
||||
dependencies = [
|
||||
"allocator-api2",
|
||||
"equivalent",
|
||||
"foldhash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
|
|
@ -2502,6 +2604,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"
|
||||
|
|
@ -2664,11 +2781,13 @@ dependencies = [
|
|||
"criterion",
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-auth-aws",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-cache",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-response",
|
||||
"litellm-cache-s3",
|
||||
"litellm-callbacks-legacy-python",
|
||||
"litellm-core",
|
||||
"litellm-core-utils",
|
||||
|
|
@ -2846,6 +2965,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"
|
||||
|
|
@ -2868,6 +2996,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"
|
||||
|
|
@ -4297,6 +4435,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"
|
||||
|
|
@ -4413,6 +4562,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"
|
||||
|
|
@ -5010,7 +5165,7 @@ dependencies = [
|
|||
"rand 0.8.7",
|
||||
"rustls 0.23.42",
|
||||
"rustls-pki-types",
|
||||
"sha1",
|
||||
"sha1 0.10.7",
|
||||
"thiserror 1.0.69",
|
||||
"utf-8",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ litellm-core-utils = { path = "crates/core-utils" }
|
|||
litellm-cache = { path = "crates/cache" }
|
||||
litellm-cache-memory = { path = "crates/cache-memory" }
|
||||
litellm-cache-redis = { path = "crates/cache-redis" }
|
||||
litellm-cache-s3 = { path = "crates/cache-s3" }
|
||||
litellm-cache-response = { path = "crates/cache-response" }
|
||||
litellm-token-counter = { path = "crates/token-counter" }
|
||||
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
|
||||
|
|
|
|||
20
litellm-rust/crates/cache-s3/Cargo.toml
Normal file
20
litellm-rust/crates/cache-s3/Cargo.toml
Normal 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.141.0", 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"] }
|
||||
32
litellm-rust/crates/cache-s3/src/auth.rs
Normal file
32
litellm-rust/crates/cache-s3/src/auth.rs
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
|
||||
use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Credentials {
|
||||
config: AwsAuthConfig,
|
||||
}
|
||||
|
||||
impl Credentials {
|
||||
pub(crate) fn new(config: AwsAuthConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
}
|
||||
|
||||
impl ProvideCredentials for Credentials {
|
||||
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
|
||||
where
|
||||
Self: 'a,
|
||||
{
|
||||
future::ProvideCredentials::new(async {
|
||||
resolve_credentials(self.config.clone(), &|name| std::env::var(name).ok())
|
||||
.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()
|
||||
}
|
||||
}
|
||||
204
litellm-rust/crates/cache-s3/src/cache.rs
Normal file
204
litellm-rust/crates/cache-s3/src/cache.rs
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
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>,
|
||||
}
|
||||
|
||||
impl<C: CacheCodec> S3Cache<C> {
|
||||
pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self {
|
||||
let mut builder = aws_sdk_s3::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new(config.region))
|
||||
.credentials_provider(Credentials::new(config.auth))
|
||||
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
|
||||
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
|
||||
if let Some(endpoint) = config.endpoint {
|
||||
builder = builder.endpoint_url(endpoint.url).force_path_style(true);
|
||||
}
|
||||
Self {
|
||||
client: aws_sdk_s3::Client::from_conf(builder.build()),
|
||||
codec,
|
||||
runtime,
|
||||
bucket: config.bucket.into(),
|
||||
key_prefix: config.key_prefix.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn bucket(&self) -> &str {
|
||||
&self.bucket
|
||||
}
|
||||
|
||||
pub fn key_prefix(&self) -> &str {
|
||||
&self.key_prefix
|
||||
}
|
||||
|
||||
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 not_found = service.err().is_no_such_key()
|
||||
|| error
|
||||
.raw_response()
|
||||
.map(|response| response.status().as_u16())
|
||||
== Some(404);
|
||||
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> {
|
||||
None
|
||||
}
|
||||
|
||||
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
4
litellm-rust/crates/cache-s3/src/lib.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
mod auth;
|
||||
mod cache;
|
||||
|
||||
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};
|
||||
264
litellm-rust/crates/cache-s3/tests/cache.rs
Normal file
264
litellm-rust/crates/cache-s3/tests/cache.rs
Normal file
|
|
@ -0,0 +1,264 @@
|
|||
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/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("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!(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/cache/src/error.rs
vendored
2
litellm-rust/crates/cache/src/error.rs
vendored
|
|
@ -6,4 +6,6 @@ pub enum Error {
|
|||
InvalidEntry,
|
||||
#[error("flushing Redis requires an explicit namespace")]
|
||||
UnscopedFlush,
|
||||
#[error("cache operation is not supported")]
|
||||
UnsupportedOperation,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.1", optional = true }
|
||||
aws-smithy-types = { version = "1.6.1", optional = true }
|
||||
bytes = "1"
|
||||
futures-util.workspace = true
|
||||
|
|
|
|||
|
|
@ -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.1"
|
||||
aws-smithy-types = "1.6.1"
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -23,9 +23,11 @@ bytes.workspace = true
|
|||
litellm-cache.workspace = true
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
litellm-cache-s3.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
|
||||
|
|
|
|||
|
|
@ -1,10 +1,12 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache::CacheType;
|
||||
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
|
||||
use pyo3::{
|
||||
exceptions::{PyTypeError, PyValueError},
|
||||
exceptions::{PyAttributeError, PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
types::{PyAny, PyDict, PyString},
|
||||
types::{PyAny, PyBool, PyDict, PyString},
|
||||
};
|
||||
|
||||
use super::{native::NativeResponseCache, request::duration};
|
||||
|
|
@ -76,6 +78,7 @@ pub(super) struct RedisCacheConfig {
|
|||
pub(super) enum CacheBackendConfig {
|
||||
Memory(MemoryCacheConfig),
|
||||
Redis(Box<RedisCacheConfig>),
|
||||
S3(Box<S3CacheConfig>),
|
||||
}
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
|
|
@ -90,6 +93,9 @@ pub(super) enum UnsupportedCacheConfig {
|
|||
RedisCredentials,
|
||||
RedisConnection,
|
||||
RedisOption,
|
||||
S3Client,
|
||||
S3Credentials,
|
||||
S3Option,
|
||||
}
|
||||
|
||||
impl UnsupportedCacheConfig {
|
||||
|
|
@ -100,6 +106,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",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -142,10 +151,16 @@ 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::RedisSemantic
|
||||
| CacheType::ValkeySemantic
|
||||
| CacheType::S3
|
||||
| CacheType::Disk
|
||||
| CacheType::QdrantSemantic
|
||||
| CacheType::AzureBlob
|
||||
|
|
@ -159,10 +174,11 @@ impl NativeCacheConfig {
|
|||
|
||||
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
|
||||
if service.default_ttl()
|
||||
!= Some(match &self.backend {
|
||||
CacheBackendConfig::Memory(config) => config.default_ttl,
|
||||
CacheBackendConfig::Redis(config) => config.default_ttl,
|
||||
})
|
||||
!= match &self.backend {
|
||||
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
|
||||
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
|
||||
CacheBackendConfig::S3(_) => None,
|
||||
}
|
||||
{
|
||||
return Some("facade and native backend default TTLs must match");
|
||||
}
|
||||
|
|
@ -185,6 +201,18 @@ 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(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -299,6 +327,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_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
|
||||
Ok(RedisTlsConfig {
|
||||
|
|
@ -380,6 +479,33 @@ 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>>> {
|
||||
let mut current = value.clone();
|
||||
for name in names {
|
||||
match optional_attribute(¤t, name)? {
|
||||
Some(next) => current = next,
|
||||
None => return Ok(None),
|
||||
}
|
||||
}
|
||||
Ok(Some(current))
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
|
||||
Ok(value
|
||||
|
|
@ -591,4 +717,145 @@ mod tests {
|
|||
assert_eq!(reason.message(), "native Redis credentials require Python");
|
||||
});
|
||||
}
|
||||
|
||||
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"));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -32,10 +32,15 @@ struct RedisPoolGuard {
|
|||
max_connections: usize,
|
||||
}
|
||||
|
||||
struct S3ClientGuard {
|
||||
reference: Py<PyAny>,
|
||||
}
|
||||
|
||||
pub(super) struct FacadeGuard {
|
||||
outer: ObjectGuard,
|
||||
backend: ObjectGuard,
|
||||
redis_pool: Option<RedisPoolGuard>,
|
||||
s3_client: Option<S3ClientGuard>,
|
||||
}
|
||||
|
||||
impl ObjectGuard {
|
||||
|
|
@ -176,6 +181,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 FacadeGuard {
|
||||
pub(super) fn capture(
|
||||
py: Python<'_>,
|
||||
|
|
@ -192,6 +213,7 @@ impl FacadeGuard {
|
|||
let (module, name, cache_kind) = match kind {
|
||||
"memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
|
||||
"redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"),
|
||||
"s3" => ("litellm.caching.s3_cache", "S3Cache", "s3"),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
let backend = facade.getattr("cache")?;
|
||||
|
|
@ -235,11 +257,16 @@ impl FacadeGuard {
|
|||
"max_size_per_item",
|
||||
"redis_kwargs",
|
||||
"redis_flush_size",
|
||||
"bucket_name",
|
||||
"key_prefix",
|
||||
],
|
||||
)?,
|
||||
redis_pool: (kind == "redis")
|
||||
.then(|| RedisPoolGuard::capture(&backend))
|
||||
.transpose()?,
|
||||
s3_client: (kind == "s3")
|
||||
.then(|| S3ClientGuard::capture(&backend))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -252,9 +279,14 @@ impl FacadeGuard {
|
|||
return Ok(false);
|
||||
}
|
||||
match &self.redis_pool {
|
||||
Some(guard) => guard.matches(py, &backend),
|
||||
None => Ok(true),
|
||||
Some(guard) if !guard.matches(py, &backend)? => return Ok(false),
|
||||
_ => {}
|
||||
}
|
||||
match &self.s3_client {
|
||||
Some(guard) if !guard.matches(py, &backend)? => return Ok(false),
|
||||
_ => {}
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
|
|
@ -263,6 +295,9 @@ impl FacadeGuard {
|
|||
if let Some(guard) = &self.redis_pool {
|
||||
guard.traverse(&visit)?;
|
||||
}
|
||||
if let Some(guard) = &self.s3_client {
|
||||
guard.traverse(&visit)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
use litellm_host_python::release_gil;
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
|
||||
use litellm_host_python::{release_gil, run_sync_value};
|
||||
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
|
||||
|
||||
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
|
||||
|
|
@ -51,6 +53,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(),
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn backend(&self) -> &'static str {
|
||||
self.service.kind()
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ mod resolver;
|
|||
|
||||
use litellm_cache::Error;
|
||||
use pyo3::{
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
};
|
||||
|
||||
|
|
@ -21,6 +21,7 @@ pub(crate) use self::{
|
|||
fn cache_error(error: Error) -> PyErr {
|
||||
match error {
|
||||
Error::InvalidEntry => PyValueError::new_err(error.to_string()),
|
||||
Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()),
|
||||
_ => PyRuntimeError::new_err(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ use litellm_cache_redis::RedisCache;
|
|||
use litellm_cache_response::{
|
||||
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
|
||||
};
|
||||
use litellm_cache_s3::{S3Cache, S3CacheConfig};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone)]
|
||||
|
|
@ -15,6 +16,7 @@ pub(super) enum NativeResponseCache {
|
|||
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
|
||||
buffer: Option<Arc<WriteBuffer>>,
|
||||
},
|
||||
S3(Arc<ResponseCache<S3Cache<ResponseCacheCodec>>>),
|
||||
}
|
||||
|
||||
impl NativeResponseCache {
|
||||
|
|
@ -43,6 +45,15 @@ 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,
|
||||
)))))
|
||||
}
|
||||
}
|
||||
|
||||
impl NativeResponseCache {
|
||||
|
|
@ -50,6 +61,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(_) => "memory",
|
||||
Self::Redis { .. } => "redis",
|
||||
Self::S3(_) => "s3",
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -57,6 +69,21 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => cache.default_ttl(),
|
||||
Self::Redis { cache, .. } => cache.default_ttl(),
|
||||
Self::S3(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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -64,20 +91,21 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(_) => None,
|
||||
Self::Redis { cache, .. } => cache.backend().namespace(),
|
||||
Self::S3(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn capacity(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
|
||||
Self::Redis { .. } => None,
|
||||
Self::Redis { .. } | Self::S3(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn max_entry_bytes(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.backend().max_entry_bytes(),
|
||||
Self::Redis { .. } => None,
|
||||
Self::Redis { .. } | Self::S3(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -87,7 +115,7 @@ impl NativeResponseCache {
|
|||
cache,
|
||||
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
|
||||
},
|
||||
memory => memory,
|
||||
cache => cache,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -99,6 +127,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),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -111,6 +140,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),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -122,6 +152,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),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -133,6 +164,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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -152,6 +184,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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -163,6 +196,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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -174,6 +208,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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -186,6 +221,7 @@ impl NativeResponseCache {
|
|||
}
|
||||
cache.async_flush().await
|
||||
}
|
||||
Self::S3(cache) => cache.async_flush().await,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -193,6 +229,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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -93,6 +93,96 @@ 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,
|
||||
) -> _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
112
tests/test_litellm_rust/support/s3_stub.py
Normal 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[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: # noqa: A002
|
||||
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[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)
|
||||
|
|
@ -6,10 +6,13 @@ import threading
|
|||
import time
|
||||
import weakref
|
||||
from collections.abc import Generator
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Final, Protocol, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import boto3
|
||||
import botocore.config
|
||||
import fakeredis
|
||||
import pytest
|
||||
import redis
|
||||
|
|
@ -17,9 +20,11 @@ import redis
|
|||
import litellm
|
||||
from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
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.isolation import rebound
|
||||
from tests.test_litellm_rust.support.s3_stub import S3Stub
|
||||
|
||||
pytestmark: Final = pytest.mark.requires_rust_extension
|
||||
|
||||
|
|
@ -393,3 +398,180 @@ 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()
|
||||
|
||||
|
||||
@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"
|
||||
|
||||
calls: Final = []
|
||||
facade.cache.s3_client.meta.events.register("before-call.s3.*", lambda **_kwargs: calls.append(1))
|
||||
binding.store(request("native"), {"answer": 1})
|
||||
assert binding.lookup(request("native")) == {"answer": 1}
|
||||
assert calls == []
|
||||
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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue