mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
wip
This commit is contained in:
Yujong Lee
parent
cc88a9479e
commit
c80617c4a6
42 changed files with 1473 additions and 1377 deletions
90
litellm-rust/Cargo.lock
generated
90
litellm-rust/Cargo.lock
generated
|
|
@ -1838,7 +1838,41 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
|
||||
|
||||
[[package]]
|
||||
name = "litellm-core"
|
||||
name = "litellm-ai-gateway"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"futures-channel",
|
||||
"futures-util",
|
||||
"litellm-config",
|
||||
"litellm-core",
|
||||
"reqwest 0.12.28",
|
||||
"rustls 0.23.42",
|
||||
"rustls-native-certs",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"subtle",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tower",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-auth"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"subtle",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"veil",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-auth-aws"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-config",
|
||||
|
|
@ -1847,13 +1881,64 @@ dependencies = [
|
|||
"aws-sigv4",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-types",
|
||||
"litellm-auth",
|
||||
"moka",
|
||||
"reqwest 0.12.28",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-auth-azure"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"azure_core",
|
||||
"azure_identity",
|
||||
"litellm-auth",
|
||||
"moka",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"strum",
|
||||
"tokio",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-auth-gcp"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"gcp_auth",
|
||||
"litellm-auth",
|
||||
"moka",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-config"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-core",
|
||||
"pyo3",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"data-url",
|
||||
"futures-util",
|
||||
"gcp_auth",
|
||||
"litellm-auth",
|
||||
"litellm-auth-aws",
|
||||
"litellm-auth-azure",
|
||||
"litellm-auth-gcp",
|
||||
"mime_guess",
|
||||
"moka",
|
||||
"rand 0.8.7",
|
||||
|
|
@ -1865,7 +1950,6 @@ dependencies = [
|
|||
"serde_json",
|
||||
"serde_path_to_error",
|
||||
"sha2 0.10.9",
|
||||
"strum",
|
||||
"subtle",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,16 @@
|
|||
[workspace]
|
||||
members = ["crates/*"]
|
||||
members = [
|
||||
"crates/auth",
|
||||
"crates/auth-aws",
|
||||
"crates/auth-azure",
|
||||
"crates/auth-gcp",
|
||||
"crates/core",
|
||||
"crates/token-counter",
|
||||
"crates/config",
|
||||
"crates/ai-gateway",
|
||||
"crates/python-interop",
|
||||
"crates/python-bridge",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
|
|
@ -11,6 +22,10 @@ repository = "https://github.com/BerriAI/litellm"
|
|||
[workspace.dependencies]
|
||||
bytes = "1"
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-auth = { path = "crates/auth" }
|
||||
litellm-auth-aws = { path = "crates/auth-aws" }
|
||||
litellm-auth-azure = { path = "crates/auth-azure" }
|
||||
litellm-auth-gcp = { path = "crates/auth-gcp" }
|
||||
litellm-token-counter = { path = "crates/token-counter" }
|
||||
litellm-python-interop = { path = "crates/python-interop" }
|
||||
pyo3 = "0.29.2"
|
||||
|
|
|
|||
24
litellm-rust/crates/auth-aws/Cargo.toml
Normal file
24
litellm-rust/crates/auth-aws/Cargo.toml
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
[package]
|
||||
name = "litellm-auth-aws"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-auth.workspace = true
|
||||
|
||||
moka = { workspace = true, features = ["sync"] }
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
|
||||
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] }
|
||||
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] }
|
||||
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] }
|
||||
aws-sigv4 = "1.5.1"
|
||||
aws-types = "1.4.0"
|
||||
aws-smithy-runtime-api = "1.13.0"
|
||||
|
||||
[dev-dependencies]
|
||||
reqwest.workspace = true
|
||||
tokio.workspace = true
|
||||
950
litellm-rust/crates/auth-aws/src/aws.rs
Normal file
950
litellm-rust/crates/auth-aws/src/aws.rs
Normal file
|
|
@ -0,0 +1,950 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use moka::sync::Cache;
|
||||
use serde_json::{Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use aws_credential_types::Credentials;
|
||||
use aws_credential_types::provider::ProvideCredentials;
|
||||
use aws_sigv4::http_request::{
|
||||
SignableBody, SignableRequest, SigningParams, SigningSettings, sign,
|
||||
};
|
||||
use aws_sigv4::sign::v4;
|
||||
use aws_smithy_runtime_api::client::identity::Identity;
|
||||
|
||||
use litellm_auth::Error;
|
||||
|
||||
use super::constants::{
|
||||
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME,
|
||||
AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
|
||||
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
|
||||
BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX,
|
||||
SIGV4_COMPUTED_HEADER_NAMES,
|
||||
};
|
||||
|
||||
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
|
||||
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
|
||||
|
||||
static STATIC_CREDENTIALS_CACHE: OnceLock<Cache<String, Credentials>> = OnceLock::new();
|
||||
static AMBIENT_CREDENTIALS_CACHE: OnceLock<Cache<String, Credentials>> = OnceLock::new();
|
||||
|
||||
fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option<Duration> {
|
||||
match flow {
|
||||
AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL),
|
||||
AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL),
|
||||
AwsAuthFlow::WebIdentity { .. }
|
||||
| AwsAuthFlow::AssumeRole { .. }
|
||||
| AwsAuthFlow::Profile { .. }
|
||||
| AwsAuthFlow::SessionToken { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct AwsAuthConfig {
|
||||
pub access_key_id: Option<String>,
|
||||
pub secret_access_key: Option<String>,
|
||||
pub session_token: Option<String>,
|
||||
pub region_name: Option<String>,
|
||||
pub session_name: Option<String>,
|
||||
pub profile_name: Option<String>,
|
||||
pub role_name: Option<String>,
|
||||
pub web_identity_token: Option<String>,
|
||||
pub sts_endpoint: Option<String>,
|
||||
pub external_id: Option<String>,
|
||||
}
|
||||
|
||||
impl AwsAuthConfig {
|
||||
fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option<String> + Sync)) -> Self {
|
||||
Self {
|
||||
access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)),
|
||||
secret_access_key: self
|
||||
.secret_access_key
|
||||
.or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)),
|
||||
session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)),
|
||||
region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)),
|
||||
session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)),
|
||||
profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)),
|
||||
role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)),
|
||||
web_identity_token: self
|
||||
.web_identity_token
|
||||
.or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)),
|
||||
sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)),
|
||||
external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum AwsAuthFlow {
|
||||
WebIdentity {
|
||||
token: String,
|
||||
role: String,
|
||||
session_name: String,
|
||||
},
|
||||
AssumeRole {
|
||||
role: String,
|
||||
session_name: Option<String>,
|
||||
},
|
||||
Profile {
|
||||
name: String,
|
||||
},
|
||||
SessionToken {
|
||||
access_key_id: String,
|
||||
secret_access_key: String,
|
||||
session_token: String,
|
||||
},
|
||||
StaticKeys {
|
||||
access_key_id: String,
|
||||
secret_access_key: String,
|
||||
region_name: String,
|
||||
},
|
||||
DefaultChain,
|
||||
}
|
||||
|
||||
fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(format!("{config:?}:{flow:?}"));
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn static_credentials_cache() -> &'static Cache<String, Credentials> {
|
||||
STATIC_CREDENTIALS_CACHE.get_or_init(|| {
|
||||
Cache::builder()
|
||||
.max_capacity(200)
|
||||
.time_to_live(STATIC_CREDENTIALS_TTL)
|
||||
.build()
|
||||
})
|
||||
}
|
||||
|
||||
fn ambient_credentials_cache() -> &'static Cache<String, Credentials> {
|
||||
AMBIENT_CREDENTIALS_CACHE.get_or_init(|| {
|
||||
Cache::builder()
|
||||
.max_capacity(200)
|
||||
.time_to_live(AMBIENT_CREDENTIALS_TTL)
|
||||
.build()
|
||||
})
|
||||
}
|
||||
|
||||
fn get_cached_credentials(key: &str) -> Option<Credentials> {
|
||||
static_credentials_cache()
|
||||
.get(key)
|
||||
.or_else(|| ambient_credentials_cache().get(key))
|
||||
}
|
||||
|
||||
fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) {
|
||||
if ttl == STATIC_CREDENTIALS_TTL {
|
||||
static_credentials_cache().insert(key, credentials);
|
||||
} else {
|
||||
ambient_credentials_cache().insert(key, credentials);
|
||||
}
|
||||
}
|
||||
|
||||
fn role_identity(arn: &str) -> Option<(&str, &str, &str)> {
|
||||
let mut parts = arn.splitn(6, ':');
|
||||
let ("arn", partition, _, _, account, resource) = (
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
) else {
|
||||
return None;
|
||||
};
|
||||
let role = if let Some(role) = resource.strip_prefix("role/") {
|
||||
role.rsplit('/').next()?
|
||||
} else {
|
||||
resource.strip_prefix("assumed-role/")?.split('/').next()?
|
||||
};
|
||||
Some((partition, account, role))
|
||||
}
|
||||
|
||||
fn same_role_arns(target: &str, caller: &str) -> bool {
|
||||
role_identity(target) == role_identity(caller)
|
||||
}
|
||||
|
||||
pub fn classify_auth(
|
||||
config: AwsAuthConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> AwsAuthFlow {
|
||||
let config = config.with_environment(env_lookup);
|
||||
if let (Some(token), Some(role), Some(session_name)) = (
|
||||
config.web_identity_token.clone(),
|
||||
config.role_name.clone(),
|
||||
config.session_name.clone(),
|
||||
) {
|
||||
return AwsAuthFlow::WebIdentity {
|
||||
token,
|
||||
role,
|
||||
session_name,
|
||||
};
|
||||
}
|
||||
if let Some(role) = config.role_name.clone() {
|
||||
return AwsAuthFlow::AssumeRole {
|
||||
role,
|
||||
session_name: config.session_name.clone(),
|
||||
};
|
||||
}
|
||||
if let Some(name) = config.profile_name {
|
||||
return AwsAuthFlow::Profile { name };
|
||||
}
|
||||
if let (Some(access_key_id), Some(secret_access_key), Some(session_token)) = (
|
||||
config.access_key_id.clone(),
|
||||
config.secret_access_key.clone(),
|
||||
config.session_token,
|
||||
) {
|
||||
return AwsAuthFlow::SessionToken {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
session_token,
|
||||
};
|
||||
}
|
||||
if let (Some(access_key_id), Some(secret_access_key), Some(region_name)) = (
|
||||
config.access_key_id,
|
||||
config.secret_access_key,
|
||||
config.region_name,
|
||||
) {
|
||||
return AwsAuthFlow::StaticKeys {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
region_name,
|
||||
};
|
||||
}
|
||||
AwsAuthFlow::DefaultChain
|
||||
}
|
||||
|
||||
pub async fn resolve_credentials(
|
||||
config: AwsAuthConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<Credentials, Error> {
|
||||
let resolved = config.clone().with_environment(env_lookup);
|
||||
let flow = classify_auth(config, env_lookup);
|
||||
match flow {
|
||||
AwsAuthFlow::SessionToken {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
session_token,
|
||||
} => Ok(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
Some(session_token),
|
||||
None,
|
||||
"litellm-static-session",
|
||||
)),
|
||||
AwsAuthFlow::StaticKeys {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
region_name,
|
||||
} => {
|
||||
let flow = AwsAuthFlow::StaticKeys {
|
||||
access_key_id: access_key_id.clone(),
|
||||
secret_access_key: secret_access_key.clone(),
|
||||
region_name,
|
||||
};
|
||||
let key = cache_key(&resolved, &flow);
|
||||
if let Some(credentials) = get_cached_credentials(&key) {
|
||||
return Ok(credentials);
|
||||
}
|
||||
let credentials = Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
None,
|
||||
None,
|
||||
"litellm-static",
|
||||
);
|
||||
set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL),
|
||||
);
|
||||
Ok(credentials)
|
||||
}
|
||||
AwsAuthFlow::Profile { name } => {
|
||||
let provider = aws_config::profile::ProfileFileCredentialsProvider::builder()
|
||||
.profile_name(name)
|
||||
.build();
|
||||
provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::AwsProfile(error.to_string()))
|
||||
}
|
||||
AwsAuthFlow::AssumeRole { role, session_name } => {
|
||||
if is_already_running_as_role(&role, &resolved).await? {
|
||||
let ambient_flow = AwsAuthFlow::DefaultChain;
|
||||
let key = cache_key(&resolved, &ambient_flow);
|
||||
if let Some(credentials) = get_cached_credentials(&key) {
|
||||
return Ok(credentials);
|
||||
}
|
||||
let provider =
|
||||
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
|
||||
.build()
|
||||
.await;
|
||||
let credentials = provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
|
||||
set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL),
|
||||
);
|
||||
return Ok(credentials);
|
||||
}
|
||||
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
|
||||
if let Some(region) = resolved.region_name.clone() {
|
||||
loader = loader.region(aws_types::region::Region::new(region));
|
||||
}
|
||||
if let Some(endpoint) = resolved.sts_endpoint.clone() {
|
||||
loader = loader.endpoint_url(endpoint);
|
||||
}
|
||||
if let (Some(access_key_id), Some(secret_access_key)) =
|
||||
(resolved.access_key_id, resolved.secret_access_key)
|
||||
{
|
||||
loader = loader.credentials_provider(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
resolved.session_token,
|
||||
None,
|
||||
"litellm-role-source",
|
||||
));
|
||||
}
|
||||
let sdk_config = loader.load().await;
|
||||
let builder = aws_config::sts::AssumeRoleProvider::builder(role);
|
||||
let builder = match session_name {
|
||||
Some(name) => builder.session_name(name),
|
||||
None => builder.session_name(default_session_name()),
|
||||
};
|
||||
let builder = match resolved.external_id {
|
||||
Some(id) => builder.external_id(id),
|
||||
None => builder,
|
||||
};
|
||||
let provider = builder.configure(&sdk_config).build().await;
|
||||
provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::AwsAssumeRole(error.to_string()))
|
||||
}
|
||||
AwsAuthFlow::WebIdentity {
|
||||
token,
|
||||
role,
|
||||
session_name,
|
||||
} => {
|
||||
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
|
||||
if let Some(region) = resolved.region_name {
|
||||
loader = loader.region(aws_types::region::Region::new(region));
|
||||
}
|
||||
if let Some(endpoint) = resolved.sts_endpoint {
|
||||
loader = loader.endpoint_url(endpoint);
|
||||
}
|
||||
let sdk_config = loader.load().await;
|
||||
let client = aws_sdk_sts::Client::new(&sdk_config);
|
||||
let response = client
|
||||
.assume_role_with_web_identity()
|
||||
.role_arn(role)
|
||||
.role_session_name(session_name)
|
||||
.web_identity_token(token)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| Error::AwsWebIdentity(error.to_string()))?;
|
||||
let credentials = response
|
||||
.credentials()
|
||||
.ok_or(Error::AwsMissingWebIdentityCredentials)?;
|
||||
let expiration = SystemTime::try_from(*credentials.expiration())
|
||||
.map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?;
|
||||
Ok(Credentials::new(
|
||||
credentials.access_key_id(),
|
||||
credentials.secret_access_key(),
|
||||
Some(credentials.session_token().to_string()),
|
||||
Some(expiration),
|
||||
"litellm-web-identity",
|
||||
))
|
||||
}
|
||||
AwsAuthFlow::DefaultChain => {
|
||||
let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain);
|
||||
if let Some(credentials) = get_cached_credentials(&key) {
|
||||
return Ok(credentials);
|
||||
}
|
||||
let provider =
|
||||
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
|
||||
.build()
|
||||
.await;
|
||||
let credentials = provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
|
||||
set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
credential_cache_ttl(&AwsAuthFlow::DefaultChain).unwrap_or(AMBIENT_CREDENTIALS_TTL),
|
||||
);
|
||||
Ok(credentials)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result<bool, Error> {
|
||||
if role_identity(role).is_none() {
|
||||
return Ok(false);
|
||||
}
|
||||
if let (Ok(current_role), Ok(token_file)) = (
|
||||
std::env::var(AWS_ROLE_ARN),
|
||||
std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE),
|
||||
) && !token_file.is_empty()
|
||||
{
|
||||
return Ok(same_role_arns(role, ¤t_role));
|
||||
}
|
||||
|
||||
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
|
||||
if let Some(region) = config.region_name.clone() {
|
||||
loader = loader.region(aws_types::region::Region::new(region));
|
||||
}
|
||||
if let Some(endpoint) = config.sts_endpoint.clone() {
|
||||
loader = loader.endpoint_url(endpoint);
|
||||
}
|
||||
let sdk_config = loader.load().await;
|
||||
let response = match aws_sdk_sts::Client::new(&sdk_config)
|
||||
.get_caller_identity()
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(_) => return Ok(false),
|
||||
};
|
||||
Ok(response
|
||||
.arn()
|
||||
.is_some_and(|caller| same_role_arns(role, caller)))
|
||||
}
|
||||
|
||||
fn default_session_name() -> String {
|
||||
let seconds = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_or(0, |duration| duration.as_secs());
|
||||
format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}")
|
||||
}
|
||||
|
||||
/// The subset of `headers` SigV4 should cover.
|
||||
///
|
||||
/// Python signs only these and reattaches the rest afterwards, so a forwarded
|
||||
/// client header cannot change the canonical request and invalidate the
|
||||
/// signature. Signing everything instead makes the request 403 on a header the
|
||||
/// caller supplied, on a deployment that works on the Python path.
|
||||
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
|
||||
headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
let name = name.to_ascii_lowercase();
|
||||
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
|
||||
|| name.starts_with("x-amz-")
|
||||
|| name.starts_with("x-amzn-")
|
||||
})
|
||||
.map(|(name, value)| (name.clone(), value.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Whether the signer produces `name` itself.
|
||||
///
|
||||
/// Python's reattach loop skips these, so a caller-supplied copy never reaches
|
||||
/// the wire next to the computed one.
|
||||
pub fn is_sigv4_computed_header(name: &str) -> bool {
|
||||
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
|
||||
}
|
||||
|
||||
pub fn sign_bedrock_post(
|
||||
url: &str,
|
||||
body: &[u8],
|
||||
headers: &BTreeMap<String, String>,
|
||||
region: &str,
|
||||
credentials: &Credentials,
|
||||
signing_time: SystemTime,
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
let identity: Identity = credentials.clone().into();
|
||||
let params = v4::SigningParams::builder()
|
||||
.identity(&identity)
|
||||
.region(region)
|
||||
.name(BEDROCK_SERVICE)
|
||||
.time(signing_time)
|
||||
.settings(SigningSettings::default())
|
||||
.build()
|
||||
.map(SigningParams::from)
|
||||
.map_err(|error| Error::AwsSigningParameters(error.to_string()))?;
|
||||
let header_refs = headers
|
||||
.iter()
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()));
|
||||
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
|
||||
.map_err(|error| Error::AwsSignableRequest(error.to_string()))?;
|
||||
let (instructions, _) = sign(request, ¶ms)
|
||||
.map_err(|error| Error::AwsSigning(error.to_string()))?
|
||||
.into_parts();
|
||||
Ok(instructions
|
||||
.headers()
|
||||
.map(|(name, value)| {
|
||||
let normalized_name = match name {
|
||||
"authorization" => "Authorization",
|
||||
"x-amz-date" => "X-Amz-Date",
|
||||
"x-amz-security-token" => "X-Amz-Security-Token",
|
||||
_ => name,
|
||||
};
|
||||
(normalized_name.to_string(), value.to_string())
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Model-id and region parsing shared by every Bedrock route.
|
||||
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
|
||||
let mut stripped = model;
|
||||
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let mut region = None;
|
||||
if let Some((candidate, remainder)) = stripped.split_once('/')
|
||||
&& is_bedrock_region(candidate)
|
||||
{
|
||||
region = Some(candidate.to_string());
|
||||
stripped = remainder;
|
||||
}
|
||||
for prefix in ["nova-2/", "nova/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if region.is_none() {
|
||||
// Python splits the whole ARN and takes field 3, the region. Stripping
|
||||
// `arn:` first shifts every field down one, so the region is field 2
|
||||
// here; field 3 is the account id.
|
||||
region = stripped
|
||||
.strip_prefix("arn:")
|
||||
.and_then(|value| value.split(':').nth(2))
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
}
|
||||
(stripped.to_string(), region)
|
||||
}
|
||||
|
||||
fn is_bedrock_region(value: &str) -> bool {
|
||||
value.len() > 3
|
||||
&& value.contains('-')
|
||||
&& value
|
||||
.chars()
|
||||
.all(|char| char.is_ascii_alphanumeric() || char == '-')
|
||||
}
|
||||
|
||||
pub fn resolve_bedrock_region(
|
||||
model_region: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
if let Some(region) = optional_params
|
||||
.get("aws_region_name")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
return region.to_string();
|
||||
}
|
||||
if let Some(region) = model_region {
|
||||
return region.to_string();
|
||||
}
|
||||
env_lookup(AWS_REGION_NAME)
|
||||
.or_else(|| env_lookup(AWS_REGION))
|
||||
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
|
||||
}
|
||||
|
||||
pub fn aws_auth_config(
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> AwsAuthConfig {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
};
|
||||
let env = |key: &str| env_lookup(key);
|
||||
AwsAuthConfig {
|
||||
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
|
||||
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
|
||||
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
|
||||
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
|
||||
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
|
||||
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
|
||||
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
|
||||
web_identity_token: value("aws_web_identity_token")
|
||||
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
|
||||
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
|
||||
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Credentials a host resolved through its own chain and handed down verbatim.
|
||||
///
|
||||
/// A host with its own resolution (LiteLLM's Python `BaseAWSLLM`, which reads
|
||||
/// profiles, STS and boto sessions) passes the result here so the core signs
|
||||
/// with exactly those. Without this the core would re-derive from ambient
|
||||
/// state, where an unrelated `AWS_ROLE_NAME` or `AWS_PROFILE_NAME` in the
|
||||
/// environment outranks explicit keys in [`classify_auth`] and the two sides
|
||||
/// would sign as different principals.
|
||||
pub fn host_supplied_credentials(optional_params: &Map<String, Value>) -> Option<Credentials> {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
};
|
||||
let access_key_id = value("aws_access_key_id")?;
|
||||
let secret_access_key = value("aws_secret_access_key")?;
|
||||
Some(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
value("aws_session_token").map(str::to_string),
|
||||
None,
|
||||
"litellm-host-supplied",
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn no_env(_: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
|
||||
(
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
|
||||
.to_string(),
|
||||
br#"{"input":"hello"}"#.to_vec(),
|
||||
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_the_region_field_of_a_model_arn_not_the_account_id() {
|
||||
// Python's `_get_aws_region_from_model_arn` splits the whole ARN and
|
||||
// takes field 3. Stripping `arn:` first shifts every field down one, so
|
||||
// the region is field 2 here. Taking field 3 after the strip returns
|
||||
// the account id, which is not a region at all.
|
||||
let (_, region) = bedrock_model_id_and_region(
|
||||
"bedrock/arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2",
|
||||
);
|
||||
assert_eq!(region.as_deref(), Some("us-west-2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classification_preserves_python_precedence() {
|
||||
let config = AwsAuthConfig {
|
||||
access_key_id: Some("ak".into()),
|
||||
secret_access_key: Some("sk".into()),
|
||||
session_token: Some("token".into()),
|
||||
region_name: Some("us-east-1".into()),
|
||||
session_name: Some("session".into()),
|
||||
profile_name: Some("profile".into()),
|
||||
role_name: Some("role".into()),
|
||||
web_identity_token: Some("oidc".into()),
|
||||
..Default::default()
|
||||
};
|
||||
assert!(matches!(
|
||||
classify_auth(config, &no_env),
|
||||
AwsAuthFlow::WebIdentity { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classification_covers_fallthroughs() {
|
||||
let env = |key: &str| match key {
|
||||
AWS_PROFILE_NAME => Some("profile".into()),
|
||||
_ => None,
|
||||
};
|
||||
assert!(matches!(
|
||||
classify_auth(AwsAuthConfig::default(), &env),
|
||||
AwsAuthFlow::Profile { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
classify_auth(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("ak".into()),
|
||||
secret_access_key: Some("sk".into()),
|
||||
session_token: Some("token".into()),
|
||||
..Default::default()
|
||||
},
|
||||
&no_env
|
||||
),
|
||||
AwsAuthFlow::SessionToken { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
classify_auth(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("ak".into()),
|
||||
secret_access_key: Some("sk".into()),
|
||||
region_name: Some("us-east-1".into()),
|
||||
..Default::default()
|
||||
},
|
||||
&no_env
|
||||
),
|
||||
AwsAuthFlow::StaticKeys { .. }
|
||||
));
|
||||
assert_eq!(
|
||||
classify_auth(AwsAuthConfig::default(), &no_env),
|
||||
AwsAuthFlow::DefaultChain
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn static_credentials_do_not_use_network() {
|
||||
let credentials = resolve_credentials(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("ak".into()),
|
||||
secret_access_key: Some("sk".into()),
|
||||
region_name: Some("us-east-1".into()),
|
||||
..Default::default()
|
||||
},
|
||||
&no_env,
|
||||
)
|
||||
.await
|
||||
.expect("static credentials");
|
||||
assert_eq!(credentials.access_key_id(), "ak");
|
||||
assert_eq!(credentials.session_token(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_policy_matches_python_flows() {
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::StaticKeys {
|
||||
access_key_id: "ak".into(),
|
||||
secret_access_key: "sk".into(),
|
||||
region_name: "us-east-1".into(),
|
||||
}),
|
||||
Some(STATIC_CREDENTIALS_TTL)
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::DefaultChain),
|
||||
Some(AMBIENT_CREDENTIALS_TTL)
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::SessionToken {
|
||||
access_key_id: "ak".into(),
|
||||
secret_access_key: "sk".into(),
|
||||
session_token: "token".into(),
|
||||
}),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::Profile {
|
||||
name: "profile".into()
|
||||
}),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::AssumeRole {
|
||||
role: "arn:aws:iam::123456789012:role/demo".into(),
|
||||
session_name: None,
|
||||
}),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::WebIdentity {
|
||||
token: "token".into(),
|
||||
role: "arn:aws:iam::123456789012:role/demo".into(),
|
||||
session_name: "session".into(),
|
||||
}),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_round_trip_preserves_credentials() {
|
||||
let key = format!("cache-test-{}", std::process::id());
|
||||
let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test");
|
||||
set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL);
|
||||
assert_eq!(
|
||||
get_cached_credentials(&key).map(|value| value.access_key_id().to_string()),
|
||||
Some("cache-ak".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_role_comparison_matches_partition_account_and_role() {
|
||||
assert!(same_role_arns(
|
||||
"arn:aws:iam::123456789012:role/path/demo",
|
||||
"arn:aws:sts::123456789012:assumed-role/demo/session"
|
||||
));
|
||||
assert!(!same_role_arns(
|
||||
"arn:aws:iam::123456789012:role/demo",
|
||||
"arn:aws:iam::999999999999:role/demo"
|
||||
));
|
||||
assert!(!same_role_arns(
|
||||
"arn:aws:iam::123456789012:role/demo",
|
||||
"arn:aws-cn:iam::123456789012:role/demo"
|
||||
));
|
||||
assert!(!same_role_arns(
|
||||
"arn:aws:iam::123456789012:user/demo",
|
||||
"arn:aws:iam::123456789012:role/demo"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_client_header_is_not_folded_into_the_signature() {
|
||||
// Python signs only the AWS header set, so a header a caller forwarded
|
||||
// cannot change the canonical request. Signing it instead makes the
|
||||
// request 403 the moment anything on the wire rewrites or drops it.
|
||||
let (url, body, mut headers) = parity_inputs();
|
||||
headers.insert("x-request-id".to_string(), "abc-123".to_string());
|
||||
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
|
||||
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
|
||||
let signable = aws_signature_headers(&headers);
|
||||
|
||||
assert!(!signable.contains_key("x-request-id"));
|
||||
assert!(!signable.contains_key("Accept-Encoding"));
|
||||
// The AWS-prefixed one is genuinely part of the signature.
|
||||
assert!(signable.contains_key("x-amzn-trace-id"));
|
||||
assert!(signable.contains_key("Content-Type"));
|
||||
|
||||
let credentials = Credentials::new(
|
||||
"AKIDEXAMPLE",
|
||||
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
|
||||
None,
|
||||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&signable,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
SystemTime::UNIX_EPOCH,
|
||||
)
|
||||
.expect("signs");
|
||||
let authorization = signed
|
||||
.get("Authorization")
|
||||
.expect("carries an authorization header");
|
||||
assert!(
|
||||
!authorization.contains("x-request-id"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
assert!(
|
||||
!authorization.contains("accept-encoding"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signing_matches_botocore_golden_vector() {
|
||||
let (url, body, headers) = parity_inputs();
|
||||
let credentials = Credentials::new(
|
||||
"AKIDEXAMPLE",
|
||||
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
|
||||
Some("session-token".to_string()),
|
||||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
|
||||
)
|
||||
.expect("golden signature");
|
||||
assert_eq!(
|
||||
signed.get("X-Amz-Date").map(String::as_str),
|
||||
Some("20240102T030405Z")
|
||||
);
|
||||
assert_eq!(
|
||||
signed.get("X-Amz-Security-Token").map(String::as_str),
|
||||
Some("session-token")
|
||||
);
|
||||
assert_eq!(
|
||||
signed.get("Authorization").map(String::as_str),
|
||||
Some(
|
||||
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signing_without_session_token_omits_security_header() {
|
||||
let (url, body, headers) = parity_inputs();
|
||||
let credentials = Credentials::new(
|
||||
"AKIDEXAMPLE",
|
||||
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
|
||||
None,
|
||||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
|
||||
)
|
||||
.expect("signature");
|
||||
assert!(!signed.contains_key("X-Amz-Security-Token"));
|
||||
}
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let access_key_id = std::env::var("AWS_BEDROCK_TEST_ACCESS_KEY_ID")?;
|
||||
let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?;
|
||||
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec();
|
||||
let headers =
|
||||
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]);
|
||||
let credentials = resolve_credentials(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some(access_key_id),
|
||||
secret_access_key: Some(secret_access_key),
|
||||
region_name: Some("us-west-2".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
&no_env,
|
||||
)
|
||||
.await?;
|
||||
let client = reqwest::Client::new();
|
||||
let mut failures = Vec::new();
|
||||
|
||||
for region in ["us-west-2", "us-east-1"] {
|
||||
let url = format!(
|
||||
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
|
||||
);
|
||||
let signed_headers = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
region,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?;
|
||||
let mut request = client.post(&url).body(body.clone());
|
||||
for (name, value) in &headers {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
for (name, value) in signed_headers {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
let response = request.send().await?;
|
||||
let status = response.status();
|
||||
let response_body = response.text().await?;
|
||||
let snippet: String = response_body.chars().take(240).collect();
|
||||
println!("region={region} status={status} response={snippet}");
|
||||
if status == reqwest::StatusCode::OK {
|
||||
return Ok(());
|
||||
}
|
||||
failures.push(format!("{region}: {status} {snippet}"));
|
||||
}
|
||||
|
||||
panic!(
|
||||
"no Bedrock region returned HTTP 200: {}",
|
||||
failures.join("; ")
|
||||
);
|
||||
}
|
||||
}
|
||||
43
litellm-rust/crates/auth-aws/src/constants.rs
Normal file
43
litellm-rust/crates/auth-aws/src/constants.rs
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID";
|
||||
pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY";
|
||||
pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN";
|
||||
pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME";
|
||||
pub const AWS_REGION: &str = "AWS_REGION";
|
||||
pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME";
|
||||
pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME";
|
||||
pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME";
|
||||
pub const AWS_WEB_IDENTITY_TOKEN: &str = "AWS_WEB_IDENTITY_TOKEN";
|
||||
pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
|
||||
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
|
||||
pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT";
|
||||
pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID";
|
||||
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";
|
||||
|
||||
/// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors
|
||||
/// Python's `_filter_headers_for_aws_signature` allowlist.
|
||||
pub const AWS_SIGNED_HEADER_NAMES: &[&str] = &[
|
||||
"host",
|
||||
"content-type",
|
||||
"date",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"x-amz-content-sha256",
|
||||
"x-amz-algorithm",
|
||||
"x-amz-credential",
|
||||
"x-amz-signedheaders",
|
||||
"x-amz-signature",
|
||||
];
|
||||
/// Headers the signer emits itself. Mirrors Python's `SIGV4_COMPUTED_HEADERS`,
|
||||
/// which the reattach loop skips so a caller's copy cannot ride alongside the
|
||||
/// computed one.
|
||||
pub const SIGV4_COMPUTED_HEADER_NAMES: &[&str] = &[
|
||||
"authorization",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"date",
|
||||
];
|
||||
pub const BEDROCK_SERVICE: &str = "bedrock";
|
||||
pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session";
|
||||
pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2";
|
||||
pub const BEDROCK_RUNTIME_ENDPOINT_TEMPLATE: &str =
|
||||
"https://bedrock-runtime.{region}.amazonaws.com";
|
||||
4
litellm-rust/crates/auth-aws/src/lib.rs
Normal file
4
litellm-rust/crates/auth-aws/src/lib.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
mod aws;
|
||||
pub mod constants;
|
||||
|
||||
pub use aws::*;
|
||||
19
litellm-rust/crates/auth-azure/Cargo.toml
Normal file
19
litellm-rust/crates/auth-azure/Cargo.toml
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
[package]
|
||||
name = "litellm-auth-azure"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
azure_core.workspace = true
|
||||
azure_identity.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
moka.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
strum.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tokio.workspace = true
|
||||
|
|
@ -4,7 +4,7 @@ use std::sync::Arc;
|
|||
use azure_core::credentials::TokenCredential;
|
||||
use moka::future::Cache;
|
||||
|
||||
use crate::AuthError;
|
||||
use litellm_auth::Error;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub(crate) struct AzureCredentialProviderCacheKey {
|
||||
|
|
@ -31,9 +31,9 @@ impl AzureCredentialProviderCache {
|
|||
&self,
|
||||
key: AzureCredentialProviderCacheKey,
|
||||
create: F,
|
||||
) -> Result<Arc<dyn TokenCredential>, AuthError>
|
||||
) -> Result<Arc<dyn TokenCredential>, Error>
|
||||
where
|
||||
F: Future<Output = Result<Arc<dyn TokenCredential>, AuthError>>,
|
||||
F: Future<Output = Result<Arc<dyn TokenCredential>, Error>>,
|
||||
{
|
||||
self.entries
|
||||
.try_get_with(key, create)
|
||||
7
litellm-rust/crates/auth-azure/src/lib.rs
Normal file
7
litellm-rust/crates/auth-azure/src/lib.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
mod credential_provider_cache;
|
||||
mod native;
|
||||
mod resolve;
|
||||
mod types;
|
||||
|
||||
pub use resolve::AzureAuthService;
|
||||
pub use types::AzureAuthInputs;
|
||||
|
|
@ -1,4 +1,3 @@
|
|||
use crate::auth::error::AuthConfigurationError;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, UNIX_EPOCH};
|
||||
|
||||
|
|
@ -13,8 +12,8 @@ use azure_identity::{
|
|||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::AuthError;
|
||||
use crate::auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
|
||||
use litellm_auth::Error;
|
||||
use litellm_auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
|
||||
|
||||
use super::credential_provider_cache::{
|
||||
AzureCredentialProviderCache, AzureCredentialProviderCacheKey,
|
||||
|
|
@ -62,7 +61,7 @@ pub(crate) struct ValidatedAzureRequest {
|
|||
}
|
||||
|
||||
impl ValidatedAzureRequest {
|
||||
pub(crate) fn new(request: NativeAzureRequest) -> Result<Self, AuthError> {
|
||||
pub(crate) fn new(request: NativeAzureRequest) -> Result<Self, Error> {
|
||||
validate_authority(&request)?;
|
||||
let credential_source = validate_sources(&request)?;
|
||||
Ok(Self {
|
||||
|
|
@ -120,7 +119,7 @@ impl NativeAzureTokenAcquirer {
|
|||
pub(crate) async fn acquire(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> Result<ResolvedCredential, AuthError> {
|
||||
) -> Result<ResolvedCredential, Error> {
|
||||
let scope = request.request.scope().to_string();
|
||||
let key = request.request.cache_key();
|
||||
let transport = self.transport.clone();
|
||||
|
|
@ -134,7 +133,7 @@ impl NativeAzureTokenAcquirer {
|
|||
let token = credential
|
||||
.get_token(&[scope.as_str()], None)
|
||||
.await
|
||||
.map_err(|error| AuthError::AzureTokenAcquisition(error.to_string()))?;
|
||||
.map_err(|error| Error::AzureTokenAcquisition(error.to_string()))?;
|
||||
let expires_on = u64::try_from(token.expires_on.unix_timestamp())
|
||||
.ok()
|
||||
.map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds));
|
||||
|
|
@ -239,7 +238,7 @@ impl NativeAzureRequest {
|
|||
}
|
||||
}
|
||||
|
||||
fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
|
||||
fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
|
||||
let authority = match request {
|
||||
NativeAzureRequest::ClientSecret { authority, .. }
|
||||
| NativeAzureRequest::ClientAssertion { authority, .. }
|
||||
|
|
@ -251,8 +250,7 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
|
|||
let Some(authority) = authority else {
|
||||
return Ok(());
|
||||
};
|
||||
let url = url::Url::parse(authority.value())
|
||||
.map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureAuthority))?;
|
||||
let url = url::Url::parse(authority.value()).map_err(|_| Error::InvalidAzureAuthority)?;
|
||||
if url.scheme() != "https"
|
||||
|| url.host_str().is_none()
|
||||
|| !url.username().is_empty()
|
||||
|
|
@ -261,14 +259,12 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
|
|||
|| url.fragment().is_some()
|
||||
|| !matches!(url.path(), "" | "/")
|
||||
{
|
||||
return Err(AuthError::Configuration(
|
||||
AuthConfigurationError::InvalidAzureAuthority,
|
||||
));
|
||||
return Err(Error::InvalidAzureAuthority);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_sources(request: &NativeAzureRequest) -> Result<InputSource, AuthError> {
|
||||
fn validate_sources(request: &NativeAzureRequest) -> Result<InputSource, Error> {
|
||||
match request {
|
||||
NativeAzureRequest::ClientSecret {
|
||||
tenant_id,
|
||||
|
|
@ -356,7 +352,7 @@ fn is_request_controlled<T>(value: &Sourced<T>, optional: Option<&Sourced<String
|
|||
|| optional.is_some_and(|value| value.source() == InputSource::Request)
|
||||
}
|
||||
|
||||
fn trusted_only(sources: &[InputSource]) -> Result<InputSource, AuthError> {
|
||||
fn trusted_only(sources: &[InputSource]) -> Result<InputSource, Error> {
|
||||
if sources.contains(&InputSource::Request) {
|
||||
return mixed_sources();
|
||||
}
|
||||
|
|
@ -371,16 +367,14 @@ fn trusted_source(sources: &[InputSource]) -> InputSource {
|
|||
}
|
||||
}
|
||||
|
||||
fn mixed_sources<T>() -> Result<T, AuthError> {
|
||||
Err(AuthError::Configuration(
|
||||
AuthConfigurationError::MixedAzureCredentialSources,
|
||||
))
|
||||
fn mixed_sources<T>() -> Result<T, Error> {
|
||||
Err(Error::MixedAzureCredentialSources)
|
||||
}
|
||||
|
||||
fn build_credential(
|
||||
request: NativeAzureRequest,
|
||||
transport: Option<azure_core::http::Transport>,
|
||||
) -> Result<Arc<dyn TokenCredential>, AuthError> {
|
||||
) -> Result<Arc<dyn TokenCredential>, Error> {
|
||||
match request {
|
||||
NativeAzureRequest::ClientSecret {
|
||||
tenant_id,
|
||||
|
|
@ -439,11 +433,7 @@ fn build_credential(
|
|||
NativeAzureRequest::DeveloperTools { .. } => DeveloperToolsCredential::new(None)
|
||||
.map(|credential| credential as Arc<dyn TokenCredential>),
|
||||
}
|
||||
.map_err(|error| {
|
||||
AuthError::Configuration(AuthConfigurationError::AzureCredentialInitialization(
|
||||
error.to_string(),
|
||||
))
|
||||
})
|
||||
.map_err(|error| Error::AzureCredentialInitialization(error.to_string()))
|
||||
}
|
||||
|
||||
fn client_options(
|
||||
|
|
@ -494,7 +484,7 @@ mod tests {
|
|||
use azure_core::{Bytes, Result};
|
||||
|
||||
use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
|
||||
use crate::auth::{InputSource, SecretValue, Sourced};
|
||||
use litellm_auth::{InputSource, SecretValue, Sourced};
|
||||
|
||||
fn deployment<T>(value: T) -> Sourced<T> {
|
||||
Sourced::new(value, InputSource::Deployment)
|
||||
|
|
@ -659,9 +649,7 @@ mod tests {
|
|||
|
||||
assert!(matches!(
|
||||
error,
|
||||
crate::AuthError::Configuration(
|
||||
crate::auth::error::AuthConfigurationError::MixedAzureCredentialSources
|
||||
)
|
||||
litellm_auth::Error::MixedAzureCredentialSources
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -691,12 +679,7 @@ mod tests {
|
|||
authority,
|
||||
))
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
error,
|
||||
crate::AuthError::Configuration(
|
||||
crate::auth::error::AuthConfigurationError::InvalidAzureAuthority
|
||||
)
|
||||
));
|
||||
assert!(matches!(error, litellm_auth::Error::InvalidAzureAuthority));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,6 +1,5 @@
|
|||
use crate::AuthError;
|
||||
use crate::auth::error::AuthConfigurationError;
|
||||
use crate::auth::{
|
||||
use litellm_auth::Error;
|
||||
use litellm_auth::{
|
||||
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
|
||||
SecretValue, Sourced, TokenProviderHandle,
|
||||
};
|
||||
|
|
@ -37,7 +36,7 @@ pub(crate) enum AzureCredentialPlan {
|
|||
}
|
||||
|
||||
/// Rust counterpart to Python's `get_azure_ad_token`, not `BaseAzureLLM`.
|
||||
pub(crate) struct AzureAuthService {
|
||||
pub struct AzureAuthService {
|
||||
native: Arc<dyn AzureTokenAcquirer>,
|
||||
}
|
||||
|
||||
|
|
@ -45,14 +44,14 @@ trait AzureTokenAcquirer: Send + Sync {
|
|||
fn acquire(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>>;
|
||||
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>;
|
||||
}
|
||||
|
||||
impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
|
||||
fn acquire(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>> {
|
||||
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>> {
|
||||
Box::pin(NativeAzureTokenAcquirer::acquire(self, request))
|
||||
}
|
||||
}
|
||||
|
|
@ -71,17 +70,17 @@ impl AzureAuthService {
|
|||
Self { native }
|
||||
}
|
||||
|
||||
pub(crate) async fn get_azure_ad_token(
|
||||
pub async fn get_azure_ad_token(
|
||||
&self,
|
||||
inputs: &AzureAuthInputs,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<Option<Sourced<ResolvedCredential>>, AuthError> {
|
||||
) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
|
||||
match select_auth_plan(inputs, env_lookup)? {
|
||||
AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)),
|
||||
AzureCredentialPlan::Caller(caller) => {
|
||||
let credential = caller.acquire().await?;
|
||||
if credential.secret().expose().is_empty() {
|
||||
return Err(AuthError::EmptyAzureToken);
|
||||
return Err(Error::EmptyAzureToken);
|
||||
}
|
||||
Ok(Some(Sourced::new(credential, InputSource::Deployment)))
|
||||
}
|
||||
|
|
@ -94,7 +93,7 @@ impl AzureAuthService {
|
|||
} => {
|
||||
let assertion = resolve_reference(inputs, env_lookup, reference.value())
|
||||
.await?
|
||||
.ok_or(AuthError::UnresolvedOidcReference)?;
|
||||
.ok_or(Error::UnresolvedOidcReference)?;
|
||||
let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion {
|
||||
tenant_id,
|
||||
client_id,
|
||||
|
|
@ -126,7 +125,7 @@ impl AzureAuthService {
|
|||
Err(error) => failures.push(error),
|
||||
}
|
||||
}
|
||||
Err(AuthError::CredentialChain(failures))
|
||||
Err(Error::CredentialChain(failures))
|
||||
}
|
||||
AzureCredentialPlan::Missing => Ok(None),
|
||||
}
|
||||
|
|
@ -136,7 +135,7 @@ impl AzureAuthService {
|
|||
pub(crate) fn select_auth_plan(
|
||||
inputs: &AzureAuthInputs,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<AzureCredentialPlan, AuthError> {
|
||||
) -> Result<AzureCredentialPlan, Error> {
|
||||
let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup);
|
||||
let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup);
|
||||
let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup);
|
||||
|
|
@ -157,7 +156,7 @@ pub(crate) fn select_auth_plan(
|
|||
.map(|selector| Sourced::new(selector, value.source()))
|
||||
})
|
||||
.transpose()
|
||||
.map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureSelector))?;
|
||||
.map_err(|_| Error::InvalidAzureSelector)?;
|
||||
let federated_token_file = configured_string(
|
||||
&inputs.federated_token_file,
|
||||
AZURE_FEDERATED_TOKEN_FILE_ENV,
|
||||
|
|
@ -229,7 +228,7 @@ fn select_native_plan(
|
|||
scope: Sourced<String>,
|
||||
authority: Option<Sourced<String>>,
|
||||
refresh_source: InputSource,
|
||||
) -> Result<AzureCredentialPlan, AuthError> {
|
||||
) -> Result<AzureCredentialPlan, Error> {
|
||||
let selected = selector.unwrap_or_else(|| {
|
||||
Sourced::new(
|
||||
{
|
||||
|
|
@ -247,9 +246,7 @@ fn select_native_plan(
|
|||
let selection_source = selected.source();
|
||||
|
||||
match selected.into_value() {
|
||||
AzureCredentialType::ClientSecretCredential => Err(AuthError::Configuration(
|
||||
AuthConfigurationError::MissingClientSecretFields,
|
||||
)),
|
||||
AzureCredentialType::ClientSecretCredential => Err(Error::MissingClientSecretFields),
|
||||
AzureCredentialType::WorkloadIdentityCredential => {
|
||||
Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
|
||||
workload_request(tenant_id, client_id, federated_token_file, scope, authority)?,
|
||||
|
|
@ -331,17 +328,11 @@ fn workload_request(
|
|||
token_file_path: Option<Sourced<String>>,
|
||||
scope: Sourced<String>,
|
||||
authority: Option<Sourced<String>>,
|
||||
) -> Result<NativeAzureRequest, AuthError> {
|
||||
) -> Result<NativeAzureRequest, Error> {
|
||||
Ok(NativeAzureRequest::WorkloadIdentity {
|
||||
tenant_id: tenant_id.ok_or(AuthError::Configuration(
|
||||
AuthConfigurationError::MissingWorkloadTenant,
|
||||
))?,
|
||||
client_id: client_id.ok_or(AuthError::Configuration(
|
||||
AuthConfigurationError::MissingWorkloadClient,
|
||||
))?,
|
||||
token_file_path: token_file_path.ok_or(AuthError::Configuration(
|
||||
AuthConfigurationError::MissingWorkloadTokenFile,
|
||||
))?,
|
||||
tenant_id: tenant_id.ok_or(Error::MissingWorkloadTenant)?,
|
||||
client_id: client_id.ok_or(Error::MissingWorkloadClient)?,
|
||||
token_file_path: token_file_path.ok_or(Error::MissingWorkloadTokenFile)?,
|
||||
scope,
|
||||
authority,
|
||||
})
|
||||
|
|
@ -383,7 +374,7 @@ async fn resolve_reference(
|
|||
inputs: &AzureAuthInputs,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
reference: &CredentialRef,
|
||||
) -> Result<Option<SecretValue>, AuthError> {
|
||||
) -> Result<Option<SecretValue>, Error> {
|
||||
let lookup = match reference {
|
||||
CredentialRef::Explicit(secret) => return Ok(Some(secret.clone())),
|
||||
CredentialRef::Env(name) => env_lookup(name)
|
||||
|
|
@ -395,9 +386,7 @@ async fn resolve_reference(
|
|||
let resolver = inputs
|
||||
.credential_resolver
|
||||
.as_ref()
|
||||
.ok_or(AuthError::Configuration(
|
||||
AuthConfigurationError::MissingHostResolver,
|
||||
))?;
|
||||
.ok_or(Error::MissingHostResolver)?;
|
||||
resolver.resolve(reference).await?
|
||||
}
|
||||
};
|
||||
|
|
@ -409,15 +398,13 @@ async fn resolve_reference(
|
|||
|
||||
fn oidc_reference(
|
||||
token: &Option<Sourced<SecretValue>>,
|
||||
) -> Result<Option<Sourced<CredentialRef>>, AuthError> {
|
||||
) -> Result<Option<Sourced<CredentialRef>>, Error> {
|
||||
let Some(token) = token.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let value = token.value().expose();
|
||||
if token.source() == InputSource::Request && value.starts_with("oidc/") {
|
||||
return Err(AuthError::Configuration(
|
||||
AuthConfigurationError::RequestAzureCredentialReference,
|
||||
));
|
||||
return Err(Error::RequestAzureCredentialReference);
|
||||
}
|
||||
if let Some(name) = value.strip_prefix("oidc/env/") {
|
||||
return non_empty_reference(name, "OIDC environment reference")
|
||||
|
|
@ -439,18 +426,14 @@ fn oidc_reference(
|
|||
)));
|
||||
}
|
||||
if value.starts_with("oidc/") {
|
||||
return Err(AuthError::Configuration(
|
||||
AuthConfigurationError::UnsupportedOidcReference,
|
||||
));
|
||||
return Err(Error::UnsupportedOidcReference);
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn non_empty_reference(value: &str, kind: &str) -> Result<String, AuthError> {
|
||||
fn non_empty_reference(value: &str, kind: &str) -> Result<String, Error> {
|
||||
if value.is_empty() {
|
||||
return Err(AuthError::Configuration(
|
||||
AuthConfigurationError::EmptyReference(kind.to_string()),
|
||||
));
|
||||
return Err(Error::EmptyReference(kind.to_string()));
|
||||
}
|
||||
Ok(value.to_string())
|
||||
}
|
||||
|
|
@ -466,14 +449,14 @@ mod tests {
|
|||
AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference,
|
||||
resolve_reference, select_auth_plan,
|
||||
};
|
||||
use crate::AuthError;
|
||||
use crate::auth::ResolvedCredential;
|
||||
use crate::auth::{
|
||||
use crate::native::ValidatedAzureRequest;
|
||||
use crate::types::AzureAuthInputs;
|
||||
use litellm_auth::Error;
|
||||
use litellm_auth::ResolvedCredential;
|
||||
use litellm_auth::{
|
||||
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef,
|
||||
CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced,
|
||||
};
|
||||
use crate::providers::azure_ai::auth::native::ValidatedAzureRequest;
|
||||
use crate::providers::azure_ai::auth::types::AzureAuthInputs;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FileResolver;
|
||||
|
|
@ -487,9 +470,8 @@ mod tests {
|
|||
fn acquire(
|
||||
&self,
|
||||
request: ValidatedAzureRequest,
|
||||
) -> std::pin::Pin<
|
||||
Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>,
|
||||
> {
|
||||
) -> std::pin::Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>
|
||||
{
|
||||
let kind = request.kind();
|
||||
self.requests.lock().unwrap().push(kind);
|
||||
Box::pin(async move {
|
||||
|
|
@ -499,7 +481,7 @@ mod tests {
|
|||
expires_on: None,
|
||||
})
|
||||
} else {
|
||||
Err(AuthError::AzureTokenAcquisition(format!("{kind} failed")))
|
||||
Err(Error::AzureTokenAcquisition(format!("{kind} failed")))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -612,12 +594,7 @@ mod tests {
|
|||
})
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
AuthError::Configuration(
|
||||
crate::auth::error::AuthConfigurationError::RequestAzureCredentialReference
|
||||
)
|
||||
));
|
||||
assert!(matches!(error, Error::RequestAzureCredentialReference));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -678,6 +655,6 @@ mod tests {
|
|||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(error, AuthError::CredentialChain(errors) if errors.len() == 2));
|
||||
assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
|
||||
}
|
||||
}
|
||||
|
|
@ -1,10 +1,9 @@
|
|||
use crate::auth::error::AuthConfigurationError;
|
||||
use serde_json::{Map, Value};
|
||||
use std::collections::BTreeMap;
|
||||
use strum::EnumString;
|
||||
|
||||
use crate::AuthError;
|
||||
use crate::auth::{
|
||||
use litellm_auth::Error;
|
||||
use litellm_auth::{
|
||||
CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle,
|
||||
};
|
||||
|
||||
|
|
@ -54,14 +53,14 @@ pub struct AzureAuthInputs {
|
|||
|
||||
impl AzureAuthInputs {
|
||||
#[cfg(test)]
|
||||
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, AuthError> {
|
||||
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, Error> {
|
||||
Self::from_sourced_optional_params(params, &BTreeMap::new())
|
||||
}
|
||||
|
||||
pub fn from_sourced_optional_params(
|
||||
params: &Map<String, Value>,
|
||||
sources: &BTreeMap<String, InputSource>,
|
||||
) -> Result<Self, AuthError> {
|
||||
) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
azure_ad_token: secret_config(params, sources, "azure_ad_token")?,
|
||||
azure_ad_token_provider: None,
|
||||
|
|
@ -88,15 +87,13 @@ fn string_config(
|
|||
params: &Map<String, Value>,
|
||||
sources: &BTreeMap<String, InputSource>,
|
||||
name: &str,
|
||||
) -> Result<ConfigValue<String>, AuthError> {
|
||||
) -> Result<ConfigValue<String>, Error> {
|
||||
let source = source_for(sources, name);
|
||||
match params.get(name) {
|
||||
None => Ok(ConfigValue::Absent),
|
||||
Some(Value::Null) => Ok(ConfigValue::ExplicitNone(source)),
|
||||
Some(Value::String(value)) => Ok(ConfigValue::Value(Sourced::new(value.clone(), source))),
|
||||
Some(_) => Err(AuthError::Configuration(
|
||||
AuthConfigurationError::InvalidFieldType(name.to_string()),
|
||||
)),
|
||||
Some(_) => Err(Error::InvalidFieldType(name.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -104,7 +101,7 @@ fn secret_config(
|
|||
params: &Map<String, Value>,
|
||||
sources: &BTreeMap<String, InputSource>,
|
||||
name: &str,
|
||||
) -> Result<ConfigValue<SecretValue>, AuthError> {
|
||||
) -> Result<ConfigValue<SecretValue>, Error> {
|
||||
Ok(match string_config(params, sources, name)? {
|
||||
ConfigValue::Absent => ConfigValue::Absent,
|
||||
ConfigValue::ExplicitNone(source) => ConfigValue::ExplicitNone(source),
|
||||
|
|
@ -123,7 +120,7 @@ mod tests {
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
|
||||
use crate::auth::{InputSource, Sourced};
|
||||
use litellm_auth::{InputSource, Sourced};
|
||||
|
||||
#[test]
|
||||
fn selector_parsing_is_exact() {
|
||||
15
litellm-rust/crates/auth-gcp/Cargo.toml
Normal file
15
litellm-rust/crates/auth-gcp/Cargo.toml
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
[package]
|
||||
name = "litellm-auth-gcp"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
gcp_auth.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
moka.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
|
|
@ -9,9 +9,8 @@ use moka::future::Cache;
|
|||
use serde_json::{Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::auth::error::AuthConfigurationError;
|
||||
use crate::auth::http::apply_credential;
|
||||
use crate::auth::{AuthError, CredentialPlacement, InputSource, SecretValue, Sourced};
|
||||
use litellm_auth::http::apply_credential;
|
||||
use litellm_auth::{CredentialPlacement, Error, InputSource, SecretValue, Sourced};
|
||||
|
||||
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
|
||||
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
|
||||
|
|
@ -24,17 +23,17 @@ const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
|
|||
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub(crate) struct VertexConfig {
|
||||
pub struct VertexConfig {
|
||||
credentials: Option<Sourced<SecretValue>>,
|
||||
project_id: Option<String>,
|
||||
location: Option<String>,
|
||||
}
|
||||
|
||||
impl VertexConfig {
|
||||
pub(crate) fn from_sourced_optional_params(
|
||||
pub fn from_sourced_optional_params(
|
||||
params: &Map<String, Value>,
|
||||
sources: &BTreeMap<String, InputSource>,
|
||||
) -> Result<Self, AuthError> {
|
||||
) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
credentials: optional_credentials(
|
||||
params,
|
||||
|
|
@ -46,16 +45,16 @@ impl VertexConfig {
|
|||
})
|
||||
}
|
||||
|
||||
pub(crate) fn project_id(&self) -> Option<&str> {
|
||||
pub fn project_id(&self) -> Option<&str> {
|
||||
self.project_id.as_deref()
|
||||
}
|
||||
|
||||
pub(crate) fn location(&self) -> Option<&str> {
|
||||
pub fn location(&self) -> Option<&str> {
|
||||
self.location.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct VertexEnvironment {
|
||||
pub struct VertexEnvironment {
|
||||
pub headers: Vec<(String, String)>,
|
||||
pub project_id: String,
|
||||
}
|
||||
|
|
@ -65,7 +64,7 @@ struct VertexAccessToken {
|
|||
project_id: String,
|
||||
}
|
||||
|
||||
pub(crate) fn get_vertex_ai_project(
|
||||
pub fn get_vertex_ai_project(
|
||||
config: &VertexConfig,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Option<String> {
|
||||
|
|
@ -75,7 +74,7 @@ pub(crate) fn get_vertex_ai_project(
|
|||
.or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV))
|
||||
}
|
||||
|
||||
pub(crate) fn get_vertex_ai_location(
|
||||
pub fn get_vertex_ai_location(
|
||||
config: &VertexConfig,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Option<String> {
|
||||
|
|
@ -87,7 +86,7 @@ pub(crate) fn get_vertex_ai_location(
|
|||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct VertexAuth {
|
||||
pub struct VertexAuth {
|
||||
providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
|
||||
loader: Arc<dyn VertexProviderLoader>,
|
||||
}
|
||||
|
|
@ -106,13 +105,14 @@ impl VertexAuth {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn validate_environment(
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
pub async fn validate_environment(
|
||||
&self,
|
||||
headers: Vec<(String, String)>,
|
||||
api_key: Option<&str>,
|
||||
config: &VertexConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<VertexEnvironment, AuthError> {
|
||||
) -> Result<VertexEnvironment, Error> {
|
||||
let has_authorization = headers
|
||||
.iter()
|
||||
.any(|(name, _)| name.eq_ignore_ascii_case("Authorization"));
|
||||
|
|
@ -160,7 +160,7 @@ impl VertexAuth {
|
|||
&self,
|
||||
config: &VertexConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<VertexAccessToken, AuthError> {
|
||||
) -> Result<VertexAccessToken, Error> {
|
||||
let provider = self.load_provider(config, env_lookup).await?;
|
||||
let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
|
||||
Ok(VertexAccessToken { token, project_id })
|
||||
|
|
@ -170,7 +170,7 @@ impl VertexAuth {
|
|||
&self,
|
||||
config: &VertexConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<Arc<dyn VertexTokenSource>, AuthError> {
|
||||
) -> Result<Arc<dyn VertexTokenSource>, Error> {
|
||||
let source = credential_source(config, env_lookup);
|
||||
let key = source.cache_key();
|
||||
self.providers
|
||||
|
|
@ -189,7 +189,7 @@ trait VertexProviderLoader: Send + Sync {
|
|||
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
|
||||
}
|
||||
|
||||
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, AuthError>> + Send + 'a>>;
|
||||
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
|
||||
struct GcpTokenSource(Arc<dyn TokenProvider>);
|
||||
|
||||
|
|
@ -249,7 +249,7 @@ impl VertexProviderLoader for GcpProviderLoader {
|
|||
}
|
||||
}
|
||||
|
||||
fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
|
||||
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
|
||||
let token_uri = serde_json::from_str::<Value>(configured)
|
||||
.ok()
|
||||
.and_then(|credentials| {
|
||||
|
|
@ -259,7 +259,7 @@ fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
|
|||
.map(str::to_string)
|
||||
});
|
||||
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
|
||||
return Err(AuthConfigurationError::RequestVertexTokenEndpoint.into());
|
||||
return Err(Error::RequestVertexTokenEndpoint);
|
||||
}
|
||||
Ok(configured)
|
||||
}
|
||||
|
|
@ -321,7 +321,7 @@ fn optional_credentials(
|
|||
params: &Map<String, Value>,
|
||||
sources: &BTreeMap<String, InputSource>,
|
||||
names: &[&str],
|
||||
) -> Result<Option<Sourced<SecretValue>>, AuthError> {
|
||||
) -> Result<Option<Sourced<SecretValue>>, Error> {
|
||||
for name in names {
|
||||
let source = source_for(sources, name);
|
||||
match params.get(*name) {
|
||||
|
|
@ -336,17 +336,10 @@ fn optional_credentials(
|
|||
.map(SecretValue::new)
|
||||
.map(|value| Sourced::new(value, source))
|
||||
.map(Some)
|
||||
.map_err(|error| {
|
||||
AuthError::Configuration(AuthConfigurationError::InvalidFieldType(format!(
|
||||
"{}: {error}",
|
||||
names[0]
|
||||
)))
|
||||
});
|
||||
.map_err(|error| Error::InvalidFieldType(format!("{}: {error}", names[0])));
|
||||
}
|
||||
Some(_) => {
|
||||
return Err(AuthError::Configuration(
|
||||
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
|
||||
));
|
||||
return Err(Error::InvalidFieldType(names[0].to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -357,19 +350,14 @@ fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSourc
|
|||
sources.get(name).copied().unwrap_or_default()
|
||||
}
|
||||
|
||||
fn optional_string(
|
||||
params: &Map<String, Value>,
|
||||
names: &[&str],
|
||||
) -> Result<Option<String>, AuthError> {
|
||||
fn optional_string(params: &Map<String, Value>, names: &[&str]) -> Result<Option<String>, Error> {
|
||||
for name in names {
|
||||
match params.get(*name) {
|
||||
None | Some(Value::Null) => continue,
|
||||
Some(Value::String(value)) if value.trim().is_empty() => continue,
|
||||
Some(Value::String(value)) => return Ok(Some(value.clone())),
|
||||
Some(_) => {
|
||||
return Err(AuthError::Configuration(
|
||||
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
|
||||
));
|
||||
return Err(Error::InvalidFieldType(names[0].to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -382,8 +370,8 @@ fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option<String>, name: &str) -> Opt
|
|||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn auth_acquisition_error(error: gcp_auth::Error) -> AuthError {
|
||||
AuthError::VertexTokenAcquisition(error.to_string())
|
||||
fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
|
||||
Error::VertexTokenAcquisition(error.to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -537,15 +525,11 @@ mod tests {
|
|||
);
|
||||
assert!(matches!(
|
||||
validate_request_credentials(r#"{"token_uri":"http://127.0.0.1/token"}"#),
|
||||
Err(AuthError::Configuration(
|
||||
AuthConfigurationError::RequestVertexTokenEndpoint
|
||||
))
|
||||
Err(Error::RequestVertexTokenEndpoint)
|
||||
));
|
||||
assert!(matches!(
|
||||
validate_request_credentials("{}"),
|
||||
Err(AuthError::Configuration(
|
||||
AuthConfigurationError::RequestVertexTokenEndpoint
|
||||
))
|
||||
Err(Error::RequestVertexTokenEndpoint)
|
||||
));
|
||||
}
|
||||
|
||||
15
litellm-rust/crates/auth/Cargo.toml
Normal file
15
litellm-rust/crates/auth/Cargo.toml
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
[package]
|
||||
name = "litellm-auth"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde.workspace = true
|
||||
subtle.workspace = true
|
||||
thiserror.workspace = true
|
||||
veil.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tokio.workspace = true
|
||||
|
|
@ -5,7 +5,7 @@ use std::sync::Arc;
|
|||
|
||||
use veil::Redact;
|
||||
|
||||
use crate::AuthError;
|
||||
use crate::Error;
|
||||
|
||||
use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
|
||||
|
||||
|
|
@ -48,7 +48,7 @@ pub enum CredentialLookup {
|
|||
}
|
||||
|
||||
pub type CredentialLookupFuture<'a> =
|
||||
Pin<Box<dyn Future<Output = Result<CredentialLookup, AuthError>> + Send + 'a>>;
|
||||
Pin<Box<dyn Future<Output = Result<CredentialLookup, Error>> + Send + 'a>>;
|
||||
|
||||
pub trait CredentialResolver: std::fmt::Debug + Send + Sync {
|
||||
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>;
|
||||
|
|
@ -62,7 +62,7 @@ impl CredentialResolverHandle {
|
|||
Self(resolver)
|
||||
}
|
||||
|
||||
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, AuthError> {
|
||||
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, Error> {
|
||||
self.0.resolve(reference).await
|
||||
}
|
||||
}
|
||||
|
|
@ -84,7 +84,7 @@ impl CredentialPlan {
|
|||
pub async fn resolve(
|
||||
&self,
|
||||
resolver: &CredentialResolverHandle,
|
||||
) -> Result<CredentialPlanResolution, AuthError> {
|
||||
) -> Result<CredentialPlanResolution, Error> {
|
||||
match self {
|
||||
Self::Static(CredentialRef::Explicit(secret)) => Ok(
|
||||
CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())),
|
||||
|
|
@ -103,7 +103,7 @@ impl CredentialPlan {
|
|||
Self::Caller(caller) => {
|
||||
let credential = caller.acquire().await?;
|
||||
if credential.secret().expose().is_empty() {
|
||||
return Err(AuthError::EmptyCallerCredential);
|
||||
return Err(Error::EmptyCallerCredential);
|
||||
}
|
||||
Ok(CredentialPlanResolution::Resolved(credential))
|
||||
}
|
||||
|
|
@ -119,8 +119,8 @@ mod tests {
|
|||
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
|
||||
CredentialRef, CredentialResolver, CredentialResolverHandle,
|
||||
};
|
||||
use crate::AuthError;
|
||||
use crate::auth::SecretValue;
|
||||
use crate::Error;
|
||||
use crate::SecretValue;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct HostResolver;
|
||||
|
|
@ -164,7 +164,7 @@ mod tests {
|
|||
|
||||
impl CredentialResolver for FailingResolver {
|
||||
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
|
||||
Box::pin(async { Err(AuthError::UnresolvedOidcReference) })
|
||||
Box::pin(async { Err(Error::UnresolvedOidcReference) })
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -178,6 +178,6 @@ mod tests {
|
|||
.await
|
||||
.expect_err("acquisition errors cannot become fallback");
|
||||
|
||||
assert_eq!(error, AuthError::UnresolvedOidcReference);
|
||||
assert_eq!(error, Error::UnresolvedOidcReference);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,15 +1,77 @@
|
|||
use thiserror::Error;
|
||||
use thiserror::Error as ThisError;
|
||||
|
||||
#[derive(Clone, Debug, Error, PartialEq, Eq)]
|
||||
pub enum AuthError {
|
||||
#[error("invalid authentication configuration: {0}")]
|
||||
Configuration(#[from] AuthConfigurationError),
|
||||
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
|
||||
pub enum Error {
|
||||
#[error("invalid authentication configuration: credential header already exists")]
|
||||
ExistingCredentialHeader,
|
||||
#[error(
|
||||
"invalid authentication configuration: credential plan is not allowed by the provider auth policy"
|
||||
)]
|
||||
DisallowedCredentialPlan,
|
||||
#[error("invalid authentication configuration: credential cannot be empty")]
|
||||
EmptyCredential,
|
||||
#[error("invalid authentication configuration: invalid Azure credential selector")]
|
||||
InvalidAzureSelector,
|
||||
#[error(
|
||||
"invalid authentication configuration: ClientSecretCredential requires tenant_id, client_id, and client_secret"
|
||||
)]
|
||||
MissingClientSecretFields,
|
||||
#[error("invalid authentication configuration: WorkloadIdentityCredential requires tenant_id")]
|
||||
MissingWorkloadTenant,
|
||||
#[error("invalid authentication configuration: WorkloadIdentityCredential requires client_id")]
|
||||
MissingWorkloadClient,
|
||||
#[error(
|
||||
"invalid authentication configuration: WorkloadIdentityCredential requires azure_federated_token_file"
|
||||
)]
|
||||
MissingWorkloadTokenFile,
|
||||
#[error(
|
||||
"invalid authentication configuration: credential reference requires a host credential resolver"
|
||||
)]
|
||||
MissingHostResolver,
|
||||
#[error(
|
||||
"invalid authentication configuration: caller credential plan requires provider-specific inputs"
|
||||
)]
|
||||
MissingCallerInputs,
|
||||
#[error("invalid authentication configuration: credential header {0} already exists")]
|
||||
DuplicateHeader(&'static str),
|
||||
#[error("invalid authentication configuration: {0} must be a string or null")]
|
||||
InvalidFieldType(String),
|
||||
#[error("invalid authentication configuration: unsupported OIDC reference")]
|
||||
UnsupportedOidcReference,
|
||||
#[error("invalid authentication configuration: {0} cannot be empty")]
|
||||
EmptyReference(String),
|
||||
#[error("invalid authentication configuration: Azure credential initialization failed: {0}")]
|
||||
AzureCredentialInitialization(String),
|
||||
#[error(
|
||||
"invalid authentication configuration: Azure authority must be an HTTPS origin without credentials, query, or fragment"
|
||||
)]
|
||||
InvalidAzureAuthority,
|
||||
#[error(
|
||||
"invalid authentication configuration: request-controlled Azure auth inputs cannot be combined with host credentials"
|
||||
)]
|
||||
MixedAzureCredentialSources,
|
||||
#[error(
|
||||
"invalid authentication configuration: request-controlled Azure credential references are not allowed"
|
||||
)]
|
||||
RequestAzureCredentialReference,
|
||||
#[error(
|
||||
"invalid authentication configuration: host credentials cannot be sent to a request-controlled Azure endpoint"
|
||||
)]
|
||||
RequestAzureCredentialDestination,
|
||||
#[error(
|
||||
"invalid authentication configuration: credentials cannot be sent to a request-controlled Vertex AI endpoint"
|
||||
)]
|
||||
RequestVertexCredentialDestination,
|
||||
#[error(
|
||||
"invalid authentication configuration: request-controlled Vertex credentials must use the canonical Google OAuth token endpoint"
|
||||
)]
|
||||
RequestVertexTokenEndpoint,
|
||||
#[error("credential acquisition failed: {0}")]
|
||||
AzureTokenAcquisition(String),
|
||||
#[error("credential acquisition failed: Vertex AI credentials: {0}")]
|
||||
VertexTokenAcquisition(String),
|
||||
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
|
||||
CredentialChain(Vec<AuthError>),
|
||||
CredentialChain(Vec<Error>),
|
||||
#[error("credential caller failed: credential caller returned an empty credential")]
|
||||
EmptyCallerCredential,
|
||||
#[error("credential caller failed: Azure AD token provider returned an empty token")]
|
||||
|
|
@ -27,102 +89,42 @@ pub enum AuthError {
|
|||
provider: &'static str,
|
||||
environment_variable: &'static str,
|
||||
},
|
||||
#[error("{0}")]
|
||||
MissingCredential(#[from] MissingCredential),
|
||||
#[error("{0}")]
|
||||
Aws(#[from] AwsAuthError),
|
||||
#[error("invalid authentication header")]
|
||||
InvalidHeader,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Error, PartialEq, Eq)]
|
||||
pub enum AuthConfigurationError {
|
||||
#[error("credential header already exists")]
|
||||
ExistingCredentialHeader,
|
||||
#[error("credential plan is not allowed by the provider auth policy")]
|
||||
DisallowedCredentialPlan,
|
||||
#[error("credential cannot be empty")]
|
||||
EmptyCredential,
|
||||
#[error("invalid Azure credential selector")]
|
||||
InvalidAzureSelector,
|
||||
#[error("ClientSecretCredential requires tenant_id, client_id, and client_secret")]
|
||||
MissingClientSecretFields,
|
||||
#[error("WorkloadIdentityCredential requires tenant_id")]
|
||||
MissingWorkloadTenant,
|
||||
#[error("WorkloadIdentityCredential requires client_id")]
|
||||
MissingWorkloadClient,
|
||||
#[error("WorkloadIdentityCredential requires azure_federated_token_file")]
|
||||
MissingWorkloadTokenFile,
|
||||
#[error("credential reference requires a host credential resolver")]
|
||||
MissingHostResolver,
|
||||
#[error("caller credential plan requires provider-specific inputs")]
|
||||
MissingCallerInputs,
|
||||
#[error("credential header {0} already exists")]
|
||||
DuplicateHeader(&'static str),
|
||||
#[error("{0} must be a string or null")]
|
||||
InvalidFieldType(String),
|
||||
#[error("unsupported OIDC reference")]
|
||||
UnsupportedOidcReference,
|
||||
#[error("{0} cannot be empty")]
|
||||
EmptyReference(String),
|
||||
#[error("Azure credential initialization failed: {0}")]
|
||||
AzureCredentialInitialization(String),
|
||||
#[error("Azure authority must be an HTTPS origin without credentials, query, or fragment")]
|
||||
InvalidAzureAuthority,
|
||||
#[error("request-controlled Azure auth inputs cannot be combined with host credentials")]
|
||||
MixedAzureCredentialSources,
|
||||
#[error("request-controlled Azure credential references are not allowed")]
|
||||
RequestAzureCredentialReference,
|
||||
#[error("host credentials cannot be sent to a request-controlled Azure endpoint")]
|
||||
RequestAzureCredentialDestination,
|
||||
#[error("credentials cannot be sent to a request-controlled Vertex AI endpoint")]
|
||||
RequestVertexCredentialDestination,
|
||||
#[error(
|
||||
"request-controlled Vertex credentials must use the canonical Google OAuth token endpoint"
|
||||
)]
|
||||
RequestVertexTokenEndpoint,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Error, PartialEq, Eq)]
|
||||
pub enum MissingCredential {
|
||||
#[error(
|
||||
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
|
||||
)]
|
||||
AnthropicApiKey,
|
||||
MissingAnthropicApiKey,
|
||||
#[error("Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable")]
|
||||
AzureApiKey,
|
||||
MissingAzureApiKey,
|
||||
#[error(
|
||||
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
|
||||
)]
|
||||
AzureApiBase,
|
||||
MissingAzureApiBase,
|
||||
#[error(
|
||||
"Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
|
||||
)]
|
||||
OpenAiRealtimeApiKey,
|
||||
MissingOpenAiRealtimeApiKey,
|
||||
#[error(
|
||||
"Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
|
||||
)]
|
||||
OpenAiResponsesApiKey,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Error, PartialEq, Eq)]
|
||||
pub enum AwsAuthError {
|
||||
MissingOpenAiResponsesApiKey,
|
||||
#[error("AWS profile credentials failed: {0}")]
|
||||
Profile(String),
|
||||
AwsProfile(String),
|
||||
#[error("AWS default credentials failed: {0}")]
|
||||
DefaultChain(String),
|
||||
AwsDefaultChain(String),
|
||||
#[error("AWS role credentials failed: {0}")]
|
||||
AssumeRole(String),
|
||||
AwsAssumeRole(String),
|
||||
#[error("AWS web identity credentials failed: {0}")]
|
||||
WebIdentity(String),
|
||||
AwsWebIdentity(String),
|
||||
#[error("AWS web identity expiration was invalid: {0}")]
|
||||
WebIdentityExpiration(String),
|
||||
AwsWebIdentityExpiration(String),
|
||||
#[error("AWS signing parameters failed: {0}")]
|
||||
SigningParameters(String),
|
||||
AwsSigningParameters(String),
|
||||
#[error("AWS signable request failed: {0}")]
|
||||
SignableRequest(String),
|
||||
AwsSignableRequest(String),
|
||||
#[error("AWS request signing failed: {0}")]
|
||||
Signing(String),
|
||||
AwsSigning(String),
|
||||
#[error("AWS web identity response had no credentials")]
|
||||
MissingWebIdentityCredentials,
|
||||
AwsMissingWebIdentityCredentials,
|
||||
#[error("invalid authentication header")]
|
||||
InvalidHeader,
|
||||
}
|
||||
|
|
@ -1,5 +1,4 @@
|
|||
use crate::AuthError;
|
||||
use crate::auth::error::AuthConfigurationError;
|
||||
use crate::Error;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum CredentialPlacement {
|
||||
|
|
@ -16,23 +15,19 @@ impl CredentialPlacement {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) fn apply_credential(
|
||||
pub fn apply_credential(
|
||||
headers: Vec<(String, String)>,
|
||||
credential: &str,
|
||||
placement: CredentialPlacement,
|
||||
) -> Result<Vec<(String, String)>, AuthError> {
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
if credential.trim().is_empty() {
|
||||
return Err(AuthError::Configuration(
|
||||
AuthConfigurationError::EmptyCredential,
|
||||
));
|
||||
return Err(Error::EmptyCredential);
|
||||
}
|
||||
if headers
|
||||
.iter()
|
||||
.any(|(name, _)| name.eq_ignore_ascii_case(placement.header_name()))
|
||||
{
|
||||
return Err(AuthError::Configuration(
|
||||
AuthConfigurationError::DuplicateHeader(placement.header_name()),
|
||||
));
|
||||
return Err(Error::DuplicateHeader(placement.header_name()));
|
||||
}
|
||||
let value = match placement {
|
||||
CredentialPlacement::Bearer => format!("Bearer {credential}"),
|
||||
56
litellm-rust/crates/auth/src/lib.rs
Normal file
56
litellm-rust/crates/auth/src/lib.rs
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
mod credential;
|
||||
mod error;
|
||||
pub mod http;
|
||||
mod policy;
|
||||
mod secret;
|
||||
mod token;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum InputSource {
|
||||
Request,
|
||||
#[default]
|
||||
Deployment,
|
||||
Environment,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Eq, PartialEq)]
|
||||
pub struct Sourced<T> {
|
||||
value: T,
|
||||
source: InputSource,
|
||||
}
|
||||
|
||||
impl<T> Sourced<T> {
|
||||
pub fn new(value: T, source: InputSource) -> Self {
|
||||
Self { value, source }
|
||||
}
|
||||
|
||||
pub fn value(&self) -> &T {
|
||||
&self.value
|
||||
}
|
||||
|
||||
pub fn source(&self) -> InputSource {
|
||||
self.source
|
||||
}
|
||||
|
||||
pub fn into_value(self) -> T {
|
||||
self.value
|
||||
}
|
||||
|
||||
pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
|
||||
Sourced::new(map(self.value), self.source)
|
||||
}
|
||||
}
|
||||
|
||||
pub use credential::{
|
||||
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
|
||||
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
|
||||
credential_default_fields, credential_index,
|
||||
};
|
||||
pub use error::Error;
|
||||
pub use http::{CredentialPlacement, RequestAuth};
|
||||
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
|
||||
pub use secret::SecretValue;
|
||||
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
|
||||
|
|
@ -1,5 +1,4 @@
|
|||
use crate::AuthError;
|
||||
use crate::auth::error::AuthConfigurationError;
|
||||
use crate::Error;
|
||||
|
||||
use super::http::apply_credential;
|
||||
use super::{CredentialPlacement, ResolvedCredential};
|
||||
|
|
@ -46,22 +45,18 @@ impl ProviderAuthPolicy {
|
|||
headers: Vec<(String, String)>,
|
||||
kind: CredentialPlanKind,
|
||||
credential: &ResolvedCredential,
|
||||
) -> Result<Vec<(String, String)>, AuthError> {
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
if self.has_existing_credential(&headers) {
|
||||
return match self.existing_header_behavior {
|
||||
ExistingHeaderBehavior::Preserve => Ok(headers),
|
||||
ExistingHeaderBehavior::Reject => Err(AuthError::Configuration(
|
||||
AuthConfigurationError::ExistingCredentialHeader,
|
||||
)),
|
||||
ExistingHeaderBehavior::Reject => Err(Error::ExistingCredentialHeader),
|
||||
};
|
||||
}
|
||||
let rule =
|
||||
self.rules
|
||||
.iter()
|
||||
.find(|rule| rule.kind == kind)
|
||||
.ok_or(AuthError::Configuration(
|
||||
AuthConfigurationError::DisallowedCredentialPlan,
|
||||
))?;
|
||||
let rule = self
|
||||
.rules
|
||||
.iter()
|
||||
.find(|rule| rule.kind == kind)
|
||||
.ok_or(Error::DisallowedCredentialPlan)?;
|
||||
apply_credential(headers, credential.secret().expose(), rule.placement)
|
||||
}
|
||||
}
|
||||
|
|
@ -69,7 +64,7 @@ impl ProviderAuthPolicy {
|
|||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
|
||||
use crate::auth::{CredentialPlacement, ResolvedCredential, SecretValue};
|
||||
use crate::{CredentialPlacement, ResolvedCredential, SecretValue};
|
||||
|
||||
const RULES: &[CredentialRule] = &[CredentialRule {
|
||||
kind: CredentialPlanKind::Static,
|
||||
|
|
@ -5,7 +5,7 @@ use std::time::SystemTime;
|
|||
|
||||
use veil::Redact;
|
||||
|
||||
use crate::AuthError;
|
||||
use crate::Error;
|
||||
|
||||
use super::secret::SecretValue;
|
||||
|
||||
|
|
@ -27,7 +27,7 @@ impl ResolvedCredential {
|
|||
}
|
||||
|
||||
pub type TokenFuture<'a> =
|
||||
Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + 'a>>;
|
||||
Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + 'a>>;
|
||||
|
||||
pub trait TokenProvider: std::fmt::Debug + Send + Sync {
|
||||
fn acquire(&self) -> TokenFuture<'_>;
|
||||
|
|
@ -41,7 +41,7 @@ impl TokenProviderHandle {
|
|||
Self(caller)
|
||||
}
|
||||
|
||||
pub async fn acquire(&self) -> Result<ResolvedCredential, AuthError> {
|
||||
pub async fn acquire(&self) -> Result<ResolvedCredential, Error> {
|
||||
self.0.acquire().await
|
||||
}
|
||||
}
|
||||
|
|
@ -10,10 +10,11 @@ autotests = false
|
|||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
base64.workspace = true
|
||||
azure_core.workspace = true
|
||||
azure_identity.workspace = true
|
||||
data-url = "0.3.2"
|
||||
gcp_auth.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
litellm-auth-aws = { workspace = true, optional = true }
|
||||
litellm-auth-azure.workspace = true
|
||||
litellm-auth-gcp.workspace = true
|
||||
moka.workspace = true
|
||||
mime_guess = "2.0.5"
|
||||
rand.workspace = true
|
||||
|
|
@ -23,7 +24,6 @@ rustls-native-certs.workspace = true
|
|||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
serde_path_to_error = "0.1"
|
||||
strum.workspace = true
|
||||
subtle.workspace = true
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
tokio-tungstenite.workspace = true
|
||||
|
|
@ -31,23 +31,10 @@ thiserror.workspace = true
|
|||
sha2.workspace = true
|
||||
url.workspace = true
|
||||
veil.workspace = true
|
||||
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
|
||||
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
|
||||
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
|
||||
aws-sigv4 = { version = "1.5.1", optional = true }
|
||||
aws-types = { version = "1.4.0", optional = true }
|
||||
aws-smithy-runtime-api = { version = "1.13.0", optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
bedrock-auth = [
|
||||
"dep:aws-config",
|
||||
"dep:aws-credential-types",
|
||||
"dep:aws-sdk-sts",
|
||||
"dep:aws-sigv4",
|
||||
"dep:aws-types",
|
||||
"dep:aws-smithy-runtime-api",
|
||||
]
|
||||
bedrock-auth = ["dep:litellm-auth-aws"]
|
||||
observability = ["dep:tracing-subscriber"]
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,57 +1 @@
|
|||
mod credential;
|
||||
pub mod error;
|
||||
pub(crate) mod vertex;
|
||||
pub use error::AuthError;
|
||||
pub(crate) mod http;
|
||||
mod policy;
|
||||
mod secret;
|
||||
mod token;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum InputSource {
|
||||
Request,
|
||||
#[default]
|
||||
Deployment,
|
||||
Environment,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Eq, PartialEq)]
|
||||
pub struct Sourced<T> {
|
||||
value: T,
|
||||
source: InputSource,
|
||||
}
|
||||
|
||||
impl<T> Sourced<T> {
|
||||
pub fn new(value: T, source: InputSource) -> Self {
|
||||
Self { value, source }
|
||||
}
|
||||
|
||||
pub fn value(&self) -> &T {
|
||||
&self.value
|
||||
}
|
||||
|
||||
pub fn source(&self) -> InputSource {
|
||||
self.source
|
||||
}
|
||||
|
||||
pub fn into_value(self) -> T {
|
||||
self.value
|
||||
}
|
||||
|
||||
pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
|
||||
Sourced::new(map(self.value), self.source)
|
||||
}
|
||||
}
|
||||
|
||||
pub use credential::{
|
||||
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
|
||||
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
|
||||
credential_default_fields, credential_index,
|
||||
};
|
||||
pub use http::{CredentialPlacement, RequestAuth};
|
||||
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
|
||||
pub use secret::SecretValue;
|
||||
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
|
||||
pub use litellm_auth::*;
|
||||
|
|
|
|||
|
|
@ -15,5 +15,5 @@ pub mod providers;
|
|||
pub mod responses;
|
||||
mod url_utils;
|
||||
|
||||
pub use auth::AuthError;
|
||||
pub use auth::Error as AuthError;
|
||||
pub use error::Error;
|
||||
|
|
|
|||
|
|
@ -7,8 +7,8 @@ use crate::ocr::document::{inline_remote_document, validate_inline_document};
|
|||
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
|
||||
use crate::ocr::prepare::{credential_env, transform_request_body};
|
||||
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
|
||||
use crate::providers::azure_ai::auth::AzureAuthInputs;
|
||||
use crate::url_utils::ApiUrl;
|
||||
use litellm_auth_azure::AzureAuthInputs;
|
||||
|
||||
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,10 @@
|
|||
use std::sync::OnceLock;
|
||||
|
||||
use crate::Error;
|
||||
use crate::auth::error::AuthConfigurationError;
|
||||
use crate::auth::{InputSource, Sourced};
|
||||
use crate::ocr::error::OcrError;
|
||||
use crate::ocr::types::OcrConnection;
|
||||
use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService};
|
||||
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
|
||||
|
||||
pub(super) async fn resolve_entra(
|
||||
config: &AzureAuthInputs,
|
||||
|
|
@ -38,10 +37,7 @@ pub(super) fn validate_destination(
|
|||
&& connection.api_base_source == InputSource::Request
|
||||
&& credential_source != InputSource::Request
|
||||
{
|
||||
return Err(Error::from(crate::AuthError::Configuration(
|
||||
AuthConfigurationError::RequestAzureCredentialDestination,
|
||||
))
|
||||
.into());
|
||||
return Err(Error::from(crate::AuthError::RequestAzureCredentialDestination).into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -621,8 +621,8 @@ use crate::ocr::OcrClient;
|
|||
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
|
||||
use crate::ocr::prepare::{credential_env, transform_request_body};
|
||||
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat};
|
||||
use crate::providers::azure_ai::auth::AzureAuthInputs;
|
||||
use crate::url_utils::ApiUrl;
|
||||
use litellm_auth_azure::AzureAuthInputs;
|
||||
|
||||
const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
|
||||
const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
|
||||
|
|
|
|||
|
|
@ -9,8 +9,8 @@ use crate::ocr::document::{inline_remote_document, validate_inline_document};
|
|||
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
|
||||
use crate::ocr::prepare::{credential_env, transform_request_body};
|
||||
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
|
||||
use crate::providers::azure_ai::auth::AzureAuthInputs;
|
||||
use crate::url_utils::ApiUrl;
|
||||
use litellm_auth_azure::AzureAuthInputs;
|
||||
|
||||
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
|
||||
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
|
||||
|
|
|
|||
|
|
@ -1,15 +1,11 @@
|
|||
use crate::Error;
|
||||
use crate::auth::InputSource;
|
||||
use crate::auth::error::AuthConfigurationError;
|
||||
use crate::ocr::error::OcrError;
|
||||
use crate::ocr::types::OcrConnection;
|
||||
|
||||
pub(super) fn validate_destination(connection: &OcrConnection) -> Result<(), OcrError> {
|
||||
if connection.api_base.is_some() && connection.api_base_source == InputSource::Request {
|
||||
return Err(Error::from(crate::AuthError::Configuration(
|
||||
AuthConfigurationError::RequestVertexCredentialDestination,
|
||||
))
|
||||
.into());
|
||||
return Err(Error::from(crate::AuthError::RequestVertexCredentialDestination).into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -213,7 +213,6 @@ pub(crate) use mapping::{transform_ocr_request, transform_ocr_response};
|
|||
|
||||
use super::common_utils::validate_destination;
|
||||
use crate::Error;
|
||||
use crate::auth::vertex::{self, VertexConfig};
|
||||
use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
|
||||
use crate::ocr::OcrClient;
|
||||
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
|
||||
|
|
@ -222,6 +221,7 @@ use crate::ocr::prepare::{
|
|||
};
|
||||
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
|
||||
use crate::url_utils::ApiUrl;
|
||||
use litellm_auth_gcp::{self as vertex, VertexConfig};
|
||||
const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com";
|
||||
const MODEL_NAMESPACE: &str = "deepseek-ai";
|
||||
const DEFAULT_LOCATION: &str = "us-central1";
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
use super::common_utils::validate_destination;
|
||||
use crate::Error;
|
||||
use crate::auth::vertex::{self, VertexConfig};
|
||||
use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
|
||||
use crate::llms::mistral::ocr::MistralOcrResponse;
|
||||
use crate::llms::mistral::ocr::transformation::MistralOCRConfig;
|
||||
|
|
@ -10,6 +9,7 @@ use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
|
|||
use crate::ocr::prepare::{credential_env, transform_request_body};
|
||||
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
|
||||
use crate::url_utils::ApiUrl;
|
||||
use litellm_auth_gcp::{self as vertex, VertexConfig};
|
||||
const DEFAULT_LOCATION: &str = "us-central1";
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
|
|
|
|||
|
|
@ -8,10 +8,10 @@ use super::error::{OcrError, OcrResponseError};
|
|||
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
|
||||
use super::wire::{DecodedOcrResponse, decode_response};
|
||||
use crate::Error;
|
||||
use crate::auth::vertex::VertexAuth;
|
||||
use crate::constants::OCR_CONNECT_TIMEOUT_SECS;
|
||||
use crate::error::TransportError;
|
||||
use crate::media::MediaFetcher;
|
||||
use litellm_auth_gcp::VertexAuth;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OcrClient {
|
||||
|
|
|
|||
|
|
@ -66,13 +66,9 @@ impl PreparedOcrCall {
|
|||
let http = match request.config {
|
||||
OcrConfigKind::Cohere => CohereParseConfig.prepare_request(&request, &client).await?,
|
||||
OcrConfigKind::Mistral => MistralOCRConfig.prepare_request(&request, &client).await?,
|
||||
OcrConfigKind::AzureAi => {
|
||||
AzureAIOCRConfig::default()
|
||||
.prepare_request(&request, &client)
|
||||
.await?
|
||||
}
|
||||
OcrConfigKind::AzureAi => AzureAIOCRConfig.prepare_request(&request, &client).await?,
|
||||
OcrConfigKind::AzureCohere => {
|
||||
AzureAICohereParseConfig::default()
|
||||
AzureAICohereParseConfig
|
||||
.prepare_request(&request, &client)
|
||||
.await?
|
||||
}
|
||||
|
|
@ -91,11 +87,7 @@ impl PreparedOcrCall {
|
|||
.prepare_request(&request, &client)
|
||||
.await?
|
||||
}
|
||||
OcrConfigKind::VertexAi => {
|
||||
VertexAIOCRConfig::default()
|
||||
.prepare_request(&request, &client)
|
||||
.await?
|
||||
}
|
||||
OcrConfigKind::VertexAi => VertexAIOCRConfig.prepare_request(&request, &client).await?,
|
||||
OcrConfigKind::VertexDeepSeek => {
|
||||
VertexAIDeepSeekOCRConfig
|
||||
.prepare_request(&request, &client)
|
||||
|
|
@ -130,12 +122,12 @@ impl PreparedOcrCall {
|
|||
.await?,
|
||||
),
|
||||
OcrConfigKind::AzureAi => OcrProviderData::AzureAi(
|
||||
AzureAIOCRConfig::default()
|
||||
AzureAIOCRConfig
|
||||
.read_response(&self.client, response, &url, &headers, &self.request)
|
||||
.await?,
|
||||
),
|
||||
OcrConfigKind::AzureCohere => OcrProviderData::AzureCohere(
|
||||
AzureAICohereParseConfig::default()
|
||||
AzureAICohereParseConfig
|
||||
.read_response(&self.client, response, &url, &headers, &self.request)
|
||||
.await?,
|
||||
),
|
||||
|
|
@ -155,7 +147,7 @@ impl PreparedOcrCall {
|
|||
.await?,
|
||||
),
|
||||
OcrConfigKind::VertexAi => OcrProviderData::VertexAi(
|
||||
VertexAIOCRConfig::default()
|
||||
VertexAIOCRConfig
|
||||
.read_response(&self.client, response, &url, &headers, &self.request)
|
||||
.await?,
|
||||
),
|
||||
|
|
@ -217,12 +209,11 @@ impl OcrProviderResponse {
|
|||
decoded.native,
|
||||
),
|
||||
OcrProviderData::AzureAi(decoded) => (
|
||||
AzureAIOCRConfig::default().transform_ocr_response(&self.request, decoded.data)?,
|
||||
AzureAIOCRConfig.transform_ocr_response(&self.request, decoded.data)?,
|
||||
decoded.native,
|
||||
),
|
||||
OcrProviderData::AzureCohere(decoded) => (
|
||||
AzureAICohereParseConfig::default()
|
||||
.transform_ocr_response(&self.request, decoded.data)?,
|
||||
AzureAICohereParseConfig.transform_ocr_response(&self.request, decoded.data)?,
|
||||
decoded.native,
|
||||
),
|
||||
OcrProviderData::AzureDocumentIntelligence(decoded) => (
|
||||
|
|
@ -239,7 +230,7 @@ impl OcrProviderResponse {
|
|||
decoded.native,
|
||||
),
|
||||
OcrProviderData::VertexAi(decoded) => (
|
||||
VertexAIOCRConfig::default().transform_ocr_response(&self.request, decoded.data)?,
|
||||
VertexAIOCRConfig.transform_ocr_response(&self.request, decoded.data)?,
|
||||
decoded.native,
|
||||
),
|
||||
OcrProviderData::VertexDeepSeek(decoded) => (
|
||||
|
|
|
|||
|
|
@ -40,16 +40,14 @@ impl OcrConfigKind {
|
|||
match self {
|
||||
Self::Cohere => CohereParseConfig.get_supported_ocr_params(model),
|
||||
Self::Mistral => MistralOCRConfig.get_supported_ocr_params(model),
|
||||
Self::AzureAi => AzureAIOCRConfig::default().get_supported_ocr_params(model),
|
||||
Self::AzureCohere => {
|
||||
AzureAICohereParseConfig::default().get_supported_ocr_params(model)
|
||||
}
|
||||
Self::AzureAi => AzureAIOCRConfig.get_supported_ocr_params(model),
|
||||
Self::AzureCohere => AzureAICohereParseConfig.get_supported_ocr_params(model),
|
||||
Self::AzureDocumentIntelligence => {
|
||||
AzureDocumentIntelligenceOCRConfig.get_supported_ocr_params(model)
|
||||
}
|
||||
Self::ReductoLegacy => ReductoParseLegacyConfig.get_supported_ocr_params(model),
|
||||
Self::ReductoV3 => ReductoParseV3Config.get_supported_ocr_params(model),
|
||||
Self::VertexAi => VertexAIOCRConfig::default().get_supported_ocr_params(model),
|
||||
Self::VertexAi => VertexAIOCRConfig.get_supported_ocr_params(model),
|
||||
Self::VertexDeepSeek => VertexAIDeepSeekOCRConfig.get_supported_ocr_params(model),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
use crate::auth::error::MissingCredential;
|
||||
use crate::error::Error;
|
||||
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
|
||||
|
||||
|
|
@ -22,7 +21,7 @@ pub fn resolve_anthropic_api_key(
|
|||
non_empty(api_key)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AnthropicApiKey)))
|
||||
.ok_or_else(|| Error::from(crate::AuthError::MissingAnthropicApiKey))
|
||||
}
|
||||
|
||||
pub fn complete_anthropic_url(
|
||||
|
|
|
|||
|
|
@ -1,7 +0,0 @@
|
|||
mod credential_provider_cache;
|
||||
mod native;
|
||||
mod resolve;
|
||||
mod types;
|
||||
|
||||
pub(crate) use resolve::AzureAuthService;
|
||||
pub(crate) use types::AzureAuthInputs;
|
||||
|
|
@ -1,4 +1,3 @@
|
|||
use crate::auth::error::MissingCredential;
|
||||
use crate::error::Error;
|
||||
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
|
||||
use crate::messages::types::{
|
||||
|
|
@ -33,7 +32,7 @@ pub fn resolve_azure_api_key(
|
|||
non_empty(api_key)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiKey)))
|
||||
.ok_or_else(|| Error::from(crate::AuthError::MissingAzureApiKey))
|
||||
}
|
||||
|
||||
pub fn complete_azure_anthropic_url(
|
||||
|
|
@ -43,7 +42,7 @@ pub fn complete_azure_anthropic_url(
|
|||
let api_base = non_empty(api_base)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiBase)))?;
|
||||
.ok_or_else(|| Error::from(crate::AuthError::MissingAzureApiBase))?;
|
||||
|
||||
let api_base = api_base.trim_end_matches('/');
|
||||
|
||||
|
|
|
|||
|
|
@ -1,2 +1 @@
|
|||
pub(crate) mod auth;
|
||||
pub mod messages;
|
||||
|
|
|
|||
|
|
@ -1,930 +1 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::time::Duration;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::caching::in_memory_cache::InMemoryCache;
|
||||
use crate::error::Error;
|
||||
use aws_credential_types::Credentials;
|
||||
use aws_credential_types::provider::ProvideCredentials;
|
||||
use aws_sigv4::http_request::{
|
||||
SignableBody, SignableRequest, SigningParams, SigningSettings, sign,
|
||||
};
|
||||
use aws_sigv4::sign::v4;
|
||||
use aws_smithy_runtime_api::client::identity::Identity;
|
||||
use serde_json::{Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::constants::{
|
||||
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME,
|
||||
AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
|
||||
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
|
||||
BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX,
|
||||
SIGV4_COMPUTED_HEADER_NAMES,
|
||||
};
|
||||
|
||||
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
|
||||
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
|
||||
|
||||
static IAM_CREDENTIALS_CACHE: OnceLock<Mutex<InMemoryCache<Credentials>>> = OnceLock::new();
|
||||
|
||||
fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option<Duration> {
|
||||
match flow {
|
||||
AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL),
|
||||
AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL),
|
||||
AwsAuthFlow::WebIdentity { .. }
|
||||
| AwsAuthFlow::AssumeRole { .. }
|
||||
| AwsAuthFlow::Profile { .. }
|
||||
| AwsAuthFlow::SessionToken { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct AwsAuthConfig {
|
||||
pub access_key_id: Option<String>,
|
||||
pub secret_access_key: Option<String>,
|
||||
pub session_token: Option<String>,
|
||||
pub region_name: Option<String>,
|
||||
pub session_name: Option<String>,
|
||||
pub profile_name: Option<String>,
|
||||
pub role_name: Option<String>,
|
||||
pub web_identity_token: Option<String>,
|
||||
pub sts_endpoint: Option<String>,
|
||||
pub external_id: Option<String>,
|
||||
}
|
||||
|
||||
impl AwsAuthConfig {
|
||||
fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option<String> + Sync)) -> Self {
|
||||
Self {
|
||||
access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)),
|
||||
secret_access_key: self
|
||||
.secret_access_key
|
||||
.or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)),
|
||||
session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)),
|
||||
region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)),
|
||||
session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)),
|
||||
profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)),
|
||||
role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)),
|
||||
web_identity_token: self
|
||||
.web_identity_token
|
||||
.or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)),
|
||||
sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)),
|
||||
external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum AwsAuthFlow {
|
||||
WebIdentity {
|
||||
token: String,
|
||||
role: String,
|
||||
session_name: String,
|
||||
},
|
||||
AssumeRole {
|
||||
role: String,
|
||||
session_name: Option<String>,
|
||||
},
|
||||
Profile {
|
||||
name: String,
|
||||
},
|
||||
SessionToken {
|
||||
access_key_id: String,
|
||||
secret_access_key: String,
|
||||
session_token: String,
|
||||
},
|
||||
StaticKeys {
|
||||
access_key_id: String,
|
||||
secret_access_key: String,
|
||||
region_name: String,
|
||||
},
|
||||
DefaultChain,
|
||||
}
|
||||
|
||||
fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(format!("{config:?}:{flow:?}"));
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn get_cached_credentials(key: &str) -> Option<Credentials> {
|
||||
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
|
||||
let mut entries = cache.lock().ok()?;
|
||||
entries.get_cache(key)
|
||||
}
|
||||
|
||||
fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) {
|
||||
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
|
||||
if let Ok(mut entries) = cache.lock() {
|
||||
entries.set_cache(key, credentials, Some(ttl));
|
||||
}
|
||||
}
|
||||
|
||||
fn role_identity(arn: &str) -> Option<(&str, &str, &str)> {
|
||||
let mut parts = arn.splitn(6, ':');
|
||||
let ("arn", partition, _, _, account, resource) = (
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
parts.next()?,
|
||||
) else {
|
||||
return None;
|
||||
};
|
||||
let role = if let Some(role) = resource.strip_prefix("role/") {
|
||||
role.rsplit('/').next()?
|
||||
} else {
|
||||
resource.strip_prefix("assumed-role/")?.split('/').next()?
|
||||
};
|
||||
Some((partition, account, role))
|
||||
}
|
||||
|
||||
fn same_role_arns(target: &str, caller: &str) -> bool {
|
||||
role_identity(target) == role_identity(caller)
|
||||
}
|
||||
|
||||
pub fn classify_auth(
|
||||
config: AwsAuthConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> AwsAuthFlow {
|
||||
let config = config.with_environment(env_lookup);
|
||||
if let (Some(token), Some(role), Some(session_name)) = (
|
||||
config.web_identity_token.clone(),
|
||||
config.role_name.clone(),
|
||||
config.session_name.clone(),
|
||||
) {
|
||||
return AwsAuthFlow::WebIdentity {
|
||||
token,
|
||||
role,
|
||||
session_name,
|
||||
};
|
||||
}
|
||||
if let Some(role) = config.role_name.clone() {
|
||||
return AwsAuthFlow::AssumeRole {
|
||||
role,
|
||||
session_name: config.session_name.clone(),
|
||||
};
|
||||
}
|
||||
if let Some(name) = config.profile_name {
|
||||
return AwsAuthFlow::Profile { name };
|
||||
}
|
||||
if let (Some(access_key_id), Some(secret_access_key), Some(session_token)) = (
|
||||
config.access_key_id.clone(),
|
||||
config.secret_access_key.clone(),
|
||||
config.session_token,
|
||||
) {
|
||||
return AwsAuthFlow::SessionToken {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
session_token,
|
||||
};
|
||||
}
|
||||
if let (Some(access_key_id), Some(secret_access_key), Some(region_name)) = (
|
||||
config.access_key_id,
|
||||
config.secret_access_key,
|
||||
config.region_name,
|
||||
) {
|
||||
return AwsAuthFlow::StaticKeys {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
region_name,
|
||||
};
|
||||
}
|
||||
AwsAuthFlow::DefaultChain
|
||||
}
|
||||
|
||||
pub async fn resolve_credentials(
|
||||
config: AwsAuthConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<Credentials, Error> {
|
||||
let resolved = config.clone().with_environment(env_lookup);
|
||||
let flow = classify_auth(config, env_lookup);
|
||||
match flow {
|
||||
AwsAuthFlow::SessionToken {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
session_token,
|
||||
} => Ok(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
Some(session_token),
|
||||
None,
|
||||
"litellm-static-session",
|
||||
)),
|
||||
AwsAuthFlow::StaticKeys {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
region_name,
|
||||
} => {
|
||||
let flow = AwsAuthFlow::StaticKeys {
|
||||
access_key_id: access_key_id.clone(),
|
||||
secret_access_key: secret_access_key.clone(),
|
||||
region_name,
|
||||
};
|
||||
let key = cache_key(&resolved, &flow);
|
||||
if let Some(credentials) = get_cached_credentials(&key) {
|
||||
return Ok(credentials);
|
||||
}
|
||||
let credentials = Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
None,
|
||||
None,
|
||||
"litellm-static",
|
||||
);
|
||||
set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL),
|
||||
);
|
||||
Ok(credentials)
|
||||
}
|
||||
AwsAuthFlow::Profile { name } => {
|
||||
let provider = aws_config::profile::ProfileFileCredentialsProvider::builder()
|
||||
.profile_name(name)
|
||||
.build();
|
||||
provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::Auth(format!("AWS profile credentials failed: {error}")))
|
||||
}
|
||||
AwsAuthFlow::AssumeRole { role, session_name } => {
|
||||
if is_already_running_as_role(&role, &resolved).await? {
|
||||
let ambient_flow = AwsAuthFlow::DefaultChain;
|
||||
let key = cache_key(&resolved, &ambient_flow);
|
||||
if let Some(credentials) = get_cached_credentials(&key) {
|
||||
return Ok(credentials);
|
||||
}
|
||||
let provider =
|
||||
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
|
||||
.build()
|
||||
.await;
|
||||
let credentials = provider.provide_credentials().await.map_err(|error| {
|
||||
Error::Auth(format!("AWS default credentials failed: {error}"))
|
||||
})?;
|
||||
set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL),
|
||||
);
|
||||
return Ok(credentials);
|
||||
}
|
||||
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
|
||||
if let Some(region) = resolved.region_name.clone() {
|
||||
loader = loader.region(aws_types::region::Region::new(region));
|
||||
}
|
||||
if let Some(endpoint) = resolved.sts_endpoint.clone() {
|
||||
loader = loader.endpoint_url(endpoint);
|
||||
}
|
||||
if let (Some(access_key_id), Some(secret_access_key)) =
|
||||
(resolved.access_key_id, resolved.secret_access_key)
|
||||
{
|
||||
loader = loader.credentials_provider(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
resolved.session_token,
|
||||
None,
|
||||
"litellm-role-source",
|
||||
));
|
||||
}
|
||||
let sdk_config = loader.load().await;
|
||||
let builder = aws_config::sts::AssumeRoleProvider::builder(role);
|
||||
let builder = match session_name {
|
||||
Some(name) => builder.session_name(name),
|
||||
None => builder.session_name(default_session_name()),
|
||||
};
|
||||
let builder = match resolved.external_id {
|
||||
Some(id) => builder.external_id(id),
|
||||
None => builder,
|
||||
};
|
||||
let provider = builder.configure(&sdk_config).build().await;
|
||||
provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::Auth(format!("AWS role credentials failed: {error}")))
|
||||
}
|
||||
AwsAuthFlow::WebIdentity {
|
||||
token,
|
||||
role,
|
||||
session_name,
|
||||
} => {
|
||||
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
|
||||
if let Some(region) = resolved.region_name {
|
||||
loader = loader.region(aws_types::region::Region::new(region));
|
||||
}
|
||||
if let Some(endpoint) = resolved.sts_endpoint {
|
||||
loader = loader.endpoint_url(endpoint);
|
||||
}
|
||||
let sdk_config = loader.load().await;
|
||||
let client = aws_sdk_sts::Client::new(&sdk_config);
|
||||
let response = client
|
||||
.assume_role_with_web_identity()
|
||||
.role_arn(role)
|
||||
.role_session_name(session_name)
|
||||
.web_identity_token(token)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| {
|
||||
Error::Auth(format!("AWS web identity credentials failed: {error}"))
|
||||
})?;
|
||||
let credentials = response.credentials().ok_or_else(|| {
|
||||
Error::Auth("AWS web identity response had no credentials".to_string())
|
||||
})?;
|
||||
let expiration = SystemTime::try_from(*credentials.expiration()).map_err(|error| {
|
||||
Error::Auth(format!("AWS web identity expiration was invalid: {error}"))
|
||||
})?;
|
||||
Ok(Credentials::new(
|
||||
credentials.access_key_id(),
|
||||
credentials.secret_access_key(),
|
||||
Some(credentials.session_token().to_string()),
|
||||
Some(expiration),
|
||||
"litellm-web-identity",
|
||||
))
|
||||
}
|
||||
AwsAuthFlow::DefaultChain => {
|
||||
let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain);
|
||||
if let Some(credentials) = get_cached_credentials(&key) {
|
||||
return Ok(credentials);
|
||||
}
|
||||
let provider =
|
||||
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
|
||||
.build()
|
||||
.await;
|
||||
let credentials = provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::Auth(format!("AWS default credentials failed: {error}")))?;
|
||||
set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
credential_cache_ttl(&AwsAuthFlow::DefaultChain).unwrap_or(AMBIENT_CREDENTIALS_TTL),
|
||||
);
|
||||
Ok(credentials)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result<bool, Error> {
|
||||
if role_identity(role).is_none() {
|
||||
return Ok(false);
|
||||
}
|
||||
if let (Ok(current_role), Ok(token_file)) = (
|
||||
std::env::var(AWS_ROLE_ARN),
|
||||
std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE),
|
||||
) && !token_file.is_empty()
|
||||
{
|
||||
return Ok(same_role_arns(role, ¤t_role));
|
||||
}
|
||||
|
||||
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
|
||||
if let Some(region) = config.region_name.clone() {
|
||||
loader = loader.region(aws_types::region::Region::new(region));
|
||||
}
|
||||
if let Some(endpoint) = config.sts_endpoint.clone() {
|
||||
loader = loader.endpoint_url(endpoint);
|
||||
}
|
||||
let sdk_config = loader.load().await;
|
||||
let response = match aws_sdk_sts::Client::new(&sdk_config)
|
||||
.get_caller_identity()
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(_) => return Ok(false),
|
||||
};
|
||||
Ok(response
|
||||
.arn()
|
||||
.is_some_and(|caller| same_role_arns(role, caller)))
|
||||
}
|
||||
|
||||
fn default_session_name() -> String {
|
||||
let seconds = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_or(0, |duration| duration.as_secs());
|
||||
format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}")
|
||||
}
|
||||
|
||||
/// The subset of `headers` SigV4 should cover.
|
||||
///
|
||||
/// Python signs only these and reattaches the rest afterwards, so a forwarded
|
||||
/// client header cannot change the canonical request and invalidate the
|
||||
/// signature. Signing everything instead makes the request 403 on a header the
|
||||
/// caller supplied, on a deployment that works on the Python path.
|
||||
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
|
||||
headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
let name = name.to_ascii_lowercase();
|
||||
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
|
||||
|| name.starts_with("x-amz-")
|
||||
|| name.starts_with("x-amzn-")
|
||||
})
|
||||
.map(|(name, value)| (name.clone(), value.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Whether the signer produces `name` itself.
|
||||
///
|
||||
/// Python's reattach loop skips these, so a caller-supplied copy never reaches
|
||||
/// the wire next to the computed one.
|
||||
pub fn is_sigv4_computed_header(name: &str) -> bool {
|
||||
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
|
||||
}
|
||||
|
||||
pub fn sign_bedrock_post(
|
||||
url: &str,
|
||||
body: &[u8],
|
||||
headers: &BTreeMap<String, String>,
|
||||
region: &str,
|
||||
credentials: &Credentials,
|
||||
signing_time: SystemTime,
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
let identity: Identity = credentials.clone().into();
|
||||
let params = v4::SigningParams::builder()
|
||||
.identity(&identity)
|
||||
.region(region)
|
||||
.name(BEDROCK_SERVICE)
|
||||
.time(signing_time)
|
||||
.settings(SigningSettings::default())
|
||||
.build()
|
||||
.map(SigningParams::from)
|
||||
.map_err(|error| Error::Auth(format!("AWS signing parameters failed: {error}")))?;
|
||||
let header_refs = headers
|
||||
.iter()
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()));
|
||||
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
|
||||
.map_err(|error| Error::Auth(format!("AWS signable request failed: {error}")))?;
|
||||
let (instructions, _) = sign(request, ¶ms)
|
||||
.map_err(|error| Error::Auth(format!("AWS request signing failed: {error}")))?
|
||||
.into_parts();
|
||||
Ok(instructions
|
||||
.headers()
|
||||
.map(|(name, value)| {
|
||||
let normalized_name = match name {
|
||||
"authorization" => "Authorization",
|
||||
"x-amz-date" => "X-Amz-Date",
|
||||
"x-amz-security-token" => "X-Amz-Security-Token",
|
||||
_ => name,
|
||||
};
|
||||
(normalized_name.to_string(), value.to_string())
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Model-id and region parsing shared by every Bedrock route.
|
||||
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
|
||||
let mut stripped = model;
|
||||
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let mut region = None;
|
||||
if let Some((candidate, remainder)) = stripped.split_once('/')
|
||||
&& is_bedrock_region(candidate)
|
||||
{
|
||||
region = Some(candidate.to_string());
|
||||
stripped = remainder;
|
||||
}
|
||||
for prefix in ["nova-2/", "nova/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if region.is_none() {
|
||||
// Python splits the whole ARN and takes field 3, the region. Stripping
|
||||
// `arn:` first shifts every field down one, so the region is field 2
|
||||
// here; field 3 is the account id.
|
||||
region = stripped
|
||||
.strip_prefix("arn:")
|
||||
.and_then(|value| value.split(':').nth(2))
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
}
|
||||
(stripped.to_string(), region)
|
||||
}
|
||||
|
||||
fn is_bedrock_region(value: &str) -> bool {
|
||||
value.len() > 3
|
||||
&& value.contains('-')
|
||||
&& value
|
||||
.chars()
|
||||
.all(|char| char.is_ascii_alphanumeric() || char == '-')
|
||||
}
|
||||
|
||||
pub fn resolve_bedrock_region(
|
||||
model_region: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
if let Some(region) = optional_params
|
||||
.get("aws_region_name")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
return region.to_string();
|
||||
}
|
||||
if let Some(region) = model_region {
|
||||
return region.to_string();
|
||||
}
|
||||
env_lookup(AWS_REGION_NAME)
|
||||
.or_else(|| env_lookup(AWS_REGION))
|
||||
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
|
||||
}
|
||||
|
||||
pub fn aws_auth_config(
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> AwsAuthConfig {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
};
|
||||
let env = |key: &str| env_lookup(key);
|
||||
AwsAuthConfig {
|
||||
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
|
||||
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
|
||||
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
|
||||
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
|
||||
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
|
||||
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
|
||||
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
|
||||
web_identity_token: value("aws_web_identity_token")
|
||||
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
|
||||
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
|
||||
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Credentials a host resolved through its own chain and handed down verbatim.
|
||||
///
|
||||
/// A host with its own resolution (LiteLLM's Python `BaseAWSLLM`, which reads
|
||||
/// profiles, STS and boto sessions) passes the result here so the core signs
|
||||
/// with exactly those. Without this the core would re-derive from ambient
|
||||
/// state, where an unrelated `AWS_ROLE_NAME` or `AWS_PROFILE_NAME` in the
|
||||
/// environment outranks explicit keys in [`classify_auth`] and the two sides
|
||||
/// would sign as different principals.
|
||||
pub fn host_supplied_credentials(optional_params: &Map<String, Value>) -> Option<Credentials> {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
};
|
||||
let access_key_id = value("aws_access_key_id")?;
|
||||
let secret_access_key = value("aws_secret_access_key")?;
|
||||
Some(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
value("aws_session_token").map(str::to_string),
|
||||
None,
|
||||
"litellm-host-supplied",
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn no_env(_: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
|
||||
(
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
|
||||
.to_string(),
|
||||
br#"{"input":"hello"}"#.to_vec(),
|
||||
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_the_region_field_of_a_model_arn_not_the_account_id() {
|
||||
// Python's `_get_aws_region_from_model_arn` splits the whole ARN and
|
||||
// takes field 3. Stripping `arn:` first shifts every field down one, so
|
||||
// the region is field 2 here. Taking field 3 after the strip returns
|
||||
// the account id, which is not a region at all.
|
||||
let (_, region) = bedrock_model_id_and_region(
|
||||
"bedrock/arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2",
|
||||
);
|
||||
assert_eq!(region.as_deref(), Some("us-west-2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classification_preserves_python_precedence() {
|
||||
let config = AwsAuthConfig {
|
||||
access_key_id: Some("ak".into()),
|
||||
secret_access_key: Some("sk".into()),
|
||||
session_token: Some("token".into()),
|
||||
region_name: Some("us-east-1".into()),
|
||||
session_name: Some("session".into()),
|
||||
profile_name: Some("profile".into()),
|
||||
role_name: Some("role".into()),
|
||||
web_identity_token: Some("oidc".into()),
|
||||
..Default::default()
|
||||
};
|
||||
assert!(matches!(
|
||||
classify_auth(config, &no_env),
|
||||
AwsAuthFlow::WebIdentity { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classification_covers_fallthroughs() {
|
||||
let env = |key: &str| match key {
|
||||
AWS_PROFILE_NAME => Some("profile".into()),
|
||||
_ => None,
|
||||
};
|
||||
assert!(matches!(
|
||||
classify_auth(AwsAuthConfig::default(), &env),
|
||||
AwsAuthFlow::Profile { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
classify_auth(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("ak".into()),
|
||||
secret_access_key: Some("sk".into()),
|
||||
session_token: Some("token".into()),
|
||||
..Default::default()
|
||||
},
|
||||
&no_env
|
||||
),
|
||||
AwsAuthFlow::SessionToken { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
classify_auth(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("ak".into()),
|
||||
secret_access_key: Some("sk".into()),
|
||||
region_name: Some("us-east-1".into()),
|
||||
..Default::default()
|
||||
},
|
||||
&no_env
|
||||
),
|
||||
AwsAuthFlow::StaticKeys { .. }
|
||||
));
|
||||
assert_eq!(
|
||||
classify_auth(AwsAuthConfig::default(), &no_env),
|
||||
AwsAuthFlow::DefaultChain
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn static_credentials_do_not_use_network() {
|
||||
let credentials = resolve_credentials(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("ak".into()),
|
||||
secret_access_key: Some("sk".into()),
|
||||
region_name: Some("us-east-1".into()),
|
||||
..Default::default()
|
||||
},
|
||||
&no_env,
|
||||
)
|
||||
.await
|
||||
.expect("static credentials");
|
||||
assert_eq!(credentials.access_key_id(), "ak");
|
||||
assert_eq!(credentials.session_token(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_policy_matches_python_flows() {
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::StaticKeys {
|
||||
access_key_id: "ak".into(),
|
||||
secret_access_key: "sk".into(),
|
||||
region_name: "us-east-1".into(),
|
||||
}),
|
||||
Some(STATIC_CREDENTIALS_TTL)
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::DefaultChain),
|
||||
Some(AMBIENT_CREDENTIALS_TTL)
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::SessionToken {
|
||||
access_key_id: "ak".into(),
|
||||
secret_access_key: "sk".into(),
|
||||
session_token: "token".into(),
|
||||
}),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::Profile {
|
||||
name: "profile".into()
|
||||
}),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::AssumeRole {
|
||||
role: "arn:aws:iam::123456789012:role/demo".into(),
|
||||
session_name: None,
|
||||
}),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
credential_cache_ttl(&AwsAuthFlow::WebIdentity {
|
||||
token: "token".into(),
|
||||
role: "arn:aws:iam::123456789012:role/demo".into(),
|
||||
session_name: "session".into(),
|
||||
}),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_round_trip_preserves_credentials() {
|
||||
let key = format!("cache-test-{}", std::process::id());
|
||||
let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test");
|
||||
set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL);
|
||||
assert_eq!(
|
||||
get_cached_credentials(&key).map(|value| value.access_key_id().to_string()),
|
||||
Some("cache-ak".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_role_comparison_matches_partition_account_and_role() {
|
||||
assert!(same_role_arns(
|
||||
"arn:aws:iam::123456789012:role/path/demo",
|
||||
"arn:aws:sts::123456789012:assumed-role/demo/session"
|
||||
));
|
||||
assert!(!same_role_arns(
|
||||
"arn:aws:iam::123456789012:role/demo",
|
||||
"arn:aws:iam::999999999999:role/demo"
|
||||
));
|
||||
assert!(!same_role_arns(
|
||||
"arn:aws:iam::123456789012:role/demo",
|
||||
"arn:aws-cn:iam::123456789012:role/demo"
|
||||
));
|
||||
assert!(!same_role_arns(
|
||||
"arn:aws:iam::123456789012:user/demo",
|
||||
"arn:aws:iam::123456789012:role/demo"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_client_header_is_not_folded_into_the_signature() {
|
||||
// Python signs only the AWS header set, so a header a caller forwarded
|
||||
// cannot change the canonical request. Signing it instead makes the
|
||||
// request 403 the moment anything on the wire rewrites or drops it.
|
||||
let (url, body, mut headers) = parity_inputs();
|
||||
headers.insert("x-request-id".to_string(), "abc-123".to_string());
|
||||
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
|
||||
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
|
||||
let signable = aws_signature_headers(&headers);
|
||||
|
||||
assert!(!signable.contains_key("x-request-id"));
|
||||
assert!(!signable.contains_key("Accept-Encoding"));
|
||||
// The AWS-prefixed one is genuinely part of the signature.
|
||||
assert!(signable.contains_key("x-amzn-trace-id"));
|
||||
assert!(signable.contains_key("Content-Type"));
|
||||
|
||||
let credentials = Credentials::new(
|
||||
"AKIDEXAMPLE",
|
||||
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
|
||||
None,
|
||||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&signable,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
SystemTime::UNIX_EPOCH,
|
||||
)
|
||||
.expect("signs");
|
||||
let authorization = signed
|
||||
.get("Authorization")
|
||||
.expect("carries an authorization header");
|
||||
assert!(
|
||||
!authorization.contains("x-request-id"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
assert!(
|
||||
!authorization.contains("accept-encoding"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signing_matches_botocore_golden_vector() {
|
||||
let (url, body, headers) = parity_inputs();
|
||||
let credentials = Credentials::new(
|
||||
"AKIDEXAMPLE",
|
||||
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
|
||||
Some("session-token".to_string()),
|
||||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
|
||||
)
|
||||
.expect("golden signature");
|
||||
assert_eq!(
|
||||
signed.get("X-Amz-Date").map(String::as_str),
|
||||
Some("20240102T030405Z")
|
||||
);
|
||||
assert_eq!(
|
||||
signed.get("X-Amz-Security-Token").map(String::as_str),
|
||||
Some("session-token")
|
||||
);
|
||||
assert_eq!(
|
||||
signed.get("Authorization").map(String::as_str),
|
||||
Some(
|
||||
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signing_without_session_token_omits_security_header() {
|
||||
let (url, body, headers) = parity_inputs();
|
||||
let credentials = Credentials::new(
|
||||
"AKIDEXAMPLE",
|
||||
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
|
||||
None,
|
||||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
|
||||
)
|
||||
.expect("signature");
|
||||
assert!(!signed.contains_key("X-Amz-Security-Token"));
|
||||
}
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let access_key_id = std::env::var("AWS_BEDROCK_TEST_ACCESS_KEY_ID")?;
|
||||
let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?;
|
||||
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec();
|
||||
let headers =
|
||||
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]);
|
||||
let credentials = resolve_credentials(
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some(access_key_id),
|
||||
secret_access_key: Some(secret_access_key),
|
||||
region_name: Some("us-west-2".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
&no_env,
|
||||
)
|
||||
.await?;
|
||||
let client = reqwest::Client::new();
|
||||
let mut failures = Vec::new();
|
||||
|
||||
for region in ["us-west-2", "us-east-1"] {
|
||||
let url = format!(
|
||||
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
|
||||
);
|
||||
let signed_headers = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
region,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?;
|
||||
let mut request = client.post(&url).body(body.clone());
|
||||
for (name, value) in &headers {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
for (name, value) in signed_headers {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
let response = request.send().await?;
|
||||
let status = response.status();
|
||||
let response_body = response.text().await?;
|
||||
let snippet: String = response_body.chars().take(240).collect();
|
||||
println!("region={region} status={status} response={snippet}");
|
||||
if status == reqwest::StatusCode::OK {
|
||||
return Ok(());
|
||||
}
|
||||
failures.push(format!("{region}: {status} {snippet}"));
|
||||
}
|
||||
|
||||
panic!(
|
||||
"no Bedrock region returned HTTP 200: {}",
|
||||
failures.join("; ")
|
||||
);
|
||||
}
|
||||
}
|
||||
pub use litellm_auth_aws::*;
|
||||
|
|
|
|||
|
|
@ -1,43 +1 @@
|
|||
pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID";
|
||||
pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY";
|
||||
pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN";
|
||||
pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME";
|
||||
pub const AWS_REGION: &str = "AWS_REGION";
|
||||
pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME";
|
||||
pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME";
|
||||
pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME";
|
||||
pub const AWS_WEB_IDENTITY_TOKEN: &str = "AWS_WEB_IDENTITY_TOKEN";
|
||||
pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
|
||||
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
|
||||
pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT";
|
||||
pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID";
|
||||
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";
|
||||
|
||||
/// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors
|
||||
/// Python's `_filter_headers_for_aws_signature` allowlist.
|
||||
pub const AWS_SIGNED_HEADER_NAMES: &[&str] = &[
|
||||
"host",
|
||||
"content-type",
|
||||
"date",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"x-amz-content-sha256",
|
||||
"x-amz-algorithm",
|
||||
"x-amz-credential",
|
||||
"x-amz-signedheaders",
|
||||
"x-amz-signature",
|
||||
];
|
||||
/// Headers the signer emits itself. Mirrors Python's `SIGV4_COMPUTED_HEADERS`,
|
||||
/// which the reattach loop skips so a caller's copy cannot ride alongside the
|
||||
/// computed one.
|
||||
pub const SIGV4_COMPUTED_HEADER_NAMES: &[&str] = &[
|
||||
"authorization",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"date",
|
||||
];
|
||||
pub const BEDROCK_SERVICE: &str = "bedrock";
|
||||
pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session";
|
||||
pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2";
|
||||
pub const BEDROCK_RUNTIME_ENDPOINT_TEMPLATE: &str =
|
||||
"https://bedrock-runtime.{region}.amazonaws.com";
|
||||
pub use litellm_auth_aws::constants::*;
|
||||
|
|
|
|||
|
|
@ -122,7 +122,7 @@ async fn configs_build_complete_requests_and_share_mistral_normalization() {
|
|||
.prepare_request(&direct, &client)
|
||||
.await
|
||||
.unwrap();
|
||||
let vertex_http = VertexAIOCRConfig::default()
|
||||
let vertex_http = VertexAIOCRConfig
|
||||
.prepare_request(&vertex, &client)
|
||||
.await
|
||||
.unwrap();
|
||||
|
|
@ -153,7 +153,7 @@ async fn configs_build_complete_requests_and_share_mistral_normalization() {
|
|||
.transform_ocr_response(&direct, serde_json::from_value(payload.clone()).unwrap())
|
||||
.unwrap()
|
||||
.into_json();
|
||||
let vertex_response = VertexAIOCRConfig::default()
|
||||
let vertex_response = VertexAIOCRConfig
|
||||
.transform_ocr_response(&vertex, serde_json::from_value(payload).unwrap())
|
||||
.unwrap()
|
||||
.into_json();
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue