diff --git a/.circleci/config.yml b/.circleci/config.yml index cc9aa7fe1c4..83acf0ac1c8 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -6,6 +6,9 @@ parameters: migration_candidate_image: type: string default: "" + migration_baseline_image: + type: string + default: "ghcr.io/berriai/litellm-database:v1.102.0" migration_source_sha: type: string default: "" @@ -2946,7 +2949,10 @@ jobs: parameters: suite: type: enum - enum: [startup, recovery, legacy] + enum: [startup, recovery, legacy, upgrade, shaped] + baseline: + type: boolean + default: false machine: image: ubuntu-2204:2024.04.1 resource_class: large @@ -2954,6 +2960,7 @@ jobs: environment: LITELLM_MIGRATION_TESTS: "1" LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci + LITELLM_MIGRATION_BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >> MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres MIGRATION_TEST_OUTPUT: /tmp/migration-results @@ -2981,6 +2988,16 @@ jobs: - wait_for_service: url: tcp://localhost:5432 timeout: "60" + - when: + condition: << parameters.baseline >> + steps: + - run: + name: Pull the baseline release the upgrade starts from + environment: + BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >> + command: | + [[ "$BASELINE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+(@sha256:[0-9a-f]{64}|:v[0-9][0-9a-z.-]*)$ ]] || exit 1 + docker pull "$BASELINE_IMAGE" - run: name: Run migration startup regressions environment: @@ -3188,6 +3205,16 @@ workflows: name: migration-legacy-and-pooling suite: legacy requires: [build_docker_database_image] + - migration_startup_tests: + name: migration-upgrade + suite: upgrade + baseline: true + requires: [build_docker_database_image] + - migration_startup_tests: + name: migration-upgrade-shaped + suite: shaped + baseline: true + requires: [build_docker_database_image] migration_startup_scheduled: triggers: - schedule: diff --git a/.circleci/scripts/run_migration_tests.py b/.circleci/scripts/run_migration_tests.py index 56029c406fb..5a73c54e3f5 100644 --- a/.circleci/scripts/run_migration_tests.py +++ b/.circleci/scripts/run_migration_tests.py @@ -13,6 +13,8 @@ SUITES: Final = { "startup": (("test_startup.py",), 12), "recovery": (("test_recovery.py",), 15), "legacy": (("test_legacy.py", "test_pooling.py"), 11), + "upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5), + "shaped": (("test_shaped_database.py",), 1), } @@ -93,6 +95,7 @@ def main() -> int: { **metadata, "suite": suite, + "baseline_image": os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE", ""), "expected_cases": expected, "passed": passed, "pytest_exit_code": result.returncode, diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 161353708f0..6f8599daf75 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -130,7 +130,7 @@ jobs: - name: Test secret manager feature combinations run: | cargo test -p litellm-auth-gcp --locked --no-default-features - for features in '' aws google cyberark aws,google aws,google,cyberark; do + for features in '' aws google azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark; do cargo test -p litellm-secrets --locked --no-default-features --features "$features" done diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ca686008de3..5c1688acac1 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -115,6 +115,28 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "async-trait" version = "0.1.91" @@ -599,6 +621,37 @@ dependencies = [ "url", ] +[[package]] +name = "azure_storage_blob" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17b10207ecf7d666df6940b50051f433b3cd5d2b9b1dd190613208d7a84e7eed" +dependencies = [ + "async-stream", + "async-trait", + "azure_core", + "azure_storage_common", + "bytes", + "futures", + "percent-encoding", + "pin-project", + "serde", + "serde_json", + "time", + "tokio", +] + +[[package]] +name = "azure_storage_common" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0af2e6aeb8d76b17fc998f453c320913f73787b944e3cc29509d19411fa0321d" +dependencies = [ + "azure_core", + "serde", + "time", +] + [[package]] name = "base64" version = "0.13.1" @@ -2470,6 +2523,23 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-azure-blob" +version = "0.1.0" +dependencies = [ + "async-trait", + "azure_core", + "azure_storage_blob", + "futures-util", + "litellm-auth-azure", + "litellm-auth-types", + "litellm-cache", + "litellm-cache-response", + "serde_json", + "tokio", + "url", +] + [[package]] name = "litellm-cache-memory" version = "0.1.0" @@ -2687,6 +2757,7 @@ dependencies = [ "litellm-auth", "litellm-auth-gcp", "litellm-cache", + "litellm-cache-azure-blob", "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-redis-semantic", @@ -2720,6 +2791,7 @@ dependencies = [ "jsonwebtoken", "litellm-core-utils", "litellm-secrets-aws", + "litellm-secrets-azure", "litellm-secrets-cyberark", "litellm-secrets-google", "litellm-secrets-types", @@ -2755,6 +2827,26 @@ dependencies = [ "wiremock", ] +[[package]] +name = "litellm-secrets-azure" +version = "0.1.0" +dependencies = [ + "litellm-auth-azure", + "litellm-auth-types", + "litellm-core-utils", + "litellm-secrets-types", + "percent-encoding", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 2.0.19", + "tokio", + "veil", + "wiremock", +] + [[package]] name = "litellm-secrets-cyberark" version = "0.1.0" @@ -3531,6 +3623,16 @@ version = "1.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" +[[package]] +name = "quick-xml" +version = "0.41.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1" +dependencies = [ + "memchr", + "serde", +] + [[package]] name = "quinn" version = "0.11.11" @@ -5076,6 +5178,7 @@ dependencies = [ "base64 0.22.1", "bytes", "futures", + "quick-xml", "serde", "serde_json", "url", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index f1d40f83ef2..91e0cc08c8f 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -22,12 +22,14 @@ litellm-secrets = { path = "crates/secrets" } litellm-secrets-types = { path = "crates/secrets-types" } litellm-secrets-aws = { path = "crates/secrets-aws" } litellm-secrets-google = { path = "crates/secrets-google" } +litellm-secrets-azure = { path = "crates/secrets-azure" } litellm-secrets-cyberark = { path = "crates/secrets-cyberark" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } litellm-core-utils = { path = "crates/core-utils" } 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-redis-semantic = { path = "crates/cache-redis-semantic" } diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs index e76227d6aa2..5c7c654b69d 100644 --- a/litellm-rust/crates/auth-azure/src/lib.rs +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -4,4 +4,4 @@ mod resolve; mod types; pub use resolve::AzureAuthService; -pub use types::AzureAuthInputs; +pub use types::{AzureAuthInputs, ConfigValue}; diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index d5a00f09751..a3a898f000f 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -51,6 +51,21 @@ pub struct AzureAuthInputs { } impl AzureAuthInputs { + pub fn default_credential_for_scope(scope: &str) -> Self { + Self { + azure_scope: ConfigValue::Value(Sourced::new( + scope.to_string(), + InputSource::Deployment, + )), + azure_credential: ConfigValue::Value(Sourced::new( + "DefaultAzureCredential".to_string(), + InputSource::Deployment, + )), + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Self::default() + } + } + pub fn or_configured_token_refresh(self, enabled: bool) -> Self { if *self.enable_azure_ad_token_refresh.value() || !enabled { return self; diff --git a/litellm-rust/crates/cache-azure-blob/Cargo.toml b/litellm-rust/crates/cache-azure-blob/Cargo.toml new file mode 100644 index 00000000000..55abaff1975 --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "litellm-cache-azure-blob" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-azure.workspace = true +litellm-auth-types.workspace = true +litellm-cache.workspace = true + +async-trait = "0.1" +azure_core = "1.1.0" +azure_storage_blob = "1.1.0" +futures-util.workspace = true +tokio.workspace = true +url.workspace = true + +[dev-dependencies] +litellm-cache-response.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/cache-azure-blob/src/cache.rs b/litellm-rust/crates/cache-azure-blob/src/cache.rs new file mode 100644 index 00000000000..6a872a0d6e6 --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/cache.rs @@ -0,0 +1,254 @@ +use std::{sync::Arc, time::Duration}; + +use azure_core::{ + credentials::TokenCredential, + error::ErrorKind, + http::{ClientOptions, RequestContent}, +}; +use azure_storage_blob::{ + BlobContainerClient, BlobContainerClientOptions, + models::{BlobClientUploadOptions, StorageErrorCode}, +}; +use futures_util::{TryStreamExt, future::try_join_all}; +use litellm_cache::{ + BaseCache, BatchCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + ExactCacheContext, FlushCache, +}; +use tokio::runtime::Handle; +use url::Url; + +use crate::credential::AzureBlobCredential; + +pub struct AzureBlobCache { + container: BlobContainerClient, + codec: C, + runtime: Handle, + account_url: String, + container_name: String, +} + +impl AzureBlobCache { + pub async fn connect( + account_url: &str, + container: &str, + codec: C, + runtime: Handle, + ) -> Result { + Self::connect_with_options( + account_url, + container, + Some(Arc::new(AzureBlobCredential::default())), + ClientOptions::default(), + codec, + runtime, + ) + .await + } + + pub async fn connect_with_options( + account_url: &str, + container: &str, + credential: Option>, + client_options: ClientOptions, + codec: C, + runtime: Handle, + ) -> Result { + let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?; + let account_url = parsed.as_str().trim_end_matches('/').to_string(); + let container_url = { + let mut url = parsed; + url.path_segments_mut() + .map_err(|()| Error::Unavailable)? + .pop_if_empty() + .push(container); + url + }; + let client = BlobContainerClient::new( + container_url, + credential, + Some(BlobContainerClientOptions { + client_options, + ..BlobContainerClientOptions::default() + }), + ) + .map_err(|_| Error::Unavailable)?; + let cache = Self { + container: client, + codec, + runtime, + account_url, + container_name: container.to_string(), + }; + cache.create_container().await?; + Ok(cache) + } + + pub fn account_url(&self) -> &str { + &self.account_url + } + + pub fn container_name(&self) -> &str { + &self.container_name + } + + async fn create_container(&self) -> Result<(), Error> { + match self.container.create(None).await { + Ok(_) => Ok(()), + Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => { + Ok(()) + } + Err(_) => Err(Error::Unavailable), + } + } + + async fn upload(&self, key: &str, value: &C::Value, overwrite: bool) -> Result<(), Error> { + let payload = self.codec.encode(value)?; + let options = (!overwrite).then(|| BlobClientUploadOptions::default().if_not_exists()); + match self + .container + .blob_client(key) + .upload(RequestContent::from(payload), options) + .await + { + Ok(_) => Ok(()), + Err(error) if !overwrite && is_already_present(&error) => Ok(()), + Err(_) => Err(Error::Unavailable), + } + } + + async fn download(&self, key: &str) -> Result, Error> { + let response = match self.container.blob_client(key).download(None).await { + Ok(response) => response, + Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => { + return Ok(None); + } + Err(_) => return Err(Error::Unavailable), + }; + let bytes = response + .body + .collect() + .await + .map_err(|_| Error::Unavailable)?; + self.codec.decode(&bytes).map(Some) + } + + async fn delete_all_blobs(&self) -> Result<(), Error> { + let mut pages = self + .container + .list_blobs(None) + .map_err(|_| Error::Unavailable)? + .into_pages(); + while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? { + let page = page.into_model().map_err(|_| Error::Unavailable)?; + for name in page.blob_items.into_iter().filter_map(|item| item.name) { + self.container + .blob_client(&name) + .delete(None) + .await + .map_err(|_| Error::Unavailable)?; + } + } + Ok(()) + } + + fn block_on(&self, future: impl Future) -> T { + self.runtime.block_on(future) + } +} + +fn is_already_present(error: &azure_core::Error) -> bool { + is_storage_error(error, StorageErrorCode::BlobAlreadyExists) + || is_storage_error(error, StorageErrorCode::ConditionNotMet) +} + +fn is_storage_error(error: &azure_core::Error, code: StorageErrorCode) -> bool { + matches!( + error.kind(), + ErrorKind::HttpResponse { + error_code: Some(error_code), + .. + } if error_code == code.as_ref() + ) +} + +impl BaseCache for AzureBlobCache { + type Value = C::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, _: &ExactCacheContext) -> Option { + None + } + + fn set_cache(&self, key: &str, value: C::Value, _: &ExactCacheContext) -> Result<(), Error> { + self.block_on(self.upload(key, &value, false)) + } + + fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { + self.block_on(self.download(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: C::Value, + _: ExactCacheContext, + ) -> Result<(), Error> { + self.upload(key, &value, true).await + } + + async fn async_get_cache( + &self, + key: &str, + _: &ExactCacheContext, + ) -> Result, Error> { + self.download(key).await + } + + async fn async_set_cache_pipeline( + &self, + entries: Vec<(String, C::Value)>, + _: ExactCacheContext, + ) -> Result<(), Error> { + try_join_all( + entries + .iter() + .map(|(key, value)| self.upload(key, value, true)), + ) + .await + .map(drop) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Ok(match self.container.get_properties(None).await { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Azure Blob cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Azure Blob connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + } +} + +impl BatchCache for AzureBlobCache {} + +impl FlushCache for AzureBlobCache { + fn flush_cache(&self) -> Result<(), Error> { + self.block_on(self.delete_all_blobs()) + } + + async fn async_flush_cache(&self) -> Result<(), Error> { + self.delete_all_blobs().await + } +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs b/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs new file mode 100644 index 00000000000..f8736ab069b --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs @@ -0,0 +1,746 @@ +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use azure_core::http::{ + AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport, + headers::{HeaderName, Headers}, +}; +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, Error, ExactCacheContext, FlushCache, +}; +use litellm_cache_response::{ + CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec, + ResponseCacheRequest, cache_key, +}; +use serde_json::json; +use tokio::runtime::Runtime; + +use super::AzureBlobCache; + +const ACCOUNT_URL: &str = "https://example.blob.core.windows.net"; +const CONTAINER: &str = "litellm-cache"; +const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match"); +const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code"); + +#[derive(Clone, Debug, PartialEq, Eq)] +struct RecordedRequest { + method: Method, + path: String, + query: String, + if_none_match: Option, +} + +#[derive(Default)] +struct FakeState { + container_exists: bool, + blobs: BTreeMap>, + requests: Vec, + failing: bool, + precondition_conflicts: bool, +} + +#[derive(Clone, Default)] +struct FakeBlobService { + state: Arc>, +} + +impl std::fmt::Debug for FakeBlobService { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("FakeBlobService") + } +} + +impl FakeBlobService { + fn with_existing_container() -> Self { + let service = Self::default(); + service.state.lock().unwrap().container_exists = true; + service + } + + fn blob(&self, name: &str) -> Option> { + self.state.lock().unwrap().blobs.get(name).cloned() + } + + fn blob_names(&self) -> Vec { + self.state.lock().unwrap().blobs.keys().cloned().collect() + } + + fn seed_blob(&self, name: &str, bytes: &[u8]) { + self.state + .lock() + .unwrap() + .blobs + .insert(name.to_string(), bytes.to_vec()); + } + + fn set_failing(&self, failing: bool) { + self.state.lock().unwrap().failing = failing; + } + + fn set_precondition_conflicts(&self, enabled: bool) { + self.state.lock().unwrap().precondition_conflicts = enabled; + } + + fn requests(&self) -> Vec { + self.state.lock().unwrap().requests.clone() + } + + fn container_exists(&self) -> bool { + self.state.lock().unwrap().container_exists + } + + fn respond(status: StatusCode, error_code: Option<&str>, body: Vec) -> AsyncRawResponse { + let mut headers = Headers::new(); + if let Some(code) = error_code { + headers.insert(ERROR_CODE, code.to_string()); + } + AsyncRawResponse::from_bytes(status, headers, body) + } + + fn list_body(state: &FakeState) -> Vec { + let mut xml = String::from( + r#""#, + ); + for name in state.blobs.keys() { + xml.push_str(&format!( + "{name}BlockBlob" + )); + } + xml.push_str(""); + xml.into_bytes() + } +} + +#[async_trait::async_trait] +impl HttpClient for FakeBlobService { + async fn execute_request(&self, request: &Request) -> azure_core::Result { + let mut state = self.state.lock().unwrap(); + let path = request.url().path().to_string(); + let query = request.url().query().unwrap_or_default().to_string(); + let if_none_match = request + .headers() + .get_optional_str(&IF_NONE_MATCH) + .map(str::to_owned); + state.requests.push(RecordedRequest { + method: request.method(), + path: path.clone(), + query: query.clone(), + if_none_match: if_none_match.clone(), + }); + if state.failing { + return Ok(Self::respond( + StatusCode::Forbidden, + Some("AuthorizationFailure"), + Vec::new(), + )); + } + let container_path = format!("/{CONTAINER}"); + let blob_name = path + .strip_prefix(&format!("{container_path}/")) + .map(str::to_owned); + let is_container = path == container_path && query.contains("restype=container"); + let response = match (request.method(), is_container, blob_name) { + (Method::Put, true, None) if state.container_exists => Self::respond( + StatusCode::Conflict, + Some("ContainerAlreadyExists"), + Vec::new(), + ), + (Method::Put, true, None) => { + state.container_exists = true; + Self::respond(StatusCode::Created, None, Vec::new()) + } + (Method::Get, true, None) if query.contains("comp=list") => { + Self::respond(StatusCode::Ok, None, Self::list_body(&state)) + } + (Method::Get, true, None) if state.container_exists => { + Self::respond(StatusCode::Ok, None, Vec::new()) + } + (Method::Get, true, None) => { + Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new()) + } + (Method::Put, false, Some(name)) => { + if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) { + if state.precondition_conflicts { + Self::respond( + StatusCode::PreconditionFailed, + Some("ConditionNotMet"), + Vec::new(), + ) + } else { + Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new()) + } + } else { + let bytes = match request.body() { + Body::Bytes(bytes) => bytes.to_vec(), + Body::SeekableStream(_) => panic!("unexpected streaming upload"), + }; + state.blobs.insert(name, bytes); + Self::respond(StatusCode::Created, None, Vec::new()) + } + } + (Method::Get, false, Some(name)) => match state.blobs.get(&name) { + Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()), + None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()), + }, + (Method::Delete, false, Some(name)) => match state.blobs.remove(&name) { + Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()), + None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()), + }, + (method, _, _) => panic!("unexpected request {method:?} {path}?{query}"), + }; + Ok(response) + } +} + +struct Fixture { + runtime: Runtime, + service: FakeBlobService, + cache: Arc>, +} + +impl Fixture { + fn new(service: FakeBlobService) -> Self { + let runtime = Runtime::new().unwrap(); + let cache = runtime + .block_on(Self::connect(&service, runtime.handle().clone())) + .unwrap(); + Self { + runtime, + service, + cache: Arc::new(cache), + } + } + + async fn connect( + service: &FakeBlobService, + handle: tokio::runtime::Handle, + ) -> Result, Error> { + AzureBlobCache::connect_with_options( + ACCOUNT_URL, + CONTAINER, + None, + ClientOptions { + transport: Some(Transport::new(Arc::new(service.clone()))), + ..ClientOptions::default() + }, + ResponseCacheCodec, + handle, + ) + .await + } + + fn response_cache(&self) -> ResponseCache> { + ResponseCache::new(self.cache.clone()) + } + + fn stored_json(&self, key: &str) -> serde_json::Value { + serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap() + } +} + +fn request(model: &str) -> ResponseCacheRequest { + ResponseCacheRequest::new(CacheKeyInput { + fields: vec![CacheKeyField { + name: "model".into(), + value: Some(model.into()), + api_parameter: true, + internal_parameter: false, + }], + preset: None, + namespace: None, + include_provider_parameters: false, + }) +} + +fn now() -> Duration { + Duration::from_secs(1_700_000_000) +} + +fn entry(value: serde_json::Value) -> CacheEntry { + CacheEntry { + timestamp: Some(1_700_000_000.5), + response: value, + } +} + +fn no_ttl() -> ExactCacheContext { + ExactCacheContext::default() +} + +fn with_ttl(seconds: u64) -> ExactCacheContext { + ExactCacheContext { + ttl: Some(Duration::from_secs(seconds)), + } +} + +#[test] +fn connect_creates_the_container_once() { + let fixture = Fixture::new(FakeBlobService::default()); + assert!(fixture.service.container_exists()); + assert_eq!( + fixture.service.requests(), + vec![RecordedRequest { + method: Method::Put, + path: format!("/{CONTAINER}"), + query: "restype=container".into(), + if_none_match: None, + }] + ); + assert_eq!(fixture.cache.account_url(), ACCOUNT_URL); + assert_eq!(fixture.cache.container_name(), CONTAINER); +} + +#[test] +fn connect_accepts_an_existing_container() { + let fixture = Fixture::new(FakeBlobService::with_existing_container()); + assert!(fixture.service.container_exists()); + assert_eq!(fixture.service.requests().len(), 1); +} + +#[test] +fn connect_accepts_account_urls_with_trailing_slash() { + let runtime = Runtime::new().unwrap(); + let service = FakeBlobService::default(); + let cache = runtime + .block_on(AzureBlobCache::connect_with_options( + "https://example.blob.core.windows.net/", + CONTAINER, + None, + ClientOptions { + transport: Some(Transport::new(Arc::new(service.clone()))), + ..ClientOptions::default() + }, + ResponseCacheCodec, + runtime.handle().clone(), + )) + .unwrap(); + assert_eq!(service.requests()[0].path, format!("/{CONTAINER}")); + assert_eq!(cache.account_url(), "https://example.blob.core.windows.net"); +} + +#[test] +fn connect_keeps_account_url_query_parameters_on_the_container_path() { + let runtime = Runtime::new().unwrap(); + let service = FakeBlobService::default(); + runtime + .block_on(AzureBlobCache::connect_with_options( + "https://example.blob.core.windows.net/?sv=2024-01-01&sig=abc", + CONTAINER, + None, + ClientOptions { + transport: Some(Transport::new(Arc::new(service.clone()))), + ..ClientOptions::default() + }, + ResponseCacheCodec, + runtime.handle().clone(), + )) + .unwrap(); + let create = &service.requests()[0]; + assert_eq!(create.path, format!("/{CONTAINER}")); + assert!(create.query.contains("sig=abc")); +} + +#[test] +fn connect_surfaces_service_failures() { + let runtime = Runtime::new().unwrap(); + let service = FakeBlobService::default(); + service.set_failing(true); + let result = runtime.block_on(Fixture::connect(&service, runtime.handle().clone())); + assert!(matches!(result, Err(Error::Unavailable))); +} + +#[test] +fn sync_set_and_get_round_trip_python_json_shape() { + let fixture = Fixture::new(FakeBlobService::default()); + let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]})); + fixture + .cache + .set_cache("key-1", value.clone(), &no_ttl()) + .unwrap(); + + assert_eq!( + fixture.stored_json("key-1"), + json!({ + "timestamp": 1_700_000_000.5, + "response": {"choices": [{"message": {"content": "héllo 🌍"}}]} + }) + ); + assert_eq!( + fixture.cache.get_cache("key-1", &no_ttl()).unwrap(), + Some(value) + ); +} + +#[test] +fn sync_set_does_not_overwrite_an_existing_blob() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture + .cache + .set_cache("key", entry(json!({"v": "first"})), &no_ttl()) + .unwrap(); + fixture + .cache + .set_cache("key", entry(json!({"v": "second"})), &no_ttl()) + .unwrap(); + + assert_eq!( + fixture.stored_json("key")["response"], + json!({"v": "first"}) + ); + let uploads: Vec<_> = fixture + .service + .requests() + .into_iter() + .filter(|request| request.method == Method::Put && request.path.ends_with("/key")) + .collect(); + assert_eq!(uploads.len(), 2); + assert!( + uploads + .iter() + .all(|request| request.if_none_match.as_deref() == Some("*")) + ); +} + +#[test] +fn sync_set_treats_a_precondition_conflict_as_an_existing_blob() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.service.set_precondition_conflicts(true); + fixture + .cache + .set_cache("key", entry(json!({"v": "first"})), &no_ttl()) + .unwrap(); + fixture + .cache + .set_cache("key", entry(json!({"v": "second"})), &no_ttl()) + .unwrap(); + + assert_eq!( + fixture.stored_json("key")["response"], + json!({"v": "first"}) + ); +} + +#[test] +fn async_set_overwrites_an_existing_blob() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.runtime.block_on(async { + fixture + .cache + .async_set_cache("key", entry(json!({"v": "first"})), no_ttl()) + .await + .unwrap(); + fixture + .cache + .async_set_cache("key", entry(json!({"v": "second"})), no_ttl()) + .await + .unwrap(); + assert_eq!( + fixture + .cache + .async_get_cache("key", &no_ttl()) + .await + .unwrap(), + Some(entry(json!({"v": "second"}))) + ); + }); + assert_eq!( + fixture.stored_json("key")["response"], + json!({"v": "second"}) + ); + assert!( + fixture + .service + .requests() + .iter() + .filter(|request| request.method == Method::Put && request.path.ends_with("/key")) + .all(|request| request.if_none_match.is_none()) + ); +} + +#[test] +fn missing_blobs_are_misses() { + let fixture = Fixture::new(FakeBlobService::default()); + assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None); + assert_eq!( + fixture + .runtime + .block_on(fixture.cache.async_get_cache("absent", &no_ttl())) + .unwrap(), + None + ); +} + +#[test] +fn ttl_is_ignored_and_entries_never_expire() { + let fixture = Fixture::new(FakeBlobService::default()); + assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None); + assert_eq!(fixture.cache.get_ttl(&no_ttl()), None); + + fixture + .cache + .set_cache("key", entry(json!("value")), &with_ttl(1)) + .unwrap(); + std::thread::sleep(Duration::from_millis(1100)); + assert_eq!( + fixture.cache.get_cache("key", &with_ttl(1)).unwrap(), + Some(entry(json!("value"))) + ); + assert!( + fixture + .service + .requests() + .iter() + .all(|request| !request.query.contains("expiry")) + ); +} + +#[test] +fn malformed_blobs_are_invalid_entries_and_response_cache_misses() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.service.seed_blob("broken-json", b"{not json"); + fixture + .service + .seed_blob("broken-utf8", &[0xff, 0xfe, 0x22]); + fixture + .service + .seed_blob("wrong-shape", br#"{"timestamp": "yesterday"}"#); + + for key in ["broken-json", "broken-utf8", "wrong-shape"] { + assert!(matches!( + fixture.cache.get_cache(key, &no_ttl()), + Err(Error::InvalidEntry) + )); + } + + let response_cache = fixture.response_cache(); + let broken = request("broken"); + fixture + .service + .seed_blob(&cache_key(&broken.key), b"{not json"); + assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None); + assert_eq!( + fixture + .runtime + .block_on(response_cache.async_lookup(&broken, now())) + .unwrap(), + None + ); +} + +#[test] +fn batch_get_preserves_order_and_marks_misses_and_invalid_entries() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture + .cache + .set_cache("a", entry(json!("A")), &no_ttl()) + .unwrap(); + fixture + .cache + .set_cache("c", entry(json!("C")), &no_ttl()) + .unwrap(); + fixture.service.seed_blob("bad", b"nope"); + let keys = ["c", "missing", "a", "bad"].map(String::from); + + let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap(); + assert_eq!( + sync, + vec![ + BatchEntry::Hit(entry(json!("C"))), + BatchEntry::Miss, + BatchEntry::Hit(entry(json!("A"))), + BatchEntry::Invalid, + ] + ); + + let asynchronous = fixture + .runtime + .block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl())) + .unwrap(); + assert_eq!(asynchronous, sync); + + let response_cache = fixture.response_cache(); + let requests = [request("hit"), request("missing"), request("bad")]; + response_cache + .store(&requests[0], json!("HIT"), now()) + .unwrap(); + fixture + .service + .seed_blob(&cache_key(&requests[2].key), b"nope"); + let hits = response_cache.lookup_batch(&requests, now()).unwrap(); + assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]); + assert_eq!(hits.missing_indices, vec![1, 2]); + let async_hits = fixture + .runtime + .block_on(response_cache.async_lookup_batch(&requests, now())) + .unwrap(); + assert_eq!(async_hits.values, hits.values); +} + +#[test] +fn async_pipeline_writes_every_entry_with_overwrite() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.service.seed_blob("k2", b"stale"); + fixture + .runtime + .block_on(fixture.cache.async_set_cache_pipeline( + vec![ + ("k1".into(), entry(json!({"n": 1}))), + ("k2".into(), entry(json!({"n": 2}))), + ("k3".into(), entry(json!({"n": 3}))), + ], + with_ttl(30), + )) + .unwrap(); + assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]); + assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2})); +} + +#[test] +fn flush_deletes_every_blob_in_the_container() { + let fixture = Fixture::new(FakeBlobService::default()); + for key in ["x", "y", "z"] { + fixture + .cache + .set_cache(key, entry(json!(key)), &no_ttl()) + .unwrap(); + } + fixture.cache.flush_cache().unwrap(); + assert!(fixture.service.blob_names().is_empty()); + assert!(fixture.service.container_exists()); + + fixture + .cache + .set_cache("again", entry(json!(1)), &no_ttl()) + .unwrap(); + fixture + .runtime + .block_on(fixture.cache.async_flush_cache()) + .unwrap(); + assert!(fixture.service.blob_names().is_empty()); +} + +#[test] +fn service_failures_map_to_unavailable() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.service.set_failing(true); + assert!(matches!( + fixture.cache.get_cache("key", &no_ttl()), + Err(Error::Unavailable) + )); + assert!(matches!( + fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()), + Err(Error::Unavailable) + )); + assert!(matches!( + fixture.cache.flush_cache(), + Err(Error::Unavailable) + )); + assert!(matches!( + fixture.runtime.block_on( + fixture + .cache + .async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl()) + ), + Err(Error::Unavailable) + )); +} + +#[test] +fn test_connection_reports_container_reachability() { + let fixture = Fixture::new(FakeBlobService::default()); + let ok = fixture + .runtime + .block_on(fixture.cache.test_connection()) + .unwrap(); + assert_eq!(ok.status, CacheConnectionStatus::Success); + assert!(ok.error.is_none()); + + fixture.service.set_failing(true); + let failed = fixture + .runtime + .block_on(fixture.cache.test_connection()) + .unwrap(); + assert_eq!(failed.status, CacheConnectionStatus::Failed); + assert!(failed.error.is_some()); +} + +#[test] +fn disconnect_is_idempotent_and_keeps_data() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture + .cache + .set_cache("key", entry(json!(1)), &no_ttl()) + .unwrap(); + fixture.runtime.block_on(async { + fixture.cache.disconnect().await.unwrap(); + fixture.cache.disconnect().await.unwrap(); + }); + assert_eq!( + fixture.cache.get_cache("key", &no_ttl()).unwrap(), + Some(entry(json!(1))) + ); +} + +#[test] +fn response_cache_stores_and_reads_through_the_backend() { + let fixture = Fixture::new(FakeBlobService::default()); + let response_cache = fixture.response_cache(); + let mut request = request("gpt"); + request.context = with_ttl(60); + let response = json!({"id": "chatcmpl-1"}); + response_cache + .store(&request, response.clone(), now()) + .unwrap(); + assert_eq!( + fixture.stored_json(&cache_key(&request.key)), + json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}}) + ); + assert_eq!( + response_cache + .lookup(&request, now() + Duration::from_secs(3600)) + .unwrap(), + Some(response.clone()) + ); + assert_eq!( + fixture + .runtime + .block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600))) + .unwrap(), + Some(response.clone()) + ); + fixture.runtime.block_on(async { + response_cache + .async_store(&request, json!("replaced"), now()) + .await + .unwrap(); + assert_eq!( + response_cache.async_lookup(&request, now()).await.unwrap(), + Some(json!("replaced")) + ); + response_cache.async_flush().await.unwrap(); + assert_eq!( + response_cache.async_lookup(&request, now()).await.unwrap(), + None + ); + }); +} + +#[test] +fn non_object_responses_are_written_serialized_like_python() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture + .cache + .set_cache("s", entry(json!("plain")), &no_ttl()) + .unwrap(); + assert_eq!( + fixture.stored_json("s"), + json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""}) + ); + assert_eq!( + fixture.cache.get_cache("s", &no_ttl()).unwrap(), + Some(entry(json!("plain"))) + ); +} diff --git a/litellm-rust/crates/cache-azure-blob/src/credential.rs b/litellm-rust/crates/cache-azure-blob/src/credential.rs new file mode 100644 index 00000000000..d1a3d0e44ec --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/credential.rs @@ -0,0 +1,84 @@ +use std::{ + fmt, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use azure_core::{ + credentials::{AccessToken, TokenCredential, TokenRequestOptions}, + error::ErrorKind, + time::OffsetDateTime, +}; +use litellm_auth_azure::{AzureAuthInputs, AzureAuthService}; +use litellm_auth_types::ResolvedCredential; + +const STATIC_TOKEN_LIFETIME: Duration = Duration::from_secs(300); +const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; + +type EnvLookup = Arc Option + Send + Sync>; + +pub struct AzureBlobCredential { + service: AzureAuthService, + env_lookup: EnvLookup, +} + +impl fmt::Debug for AzureBlobCredential { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("AzureBlobCredential") + } +} + +impl Default for AzureBlobCredential { + fn default() -> Self { + Self::new( + AzureAuthService::default(), + Arc::new(|name| std::env::var(name).ok()), + ) + } +} + +impl AzureBlobCredential { + pub fn new(service: AzureAuthService, env_lookup: EnvLookup) -> Self { + Self { + service, + env_lookup, + } + } +} + +#[async_trait::async_trait] +impl TokenCredential for AzureBlobCredential { + async fn get_token( + &self, + scopes: &[&str], + _options: Option>, + ) -> azure_core::Result { + let env_lookup = &self.env_lookup; + let lookup = move |name: &str| (name != LLM_TOKEN_ENV).then(|| env_lookup(name)).flatten(); + let credential = self + .service + .get_azure_ad_token( + &AzureAuthInputs::default_credential_for_scope(&scopes.join(" ")), + &lookup, + ) + .await + .map_err(|error| { + azure_core::Error::with_message(ErrorKind::Credential, error.to_string()) + })? + .ok_or_else(|| { + azure_core::Error::with_message( + ErrorKind::Credential, + "no Azure credential is available for blob storage", + ) + })?; + let (token, expires_on) = match credential.into_value() { + ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on), + ResolvedCredential::Static(token) => (token, None), + }; + let expires_on = expires_on.unwrap_or_else(|| SystemTime::now() + STATIC_TOKEN_LIFETIME); + Ok(AccessToken::new( + token.expose().to_string(), + OffsetDateTime::from(expires_on), + )) + } +} diff --git a/litellm-rust/crates/cache-azure-blob/src/lib.rs b/litellm-rust/crates/cache-azure-blob/src/lib.rs new file mode 100644 index 00000000000..5ae752c111d --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod credential; + +pub use cache::AzureBlobCache; +pub use credential::AzureBlobCredential; diff --git a/litellm-rust/crates/cache-azure-blob/src/tests.rs b/litellm-rust/crates/cache-azure-blob/src/tests.rs new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 93ce5828489..507bd20348a 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -21,6 +21,7 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true litellm-cache.workspace = true +litellm-cache-azure-blob.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true litellm-cache-redis-semantic.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 3c20dc6f75c..3db49f9dfc5 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -75,6 +75,11 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +pub(super) struct AzureBlobCacheConfig { + pub(super) account_url: String, + pub(super) container: String, +} + #[allow( dead_code, reason = "embedding settings are projected so drift falls back to Python" @@ -102,6 +107,7 @@ const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31; pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + AzureBlob(AzureBlobCacheConfig), RedisSemantic(Box), } @@ -169,6 +175,12 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::AzureBlob(backend), + })) + }), Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| { CacheConfigProjection::Native(Box::new(Self { policy, @@ -180,7 +192,6 @@ impl NativeCacheConfig { | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic - | CacheType::AzureBlob | CacheType::Gcs, ) | None => Ok(CacheConfigProjection::Unsupported( @@ -190,13 +201,12 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - if service.default_ttl() - != match &self.backend { - CacheBackendConfig::Memory(config) => Some(config.default_ttl), - CacheBackendConfig::Redis(config) => Some(config.default_ttl), - CacheBackendConfig::RedisSemantic(_) => None, - } - { + let default_ttl = match &self.backend { + CacheBackendConfig::Memory(config) => Some(config.default_ttl), + CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::AzureBlob(_) | CacheBackendConfig::RedisSemantic(_) => None, + }; + if service.default_ttl() != default_ttl { return Some("facade and native backend default TTLs must match"); } match &self.backend { @@ -235,10 +245,34 @@ impl NativeCacheConfig { Some("facade and native backend similarity thresholds must match") } CacheBackendConfig::RedisSemantic(_) => None, + CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() { + None => Some("facade and native backend types must match"), + Some((account_url, container)) + if account_url != config.account_url || container != config.container => + { + Some("facade and native backend containers must match") + } + Some(_) => None, + }, } } } +#[inline(never)] +fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult { + let client = backend.getattr("container_client")?; + let container = client.getattr("container_name")?.extract::()?; + let url = client.getattr("url")?.extract::()?; + let account_url = url + .strip_suffix(container.as_str()) + .and_then(|url| url.strip_suffix('/')) + .ok_or_else(|| PyValueError::new_err("Azure Blob container URL is malformed"))?; + Ok(AzureBlobCacheConfig { + account_url: account_url.to_string(), + container, + }) +} + #[inline(never)] pub(super) fn project_redis_semantic( backend: &Bound<'_, PyAny>, diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 398f46d0d1e..bf1b69cfbf5 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -34,6 +34,19 @@ struct RedisPoolGuard { attributes: RedisPoolAttributes, } +struct AzureBlobClientGuard { + sync_client: Py, + async_client: Py, + url: String, + container_name: String, +} + +enum ConnectionGuard { + None, + RedisPool(RedisPoolGuard), + AzureBlob(AzureBlobClientGuard), +} + struct RedisPoolAttributes { pool: &'static str, connection_class: &'static str, @@ -55,7 +68,7 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes { pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, - redis_pool: Option, + connection: ConnectionGuard, } impl ObjectGuard { @@ -208,6 +221,61 @@ impl RedisPoolGuard { } } +impl AzureBlobClientGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + let sync_client = backend.getattr("container_client")?; + Ok(Self { + url: sync_client.getattr("url")?.extract::()?, + container_name: sync_client.getattr("container_name")?.extract::()?, + sync_client: sync_client.unbind(), + async_client: backend.getattr("async_container_client")?.unbind(), + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + let sync_client = backend.getattr("container_client")?; + Ok(self.sync_client.bind(py).is(&sync_client) + && self + .async_client + .bind(py) + .is(&backend.getattr("async_container_client")?) + && self.url == sync_client.getattr("url")?.extract::()? + && self.container_name == sync_client.getattr("container_name")?.extract::()?) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.sync_client)?; + visit.call(&self.async_client) + } +} + +impl ConnectionGuard { + fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult { + Ok(match (kind, cluster) { + ("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?), + ("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?), + ("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?), + _ => Self::None, + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + match self { + Self::None => Ok(true), + Self::RedisPool(guard) => guard.matches(py, backend), + Self::AzureBlob(guard) => guard.matches(py, backend), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + match self { + Self::None => Ok(()), + Self::RedisPool(guard) => guard.traverse(visit), + Self::AzureBlob(guard) => guard.traverse(visit), + } + } +} + impl FacadeGuard { pub(super) fn capture( py: Python<'_>, @@ -235,6 +303,11 @@ impl FacadeGuard { "RedisClusterCache", "redis", ), + ("azure-blob", _) => ( + "litellm.caching.azure_blob_cache", + "AzureBlobCache", + "azure-blob", + ), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -296,11 +369,7 @@ impl FacadeGuard { "_redis_url", ], )?, - redis_pool: match (kind, cluster) { - ("redis", false) => Some(RedisPoolGuard::capture(&backend, STANDALONE_POOL)?), - ("redis", true) => Some(RedisPoolGuard::capture(&backend, CLUSTER_POOL)?), - _ => None, - }, + connection: ConnectionGuard::capture(kind, cluster, &backend)?, }) } @@ -312,19 +381,13 @@ impl FacadeGuard { if !self.backend.matches(py, &backend)? { return Ok(false); } - match &self.redis_pool { - Some(guard) => guard.matches(py, &backend), - None => Ok(true), - } + self.connection.matches(py, &backend) } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.outer.traverse(&visit)?; self.backend.traverse(&visit)?; - if let Some(guard) = &self.redis_pool { - guard.traverse(&visit)?; - } - Ok(()) + self.connection.traverse(&visit) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 7b5300b52f7..e124f489644 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,6 +1,6 @@ use litellm_cache_redis::{RedisNode, RedisTopology}; use litellm_cache_redis_semantic::RedisSemanticConfig; -use litellm_host_python::release_gil; +use litellm_host_python::{release_gil, run_sync_value}; use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyRuntimeError, PyTypeError}, @@ -72,6 +72,21 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (account_url, container))] + fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult { + let service = run_sync_value(py, async move { + NativeResponseCache::azure_blob(&account_url, &container) + .await + .map_err(cache_error) + })?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[staticmethod] fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult { let class = py diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index de9c4afa236..bad2c30f189 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,6 +1,7 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error, ExactCacheContext}; +use litellm_cache_azure_blob::AzureBlobCache; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; @@ -19,6 +20,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + AzureBlob(Arc>>), RedisSemantic(Arc>>), } @@ -50,7 +52,6 @@ impl NativeResponseCache { buffer: None, }) } - pub fn redis_semantic( url: &str, embedder: PythonEmbedder, @@ -61,6 +62,29 @@ impl NativeResponseCache { backend, ))))) } + + pub async fn azure_blob(account_url: &str, container: &str) -> Result { + let backend = AzureBlobCache::connect( + account_url, + container, + ResponseCacheCodec, + tokio::runtime::Handle::current(), + ) + .await?; + Ok(Self::AzureBlob(Arc::new(ResponseCache::new(Arc::new( + backend, + ))))) + } + + pub fn azure_blob_identity(&self) -> Option<(&str, &str)> { + match self { + Self::AzureBlob(cache) => Some(( + cache.backend().account_url(), + cache.backend().container_name(), + )), + Self::Memory(_) | Self::Redis { .. } | Self::RedisSemantic(_) => None, + } + } } impl NativeResponseCache { @@ -69,6 +93,7 @@ impl NativeResponseCache { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", Self::RedisSemantic(_) => "redis_semantic", + Self::AzureBlob(_) => "azure-blob", } } @@ -77,19 +102,20 @@ impl NativeResponseCache { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), Self::RedisSemantic(cache) => cache.default_ttl(), + Self::AzureBlob(cache) => cache.default_ttl(), } } pub fn namespace(&self) -> Option<&str> { match self { - Self::Memory(_) | Self::RedisSemantic(_) => None, + Self::Memory(_) | Self::AzureBlob(_) | Self::RedisSemantic(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), } } pub fn topology(&self) -> Option<&RedisTopology> { match self { - Self::Memory(_) | Self::RedisSemantic(_) => None, + Self::Memory(_) | Self::AzureBlob(_) | Self::RedisSemantic(_) => None, Self::Redis { cache, .. } => Some(cache.backend().topology()), } } @@ -97,42 +123,42 @@ impl NativeResponseCache { pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } | Self::RedisSemantic(_) => None, + Self::Redis { .. } | Self::AzureBlob(_) | Self::RedisSemantic(_) => None, } } pub fn max_entry_bytes(&self) -> Option { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } | Self::RedisSemantic(_) => None, + Self::Redis { .. } | Self::AzureBlob(_) | Self::RedisSemantic(_) => None, } } pub fn index_name(&self) -> Option<&str> { match self { Self::RedisSemantic(cache) => Some(cache.backend().index_name()), - Self::Memory(_) | Self::Redis { .. } => None, + Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) => None, } } pub fn similarity_threshold(&self) -> Option { match self { Self::RedisSemantic(cache) => Some(cache.backend().similarity_threshold()), - Self::Memory(_) | Self::Redis { .. } => None, + Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) => None, } } pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> { match self { Self::RedisSemantic(cache) => Some(cache.backend().embedder()), - Self::Memory(_) | Self::Redis { .. } => None, + Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) => None, } } pub fn embedder_object(&self) -> Option<&Py> { match self { Self::RedisSemantic(cache) => Some(cache.backend().embedder().object()), - Self::Memory(_) | Self::Redis { .. } => None, + Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) => None, } } @@ -162,6 +188,7 @@ impl NativeResponseCache { Self::Memory(cache) => cache.lookup(&request.exact(), now), Self::Redis { cache, .. } => cache.lookup(&request.exact(), now), Self::RedisSemantic(cache) => cache.lookup(&request.semantic(), now), + Self::AzureBlob(cache) => cache.lookup(&request.exact(), now), } } @@ -175,6 +202,7 @@ impl NativeResponseCache { Self::Memory(cache) => cache.store(&request.exact(), response, now), Self::Redis { cache, .. } => cache.store(&request.exact(), response, now), Self::RedisSemantic(cache) => cache.store(&request.semantic(), response, now), + Self::AzureBlob(cache) => cache.store(&request.exact(), response, now), } } @@ -187,6 +215,7 @@ impl NativeResponseCache { Self::Memory(cache) => cache.lookup_batch(&Self::exact_requests(requests), now), Self::Redis { cache, .. } => cache.lookup_batch(&Self::exact_requests(requests), now), Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), + Self::AzureBlob(cache) => cache.lookup_batch(&Self::exact_requests(requests), now), } } @@ -199,6 +228,7 @@ impl NativeResponseCache { Self::Memory(cache) => cache.async_lookup(&request.exact(), now).await, Self::Redis { cache, .. } => cache.async_lookup(&request.exact(), now).await, Self::RedisSemantic(cache) => cache.async_lookup(&request.semantic(), now).await, + Self::AzureBlob(cache) => cache.async_lookup(&request.exact(), now).await, } } @@ -225,6 +255,7 @@ impl NativeResponseCache { Self::RedisSemantic(cache) => { cache.async_store(&request.semantic(), response, now).await } + Self::AzureBlob(cache) => cache.async_store(&request.exact(), response, now).await, } } @@ -245,6 +276,11 @@ impl NativeResponseCache { .await } Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), + Self::AzureBlob(cache) => { + cache + .async_lookup_batch(&Self::exact_requests(requests), now) + .await + } } } @@ -287,6 +323,17 @@ impl NativeResponseCache { ) .await } + Self::AzureBlob(cache) => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (request.exact(), value)) + .collect(), + now, + ) + .await + } } } @@ -300,6 +347,7 @@ impl NativeResponseCache { cache.async_flush().await } Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), + Self::AzureBlob(cache) => cache.async_flush().await, } } @@ -308,6 +356,7 @@ impl NativeResponseCache { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), + Self::AzureBlob(cache) => cache.test_connection().await, } } } diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml new file mode 100644 index 00000000000..96db7f235ef --- /dev/null +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-secrets-azure" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-azure.workspace = true +litellm-auth-types.workspace = true +litellm-secrets-types.workspace = true +litellm-core-utils.workspace = true +reqwest.workspace = true +serde.workspace = true +thiserror.workspace = true +veil.workspace = true +percent-encoding = "2.3" + +[dev-dependencies] +tokio.workspace = true +wiremock = "0.6.5" +rstest.workspace = true +serde_json.workspace = true +sha2.workspace = true diff --git a/litellm-rust/crates/secrets-azure/src/error.rs b/litellm-rust/crates/secrets-azure/src/error.rs new file mode 100644 index 00000000000..9b20efe4f7c --- /dev/null +++ b/litellm-rust/crates/secrets-azure/src/error.rs @@ -0,0 +1,25 @@ +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("{0} environment variable is missing")] + MissingEnvironment(&'static str), + #[error("AZURE_KEY_VAULT_URI is not a valid https vault URL")] + VaultUri, + #[error("Azure Key Vault credentials are not configured")] + MissingCredentials, + #[error(transparent)] + Auth( + #[from] + #[redact] + litellm_auth_types::Error, + ), + #[error("Azure Key Vault request failed")] + Http( + #[source] + #[redact] + reqwest::Error, + ), + #[error("Azure Key Vault returned HTTP {0}")] + Status(u16), + #[error("Azure Key Vault response is missing the secret value")] + MissingValue, +} diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs new file mode 100644 index 00000000000..e12289b83f5 --- /dev/null +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -0,0 +1,118 @@ +use std::sync::Arc; + +use litellm_auth_azure::{AzureAuthInputs, AzureAuthService, ConfigValue}; +use litellm_auth_types::{InputSource, Sourced}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{Secret, SecretValue}; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC}; +use serde::Deserialize; + +use crate::Error; + +const AZURE_KEY_VAULT_URI: &str = "AZURE_KEY_VAULT_URI"; +const API_VERSION: &str = "7.4"; +const PATH_SEGMENT: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'.') + .remove(b'_') + .remove(b'~'); + +#[derive(Clone)] +pub struct AzureKeyVault { + client: reqwest::Client, + vault: reqwest::Url, + auth: Arc, + inputs: Arc, + environment: Arc, +} + +#[derive(Deserialize)] +struct SecretResponse { + value: Option, +} + +impl AzureKeyVault { + pub fn with_client( + client: reqwest::Client, + vault: reqwest::Url, + environment: Arc, + ) -> Result { + if vault.host_str().is_none() { + return Err(Error::VaultUri); + } + let inputs = AzureAuthInputs { + azure_scope: ConfigValue::Value(Sourced::new( + scope_for(&vault), + InputSource::Deployment, + )), + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..AzureAuthInputs::default() + }; + Ok(Self { + client, + vault, + auth: Arc::new(AzureAuthService::default()), + inputs: Arc::new(inputs), + environment, + }) + } + + pub fn new(environment: Arc) -> Result { + let value = environment + .get(AZURE_KEY_VAULT_URI) + .ok_or(Error::MissingEnvironment(AZURE_KEY_VAULT_URI))?; + let vault = reqwest::Url::parse(&value).map_err(|_| Error::VaultUri)?; + if vault.scheme() != "https" || vault.host_str().is_none() { + return Err(Error::VaultUri); + } + Self::with_client(reqwest::Client::new(), vault, environment) + } + + pub fn scope(&self) -> &str { + self.inputs + .azure_scope + .as_value() + .map(|value| value.value().as_str()) + .unwrap_or_default() + } + + pub async fn get_secret_from_azure_key_vault( + &self, + name: &str, + ) -> Result, Error> { + let token = self + .auth + .get_azure_ad_token(&self.inputs, &|key| self.environment.get(key)) + .await? + .ok_or(Error::MissingCredentials)?; + let encoded_name = percent_encoding::utf8_percent_encode(name, PATH_SEGMENT); + let url = self + .vault + .join(&format!("secrets/{encoded_name}?api-version={API_VERSION}")) + .map_err(|_| Error::VaultUri)?; + let response = self + .client + .get(url) + .bearer_auth(token.value().secret().expose()) + .send() + .await + .map_err(Error::Http)?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if response.status() != reqwest::StatusCode::OK { + return Err(Error::Status(response.status().as_u16())); + } + let payload: SecretResponse = response.json().await.map_err(Error::Http)?; + let value = payload.value.ok_or(Error::MissingValue)?; + Ok(Some(Secret::String(SecretValue::new(value)))) + } +} + +fn scope_for(vault: &reqwest::Url) -> String { + let host = vault.host_str().unwrap_or_default(); + let resource = host + .split_once('.') + .map_or(host, |(_, remainder)| remainder); + format!("https://{resource}/.default") +} diff --git a/litellm-rust/crates/secrets-azure/src/lib.rs b/litellm-rust/crates/secrets-azure/src/lib.rs new file mode 100644 index 00000000000..c0094fc033b --- /dev/null +++ b/litellm-rust/crates/secrets-azure/src/lib.rs @@ -0,0 +1,7 @@ +#![forbid(unsafe_code)] + +mod error; +mod key_vault; + +pub use error::Error; +pub use key_vault::AzureKeyVault; diff --git a/litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json b/litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json new file mode 100644 index 00000000000..c4a83cd150a --- /dev/null +++ b/litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json @@ -0,0 +1,8 @@ +{ + "cases": [ + {"name": "plain_value", "secret_name": "OPENAI-API-KEY", "response": {"status": 200, "body": {"value": "sk-parity-1", "id": "https://example.vault.azure.net/secrets/OPENAI-API-KEY/abc"}}, "expected": {"value": "sk-parity-1"}}, + {"name": "json_value_is_kept_as_string", "secret_name": "JSON-SECRET", "response": {"status": 200, "body": {"value": "{\"api_key\": \"nested\"}", "id": "https://example.vault.azure.net/secrets/JSON-SECRET/abc"}}, "expected": {"value": "{\"api_key\": \"nested\"}"}}, + {"name": "missing_secret", "secret_name": "MISSING", "response": {"status": 404, "body": {"error": {"code": "SecretNotFound", "message": "not found"}}}, "expected": {"missing": true}}, + {"name": "forbidden", "secret_name": "FORBIDDEN", "response": {"status": 403, "body": {"error": {"code": "Forbidden", "message": "denied"}}}, "expected": {"error": true}} + ] +} diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs new file mode 100644 index 00000000000..cf9102d0b45 --- /dev/null +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -0,0 +1,222 @@ +use std::sync::Arc; + +use litellm_secrets_azure::{AzureKeyVault, Error}; +use litellm_secrets_types::{Secret, SecretValue}; +use serde::Deserialize; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, path, query_param}, +}; + +fn manager(server: &MockServer) -> AzureKeyVault { + AzureKeyVault::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), + ) + .unwrap() +} + +#[tokio::test] +async fn reads_secret_with_bearer_token_and_api_version() { + let server = MockServer::start().await; + Mock::given(path("/secrets/OPENAI-API-KEY")) + .and(query_param("api-version", "7.4")) + .and(header("authorization", "Bearer fake")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"value": "s3cret", "id": "secret-id"})), + ) + .expect(1) + .mount(&server) + .await; + + let secret = manager(&server) + .get_secret_from_azure_key_vault("OPENAI-API-KEY") + .await + .unwrap() + .unwrap(); + + assert_eq!(secret, Secret::String(SecretValue::new("s3cret"))); +} + +#[tokio::test] +async fn percent_encodes_secret_name_path_segment() { + let server = MockServer::start().await; + Mock::given(path("/secrets/name%2Fwith%20spaces")) + .and(query_param("api-version", "7.4")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"value": "value"})), + ) + .expect(1) + .mount(&server) + .await; + + let secret = manager(&server) + .get_secret_from_azure_key_vault("name/with spaces") + .await + .unwrap() + .unwrap(); + + assert_eq!(secret.as_str(), Some("value")); +} + +#[rstest::rstest] +#[case::not_found(404, None)] +#[case::forbidden(403, Some(403))] +#[tokio::test] +async fn handles_statuses(#[case] status: u16, #[case] expected_status: Option) { + let server = MockServer::start().await; + Mock::given(path("/secrets/NAME")) + .respond_with(ResponseTemplate::new(status)) + .expect(1) + .mount(&server) + .await; + + let result = manager(&server) + .get_secret_from_azure_key_vault("NAME") + .await; + + match expected_status { + None => assert_eq!(result.unwrap(), None), + Some(status) => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)), + } +} + +#[tokio::test] +async fn missing_value_is_an_error() { + let server = MockServer::start().await; + Mock::given(path("/secrets/NAME")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(1) + .mount(&server) + .await; + + assert!(matches!( + manager(&server) + .get_secret_from_azure_key_vault("NAME") + .await, + Err(Error::MissingValue) + )); +} + +#[test] +fn new_validates_vault_environment() { + assert!(matches!( + AzureKeyVault::new(Arc::new(|_: &str| None)), + Err(Error::MissingEnvironment("AZURE_KEY_VAULT_URI")) + )); + assert!(matches!( + AzureKeyVault::new(Arc::new(|name: &str| { + (name == "AZURE_KEY_VAULT_URI").then(|| "http://vault.example".to_owned()) + })), + Err(Error::VaultUri) + )); + assert!(matches!( + AzureKeyVault::new(Arc::new(|name: &str| { + (name == "AZURE_KEY_VAULT_URI").then(|| "vault.example".to_owned()) + })), + Err(Error::VaultUri) + )); +} + +#[rstest::rstest] +#[case("https://myvault.vault.azure.net", "https://vault.azure.net/.default")] +#[case( + "https://v.vault.usgovcloudapi.net/", + "https://vault.usgovcloudapi.net/.default" +)] +#[case("http://localhost:8080", "https://localhost/.default")] +#[test] +fn derives_scope_from_vault_host(#[case] uri: &str, #[case] expected: &str) { + let manager = AzureKeyVault::with_client( + reqwest::Client::new(), + uri.parse().unwrap(), + Arc::new(|_: &str| None), + ) + .unwrap(); + + assert_eq!(manager.scope(), expected); +} + +#[tokio::test] +async fn missing_credentials_do_not_request_vault() { + let server = MockServer::start().await; + Mock::given(path("/secrets/NAME")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + + assert!( + manager_without_credentials(&server) + .get_secret_from_azure_key_vault("NAME") + .await + .is_err() + ); +} + +fn manager_without_credentials(server: &MockServer) -> AzureKeyVault { + AzureKeyVault::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + Arc::new(|name: &str| { + (name == "AZURE_CREDENTIAL").then(|| "ClientSecretCredential".to_owned()) + }), + ) + .unwrap() +} + +#[derive(Deserialize)] +struct Fixture { + cases: Vec, +} + +#[derive(Deserialize)] +struct FixtureCase { + secret_name: String, + response: FixtureResponse, + expected: FixtureExpected, +} + +#[derive(Deserialize)] +struct FixtureResponse { + status: u16, + body: serde_json::Value, +} + +#[derive(Deserialize)] +struct FixtureExpected { + value: Option, + missing: Option, + error: Option, +} + +#[tokio::test] +async fn parity_fixture_matches_python_backend_contract() { + let fixture: Fixture = + serde_json::from_str(include_str!("fixtures/key_vault_parity.json")).unwrap(); + for case in fixture.cases { + let server = MockServer::start().await; + Mock::given(path(format!("/secrets/{}", case.secret_name))) + .respond_with( + ResponseTemplate::new(case.response.status).set_body_json(case.response.body), + ) + .expect(1) + .mount(&server) + .await; + let result = manager(&server) + .get_secret_from_azure_key_vault(&case.secret_name) + .await; + if case.expected.missing == Some(true) { + assert_eq!(result.unwrap(), None); + } else if case.expected.error == Some(true) { + assert!(result.is_err()); + } else { + assert_eq!( + result.unwrap().unwrap().as_str(), + case.expected.value.as_deref() + ); + } + } +} diff --git a/litellm-rust/crates/secrets-azure/tests/live.rs b/litellm-rust/crates/secrets-azure/tests/live.rs new file mode 100644 index 00000000000..a062ba95070 --- /dev/null +++ b/litellm-rust/crates/secrets-azure/tests/live.rs @@ -0,0 +1,30 @@ +use std::sync::Arc; + +use litellm_core_utils::settings::ProcessEnvironment; +use litellm_secrets_azure::AzureKeyVault; +use litellm_secrets_types::Secret; + +#[tokio::test] +#[ignore] +async fn reads_a_real_secret() { + let environment = Arc::new(ProcessEnvironment); + let manager = AzureKeyVault::new(environment).unwrap(); + let name = std::env::var("AZURE_KEY_VAULT_LIVE_SECRET_NAME").unwrap(); + let secret = manager + .get_secret_from_azure_key_vault(&name) + .await + .unwrap() + .unwrap(); + assert!(matches!(&secret, Secret::String(_))); + let host = std::env::var("AZURE_KEY_VAULT_URI") + .unwrap() + .parse::() + .unwrap() + .host_str() + .unwrap() + .to_owned(); + let value_len = secret.as_str().unwrap().len(); + println!( + "native provider=litellm-secrets-azure vault_host={host} secret={name} value_len={value_len}" + ); +} diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 1298cbd33fe..962f66c92d1 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -9,12 +9,14 @@ repository.workspace = true default = [] aws = ["dep:litellm-secrets-aws"] google = ["dep:litellm-secrets-google"] +azure = ["dep:litellm-secrets-azure"] cyberark = ["dep:litellm-secrets-cyberark"] [dependencies] litellm-secrets-types.workspace = true litellm-secrets-aws = { workspace = true, optional = true } litellm-secrets-google = { workspace = true, optional = true } +litellm-secrets-azure = { workspace = true, optional = true } litellm-secrets-cyberark = { workspace = true, optional = true } litellm-core-utils.workspace = true base64.workspace = true diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index cf08c7c2abb..7e03f1f8cbf 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -30,6 +30,9 @@ pub enum Error { #[cfg(feature = "google")] #[error(transparent)] Google(#[from] litellm_secrets_google::Error), + #[cfg(feature = "azure")] + #[error(transparent)] + Azure(#[from] litellm_secrets_azure::Error), #[cfg(feature = "cyberark")] #[error(transparent)] Cyberark(#[from] litellm_secrets_cyberark::Error), diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 190a67b7671..71214ffc9ce 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -13,6 +13,8 @@ pub enum SecretManager { GoogleKms(crate::google::GoogleKms), #[cfg(feature = "google")] GoogleSecretManager(crate::google::GoogleSecretManager), + #[cfg(feature = "azure")] + AzureKeyVault(crate::azure::AzureKeyVault), #[cfg(feature = "cyberark")] Cyberark(crate::cyberark::CyberArkSecretManager), } @@ -29,6 +31,8 @@ impl SecretManager { Self::GoogleKms(_) => KeyManagementSystem::GoogleKms, #[cfg(feature = "google")] Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager, + #[cfg(feature = "azure")] + Self::AzureKeyVault(_) => KeyManagementSystem::AzureKeyVault, #[cfg(feature = "cyberark")] Self::Cyberark(_) => KeyManagementSystem::Cyberark, } @@ -82,6 +86,11 @@ pub async fn get_secret_from_manager( .get_secret_from_google_secret_manager(secret_name) .await .map_err(Error::from), + #[cfg(feature = "azure")] + SecretManager::AzureKeyVault(client) => client + .get_secret_from_azure_key_vault(secret_name) + .await + .map_err(Error::from), #[cfg(feature = "cyberark")] SecretManager::Cyberark(client) => client .async_read_secret(secret_name) diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index e132627345f..dec924abdd0 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -17,6 +17,8 @@ pub use state::{SecretManagerState, secret_manager_would_be_consulted}; #[cfg(feature = "aws")] pub use litellm_secrets_aws as aws; +#[cfg(feature = "azure")] +pub use litellm_secrets_azure as azure; #[cfg(feature = "cyberark")] pub use litellm_secrets_cyberark as cyberark; #[cfg(feature = "google")] diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs index 8f996cad63f..fbe13721491 100644 --- a/litellm-rust/crates/secrets/tests/handler.rs +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -106,6 +106,67 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites )); } +#[cfg(feature = "azure")] +#[tokio::test] +async fn azure_handler_reads_missing_and_failed_secrets() { + use litellm_secrets::{ + Error, KeyManagementSettings, KeyManagementSystem, SecretManager, azure::AzureKeyVault, + get_secret_from_manager, + }; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{path, query_param}, + }; + + let server = MockServer::start().await; + Mock::given(path("/secrets/KEY")) + .and(query_param("api-version", "7.4")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"value": "value"})), + ) + .expect(1) + .mount(&server) + .await; + let manager = SecretManager::AzureKeyVault( + AzureKeyVault::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + std::sync::Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), + ) + .unwrap(), + ); + assert_eq!(manager.system(), KeyManagementSystem::AzureKeyVault); + let settings = KeyManagementSettings::default(); + let value = get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None) + .await + .unwrap() + .unwrap(); + assert_eq!(value.as_str(), Some("value")); + + let not_found = Mock::given(path("/secrets/MISSING")) + .respond_with(ResponseTemplate::new(404)) + .expect(1) + .mount_as_scoped(&server) + .await; + assert_eq!( + get_secret_from_manager(&manager, "MISSING", &settings, &|_: &str| None) + .await + .unwrap(), + None + ); + drop(not_found); + + Mock::given(path("/secrets/FAILED")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&server) + .await; + assert!(matches!( + get_secret_from_manager(&manager, "FAILED", &settings, &|_: &str| None).await, + Err(Error::Azure(_)) + )); +} + #[cfg(feature = "cyberark")] #[tokio::test] async fn cyberark_handler_reads_values_and_surfaces_errors() { diff --git a/litellm/integrations/SlackAlerting/batching_handler.py b/litellm/integrations/SlackAlerting/batching_handler.py index 1c35a15d5a1..d152985a2c5 100644 --- a/litellm/integrations/SlackAlerting/batching_handler.py +++ b/litellm/integrations/SlackAlerting/batching_handler.py @@ -1,14 +1,18 @@ """ Handles Batching + sending Httpx Post requests to slack -Slack alerts are sent every 10s or when events are greater than X events +Slack alerts are sent every DEFAULT_FLUSH_INTERVAL_SECONDS or when events are greater than X events see custom_batch_logger.py for more details / defaults """ +from collections import Counter +from collections.abc import Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger +from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType from .ms_teams import MS_TEAMS_ALERTING_DESTINATION, build_ms_teams_payload @@ -20,26 +24,20 @@ else: SlackAlertingType = Any -def squash_payloads(queue): - squashed: Final = {} - if len(queue) == 0: - return squashed - if len(queue) == 1: - return {"key": {"item": queue[0], "count": 1}} +@dataclass(frozen=True, slots=True) +class SquashedAlert: + item: AlertQueueItem + count: int - for item in queue: - url = item["url"] - alert_type = item["alert_type"] - _key = (url, alert_type) - if _key in squashed: - squashed[_key]["count"] += 1 - # Merge the payloads +def _squash_key(item: AlertQueueItem) -> tuple[str, AlertType | str, str]: + return (item["url"], item["alert_type"], item["payload"]["text"]) - else: - squashed[_key] = {"item": item, "count": 1} - return squashed +def squash_payloads(queue: Sequence[AlertQueueItem]) -> tuple[SquashedAlert, ...]: + counts: Final = Counter(_squash_key(item) for item in queue) + first_item_by_key: Final = {_squash_key(item): item for item in reversed(queue)} + return tuple(SquashedAlert(item=first_item_by_key[key], count=count) for key, count in counts.items()) def _print_alerting_payload_warning(payload: dict, slackAlertingInstance: SlackAlertingType): @@ -53,17 +51,15 @@ def _print_alerting_payload_warning(payload: dict, slackAlertingInstance: SlackA verbose_proxy_logger.warning(payload) -async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count): +async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item: AlertQueueItem, count: int) -> None: """ Send a single slack alert to the webhook """ import json - payload: Final = item.get("payload", {}) + text: Final = item["payload"]["text"] + payload: Final = {"text": text if count == 1 else f"[Num Alerts: {count}]\n\n{text}"} try: - if count > 1: - payload["text"] = f"[Num Alerts: {count}]\n\n{payload['text']}" - request_body: Final = ( build_ms_teams_payload(payload["text"]) if item.get("format") == MS_TEAMS_ALERTING_DESTINATION else payload ) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 66e2754d5ad..8d0d044ff93 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import ( _add_key_name_and_team_to_alert, ) from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) @@ -99,6 +100,7 @@ class SlackAlerting(CustomBatchLogger): alerting_args={}, default_webhook_url: str | None = None, alert_type_config: dict[str, dict] | None = None, + async_http_handler: AsyncHTTPHandler | None = None, **kwargs, ): if alerting_threshold is None: @@ -107,7 +109,9 @@ class SlackAlerting(CustomBatchLogger): self.alerting = alerting self.alert_types = alert_types self.internal_usage_cache = internal_usage_cache or DualCache() - self.async_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + self.async_http_handler = async_http_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) self.alert_to_webhook_url = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) self.is_running = False self.alerting_args = SlackAlertingArgs(**alerting_args) @@ -1583,12 +1587,12 @@ Model Info: if not self.log_queue: return - squashed_queue: Final = squash_payloads(self.log_queue) - tasks: Final = [ - send_to_webhook(slackAlertingInstance=self, item=item["item"], count=item["count"]) - for item in squashed_queue.values() - ] - await asyncio.gather(*tasks) + await asyncio.gather( + *( + send_to_webhook(slackAlertingInstance=self, item=squashed.item, count=squashed.count) + for squashed in squash_payloads(self.log_queue) + ) + ) self.log_queue.clear() async def _flush_digest_buckets(self): diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py index 390abf41955..70b2cd08b4c 100644 --- a/litellm/litellm_core_utils/secret_redaction.py +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -65,7 +65,7 @@ def _build_secret_patterns() -> "re.Pattern[str]": # private_key with PEM-aware value capture r"""private_key['\"]?\s*[:=]\s*['\"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'\"})\]{}>]+)""", r"(?:master_key|xai_key|database_url|db_url|connection_string|" - r"aws_secret_access_key|aws_session_token|aws_access_key_id|" + r"aws_secret_access_key|aws_session_token|aws_access_key_id|s3_secret_access_key|s3_access_key_id|" r"signing_key|encryption_key|" r"auth_token|access_token|refresh_token|" r"slack_webhook_url|webhook_url|" diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9d1a4065f31..c2279fb2fe1 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2297,23 +2297,19 @@ async def _load_team_membership_on_cache_miss( parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, ) -> LiteLLM_TeamMembership | None: - try: - redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) - redis_membership: Final = _membership_from_cached_payload(redis_cached) - if not isinstance(redis_membership, _TeamMembershipCacheMiss): - return redis_membership + redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + redis_membership: Final = _membership_from_cached_payload(redis_cached) + if not isinstance(redis_membership, _TeamMembershipCacheMiss): + return redis_membership - return await _fetch_team_membership_from_db( - user_id=user_id, - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception: - verbose_proxy_logger.exception("Error getting team membership") - return None + return await _fetch_team_membership_from_db( + user_id=user_id, + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) async def get_team_membership( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 1c986305c21..c6b89096ca9 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2518,6 +2518,7 @@ async def delete_user( ## DELETE USERS deleted_users: Final = await _user_table(prisma_client).delete_many(where={"user_id": {"in": data.user_ids}}) + await evict_and_broadcast(cache_keys=tuple(data.user_ids), user_api_key_cache=user_api_key_cache) return deleted_users diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index b676c0ddb82..3292a0141d1 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1871,6 +1871,10 @@ async def delete_user( # Delete user await _table(UserRepository(prisma_client)).delete(where={"user_id": user_id}) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache) + return Response(status_code=204) except Exception as e: raise handle_exception_on_proxy(e) diff --git a/litellm/types/integrations/slack_alerting.py b/litellm/types/integrations/slack_alerting.py index 64c0c530e9b..33bb446364e 100644 --- a/litellm/types/integrations/slack_alerting.py +++ b/litellm/types/integrations/slack_alerting.py @@ -1,11 +1,12 @@ import os import time +from collections.abc import Mapping from datetime import datetime as dt from enum import Enum from typing import Any, Final, Literal, Optional, Union from pydantic import BaseModel, Field -from typing_extensions import TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.types.utils import LiteLLMPydanticObjectBase @@ -235,6 +236,18 @@ DEFAULT_ALERT_TYPES: Final[list[AlertType]] = [ ] +class AlertText(TypedDict): + text: ReadOnly[str] + + +class AlertQueueItem(TypedDict): + url: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + payload: ReadOnly[AlertText] + alert_type: ReadOnly[AlertType | str] + format: NotRequired[ReadOnly[str]] + + class HangingRequestData(BaseModel): request_id: str model: str diff --git a/litellm/types/router.py b/litellm/types/router.py index fd426835d65..8ad01206811 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -305,6 +305,8 @@ class CredentialLiteLLMParams(BaseModel): s3_bucket_name: str | None = None s3_endpoint_url: str | None = None s3_region_name: str | None = None + s3_access_key_id: str | None = None + s3_secret_access_key: str | None = None s3_encryption_key_id: str | None = None s3_bucket_owner: str | None = None aws_batch_role_arn: str | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 142304223d0..e1d43b7fccb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3840,6 +3840,9 @@ bedrock_batch_litellm_params: Final = ( "s3_endpoint_url", "s3_output_bucket_name", "s3_bucket_owner", + "s3_access_key_id", + "s3_secret_access_key", + "s3_encryption_key_id", "bedrock_tags", ) diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 36fbd39154d..49d4d92ff0b 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -31,6 +31,7 @@ - {id: llm.chat_completions.bedrock_converse.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Anthropic thinking on Bedrock"} - {id: llm.chat_completions.bedrock_converse.response_headers.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: response_headers, streaming: nonstream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py:248", rationale: "Bedrock request ids must surface as llm_provider-* response headers on /chat/completions so callers can correlate calls with AWS-side logs (#37003)", fail_before_fix: proven} - {id: llm.chat_completions.bedrock_converse.response_headers.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: response_headers, streaming: stream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py:154", rationale: "The llm_provider-* headers must also surface on streaming /chat/completions, where CustomStreamWrapper carries them instead of the nonstream setter"} +- {id: llm.chat_completions.bedrock_converse.batch_deployment.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: batch_deployment, streaming: nonstream, assertions: [works], source: "types/utils.py bedrock_batch_litellm_params", rationale: "A deployment carrying the documented batch-only S3 keys (s3_access_key_id, s3_secret_access_key, s3_encryption_key_id) must still serve ordinary chat; unregistered keys fall into optional_params and are forwarded as additionalModelRequestFields, which Bedrock 400s and which puts the S3 secret in the request body and debug log (LIT-8290)", fail_before_fix: proven} - {id: llm.chat_completions.bedrock_invoke.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_invoke, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "Regional inference-profile ids (us.anthropic.*) over the invoke route, the deployment shape behind a customer timeout report on v1.90.0"} - {id: llm.chat_completions.bedrock_invoke.basic.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_invoke, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:8455", rationale: "Streaming with regional inference-profile ids over the invoke route"} - {id: llm.chat_completions.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "P0 route; Vertex AI"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index fa6dad90126..8b0d38a083c 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -63,6 +63,7 @@ LlmRoute = Literal[ LlmCapability = Literal[ "assume_role", "basic", + "batch_deployment", "count_tokens", "govcloud_partition", "input_validation", diff --git a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py index 3c6aaa75ab3..5f0a931109c 100644 --- a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py +++ b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py @@ -131,6 +131,50 @@ class TestBedrockResponseHeaders: _assert_request_id_header(result) +def _register_bedrock_batch_deployment(client: PassthroughClient, resources: ResourceManager) -> str: + model = f"e2e-bedrock-batch-chat-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=CONVERSE_REGIONAL_BACKEND, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + s3_bucket_name="os.environ/AWS_BATCH_S3_BUCKET", + s3_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + s3_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + s3_encryption_key_id=f"alias/e2e-unused-{unique_marker()}", + aws_batch_role_arn="os.environ/AWS_BATCH_ROLE_ARN", + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + +class TestBedrockBatchDeploymentServesChat: + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.batch_deployment.nonstream.works", + exercised_on=[], + ) + def test_batch_s3_keys_do_not_break_chat( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = _register_bedrock_batch_deployment(client, resources) + key = resources.key() + + result = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ChatBody(model=model, messages=_prompt(), max_tokens=64), + ) + + assert result.ok, ( + f"chat on a batch-configured deployment failed: {result.status_code} {result.body[:300]}; " + "batch-only S3 keys were forwarded to Bedrock as additionalModelRequestFields" + ) + _assert_completion(ChatResponse.model_validate_json(result.body)) + + class TestBedrockInvokeRegionalModelIds: @pytest.mark.covers("llm.chat_completions.bedrock_invoke.basic.nonstream.works", exercised_on=[]) def test_invoke_regional_id_completes( diff --git a/tests/e2e/migrations/conftest.py b/tests/e2e/migrations/conftest.py index 735adeedbdb..b7604a4fdda 100644 --- a/tests/e2e/migrations/conftest.py +++ b/tests/e2e/migrations/conftest.py @@ -60,3 +60,32 @@ def containers(migration_image: str, tmp_path: Path, request: SubRequest) -> Con output: Final = Path(configured) / request.node.name if configured else tmp_path output.mkdir(parents=True, exist_ok=True) return Containers(migration_image, output) + + +@pytest.fixture(scope="session") +def baseline_image(tmp_path_factory: pytest.TempPathFactory) -> str: + configured: Final = os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE") + assert configured, "LITELLM_MIGRATION_BASELINE_IMAGE must name the released image the upgrade starts from" + image: Final = docker("image", "inspect", configured, "--format", "{{.Id}}") + assert image.startswith("sha256:"), "Unable to identify the baseline image" + output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) + output.mkdir(parents=True, exist_ok=True) + (output / "baseline-image.json").write_text(json.dumps({"requested": configured, "image_id": image})) + return image + + +@pytest.fixture(scope="session") +def baseline_template( + databases: Databases, baseline_image: str, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Database]: + output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) / "baseline-seed" + with databases.create() as database: + with Containers(baseline_image, output).start(database) as replica: + ready((replica,), database) + yield database + + +@pytest.fixture +def baseline_database(databases: Databases, baseline_template: Database) -> Iterator[Database]: + with databases.create(baseline_template) as database: + yield database diff --git a/tests/e2e/migrations/containers.py b/tests/e2e/migrations/containers.py index 0f5793b81dd..dd126b994d3 100644 --- a/tests/e2e/migrations/containers.py +++ b/tests/e2e/migrations/containers.py @@ -5,7 +5,7 @@ import subprocess import time from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Final from uuid import uuid4 @@ -123,6 +123,9 @@ class Containers: image: str output: Path + def using(self, image: str) -> "Containers": + return replace(self, image=image) + @contextmanager def start( self, diff --git a/tests/e2e/migrations/test_rolling_upgrade.py b/tests/e2e/migrations/test_rolling_upgrade.py new file mode 100644 index 00000000000..5ad74e0ba8c --- /dev/null +++ b/tests/e2e/migrations/test_rolling_upgrade.py @@ -0,0 +1,57 @@ +from typing import Final + +import pytest + +from .containers import Containers, ready +from .database import Database +from .upgrade import ( + CACHED_PLAN, + assert_history_clean, + assert_upgraded, + auth_traffic, + confirm, + keep_serving, + migration_names, + provision, +) + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +class TestRollingUpgrade: + def test_baseline_replica_keeps_serving_while_the_candidate_migrates( + self, containers: Containers, baseline_image: str, baseline_database: Database + ) -> None: + with containers.using(baseline_image).start(baseline_database) as old: + ready((old,), baseline_database) + key, _ = provision(old) + before: Final = migration_names(baseline_database) + with auth_traffic(old, key) as traffic: + keep_serving(traffic, "the baseline replica authenticating before the upgrade") + with containers.start(baseline_database) as new: + ready((new,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + keep_serving(traffic, "the baseline replica authenticating after the schema moved") + with auth_traffic(old, provision(new)[0]) as uncached: + keep_serving(uncached, "the baseline replica resolving a key minted after the schema moved") + assert_history_clean(baseline_database) + assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement" + assert old.state().Running, "The baseline replica died during the upgrade" + + def test_both_releases_serve_and_share_keys_during_the_overlap( + self, containers: Containers, baseline_image: str, baseline_database: Database + ) -> None: + with containers.using(baseline_image).start(baseline_database) as old: + ready((old,), baseline_database) + old_key, old_alias = provision(old) + before: Final = migration_names(baseline_database) + with containers.start(baseline_database) as new: + ready((new,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + new_key, new_alias = provision(new) + with auth_traffic(old, old_key) as old_traffic, auth_traffic(new, new_key) as new_traffic: + keep_serving(old_traffic, "the baseline replica serving through the overlap") + keep_serving(new_traffic, "the candidate replica serving through the overlap") + confirm(old, new_key, new_alias) + confirm(new, old_key, old_alias) + assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement" diff --git a/tests/e2e/migrations/test_shaped_database.py b/tests/e2e/migrations/test_shaped_database.py new file mode 100644 index 00000000000..20c4368ae33 --- /dev/null +++ b/tests/e2e/migrations/test_shaped_database.py @@ -0,0 +1,43 @@ +from typing import Final + +import pytest + +from .containers import Containers, ready +from .database import Database +from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision + +SPEND_ROWS: Final = 20_000 + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +def seed_spend_logs(database: Database, rows: int) -> None: + database.execute( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime") ' + "SELECT 'upgrade-shape-' || g, 'acompletion', now() - (g || ' seconds')::interval, " + "now() - (g || ' seconds')::interval FROM generate_series(1, %s) AS g", + (rows,), + ) + assert database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((rows,),) + + +class TestPopulatedDatabaseUpgrade: + def test_upgrade_completes_and_preserves_a_populated_spend_log( + self, containers: Containers, baseline_image: str, baseline_database: Database + ) -> None: + with containers.using(baseline_image).start(baseline_database) as old: + ready((old,), baseline_database) + key, alias = provision(old) + seed_spend_logs(baseline_database, SPEND_ROWS) + before: Final = migration_names(baseline_database) + with containers.start(baseline_database) as new: + ready((new,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + confirm(new, key, alias) + assert_history_clean(baseline_database) + assert baseline_database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((SPEND_ROWS,),), ( + "The upgrade lost spend rows" + ) + assert baseline_database.query( + 'SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE "startTime" IS NULL OR "endTime" IS NULL' + ) == ((0,),), "The upgrade nulled timestamps on existing spend rows" diff --git a/tests/e2e/migrations/test_upgrade.py b/tests/e2e/migrations/test_upgrade.py new file mode 100644 index 00000000000..23f0bbe9124 --- /dev/null +++ b/tests/e2e/migrations/test_upgrade.py @@ -0,0 +1,47 @@ +from contextlib import ExitStack +from typing import Final + +import pytest + +from .checks import start_replicas +from .containers import Containers, ready +from .database import Database +from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +class TestReleaseUpgrade: + def test_candidate_applies_the_pending_release_migrations( + self, containers: Containers, baseline_database: Database + ) -> None: + before: Final = migration_names(baseline_database) + with containers.start(baseline_database) as replica: + ready((replica,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + assert_history_clean(baseline_database) + + def test_upgrade_preserves_keys_minted_by_the_baseline_release( + self, containers: Containers, baseline_image: str, baseline_database: Database + ) -> None: + with containers.using(baseline_image).start(baseline_database) as old: + ready((old,), baseline_database) + key, alias = provision(old) + confirm(old, key, alias) + before: Final = migration_names(baseline_database) + with containers.start(baseline_database) as new: + ready((new,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + confirm(new, key, alias) + + def test_concurrent_replicas_upgrade_a_baseline_database_once( + self, containers: Containers, baseline_database: Database + ) -> None: + before: Final = migration_names(baseline_database) + with ExitStack() as stack: + ready(start_replicas(stack, containers, baseline_database), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + assert_history_clean(baseline_database) + assert baseline_database.query("SELECT count(*) FROM _prisma_migrations WHERE applied_steps_count > 1") == ( + (0,), + ), "A migration was executed more than once across the upgrading replicas" diff --git a/tests/e2e/migrations/upgrade.py b/tests/e2e/migrations/upgrade.py new file mode 100644 index 00000000000..2123f86450a --- /dev/null +++ b/tests/e2e/migrations/upgrade.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +import threading +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Final +from uuid import uuid4 + +from e2e_http import Result, Success, unwrap +from models import ( + KeyGenerateBody, + KeyGenerateResponse, + KeyInfoParams, + KeyInfoResponse, + ModelsListParams, + ModelsListResponse, +) +from pydantic import BaseModel + +from .containers import Replica, until +from .database import Database + +CACHED_PLAN: Final = "cached plan must not change result type" + + +def provision(replica: Replica) -> tuple[str, str]: + alias: Final = f"upgrade-{uuid4().hex}" + key: Final = unwrap( + replica.transport.post( + "/key/generate", + headers=replica.transport.master, + json=KeyGenerateBody(key_alias=alias), + response_type=KeyGenerateResponse, + ) + ).key + return key, alias + + +def confirm(replica: Replica, key: str, alias: str) -> None: + info: Final = unwrap( + replica.transport.get( + "/key/info", + headers=replica.transport.master, + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + ) + assert info.info.key_alias == alias, "Key minted on one release did not resolve on the other" + + +@dataclass(slots=True) +class Outcomes: + served: int = 0 + failures: list[str] = field(default_factory=list) + + def record(self, result: Result[BaseModel]) -> None: + match result: + case Success(): + self.served += 1 + case _: + self.failures.append(result.model_dump_json()) + + +@contextmanager +def auth_traffic(replica: Replica, key: str, interval: float = 0.05) -> Generator[Outcomes]: + outcomes: Final = Outcomes() + stop: Final = threading.Event() + + def drive() -> None: + while not stop.is_set(): + outcomes.record( + replica.transport.get( + "/v1/models", + headers=replica.transport.bearer(key), + params=ModelsListParams(), + response_type=ModelsListResponse, + timeout=10, + ) + ) + stop.wait(interval) + + thread: Final = threading.Thread(target=drive, name="upgrade-auth-traffic", daemon=True) + thread.start() + try: + yield outcomes + finally: + stop.set() + thread.join(30) + assert not thread.is_alive(), "Auth traffic thread did not stop" + assert not outcomes.failures, ( + f"Virtual-key auth failed on {replica.name} after the traffic window closed: {outcomes.failures[:5]}" + ) + + +def keep_serving(outcomes: Outcomes, description: str, calls: int = 20) -> int: + target: Final = outcomes.served + calls + until(description, lambda: outcomes.served >= target or bool(outcomes.failures)) + assert not outcomes.failures, f"Virtual-key auth failed during {description}: {outcomes.failures[:5]}" + return outcomes.served + + +def migration_names(database: Database) -> frozenset[str]: + return frozenset(str(row[0]) for row in database.query("SELECT migration_name FROM _prisma_migrations")) + + +def assert_history_clean(database: Database) -> None: + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL" + ) == ((0,),), "The upgrade left an unfinished or rolled-back migration behind" + assert database.query( + "SELECT count(*) FROM (SELECT migration_name FROM _prisma_migrations GROUP BY migration_name " + "HAVING count(*) > 1) duplicated" + ) == ((0,),), "A migration was recorded more than once, so it ran on more than one replica" + + +def assert_upgraded(before: frozenset[str], after: frozenset[str]) -> frozenset[str]: + applied: Final = after - before + assert applied, ( + "The candidate applied no migrations the baseline release had not: the pinned " + "LITELLM_MIGRATION_BASELINE_IMAGE is at or ahead of the candidate, so this suite proves nothing" + ) + assert not before - after, "The upgrade removed migration history the baseline release had already applied" + return applied diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9440505e58c..355329585fb 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1011,6 +1011,7 @@ class LiteLLMParamsBody(BaseModel): s3_region_name: str | None = None s3_access_key_id: str | None = None s3_secret_access_key: str | None = None + s3_encryption_key_id: str | None = None aws_batch_role_arn: str | None = None aws_role_name: str | None = None aws_session_name: str | None = None diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index a8fce58c60b..f5e8d861d79 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -201,6 +201,14 @@ async def test_returned_user_api_key_auth(user_role, expected_role): assert new_obj.user_role == expected_role +class _NoMembershipRowPrisma: + class db: + class litellm_teammembership: + @staticmethod + async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None: + return None + + @pytest.mark.parametrize("key_ownership", ["user_key", "team_key"]) @pytest.mark.asyncio async def test_aaauser_personal_budgets(key_ownership): @@ -253,7 +261,7 @@ async def test_aaauser_personal_budgets(key_ownership): setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") - setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world") + setattr(litellm.proxy.proxy_server, "prisma_client", _NoMembershipRowPrisma()) request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") diff --git a/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py b/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py index 41b7f3b969b..44426c00628 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py @@ -2,18 +2,39 @@ import json from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +from pydantic import TypeAdapter from litellm.integrations.SlackAlerting.batching_handler import send_to_webhook from litellm.integrations.SlackAlerting.ms_teams import ( MS_TEAMS_ALERTING_DESTINATION, MS_TEAMS_WEBHOOK_URL_ENV, + MSTeamsMessage, build_ms_teams_payload, get_ms_teams_webhook_url, ) from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import AlertType +_MS_TEAMS_MESSAGE: Final = TypeAdapter(MSTeamsMessage) + + +def _webhook_accepting_posts() -> AsyncMock: + response: Final = MagicMock(spec=httpx.Response) + response.status_code = 200 + http_handler: Final = AsyncMock(spec=AsyncHTTPHandler) + http_handler.post.return_value = response + return http_handler + + +def _posted_card_texts(http_handler: AsyncMock) -> tuple[str, ...]: + return tuple( + _MS_TEAMS_MESSAGE.validate_json(call.kwargs["data"])["attachments"][0]["content"]["body"][0]["text"] + for call in http_handler.post.call_args_list + ) + def test_build_ms_teams_payload_wraps_text_in_adaptive_card(): payload: Final = build_ms_teams_payload("hello alert") @@ -80,11 +101,8 @@ async def test_send_alert_slack_and_ms_teams_enqueue_both(monkeypatch): @pytest.mark.asyncio async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items(): - slack_alerting: Final = SlackAlerting(alerting=["ms_teams"]) - mock_response: Final = MagicMock() - mock_response.status_code = 200 - slack_alerting.async_http_handler = MagicMock() - slack_alerting.async_http_handler.post = AsyncMock(return_value=mock_response) + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting(alerting=["ms_teams"], async_http_handler=http_handler) item: Final = { "url": "https://teams.example/webhook", @@ -95,7 +113,7 @@ async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items(): } await send_to_webhook(slackAlertingInstance=slack_alerting, item=item, count=1) - call_kwargs: Final = slack_alerting.async_http_handler.post.call_args.kwargs + call_kwargs: Final = http_handler.post.call_args.kwargs assert call_kwargs["url"] == "https://teams.example/webhook" sent_body: Final = json.loads(call_kwargs["data"]) assert sent_body["type"] == "message" @@ -104,11 +122,8 @@ async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items(): @pytest.mark.asyncio async def test_send_to_webhook_keeps_slack_payload_shape(): - slack_alerting: Final = SlackAlerting(alerting=["slack"]) - mock_response: Final = MagicMock() - mock_response.status_code = 200 - slack_alerting.async_http_handler = MagicMock() - slack_alerting.async_http_handler.post = AsyncMock(return_value=mock_response) + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting(alerting=["slack"], async_http_handler=http_handler) item: Final = { "url": "https://hooks.slack.com/services/test", @@ -118,5 +133,27 @@ async def test_send_to_webhook_keeps_slack_payload_shape(): } await send_to_webhook(slackAlertingInstance=slack_alerting, item=item, count=1) - call_kwargs: Final = slack_alerting.async_http_handler.post.call_args.kwargs + call_kwargs: Final = http_handler.post.call_args.kwargs assert json.loads(call_kwargs["data"]) == {"text": "alert body"} + + +@pytest.mark.asyncio +async def test_async_send_batch_delivers_every_distinct_ms_teams_alert(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(MS_TEAMS_WEBHOOK_URL_ENV, "https://teams.example/webhook") + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting(alerting=["ms_teams"], async_http_handler=http_handler) + slack_alerting.periodic_started = True + for message in ("User Budget: 15% or less of budget remaining", "User Budget: Budget Crossed"): + await slack_alerting.send_alert( + message=message, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + await slack_alerting.async_send_batch() + + card_texts: Final = _posted_card_texts(http_handler) + assert len(card_texts) == 2 + assert "User Budget: 15% or less of budget remaining" in card_texts[0] + assert "User Budget: Budget Crossed" in card_texts[1] diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 4bc6c08bd63..2d5eb78950c 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -6,14 +6,18 @@ import unittest from typing import Final, List, Optional, Tuple from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch +import httpx import pytest +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.caching.caching import DualCache from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import CallInfo, Litellm_EntityType -from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingCacheKeys +from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType, SlackAlertingCacheKeys class TestSlackAlerting(unittest.TestCase): @@ -434,3 +438,91 @@ async def test_send_alert_raises_when_no_webhook_url_configured(monkeypatch): alert_type=AlertType.budget_alerts, alerting_metadata={}, ) + + +SLACK_WEBHOOK_URL: Final = "https://hooks.slack.com/services/test" +THRESHOLD_ALERT: Final = "User Budget: 15% or less of budget remaining\n\n*user_id:* `user-a`" +CROSSED_ALERT: Final = "User Budget: Budget Crossed\n\n*user_id:* `user-b`" + + +class _SlackWebhookBody(TypedDict): + text: ReadOnly[str] + + +_SLACK_WEBHOOK_BODY: Final = TypeAdapter(_SlackWebhookBody) + + +def _webhook_accepting_posts() -> AsyncMock: + response: Final = MagicMock(spec=httpx.Response) + response.status_code = 200 + http_handler: Final = AsyncMock(spec=AsyncHTTPHandler) + http_handler.post.return_value = response + return http_handler + + +def _slack_alerting_flushing_to(http_handler: AsyncHTTPHandler) -> SlackAlerting: + slack_alerting: Final = SlackAlerting(alerting=["slack"], async_http_handler=http_handler) + slack_alerting.periodic_started = True + return slack_alerting + + +def _queued_slack_alert(text: str) -> AlertQueueItem: + return { + "url": SLACK_WEBHOOK_URL, + "headers": {"Content-type": "application/json"}, + "payload": {"text": text}, + "alert_type": AlertType.budget_alerts, + } + + +def _posted_slack_bodies(http_handler: AsyncMock) -> tuple[_SlackWebhookBody, ...]: + return tuple(_SLACK_WEBHOOK_BODY.validate_json(call.kwargs["data"]) for call in http_handler.post.call_args_list) + + +async def _send_budget_alert(slack_alerting: SlackAlerting, message: str) -> None: + await slack_alerting.send_alert( + message=message, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + +@pytest.mark.asyncio +async def test_async_send_batch_delivers_every_distinct_alert_queued_in_one_flush( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SLACK_WEBHOOK_URL", SLACK_WEBHOOK_URL) + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = _slack_alerting_flushing_to(http_handler) + await _send_budget_alert(slack_alerting, THRESHOLD_ALERT) + await _send_budget_alert(slack_alerting, CROSSED_ALERT) + + await slack_alerting.async_send_batch() + + posted_texts: Final = tuple(body["text"] for body in _posted_slack_bodies(http_handler)) + assert len(posted_texts) == 2 + assert THRESHOLD_ALERT in posted_texts[0] + assert CROSSED_ALERT in posted_texts[1] + assert not any(text.startswith("[Num Alerts") for text in posted_texts) + assert slack_alerting.log_queue == [] + + +@pytest.mark.asyncio +async def test_async_send_batch_collapses_only_identical_alerts() -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = _slack_alerting_flushing_to(http_handler) + slack_alerting.log_queue.extend( + ( + _queued_slack_alert(THRESHOLD_ALERT), + _queued_slack_alert(CROSSED_ALERT), + _queued_slack_alert(THRESHOLD_ALERT), + ) + ) + + await slack_alerting.async_send_batch() + + assert _posted_slack_bodies(http_handler) == ( + {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"}, + {"text": CROSSED_ALERT}, + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index b0cda30dfe5..c1d5cedeba0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11515,7 +11515,9 @@ def jwt_oauth_identity(monkeypatch: pytest.MonkeyPatch) -> tuple["JWTHandler", " monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True}) monkeypatch.setattr(proxy_server, "premium_user", True) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) - monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + prisma: Final = MagicMock() + prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) return handler, signing_key diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 2824708d502..44c1d6a3c6b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7300,7 +7300,7 @@ async def test_common_checks_skips_membership_load_when_no_check_reads_it(): @pytest.mark.asyncio -async def test_get_team_membership_db_error_returns_none_and_retries_next_call(): +async def test_get_team_membership_db_error_surfaces_and_retries_next_call(): from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key @@ -7312,12 +7312,13 @@ async def test_get_team_membership_db_error_returns_none_and_retries_next_call() ) cache = UserApiKeyCache() - failed = await get_team_membership( - user_id="u-fail", - team_id="t-fail", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ) + with pytest.raises(RuntimeError, match="db down"): + await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) cached_after_failure = await cache.async_get_cache( key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") ) @@ -7328,24 +7329,52 @@ async def test_get_team_membership_db_error_returns_none_and_retries_next_call() user_api_key_cache=cache, ) - assert failed is None assert cached_after_failure is None assert recovered is not None assert recovered.user_id == "u-fail" assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 -@pytest.mark.asyncio -async def test_get_team_membership_string_prisma_client_returns_none(): - from litellm.proxy.auth.auth_checks import get_team_membership +class _UnreachableMembershipPrisma: + class db: + class litellm_teammembership: + @staticmethod + async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None: + raise httpx.ConnectError("All connection attempts failed") - result = await get_team_membership( - user_id="u-str", - team_id="t-str", - prisma_client="hello-world", - user_api_key_cache=UserApiKeyCache(), - ) - assert result is None + +def _restricted_member_check_deps() -> dict[str, object]: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + return { + "team_object": LiteLLM_TeamTable(team_id="team-outage", models=["claude-sonnet-5"]), + "valid_token": UserAPIKeyAuth(token="hashed-fake", user_id="bob", team_id="team-outage"), + "prisma_client": _UnreachableMembershipPrisma(), + "user_api_key_cache": cache, + "proxy_logging_obj": ProxyLogging(user_api_key_cache=cache), + } + + +@pytest.mark.asyncio +async def test_check_team_member_model_access_fails_closed_when_the_membership_read_hits_a_db_outage(): + from litellm.proxy.auth.auth_checks import _check_team_member_model_access + from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception + + with pytest.raises(httpx.ConnectError) as raised: + await _check_team_member_model_access( + model="claude-sonnet-5", llm_router=None, **_restricted_member_check_deps() + ) + + surfaced = _as_proxy_exception(raised.value) + assert (surfaced.code, surfaced.type) == ("503", ProxyErrorTypes.no_db_connection) + + +@pytest.mark.asyncio +async def test_check_team_member_budget_fails_closed_when_the_membership_read_hits_a_db_outage(): + with pytest.raises(httpx.ConnectError): + await _check_team_member_budget(user_object=None, **_restricted_member_check_deps()) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_resolvers_grants.py b/tests/test_litellm/proxy/auth/test_resolvers_grants.py index 3f6d943bf98..e61269ec1c7 100644 --- a/tests/test_litellm/proxy/auth/test_resolvers_grants.py +++ b/tests/test_litellm/proxy/auth/test_resolvers_grants.py @@ -1,4 +1,5 @@ from fastapi import HTTPException +import httpx import pytest from litellm.proxy._types import ( @@ -8,6 +9,7 @@ from litellm.proxy._types import ( ProxyException, ) from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.auth.resolvers.grants import ( GrantResolver, LookupDegraded, @@ -172,6 +174,29 @@ async def test_resolve_identity_lets_loader_errors_surface(): await loaders.resolver().resolve_identity(UserLookup(user_id=USER_ID), team_id=None) +class _UnreachableMembershipPrisma: + class db: + class litellm_teammembership: + @staticmethod + async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None: + raise httpx.ConnectError("All connection attempts failed") + + +async def test_resolve_marks_a_membership_read_that_hits_a_db_outage_as_degraded(): + loaders = _Loaders(user=_user(), team=_team()) + resolver = GrantResolver( + _UnreachableMembershipPrisma(), + UserApiKeyCache(), + load_user=loaders.load_user, + load_team=loaders.load_team, + ) + + outcome = await resolver.resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID) + + assert isinstance(outcome, LookupDegraded) + assert isinstance(outcome.error, httpx.ConnectError) + + def test_raise_public_maps_a_deleted_user_to_401(): with pytest.raises(ProxyException) as exc_info: raise_public(UserGone(user_id=USER_ID)) diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py index 2cfbfbbb3cb..d35a676a28b 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py @@ -601,3 +601,42 @@ async def test_scim_status_write_refreshes_user_cache( else: assert cached is None broadcast.assert_awaited_once_with(cache_key=user_id) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [None, "delete"]) +async def test_scim_delete_user_evicts_cached_user_row(failure: str | None) -> None: + from typing import Final + + from litellm.proxy._types import ProxyException + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + user_id: Final = "scim-deleted-user" + saved: Final = LiteLLM_UserTable(user_id=user_id, user_email="x@example.com", teams=[], metadata={}) + client, db = _build_prisma_with_keys([], mock_user=saved.model_copy(deep=True)) + if failure == "delete": + db.litellm_usertable.delete.side_effect = RuntimeError("user delete failed") + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=user_id, value=saved, model_type=LiteLLM_UserTable) + with ( + patch("litellm.proxy.proxy_server.prisma_client", client), # test-quality-ok: substitute the database dependency + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: exercise a real isolated cache + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: isolate the logging dependency + patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=AsyncMock, + ) as broadcast, + ): + if failure == "delete": + with pytest.raises(ProxyException, match="user delete failed"): + await delete_user(user_id=user_id) + else: + response: Final = await delete_user(user_id=user_id) + assert response.status_code == 204 + cached: Final = await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable) + if failure == "delete": + assert cached == saved + broadcast.assert_not_awaited() + else: + assert cached is None + broadcast.assert_awaited_once_with(cache_key=user_id) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index b8f1aa0330b..ede0ecb2790 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4691,3 +4691,42 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo written_data = mock_prisma_client.update_data.call_args.kwargs["data"] assert written_data.get("password") is not None assert written_data["password"] != strong_password + + +@pytest.mark.asyncio +async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> None: + from litellm.proxy._types import DeleteUserRequest, LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user + + deleted: Final = LiteLLM_UserTable(user_id="user-gone", user_email="gone@example.test", teams=[]) + survivor: Final = LiteLLM_UserTable(user_id="user-stays", user_email="stays@example.test", teams=[]) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=deleted) + prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_jwtkeymapping.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + cache: Final = UserApiKeyCache() + for row in (deleted, survivor): + await cache.async_set_cache(key=row.user_id, value=row, model_type=LiteLLM_UserTable) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + + await delete_user( + data=DeleteUserRequest(user_ids=[deleted.user_id]), + user_api_key_dict=UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert await cache.async_get_cache(key=deleted.user_id, model_type=LiteLLM_UserTable) is None + assert await cache.async_get_cache(key=survivor.user_id, model_type=LiteLLM_UserTable) == survivor + broadcast.assert_awaited_once_with(cache_key=deleted.user_id) diff --git a/tests/test_litellm/secret_managers/test_secret_manager_handler.py b/tests/test_litellm/secret_managers/test_secret_manager_handler.py new file mode 100644 index 00000000000..b4838912d27 --- /dev/null +++ b/tests/test_litellm/secret_managers/test_secret_manager_handler.py @@ -0,0 +1,107 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +from pydantic import BaseModel, ConfigDict + +from litellm.secret_managers.secret_manager_handler import get_secret_from_manager +from litellm.types.secret_managers.main import KeyManagementSystem + + +def _azure_exception_types() -> tuple[type[Exception], type[Exception]]: + try: + from azure.core.exceptions import ( + HttpResponseError, + ResourceNotFoundError, + ) + except ImportError: + return Exception, Exception + return HttpResponseError, ResourceNotFoundError + + +_AZURE_EXCEPTION_TYPES: Final[tuple[type[Exception], type[Exception]]] = _azure_exception_types() +AzureHttpResponseError: Final[type[Exception]] = _AZURE_EXCEPTION_TYPES[0] +AzureResourceNotFoundError: Final[type[Exception]] = _AZURE_EXCEPTION_TYPES[1] + + +class FixtureResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + status: int + body: dict[str, object] + + +class FixtureExpected(BaseModel): + model_config = ConfigDict(frozen=True) + + value: str | None = None + missing: bool = False + error: bool = False + + +class FixtureCase(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + secret_name: str + response: FixtureResponse + expected: FixtureExpected + + +class Fixture(BaseModel): + model_config = ConfigDict(frozen=True) + + cases: tuple[FixtureCase, ...] + + +@dataclass(frozen=True, slots=True) +class FakeSecret: + value: str | None + + +@dataclass(frozen=True, slots=True) +class FakeAzureKeyVaultClient: + status: int + value: str | None + + def get_secret(self, name: str) -> FakeSecret: + if self.status == 404: + raise AzureResourceNotFoundError() + if self.status != 200: + raise AzureHttpResponseError() + return FakeSecret(value=self.value) + + +FIXTURE_PATH: Path = ( + Path(__file__).parents[3] + / "litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json" +) + + +def test_azure_key_vault_matches_rust_parity_fixture() -> None: + fixture: Fixture = Fixture.model_validate_json(FIXTURE_PATH.read_text()) + for case in fixture.cases: + value: object = case.response.body.get("value") + secret: str | None = value if isinstance(value, str) else None + client: FakeAzureKeyVaultClient = FakeAzureKeyVaultClient( + status=case.response.status, + value=secret, + ) + if case.expected.missing or case.expected.error: + with pytest.raises( + AzureResourceNotFoundError if case.expected.missing else AzureHttpResponseError + ): + get_secret_from_manager( + secret_name=case.secret_name, + key_manager=KeyManagementSystem.AZURE_KEY_VAULT.value, + client=client, + ) + continue + + result: str | None = get_secret_from_manager( + secret_name=case.secret_name, + key_manager=KeyManagementSystem.AZURE_KEY_VAULT.value, + client=client, + ) + assert result == case.expected.value diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index 85933fbf9e8..f58ade11d1c 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -629,6 +629,25 @@ def test_aws_credential_redaction_catches_quoted_values(): assert redact_string(safe) == safe +def test_bedrock_batch_s3_credential_redaction_in_deployment_dump(): + """The router logs each deployment's litellm_params at DEBUG. A Bedrock batch + deployment carries s3_secret_access_key there, which the aws_* key-name rule + did not cover, so the S3 secret was printed verbatim (LIT-8290).""" + cases = ( + "{'s3_secret_access_key': 'wJalrXUtnFEMIK7MDENGbPxRfiCYEXAMPLEKEY'}", + "s3_secret_access_key=wJalrXUtnFEMIK7MDENGbPxRfiCYEXAMPLEKEY", + "{'s3_access_key_id': 'not-an-akia-shaped-value'}", + ) + for secret_line in cases: + result = redact_string(secret_line) + assert "REDACTED" in result, f"S3 credential redaction missed: {secret_line!r}" + assert "wJalrXUtnFEMIK7MDENGbPxRfiCYEXAMPLEKEY" not in result + assert "not-an-akia-shaped-value" not in result + + safe = "'s3_bucket_name': 'my-batch-bucket'" + assert redact_string(safe) == safe + + @pytest.mark.parametrize( "extra", ( diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 3eba1bcfced..1e59f4d878e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4682,6 +4682,33 @@ def test_bedrock_batch_params_never_reach_the_provider(): ) +def test_documented_batch_s3_credentials_never_reach_the_provider(): + """The Bedrock batch docs tell users to put s3_access_key_id, s3_secret_access_key + and s3_encryption_key_id on the deployment. Left unregistered they are swept into + additionalModelRequestFields, Bedrock 400s ordinary chat on that deployment with + `s3_secret_access_key: Extra inputs are not permitted`, and the S3 secret is sent + to the provider and printed in the debug log (LIT-8290). + """ + configured = { + "s3_access_key_id": "configured-access-key-id", + "s3_secret_access_key": "configured-secret-access-key", + "s3_encryption_key_id": "arn:aws:kms:us-east-1:000000000000:key/configured", + } + kwargs = {"a_real_provider_specific_param": 1, **configured} + + non_default = get_non_default_completion_params(dict(kwargs)) + + assert non_default == {"a_real_provider_specific_param": 1}, ( + "documented batch S3 credentials leaked into the provider params: " + f"{sorted(set(non_default) - {'a_real_provider_specific_param'})}" + ) + + batch_params = dict(GenericLiteLLMParams(**kwargs)) + assert {field: batch_params.get(field) for field in configured} == configured, ( + "registering these must not strip them from the batch path" + ) + + def test_client_side_timeout_marker_never_reaches_the_provider(): """The proxy stamps kwargs["client_side_timeout"] = True whenever a request carries a caller-supplied timeout (body timeout / request_timeout / stream_timeout or the diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 9312fdde075..fd1610dff6b 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -7,6 +7,7 @@ import math import os import threading import time +import uuid import weakref from collections.abc import Callable, Generator from contextlib import ExitStack @@ -18,8 +19,10 @@ from uuid import uuid4 import fakeredis import pytest import redis +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.in_memory_cache import InMemoryCache from litellm.caching.redis_cluster_cache import RedisClusterCache @@ -57,6 +60,36 @@ def redis_url() -> Generator[str]: worker.join(timeout=5) +@pytest.fixture +def azure_blob_facade() -> Generator[Cache]: + account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") + if account_url is None: + pytest.skip( + "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" + ) + facade: Final = Cache( + type=LiteLLMCacheType.AZURE_BLOB, + azure_account_url=account_url, + azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", + ) + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + try: + yield facade + finally: + backend.container_client.delete_container() + asyncio.run(backend.disconnect()) + + +def azure_blob_handle(facade: Cache) -> _native._CacheTestHandle: + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + return _native._CacheTestHandle.azure_blob( + backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}"), + backend.container_client.container_name, + ) + + @pytest.fixture def cluster_nodes() -> tuple[tuple[str, int], ...]: configured: Final = os.environ.get("LITELLM_TEST_REDIS_CLUSTER_NODES") @@ -381,6 +414,89 @@ def test_facade_registration_rejects_mismatched_capacity() -> None: _CacheTestHandle.memory(capacity=7)._bind_facade(facade) +def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: + backend: Final = azure_blob_facade.cache + assert isinstance(backend, AzureBlobCache) + handle: Final = azure_blob_handle(azure_blob_facade) + 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 + ) + 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} + native.store({**request("sync"), "ttl_seconds": 0.001}, response) + native.store(request("sync"), {"choices": [{"text": "second"}]}) + time.sleep(0.01) + stored: Final = json.loads(backend.container_client.download_blob("sync").readall()) + assert stored["response"] == response + assert isinstance(stored["timestamp"], float) + assert native.lookup(request("sync")) == response + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response + + backend.set_cache("python", {"timestamp": time.time(), "response": response}) + backend.set_cache("legacy", "bare legacy value") + backend.container_client.upload_blob("invalid", b"{not json", overwrite=True) + assert native.lookup(request("python")) == response + assert native.lookup(request("legacy")) == cast(CacheLookup, azure_blob_facade).get_cache(cache_key="legacy") + assert native.lookup_batch([request("python"), request("missing"), request("invalid"), request("sync")]) == { + "values": [response, None, None, response], + "missing_indices": [1, 2], + } + + with rebound(azure_blob_facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + with rebound(backend, "container_client", ContainerClient.from_container_url(backend.container_client.url)): + assert resolver.resolve().kind == "python_callback" + + def custom_get(*_args: object, **_kwargs: object) -> None: + return None + + with rebound(backend, "get_cache", custom_get): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "python_callback" + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response + + class CustomBlobCache(AzureBlobCache): + pass + + with rebound(azure_blob_facade, "cache", CustomBlobCache(account_url, backend.container_client.container_name)): + assert resolver.resolve().kind == "python_callback" + with pytest.raises(TypeError): + azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) + + +async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_python(azure_blob_facade: Cache) -> None: + backend: Final = azure_blob_facade.cache + assert isinstance(backend, AzureBlobCache) + azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)).resolve() + assert binding.kind == "native" + ping: Final = cast(dict[str, object], await binding.ping()) + assert ping["status"] == "success", ping + + await binding.async_store(request("async"), {"value": 1}) + 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 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}]) + assert await binding.async_lookup_batch([request("second"), request("missing"), request("first")]) == { + "values": [{"value": 4}, None, {"value": 3}], + "missing_indices": [1], + } + await binding.async_flush() + assert [blob.name for blob in backend.container_client.list_blobs()] == [] + assert await binding.async_lookup(request("async")) is None + + async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: parsed: Final = urlparse(redis_url) with rebound(litellm, "default_redis_ttl", 60): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d9c7e370d4f..0b09f3654dc 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31245,6 +31245,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; + /** S3 Access Key Id */ + s3_access_key_id?: string | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Bucket Owner */ @@ -31257,6 +31259,8 @@ export interface components { s3_output_bucket_name?: string | null; /** S3 Region Name */ s3_region_name?: string | null; + /** S3 Secret Access Key */ + s3_secret_access_key?: string | null; /** Search Context Cost Per Query */ search_context_cost_per_query?: { [key: string]: unknown; @@ -42045,6 +42049,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; + /** S3 Access Key Id */ + s3_access_key_id?: string | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Bucket Owner */ @@ -42057,6 +42063,8 @@ export interface components { s3_output_bucket_name?: string | null; /** S3 Region Name */ s3_region_name?: string | null; + /** S3 Secret Access Key */ + s3_secret_access_key?: string | null; /** Search Context Cost Per Query */ search_context_cost_per_query?: { [key: string]: unknown;