scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

chore: merge main into litellm_gcs_native_cache

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 22:08:55 +00:00
commit cb3e212e29
62 changed files with 2948 additions and 111 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
.circleci/config.yml
View file

@ -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:

3
.circleci/scripts/run_migration_tests.py
View file

@ -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
View file

@ -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
View file

@ -115,6 +115,28 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "async-stream"
version = "0.3.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
dependencies = [
"async-stream-impl",
"futures-core",
"pin-project-lite",
]
[[package]]
name = "async-stream-impl"
version = "0.3.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "async-trait"
version = "0.1.91"
@ -599,6 +621,37 @@ dependencies = [
"url",
]
[[package]]
name = "azure_storage_blob"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "17b10207ecf7d666df6940b50051f433b3cd5d2b9b1dd190613208d7a84e7eed"
dependencies = [
"async-stream",
"async-trait",
"azure_core",
"azure_storage_common",
"bytes",
"futures",
"percent-encoding",
"pin-project",
"serde",
"serde_json",
"time",
"tokio",
]
[[package]]
name = "azure_storage_common"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0af2e6aeb8d76b17fc998f453c320913f73787b944e3cc29509d19411fa0321d"
dependencies = [
"azure_core",
"serde",
"time",
]
[[package]]
name = "base64"
version = "0.13.1"
@ -2470,6 +2523,23 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-azure-blob"
version = "0.1.0"
dependencies = [
"async-trait",
"azure_core",
"azure_storage_blob",
"futures-util",
"litellm-auth-azure",
"litellm-auth-types",
"litellm-cache",
"litellm-cache-response",
"serde_json",
"tokio",
"url",
]
[[package]]
name = "litellm-cache-gcs"
version = "0.1.0"
@ -2687,6 +2757,7 @@ dependencies = [
"litellm-auth",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
"litellm-cache-gcs",
"litellm-cache-memory",
"litellm-cache-redis",
@ -2720,6 +2791,7 @@ dependencies = [
"jsonwebtoken",
"litellm-core-utils",
"litellm-secrets-aws",
"litellm-secrets-azure",
"litellm-secrets-cyberark",
"litellm-secrets-google",
"litellm-secrets-types",
@ -2755,6 +2827,26 @@ dependencies = [
"wiremock",
]
[[package]]
name = "litellm-secrets-azure"
version = "0.1.0"
dependencies = [
"litellm-auth-azure",
"litellm-auth-types",
"litellm-core-utils",
"litellm-secrets-types",
"percent-encoding",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"veil",
"wiremock",
]
[[package]]
name = "litellm-secrets-cyberark"
version = "0.1.0"
@ -3531,6 +3623,16 @@ version = "1.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
[[package]]
name = "quick-xml"
version = "0.41.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1"
dependencies = [
"memchr",
"serde",
]
[[package]]
name = "quinn"
version = "0.11.11"
@ -5076,6 +5178,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"futures",
"quick-xml",
"serde",
"serde_json",
"url",

2
litellm-rust/Cargo.toml
View file

@ -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-gcs = { path = "crates/cache-gcs" }

2
litellm-rust/crates/auth-azure/src/lib.rs
View file

@ -4,4 +4,4 @@ mod resolve;
mod types;
pub use resolve::AzureAuthService;
pub use types::AzureAuthInputs;
pub use types::{AzureAuthInputs, ConfigValue};

15
litellm-rust/crates/auth-azure/src/types.rs
View file

@ -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
View 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
View 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
View 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
View 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
View 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
View file

1
litellm-rust/crates/python-bridge/Cargo.toml
View 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-gcs.workspace = true

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

@ -82,6 +82,11 @@ pub(super) struct GcsCacheConfig {
pub(super) path_service_account: Option<String>,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
struct RedisClientProjection<'py> {
topology: RedisTopology,
host: String,
@ -97,6 +102,7 @@ pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
Gcs(GcsCacheConfig),
AzureBlob(AzureBlobCacheConfig),
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
@ -172,13 +178,18 @@ 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::Disk
| CacheType::QdrantSemantic
| CacheType::AzureBlob,
| CacheType::QdrantSemantic,
)
| None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
@ -187,12 +198,12 @@ impl NativeCacheConfig {
}
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
let expected = match &self.backend {
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::AzureBlob(_) | CacheBackendConfig::Gcs(_) => None,
};
if service.default_ttl() != expected {
if service.default_ttl() != default_ttl {
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
@ -242,10 +253,34 @@ impl NativeCacheConfig {
Some("facade and native backend credentials must match")
}
CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
None => Some("facade and native backend types must match"),
Some((account_url, container))
if account_url != config.account_url || container != config.container =>
{
Some("facade and native backend containers must match")
}
Some(_) => None,
},
}
}
}
#[inline(never)]
fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<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>()?;

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

@ -34,6 +34,19 @@ struct RedisPoolGuard {
attributes: RedisPoolAttributes,
}
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,
@ -55,7 +68,7 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
redis_pool: Option<RedisPoolGuard>,
connection: ConnectionGuard,
}
impl ObjectGuard {
@ -205,6 +218,61 @@ impl RedisPoolGuard {
}
}
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<'_>,
@ -228,6 +296,11 @@ impl FacadeGuard {
"redis",
),
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
"azure-blob",
),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -276,11 +349,7 @@ impl FacadeGuard {
"path_service_account",
],
)?,
redis_pool: match (kind, cluster) {
("redis", false) => Some(RedisPoolGuard::capture(&backend, STANDALONE_POOL)?),
("redis", true) => Some(RedisPoolGuard::capture(&backend, CLUSTER_POOL)?),
_ => None,
},
connection: ConnectionGuard::capture(kind, cluster, &backend)?,
})
}
@ -292,19 +361,13 @@ impl FacadeGuard {
if !self.backend.matches(py, &backend)? {
return Ok(false);
}
match &self.redis_pool {
Some(guard) => guard.matches(py, &backend),
None => Ok(true),
}
self.connection.matches(py, &backend)
}
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
self.outer.traverse(&visit)?;
self.backend.traverse(&visit)?;
if let Some(guard) = &self.redis_pool {
guard.traverse(&visit)?;
}
Ok(())
self.connection.traverse(&visit)
}
}

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

@ -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 litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
@ -91,6 +91,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()

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

@ -1,6 +1,7 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
@ -17,6 +18,7 @@ pub(super) enum NativeResponseCache {
buffer: Option<Arc<WriteBuffer>>,
},
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
impl NativeResponseCache {
@ -59,6 +61,29 @@ impl NativeResponseCache {
};
Ok(Self::Gcs(Arc::new(ResponseCache::new(Arc::new(backend)))))
}
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::Gcs(_) => None,
}
}
}
impl NativeResponseCache {
@ -67,6 +92,7 @@ impl NativeResponseCache {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::Gcs(_) => "gcs",
Self::AzureBlob(_) => "azure-blob",
}
}
@ -75,12 +101,13 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
Self::Gcs(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::Gcs(_) => None,
}
@ -88,7 +115,7 @@ impl NativeResponseCache {
pub fn topology(&self) -> Option<&RedisTopology> {
match self {
Self::Memory(_) | Self::Gcs(_) => None,
Self::Memory(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
}
}
@ -96,16 +123,14 @@ impl NativeResponseCache {
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } => None,
Self::Gcs(_) => None,
Self::Redis { .. } | Self::AzureBlob(_) | Self::Gcs(_) => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } => None,
Self::Gcs(_) => None,
Self::Redis { .. } | Self::AzureBlob(_) | Self::Gcs(_) => None,
}
}
@ -115,7 +140,7 @@ impl NativeResponseCache {
cache,
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
},
memory => memory,
other => other,
}
}
@ -128,6 +153,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::Gcs(cache) => cache.lookup(request, now),
Self::AzureBlob(cache) => cache.lookup(request, now),
}
}
@ -141,6 +167,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
Self::Gcs(cache) => cache.store(request, response, now),
Self::AzureBlob(cache) => cache.store(request, response, now),
}
}
@ -153,6 +180,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
Self::Gcs(cache) => cache.lookup_batch(requests, now),
Self::AzureBlob(cache) => cache.lookup_batch(requests, now),
}
}
@ -165,6 +193,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
Self::Gcs(cache) => cache.async_lookup(request, now).await,
Self::AzureBlob(cache) => cache.async_lookup(request, now).await,
}
}
@ -185,6 +214,7 @@ impl NativeResponseCache {
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
Self::Gcs(cache) => cache.async_store(request, response, now).await,
Self::AzureBlob(cache) => cache.async_store(request, response, now).await,
}
}
@ -197,6 +227,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await,
Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await,
}
}
@ -209,6 +240,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
Self::Gcs(cache) => cache.async_store_batch(entries, now).await,
Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await,
}
}
@ -222,6 +254,7 @@ impl NativeResponseCache {
cache.async_flush().await
}
Self::Gcs(cache) => cache.async_flush().await,
Self::AzureBlob(cache) => cache.async_flush().await,
}
}
@ -230,6 +263,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::Gcs(cache) => cache.test_connection().await,
Self::AzureBlob(cache) => cache.test_connection().await,
}
}

24
litellm-rust/crates/secrets-azure/Cargo.toml Normal file
View 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
View 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
View 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
View 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
View 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
View 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
View 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}"
);
}

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

@ -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

3
litellm-rust/crates/secrets/src/error.rs
View file

@ -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),

9
litellm-rust/crates/secrets/src/handler.rs
View file

@ -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)

2
litellm-rust/crates/secrets/src/lib.rs
View file

@ -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")]

61
litellm-rust/crates/secrets/tests/handler.rs
View file

@ -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() {

40
litellm/integrations/SlackAlerting/batching_handler.py
View file

@ -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
)

18
litellm/integrations/SlackAlerting/slack_alerting.py
View file

@ -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):

2
litellm/litellm_core_utils/secret_redaction.py
View file

@ -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|"

28
litellm/proxy/auth/auth_checks.py
View file

@ -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(

1
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -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

4
litellm/proxy/management_endpoints/scim/scim_v2.py
View file

@ -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)

15
litellm/types/integrations/slack_alerting.py
View file

@ -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

2
litellm/types/router.py
View file

@ -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

3
litellm/types/utils.py
View file

@ -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",
)

1
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -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"}

1
tests/e2e/coverage_registry/schema.py
View file

@ -63,6 +63,7 @@ LlmRoute = Literal[
LlmCapability = Literal[
"assume_role",
"basic",
"batch_deployment",
"count_tokens",
"govcloud_partition",
"input_validation",

44
tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py
View file

@ -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(

29
tests/e2e/migrations/conftest.py
View file

@ -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
tests/e2e/migrations/containers.py
View file

@ -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
View 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
View 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
View 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
View 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

1
tests/e2e/models.py
View file

@ -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

10
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -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")

61
tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py
View file

@ -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]

94
tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py
View file

@ -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},
)

4
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -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

65
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -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

25
tests/test_litellm/proxy/auth/test_resolvers_grants.py
View file

@ -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))

39
tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py
View file

@ -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)

39
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -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)

107
tests/test_litellm/secret_managers/test_secret_manager_handler.py Normal file
View file

@ -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

19
tests/test_litellm/test_secret_redaction.py
View file

@ -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",
(

27
tests/test_litellm/test_utils.py
View file

@ -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

116
tests/test_litellm_rust/test_cache.py
View file

@ -5,6 +5,7 @@ import json
import os
import threading
import time
import uuid
import weakref
from collections.abc import Generator
from types import SimpleNamespace
@ -14,8 +15,10 @@ from urllib.parse import urlparse
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.gcs_cache import GCSCache
from litellm.caching.in_memory_cache import InMemoryCache
@ -59,6 +62,36 @@ def fake_gcs() -> Generator[FakeGcs]:
server.close()
@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")
@ -383,6 +416,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
View file

@ -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;