scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge pull request #41464 from BerriAI/litellm_rust_extract_auth_cache_crates

Browse source 
refactor(rust): extract auth and cache crates
This commit is contained in:
yujonglee 2026-09-16 11:45:04 -07:00 • committed by GitHub
commit 560200da21
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
129 changed files with 3108 additions and 2152 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.github/workflows/test-rust.yml vendored
View file

@ -95,8 +95,6 @@ jobs:
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
- run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
rust-test:
runs-on: ubuntu-latest
timeout-minutes: 30
@ -129,9 +127,6 @@ jobs:
- run: cargo test --workspace --locked
working-directory: litellm-rust
- run: cargo test -p litellm-core --features bedrock-auth --locked
working-directory: litellm-rust
- run: uv build --wheel --out-dir dist
- run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl

78
litellm-rust/Cargo.lock generated
View file

@ -1838,7 +1838,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "litellm-core"
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 +1858,75 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"litellm-auth",
"moka",
"reqwest 0.12.28",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"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",
]
[[package]]
name = "litellm-cache"
version = "0.1.0"
dependencies = [
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-cache-memory"
version = "0.1.0"
dependencies = [
"litellm-cache",
"rstest",
"serde_json",
"tokio",
]
[[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",
@ -1880,6 +1953,7 @@ version = "0.1.0"
dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-core",
"litellm-python-interop",
"litellm-token-counter",

9
litellm-rust/Cargo.toml
View file

@ -11,6 +11,12 @@ 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-cache = { path = "crates/cache" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-python-interop = { path = "crates/python-interop" }
pyo3 = "0.29.2"
@ -30,9 +36,6 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
gcp_auth = "0.12.7"
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"

25
litellm-rust/crates/auth-aws/Cargo.toml Normal file
View file

@ -0,0 +1,25 @@
[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
thiserror.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

949
litellm-rust/crates/auth-aws/src/aws.rs Normal file
View file

@ -0,0 +1,949 @@
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 super::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, &current_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, &params)
.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
View 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";

46
litellm-rust/crates/auth-aws/src/error.rs Normal file
View file

@ -0,0 +1,46 @@
use thiserror::Error as ThisError;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum Error {
#[error("AWS profile credentials failed: {0}")]
AwsProfile(String),
#[error("AWS default credentials failed: {0}")]
AwsDefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AwsAssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
AwsWebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
AwsWebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
AwsSigningParameters(String),
#[error("AWS signable request failed: {0}")]
AwsSignableRequest(String),
#[error("AWS request signing failed: {0}")]
AwsSigning(String),
#[error("AWS web identity response had no credentials")]
AwsMissingWebIdentityCredentials,
}
impl From<Error> for litellm_auth::Error {
fn from(error: Error) -> Self {
Self::ProviderAuthentication(error.to_string())
}
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn converts_to_shared_auth_error_without_losing_context() {
let error = litellm_auth::Error::from(Error::AwsProfile("profile not found".into()));
assert_eq!(
error,
litellm_auth::Error::ProviderAuthentication(
"AWS profile credentials failed: profile not found".into()
)
);
}
}

6
litellm-rust/crates/auth-aws/src/lib.rs Normal file
View file

@ -0,0 +1,6 @@
mod aws;
pub mod constants;
mod error;
pub use aws::*;
pub use error::Error;

21
litellm-rust/crates/auth-azure/Cargo.toml Normal file
View file

@ -0,0 +1,21 @@
[package]
name = "litellm-auth-azure"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
moka.workspace = true
serde_json.workspace = true
sha2.workspace = true
strum.workspace = true
url.workspace = true
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
[dev-dependencies]
tokio.workspace = true

6
litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs → litellm-rust/crates/auth-azure/src/credential_provider_cache.rs
View file

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

@ -0,0 +1,7 @@
mod credential_provider_cache;
mod native;
mod resolve;
mod types;
pub use resolve::AzureAuthService;
pub use types::AzureAuthInputs;

51
litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs → litellm-rust/crates/auth-azure/src/native.rs
View file

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

93
litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs → litellm-rust/crates/auth-azure/src/resolve.rs
View file

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

19
litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs → litellm-rust/crates/auth-azure/src/types.rs
View file

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

16
litellm-rust/crates/auth-gcp/Cargo.toml Normal file
View file

@ -0,0 +1,16 @@
[package]
name = "litellm-auth-gcp"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
moka.workspace = true
serde_json.workspace = true
sha2.workspace = true
tokio.workspace = true
gcp_auth = "0.12.7"

71
litellm-rust/crates/core/src/auth/vertex.rs → litellm-rust/crates/auth-gcp/src/lib.rs
View file

@ -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,13 @@ impl VertexAuth {
}
}
pub(crate) async fn validate_environment(
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 +159,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 +169,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 +188,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 +248,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 +258,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 +320,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 +335,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 +349,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 +369,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 +524,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
View 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

18
litellm-rust/crates/core/src/auth/credential.rs → litellm-rust/crates/auth/src/credential.rs
View file

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

120
litellm-rust/crates/auth/src/error.rs Normal file
View file

@ -0,0 +1,120 @@
use thiserror::Error as ThisError;
#[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("{0}")]
ProviderAuthentication(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
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")]
EmptyAzureToken,
#[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")]
UnresolvedOidcReference,
#[error(
"Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable"
)]
MissingApiKey {
provider: &'static str,
environment_variable: &'static str,
},
#[error(
"Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter"
)]
MissingApiBase {
provider: &'static str,
environment_variable: &'static str,
},
#[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"
)]
MissingAzureApiBase,
#[error("invalid authentication header")]
InvalidHeader,
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn missing_api_key_names_provider_and_environment_variable() {
assert_eq!(
Error::MissingApiKey {
provider: "Anthropic",
environment_variable: "ANTHROPIC_API_KEY",
}
.to_string(),
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
);
}
}

15
litellm-rust/crates/core/src/auth/http.rs → litellm-rust/crates/auth/src/http.rs
View file

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

7
litellm-rust/crates/core/src/auth/mod.rs → litellm-rust/crates/auth/src/lib.rs
View file

@ -1,8 +1,6 @@
mod credential;
pub mod error;
pub(crate) mod vertex;
pub use error::AuthError;
pub(crate) mod http;
mod error;
pub mod http;
mod policy;
mod secret;
mod token;
@ -51,6 +49,7 @@ pub use credential::{
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;

23
litellm-rust/crates/core/src/auth/policy.rs → litellm-rust/crates/auth/src/policy.rs
View file

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

0
litellm-rust/crates/core/src/auth/secret.rs → litellm-rust/crates/auth/src/secret.rs
View file

6
litellm-rust/crates/core/src/auth/token.rs → litellm-rust/crates/auth/src/token.rs
View file

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

14
litellm-rust/crates/cache-memory/Cargo.toml Normal file
View file

@ -0,0 +1,14 @@
[package]
name = "litellm-cache-memory"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
serde_json.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

254
litellm-rust/crates/cache-memory/src/cache.rs Normal file
View file

@ -0,0 +1,254 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
type ValueMeasure<V> = Arc<dyn Fn(&V) -> Result<usize, Error> + Send + Sync>;
type ValueValidator<V> = Arc<dyn Fn(&V) -> Result<(), Error> + Send + Sync>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CacheWrite {
Stored,
Disabled,
TooLarge,
}
struct CacheState<V> {
values: HashMap<String, V>,
expirations: HashMap<String, Duration>,
expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
}
pub struct InMemoryCache<V: Clone> {
state: Mutex<CacheState<V>>,
max_size_in_memory: usize,
default_ttl: Duration,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
validate_value: Option<ValueValidator<V>>,
now: Arc<dyn Fn() -> Duration + Send + Sync>,
}
impl<V: Clone> Default for InMemoryCache<V> {
fn default() -> Self {
Self::new(None, None)
}
}
impl<V: Clone> InMemoryCache<V> {
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
Self::with_clock(max_size_in_memory, default_ttl, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn with_clock(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self::with_clock_and_size_measurement(max_size_in_memory, default_ttl, None, None, now)
}
pub fn with_clock_and_size_measurement(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self {
state: Mutex::new(CacheState {
values: HashMap::new(),
expirations: HashMap::new(),
expiration_heap: BinaryHeap::new(),
}),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
max_entry_bytes,
measure_value,
validate_value: None,
now: Arc::new(now),
}
}
pub fn set_cache(
&self,
key: impl Into<String>,
value: V,
ttl: Option<Duration>,
) -> Result<CacheWrite, Error> {
if self.max_size_in_memory == 0 {
return Ok(CacheWrite::Disabled);
}
if let Some(validate) = &self.validate_value {
validate(&value)?;
}
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&value)? > limit
{
return Ok(CacheWrite::TooLarge);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now);
let key = key.into();
state.values.insert(key.clone(), value);
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
let expiration = now + ttl.unwrap_or(self.default_ttl);
state.expirations.insert(key.clone(), expiration);
state.expiration_heap.push(Reverse((expiration, key)));
}
Ok(CacheWrite::Stored)
}
pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
if state
.expirations
.get(key)
.is_some_and(|expiration| *expiration < now)
{
Self::remove(&mut state, key);
}
Ok(state.values.get(key).cloned())
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expirations
.get(key)
.copied())
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::remove(&mut state, key);
Ok(())
}
pub fn flush_cache(&self) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
state.values.clear();
state.expirations.clear();
state.expiration_heap.clear();
Ok(())
}
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration) {
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
if state.expirations.get(&key).copied() != Some(expiration) {
state.expiration_heap.pop();
} else if expiration <= now {
state.expiration_heap.pop();
Self::remove(state, &key);
} else {
break;
}
}
while state.values.len() >= capacity {
let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else {
break;
};
if state.expirations.get(&key).copied() == Some(expiration) {
Self::remove(state, &key);
}
}
}
fn remove(state: &mut CacheState<V>, key: &str) {
state.values.remove(key);
state.expirations.remove(key);
}
}
impl InMemoryCache<CacheEntry> {
pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn response_cache_with_clock(
capacity: usize,
ttl: Duration,
max_entry_bytes: usize,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
let mut cache = Self::with_clock_and_size_measurement(
Some(capacity),
Some(ttl),
Some(max_entry_bytes),
Some(Arc::new(|entry: &CacheEntry| {
serde_json::to_vec(entry)
.map(|bytes| bytes.len())
.map_err(|_| Error::InvalidEntry)
})),
now,
);
cache.validate_value = Some(Arc::new(|entry: &CacheEntry| {
entry
.timestamp
.is_finite()
.then_some(())
.ok_or(Error::InvalidEntry)
}));
cache
}
}
impl BaseCache for InMemoryCache<CacheEntry> {
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let ttl = self.get_ttl(&kwargs);
self.set_cache(key, value, Some(ttl)).map(|_| ())
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
self.get_cache(key)
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.delete_cache(key)
}
fn flush_cache(&self) -> Result<(), Error> {
self.flush_cache()
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "In-memory cache connection test successful".into(),
error: None,
})
})
}
}

3
litellm-rust/crates/cache-memory/src/lib.rs Normal file
View file

@ -0,0 +1,3 @@
mod cache;
pub use cache::{CacheWrite, InMemoryCache};

158
litellm-rust/crates/cache-memory/tests/cache.rs Normal file
View file

@ -0,0 +1,158 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
#[fixture]
fn clock() -> Arc<AtomicU64> {
Arc::new(AtomicU64::new(100))
}
fn cache(clock: Arc<AtomicU64>, capacity: usize) -> InMemoryCache<String> {
InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || {
Duration::from_secs(clock.load(Ordering::SeqCst))
})
}
#[rstest]
fn default_explicit_and_override_ttls_follow_python_rules(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
cache.set_cache("key", "first".into(), None).unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
clock.store(160, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
clock.store(161, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), None);
cache
.set_cache("key", "third".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(171))
);
}
#[rstest]
fn write_at_expiry_boundary_refreshes_ttl(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
cache
.set_cache("key", "first".into(), Some(Duration::from_secs(10)))
.unwrap();
clock.store(110, Ordering::SeqCst);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(120))
);
clock.store(115, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
}
#[rstest]
fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("early", "a".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("late", "b".into(), Some(Duration::from_secs(20)))
.unwrap();
cache.delete_cache("early").unwrap();
cache
.set_cache("new", "c".into(), Some(Duration::from_secs(30)))
.unwrap();
assert_eq!(cache.get_cache("late").unwrap(), Some("b".into()));
cache
.set_cache("last", "d".into(), Some(Duration::from_secs(40)))
.unwrap();
assert_eq!(cache.get_cache("late").unwrap(), None);
}
#[test]
fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
let disabled = InMemoryCache::<CacheEntry>::response_cache(0, Duration::from_secs(60), 80);
assert_eq!(
disabled
.set_cache(
"a",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("x")
},
None
)
.unwrap(),
CacheWrite::Disabled
);
let cache = InMemoryCache::<CacheEntry>::response_cache(2, Duration::from_secs(60), 80);
assert_eq!(
cache
.set_cache(
"large",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("x".repeat(100))
},
None
)
.unwrap(),
CacheWrite::TooLarge
);
cache
.set_cache(
"small",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("ok"),
},
None,
)
.unwrap();
assert!(cache.get_cache("small").unwrap().is_some());
assert_eq!(
cache
.set_cache(
"invalid",
CacheEntry {
timestamp: f64::NAN,
response: serde_json::json!("bad"),
},
None,
)
.unwrap_err(),
Error::InvalidEntry
);
cache.delete_cache("small").unwrap();
cache.flush_cache().unwrap();
}
#[tokio::test]
async fn connection_test_matches_python_result_contract() {
let cache = InMemoryCache::<CacheEntry>::default();
let result = BaseCache::test_connection(&cache).await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "In-memory cache connection test successful");
assert_eq!(result.error, None);
assert_eq!(
serde_json::to_value(result).unwrap(),
serde_json::json!({
"status": "success",
"message": "In-memory cache connection test successful"
})
);
}

15
litellm-rust/crates/cache/Cargo.toml vendored Normal file
View file

@ -0,0 +1,15 @@
[package]
name = "litellm-cache"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true

98
litellm-rust/crates/cache/src/base_cache.rs vendored Normal file
View file

@ -0,0 +1,98 @@
use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::Error;
pub type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
#[derive(Clone, Debug, Default, PartialEq)]
pub struct CacheKwargs {
pub ttl: Option<Duration>,
pub extras: Map<String, Value>,
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum CacheConnectionStatus {
Success,
Failed,
}
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub struct CacheConnectionResult {
pub status: CacheConnectionStatus,
pub message: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
pub trait BaseCache: Send + Sync {
type Value: Clone + Send + Sync + 'static;
fn default_ttl(&self) -> Duration {
Duration::from_secs(60)
}
fn get_ttl(&self, kwargs: &CacheKwargs) -> Duration {
kwargs.ttl.unwrap_or_else(|| self.default_ttl())
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error>;
fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<Self::Value>, Error>;
fn async_set_cache<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
Box::pin(async move { self.set_cache(key, value, kwargs) })
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
kwargs: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
Box::pin(async move { self.get_cache(key, kwargs) })
}
fn async_set_cache_pipeline<'a>(
&'a self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
Box::pin(async move {
for (key, value) in cache_list {
self.set_cache(&key, value, kwargs.clone())?;
}
Ok(())
})
}
fn batch_cache_write<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
self.async_set_cache(key, value, kwargs)
}
fn delete_cache(&self, key: &str) -> Result<(), Error>;
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
Box::pin(async move { self.delete_cache(key) })
}
fn flush_cache(&self) -> Result<(), Error>;
fn disconnect(&self) -> CacheFuture<'_, ()>;
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult>;
}

166
litellm-rust/crates/cache/src/caching.rs vendored Normal file
View file

@ -0,0 +1,166 @@
use std::sync::Arc;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::{BaseCache, CacheKwargs, Error};
pub use crate::BaseCache as Cache;
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
pub enum CacheMode {
#[default]
#[serde(rename = "default_on")]
DefaultOn,
#[serde(rename = "default_off")]
DefaultOff,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct CacheKeyField {
pub name: String,
pub value: Option<String>,
pub api_parameter: bool,
pub internal_parameter: bool,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
pub struct CacheKeyInput {
pub fields: Vec<CacheKeyField>,
pub preset: Option<String>,
pub namespace: Option<String>,
pub include_provider_parameters: bool,
}
#[derive(Default)]
pub struct CacheKeyContext {
pub model_group: Option<String>,
pub caching_groups: Vec<(Vec<String>, String)>,
pub file_checksum: Option<String>,
pub file_object_name: Option<String>,
pub metadata_file_name: Option<String>,
pub parameters_file_name: Option<String>,
}
impl CacheKeyContext {
pub fn apply(self, input: &mut CacheKeyInput) {
let group = self.model_group.as_ref().and_then(|model| {
self.caching_groups
.iter()
.find(|(models, _)| models.contains(model))
});
for field in &mut input.fields {
match field.name.as_str() {
"model" => {
field.value = group
.map(|(_, formatted)| formatted.clone())
.or_else(|| self.model_group.clone())
.or_else(|| field.value.take())
}
"file" => {
field.value = self
.file_checksum
.clone()
.or_else(|| self.file_object_name.clone())
.or_else(|| self.metadata_file_name.clone())
.or_else(|| self.parameters_file_name.clone())
}
_ => {}
}
}
}
}
pub fn get_cache_key(input: &CacheKeyInput) -> String {
cache_key(input)
}
pub fn cache_key(input: &CacheKeyInput) -> String {
if let Some(preset) = &input.preset {
return preset.clone();
}
let mut digest = Sha256::new();
for field in &input.fields {
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
&& let Some(value) = &field.value
{
digest.update(field.name.as_bytes());
digest.update(b": ");
digest.update(value.as_bytes());
}
}
let hash = format!("{:x}", digest.finalize());
input
.namespace
.as_deref()
.filter(|namespace| !namespace.is_empty())
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
pub struct CacheControls {
pub supported_call_type: bool,
pub configured: bool,
pub native_backend: bool,
pub default_on: bool,
pub caching: Option<bool>,
pub no_cache: bool,
pub no_store: bool,
#[serde(default)]
pub use_cache: bool,
}
impl CacheControls {
pub fn reads(self) -> bool {
self.supported_call_type
&& self.configured
&& self.caching.unwrap_or(true)
&& !self.no_cache
&& (self.default_on || self.use_cache)
}
pub fn writes(self) -> bool {
self.supported_call_type
&& self.configured
&& !self.no_store
&& (self.default_on || self.use_cache)
}
}
pub fn should_use_cache(controls: CacheControls) -> bool {
controls.reads() || controls.writes()
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CacheEntry {
pub timestamp: f64,
pub response: Value,
}
impl CacheEntry {
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
self.timestamp.is_finite()
&& max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64())
}
}
pub fn get_cache(
cache: &dyn BaseCache<Value = CacheEntry>,
key: &str,
kwargs: &CacheKwargs,
) -> Result<Option<CacheEntry>, Error> {
cache.get_cache(key, kwargs)
}
pub fn set_cache(
cache: &dyn BaseCache<Value = CacheEntry>,
key: &str,
entry: CacheEntry,
kwargs: CacheKwargs,
) -> Result<(), Error> {
cache.set_cache(key, entry, kwargs)
}
pub type CacheBackend = Arc<dyn BaseCache<Value = CacheEntry>>;

7
litellm-rust/crates/cache/src/error.rs vendored Normal file
View file

@ -0,0 +1,7 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("cache is unavailable")]
Unavailable,
#[error("invalid cache entry")]
InvalidEntry,
}

12
litellm-rust/crates/cache/src/lib.rs vendored Normal file
View file

@ -0,0 +1,12 @@
mod base_cache;
mod caching;
mod error;
pub use base_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs,
};
pub use caching::{
Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput,
CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache,
};
pub use error::Error;

139
litellm-rust/crates/cache/tests/caching.rs vendored Normal file
View file

@ -0,0 +1,139 @@
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext,
CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key,
};
use sha2::{Digest, Sha256};
use std::time::Duration;
struct TestCache {
default_ttl: Duration,
}
impl BaseCache for TestCache {
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> {
Ok(())
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
Ok(None)
}
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Ok(())
}
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
unreachable!()
}
}
#[test]
fn ttl_uses_default_and_allows_per_call_override() {
let cache = TestCache {
default_ttl: Duration::from_secs(60),
};
assert_eq!(
cache.get_ttl(&CacheKwargs::default()),
Duration::from_secs(60)
);
assert_eq!(
cache.get_ttl(&CacheKwargs {
ttl: Some(Duration::from_secs(5)),
..Default::default()
}),
Duration::from_secs(5)
);
}
#[test]
fn keys_match_python_order_groups_files_presets_and_namespaces() {
let mut input = CacheKeyInput {
fields: vec![
CacheKeyField {
name: "model".into(),
value: Some("deployment".into()),
api_parameter: true,
internal_parameter: false,
},
CacheKeyField {
name: "file".into(),
value: None,
api_parameter: true,
internal_parameter: false,
},
],
namespace: Some("team".into()),
..Default::default()
};
CacheKeyContext {
model_group: Some("group".into()),
caching_groups: vec![(vec!["group".into()], "['group']".into())],
file_checksum: Some("checksum".into()),
..Default::default()
}
.apply(&mut input);
assert_eq!(
cache_key(&input),
format!(
"team:{:x}",
Sha256::digest(b"model: ['group']file: checksum")
)
);
input.preset = Some("preset".into());
assert_eq!(get_cache_key(&input), "preset");
}
#[test]
fn cache_controls_honor_default_modes_and_directives() {
let enabled = CacheControls {
supported_call_type: true,
configured: true,
default_on: true,
..Default::default()
};
assert!(enabled.reads());
assert!(enabled.writes());
assert!(
!CacheControls {
default_on: false,
..enabled
}
.reads()
);
assert!(
CacheControls {
default_on: false,
use_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_store: true,
..enabled
}
.writes()
);
}

24
litellm-rust/crates/core/Cargo.toml
View file

@ -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
litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
rand.workspace = true
@ -31,23 +32,6 @@ 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",
]
[dev-dependencies]
rstest.workspace = true

26
litellm-rust/crates/core/src/audio_transcription/error.rs Normal file
View file

@ -0,0 +1,26 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
InvalidType {
expected: &'static str,
actual: &'static str,
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("unsupported by the rust path: {0}")]
Unsupported(&'static str),
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

26
litellm-rust/crates/core/src/audio_transcription/handler.rs
View file

@ -1,6 +1,6 @@
use serde_json::Value;
use crate::error::Error;
use super::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
@ -21,17 +21,17 @@ pub async fn execute_audio_transcription_provider_call(
}
let response = http_request(request_builder)
.await
.map_err(|error| Error::Network(error.to_string()))?;
.map_err(|error| Error::Transport(crate::transport::Error::Network(error.to_string())))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
.map_err(|error| Error::Transport(crate::transport::Error::Network(error.to_string())))?;
if !status.is_success() {
return Err(Error::Http {
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}));
}
let response_json = serde_json::from_str(&text)
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
@ -41,7 +41,6 @@ pub async fn execute_audio_transcription_provider_call(
.into_json())
}
#[cfg(feature = "bedrock-auth")]
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
body: &[u8],
@ -73,18 +72,3 @@ async fn signed_headers(
)?;
Ok(unsigned.into_iter().chain(signature).collect())
}
#[cfg(not(feature = "bedrock-auth"))]
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
_body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
match request.auth {
AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
"AWS SigV4 requires the bedrock-auth feature",
)),
AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()),
}
}

3
litellm-rust/crates/core/src/audio_transcription/mod.rs
View file

@ -1,4 +1,5 @@
use crate::Error;
mod error;
pub use error::Error;
mod client;
mod handler;
mod prepare;

5
litellm-rust/crates/core/src/audio_transcription/prepare.rs
View file

@ -1,6 +1,5 @@
use crate::error::Error;
use super::Error;
use crate::http_utils::{has_header, string_headers};
#[cfg(feature = "bedrock-auth")]
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
@ -8,7 +7,6 @@ use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderCo
use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
#[cfg(feature = "bedrock-auth")]
if provider == "bedrock" {
return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);
}
@ -65,7 +63,6 @@ pub fn prepare_audio_transcription_provider_call(
body: transformed.body,
upstream_headers: headers,
auth,
#[cfg(feature = "bedrock-auth")]
optional_params: request.optional_params,
timeout: request.timeout,
})

2
litellm-rust/crates/core/src/audio_transcription/transformation.rs
View file

@ -1,4 +1,4 @@
use crate::Error;
use super::Error;
use serde_json::{Map, Value};
use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData};

1
litellm-rust/crates/core/src/audio_transcription/types.rs
View file

@ -25,7 +25,6 @@ pub struct ProviderAudioTranscriptionRequest {
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: AudioTranscriptionAuth,
#[cfg(feature = "bedrock-auth")]
pub(super) optional_params: Map<String, Value>,
pub(super) timeout: Option<Duration>,
}

128
litellm-rust/crates/core/src/auth/error.rs
View file

@ -1,128 +0,0 @@
use thiserror::Error;
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthError {
#[error("invalid authentication configuration: {0}")]
Configuration(#[from] AuthConfigurationError),
#[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>),
#[error("credential caller failed: credential caller returned an empty credential")]
EmptyCallerCredential,
#[error("credential caller failed: Azure AD token provider returned an empty token")]
EmptyAzureToken,
#[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")]
UnresolvedOidcReference,
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter"
)]
MissingApiBase {
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,
#[error("Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable")]
AzureApiKey,
#[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,
#[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,
#[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 {
#[error("AWS profile credentials failed: {0}")]
Profile(String),
#[error("AWS default credentials failed: {0}")]
DefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
WebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
WebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
SigningParameters(String),
#[error("AWS signable request failed: {0}")]
SignableRequest(String),
#[error("AWS request signing failed: {0}")]
Signing(String),
#[error("AWS web identity response had no credentials")]
MissingWebIdentityCredentials,
}

258
litellm-rust/crates/core/src/caching/in_memory_cache.rs
View file

@ -1,258 +0,0 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
pub struct InMemoryCache<V: Clone> {
pub cache_dict: HashMap<String, V>,
pub ttl_dict: HashMap<String, Duration>,
pub expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
pub max_size_in_memory: usize,
pub default_ttl: Duration,
now: Box<dyn Fn() -> Duration + Send + Sync>,
}
impl<V: Clone> Default for InMemoryCache<V> {
fn default() -> Self {
Self::new(None, None)
}
}
impl<V: Clone> InMemoryCache<V> {
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
Self::with_clock(max_size_in_memory, default_ttl, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn with_clock(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self {
cache_dict: HashMap::new(),
ttl_dict: HashMap::new(),
expiration_heap: BinaryHeap::new(),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
now: Box::new(now),
}
}
pub fn evict_cache(&mut self) {
if self.max_size_in_memory == 0 {
return;
}
let current_time = (self.now)();
while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() {
if self.ttl_dict.get(&key).copied() != Some(expiration_time) {
self.expiration_heap.pop();
} else if expiration_time <= current_time {
self.expiration_heap.pop();
self.remove_key(&key);
} else {
break;
}
}
while self.cache_dict.len() >= self.max_size_in_memory {
let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else {
break;
};
if self.ttl_dict.get(&key).copied() == Some(expiration_time) {
self.remove_key(&key);
}
}
}
pub fn allow_ttl_override(&self, key: &str) -> bool {
match self.ttl_dict.get(key).copied() {
None => true,
Some(expiration_time) => expiration_time < (self.now)(),
}
}
pub fn set_cache(&mut self, key: impl Into<String>, value: V, ttl: Option<Duration>) {
if self.max_size_in_memory == 0 {
return;
}
self.evict_cache();
let key = key.into();
self.cache_dict.insert(key.clone(), value);
if self.allow_ttl_override(&key) {
let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl);
self.ttl_dict.insert(key.clone(), expiration_time);
self.expiration_heap.push(Reverse((expiration_time, key)));
}
}
// Generic values intentionally omit Python's per-item size check.
pub fn get_cache(&mut self, key: &str) -> Option<V> {
if self.cache_dict.contains_key(key) {
if self.is_key_expired(key) {
self.remove_key(key);
return None;
}
return self.cache_dict.get(key).cloned();
}
None
}
pub fn get_ttl(&self, key: &str) -> Option<Duration> {
self.ttl_dict.get(key).copied()
}
pub fn delete_cache(&mut self, key: &str) {
self.remove_key(key);
}
pub fn flush_cache(&mut self) {
self.cache_dict.clear();
self.ttl_dict.clear();
self.expiration_heap.clear();
}
fn is_key_expired(&self, key: &str) -> bool {
self.ttl_dict
.get(key)
.is_some_and(|expiration_time| *expiration_time < (self.now)())
}
fn remove_key(&mut self, key: &str) {
self.cache_dict.remove(key);
self.ttl_dict.remove(key);
}
}
#[cfg(test)]
mod tests {
use std::sync::{
Arc,
atomic::{AtomicU64, Ordering},
};
use super::InMemoryCache;
use std::time::Duration;
fn cache(now: Arc<AtomicU64>, max_size: usize, default_ttl: Duration) -> InMemoryCache<String> {
InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || {
Duration::from_secs(now.load(Ordering::Relaxed))
})
}
#[test]
fn ttl_expiry_is_deterministic() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("key", "value".to_string(), None);
assert_eq!(cache.get_cache("key"), Some("value".to_string()));
now.store(161, Ordering::Relaxed);
assert_eq!(cache.get_cache("key"), None);
assert_eq!(cache.get_ttl("key"), None);
}
#[test]
fn default_and_per_set_ttl_are_applied() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("default", "value".to_string(), None);
cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20)));
assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160)));
assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120)));
}
#[test]
fn unexpired_entries_do_not_allow_ttl_override() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80)));
assert_eq!(cache.get_cache("key"), Some("second".to_string()));
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120)));
now.store(121, Ordering::Relaxed);
cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80)));
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201)));
}
#[test]
fn max_size_evicts_earliest_expiration() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 2, Duration::from_secs(60));
cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10)));
cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
assert_eq!(cache.get_cache("early"), None);
assert!(cache.get_cache("late").is_some());
assert!(cache.get_cache("new").is_some());
}
#[test]
fn expired_entries_are_evicted_before_live_entries() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 3, Duration::from_secs(60));
cache.set_cache(
"expired-one",
"value".to_string(),
Some(Duration::from_secs(10)),
);
cache.set_cache(
"expired-two",
"value".to_string(),
Some(Duration::from_secs(20)),
);
cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100)));
now.store(121, Ordering::Relaxed);
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100)));
assert_eq!(cache.get_cache("expired-one"), None);
assert_eq!(cache.get_cache("expired-two"), None);
assert!(cache.get_cache("live").is_some());
assert!(cache.get_cache("new").is_some());
}
#[test]
fn stale_heap_entries_are_skipped() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 1, Duration::from_secs(60));
cache.set_cache(
"removed",
"value".to_string(),
Some(Duration::from_secs(10)),
);
cache.delete_cache("removed");
cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
assert_eq!(cache.get_cache("removed"), None);
assert_eq!(cache.get_cache("kept"), None);
assert!(cache.get_cache("new").is_some());
}
#[test]
fn delete_and_flush_remove_values_and_ttls() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 10, Duration::from_secs(60));
cache.set_cache("one", "value".to_string(), None);
cache.set_cache("two", "value".to_string(), None);
cache.delete_cache("one");
assert_eq!(cache.get_cache("one"), None);
cache.flush_cache();
assert!(cache.cache_dict.is_empty());
assert!(cache.ttl_dict.is_empty());
assert!(cache.expiration_heap.is_empty());
}
#[test]
fn zero_max_size_does_not_cache() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 0, Duration::from_secs(60));
cache.set_cache("key", "value".to_string(), None);
assert_eq!(cache.get_cache("key"), None);
assert!(cache.cache_dict.is_empty());
}
}

1
litellm-rust/crates/core/src/caching/mod.rs
View file

@ -1 +0,0 @@
pub mod in_memory_cache;

19
litellm-rust/crates/core/src/call_lifecycle/host.rs
View file

@ -6,10 +6,11 @@ pub enum HostCallStep<O, C> {
Complete(C),
}
pub type HostCallFuture<'a, O, C> =
Pin<Box<dyn Future<Output = Result<HostCallStep<O, C>, crate::Error>> + Send + 'a>>;
pub type HostCallFuture<'a, O, C, E> =
Pin<Box<dyn Future<Output = Result<HostCallStep<O, C>, E>> + Send + 'a>>;
pub trait HostCall: Send + Sync {
type Error: Send + Sync + 'static;
type Operation: Send + 'static;
type Result: Send + 'static;
type Complete: Send + 'static;
@ -17,12 +18,12 @@ pub trait HostCall: Send + Sync {
fn resume(
&mut self,
result: Option<Self::Result>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>;
fn interrupt(
&mut self,
failure: HostFailure,
) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
failure: HostFailure<Self::Error>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>;
}
pub enum HostStep<V, S> {
@ -48,9 +49,9 @@ pub enum HostPhase {
}
#[derive(Clone, Debug)]
pub enum HostFailure {
Error(crate::Error),
Cancelled(crate::Error),
pub enum HostFailure<E> {
Error(E),
Cancelled(E),
}
pub struct HostLifecycle {
@ -70,7 +71,7 @@ impl HostLifecycle {
self.phase
}
pub fn accept(&mut self, result: Result<(), HostFailure>) -> Option<crate::Error> {
pub fn accept<E>(&mut self, result: Result<(), HostFailure<E>>) -> Option<E> {
if let Err(failure) = result {
if self.phase == HostPhase::DeploymentFailure {
self.phase = HostPhase::Failure;

44
litellm-rust/crates/core/src/call_lifecycle/mod.rs
View file

@ -1,8 +1,6 @@
use std::future::Future;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use crate::Error;
pub mod host;
#[cfg(test)]
#[path = "../../tests/host_lifecycle.rs"]
@ -15,14 +13,15 @@ pub use types::{
};
pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
type PreCallFuture<'a>: Future<Output = Result<InitialReq, Error>> + Send + 'a
type Error: Send + Sync;
type PreCallFuture<'a>: Future<Output = Result<InitialReq, Self::Error>> + Send + 'a
where
Self: 'a,
InitialReq: 'a,
ProviderReq: 'a,
Resp: 'a;
type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Error>> + Send + 'a
type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Self::Error>> + Send + 'a
where
Self: 'a,
InitialReq: 'a,
@ -60,7 +59,7 @@ pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a Error,
error: &'a Self::Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a>;
}
@ -90,12 +89,12 @@ impl<'a> CallLifecycle<'a> {
request: InitialReq,
hooks: &Hooks,
provider_call: ProviderCall,
) -> Result<Resp, Error>
) -> Result<Resp, Hooks::Error>
where
InitialReq: CallLifecycleRequest,
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
ProviderFuture: Future<Output = Result<Resp, Hooks::Error>>,
{
let context = request.lifecycle_context();
self.run(context, request, hooks, provider_call).await
@ -107,11 +106,11 @@ impl<'a> CallLifecycle<'a> {
request: InitialReq,
hooks: &Hooks,
provider_call: ProviderCall,
) -> Result<Resp, Error>
) -> Result<Resp, Hooks::Error>
where
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
ProviderFuture: Future<Output = Result<Resp, Hooks::Error>>,
{
let call_start = epoch_seconds();
let mut phases = Vec::new();
@ -170,7 +169,7 @@ impl<'a> CallLifecycle<'a> {
&self,
context: &CallLifecycleContext,
hooks: &Hooks,
error: &Error,
error: &Hooks::Error,
call_start: f64,
phases: &mut Vec<CallLifecyclePhaseTiming>,
) where
@ -255,8 +254,9 @@ mod tests {
}
impl CallLifecycleHooks<String, String, String> for RecordingHooks {
type PreCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type Error = crate::messages::Error;
type PreCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type SuccessFuture<'a> = BoxFuture<'a, ()>;
type FailureFuture<'a> = BoxFuture<'a, ()>;
@ -298,7 +298,7 @@ mod tests {
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a Error,
_error: &'a crate::messages::Error,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
@ -308,8 +308,9 @@ mod tests {
}
impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type Error = crate::messages::Error;
type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, crate::messages::Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type SuccessFuture<'a> = BoxFuture<'a, ()>;
type FailureFuture<'a> = BoxFuture<'a, ()>;
@ -349,7 +350,7 @@ mod tests {
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a Error,
_error: &'a crate::messages::Error,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
@ -387,13 +388,20 @@ mod tests {
"request".to_string(),
&hooks,
|_request| async move {
Err::<String, Error>(Error::Network("provider down".to_string()))
Err::<String, crate::messages::Error>(crate::messages::Error::Transport(
crate::transport::Error::Network("provider down".to_string()),
))
},
)
.await
.expect_err("call fails");
assert_eq!(error, Error::Network("provider down".to_string()));
assert_eq!(
error,
crate::messages::Error::Transport(crate::transport::Error::Network(
"provider down".to_string()
))
);
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]);
}

5
litellm-rust/crates/core/src/chat_completions/common_utils.rs
View file

@ -1,4 +1,4 @@
use crate::Error;
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use serde_json::{Map, Value};
@ -12,7 +12,6 @@ pub(super) fn chat_completions_provider_config(
) -> Option<&'static dyn ChatCompletionsProviderConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG),
#[cfg(feature = "bedrock-auth")]
"bedrock" => Some(
&crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
),
@ -23,5 +22,5 @@ pub(super) fn chat_completions_provider_config(
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers)
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
}

26
litellm-rust/crates/core/src/chat_completions/error.rs Normal file
View file

@ -0,0 +1,26 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
InvalidType {
expected: &'static str,
actual: &'static str,
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("unsupported by the rust path: {0}")]
Unsupported(&'static str),
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

29
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -1,6 +1,6 @@
use serde_json::Value;
use crate::error::Error;
use super::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
@ -35,9 +35,9 @@ pub(super) async fn execute_chat_completions_provider_call(
// so the host can still serve it. Everything else here, a timeout
// above all, may have reached the provider and been answered.
if err.is_connect() || err.is_builder() {
Error::Connect(err.to_string())
Error::Transport(crate::transport::Error::Connect(err.to_string()))
} else {
Error::Network(err.to_string())
Error::Transport(crate::transport::Error::Network(err.to_string()))
}
})?;
@ -45,13 +45,13 @@ pub(super) async fn execute_chat_completions_provider_call(
let text = response
.text()
.await
.map_err(|err| Error::Network(err.to_string()))?;
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
if !status.is_success() {
return Err(Error::Http {
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}));
}
let body: Value = serde_json::from_str(&text).map_err(|err| {
@ -74,12 +74,12 @@ pub(super) async fn execute_chat_completions_provider_call(
/// can only mean the provider was already called.
pub(super) fn as_response_error(err: Error) -> Error {
match err {
already @ (Error::InvalidResponse(_) | Error::Http { .. }) => already,
already @ (Error::InvalidResponse(_)
| Error::Transport(crate::transport::Error::Http { .. })) => already,
other => Error::InvalidResponse(other.to_string()),
}
}
#[cfg(feature = "bedrock-auth")]
pub(super) async fn signed_headers(
request: &ProviderChatCompletionsRequest,
body: &[u8],
@ -135,16 +135,3 @@ pub(super) async fn signed_headers(
// that would collide, so no name appears twice.
Ok(unsigned.into_iter().chain(signature).collect())
}
#[cfg(not(feature = "bedrock-auth"))]
pub(super) async fn signed_headers(
request: &ProviderChatCompletionsRequest,
_body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
match &request.auth {
ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
"AWS SigV4 requires the bedrock-auth feature",
)),
_ => Ok(request.upstream_headers.clone()),
}
}

3
litellm-rust/crates/core/src/chat_completions/mod.rs
View file

@ -6,7 +6,8 @@
//! credentials, and it resolves the provider, translates the conversation,
//! calls the provider, and returns a typed OpenAI-shaped response.
use crate::Error;
mod error;
pub use error::Error;
mod client;
mod common_utils;
pub mod conversation;

2
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -1,6 +1,6 @@
use serde_json::Value;
use crate::error::Error;
use super::Error;
use crate::http_utils::has_header;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};

30
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -1,6 +1,6 @@
use serde_json::{Map, Value, json};
use crate::error::Error;
use super::Error;
use super::prepare::{prepare_provider_request, resolve_request};
use super::transformation::ChatCompletionsAuth;
@ -264,13 +264,14 @@ fn rejects_non_string_extra_headers() {
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
assert_eq!(
decline(call),
Error::InvalidRequest(
"chat completions extra_headers.x-trace must be a string, got number".to_string()
)
Error::Headers(crate::http_utils::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
})
);
}
#[cfg(feature = "bedrock-auth")]
#[test]
fn prepares_a_bedrock_call_without_resolving_credentials() {
let mut call = request(
@ -302,7 +303,6 @@ fn prepares_a_bedrock_call_without_resolving_credentials() {
assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16}));
}
#[cfg(feature = "bedrock-auth")]
#[tokio::test]
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
// Python signs only the AWS header set and reattaches the rest, so a header
@ -351,7 +351,6 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
);
}
#[cfg(feature = "bedrock-auth")]
#[tokio::test]
async fn a_forwarded_header_the_signer_computes_declines_to_python() {
// Reattaching the caller's copy next to the computed one puts the name on
@ -386,7 +385,6 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() {
}
}
#[cfg(feature = "bedrock-auth")]
#[test]
fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
// `get_request_headers` assigns `headers["Authorization"]` unconditionally
@ -453,7 +451,6 @@ fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
);
}
#[cfg(feature = "bedrock-auth")]
#[test]
fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
// The configured bearer identity has its own account and quota boundary,
@ -769,7 +766,10 @@ mod round_trip {
.expect_err("upstream rejects");
handle.await.expect("server task");
assert!(
matches!(err, Error::Http { status: 429, .. }),
matches!(
err,
Error::Transport(crate::transport::Error::Http { status: 429, .. })
),
"expected a 429, got {err:?}"
);
}
@ -793,7 +793,7 @@ mod round_trip {
.await
.expect_err("nothing is listening");
assert!(
matches!(err, Error::Connect(_)),
matches!(err, Error::Transport(crate::transport::Error::Connect(_))),
"expected a pre-send connect failure, got {err:?}"
);
}
@ -806,7 +806,7 @@ mod round_trip {
Error::MissingField("usage"),
Error::Unsupported("non-text response content block"),
Error::InvalidRequest("whatever".to_string()),
Error::Auth("whatever".to_string()),
Error::Auth(litellm_auth::Error::InvalidHeader),
] {
let label = format!("{original:?}");
assert!(
@ -816,11 +816,11 @@ mod round_trip {
}
// An upstream status is already unambiguous, so it survives intact.
assert!(matches!(
as_response_error(Error::Http {
as_response_error(Error::Transport(crate::transport::Error::Http {
status: 500,
body: "boom".to_string()
}),
Error::Http { status: 500, .. }
})),
Error::Transport(crate::transport::Error::Http { status: 500, .. })
));
}
}

2
litellm-rust/crates/core/src/chat_completions/transformation.rs
View file

@ -1,4 +1,4 @@
use crate::Error;
use super::Error;
use serde_json::{Map, Value};
use super::types::{

1
litellm-rust/crates/core/src/chat_completions/types.rs
View file

@ -40,7 +40,6 @@ pub(super) struct ProviderChatCompletionsRequest {
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: ChatCompletionsAuth,
#[cfg_attr(not(feature = "bedrock-auth"), allow(dead_code))]
pub(super) optional_params: Map<String, Value>,
pub(super) timeout: Option<Duration>,
}

229
litellm-rust/crates/core/src/error.rs
View file

@ -1,220 +1,13 @@
use thiserror::Error as ThisError;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
InvalidType {
expected: &'static str,
actual: &'static str,
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("Document URL is required")]
MissingDocumentUrl,
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("{0}")]
Auth(String),
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID"
)]
MissingAzureAiCredentials,
#[error(
"invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID"
)]
MissingAzureDocumentIntelligenceCredentials,
#[error(
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
)]
MissingReductoApiKey,
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
Network(String),
/// The provider was never reached: DNS, TCP, TLS or proxy setup failed
/// before any byte of the request went out. Nothing was billed, so a host
/// that keeps a reference implementation can serve the request itself.
/// A timeout is deliberately not this, since the provider may have received
/// and answered the request already.
#[error("could not reach the provider: {0}")]
Connect(String),
#[error("routing error: {0}")]
Routing(String),
/// The request is outside the surface this route covers in Rust. Hosts that
/// keep a reference implementation treat this as "fall back", not "fail".
#[error("unsupported by the rust path: {0}")]
Unsupported(&'static str),
}
impl Error {
pub const fn http_status_code(&self) -> Option<u16> {
match self {
Self::InvalidRequest(_) => Some(400),
Self::MissingDocumentUrl => Some(500),
Self::Http { status, .. } => Some(*status),
_ => None,
}
}
}
#[derive(Debug, ThisError)]
pub(crate) enum MediaError {
#[error("media URL rejected by network policy")]
BlockedUrl,
#[error("media download is disabled")]
DownloadDisabled,
#[error("media download exceeds the maximum size")]
DownloadTooLarge,
#[error("too many redirects while fetching media")]
TooManyRedirects,
#[error("media redirect is missing a Location header")]
MissingRedirectLocation,
#[error("invalid media redirect")]
InvalidRedirect,
#[error("media download failed with status {0}")]
Http(u16),
#[error("media download timed out")]
Timeout,
#[error("{0}")]
Transport(#[from] TransportError),
}
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum TransportError {
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
Network(String),
#[error("could not reach the provider: {0}")]
Connect(String),
}
impl TransportError {
pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self {
let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder());
let message = error.without_url().to_string();
if before_dispatch {
Self::Connect(message)
} else {
Self::Network(message)
}
}
}
impl From<reqwest::Error> for TransportError {
fn from(error: reqwest::Error) -> Self {
Self::Network(error.without_url().to_string())
}
}
impl From<crate::ocr::error::OcrRequestError> for Error {
fn from(error: crate::ocr::error::OcrRequestError) -> Self {
match error {
crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field),
crate::ocr::error::OcrRequestError::MissingDocumentUrl => Self::MissingDocumentUrl,
error => Self::InvalidRequest(error.to_string()),
}
}
}
impl From<crate::ocr::error::OcrResponseError> for Error {
fn from(error: crate::ocr::error::OcrResponseError) -> Self {
Self::InvalidResponse(error.to_string())
}
}
impl From<TransportError> for Error {
fn from(error: TransportError) -> Self {
match error {
TransportError::Http { status, body } => Self::Http { status, body },
TransportError::Network(message) => Self::Network(message),
TransportError::Connect(message) => Self::Connect(message),
}
}
}
impl From<crate::AuthError> for Error {
fn from(error: crate::AuthError) -> Self {
match error {
crate::AuthError::MissingApiKey { provider } => Self::MissingApiKey { provider },
error => Self::Auth(error.to_string()),
}
}
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "bool",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => "object",
}
}
#[cfg(test)]
mod transport_tests {
use super::*;
#[test]
fn missing_auth_key_preserves_provider_in_public_error() {
assert_eq!(
Error::from(crate::AuthError::MissingApiKey { provider: "Vertex" }),
Error::MissingApiKey { provider: "Vertex" }
);
}
#[tokio::test]
async fn transport_errors_remove_urls_and_keep_dispatch_context() {
let error = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get("http://localhost:invalid/private?api_key=secret")
.send()
.await
.expect_err("invalid port");
let error = TransportError::from_reqwest_before_dispatch(error);
assert!(matches!(error, TransportError::Connect(_)));
assert!(!error.to_string().contains("secret"));
assert!(!error.to_string().contains("private"));
}
#[tokio::test]
async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() {
use std::time::Duration;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let address = listener.local_addr().expect("address");
let request = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}"))
.timeout(Duration::from_millis(200))
.send();
let (response, accepted) = tokio::join!(
request,
tokio::time::timeout(Duration::from_secs(2), listener.accept())
);
let _connection = accepted
.expect("accept deadline")
.expect("accepted connection");
let error = response.expect_err("server does not respond");
assert!(error.is_timeout());
assert!(matches!(
TransportError::from_reqwest_before_dispatch(error),
TransportError::Network(_)
));
}
#[error(transparent)]
Ocr(#[from] crate::ocr::Error),
#[error(transparent)]
Messages(#[from] crate::messages::Error),
#[error(transparent)]
ChatCompletions(#[from] crate::chat_completions::Error),
#[error(transparent)]
AudioTranscription(#[from] crate::audio_transcription::Error),
#[error(transparent)]
Responses(#[from] crate::responses::Error),
}

46
litellm-rust/crates/core/src/http_utils.rs
View file

@ -1,7 +1,14 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("invalid request: {context} extra_headers.{name} must be a string, got {actual}")]
pub struct HeaderError {
pub context: &'static str,
pub name: String,
pub actual: &'static str,
}
use serde_json::{Map, Value};
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
use crate::error::{Error, json_type_name};
#[allow(
dead_code,
@ -44,6 +51,13 @@ pub async fn http_request(
request.send().await
}
pub async fn execute_http_request(
client: &reqwest::Client,
request: reqwest::Request,
) -> Result<reqwest::Response, reqwest::Error> {
client.execute(request).await
}
pub fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS {
return body.to_string();
@ -55,7 +69,7 @@ pub fn truncate_error_body(body: &str) -> String {
pub fn string_headers(
context: &'static str,
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
) -> Result<Vec<(String, String)>, HeaderError> {
extra_headers
.unwrap_or_default()
.into_iter()
@ -63,11 +77,10 @@ pub fn string_headers(
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
Error::InvalidRequest(format!(
"{context} extra_headers.{key} must be a string, got {}",
json_type_name(&value)
))
.ok_or_else(|| HeaderError {
context,
name: key,
actual: json_type_name(&value),
})
})
.collect()
@ -105,6 +118,17 @@ where
<Option<T> as serde::Deserialize>::deserialize(deserializer).map(Some)
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "bool",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => "object",
}
}
#[cfg(test)]
mod tests {
use super::*;
@ -184,9 +208,11 @@ mod tests {
let err = string_headers("chat completions", Some(headers)).expect_err("non-string value");
assert_eq!(
err,
Error::InvalidRequest(
"chat completions extra_headers.x-trace must be a string, got number".to_string()
)
HeaderError {
context: "chat completions",
name: "x-trace".into(),
actual: "number"
}
);
}

4
litellm-rust/crates/core/src/lib.rs
View file

@ -1,6 +1,4 @@
pub mod audio_transcription;
pub mod auth;
pub mod caching;
pub mod call_lifecycle;
pub mod chat_completions;
pub mod constants;
@ -11,7 +9,7 @@ pub mod messages;
pub mod ocr;
pub mod providers;
pub mod responses;
pub mod transport;
mod url_utils;
pub use auth::AuthError;
pub use error::Error;

81
litellm-rust/crates/core/src/media.rs
View file

@ -9,7 +9,28 @@ use reqwest::Url;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use crate::constants::MEDIA_CONNECT_TIMEOUT_SECS;
use crate::error::{MediaError, TransportError};
#[derive(Debug, thiserror::Error)]
pub(crate) enum Error {
#[error("media URL rejected by network policy")]
BlockedUrl,
#[error("media download is disabled")]
DownloadDisabled,
#[error("media download exceeds the maximum size")]
DownloadTooLarge,
#[error("too many redirects while fetching media")]
TooManyRedirects,
#[error("media redirect is missing a Location header")]
MissingRedirectLocation,
#[error("invalid media redirect")]
InvalidRedirect,
#[error("media download failed with status {0}")]
Http(u16),
#[error("media download timed out")]
Timeout,
#[error("{0}")]
Transport(#[from] crate::transport::Error),
}
#[derive(Clone)]
pub(crate) struct MediaFetcher {
@ -75,20 +96,20 @@ impl MediaFetcher {
&self,
url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, MediaError> {
) -> Result<DownloadedMedia, Error> {
if policy.max_bytes == 0 {
return Err(MediaError::DownloadDisabled);
return Err(Error::DownloadDisabled);
}
tokio::time::timeout(policy.timeout, self.fetch_before_deadline(url, policy))
.await
.map_err(|_| MediaError::Timeout)?
.map_err(|_| Error::Timeout)?
}
async fn fetch_before_deadline(
&self,
mut url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, MediaError> {
) -> Result<DownloadedMedia, Error> {
let mut redirects_followed = 0;
loop {
self.validate_url(&url).await?;
@ -97,24 +118,22 @@ impl MediaFetcher {
.get(url.clone())
.send()
.await
.map_err(TransportError::from)?;
.map_err(crate::transport::Error::from)?;
if response.status().is_redirection() {
if redirects_followed == policy.max_redirects {
return Err(MediaError::TooManyRedirects);
return Err(Error::TooManyRedirects);
}
let location = response
.headers()
.get(reqwest::header::LOCATION)
.and_then(|value| value.to_str().ok())
.ok_or(MediaError::MissingRedirectLocation)?;
url = url
.join(location)
.map_err(|_| MediaError::InvalidRedirect)?;
.ok_or(Error::MissingRedirectLocation)?;
url = url.join(location).map_err(|_| Error::InvalidRedirect)?;
redirects_followed += 1;
continue;
}
if !response.status().is_success() {
return Err(MediaError::Http(response.status().as_u16()));
return Err(Error::Http(response.status().as_u16()));
}
enforce_download_size(response.content_length().unwrap_or(0), policy.max_bytes)?;
let content_type = response
@ -127,7 +146,11 @@ impl MediaFetcher {
.unwrap_or("application/octet-stream")
.to_string();
let mut bytes = Vec::new();
while let Some(chunk) = response.chunk().await.map_err(TransportError::from)? {
while let Some(chunk) = response
.chunk()
.await
.map_err(crate::transport::Error::from)?
{
enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?;
bytes.extend_from_slice(&chunk);
}
@ -138,42 +161,40 @@ impl MediaFetcher {
}
}
async fn validate_url(&self, url: &Url) -> Result<(), MediaError> {
async fn validate_url(&self, url: &Url) -> Result<(), Error> {
if !matches!(url.scheme(), "http" | "https")
|| !url.username().is_empty()
|| url.password().is_some()
{
return Err(MediaError::BlockedUrl);
return Err(Error::BlockedUrl);
}
let host = url.host_str().ok_or(MediaError::BlockedUrl)?;
let host = url.host_str().ok_or(Error::BlockedUrl)?;
if self.allow_private_network {
return Ok(());
}
if let Ok(ip) = host.parse::<IpAddr>() {
return (!is_blocked_ip(ip))
.then_some(())
.ok_or(MediaError::BlockedUrl);
return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl);
}
let port = url.port_or_known_default().ok_or(MediaError::BlockedUrl)?;
let port = url.port_or_known_default().ok_or(Error::BlockedUrl)?;
let addresses = self
.address_resolver
.resolve(host, port)
.await
.map_err(|error| TransportError::Network(error.to_string()))?;
.map_err(|error| crate::transport::Error::Network(error.to_string()))?;
validate_addresses(&addresses)
}
}
fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), MediaError> {
fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), Error> {
if length > max_bytes {
return Err(MediaError::DownloadTooLarge);
return Err(Error::DownloadTooLarge);
}
Ok(())
}
fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), MediaError> {
fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), Error> {
if addresses.is_empty() || addresses.iter().any(|address| is_blocked_ip(address.ip())) {
return Err(MediaError::BlockedUrl);
return Err(Error::BlockedUrl);
}
Ok(())
}
@ -415,7 +436,7 @@ mod tests {
.await
.expect_err("oversize body is rejected");
server.await.expect("server completes");
assert!(matches!(error, MediaError::DownloadTooLarge));
assert!(matches!(error, Error::DownloadTooLarge));
}
#[tokio::test]
@ -433,7 +454,7 @@ mod tests {
.await
.expect_err("stream crossing limit is rejected");
server.await.expect("server completes");
assert!(matches!(error, MediaError::DownloadTooLarge));
assert!(matches!(error, Error::DownloadTooLarge));
}
#[tokio::test]
@ -469,7 +490,7 @@ mod tests {
.expect_err("private redirect is rejected");
let requests = server.await.expect("server completes");
assert_eq!(requests.len(), 1);
assert!(matches!(error, MediaError::BlockedUrl));
assert!(matches!(error, Error::BlockedUrl));
}
#[tokio::test]
@ -496,7 +517,7 @@ mod tests {
.await
.expect_err("fetch times out");
server.await.expect("server completes");
assert!(matches!(error, MediaError::Timeout));
assert!(matches!(error, Error::Timeout));
}
#[tokio::test]
@ -522,7 +543,7 @@ mod tests {
Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses");
assert!(matches!(
fetcher.validate_url(&url).await,
Err(MediaError::BlockedUrl)
Err(Error::BlockedUrl)
));
}
}

4
litellm-rust/crates/core/src/messages/common_utils.rs
View file

@ -1,4 +1,4 @@
use crate::Error;
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
@ -23,5 +23,5 @@ pub(super) fn messages_provider_config(
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers)
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
}

17
litellm-rust/crates/core/src/messages/error.rs Normal file
View file

@ -0,0 +1,17 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("routing error: {0}")]
Routing(String),
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
}

18
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -1,5 +1,5 @@
use super::Error;
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
use crate::error::Error;
use crate::http_utils::http_request;
use super::client::http_client;
@ -21,19 +21,19 @@ pub(super) async fn execute_messages_provider_call(
let response = http_request(request_builder)
.await
.map_err(|err| Error::Network(err.to_string()))?;
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| Error::Network(err.to_string()))?;
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
if !status.is_success() {
return Err(Error::Http {
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}));
}
let response = serde_json::from_str(&text)
@ -61,17 +61,17 @@ pub(super) async fn execute_messages_provider_stream(
let response = http_request(request_builder)
.await
.map_err(|err| Error::Network(err.to_string()))?;
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
let status = response.status();
if !status.is_success() {
let text = response
.text()
.await
.map_err(|err| Error::Network(err.to_string()))?;
return Err(Error::Http {
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}));
}
Ok(response)
}

3
litellm-rust/crates/core/src/messages/mod.rs
View file

@ -7,7 +7,8 @@
//! is the streaming variant; it hands the raw upstream response back so a host
//! can splice the event stream to its own caller.
use crate::Error;
mod error;
pub use error::Error;
mod client;
mod common_utils;
mod handler;

2
litellm-rust/crates/core/src/messages/prepare.rs
View file

@ -1,4 +1,4 @@
use crate::error::Error;
use super::Error;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};

16
litellm-rust/crates/core/src/messages/tests.rs
View file

@ -4,7 +4,7 @@ use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use crate::error::Error;
use super::Error;
use super::common_utils::{
has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body,
@ -77,7 +77,14 @@ fn truncate_error_body_caps_long_payloads() {
fn string_headers_rejects_non_string_values() {
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert!(matches!(err, Error::InvalidRequest(_)));
assert_eq!(
err,
Error::Headers(crate::http_utils::HeaderError {
context: "messages",
name: "x-count".to_string(),
actual: "number",
})
);
}
#[test]
@ -420,7 +427,10 @@ async fn messages_maps_provider_error_status_to_http_error() {
.await
.expect_err("provider error propagates");
assert!(matches!(err, Error::Http { status: 401, .. }));
assert!(matches!(
err,
Error::Transport(crate::transport::Error::Http { status: 401, .. })
));
}
#[tokio::test]

2
litellm-rust/crates/core/src/messages/transformation.rs
View file

@ -1,5 +1,5 @@
use super::Error;
use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
use crate::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MessagesAuthStrategy {

4
litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs
View file

@ -1,5 +1,5 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::cohere::{
CohereParams, CohereResponse, transform_request, transform_response, validate_document,
@ -9,8 +9,8 @@ use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
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";

6
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs
View file

@ -1,7 +1,6 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::auth::{InputSource, Sourced};
use crate::constants::{AZURE_DI_API_VERSION, AZURE_DI_SUBSCRIPTION_HEADER};
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::document_intelligence::{
self, AzureDocumentIntelligenceOperation, DocumentIntelligenceParams,
@ -10,8 +9,9 @@ use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat};
use crate::providers::azure_ai::auth::AzureAuthInputs;
use crate::url_utils::ApiUrl;
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::AzureAuthInputs;
mod polling;

2
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs
View file

@ -77,7 +77,7 @@ async fn poll_operation(
let response = tokio::time::timeout_at(deadline, crate::http_utils::http_request(builder))
.await
.map_err(|_| OcrPollingError::PollTimeout)?
.map_err(crate::error::TransportError::from)?;
.map_err(crate::transport::Error::from)?;
let retry = response
.headers()
.get(reqwest::header::RETRY_AFTER)

6
litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs
View file

@ -1,7 +1,6 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::auth::{InputSource, Sourced};
use crate::constants::AZURE_AI_OCR_PATH;
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
@ -11,8 +10,9 @@ use crate::ocr::prepare::{
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
use crate::providers::azure_ai::auth::AzureAuthInputs;
use crate::url_utils::ApiUrl;
use litellm_auth::{InputSource, Sourced};
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";

15
litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs
View file

@ -4,12 +4,12 @@ mod mistral;
use std::sync::OnceLock;
use crate::Error;
use crate::auth::error::AuthConfigurationError;
use crate::auth::{InputSource, Sourced};
use crate::ocr::Error;
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService};
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
pub(crate) use cohere::AzureCohereAdapter;
pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter;
@ -26,7 +26,7 @@ async fn resolve_entra(
.get_azure_ad_token(config, env_lookup)
.await
.or_else(|error| match error {
crate::AuthError::EmptyAzureToken => Ok(None),
litellm_auth::Error::EmptyAzureToken => Ok(None),
other => Err(other),
})
.map(|credential| {
@ -47,10 +47,7 @@ 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(litellm_auth::Error::RequestAzureCredentialDestination).into());
}
Ok(())
}

2
litellm-rust/crates/core/src/ocr/adapters/cohere.rs
View file

@ -1,6 +1,6 @@
use super::OcrAdapter;
use crate::Error;
use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE};
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::cohere::{
CohereParams, CohereResponse, transform_request, transform_response, validate_document,

2
litellm-rust/crates/core/src/ocr/adapters/mistral.rs
View file

@ -1,6 +1,6 @@
use super::OcrAdapter;
use crate::Error;
use crate::constants::MISTRAL_OCR_API_BASE;
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};

4
litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs
View file

@ -1,8 +1,8 @@
mod legacy;
mod v3;
use crate::Error;
use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX};
use crate::ocr::Error;
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::types::{OcrConnection, OcrDocument};
@ -90,7 +90,7 @@ pub(super) async fn prepare_document(
);
let response = crate::http_utils::http_request(builder)
.await
.map_err(crate::error::TransportError::from)?;
.map_err(crate::transport::Error::from)?;
let uploaded = crate::ocr::client::read_json_response::<
crate::ocr::codecs::reducto::ReductoUploadResponse,
>(response, false, connection.max_response_bytes)

4
litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs
View file

@ -1,7 +1,6 @@
use super::super::OcrAdapter;
use super::validate_destination;
use crate::Error;
use crate::auth::vertex::{self, VertexConfig};
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::deepseek::{self, DeepSeekOcrParams, DeepSeekOcrResponse};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
@ -11,6 +10,7 @@ use crate::ocr::prepare::{
use crate::ocr::registry::OcrProvider;
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";

4
litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs
View file

@ -1,7 +1,6 @@
use super::super::OcrAdapter;
use super::validate_destination;
use crate::Error;
use crate::auth::vertex::{self, VertexConfig};
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
@ -12,6 +11,7 @@ use crate::ocr::prepare::{
use crate::ocr::registry::OcrProvider;
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)]

11
litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs
View file

@ -1,9 +1,9 @@
mod deepseek;
mod mistral;
use crate::Error;
use crate::auth::InputSource;
use crate::auth::error::AuthConfigurationError;
use crate::ocr::Error;
use litellm_auth::InputSource;
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
@ -12,10 +12,7 @@ pub(crate) use mistral::VertexMistralAdapter;
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(litellm_auth::Error::RequestVertexCredentialDestination).into());
}
Ok(())
}

11
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -4,14 +4,13 @@ use std::time::Duration;
use bytes::{Bytes, BytesMut};
use serde::de::DeserializeOwned;
use super::error::{OcrError, OcrResponseError};
use super::error::{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 crate::transport::Error as TransportError;
use litellm_auth_gcp::VertexAuth;
#[derive(Clone)]
pub struct OcrClient {
@ -158,7 +157,7 @@ pub(crate) async fn read_response_bytes(
}
}
if !status.is_success() {
return Err(crate::error::TransportError::Http {
return Err(crate::transport::Error::Http {
status: status.as_u16(),
body: crate::http_utils::truncate_error_body(&String::from_utf8_lossy(&bytes)),
}
@ -174,7 +173,7 @@ pub(crate) fn transport_error(error: reqwest::Error) -> Error {
body: "OCR request timed out".into(),
};
}
crate::error::TransportError::from(error).into()
crate::transport::Error::from(error).into()
}
#[cfg(test)]

3
litellm-rust/crates/core/src/ocr/document.rs
View file

@ -7,8 +7,9 @@ use serde_json::Map;
use super::error::{OcrError, OcrRequestError, OcrResponseError};
use super::types::{OcrConnection, OcrDocument};
use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS};
use crate::error::{MediaError, TransportError};
use crate::media::Error as MediaError;
use crate::media::{DownloadPolicy, MediaFetcher};
use crate::transport::Error as TransportError;
pub fn encode_file_document(
bytes: &[u8],

108
litellm-rust/crates/core/src/ocr/error.rs
View file

@ -1,6 +1,106 @@
use thiserror::Error;
use crate::error::TransportError;
use crate::transport::Error as TransportError;
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
InvalidType {
expected: &'static str,
actual: &'static str,
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("Document URL is required")]
MissingDocumentUrl,
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("{0}")]
Auth(String),
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID"
)]
MissingAzureAiCredentials,
#[error(
"invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID"
)]
MissingAzureDocumentIntelligenceCredentials,
#[error(
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
)]
MissingReductoApiKey,
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
Network(String),
/// The provider was never reached: DNS, TCP, TLS or proxy setup failed
/// before any byte of the request went out. Nothing was billed, so a host
/// that keeps a reference implementation can serve the request itself.
/// A timeout is deliberately not this, since the provider may have received
/// and answered the request already.
#[error("could not reach the provider: {0}")]
Connect(String),
#[error("routing error: {0}")]
Routing(String),
/// The request is outside the surface this route covers in Rust. Hosts that
/// keep a reference implementation treat this as "fall back", not "fail".
#[error("unsupported by the rust path: {0}")]
Unsupported(&'static str),
}
impl Error {
pub const fn http_status_code(&self) -> Option<u16> {
match self {
Self::InvalidRequest(_) => Some(400),
Self::MissingDocumentUrl => Some(500),
Self::Http { status, .. } => Some(*status),
_ => None,
}
}
}
impl From<OcrRequestError> for Error {
fn from(error: OcrRequestError) -> Self {
match error {
OcrRequestError::MissingField(field) => Self::MissingField(field),
OcrRequestError::MissingDocumentUrl => Self::MissingDocumentUrl,
error => Self::InvalidRequest(error.to_string()),
}
}
}
impl From<OcrResponseError> for Error {
fn from(error: OcrResponseError) -> Self {
Self::InvalidResponse(error.to_string())
}
}
impl From<TransportError> for Error {
fn from(error: TransportError) -> Self {
match error {
TransportError::Http { status, body } => Self::Http { status, body },
TransportError::Network(message) => Self::Network(message),
TransportError::Connect(message) => Self::Connect(message),
}
}
}
impl From<litellm_auth::Error> for Error {
fn from(error: litellm_auth::Error) -> Self {
match error {
litellm_auth::Error::MissingApiKey { provider, .. } => Self::MissingApiKey { provider },
error => Self::Auth(error.to_string()),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum OcrRequestError {
@ -83,16 +183,16 @@ pub enum OcrError {
#[error("{0}")]
Polling(#[from] OcrPollingError),
#[error("{0}")]
Public(#[from] crate::Error),
Public(#[from] Error),
}
impl From<OcrError> for crate::Error {
impl From<OcrError> for Error {
fn from(error: OcrError) -> Self {
match error {
OcrError::Request(error) => error.into(),
OcrError::Response(error) => error.into(),
OcrError::Transport(error) => error.into(),
OcrError::Polling(error) => crate::Error::InvalidResponse(error.to_string()),
OcrError::Polling(error) => Error::InvalidResponse(error.to_string()),
OcrError::Public(error) => error,
}
}

2
litellm-rust/crates/core/src/ocr/handler.rs
View file

@ -3,8 +3,8 @@ use super::adapters::OcrAdapter;
use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest};
use super::registry::OcrAdapterKind;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::Error;
use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext};
use crate::ocr::Error;
use std::sync::Arc;
pub(crate) async fn perform_ocr_request(

3
litellm-rust/crates/core/src/ocr/hooks.rs
View file

@ -3,8 +3,8 @@ use std::pin::Pin;
use std::sync::Arc;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument};
use crate::Error;
use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use crate::ocr::Error;
use serde::Serialize;
use serde_json::Value;
@ -80,6 +80,7 @@ pub(crate) struct OcrLifecycleHooks {
impl CallLifecycleHooks<LiteLLMOcrRequest, LiteLLMOcrRequest, LiteLLMOcrResponse>
for OcrLifecycleHooks
{
type Error = crate::ocr::Error;
type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>;
type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>;
type SuccessFuture<'a> = OcrLogFuture<'a>;

19
litellm-rust/crates/core/src/ocr/lifecycle.rs
View file

@ -10,13 +10,13 @@ use super::hooks::{
OcrPreCallRequest,
};
use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
use crate::AuthError;
use crate::Error;
use crate::auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
use crate::call_lifecycle::host::{
HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase,
};
use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
use crate::ocr::Error;
use litellm_auth::Error as AuthError;
use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
pub type NativeResult<T> = Result<NativeOutcome<T>, Error>;
@ -84,7 +84,7 @@ impl OcrHostOperation {
pub enum OcrHostResult {
Request(Result<(Box<LiteLLMOcrRequest>, bool), Error>),
Lifecycle(Result<(), HostFailure>),
Lifecycle(Result<(), HostFailure<Error>>),
AzureAdToken(Result<ResolvedCredential, AuthError>),
PreCall(Result<OcrPreCallRequest, Error>),
DuringCall(Result<OcrDuringCallRequest, Error>),
@ -256,7 +256,7 @@ impl OcrCall {
Ok(self.host_step(operation))
}
fn accept(&mut self, result: Result<(), HostFailure>) {
fn accept(&mut self, result: Result<(), HostFailure<Error>>) {
let cancelled = matches!(&result, Err(HostFailure::Cancelled(_)));
if let Some(error) = self.lifecycle.accept(result) {
if cancelled {
@ -268,7 +268,7 @@ impl OcrCall {
}
}
pub async fn interrupt(&mut self, failure: HostFailure) -> Result<OcrCallStep, Error> {
pub async fn interrupt(&mut self, failure: HostFailure<Error>) -> Result<OcrCallStep, Error> {
if self.completed {
return Err(Error::InvalidRequest(
"OCR call cannot be interrupted after completion".into(),
@ -286,6 +286,7 @@ impl OcrCall {
}
impl HostCall for OcrCall {
type Error = crate::ocr::Error;
type Operation = OcrHostOperation;
type Result = OcrHostResult;
type Complete = LiteLLMOcrResponse;
@ -293,14 +294,14 @@ impl HostCall for OcrCall {
fn resume(
&mut self,
result: Option<Self::Result>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete> {
) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error> {
Box::pin(OcrCall::resume(self, result))
}
fn interrupt(
&mut self,
failure: HostFailure,
) -> HostCallFuture<'_, Self::Operation, Self::Complete> {
failure: HostFailure<Self::Error>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error> {
Box::pin(OcrCall::interrupt(self, failure))
}
}

1
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -3,6 +3,7 @@ pub mod client;
mod codecs;
mod document;
pub mod error;
pub use error::Error;
mod handler;
pub mod hooks;
mod lifecycle;

2
litellm-rust/crates/core/src/ocr/prepare.rs
View file

@ -119,7 +119,7 @@ pub(crate) fn build_http_request<B: Serialize>(
.timeout(request.connection.timeout);
crate::http_utils::with_headers(builder, headers, crate::http_utils::HeaderPolicy::All)
.build()
.map_err(crate::error::TransportError::from)
.map_err(crate::transport::Error::from)
.map_err(OcrError::from)
}

2
litellm-rust/crates/core/src/ocr/registry.rs
View file

@ -1,5 +1,5 @@
use super::adapters::OcrAdapter;
use crate::Error;
use crate::ocr::Error;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
macro_rules! define_adapter_types {

4
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -7,9 +7,9 @@ use serde_json::{Map, Value};
use super::hooks::{NoopOcrHooks, OcrHooks};
use super::registry::{OcrAdapterKind, resolve_wire_adapter};
use crate::Error;
use crate::auth::{InputSource, TokenProviderHandle};
use crate::constants::OCR_HTTP_TIMEOUT_SECS;
use crate::ocr::Error;
use litellm_auth::{InputSource, TokenProviderHandle};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]

4
litellm-rust/crates/core/src/ocr/wire.rs
View file

@ -4,8 +4,8 @@ use std::collections::BTreeMap;
use std::time::Duration;
use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument};
use crate::Error;
use crate::auth::InputSource;
use crate::ocr::Error;
use litellm_auth::InputSource;
use serde::{
Deserialize,
de::{DeserializeOwned, IntoDeserializer},

2
litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs
View file

@ -1,5 +1,5 @@
use super::*;
use crate::Error;
use crate::chat_completions::Error;
use serde_json::json;
fn messages(value: Value) -> Vec<ChatMessage> {

2
litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs
View file

@ -1,5 +1,6 @@
use serde_json::{Map, Value, json};
use crate::chat_completions::Error;
use crate::chat_completions::conversation::{Conversation, build_conversation};
use crate::chat_completions::transformation::{
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
@ -10,7 +11,6 @@ use crate::chat_completions::types::{
ProviderChatRequestData, ProviderChatResponseData,
};
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
use crate::error::Error;
use crate::providers::anthropic::messages::transformation::{
complete_anthropic_url, resolve_anthropic_api_key,
};

22
litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs
View file

@ -1,5 +1,4 @@
use crate::auth::error::MissingCredential;
use crate::error::Error;
use crate::messages::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
@ -18,11 +17,14 @@ pub fn non_empty(value: Option<&str>) -> Option<&str> {
pub fn resolve_anthropic_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
) -> Result<String, litellm_auth::Error> {
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(litellm_auth::Error::MissingApiKey {
provider: "Anthropic",
environment_variable: ANTHROPIC_API_KEY_ENV,
})
}
pub fn complete_anthropic_url(
@ -56,7 +58,7 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_anthropic_api_key(api_key, env_lookup)
resolve_anthropic_api_key(api_key, env_lookup).map_err(Error::from)
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
@ -114,10 +116,12 @@ mod tests {
resolve_anthropic_api_key(Some(" "), &with_env).unwrap(),
"sk-env"
);
assert!(matches!(
resolve_anthropic_api_key(None, &|_| None).expect_err("missing key"),
Error::Auth(_)
));
assert_eq!(
resolve_anthropic_api_key(None, &|_| None)
.expect_err("missing key")
.to_string(),
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
);
}
#[test]

7
litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs
View file

@ -1,7 +0,0 @@
mod credential_provider_cache;
mod native;
mod resolve;
mod types;
pub(crate) use resolve::AzureAuthService;
pub(crate) use types::AzureAuthInputs;

12
litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs
View file

@ -1,5 +1,4 @@
use crate::auth::error::MissingCredential;
use crate::error::Error;
use crate::messages::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use crate::messages::types::{
AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock,
@ -33,7 +32,12 @@ 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(litellm_auth::Error::MissingApiKey {
provider: "Azure",
environment_variable: AZURE_API_KEY_ENV,
})
})
}
pub fn complete_azure_anthropic_url(
@ -43,7 +47,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(litellm_auth::Error::MissingAzureApiBase))?;
let api_base = api_base.trim_end_matches('/');

1
litellm-rust/crates/core/src/providers/azure_ai/mod.rs
View file

@ -1,2 +1 @@
pub(crate) mod auth;
pub mod messages;

3
litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs
View file

@ -1,12 +1,13 @@
use serde_json::{Map, Value, json};
use crate::audio_transcription::Error;
use crate::audio_transcription::transformation::{
AudioTranscriptionAuth, AudioTranscriptionProviderConfig,
};
use crate::audio_transcription::types::{
AudioTranscriptionRequestData, AudioTranscriptionResponseData,
};
use crate::error::{Error, json_type_name};
use crate::http_utils::json_type_name;
pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region};
use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE};

931
litellm-rust/crates/core/src/providers/bedrock/aws_base.rs
View file

@ -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, &current_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, &params)
.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::*;

2
litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs
View file

@ -1,5 +1,5 @@
use super::*;
use crate::Error;
use crate::chat_completions::Error;
use serde_json::json;
fn messages(value: Value) -> Vec<ChatMessage> {

2
litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs
View file

@ -1,5 +1,6 @@
use serde_json::{Map, Value, json};
use crate::chat_completions::Error;
use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation};
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
use crate::chat_completions::transformation::{
@ -11,7 +12,6 @@ use crate::chat_completions::types::{
ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData,
ProviderChatResponseData,
};
use crate::error::Error;
use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region};
use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE};

44
litellm-rust/crates/core/src/providers/bedrock/constants.rs
View file

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

1
litellm-rust/crates/core/src/providers/bedrock/mod.rs
View file

@ -2,7 +2,6 @@
//! with Python's `BaseAWSLLM`; the broader core purity guidance is reconciled
//! separately.
#[cfg(feature = "bedrock-auth")]
pub mod audio_transcription;
pub mod aws_base;
pub mod chat_completions;

1
litellm-rust/crates/core/src/providers/mod.rs
View file

@ -1,6 +1,5 @@
pub mod anthropic;
pub mod azure_ai;
#[cfg(feature = "bedrock-auth")]
pub mod bedrock;
pub mod custom_llm_provider;
pub mod openai;

Some files were not shown because too many files have changed in this diff Show more