mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge origin/main into litellm_native_disk_cache
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
commit
f7b96c9ced
62 changed files with 2949 additions and 107 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
2
.github/workflows/test-rust.yml
vendored
2
.github/workflows/test-rust.yml
vendored
|
|
@ -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
|
||||
|
||||
|
|
|
|||
103
litellm-rust/Cargo.lock
generated
103
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
@ -2526,6 +2579,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-disk"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2743,6 +2813,7 @@ dependencies = [
|
|||
"litellm-auth",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-cache",
|
||||
"litellm-cache-azure-blob",
|
||||
"litellm-cache-disk",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
|
|
@ -2776,6 +2847,7 @@ dependencies = [
|
|||
"jsonwebtoken",
|
||||
"litellm-core-utils",
|
||||
"litellm-secrets-aws",
|
||||
"litellm-secrets-azure",
|
||||
"litellm-secrets-cyberark",
|
||||
"litellm-secrets-google",
|
||||
"litellm-secrets-types",
|
||||
|
|
@ -2811,6 +2883,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"
|
||||
|
|
@ -3587,6 +3679,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"
|
||||
|
|
@ -5182,6 +5284,7 @@ dependencies = [
|
|||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures",
|
||||
"quick-xml",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"url",
|
||||
|
|
|
|||
|
|
@ -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-disk = { path = "crates/cache-disk" }
|
||||
|
|
|
|||
|
|
@ -4,4 +4,4 @@ mod resolve;
|
|||
mod types;
|
||||
|
||||
pub use resolve::AzureAuthService;
|
||||
pub use types::AzureAuthInputs;
|
||||
pub use types::{AzureAuthInputs, ConfigValue};
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
22
litellm-rust/crates/cache-azure-blob/Cargo.toml
Normal file
22
litellm-rust/crates/cache-azure-blob/Cargo.toml
Normal file
|
|
@ -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
|
||||
254
litellm-rust/crates/cache-azure-blob/src/cache.rs
Normal file
254
litellm-rust/crates/cache-azure-blob/src/cache.rs
Normal file
|
|
@ -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<C> {
|
||||
container: BlobContainerClient,
|
||||
codec: C,
|
||||
runtime: Handle,
|
||||
account_url: String,
|
||||
container_name: String,
|
||||
}
|
||||
|
||||
impl<C: CacheCodec> AzureBlobCache<C> {
|
||||
pub async fn connect(
|
||||
account_url: &str,
|
||||
container: &str,
|
||||
codec: C,
|
||||
runtime: Handle,
|
||||
) -> Result<Self, Error> {
|
||||
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<Arc<dyn TokenCredential>>,
|
||||
client_options: ClientOptions,
|
||||
codec: C,
|
||||
runtime: Handle,
|
||||
) -> Result<Self, Error> {
|
||||
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<Option<C::Value>, 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<T>(&self, future: impl Future<Output = T>) -> 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<C: CacheCodec> BaseCache for AzureBlobCache<C> {
|
||||
type Value = C::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, _: &ExactCacheContext) -> Option<Duration> {
|
||||
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<Option<C::Value>, 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<Option<C::Value>, 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<CacheConnectionResult, Error> {
|
||||
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<C: CacheCodec> BatchCache for AzureBlobCache<C> {}
|
||||
|
||||
impl<C: CacheCodec> FlushCache for AzureBlobCache<C> {
|
||||
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;
|
||||
746
litellm-rust/crates/cache-azure-blob/src/cache/tests.rs
vendored
Normal file
746
litellm-rust/crates/cache-azure-blob/src/cache/tests.rs
vendored
Normal file
|
|
@ -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<String>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct FakeState {
|
||||
container_exists: bool,
|
||||
blobs: BTreeMap<String, Vec<u8>>,
|
||||
requests: Vec<RecordedRequest>,
|
||||
failing: bool,
|
||||
precondition_conflicts: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct FakeBlobService {
|
||||
state: Arc<Mutex<FakeState>>,
|
||||
}
|
||||
|
||||
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<Vec<u8>> {
|
||||
self.state.lock().unwrap().blobs.get(name).cloned()
|
||||
}
|
||||
|
||||
fn blob_names(&self) -> Vec<String> {
|
||||
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<RecordedRequest> {
|
||||
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<u8>) -> 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<u8> {
|
||||
let mut xml = String::from(
|
||||
r#"<?xml version="1.0" encoding="utf-8"?><EnumerationResults ServiceEndpoint="https://example.blob.core.windows.net/" ContainerName="litellm-cache"><Blobs>"#,
|
||||
);
|
||||
for name in state.blobs.keys() {
|
||||
xml.push_str(&format!(
|
||||
"<Blob><Name>{name}</Name><Properties><BlobType>BlockBlob</BlobType></Properties></Blob>"
|
||||
));
|
||||
}
|
||||
xml.push_str("</Blobs><NextMarker /></EnumerationResults>");
|
||||
xml.into_bytes()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl HttpClient for FakeBlobService {
|
||||
async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
|
||||
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<AzureBlobCache<ResponseCacheCodec>>,
|
||||
}
|
||||
|
||||
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<AzureBlobCache<ResponseCacheCodec>, 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<AzureBlobCache<ResponseCacheCodec>> {
|
||||
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")))
|
||||
);
|
||||
}
|
||||
84
litellm-rust/crates/cache-azure-blob/src/credential.rs
Normal file
84
litellm-rust/crates/cache-azure-blob/src/credential.rs
Normal file
|
|
@ -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<dyn Fn(&str) -> Option<String> + 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<TokenRequestOptions<'_>>,
|
||||
) -> azure_core::Result<AccessToken> {
|
||||
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),
|
||||
))
|
||||
}
|
||||
}
|
||||
5
litellm-rust/crates/cache-azure-blob/src/lib.rs
Normal file
5
litellm-rust/crates/cache-azure-blob/src/lib.rs
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
mod cache;
|
||||
mod credential;
|
||||
|
||||
pub use cache::AzureBlobCache;
|
||||
pub use credential::AzureBlobCredential;
|
||||
0
litellm-rust/crates/cache-azure-blob/src/tests.rs
Normal file
0
litellm-rust/crates/cache-azure-blob/src/tests.rs
Normal file
|
|
@ -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-disk.workspace = true
|
||||
|
|
|
|||
|
|
@ -79,6 +79,11 @@ pub(super) struct RedisCacheConfig {
|
|||
pub(super) connection: RedisConnectionConfig,
|
||||
}
|
||||
|
||||
pub(super) struct AzureBlobCacheConfig {
|
||||
pub(super) account_url: String,
|
||||
pub(super) container: String,
|
||||
}
|
||||
|
||||
struct RedisClientProjection<'py> {
|
||||
topology: RedisTopology,
|
||||
host: String,
|
||||
|
|
@ -94,6 +99,7 @@ pub(super) enum CacheBackendConfig {
|
|||
Memory(MemoryCacheConfig),
|
||||
Redis(Box<RedisCacheConfig>),
|
||||
Disk(DiskCacheConfig),
|
||||
AzureBlob(AzureBlobCacheConfig),
|
||||
}
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
|
|
@ -169,12 +175,17 @@ 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
|
||||
| CacheType::ValkeySemantic
|
||||
| CacheType::S3
|
||||
| CacheType::QdrantSemantic
|
||||
| CacheType::AzureBlob
|
||||
| CacheType::Gcs,
|
||||
)
|
||||
| None => Ok(CacheConfigProjection::Unsupported(
|
||||
|
|
@ -188,6 +199,7 @@ impl NativeCacheConfig {
|
|||
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
|
||||
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
|
||||
CacheBackendConfig::Disk(_) => None,
|
||||
CacheBackendConfig::AzureBlob(_) => None,
|
||||
};
|
||||
if service.default_ttl() != default_ttl {
|
||||
return Some("facade and native backend default TTLs must match");
|
||||
|
|
@ -225,10 +237,34 @@ impl NativeCacheConfig {
|
|||
let facade = std::fs::canonicalize(&config.directory).ok();
|
||||
(native != facade).then_some("facade and native backend directories must match")
|
||||
}
|
||||
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<AzureBlobCacheConfig> {
|
||||
let client = backend.getattr("container_client")?;
|
||||
let container = client.getattr("container_name")?.extract::<String>()?;
|
||||
let url = client.getattr("url")?.extract::<String>()?;
|
||||
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)]
|
||||
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
|
||||
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
|
||||
|
|
|
|||
|
|
@ -39,6 +39,18 @@ struct DiskStoreGuard {
|
|||
directory: String,
|
||||
}
|
||||
|
||||
struct AzureBlobClientGuard {
|
||||
sync_client: Py<PyAny>,
|
||||
async_client: Py<PyAny>,
|
||||
url: String,
|
||||
container_name: String,
|
||||
}
|
||||
|
||||
enum ConnectionGuard {
|
||||
None,
|
||||
RedisPool(RedisPoolGuard),
|
||||
AzureBlob(AzureBlobClientGuard),
|
||||
}
|
||||
struct RedisPoolAttributes {
|
||||
pool: &'static str,
|
||||
connection_class: &'static str,
|
||||
|
|
@ -60,8 +72,8 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
|
|||
pub(super) struct FacadeGuard {
|
||||
outer: ObjectGuard,
|
||||
backend: ObjectGuard,
|
||||
redis_pool: Option<RedisPoolGuard>,
|
||||
disk_store: Option<DiskStoreGuard>,
|
||||
connection: ConnectionGuard,
|
||||
}
|
||||
|
||||
impl ObjectGuard {
|
||||
|
|
@ -231,6 +243,61 @@ impl DiskStoreGuard {
|
|||
}
|
||||
}
|
||||
|
||||
impl AzureBlobClientGuard {
|
||||
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
let sync_client = backend.getattr("container_client")?;
|
||||
Ok(Self {
|
||||
url: sync_client.getattr("url")?.extract::<String>()?,
|
||||
container_name: sync_client.getattr("container_name")?.extract::<String>()?,
|
||||
sync_client: sync_client.unbind(),
|
||||
async_client: backend.getattr("async_container_client")?.unbind(),
|
||||
})
|
||||
}
|
||||
|
||||
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
|
||||
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::<String>()?
|
||||
&& self.container_name == sync_client.getattr("container_name")?.extract::<String>()?)
|
||||
}
|
||||
|
||||
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<Self> {
|
||||
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<bool> {
|
||||
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<'_>,
|
||||
|
|
@ -254,6 +321,11 @@ impl FacadeGuard {
|
|||
"redis",
|
||||
),
|
||||
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
|
||||
("azure-blob", _) => (
|
||||
"litellm.caching.azure_blob_cache",
|
||||
"AzureBlobCache",
|
||||
"azure-blob",
|
||||
),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
let backend = facade.getattr("cache")?;
|
||||
|
|
@ -299,14 +371,10 @@ impl FacadeGuard {
|
|||
"redis_flush_size",
|
||||
],
|
||||
)?,
|
||||
redis_pool: match (kind, cluster) {
|
||||
("redis", false) => Some(RedisPoolGuard::capture(&backend, STANDALONE_POOL)?),
|
||||
("redis", true) => Some(RedisPoolGuard::capture(&backend, CLUSTER_POOL)?),
|
||||
_ => None,
|
||||
},
|
||||
disk_store: (kind == "disk")
|
||||
.then(|| DiskStoreGuard::capture(&backend))
|
||||
.transpose()?,
|
||||
connection: ConnectionGuard::capture(kind, cluster, &backend)?,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -318,25 +386,21 @@ impl FacadeGuard {
|
|||
if !self.backend.matches(py, &backend)? {
|
||||
return Ok(false);
|
||||
}
|
||||
match &self.redis_pool {
|
||||
Some(guard) => guard.matches(py, &backend),
|
||||
None => match &self.disk_store {
|
||||
Some(guard) => guard.matches(py, &backend),
|
||||
None => Ok(true),
|
||||
},
|
||||
if let Some(guard) = &self.disk_store
|
||||
&& !guard.matches(py, &backend)?
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
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)?;
|
||||
}
|
||||
if let Some(guard) = &self.disk_store {
|
||||
guard.traverse(&visit)?;
|
||||
}
|
||||
Ok(())
|
||||
self.connection.traverse(&visit)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
use litellm_host_python::release_gil;
|
||||
use litellm_host_python::{release_gil, run_sync_value};
|
||||
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
|
||||
|
||||
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
|
||||
|
|
@ -76,6 +76,21 @@ impl CacheTestHandle {
|
|||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (account_url, container))]
|
||||
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {
|
||||
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(),
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn backend(&self) -> &'static str {
|
||||
self.service.kind()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
use std::{path::Path, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
|
||||
use litellm_cache_azure_blob::AzureBlobCache;
|
||||
use litellm_cache_disk::DiskCache;
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::{RedisCache, RedisTopology};
|
||||
|
|
@ -17,6 +18,7 @@ pub(super) enum NativeResponseCache {
|
|||
buffer: Option<Arc<WriteBuffer>>,
|
||||
},
|
||||
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
|
||||
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
|
||||
}
|
||||
|
||||
impl NativeResponseCache {
|
||||
|
|
@ -47,11 +49,33 @@ impl NativeResponseCache {
|
|||
buffer: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn disk(directory: &str) -> Result<Self, Error> {
|
||||
let cache = DiskCache::open(directory, ResponseCacheCodec)?;
|
||||
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
|
||||
}
|
||||
|
||||
pub async fn azure_blob(account_url: &str, container: &str) -> Result<Self, Error> {
|
||||
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::Disk(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NativeResponseCache {
|
||||
|
|
@ -60,6 +84,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(_) => "memory",
|
||||
Self::Redis { .. } => "redis",
|
||||
Self::Disk(_) => "disk",
|
||||
Self::AzureBlob(_) => "azure-blob",
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -68,12 +93,13 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.default_ttl(),
|
||||
Self::Redis { cache, .. } => cache.default_ttl(),
|
||||
Self::Disk(cache) => cache.default_ttl(),
|
||||
Self::AzureBlob(cache) => cache.default_ttl(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn namespace(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Memory(_) => None,
|
||||
Self::Memory(_) | Self::AzureBlob(_) => None,
|
||||
Self::Redis { cache, .. } => cache.backend().namespace(),
|
||||
Self::Disk(_) => None,
|
||||
}
|
||||
|
|
@ -81,7 +107,7 @@ impl NativeResponseCache {
|
|||
|
||||
pub fn topology(&self) -> Option<&RedisTopology> {
|
||||
match self {
|
||||
Self::Memory(_) => None,
|
||||
Self::Memory(_) | Self::AzureBlob(_) => None,
|
||||
Self::Redis { cache, .. } => Some(cache.backend().topology()),
|
||||
Self::Disk(_) => None,
|
||||
}
|
||||
|
|
@ -92,6 +118,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
|
||||
Self::Redis { .. } => None,
|
||||
Self::Disk(_) => None,
|
||||
Self::AzureBlob(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -100,6 +127,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.backend().max_entry_bytes(),
|
||||
Self::Redis { .. } => None,
|
||||
Self::Disk(_) => None,
|
||||
Self::AzureBlob(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -110,14 +138,14 @@ impl NativeResponseCache {
|
|||
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
|
||||
},
|
||||
disk @ Self::Disk(_) => disk,
|
||||
memory => memory,
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn directory(&self) -> Option<&Path> {
|
||||
match self {
|
||||
Self::Disk(cache) => Some(cache.backend().directory()),
|
||||
Self::Memory(_) | Self::Redis { .. } => None,
|
||||
Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -130,6 +158,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.lookup(request, now),
|
||||
Self::Redis { cache, .. } => cache.lookup(request, now),
|
||||
Self::Disk(cache) => cache.lookup(request, now),
|
||||
Self::AzureBlob(cache) => cache.lookup(request, now),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -143,6 +172,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.store(request, response, now),
|
||||
Self::Redis { cache, .. } => cache.store(request, response, now),
|
||||
Self::Disk(cache) => cache.store(request, response, now),
|
||||
Self::AzureBlob(cache) => cache.store(request, response, now),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -155,6 +185,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.lookup_batch(requests, now),
|
||||
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
|
||||
Self::Disk(cache) => cache.lookup_batch(requests, now),
|
||||
Self::AzureBlob(cache) => cache.lookup_batch(requests, now),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -167,6 +198,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.async_lookup(request, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
|
||||
Self::Disk(cache) => cache.async_lookup(request, now).await,
|
||||
Self::AzureBlob(cache) => cache.async_lookup(request, now).await,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -187,6 +219,7 @@ impl NativeResponseCache {
|
|||
buffer: Some(buffer),
|
||||
} => buffer.async_store(cache, request, response, now).await,
|
||||
Self::Disk(cache) => cache.async_store(request, response, now).await,
|
||||
Self::AzureBlob(cache) => cache.async_store(request, response, now).await,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -199,6 +232,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
|
||||
Self::Disk(cache) => cache.async_lookup_batch(requests, now).await,
|
||||
Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -211,6 +245,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
|
||||
Self::Disk(cache) => cache.async_store_batch(entries, now).await,
|
||||
Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -224,6 +259,7 @@ impl NativeResponseCache {
|
|||
cache.async_flush().await
|
||||
}
|
||||
Self::Disk(cache) => cache.async_flush().await,
|
||||
Self::AzureBlob(cache) => cache.async_flush().await,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -232,6 +268,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.test_connection().await,
|
||||
Self::Redis { cache, .. } => cache.test_connection().await,
|
||||
Self::Disk(cache) => cache.test_connection().await,
|
||||
Self::AzureBlob(cache) => cache.test_connection().await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
24
litellm-rust/crates/secrets-azure/Cargo.toml
Normal file
24
litellm-rust/crates/secrets-azure/Cargo.toml
Normal file
|
|
@ -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
|
||||
25
litellm-rust/crates/secrets-azure/src/error.rs
Normal file
25
litellm-rust/crates/secrets-azure/src/error.rs
Normal file
|
|
@ -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,
|
||||
}
|
||||
118
litellm-rust/crates/secrets-azure/src/key_vault.rs
Normal file
118
litellm-rust/crates/secrets-azure/src/key_vault.rs
Normal file
|
|
@ -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<AzureAuthService>,
|
||||
inputs: Arc<AzureAuthInputs>,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SecretResponse {
|
||||
value: Option<String>,
|
||||
}
|
||||
|
||||
impl AzureKeyVault {
|
||||
pub fn with_client(
|
||||
client: reqwest::Client,
|
||||
vault: reqwest::Url,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
) -> Result<Self, Error> {
|
||||
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<dyn Lookup + Send + Sync>) -> Result<Self, Error> {
|
||||
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<Option<Secret>, 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")
|
||||
}
|
||||
7
litellm-rust/crates/secrets-azure/src/lib.rs
Normal file
7
litellm-rust/crates/secrets-azure/src/lib.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
mod error;
|
||||
mod key_vault;
|
||||
|
||||
pub use error::Error;
|
||||
pub use key_vault::AzureKeyVault;
|
||||
8
litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json
vendored
Normal file
8
litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json
vendored
Normal file
|
|
@ -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}}
|
||||
]
|
||||
}
|
||||
222
litellm-rust/crates/secrets-azure/tests/key_vault.rs
Normal file
222
litellm-rust/crates/secrets-azure/tests/key_vault.rs
Normal file
|
|
@ -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<u16>) {
|
||||
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<FixtureCase>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
missing: Option<bool>,
|
||||
error: Option<bool>,
|
||||
}
|
||||
|
||||
#[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()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
30
litellm-rust/crates/secrets-azure/tests/live.rs
Normal file
30
litellm-rust/crates/secrets-azure/tests/live.rs
Normal file
|
|
@ -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::<reqwest::Url>()
|
||||
.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}"
|
||||
);
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -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() {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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|"
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -63,6 +63,7 @@ LlmRoute = Literal[
|
|||
LlmCapability = Literal[
|
||||
"assume_role",
|
||||
"basic",
|
||||
"batch_deployment",
|
||||
"count_tokens",
|
||||
"govcloud_partition",
|
||||
"input_validation",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
57
tests/e2e/migrations/test_rolling_upgrade.py
Normal file
57
tests/e2e/migrations/test_rolling_upgrade.py
Normal file
|
|
@ -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"
|
||||
43
tests/e2e/migrations/test_shaped_database.py
Normal file
43
tests/e2e/migrations/test_shaped_database.py
Normal file
|
|
@ -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"
|
||||
47
tests/e2e/migrations/test_upgrade.py
Normal file
47
tests/e2e/migrations/test_upgrade.py
Normal file
|
|
@ -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"
|
||||
124
tests/e2e/migrations/upgrade.py
Normal file
124
tests/e2e/migrations/upgrade.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import json
|
|||
import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import weakref
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
|
|
@ -16,8 +17,10 @@ import diskcache
|
|||
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.disk_cache import DiskCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
|
|
@ -50,6 +53,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")
|
||||
|
|
@ -374,6 +407,89 @@ def test_facade_registration_rejects_mismatched_capacity() -> None:
|
|||
_native._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):
|
||||
|
|
|
|||
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue