mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge origin/main into litellm_valkey_semantic_native_cache
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
commit
2eee880aba
23 changed files with 1911 additions and 159 deletions
176
litellm-rust/Cargo.lock
generated
176
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"
|
||||
|
|
@ -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-cache-valkey-semantic"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2842,6 +2950,7 @@ dependencies = [
|
|||
"criterion",
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-auth-aws",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-cache",
|
||||
"litellm-cache-azure-blob",
|
||||
|
|
@ -2850,6 +2959,7 @@ dependencies = [
|
|||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-response",
|
||||
"litellm-cache-s3",
|
||||
"litellm-cache-valkey-semantic",
|
||||
"litellm-callbacks-legacy-python",
|
||||
"litellm-core",
|
||||
|
|
@ -3094,6 +3204,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"
|
||||
|
|
@ -3116,6 +3235,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"
|
||||
|
|
@ -3758,7 +3887,7 @@ dependencies = [
|
|||
"quinn-udp",
|
||||
"rustc-hash",
|
||||
"rustls 0.23.42",
|
||||
"socket2 0.6.5",
|
||||
"socket2 0.5.10",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
|
|
@ -3797,9 +3926,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]]
|
||||
|
|
@ -4285,7 +4414,7 @@ dependencies = [
|
|||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4356,7 +4485,7 @@ dependencies = [
|
|||
"security-framework",
|
||||
"security-framework-sys",
|
||||
"webpki-root-certs",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4629,6 +4758,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"
|
||||
|
|
@ -4745,6 +4885,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"
|
||||
|
|
@ -4911,10 +5057,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]]
|
||||
|
|
@ -5378,7 +5524,7 @@ dependencies = [
|
|||
"rand 0.8.7",
|
||||
"rustls 0.23.42",
|
||||
"rustls-pki-types",
|
||||
"sha1",
|
||||
"sha1 0.10.7",
|
||||
"thiserror 1.0.69",
|
||||
"utf-8",
|
||||
]
|
||||
|
|
@ -5787,7 +5933,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]]
|
||||
|
|
|
|||
|
|
@ -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
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.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
101
litellm-rust/crates/cache-s3/src/auth.rs
Normal 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
220
litellm-rust/crates/cache-s3/src/cache.rs
Normal 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
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};
|
||||
278
litellm-rust/crates/cache-s3/tests/cache.rs
Normal file
278
litellm-rust/crates/cache-s3/tests/cache.rs
Normal 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})));
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -24,12 +24,14 @@ 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
|
||||
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
|
||||
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,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};
|
||||
|
|
@ -113,6 +115,7 @@ pub(super) struct AzureBlobCacheConfig {
|
|||
pub(super) enum CacheBackendConfig {
|
||||
Memory(MemoryCacheConfig),
|
||||
Redis(Box<RedisCacheConfig>),
|
||||
S3(Box<S3CacheConfig>),
|
||||
Gcs(GcsCacheConfig),
|
||||
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
|
||||
Disk(DiskCacheConfig),
|
||||
|
|
@ -131,6 +134,9 @@ pub(super) enum UnsupportedCacheConfig {
|
|||
RedisCredentials,
|
||||
RedisConnection,
|
||||
RedisOption,
|
||||
S3Client,
|
||||
S3Credentials,
|
||||
S3Option,
|
||||
GcsBucket,
|
||||
DiskStore,
|
||||
}
|
||||
|
|
@ -143,6 +149,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",
|
||||
}
|
||||
|
|
@ -187,6 +196,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,
|
||||
|
|
@ -214,11 +230,9 @@ impl NativeCacheConfig {
|
|||
backend: CacheBackendConfig::AzureBlob(backend),
|
||||
}))
|
||||
}),
|
||||
Some(CacheType::RedisSemantic | CacheType::S3 | CacheType::QdrantSemantic) | None => {
|
||||
Ok(CacheConfigProjection::Unsupported(
|
||||
UnsupportedCacheConfig::Backend,
|
||||
))
|
||||
}
|
||||
Some(CacheType::RedisSemantic | CacheType::QdrantSemantic) | None => Ok(
|
||||
CacheConfigProjection::Unsupported(UnsupportedCacheConfig::Backend),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -226,6 +240,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::ValkeySemantic(_) => Some(Duration::ZERO),
|
||||
CacheBackendConfig::Disk(_)
|
||||
| CacheBackendConfig::AzureBlob(_)
|
||||
|
|
@ -258,6 +273,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")
|
||||
}
|
||||
|
|
@ -470,6 +509,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>,
|
||||
|
|
@ -738,6 +848,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(¤t, name),
|
||||
None => Ok(None),
|
||||
})
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
|
||||
Ok(value
|
||||
|
|
|
|||
|
|
@ -47,10 +47,15 @@ struct AzureBlobClientGuard {
|
|||
container_name: String,
|
||||
}
|
||||
|
||||
struct S3ClientGuard {
|
||||
reference: Py<PyAny>,
|
||||
}
|
||||
|
||||
enum ConnectionGuard {
|
||||
None,
|
||||
RedisPool(RedisPoolGuard),
|
||||
AzureBlob(AzureBlobClientGuard),
|
||||
S3(S3ClientGuard),
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
|
|
@ -281,6 +286,22 @@ impl AzureBlobClientGuard {
|
|||
}
|
||||
}
|
||||
|
||||
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 ConnectionGuard {
|
||||
fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
Ok(match (kind, cluster) {
|
||||
|
|
@ -301,6 +322,7 @@ impl ConnectionGuard {
|
|||
)?),
|
||||
("disk", _) => Self::None,
|
||||
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
|
||||
("s3", _) => Self::S3(S3ClientGuard::capture(backend)?),
|
||||
_ => Self::None,
|
||||
})
|
||||
}
|
||||
|
|
@ -310,6 +332,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),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -318,6 +341,7 @@ impl ConnectionGuard {
|
|||
Self::None => Ok(()),
|
||||
Self::RedisPool(guard) => guard.traverse(visit),
|
||||
Self::AzureBlob(guard) => guard.traverse(visit),
|
||||
Self::S3(guard) => guard.traverse(visit),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -356,6 +380,7 @@ impl FacadeGuard {
|
|||
"AzureBlobCache",
|
||||
"azure-blob",
|
||||
),
|
||||
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
let backend = facade.getattr("cache")?;
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
|
||||
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::*};
|
||||
|
||||
|
|
@ -68,6 +70,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(
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ use litellm_cache_response::{
|
|||
CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec,
|
||||
ResponseCacheRequest, WriteBuffer,
|
||||
};
|
||||
use litellm_cache_s3::{S3Cache, S3CacheConfig};
|
||||
use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig};
|
||||
use pyo3::prelude::*;
|
||||
use serde_json::Value;
|
||||
|
|
@ -78,6 +79,7 @@ pub(super) enum NativeResponseCache {
|
|||
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
|
||||
buffer: Option<Arc<WriteBuffer>>,
|
||||
},
|
||||
S3(Arc<ResponseCache<S3Cache<ResponseCacheCodec>>>),
|
||||
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
|
||||
ValkeySemantic {
|
||||
cache: Arc<ResponseCache<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>>,
|
||||
|
|
@ -117,6 +119,15 @@ impl NativeResponseCache {
|
|||
})
|
||||
}
|
||||
|
||||
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 valkey_semantic(
|
||||
url: &str,
|
||||
similarity_threshold: f64,
|
||||
|
|
@ -177,6 +188,7 @@ impl NativeResponseCache {
|
|||
)),
|
||||
Self::Memory(_)
|
||||
| Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::Disk(_)
|
||||
| Self::Gcs(_) => None,
|
||||
|
|
@ -237,6 +249,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(_) => "memory",
|
||||
Self::Redis { .. } => "redis",
|
||||
Self::S3(_) => "s3",
|
||||
Self::Gcs(_) => "gcs",
|
||||
Self::ValkeySemantic { .. } => "valkey-semantic",
|
||||
Self::Disk(_) => "disk",
|
||||
|
|
@ -248,6 +261,7 @@ 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::ValkeySemantic { cache, .. } => cache.default_ttl(),
|
||||
Self::Disk(cache) => cache.default_ttl(),
|
||||
|
|
@ -255,9 +269,38 @@ impl NativeResponseCache {
|
|||
}
|
||||
}
|
||||
|
||||
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::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
|
|
@ -269,6 +312,7 @@ impl NativeResponseCache {
|
|||
pub fn topology(&self) -> Option<&RedisTopology> {
|
||||
match self {
|
||||
Self::Memory(_)
|
||||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
|
|
@ -281,6 +325,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
|
||||
Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
|
|
@ -292,6 +337,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => cache.backend().max_entry_bytes(),
|
||||
Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
|
|
@ -304,6 +350,7 @@ impl NativeResponseCache {
|
|||
Self::Disk(cache) => Some(cache.backend().directory()),
|
||||
Self::Memory(_)
|
||||
| Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
|
|
@ -324,6 +371,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => cache.lookup(&Self::exact(request), now),
|
||||
Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now),
|
||||
Self::S3(cache) => cache.lookup(&Self::exact(request), now),
|
||||
Self::ValkeySemantic { cache, scope, .. } => {
|
||||
cache.lookup(&Self::semantic(request, scope), now)
|
||||
}
|
||||
|
|
@ -342,6 +390,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => cache.store(&Self::exact(request), response, now),
|
||||
Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now),
|
||||
Self::S3(cache) => cache.store(&Self::exact(request), response, now),
|
||||
Self::ValkeySemantic { cache, scope, .. } => {
|
||||
cache.store(&Self::semantic(request, scope), response, now)
|
||||
}
|
||||
|
|
@ -365,6 +414,9 @@ impl NativeResponseCache {
|
|||
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
|
||||
cache.lookup_batch(&requests, now)
|
||||
}
|
||||
Self::S3(cache) => {
|
||||
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
|
||||
}
|
||||
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
|
||||
Self::Gcs(cache) => {
|
||||
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
|
||||
|
|
@ -386,6 +438,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await,
|
||||
Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await,
|
||||
Self::S3(cache) => cache.async_lookup(&Self::exact(request), now).await,
|
||||
Self::ValkeySemantic { cache, scope, .. } => {
|
||||
cache
|
||||
.async_lookup(&Self::semantic(request, scope), now)
|
||||
|
|
@ -405,6 +458,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(_)
|
||||
| Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => {
|
||||
|
|
@ -458,6 +512,11 @@ impl NativeResponseCache {
|
|||
.async_store(cache, &Self::exact(request), response, now)
|
||||
.await
|
||||
}
|
||||
Self::S3(cache) => {
|
||||
cache
|
||||
.async_store(&Self::exact(request), response, now)
|
||||
.await
|
||||
}
|
||||
Self::ValkeySemantic { cache, scope, .. } => {
|
||||
cache
|
||||
.async_store(&Self::semantic(request, scope), response, now)
|
||||
|
|
@ -490,6 +549,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(_)
|
||||
| Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => {
|
||||
|
|
@ -534,6 +594,11 @@ impl NativeResponseCache {
|
|||
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
|
||||
cache.async_lookup_batch(&requests, now).await
|
||||
}
|
||||
Self::S3(cache) => {
|
||||
cache
|
||||
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
|
||||
.await
|
||||
}
|
||||
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
|
||||
Self::Gcs(cache) => {
|
||||
cache
|
||||
|
|
@ -573,6 +638,13 @@ impl NativeResponseCache {
|
|||
.collect();
|
||||
cache.async_store_batch(entries, now).await
|
||||
}
|
||||
Self::S3(cache) => {
|
||||
let entries = entries
|
||||
.into_iter()
|
||||
.map(|(request, value)| (Self::exact(&request), value))
|
||||
.collect();
|
||||
cache.async_store_batch(entries, now).await
|
||||
}
|
||||
Self::ValkeySemantic { cache, scope, .. } => {
|
||||
let entries = entries
|
||||
.into_iter()
|
||||
|
|
@ -612,6 +684,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(_)
|
||||
| Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => {
|
||||
|
|
@ -657,6 +730,7 @@ impl NativeResponseCache {
|
|||
}
|
||||
cache.async_flush().await
|
||||
}
|
||||
Self::S3(cache) => cache.async_flush().await,
|
||||
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
|
||||
Self::Gcs(cache) => cache.async_flush().await,
|
||||
Self::Disk(cache) => cache.async_flush().await,
|
||||
|
|
@ -668,6 +742,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::ValkeySemantic { cache, .. } => cache.test_connection().await,
|
||||
Self::Gcs(cache) => cache.test_connection().await,
|
||||
Self::Disk(cache) => cache.test_connection().await,
|
||||
|
|
|
|||
|
|
@ -43037,21 +43037,21 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/deepseek/deepseek-v4-pro": {
|
||||
"input_cost_per_token": 8.92272e-07,
|
||||
"input_cost_per_token": 8.87226e-07,
|
||||
"input_cost_per_token_cache_hit": 4.4e-08,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 384000,
|
||||
"max_tokens": 384000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.784544e-06,
|
||||
"output_cost_per_token": 1.774452e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 7.4356e-08,
|
||||
"cache_read_input_token_cost": 7.39355e-08,
|
||||
"supports_audio_input": false,
|
||||
"supports_pdf_input": false,
|
||||
"supports_vision": false,
|
||||
|
|
@ -73794,6 +73794,7 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/bytedance-seed/seed-1.6": {
|
||||
"deprecation_date": "2026-11-11",
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
"input_cost_per_token_above_128k_tokens": 5e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
|
|
@ -73815,6 +73816,7 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/bytedance-seed/seed-1.6-flash": {
|
||||
"deprecation_date": "2026-11-11",
|
||||
"input_cost_per_token": 7.5e-08,
|
||||
"input_cost_per_token_above_128k_tokens": 1e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
|
|
@ -73855,6 +73857,7 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/bytedance-seed/seed-2.0-code": {
|
||||
"deprecation_date": "2026-11-11",
|
||||
"input_cost_per_token": 5e-07,
|
||||
"input_cost_per_token_above_128k_tokens": 1e-06,
|
||||
"litellm_provider": "openrouter",
|
||||
|
|
|
|||
|
|
@ -1078,6 +1078,9 @@ if MCP_AVAILABLE:
|
|||
return {"servers": registry_servers}
|
||||
|
||||
## FastAPI Routes
|
||||
def _mcp_server_display_order(server: LiteLLM_MCPServerTable) -> tuple[str, str]:
|
||||
return ((server.server_name or server.alias or server.server_id).lower(), server.server_id)
|
||||
|
||||
def _get_user_mcp_management_mode() -> UserMCPManagementMode:
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings as proxy_general_settings,
|
||||
|
|
@ -1228,10 +1231,12 @@ if MCP_AVAILABLE:
|
|||
detail="You do not have permission to view MCP servers for this team.",
|
||||
)
|
||||
|
||||
redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id)
|
||||
redacted_mcp_servers = sorted(
|
||||
await _get_team_scoped_mcp_server_list(sanitized_team_id), key=_mcp_server_display_order
|
||||
)
|
||||
else:
|
||||
servers: Final = await _resolve_accessible_mcp_servers(user_api_key_dict)
|
||||
redacted_mcp_servers = _redact_mcp_credentials_list(servers)
|
||||
redacted_mcp_servers = sorted(_redact_mcp_credentials_list(servers), key=_mcp_server_display_order)
|
||||
|
||||
if connected_app_view is True and is_ui_session_credential(user_api_key_dict):
|
||||
reachable_ids: Final = await _connected_app_reachable_server_ids(user_api_key_dict)
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -43037,21 +43037,21 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/deepseek/deepseek-v4-pro": {
|
||||
"input_cost_per_token": 8.92272e-07,
|
||||
"input_cost_per_token": 8.87226e-07,
|
||||
"input_cost_per_token_cache_hit": 4.4e-08,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 384000,
|
||||
"max_tokens": 384000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.784544e-06,
|
||||
"output_cost_per_token": 1.774452e-06,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 7.4356e-08,
|
||||
"cache_read_input_token_cost": 7.39355e-08,
|
||||
"supports_audio_input": false,
|
||||
"supports_pdf_input": false,
|
||||
"supports_vision": false,
|
||||
|
|
@ -73794,6 +73794,7 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/bytedance-seed/seed-1.6": {
|
||||
"deprecation_date": "2026-11-11",
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
"input_cost_per_token_above_128k_tokens": 5e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
|
|
@ -73815,6 +73816,7 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/bytedance-seed/seed-1.6-flash": {
|
||||
"deprecation_date": "2026-11-11",
|
||||
"input_cost_per_token": 7.5e-08,
|
||||
"input_cost_per_token_above_128k_tokens": 1e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
|
|
@ -73855,6 +73857,7 @@
|
|||
"supports_web_search": false
|
||||
},
|
||||
"openrouter/bytedance-seed/seed-2.0-code": {
|
||||
"deprecation_date": "2026-11-11",
|
||||
"input_cost_per_token": 5e-07,
|
||||
"input_cost_per_token_above_128k_tokens": 1e-06,
|
||||
"litellm_provider": "openrouter",
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import logging
|
|||
from contextlib import ExitStack
|
||||
from datetime import datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import List, Optional, cast
|
||||
from typing import Final, List, Optional, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -1614,6 +1614,60 @@ class TestTeamScopedMCPServerAccess:
|
|||
result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id")
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
class TestFetchAllMCPServersOrdering:
|
||||
def test_display_order_is_case_insensitive_name_then_id(self) -> None:
|
||||
servers: Final = (
|
||||
LiteLLM_MCPServerTable(server_id="s-2", server_name="GitHub", alias="aaa", transport=MCPTransport.http),
|
||||
LiteLLM_MCPServerTable(server_id="s-1", alias="github", transport=MCPTransport.http),
|
||||
LiteLLM_MCPServerTable(server_id="s-0", server_name="Slack", alias="zzz", transport=MCPTransport.http),
|
||||
LiteLLM_MCPServerTable(server_id="confluence", server_name="", alias="", transport=MCPTransport.http),
|
||||
)
|
||||
|
||||
ordered: Final = sorted(servers, key=mgmt_endpoints._mcp_server_display_order)
|
||||
assert [s.server_id for s in ordered] == ["confluence", "s-1", "s-2", "s-0"]
|
||||
|
||||
@pytest.mark.parametrize("team_id", [None, "team-1"])
|
||||
@pytest.mark.parametrize("reverse", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_is_sorted_by_display_name_regardless_of_resolution_order(
|
||||
self, team_id: str | None, reverse: bool
|
||||
) -> None:
|
||||
mock_user_auth: Final = generate_mock_user_api_key_auth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="admin_user",
|
||||
)
|
||||
servers: Final = (
|
||||
generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"),
|
||||
generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"),
|
||||
generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"),
|
||||
)
|
||||
resolved: Final = list(reversed(servers) if reverse else servers)
|
||||
mock_manager: Final = MagicMock()
|
||||
mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved)
|
||||
with (
|
||||
patch( # test-quality-ok: the route reads a module-global manager with no injection seam
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
|
||||
mock_manager,
|
||||
),
|
||||
patch( # test-quality-ok: admin view is derived from module-global proxy settings
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
|
||||
return_value=True,
|
||||
),
|
||||
patch( # test-quality-ok: auth contexts need a live prisma client
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[mock_user_auth]),
|
||||
),
|
||||
patch( # test-quality-ok: isolate the route's ordering from team database resolution
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list",
|
||||
AsyncMock(return_value=resolved),
|
||||
),
|
||||
):
|
||||
result: Final = await mgmt_endpoints.fetch_all_mcp_servers(
|
||||
user_api_key_dict=mock_user_auth, team_id=team_id
|
||||
)
|
||||
assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restricted_virtual_key_cannot_use_team_id_filter(self):
|
||||
"""Restricted virtual keys must not bypass access limits via team_id."""
|
||||
|
|
|
|||
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 = {**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)
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@ import { render, waitFor, screen, act, within } from "@testing-library/react";
|
|||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import MCPServers from "./mcp_servers";
|
||||
import MCPServers, { compareServers, type SortKey } from "./mcp_servers";
|
||||
import type { MCPServer } from "@/components/mcp_tools/types";
|
||||
import * as networking from "@/components/networking";
|
||||
|
||||
// Mock the networking module
|
||||
|
|
@ -31,6 +32,98 @@ const createQueryClient = () =>
|
|||
},
|
||||
});
|
||||
|
||||
describe("compareServers", () => {
|
||||
const server = (server_id: string, name: string, created_at = ""): MCPServer => ({
|
||||
server_id,
|
||||
server_name: name,
|
||||
created_at,
|
||||
updated_at: created_at,
|
||||
created_by: "user",
|
||||
updated_by: "user",
|
||||
});
|
||||
|
||||
const shuffled = [server("c", "github"), server("a", "slack"), server("b", "Jira")];
|
||||
|
||||
it("orders servers without timestamps by name so config.yaml servers render in a stable order", () => {
|
||||
const byCreated = [...shuffled].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id);
|
||||
const byUpdated = [...shuffled].sort((a, b) => compareServers(a, b, "updated_desc")).map((s) => s.server_id);
|
||||
const byHealth = [...shuffled].sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id);
|
||||
|
||||
expect(byCreated).toEqual(["c", "b", "a"]);
|
||||
expect(byUpdated).toEqual(["c", "b", "a"]);
|
||||
expect(byHealth).toEqual(["c", "b", "a"]);
|
||||
});
|
||||
|
||||
it("keeps newest-first when timestamps differ", () => {
|
||||
const newest = server("new", "zzz", "2026-02-01T00:00:00Z");
|
||||
const oldest = server("old", "aaa", "2026-01-01T00:00:00Z");
|
||||
expect([oldest, newest].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id)).toEqual([
|
||||
"new",
|
||||
"old",
|
||||
]);
|
||||
});
|
||||
|
||||
it.each<SortKey>(["created_desc", "updated_desc", "name_asc", "health"])(
|
||||
"breaks equal timestamps and names by ID for %s regardless of input order",
|
||||
(sort) => {
|
||||
const servers = [
|
||||
server("b", "GitHub", "2026-01-01T00:00:00Z"),
|
||||
server("c", "Slack", "2026-01-01T00:00:00Z"),
|
||||
server("a", "github", "2026-01-01T00:00:00Z"),
|
||||
];
|
||||
for (const input of [servers, [...servers].reverse()]) {
|
||||
expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual(["a", "b", "c"]);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
it("uses the display name before alias, then falls back to alias and ID", () => {
|
||||
const servers: MCPServer[] = [
|
||||
{ ...server("s-slack", "Slack"), alias: "aaa" },
|
||||
{ ...server("s-github", ""), server_name: null, alias: "GitHub" },
|
||||
{ ...server("confluence", ""), alias: "" },
|
||||
];
|
||||
for (const input of [servers, [...servers].reverse()]) {
|
||||
expect([...input].sort((a, b) => compareServers(a, b, "name_asc")).map((s) => s.server_id)).toEqual([
|
||||
"confluence",
|
||||
"s-github",
|
||||
"s-slack",
|
||||
]);
|
||||
}
|
||||
});
|
||||
|
||||
it.each<SortKey>(["created_desc", "updated_desc", "health"])(
|
||||
"keeps timestamped servers before missing timestamps for %s",
|
||||
(sort) => {
|
||||
const servers = [
|
||||
server("config", "aaa"),
|
||||
server("older", "bbb", "2026-01-01T00:00:00Z"),
|
||||
server("newer", "zzz", "2026-02-01T00:00:00Z"),
|
||||
];
|
||||
for (const input of [servers, [...servers].reverse()]) {
|
||||
expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual([
|
||||
"newer",
|
||||
"older",
|
||||
"config",
|
||||
]);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
it("sorts health before recency and display name", () => {
|
||||
const servers: MCPServer[] = [
|
||||
{ ...server("healthy", "aaa", "2026-03-01T00:00:00Z"), status: "healthy" },
|
||||
{ ...server("unknown", "bbb", "2026-02-01T00:00:00Z"), status: "unknown" },
|
||||
{ ...server("unhealthy", "zzz", "2026-01-01T00:00:00Z"), status: "unhealthy" },
|
||||
];
|
||||
expect(servers.sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id)).toEqual([
|
||||
"unhealthy",
|
||||
"unknown",
|
||||
"healthy",
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("MCPServers", () => {
|
||||
const defaultProps = {
|
||||
accessToken: "123",
|
||||
|
|
@ -74,6 +167,134 @@ describe("MCPServers", () => {
|
|||
const myConnections = await screen.findByRole("link", { name: "My Connections" });
|
||||
expect(myConnections).toBeVisible();
|
||||
expect(myConnections).toHaveAttribute("href", "/ui/connect");
|
||||
for (const name of ["Semantic Filter", "Tool Search", "Network Settings", "Submitted MCPs"]) {
|
||||
const tab = screen.queryByRole("tab", { name });
|
||||
if (userRole === "Admin") {
|
||||
expect(tab).toBeVisible();
|
||||
} else {
|
||||
expect(tab).not.toBeInTheDocument();
|
||||
}
|
||||
}
|
||||
expect(
|
||||
screen.getByRole("button", {
|
||||
name: userRole === "Admin" ? "+ Add New MCP Server" : "+ Submit MCP Server",
|
||||
}),
|
||||
).toBeVisible();
|
||||
});
|
||||
|
||||
it.each(["cancel", "success", "failure", "unnamed"])("preserves delete confirmation on %s", async (outcome) => {
|
||||
const server: MCPServer = {
|
||||
created_at: "",
|
||||
updated_at: "",
|
||||
server_id: "delete-server",
|
||||
server_name: outcome === "unnamed" ? null : "Delete fixture",
|
||||
alias: "delete-alias",
|
||||
url: outcome === "unnamed" ? null : "https://example.com/mcp",
|
||||
created_by: "user",
|
||||
updated_by: "user",
|
||||
};
|
||||
vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]);
|
||||
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
|
||||
let finishDelete: () => void = () => {};
|
||||
vi.mocked(networking.deleteMCPServer).mockImplementation(
|
||||
() =>
|
||||
new Promise((resolve, reject) => {
|
||||
finishDelete = () => (outcome === "failure" ? reject(new Error("Delete failed")) : resolve(undefined));
|
||||
}),
|
||||
);
|
||||
render(
|
||||
<QueryClientProvider client={createQueryClient()}>
|
||||
<MCPServers {...defaultProps} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
await userEvent.click(await screen.findByRole("button", { name: "Server actions" }));
|
||||
await userEvent.click(await screen.findByRole("menuitem", { name: "Delete" }));
|
||||
const dialog = await screen.findByRole("alertdialog", { name: "Delete MCP Server?" });
|
||||
expect(within(dialog).getByText("delete-server")).toBeVisible();
|
||||
if (outcome === "unnamed") {
|
||||
expect(within(dialog).queryByText("Name")).not.toBeInTheDocument();
|
||||
expect(within(dialog).queryByText("URL")).not.toBeInTheDocument();
|
||||
} else {
|
||||
expect(within(dialog).getByText("Delete fixture")).toBeVisible();
|
||||
expect(within(dialog).getByText("https://example.com/mcp")).toBeVisible();
|
||||
}
|
||||
if (outcome === "cancel") {
|
||||
await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" }));
|
||||
expect(networking.deleteMCPServer).not.toHaveBeenCalled();
|
||||
} else {
|
||||
await userEvent.click(within(dialog).getByRole("button", { name: "Delete" }));
|
||||
expect(within(dialog).getByRole("button", { name: "Deleting..." })).toBeDisabled();
|
||||
expect(within(dialog).getByRole("button", { name: "Cancel" })).toBeDisabled();
|
||||
expect(networking.deleteMCPServer).toHaveBeenCalledWith("123", "delete-server");
|
||||
await act(async () => finishDelete());
|
||||
}
|
||||
await waitFor(() => expect(screen.queryByRole("alertdialog")).not.toBeInTheDocument());
|
||||
});
|
||||
|
||||
it("filters servers by access group", async () => {
|
||||
const server = { created_by: "user", updated_by: "user" };
|
||||
vi.mocked(networking.fetchMCPServers).mockResolvedValue([
|
||||
{
|
||||
...server,
|
||||
server_id: "string-group",
|
||||
server_name: "String group",
|
||||
alias: "string-alias",
|
||||
mcp_access_groups: ["shared"],
|
||||
},
|
||||
{
|
||||
...server,
|
||||
server_id: "legacy-group",
|
||||
server_name: "Legacy group",
|
||||
alias: "legacy-alias",
|
||||
mcp_access_groups: ["shared"],
|
||||
},
|
||||
{
|
||||
...server,
|
||||
server_id: "other-group",
|
||||
server_name: "Other group",
|
||||
alias: "other-alias",
|
||||
mcp_access_groups: ["different"],
|
||||
},
|
||||
]);
|
||||
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
|
||||
render(
|
||||
<QueryClientProvider client={createQueryClient()}>
|
||||
<MCPServers {...defaultProps} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
await screen.findByText("String group");
|
||||
await userEvent.click(screen.getByRole("combobox", { name: "Access Group" }));
|
||||
await userEvent.click(await screen.findByRole("option", { name: "shared" }));
|
||||
expect(screen.getByText("String group")).toBeVisible();
|
||||
expect(screen.getByText("Legacy group")).toBeVisible();
|
||||
expect(screen.queryByText("Other group")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it.each(["server_name", "alias", "url", "server_id"] as const)("searches by %s case-insensitively", async (field) => {
|
||||
const server: MCPServer = {
|
||||
created_at: "",
|
||||
updated_at: "",
|
||||
server_id: "search-server",
|
||||
server_name: "Search fixture",
|
||||
created_by: "user",
|
||||
updated_by: "user",
|
||||
[field]: "Needle",
|
||||
};
|
||||
vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]);
|
||||
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
|
||||
render(
|
||||
<QueryClientProvider client={createQueryClient()}>
|
||||
<MCPServers {...defaultProps} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
await screen.findByTestId("mcp-servers-grid");
|
||||
const search = screen.getByPlaceholderText("Search by name, alias, URL, or ID");
|
||||
await userEvent.type(search, " NEEDLE ");
|
||||
expect(screen.getByTestId("mcp-servers-grid")).toBeVisible();
|
||||
await userEvent.clear(search);
|
||||
await userEvent.type(search, "no-match");
|
||||
expect(screen.queryByTestId("mcp-servers-grid")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("No servers match the current filters or search.")).toBeVisible();
|
||||
});
|
||||
|
||||
it("should render mocked MCP servers data in the table", async () => {
|
||||
|
|
@ -316,9 +537,7 @@ describe("MCPServers", () => {
|
|||
expect(screen.getByText("Team B Server")).toBeInTheDocument();
|
||||
expect(screen.getByText("Team A Server 2")).toBeInTheDocument();
|
||||
|
||||
// Find the team select by its "Team" label, then the combobox it labels
|
||||
const teamLabel = screen.getByText("Team");
|
||||
const teamSelect = within(teamLabel.parentElement!).getByRole("combobox");
|
||||
const teamSelect = screen.getByRole("combobox", { name: "Team" });
|
||||
|
||||
await userEvent.click(teamSelect);
|
||||
|
||||
|
|
|
|||
|
|
@ -49,7 +49,7 @@ import { cn } from "@/lib/cva.config";
|
|||
import UserEnvVarsModal from "./UserEnvVarsModal";
|
||||
import { listMCPUserEnvVarStatus } from "@/components/networking";
|
||||
|
||||
type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health";
|
||||
export type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health";
|
||||
|
||||
const SORT_OPTIONS: { value: SortKey; label: string }[] = [
|
||||
{ value: "created_desc", label: "Recently created" },
|
||||
|
|
@ -64,32 +64,33 @@ const HEALTH_RANK: Record<string, number> = {
|
|||
healthy: 2,
|
||||
};
|
||||
|
||||
const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => {
|
||||
const compareByName = (a: MCPServer, b: MCPServer): number => {
|
||||
const nameA = (a.server_name || a.alias || a.server_id).toLowerCase();
|
||||
const nameB = (b.server_name || b.alias || b.server_id).toLowerCase();
|
||||
return nameA.localeCompare(nameB) || a.server_id.localeCompare(b.server_id);
|
||||
};
|
||||
|
||||
const compareByTimestampDesc = (a: string | null | undefined, b: string | null | undefined): number => {
|
||||
const ta = a ? new Date(a).getTime() : 0;
|
||||
const tb = b ? new Date(b).getTime() : 0;
|
||||
return tb - ta;
|
||||
};
|
||||
|
||||
export const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => {
|
||||
switch (sort) {
|
||||
case "name_asc": {
|
||||
const nameA = (a.server_name || a.alias || a.server_id).toLowerCase();
|
||||
const nameB = (b.server_name || b.alias || b.server_id).toLowerCase();
|
||||
return nameA.localeCompare(nameB);
|
||||
}
|
||||
case "updated_desc": {
|
||||
const ta = a.updated_at ? new Date(a.updated_at).getTime() : 0;
|
||||
const tb = b.updated_at ? new Date(b.updated_at).getTime() : 0;
|
||||
return tb - ta;
|
||||
}
|
||||
case "name_asc":
|
||||
return compareByName(a, b);
|
||||
case "updated_desc":
|
||||
return compareByTimestampDesc(a.updated_at, b.updated_at) || compareByName(a, b);
|
||||
case "health": {
|
||||
const ra = HEALTH_RANK[a.status ?? "unknown"] ?? 1;
|
||||
const rb = HEALTH_RANK[b.status ?? "unknown"] ?? 1;
|
||||
if (ra !== rb) return ra - rb;
|
||||
const ta = a.created_at ? new Date(a.created_at).getTime() : 0;
|
||||
const tb = b.created_at ? new Date(b.created_at).getTime() : 0;
|
||||
return tb - ta;
|
||||
return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b);
|
||||
}
|
||||
case "created_desc":
|
||||
default: {
|
||||
const ta = a.created_at ? new Date(a.created_at).getTime() : 0;
|
||||
const tb = b.created_at ? new Date(b.created_at).getTime() : 0;
|
||||
return tb - ta;
|
||||
}
|
||||
default:
|
||||
return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b);
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -112,6 +113,62 @@ const readToolsOAuthServerId = (): string | null => {
|
|||
}
|
||||
};
|
||||
|
||||
function DeleteServerDialog({
|
||||
open,
|
||||
onOpenChange,
|
||||
server,
|
||||
isDeleting,
|
||||
onConfirm,
|
||||
}: {
|
||||
open: boolean;
|
||||
onOpenChange: (open: boolean) => void;
|
||||
server: MCPServer | undefined;
|
||||
isDeleting: boolean;
|
||||
onConfirm: () => Promise<void>;
|
||||
}) {
|
||||
return (
|
||||
<AlertDialog open={open} onOpenChange={onOpenChange}>
|
||||
<AlertDialogContent>
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>Delete MCP Server?</AlertDialogTitle>
|
||||
</AlertDialogHeader>
|
||||
<div className="space-y-4">
|
||||
<p className="text-sm text-muted-foreground">
|
||||
This action is permanent and cannot be undone. All associated configurations will be removed.
|
||||
</p>
|
||||
|
||||
{server && (
|
||||
<dl className="mt-3 space-y-1 rounded-lg border border-border bg-muted p-4">
|
||||
{server.server_name && (
|
||||
<div className="flex gap-2">
|
||||
<dt className="text-sm text-muted-foreground">Name</dt>
|
||||
<dd className="text-sm font-semibold">{server.server_name}</dd>
|
||||
</div>
|
||||
)}
|
||||
<div className="flex gap-2">
|
||||
<dt className="text-sm text-muted-foreground">ID</dt>
|
||||
<dd className="font-mono text-xs">{server.server_id}</dd>
|
||||
</div>
|
||||
{server.url && (
|
||||
<div className="flex gap-2">
|
||||
<dt className="text-sm text-muted-foreground">URL</dt>
|
||||
<dd className="font-mono text-xs break-all">{server.url}</dd>
|
||||
</div>
|
||||
)}
|
||||
</dl>
|
||||
)}
|
||||
</div>
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel disabled={isDeleting}>Cancel</AlertDialogCancel>
|
||||
<Button variant="destructive" disabled={isDeleting} onClick={onConfirm}>
|
||||
{isDeleting ? "Deleting..." : "Delete"}
|
||||
</Button>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>
|
||||
);
|
||||
}
|
||||
|
||||
const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, isViewOnly = false }) => {
|
||||
const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers();
|
||||
|
||||
|
|
@ -298,16 +355,12 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
}
|
||||
if (group !== "all") {
|
||||
filtered = filtered.filter((server) =>
|
||||
server.mcp_access_groups?.some((g: any) => (typeof g === "string" ? g === group : g && g.name === group)),
|
||||
server.mcp_access_groups?.some((g: string | { name?: string } | null) =>
|
||||
typeof g === "string" ? g === group : g?.name === group,
|
||||
),
|
||||
);
|
||||
}
|
||||
const sorted = [...filtered].sort((a, b) => {
|
||||
if (!a.created_at && !b.created_at) return 0;
|
||||
if (!a.created_at) return 1;
|
||||
if (!b.created_at) return -1;
|
||||
return new Date(b.created_at).getTime() - new Date(a.created_at).getTime();
|
||||
});
|
||||
setFilteredServers(sorted);
|
||||
setFilteredServers(filtered);
|
||||
},
|
||||
[serversWithHealth],
|
||||
);
|
||||
|
|
@ -338,7 +391,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
const alias = (s.alias || "").toLowerCase();
|
||||
const url = (s.url || "").toLowerCase();
|
||||
const id = s.server_id.toLowerCase();
|
||||
return name.includes(q) || alias.includes(q) || url.includes(q) || id.includes(q);
|
||||
return [name, alias, url, id].some((value) => value.includes(q));
|
||||
})
|
||||
: filteredServers;
|
||||
return [...matches].sort((a, b) => compareServers(a, b, sortKey));
|
||||
|
|
@ -381,9 +434,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
};
|
||||
|
||||
// Find the server to delete from the servers list
|
||||
const serverToDelete = serverIdToDelete
|
||||
? (mcpServers || []).find((server) => server.server_id === serverIdToDelete)
|
||||
: null;
|
||||
const serverToDelete = mcpServers?.find((server) => server.server_id === serverIdToDelete);
|
||||
|
||||
const handleCreateSuccess = (newMcpServer: MCPServer) => {
|
||||
setFilteredServers((prev) => [...prev, newMcpServer]);
|
||||
|
|
@ -425,45 +476,13 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
return (
|
||||
<TooltipProvider>
|
||||
<div className="h-full w-full p-6">
|
||||
<AlertDialog open={isDeleteModalOpen} onOpenChange={(open) => !open && cancelDelete()}>
|
||||
<AlertDialogContent>
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>Delete MCP Server?</AlertDialogTitle>
|
||||
</AlertDialogHeader>
|
||||
<div className="space-y-4">
|
||||
<p className="text-sm text-muted-foreground">
|
||||
This action is permanent and cannot be undone. All associated configurations will be removed.
|
||||
</p>
|
||||
|
||||
{serverToDelete && (
|
||||
<dl className="mt-3 space-y-1 rounded-lg border border-border bg-muted p-4">
|
||||
{serverToDelete.server_name && (
|
||||
<div className="flex gap-2">
|
||||
<dt className="text-sm text-muted-foreground">Name</dt>
|
||||
<dd className="text-sm font-semibold">{serverToDelete.server_name}</dd>
|
||||
</div>
|
||||
)}
|
||||
<div className="flex gap-2">
|
||||
<dt className="text-sm text-muted-foreground">ID</dt>
|
||||
<dd className="font-mono text-xs">{serverToDelete.server_id}</dd>
|
||||
</div>
|
||||
{serverToDelete.url && (
|
||||
<div className="flex gap-2">
|
||||
<dt className="text-sm text-muted-foreground">URL</dt>
|
||||
<dd className="font-mono text-xs break-all">{serverToDelete.url}</dd>
|
||||
</div>
|
||||
)}
|
||||
</dl>
|
||||
)}
|
||||
</div>
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel disabled={isDeletingServer}>Cancel</AlertDialogCancel>
|
||||
<Button variant="destructive" disabled={isDeletingServer} onClick={confirmDelete}>
|
||||
{isDeletingServer ? "Deleting..." : "Delete"}
|
||||
</Button>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>
|
||||
<DeleteServerDialog
|
||||
open={isDeleteModalOpen}
|
||||
onOpenChange={(open) => !open && cancelDelete()}
|
||||
server={serverToDelete}
|
||||
isDeleting={isDeletingServer}
|
||||
onConfirm={confirmDelete}
|
||||
/>
|
||||
<CreateMCPServer
|
||||
userRole={userRole}
|
||||
userID={userID}
|
||||
|
|
@ -492,7 +511,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
<Plug />
|
||||
My Connections
|
||||
</Link>
|
||||
{isAdminRole(userRole) && (
|
||||
{isAdminRole(userRole) ? (
|
||||
<>
|
||||
<Button className="shrink-0" variant="secondary" onClick={() => setImportVisible(true)}>
|
||||
Import from JSON
|
||||
|
|
@ -501,8 +520,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
+ Add New MCP Server
|
||||
</Button>
|
||||
</>
|
||||
)}
|
||||
{!isAdminRole(userRole) && (
|
||||
) : (
|
||||
<Button
|
||||
className="shrink-0"
|
||||
onClick={() => {
|
||||
|
|
@ -549,24 +567,23 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
Connect
|
||||
</TabsTrigger>
|
||||
{isAdminRole(userRole) && (
|
||||
<TabsTrigger value="semantic-filter" className="flex-none rounded-none px-4 py-2">
|
||||
Semantic Filter
|
||||
</TabsTrigger>
|
||||
)}
|
||||
{isAdminRole(userRole) && (
|
||||
<TabsTrigger value="tool-search" className="flex-none rounded-none px-4 py-2">
|
||||
Tool Search
|
||||
</TabsTrigger>
|
||||
)}
|
||||
{isAdminRole(userRole) && (
|
||||
<TabsTrigger value="network-settings" className="flex-none rounded-none px-4 py-2">
|
||||
Network Settings
|
||||
</TabsTrigger>
|
||||
)}
|
||||
{isAdminRole(userRole) && (
|
||||
<TabsTrigger value="submitted" className="flex-none rounded-none px-4 py-2">
|
||||
Submitted MCPs
|
||||
</TabsTrigger>
|
||||
<>
|
||||
<TabsTrigger value="semantic-filter" className="flex-none rounded-none px-4 py-2">
|
||||
Semantic Filter
|
||||
</TabsTrigger>
|
||||
|
||||
<TabsTrigger value="tool-search" className="flex-none rounded-none px-4 py-2">
|
||||
Tool Search
|
||||
</TabsTrigger>
|
||||
|
||||
<TabsTrigger value="network-settings" className="flex-none rounded-none px-4 py-2">
|
||||
Network Settings
|
||||
</TabsTrigger>
|
||||
|
||||
<TabsTrigger value="submitted" className="flex-none rounded-none px-4 py-2">
|
||||
Submitted MCPs
|
||||
</TabsTrigger>
|
||||
</>
|
||||
)}
|
||||
{isProxyAdminTierRole(userRole) && (
|
||||
<TabsTrigger value="connections" className="flex-none rounded-none px-4 py-2">
|
||||
|
|
@ -601,13 +618,11 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
value={selectedTeam}
|
||||
onValueChange={(v: string | null) => handleTeamChange(v ?? "all")}
|
||||
>
|
||||
<SelectTrigger className="w-55">
|
||||
<SelectTrigger className="w-55" aria-label="Team">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="all">
|
||||
{isInternalUser ? "All Available Servers" : "All Servers"}
|
||||
</SelectItem>
|
||||
<SelectItem value="all">{teamSelectItems.all}</SelectItem>
|
||||
<SelectItem value="personal">Personal</SelectItem>
|
||||
{uniqueTeams.map((team) => (
|
||||
<SelectItem key={team.team_id} value={team.team_id}>
|
||||
|
|
@ -641,7 +656,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
value={selectedMcpAccessGroup}
|
||||
onValueChange={(v: string | null) => handleMcpAccessGroupChange(v ?? "all")}
|
||||
>
|
||||
<SelectTrigger className="w-55">
|
||||
<SelectTrigger className="w-55" aria-label="Access Group">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
|
|
@ -742,24 +757,23 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
|
|||
<MCPConnect />
|
||||
</TabsContent>
|
||||
{isAdminRole(userRole) && (
|
||||
<TabsContent value="semantic-filter" keepMounted>
|
||||
<MCPSemanticFilterSettings accessToken={accessToken} />
|
||||
</TabsContent>
|
||||
)}
|
||||
{isAdminRole(userRole) && (
|
||||
<TabsContent value="tool-search" keepMounted>
|
||||
<MCPToolSearchSettings accessToken={accessToken} />
|
||||
</TabsContent>
|
||||
)}
|
||||
{isAdminRole(userRole) && (
|
||||
<TabsContent value="network-settings" keepMounted>
|
||||
<MCPNetworkSettings accessToken={accessToken} />
|
||||
</TabsContent>
|
||||
)}
|
||||
{isAdminRole(userRole) && (
|
||||
<TabsContent value="submitted" keepMounted>
|
||||
<MCPSubmissionsTab accessToken={accessToken} />
|
||||
</TabsContent>
|
||||
<>
|
||||
<TabsContent value="semantic-filter" keepMounted>
|
||||
<MCPSemanticFilterSettings accessToken={accessToken} />
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="tool-search" keepMounted>
|
||||
<MCPToolSearchSettings accessToken={accessToken} />
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="network-settings" keepMounted>
|
||||
<MCPNetworkSettings accessToken={accessToken} />
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="submitted" keepMounted>
|
||||
<MCPSubmissionsTab accessToken={accessToken} />
|
||||
</TabsContent>
|
||||
</>
|
||||
)}
|
||||
{isProxyAdminTierRole(userRole) && (
|
||||
<TabsContent value="connections">
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue