From 23633445d74c3e347eb61914ce2b48245d56adfe Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 18:12:11 +0000 Subject: [PATCH] refactor(rust): extract auth and cache crates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 78 +- litellm-rust/Cargo.toml | 9 +- litellm-rust/crates/auth-aws/Cargo.toml | 25 + litellm-rust/crates/auth-aws/src/aws.rs | 949 ++++++++++++++++++ litellm-rust/crates/auth-aws/src/constants.rs | 43 + litellm-rust/crates/auth-aws/src/error.rs | 46 + litellm-rust/crates/auth-aws/src/lib.rs | 6 + litellm-rust/crates/auth-azure/Cargo.toml | 21 + .../src}/credential_provider_cache.rs | 6 +- litellm-rust/crates/auth-azure/src/lib.rs | 7 + .../auth => auth-azure/src}/native.rs | 51 +- .../auth => auth-azure/src}/resolve.rs | 93 +- .../azure_ai/auth => auth-azure/src}/types.rs | 19 +- litellm-rust/crates/auth-gcp/Cargo.toml | 16 + .../auth/vertex.rs => auth-gcp/src/lib.rs} | 71 +- litellm-rust/crates/auth/Cargo.toml | 15 + .../{core/src/auth => auth/src}/credential.rs | 18 +- litellm-rust/crates/auth/src/error.rs | 120 +++ .../{core/src/auth => auth/src}/http.rs | 15 +- .../{core/src/auth/mod.rs => auth/src/lib.rs} | 7 +- .../{core/src/auth => auth/src}/policy.rs | 23 +- .../{core/src/auth => auth/src}/secret.rs | 0 .../{core/src/auth => auth/src}/token.rs | 6 +- litellm-rust/crates/cache-memory/Cargo.toml | 14 + litellm-rust/crates/cache-memory/src/cache.rs | 254 +++++ litellm-rust/crates/cache-memory/src/lib.rs | 3 + .../crates/cache-memory/tests/cache.rs | 140 +++ litellm-rust/crates/cache/Cargo.toml | 15 + litellm-rust/crates/cache/src/base_cache.rs | 98 ++ litellm-rust/crates/cache/src/caching.rs | 166 +++ litellm-rust/crates/cache/src/error.rs | 7 + litellm-rust/crates/cache/src/lib.rs | 12 + litellm-rust/crates/cache/tests/caching.rs | 139 +++ litellm-rust/crates/core/Cargo.toml | 24 +- .../core/src/audio_transcription/error.rs | 26 + .../core/src/audio_transcription/handler.rs | 26 +- .../core/src/audio_transcription/mod.rs | 3 +- .../core/src/audio_transcription/prepare.rs | 5 +- .../src/audio_transcription/transformation.rs | 2 +- .../core/src/audio_transcription/types.rs | 1 - litellm-rust/crates/core/src/auth/error.rs | 128 --- .../core/src/caching/in_memory_cache.rs | 258 ----- litellm-rust/crates/core/src/caching/mod.rs | 1 - .../crates/core/src/call_lifecycle/host.rs | 19 +- .../crates/core/src/call_lifecycle/mod.rs | 44 +- .../core/src/chat_completions/common_utils.rs | 5 +- .../crates/core/src/chat_completions/error.rs | 26 + .../core/src/chat_completions/handler.rs | 29 +- .../crates/core/src/chat_completions/mod.rs | 3 +- .../core/src/chat_completions/prepare.rs | 2 +- .../crates/core/src/chat_completions/tests.rs | 30 +- .../src/chat_completions/transformation.rs | 2 +- .../crates/core/src/chat_completions/types.rs | 1 - litellm-rust/crates/core/src/error.rs | 229 +---- litellm-rust/crates/core/src/http_utils.rs | 46 +- litellm-rust/crates/core/src/lib.rs | 4 +- litellm-rust/crates/core/src/media.rs | 81 +- .../crates/core/src/messages/common_utils.rs | 4 +- .../crates/core/src/messages/error.rs | 17 + .../crates/core/src/messages/handler.rs | 18 +- litellm-rust/crates/core/src/messages/mod.rs | 3 +- .../crates/core/src/messages/prepare.rs | 2 +- .../crates/core/src/messages/tests.rs | 16 +- .../core/src/messages/transformation.rs | 2 +- .../core/src/ocr/adapters/azure/cohere.rs | 4 +- .../azure/document_intelligence/mod.rs | 6 +- .../azure/document_intelligence/polling.rs | 2 +- .../core/src/ocr/adapters/azure/mistral.rs | 6 +- .../crates/core/src/ocr/adapters/azure/mod.rs | 15 +- .../crates/core/src/ocr/adapters/cohere.rs | 2 +- .../crates/core/src/ocr/adapters/mistral.rs | 2 +- .../core/src/ocr/adapters/reducto/mod.rs | 4 +- .../core/src/ocr/adapters/vertex/deepseek.rs | 4 +- .../core/src/ocr/adapters/vertex/mistral.rs | 4 +- .../core/src/ocr/adapters/vertex/mod.rs | 11 +- litellm-rust/crates/core/src/ocr/client.rs | 11 +- litellm-rust/crates/core/src/ocr/document.rs | 3 +- litellm-rust/crates/core/src/ocr/error.rs | 108 +- litellm-rust/crates/core/src/ocr/handler.rs | 2 +- litellm-rust/crates/core/src/ocr/hooks.rs | 3 +- litellm-rust/crates/core/src/ocr/lifecycle.rs | 19 +- litellm-rust/crates/core/src/ocr/mod.rs | 1 + litellm-rust/crates/core/src/ocr/prepare.rs | 2 +- litellm-rust/crates/core/src/ocr/registry.rs | 2 +- litellm-rust/crates/core/src/ocr/types.rs | 4 +- litellm-rust/crates/core/src/ocr/wire.rs | 4 +- .../anthropic/chat_completions/tests.rs | 2 +- .../chat_completions/transformation.rs | 2 +- .../anthropic/messages/transformation.rs | 22 +- .../core/src/providers/azure_ai/auth/mod.rs | 7 - .../azure_ai/messages/transformation.rs | 12 +- .../crates/core/src/providers/azure_ai/mod.rs | 1 - .../providers/bedrock/audio_transcription.rs | 3 +- .../core/src/providers/bedrock/aws_base.rs | 931 +---------------- .../bedrock/chat_completions/tests.rs | 2 +- .../chat_completions/transformation.rs | 2 +- .../core/src/providers/bedrock/constants.rs | 44 +- .../crates/core/src/providers/bedrock/mod.rs | 1 - litellm-rust/crates/core/src/providers/mod.rs | 1 - .../openai/responses/transformation.rs | 2 +- .../crates/core/src/responses/error.rs | 17 + .../core/src/responses/instrumentation.rs | 3 +- litellm-rust/crates/core/src/responses/mod.rs | 2 + .../crates/core/src/responses/websocket.rs | 45 +- .../crates/core/src/transport/error.rs | 75 ++ litellm-rust/crates/core/src/transport/mod.rs | 2 + litellm-rust/crates/core/src/url_utils.rs | 3 +- .../crates/core/tests/host_lifecycle.rs | 6 +- litellm-rust/crates/core/tests/ocr.rs | 35 +- litellm-rust/crates/core/tests/ocr/support.rs | 2 +- .../core/tests/vertex_ai_deepseek_ocr.rs | 2 +- .../crates/core/tests/vertex_ai_ocr.rs | 2 +- litellm-rust/crates/python-bridge/Cargo.toml | 3 +- litellm-rust/crates/python-bridge/src/auth.rs | 2 +- .../crates/python-bridge/src/errors.rs | 157 ++- .../crates/python-bridge/src/execution.rs | 2 +- litellm-rust/crates/python-bridge/src/lib.rs | 13 +- .../crates/python-bridge/src/lifecycle/mod.rs | 30 +- .../src/lifecycle/preparation.rs | 2 +- .../crates/python-bridge/src/marshal.rs | 2 +- .../src/routes/audio_transcription/value.rs | 6 +- .../src/routes/chat_completions/value.rs | 2 +- .../python-bridge/src/routes/definition.rs | 2 +- .../src/routes/messages/value.rs | 6 +- .../python-bridge/src/routes/ocr/errors.rs | 4 +- .../python-bridge/src/routes/ocr/lifecycle.rs | 8 +- .../python-bridge/src/routes/ocr/project.rs | 2 +- .../python-bridge/src/routes/ocr/value.rs | 2 +- 128 files changed, 3090 insertions(+), 2147 deletions(-) create mode 100644 litellm-rust/crates/auth-aws/Cargo.toml create mode 100644 litellm-rust/crates/auth-aws/src/aws.rs create mode 100644 litellm-rust/crates/auth-aws/src/constants.rs create mode 100644 litellm-rust/crates/auth-aws/src/error.rs create mode 100644 litellm-rust/crates/auth-aws/src/lib.rs create mode 100644 litellm-rust/crates/auth-azure/Cargo.toml rename litellm-rust/crates/{core/src/providers/azure_ai/auth => auth-azure/src}/credential_provider_cache.rs (87%) create mode 100644 litellm-rust/crates/auth-azure/src/lib.rs rename litellm-rust/crates/{core/src/providers/azure_ai/auth => auth-azure/src}/native.rs (94%) rename litellm-rust/crates/{core/src/providers/azure_ai/auth => auth-azure/src}/resolve.rs (89%) rename litellm-rust/crates/{core/src/providers/azure_ai/auth => auth-azure/src}/types.rs (93%) create mode 100644 litellm-rust/crates/auth-gcp/Cargo.toml rename litellm-rust/crates/{core/src/auth/vertex.rs => auth-gcp/src/lib.rs} (90%) create mode 100644 litellm-rust/crates/auth/Cargo.toml rename litellm-rust/crates/{core/src/auth => auth/src}/credential.rs (91%) create mode 100644 litellm-rust/crates/auth/src/error.rs rename litellm-rust/crates/{core/src/auth => auth/src}/http.rs (84%) rename litellm-rust/crates/{core/src/auth/mod.rs => auth/src/lib.rs} (94%) rename litellm-rust/crates/{core/src/auth => auth/src}/policy.rs (82%) rename litellm-rust/crates/{core/src/auth => auth/src}/secret.rs (100%) rename litellm-rust/crates/{core/src/auth => auth/src}/token.rs (83%) create mode 100644 litellm-rust/crates/cache-memory/Cargo.toml create mode 100644 litellm-rust/crates/cache-memory/src/cache.rs create mode 100644 litellm-rust/crates/cache-memory/src/lib.rs create mode 100644 litellm-rust/crates/cache-memory/tests/cache.rs create mode 100644 litellm-rust/crates/cache/Cargo.toml create mode 100644 litellm-rust/crates/cache/src/base_cache.rs create mode 100644 litellm-rust/crates/cache/src/caching.rs create mode 100644 litellm-rust/crates/cache/src/error.rs create mode 100644 litellm-rust/crates/cache/src/lib.rs create mode 100644 litellm-rust/crates/cache/tests/caching.rs create mode 100644 litellm-rust/crates/core/src/audio_transcription/error.rs delete mode 100644 litellm-rust/crates/core/src/auth/error.rs delete mode 100644 litellm-rust/crates/core/src/caching/in_memory_cache.rs delete mode 100644 litellm-rust/crates/core/src/caching/mod.rs create mode 100644 litellm-rust/crates/core/src/chat_completions/error.rs create mode 100644 litellm-rust/crates/core/src/messages/error.rs delete mode 100644 litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs create mode 100644 litellm-rust/crates/core/src/responses/error.rs create mode 100644 litellm-rust/crates/core/src/transport/error.rs create mode 100644 litellm-rust/crates/core/src/transport/mod.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 500a12035cb..e2a3af77594 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 2a0a7687e0e..879090870d8 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml new file mode 100644 index 00000000000..d998b647960 --- /dev/null +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs new file mode 100644 index 00000000000..3b6b73bc6a9 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -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> = OnceLock::new(); +static AMBIENT_CREDENTIALS_CACHE: OnceLock> = OnceLock::new(); + +fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { + 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, + pub secret_access_key: Option, + pub session_token: Option, + pub region_name: Option, + pub session_name: Option, + pub profile_name: Option, + pub role_name: Option, + pub web_identity_token: Option, + pub sts_endpoint: Option, + pub external_id: Option, +} + +impl AwsAuthConfig { + fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option + 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, + }, + 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 { + STATIC_CREDENTIALS_CACHE.get_or_init(|| { + Cache::builder() + .max_capacity(200) + .time_to_live(STATIC_CREDENTIALS_TTL) + .build() + }) +} + +fn ambient_credentials_cache() -> &'static Cache { + 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 { + 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 + 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 + Sync), +) -> Result { + 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 { + if role_identity(role).is_none() { + return Ok(false); + } + if let (Ok(current_role), Ok(token_file)) = ( + std::env::var(AWS_ROLE_ARN), + std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), + ) && !token_file.is_empty() + { + return Ok(same_role_arns(role, ¤t_role)); + } + + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = config.region_name.clone() { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = config.sts_endpoint.clone() { + loader = loader.endpoint_url(endpoint); + } + let sdk_config = loader.load().await; + let response = match aws_sdk_sts::Client::new(&sdk_config) + .get_caller_identity() + .send() + .await + { + Ok(response) => response, + Err(_) => return Ok(false), + }; + Ok(response + .arn() + .is_some_and(|caller| same_role_arns(role, caller))) +} + +fn default_session_name() -> String { + let seconds = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()); + format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}") +} + +/// The subset of `headers` SigV4 should cover. +/// +/// Python signs only these and reattaches the rest afterwards, so a forwarded +/// client header cannot change the canonical request and invalidate the +/// signature. Signing everything instead makes the request 403 on a header the +/// caller supplied, on a deployment that works on the Python path. +pub fn aws_signature_headers(headers: &BTreeMap) -> BTreeMap { + 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, + region: &str, + credentials: &Credentials, + signing_time: SystemTime, +) -> Result, Error> { + let identity: Identity = credentials.clone().into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(region) + .name(BEDROCK_SERVICE) + .time(signing_time) + .settings(SigningSettings::default()) + .build() + .map(SigningParams::from) + .map_err(|error| Error::AwsSigningParameters(error.to_string()))?; + let header_refs = headers + .iter() + .map(|(name, value)| (name.as_str(), value.as_str())); + let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) + .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; + let (instructions, _) = sign(request, ¶ms) + .map_err(|error| Error::AwsSigning(error.to_string()))? + .into_parts(); + Ok(instructions + .headers() + .map(|(name, value)| { + let normalized_name = match name { + "authorization" => "Authorization", + "x-amz-date" => "X-Amz-Date", + "x-amz-security-token" => "X-Amz-Security-Token", + _ => name, + }; + (normalized_name.to_string(), value.to_string()) + }) + .collect()) +} + +/// Model-id and region parsing shared by every Bedrock route. +pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { + 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, + env_lookup: &dyn Fn(&str) -> Option, +) -> 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, + env_lookup: &dyn Fn(&str) -> Option, +) -> 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) -> Option { + 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 { + None + } + + fn parity_inputs() -> (String, Vec, BTreeMap) { + ( + "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> { + 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("; ") + ); + } +} diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs new file mode 100644 index 00000000000..be215cc9016 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -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"; diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs new file mode 100644 index 00000000000..f80fbce456e --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -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 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() + ) + ); + } +} diff --git a/litellm-rust/crates/auth-aws/src/lib.rs b/litellm-rust/crates/auth-aws/src/lib.rs new file mode 100644 index 00000000000..264592ccb2e --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/lib.rs @@ -0,0 +1,6 @@ +mod aws; +pub mod constants; +mod error; + +pub use aws::*; +pub use error::Error; diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml new file mode 100644 index 00000000000..9f8260c7b3f --- /dev/null +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs b/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs similarity index 87% rename from litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs rename to litellm-rust/crates/auth-azure/src/credential_provider_cache.rs index 297e4cc6502..ab9ffc719df 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs +++ b/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs @@ -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, AuthError> + ) -> Result, Error> where - F: Future, AuthError>>, + F: Future, Error>>, { self.entries .try_get_with(key, create) diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs new file mode 100644 index 00000000000..e76227d6aa2 --- /dev/null +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -0,0 +1,7 @@ +mod credential_provider_cache; +mod native; +mod resolve; +mod types; + +pub use resolve::AzureAuthService; +pub use types::AzureAuthInputs; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs b/litellm-rust/crates/auth-azure/src/native.rs similarity index 94% rename from litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs rename to litellm-rust/crates/auth-azure/src/native.rs index b8f19818d16..5f913a8ad01 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -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 { + pub(crate) fn new(request: NativeAzureRequest) -> Result { 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 { + ) -> Result { 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 { +fn validate_sources(request: &NativeAzureRequest) -> Result { match request { NativeAzureRequest::ClientSecret { tenant_id, @@ -356,7 +352,7 @@ fn is_request_controlled(value: &Sourced, optional: Option<&Sourced Result { +fn trusted_only(sources: &[InputSource]) -> Result { if sources.contains(&InputSource::Request) { return mixed_sources(); } @@ -371,16 +367,14 @@ fn trusted_source(sources: &[InputSource]) -> InputSource { } } -fn mixed_sources() -> Result { - Err(AuthError::Configuration( - AuthConfigurationError::MixedAzureCredentialSources, - )) +fn mixed_sources() -> Result { + Err(Error::MixedAzureCredentialSources) } fn build_credential( request: NativeAzureRequest, transport: Option, -) -> Result, AuthError> { +) -> Result, Error> { match request { NativeAzureRequest::ClientSecret { tenant_id, @@ -439,11 +433,7 @@ fn build_credential( NativeAzureRequest::DeveloperTools { .. } => DeveloperToolsCredential::new(None) .map(|credential| credential as Arc), } - .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(value: T) -> Sourced { 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)); } } } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs similarity index 89% rename from litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs rename to litellm-rust/crates/auth-azure/src/resolve.rs index 025dd4f8740..660a95b79d8 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -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, } @@ -45,14 +44,14 @@ trait AzureTokenAcquirer: Send + Sync { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>>; + ) -> Pin> + Send + '_>>; } impl AzureTokenAcquirer for NativeAzureTokenAcquirer { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>> { + ) -> Pin> + 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 + Sync), - ) -> Result>, AuthError> { + ) -> Result>, 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, -) -> Result { +) -> Result { 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, authority: Option>, refresh_source: InputSource, -) -> Result { +) -> Result { 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>, scope: Sourced, authority: Option>, -) -> Result { +) -> Result { 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 + Sync), reference: &CredentialRef, -) -> Result, AuthError> { +) -> Result, 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>, -) -> Result>, AuthError> { +) -> Result>, 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 { +fn non_empty_reference(value: &str, kind: &str) -> Result { 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> + Send + '_>, - > { + ) -> std::pin::Pin> + 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)); } } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs b/litellm-rust/crates/auth-azure/src/types.rs similarity index 93% rename from litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs rename to litellm-rust/crates/auth-azure/src/types.rs index f15d526d945..2a510de1f43 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -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) -> Result { + pub fn from_optional_params(params: &Map) -> Result { Self::from_sourced_optional_params(params, &BTreeMap::new()) } pub fn from_sourced_optional_params( params: &Map, sources: &BTreeMap, - ) -> Result { + ) -> Result { 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, sources: &BTreeMap, name: &str, -) -> Result, AuthError> { +) -> Result, 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, sources: &BTreeMap, name: &str, -) -> Result, AuthError> { +) -> Result, 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() { diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml new file mode 100644 index 00000000000..f24582db13e --- /dev/null +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/core/src/auth/vertex.rs b/litellm-rust/crates/auth-gcp/src/lib.rs similarity index 90% rename from litellm-rust/crates/core/src/auth/vertex.rs rename to litellm-rust/crates/auth-gcp/src/lib.rs index e13b57abdfe..f8402624edc 100644 --- a/litellm-rust/crates/core/src/auth/vertex.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -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>, project_id: Option, location: Option, } impl VertexConfig { - pub(crate) fn from_sourced_optional_params( + pub fn from_sourced_optional_params( params: &Map, sources: &BTreeMap, - ) -> Result { + ) -> Result { 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, ) -> Option { @@ -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, ) -> Option { @@ -87,7 +86,7 @@ pub(crate) fn get_vertex_ai_location( } #[derive(Clone)] -pub(crate) struct VertexAuth { +pub struct VertexAuth { providers: Cache>, loader: Arc, } @@ -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 + Sync), - ) -> Result { + ) -> Result { 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 + Sync), - ) -> Result { + ) -> Result { 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 + Sync), - ) -> Result, AuthError> { + ) -> Result, 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>; } -type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; +type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; struct GcpTokenSource(Arc); @@ -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::(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, sources: &BTreeMap, names: &[&str], -) -> Result>, AuthError> { +) -> Result>, 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, name: &str) -> InputSourc sources.get(name).copied().unwrap_or_default() } -fn optional_string( - params: &Map, - names: &[&str], -) -> Result, AuthError> { +fn optional_string(params: &Map, names: &[&str]) -> Result, 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, 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) )); } diff --git a/litellm-rust/crates/auth/Cargo.toml b/litellm-rust/crates/auth/Cargo.toml new file mode 100644 index 00000000000..128a05c1a25 --- /dev/null +++ b/litellm-rust/crates/auth/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/core/src/auth/credential.rs b/litellm-rust/crates/auth/src/credential.rs similarity index 91% rename from litellm-rust/crates/core/src/auth/credential.rs rename to litellm-rust/crates/auth/src/credential.rs index c64d331b877..6721eb67a35 100644 --- a/litellm-rust/crates/core/src/auth/credential.rs +++ b/litellm-rust/crates/auth/src/credential.rs @@ -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> + Send + 'a>>; + Pin> + 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 { + pub async fn resolve(&self, reference: &CredentialRef) -> Result { self.0.resolve(reference).await } } @@ -84,7 +84,7 @@ impl CredentialPlan { pub async fn resolve( &self, resolver: &CredentialResolverHandle, - ) -> Result { + ) -> Result { 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); } } diff --git a/litellm-rust/crates/auth/src/error.rs b/litellm-rust/crates/auth/src/error.rs new file mode 100644 index 00000000000..914265ffb32 --- /dev/null +++ b/litellm-rust/crates/auth/src/error.rs @@ -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::>().join("; "))] + CredentialChain(Vec), + #[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://.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" + ); + } +} diff --git a/litellm-rust/crates/core/src/auth/http.rs b/litellm-rust/crates/auth/src/http.rs similarity index 84% rename from litellm-rust/crates/core/src/auth/http.rs rename to litellm-rust/crates/auth/src/http.rs index 83931311550..7d20991d838 100644 --- a/litellm-rust/crates/core/src/auth/http.rs +++ b/litellm-rust/crates/auth/src/http.rs @@ -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, AuthError> { +) -> Result, 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}"), diff --git a/litellm-rust/crates/core/src/auth/mod.rs b/litellm-rust/crates/auth/src/lib.rs similarity index 94% rename from litellm-rust/crates/core/src/auth/mod.rs rename to litellm-rust/crates/auth/src/lib.rs index 2940a983fb9..7a24d2acf70 100644 --- a/litellm-rust/crates/core/src/auth/mod.rs +++ b/litellm-rust/crates/auth/src/lib.rs @@ -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; diff --git a/litellm-rust/crates/core/src/auth/policy.rs b/litellm-rust/crates/auth/src/policy.rs similarity index 82% rename from litellm-rust/crates/core/src/auth/policy.rs rename to litellm-rust/crates/auth/src/policy.rs index b796dedf0d8..4a1f5eeecf9 100644 --- a/litellm-rust/crates/core/src/auth/policy.rs +++ b/litellm-rust/crates/auth/src/policy.rs @@ -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, AuthError> { + ) -> Result, 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, diff --git a/litellm-rust/crates/core/src/auth/secret.rs b/litellm-rust/crates/auth/src/secret.rs similarity index 100% rename from litellm-rust/crates/core/src/auth/secret.rs rename to litellm-rust/crates/auth/src/secret.rs diff --git a/litellm-rust/crates/core/src/auth/token.rs b/litellm-rust/crates/auth/src/token.rs similarity index 83% rename from litellm-rust/crates/core/src/auth/token.rs rename to litellm-rust/crates/auth/src/token.rs index cfc6b8f0d6b..94da5f259fb 100644 --- a/litellm-rust/crates/core/src/auth/token.rs +++ b/litellm-rust/crates/auth/src/token.rs @@ -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> + Send + 'a>>; + Pin> + 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 { + pub async fn acquire(&self) -> Result { self.0.acquire().await } } diff --git a/litellm-rust/crates/cache-memory/Cargo.toml b/litellm-rust/crates/cache-memory/Cargo.toml new file mode 100644 index 00000000000..d4487573a9a --- /dev/null +++ b/litellm-rust/crates/cache-memory/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs new file mode 100644 index 00000000000..3b62c80b125 --- /dev/null +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -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 = Arc Result + Send + Sync>; +type ValueValidator = Arc Result<(), Error> + Send + Sync>; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CacheWrite { + Stored, + Disabled, + TooLarge, +} + +struct CacheState { + values: HashMap, + expirations: HashMap, + expiration_heap: BinaryHeap>, +} + +pub struct InMemoryCache { + state: Mutex>, + max_size_in_memory: usize, + default_ttl: Duration, + max_entry_bytes: Option, + measure_value: Option>, + validate_value: Option>, + now: Arc Duration + Send + Sync>, +} + +impl Default for InMemoryCache { + fn default() -> Self { + Self::new(None, None) + } +} + +impl InMemoryCache { + pub fn new(max_size_in_memory: Option, default_ttl: Option) -> 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, + default_ttl: Option, + 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, + default_ttl: Option, + max_entry_bytes: Option, + measure_value: Option>, + 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, + value: V, + ttl: Option, + ) -> Result { + 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, 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, 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, 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, key: &str) { + state.values.remove(key); + state.expirations.remove(key); + } +} + +impl InMemoryCache { + 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 { + 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, 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, + }) + }) + } +} diff --git a/litellm-rust/crates/cache-memory/src/lib.rs b/litellm-rust/crates/cache-memory/src/lib.rs new file mode 100644 index 00000000000..c5b7fb6cb54 --- /dev/null +++ b/litellm-rust/crates/cache-memory/src/lib.rs @@ -0,0 +1,3 @@ +mod cache; + +pub use cache::{CacheWrite, InMemoryCache}; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs new file mode 100644 index 00000000000..423b8c959f2 --- /dev/null +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -0,0 +1,140 @@ +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 { + Arc::new(AtomicU64::new(100)) +} + +fn cache(clock: Arc, capacity: usize) -> InMemoryCache { + 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) { + 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 capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc) { + 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::::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::::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::::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" + }) + ); +} diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml new file mode 100644 index 00000000000..a14c4294aa0 --- /dev/null +++ b/litellm-rust/crates/cache/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs new file mode 100644 index 00000000000..2ba8ff92ebd --- /dev/null +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -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> + Send + 'a>>; + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct CacheKwargs { + pub ttl: Option, + pub extras: Map, +} + +#[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, +} + +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, 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> { + 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>; +} diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs new file mode 100644 index 00000000000..1aab6ee8e91 --- /dev/null +++ b/litellm-rust/crates/cache/src/caching.rs @@ -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, + pub api_parameter: bool, + pub internal_parameter: bool, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct CacheKeyInput { + pub fields: Vec, + pub preset: Option, + pub namespace: Option, + pub include_provider_parameters: bool, +} + +#[derive(Default)] +pub struct CacheKeyContext { + pub model_group: Option, + pub caching_groups: Vec<(Vec, String)>, + pub file_checksum: Option, + pub file_object_name: Option, + pub metadata_file_name: Option, + pub parameters_file_name: Option, +} + +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, + 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) -> 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, + key: &str, + kwargs: &CacheKwargs, +) -> Result, Error> { + cache.get_cache(key, kwargs) +} + +pub fn set_cache( + cache: &dyn BaseCache, + key: &str, + entry: CacheEntry, + kwargs: CacheKwargs, +) -> Result<(), Error> { + cache.set_cache(key, entry, kwargs) +} + +pub type CacheBackend = Arc>; diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs new file mode 100644 index 00000000000..d447c80f62d --- /dev/null +++ b/litellm-rust/crates/cache/src/error.rs @@ -0,0 +1,7 @@ +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum Error { + #[error("cache is unavailable")] + Unavailable, + #[error("invalid cache entry")] + InvalidEntry, +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs new file mode 100644 index 00000000000..d0fe3de15cd --- /dev/null +++ b/litellm-rust/crates/cache/src/lib.rs @@ -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; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs new file mode 100644 index 00000000000..1192fc9a2b0 --- /dev/null +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -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, 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() + ); +} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 2b3b4c6a03f..ededfeef8af 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/core/src/audio_transcription/error.rs b/litellm-rust/crates/core/src/audio_transcription/error.rs new file mode 100644 index 00000000000..f9ffb12d349 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/error.rs @@ -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), +} diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index db665cf1faf..bd1740a8b93 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -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, 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()), - } -} diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index b71748082bf..87f6c41d80f 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -1,4 +1,5 @@ -use crate::Error; +mod error; +pub use error::Error; mod client; mod handler; mod prepare; diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 9f45b8804be..82f85ba85ce 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -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, }) diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs index 16a28fbcac0..a849f052e12 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -1,4 +1,4 @@ -use crate::Error; +use super::Error; use serde_json::{Map, Value}; use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData}; diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index 559d7837027..1f90f61c0da 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -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, pub(super) timeout: Option, } diff --git a/litellm-rust/crates/core/src/auth/error.rs b/litellm-rust/crates/core/src/auth/error.rs deleted file mode 100644 index e7027c0df10..00000000000 --- a/litellm-rust/crates/core/src/auth/error.rs +++ /dev/null @@ -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::>().join("; "))] - CredentialChain(Vec), - #[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://.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, -} diff --git a/litellm-rust/crates/core/src/caching/in_memory_cache.rs b/litellm-rust/crates/core/src/caching/in_memory_cache.rs deleted file mode 100644 index 45d4bd69b79..00000000000 --- a/litellm-rust/crates/core/src/caching/in_memory_cache.rs +++ /dev/null @@ -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 { - pub cache_dict: HashMap, - pub ttl_dict: HashMap, - pub expiration_heap: BinaryHeap>, - pub max_size_in_memory: usize, - pub default_ttl: Duration, - now: Box Duration + Send + Sync>, -} - -impl Default for InMemoryCache { - fn default() -> Self { - Self::new(None, None) - } -} - -impl InMemoryCache { - pub fn new(max_size_in_memory: Option, default_ttl: Option) -> 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, - default_ttl: Option, - 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, value: V, ttl: Option) { - 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 { - 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 { - 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, max_size: usize, default_ttl: Duration) -> InMemoryCache { - 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()); - } -} diff --git a/litellm-rust/crates/core/src/caching/mod.rs b/litellm-rust/crates/core/src/caching/mod.rs deleted file mode 100644 index 5fb8a0e5174..00000000000 --- a/litellm-rust/crates/core/src/caching/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod in_memory_cache; diff --git a/litellm-rust/crates/core/src/call_lifecycle/host.rs b/litellm-rust/crates/core/src/call_lifecycle/host.rs index ac6ddf99b9e..97eb9c4c650 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/host.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/host.rs @@ -6,10 +6,11 @@ pub enum HostCallStep { Complete(C), } -pub type HostCallFuture<'a, O, C> = - Pin, crate::Error>> + Send + 'a>>; +pub type HostCallFuture<'a, O, C, E> = + Pin, 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, - ) -> 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, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>; } pub enum HostStep { @@ -48,9 +49,9 @@ pub enum HostPhase { } #[derive(Clone, Debug)] -pub enum HostFailure { - Error(crate::Error), - Cancelled(crate::Error), +pub enum HostFailure { + Error(E), + Cancelled(E), } pub struct HostLifecycle { @@ -70,7 +71,7 @@ impl HostLifecycle { self.phase } - pub fn accept(&mut self, result: Result<(), HostFailure>) -> Option { + pub fn accept(&mut self, result: Result<(), HostFailure>) -> Option { if let Err(failure) = result { if self.phase == HostPhase::DeploymentFailure { self.phase = HostPhase::Failure; diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index 5c752a73899..dce240c3d2b 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -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: Send + Sync { - type PreCallFuture<'a>: Future> + Send + 'a + type Error: Send + Sync; + type PreCallFuture<'a>: Future> + Send + 'a where Self: 'a, InitialReq: 'a, ProviderReq: 'a, Resp: 'a; - type DuringCallFuture<'a>: Future> + Send + 'a + type DuringCallFuture<'a>: Future> + Send + 'a where Self: 'a, InitialReq: 'a, @@ -60,7 +59,7 @@ pub trait CallLifecycleHooks: 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 + ) -> Result where InitialReq: CallLifecycleRequest, Hooks: CallLifecycleHooks, ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, + ProviderFuture: Future>, { 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 + ) -> Result where Hooks: CallLifecycleHooks, ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, + ProviderFuture: Future>, { 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, ) where @@ -255,8 +254,9 @@ mod tests { } impl CallLifecycleHooks for RecordingHooks { - type PreCallFuture<'a> = BoxFuture<'a, Result>; - type DuringCallFuture<'a> = BoxFuture<'a, Result>; + type Error = crate::messages::Error; + type PreCallFuture<'a> = BoxFuture<'a, Result>; + type DuringCallFuture<'a> = BoxFuture<'a, Result>; 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 for RecordingHooks { - type PreCallFuture<'a> = BoxFuture<'a, Result>; - type DuringCallFuture<'a> = BoxFuture<'a, Result>; + type Error = crate::messages::Error; + type PreCallFuture<'a> = BoxFuture<'a, Result>; + type DuringCallFuture<'a> = BoxFuture<'a, Result>; 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::(Error::Network("provider down".to_string())) + Err::(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"]); } diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index ca51471eb7c..9ebc5ae0efa 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -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>, ) -> Result, Error> { - shared_string_headers(HEADER_CONTEXT, extra_headers) + shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from) } diff --git a/litellm-rust/crates/core/src/chat_completions/error.rs b/litellm-rust/crates/core/src/chat_completions/error.rs new file mode 100644 index 00000000000..f9ffb12d349 --- /dev/null +++ b/litellm-rust/crates/core/src/chat_completions/error.rs @@ -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), +} diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 799dd3ceeb0..d4527e99a10 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -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, Error> { - match &request.auth { - ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported( - "AWS SigV4 requires the bedrock-auth feature", - )), - _ => Ok(request.upstream_headers.clone()), - } -} diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 8117663a23b..401eef609f2 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -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; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index e965a0fe5d7..e8d8d70f271 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -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}; diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index f8594dee447..39fabe27f44 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -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, .. }) )); } } diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index d7b9704c46c..1000dbaa673 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -1,4 +1,4 @@ -use crate::Error; +use super::Error; use serde_json::{Map, Value}; use super::types::{ diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 3238d09b6b5..7178d594870 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -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, pub(super) timeout: Option, } diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 359ad56c336..15d27602052 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -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 { - 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 for TransportError { - fn from(error: reqwest::Error) -> Self { - Self::Network(error.without_url().to_string()) - } -} - -impl From 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 for Error { - fn from(error: crate::ocr::error::OcrResponseError) -> Self { - Self::InvalidResponse(error.to_string()) - } -} - -impl From 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 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), } diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 7b9d6f6a648..53d2f961bd5 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -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 { + 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>, -) -> Result, Error> { +) -> Result, 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 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" + } ); } diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index f1d888ec12c..b028b7bc9b1 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -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; diff --git a/litellm-rust/crates/core/src/media.rs b/litellm-rust/crates/core/src/media.rs index 5f9a43794c2..ba26f431e57 100644 --- a/litellm-rust/crates/core/src/media.rs +++ b/litellm-rust/crates/core/src/media.rs @@ -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 { + ) -> Result { 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 { + ) -> Result { 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::() { - 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) )); } } diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 8dfdb2e361a..cbaf92b4986 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -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>, ) -> Result, Error> { - shared_string_headers(HEADER_CONTEXT, extra_headers) + shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from) } diff --git a/litellm-rust/crates/core/src/messages/error.rs b/litellm-rust/crates/core/src/messages/error.rs new file mode 100644 index 00000000000..8bea035f0b0 --- /dev/null +++ b/litellm-rust/crates/core/src/messages/error.rs @@ -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), +} diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 4d2627c6383..d7d593f2d57 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -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) } diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 0083272bcb4..156f42056f1 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -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; diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 3c2f6b9d372..b10e03ea9c0 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -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}; diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index df9f7051011..f454effd7b5 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -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] diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index 673a5728aca..2719e62d280 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -1,5 +1,5 @@ +use super::Error; use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse}; -use crate::Error; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MessagesAuthStrategy { diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs index 4c8455a171c..3691e9e1809 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs @@ -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"; diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs index 22034ddc3f6..eba300908f1 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs @@ -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; diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs index 6ed1e4441d4..87378dccdb7 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs @@ -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) diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs index 8639590b05c..28e09cdc80f 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs @@ -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"; diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs index 3d30ae6d6bd..0b2fcb0f4cb 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs @@ -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(()) } diff --git a/litellm-rust/crates/core/src/ocr/adapters/cohere.rs b/litellm-rust/crates/core/src/ocr/adapters/cohere.rs index 933ead7f7f7..d1faeeb7b1d 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/cohere.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/cohere.rs @@ -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, diff --git a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs index cdbc2c3effc..c379462c089 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs @@ -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}; diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs index 2dafe291674..40cefa05373 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs @@ -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) diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs index d16b3e7f386..fc24dbe489c 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs @@ -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"; diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs index 88c61725cee..3a1abf47ddf 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs @@ -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)] diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs index 270c41e647d..798510e7405 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs @@ -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(()) } diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 041c5c6fd21..00bfeb2b7b2 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -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)] diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index 82a32ac1ab5..a7afdaf8793 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -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], diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 55ea2cbcdae..1c21edb6c91 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -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 { + match self { + Self::InvalidRequest(_) => Some(400), + Self::MissingDocumentUrl => Some(500), + Self::Http { status, .. } => Some(*status), + _ => None, + } + } +} + +impl From 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 for Error { + fn from(error: OcrResponseError) -> Self { + Self::InvalidResponse(error.to_string()) + } +} + +impl From 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 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 for crate::Error { +impl From 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, } } diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index cd1d538aaa8..1ec02f3b622 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -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( diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs index 4d4a387e553..1d8c5953fa7 100644 --- a/litellm-rust/crates/core/src/ocr/hooks.rs +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -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 for OcrLifecycleHooks { + type Error = crate::ocr::Error; type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; type SuccessFuture<'a> = OcrLogFuture<'a>; diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index 92c9d4b717c..efa2b1f2873 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -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 = Result, Error>; @@ -84,7 +84,7 @@ impl OcrHostOperation { pub enum OcrHostResult { Request(Result<(Box, bool), Error>), - Lifecycle(Result<(), HostFailure>), + Lifecycle(Result<(), HostFailure>), AzureAdToken(Result), PreCall(Result), DuringCall(Result), @@ -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>) { 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 { + pub async fn interrupt(&mut self, failure: HostFailure) -> Result { 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, - ) -> 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, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error> { Box::pin(OcrCall::interrupt(self, failure)) } } diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index e29fd6ac572..3b51ff98356 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -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; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 241cb09529a..5a48206d53c 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -119,7 +119,7 @@ pub(crate) fn build_http_request( .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) } diff --git a/litellm-rust/crates/core/src/ocr/registry.rs b/litellm-rust/crates/core/src/ocr/registry.rs index 54f57cc3ca2..17185a02020 100644 --- a/litellm-rust/crates/core/src/ocr/registry.rs +++ b/litellm-rust/crates/core/src/ocr/registry.rs @@ -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 { diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 76df8b42806..69e6982414b 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -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")] diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 6dc6b34b73d..93816effcb1 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -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}, diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs index b22de6c47de..2cc94751fb4 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs @@ -1,5 +1,5 @@ use super::*; -use crate::Error; +use crate::chat_completions::Error; use serde_json::json; fn messages(value: Value) -> Vec { diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index 56175322442..ba1a1e1d350 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -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, }; diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index fe96155bbbb..080f11c8cac 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -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, -) -> Result { +) -> Result { 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, ) -> Result { - 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] diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs deleted file mode 100644 index 33d007c1945..00000000000 --- a/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -mod credential_provider_cache; -mod native; -mod resolve; -mod types; - -pub(crate) use resolve::AzureAuthService; -pub(crate) use types::AzureAuthInputs; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 8edda1901e5..182aea84ab2 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -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('/'); diff --git a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs index 4f41d1d6abb..ba63992f3cb 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs @@ -1,2 +1 @@ -pub(crate) mod auth; pub mod messages; diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index bb4f6afe5f9..a418e860b92 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -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}; diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index e5e52bfce95..b51cef7545c 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -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>> = OnceLock::new(); - -fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { - 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, - pub secret_access_key: Option, - pub session_token: Option, - pub region_name: Option, - pub session_name: Option, - pub profile_name: Option, - pub role_name: Option, - pub web_identity_token: Option, - pub sts_endpoint: Option, - pub external_id: Option, -} - -impl AwsAuthConfig { - fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option + 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, - }, - 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 { - 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 + 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 + Sync), -) -> Result { - 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 { - if role_identity(role).is_none() { - return Ok(false); - } - if let (Ok(current_role), Ok(token_file)) = ( - std::env::var(AWS_ROLE_ARN), - std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), - ) && !token_file.is_empty() - { - return Ok(same_role_arns(role, ¤t_role)); - } - - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = config.region_name.clone() { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = config.sts_endpoint.clone() { - loader = loader.endpoint_url(endpoint); - } - let sdk_config = loader.load().await; - let response = match aws_sdk_sts::Client::new(&sdk_config) - .get_caller_identity() - .send() - .await - { - Ok(response) => response, - Err(_) => return Ok(false), - }; - Ok(response - .arn() - .is_some_and(|caller| same_role_arns(role, caller))) -} - -fn default_session_name() -> String { - let seconds = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_or(0, |duration| duration.as_secs()); - format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}") -} - -/// The subset of `headers` SigV4 should cover. -/// -/// Python signs only these and reattaches the rest afterwards, so a forwarded -/// client header cannot change the canonical request and invalidate the -/// signature. Signing everything instead makes the request 403 on a header the -/// caller supplied, on a deployment that works on the Python path. -pub fn aws_signature_headers(headers: &BTreeMap) -> BTreeMap { - 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, - region: &str, - credentials: &Credentials, - signing_time: SystemTime, -) -> Result, Error> { - let identity: Identity = credentials.clone().into(); - let params = v4::SigningParams::builder() - .identity(&identity) - .region(region) - .name(BEDROCK_SERVICE) - .time(signing_time) - .settings(SigningSettings::default()) - .build() - .map(SigningParams::from) - .map_err(|error| Error::Auth(format!("AWS signing parameters failed: {error}")))?; - let header_refs = headers - .iter() - .map(|(name, value)| (name.as_str(), value.as_str())); - let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) - .map_err(|error| Error::Auth(format!("AWS signable request failed: {error}")))?; - let (instructions, _) = sign(request, ¶ms) - .map_err(|error| Error::Auth(format!("AWS request signing failed: {error}")))? - .into_parts(); - Ok(instructions - .headers() - .map(|(name, value)| { - let normalized_name = match name { - "authorization" => "Authorization", - "x-amz-date" => "X-Amz-Date", - "x-amz-security-token" => "X-Amz-Security-Token", - _ => name, - }; - (normalized_name.to_string(), value.to_string()) - }) - .collect()) -} - -/// Model-id and region parsing shared by every Bedrock route. -pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { - 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, - env_lookup: &dyn Fn(&str) -> Option, -) -> 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, - env_lookup: &dyn Fn(&str) -> Option, -) -> 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) -> Option { - 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 { - None - } - - fn parity_inputs() -> (String, Vec, BTreeMap) { - ( - "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> { - 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::*; diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs index c86f061b9ca..74716a2200b 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs @@ -1,5 +1,5 @@ use super::*; -use crate::Error; +use crate::chat_completions::Error; use serde_json::json; fn messages(value: Value) -> Vec { diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index 02e6797b338..19efaf833bd 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -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}; diff --git a/litellm-rust/crates/core/src/providers/bedrock/constants.rs b/litellm-rust/crates/core/src/providers/bedrock/constants.rs index be215cc9016..663f887c1fd 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/constants.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/constants.rs @@ -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::*; diff --git a/litellm-rust/crates/core/src/providers/bedrock/mod.rs b/litellm-rust/crates/core/src/providers/bedrock/mod.rs index d9cd3efcb74..5c849064989 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/mod.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/mod.rs @@ -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; diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index 6b1881d56ca..70ca4386fff 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -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; diff --git a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs index be86bb90311..6203b195d5e 100644 --- a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs @@ -1,4 +1,4 @@ -use crate::Error; +use crate::responses::Error; use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model}; diff --git a/litellm-rust/crates/core/src/responses/error.rs b/litellm-rust/crates/core/src/responses/error.rs new file mode 100644 index 00000000000..8bea035f0b0 --- /dev/null +++ b/litellm-rust/crates/core/src/responses/error.rs @@ -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), +} diff --git a/litellm-rust/crates/core/src/responses/instrumentation.rs b/litellm-rust/crates/core/src/responses/instrumentation.rs index b1098f4d386..b1cf5ae09d8 100644 --- a/litellm-rust/crates/core/src/responses/instrumentation.rs +++ b/litellm-rust/crates/core/src/responses/instrumentation.rs @@ -5,7 +5,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use serde_json::Value; -use crate::Error; +use super::Error; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType}; @@ -208,6 +208,7 @@ impl ResponsesWsInstrumentation { type LifecycleFuture<'a, T> = Pin> + Send + 'a>>; impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation { + type Error = Error; type PreCallFuture<'a> = LifecycleFuture<'a, ()>; type DuringCallFuture<'a> = LifecycleFuture<'a, ()>; type SuccessFuture<'a> = Pin + Send + 'a>>; diff --git a/litellm-rust/crates/core/src/responses/mod.rs b/litellm-rust/crates/core/src/responses/mod.rs index 5ec5a2caef8..f8b6d27ffab 100644 --- a/litellm-rust/crates/core/src/responses/mod.rs +++ b/litellm-rust/crates/core/src/responses/mod.rs @@ -1,3 +1,5 @@ +mod error; +pub use error::Error; pub mod instrumentation; pub mod types; pub mod websocket; diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index 34213e5f6c4..ab7738e81b9 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -16,7 +16,7 @@ use tokio_tungstenite::{ Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, }; -use crate::Error; +use super::Error; use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult}; @@ -204,9 +204,9 @@ impl ResponsesWebSocketConnection { headers: &HashMap, timeout: Option, ) -> Result { - let mut request = url - .into_client_request() - .map_err(|error| Error::Network(error.to_string()))?; + let mut request = url.into_client_request().map_err(|error| { + Error::Transport(crate::transport::Error::Network(error.to_string())) + })?; for (name, value) in headers { let header_name = name .parse::() @@ -217,17 +217,21 @@ impl ResponsesWebSocketConnection { } let connect = connect_upstream(request); let result = match timeout { - Some(timeout) => tokio::time::timeout(timeout, connect) - .await - .map_err(|_| Error::Network("Responses WebSocket connection timed out".into()))?, + Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { + Error::Transport(crate::transport::Error::Network( + "Responses WebSocket connection timed out".into(), + )) + })?, None => connect.await, }; let (socket, _) = result.map_err(|error| match *error { - tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { - status: response.status().as_u16(), - body: String::new(), - }, - other => Error::Network(other.to_string()), + tokio_tungstenite::tungstenite::Error::Http(response) => { + Error::Transport(crate::transport::Error::Http { + status: response.status().as_u16(), + body: String::new(), + }) + } + other => Error::Transport(crate::transport::Error::Network(other.to_string())), })?; Ok(Self { socket: Arc::new(Mutex::new(Some(socket))), @@ -237,12 +241,14 @@ impl ResponsesWebSocketConnection { pub async fn send_text(&self, text: String) -> Result<(), Error> { let mut socket = self.socket.lock().await; let Some(socket) = socket.as_mut() else { - return Err(Error::Network("Responses WebSocket is closed".into())); + return Err(Error::Transport(crate::transport::Error::Network( + "Responses WebSocket is closed".into(), + ))); }; socket .send(Message::Text(text)) .await - .map_err(|error| Error::Network(error.to_string())) + .map_err(|error| Error::Transport(crate::transport::Error::Network(error.to_string()))) } pub async fn recv_text(&self) -> Result, Error> { @@ -257,17 +263,18 @@ impl ResponsesWebSocketConnection { .map_err(|error| Error::InvalidResponse(error.to_string())), Some(Ok(Message::Close(_))) | None => Ok(None), Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(Error::Network(error.to_string())), + Some(Err(error)) => Err(Error::Transport(crate::transport::Error::Network( + error.to_string(), + ))), } } pub async fn close(&self) -> Result<(), Error> { let mut socket = self.socket.lock().await; if let Some(socket) = socket.as_mut() { - socket - .close(None) - .await - .map_err(|error| Error::Network(error.to_string()))?; + socket.close(None).await.map_err(|error| { + Error::Transport(crate::transport::Error::Network(error.to_string())) + })?; } *socket = None; Ok(()) diff --git a/litellm-rust/crates/core/src/transport/error.rs b/litellm-rust/crates/core/src/transport/error.rs new file mode 100644 index 00000000000..eff15365ea8 --- /dev/null +++ b/litellm-rust/crates/core/src/transport/error.rs @@ -0,0 +1,75 @@ +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum Error { + #[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 Error { + 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 for Error { + fn from(error: reqwest::Error) -> Self { + Self::Network(error.without_url().to_string()) + } +} + +#[cfg(test)] +mod tests { + #[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 = crate::transport::Error::from_reqwest_before_dispatch(error); + assert!(matches!(error, crate::transport::Error::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!( + crate::transport::Error::from_reqwest_before_dispatch(error), + crate::transport::Error::Network(_) + )); + } +} diff --git a/litellm-rust/crates/core/src/transport/mod.rs b/litellm-rust/crates/core/src/transport/mod.rs new file mode 100644 index 00000000000..0405e9de3c3 --- /dev/null +++ b/litellm-rust/crates/core/src/transport/mod.rs @@ -0,0 +1,2 @@ +mod error; +pub use error::Error; diff --git a/litellm-rust/crates/core/src/url_utils.rs b/litellm-rust/crates/core/src/url_utils.rs index 1150f93a5c7..b8d82b7a04a 100644 --- a/litellm-rust/crates/core/src/url_utils.rs +++ b/litellm-rust/crates/core/src/url_utils.rs @@ -1,9 +1,8 @@ use std::marker::PhantomData; -use thiserror::Error; use url::Url; -#[derive(Debug, Error)] +#[derive(Debug, thiserror::Error)] pub(crate) enum ApiUrlError { #[error("invalid URL: {0}")] Parse(#[from] url::ParseError), diff --git a/litellm-rust/crates/core/tests/host_lifecycle.rs b/litellm-rust/crates/core/tests/host_lifecycle.rs index 19fb946afde..0e58462af1a 100644 --- a/litellm-rust/crates/core/tests/host_lifecycle.rs +++ b/litellm-rust/crates/core/tests/host_lifecycle.rs @@ -1,5 +1,5 @@ -use crate::Error; use crate::call_lifecycle::host::{HostFailure, HostLifecycle, HostPhase}; +use crate::ocr::Error; fn run(fail_at: Option, asynchronous: bool) -> (Vec, Vec) { let mut lifecycle = HostLifecycle::new(asynchronous); @@ -80,14 +80,14 @@ fn only_provider_and_response_construction_failures_use_provider_mapping() { fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() { let mut lifecycle = HostLifecycle::new(true); while lifecycle.phase() != HostPhase::Execute { - lifecycle.accept(Ok(())); + lifecycle.accept::(Ok(())); } let selected = Error::InvalidRequest("provider".into()); assert_eq!( lifecycle.accept(Err(HostFailure::Error(selected.clone()))), Some(selected) ); - lifecycle.accept(Ok(())); + lifecycle.accept::(Ok(())); for phase in [ HostPhase::DeploymentFailure, HostPhase::Failure, diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 55f8713d76e..373972cf68b 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -140,7 +140,7 @@ impl OcrHooks for RecordingHooks { Box::pin(async move { self.events.lock().unwrap().push("pre"); if self.block { - return Err(crate::Error::InvalidRequest("blocked".into())); + return Err(crate::ocr::Error::InvalidRequest("blocked".into())); } Ok(request) }) @@ -177,7 +177,7 @@ impl OcrHooks for RecordingHooks { fn failure<'a>( &'a self, _context: &'a CallLifecycleContext, - _error: &'a crate::Error, + _error: &'a crate::ocr::Error, _timing: &'a CallLifecycleTiming, ) -> OcrLogFuture<'a> { Box::pin(async move { @@ -251,7 +251,7 @@ async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { ..request }; let error = perform_ocr(request).await.unwrap_err(); - assert!(matches!(error, crate::Error::InvalidRequest(_))); + assert!(matches!(error, crate::ocr::Error::InvalidRequest(_))); assert_eq!(*events.lock().unwrap(), ["pre", "failure"]); } @@ -358,7 +358,7 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() { OcrHostOperation::PreCall(request) => { phases.push("pre"); result = Some(OcrHostResult::PreCall(if failure_phase == "pre" { - Err(crate::Error::InvalidRequest("pre failed".into())) + Err(crate::ocr::Error::InvalidRequest("pre failed".into())) } else { Ok(request) })); @@ -366,7 +366,7 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() { OcrHostOperation::DuringCall(request) => { phases.push("during"); result = Some(OcrHostResult::DuringCall(if failure_phase == "during" { - Err(crate::Error::InvalidRequest("during failed".into())) + Err(crate::ocr::Error::InvalidRequest("during failed".into())) } else { Ok(request) })); @@ -377,7 +377,7 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() { Ok(OcrCallStep::Complete(_)) => panic!("failed call completed"), } }; - assert!(matches!(error, crate::Error::InvalidRequest(_))); + assert!(matches!(error, crate::ocr::Error::InvalidRequest(_))); assert_eq!( phases .iter() @@ -420,7 +420,7 @@ async fn invalid_provider_response_runs_post_call_before_normalization_failure() } }; server.await.unwrap(); - assert!(matches!(error, crate::Error::InvalidResponse(_))); + assert!(matches!(error, crate::ocr::Error::InvalidResponse(_))); assert_eq!(seen.lock().unwrap().len(), 1); assert_eq!(post_calls, [json!(r#"{"pages":"invalid"}"#)]); } @@ -497,7 +497,7 @@ async fn direct_native_host_drives_the_same_state_machine() { ); assert!(matches!( call.resume(None).await, - Err(crate::Error::InvalidRequest(_)) + Err(crate::ocr::Error::InvalidRequest(_)) )); } @@ -516,7 +516,7 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide ) else { panic!("supported call declined") }; - let selected = crate::Error::InvalidRequest("public metadata failed".into()); + let selected = crate::ocr::Error::InvalidRequest("public metadata failed".into()); let host = NoopOcrHost; let mut result = None; let mut failures = Vec::new(); @@ -531,7 +531,7 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide assert_eq!(error, selected); failures.push("sync"); OcrHostResult::Lifecycle(Err(HostFailure::Error( - crate::Error::InvalidRequest("failure callback failed".into()), + crate::ocr::Error::InvalidRequest("failure callback failed".into()), ))) } OcrHostOperation::Lifecycle(HostPhase::AsyncFailure) => { @@ -590,7 +590,7 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption OcrCallStep::Complete(_) => panic!("provider executed before pre-call result"), } } - let selected = crate::Error::InvalidRequest("cancelled".into()); + let selected = crate::ocr::Error::InvalidRequest("cancelled".into()); assert!(matches!( call.interrupt(HostFailure::Cancelled(selected.clone())).await, Err(error) if error == selected @@ -694,10 +694,7 @@ async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_dra .await .unwrap_err(); match error { - super::error::OcrError::Transport(crate::error::TransportError::Http { - status, - body, - }) => { + super::error::OcrError::Transport(crate::transport::Error::Http { status, body }) => { assert_eq!(status, 429); assert_eq!( body, @@ -755,8 +752,8 @@ impl Drop for TokenFutureDrop { } } -impl crate::auth::TokenProvider for PendingToken { - fn acquire(&self) -> crate::auth::TokenFuture<'_> { +impl litellm_auth::TokenProvider for PendingToken { + fn acquire(&self) -> litellm_auth::TokenFuture<'_> { Box::pin(async move { let _guard = TokenFutureDrop(self.dropped.clone()); self.entered.notify_one(); @@ -781,7 +778,7 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_ extra_headers: vec![("authorization".into(), "Bearer test-key".into())], ..request.connection }, - azure_ad_token_provider: Some(crate::auth::TokenProviderHandle::new(Arc::new( + azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( PendingToken { entered: entered.clone(), dropped: dropped.clone(), @@ -811,7 +808,7 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_ } }).await.unwrap(); assert!(!dropped.load(Ordering::SeqCst)); - let selected = crate::Error::InvalidRequest("cancelled".into()); + let selected = crate::ocr::Error::InvalidRequest("cancelled".into()); if interrupt_acknowledgement { let mut acknowledgement = Box::pin(call.interrupt(HostFailure::Cancelled(selected.clone()))); diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index a2e67dffc7d..c7b64e300f0 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -17,7 +17,7 @@ pub(crate) fn ocr_client() -> OcrClient { pub(crate) async fn perform_ocr( request: LiteLLMOcrRequest, -) -> Result { +) -> Result { ocr_client().perform(request).await } diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs index 676799eb2fe..a73c1e7710a 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs @@ -1,7 +1,7 @@ use serde_json::{Value, json}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; -use crate::auth::InputSource; +use litellm_auth::InputSource; fn request_body(request: &str) -> Value { serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index 96a19dd62b4..93e9efca849 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -1,7 +1,7 @@ use serde_json::{Value, json}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; -use crate::auth::InputSource; +use litellm_auth::InputSource; fn request_body(request: &str) -> Value { serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index c112502b76f..1562d4c1021 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -17,7 +17,8 @@ panic-test = [] [dependencies] futures-util.workspace = true -litellm-core = { workspace = true, features = ["bedrock-auth"] } +litellm-core.workspace = true +litellm-auth.workspace = true litellm-token-counter.workspace = true litellm-python-interop.workspace = true pyo3.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/auth.rs b/litellm-rust/crates/python-bridge/src/auth.rs index 8dc0b7aabf0..dcc1a60e9f0 100644 --- a/litellm-rust/crates/python-bridge/src/auth.rs +++ b/litellm-rust/crates/python-bridge/src/auth.rs @@ -1,4 +1,4 @@ -use litellm_core::auth::{ResolvedCredential, SecretValue}; +use litellm_auth::{ResolvedCredential, SecretValue}; use pyo3::exceptions::{PyException, PyRuntimeError, PyTypeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 701c6abb68c..7ca86b3ccfa 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -1,4 +1,5 @@ -use litellm_core::error::Error; +use litellm_core::transport::Error as TransportError; +use litellm_core::{Error, audio_transcription, chat_completions, messages, ocr, responses}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; @@ -16,43 +17,99 @@ pyo3::create_exception!( "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." ); -pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::Auth(message) => PyValueError::new_err(message), - Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) - | Error::MissingDocumentUrl => PyValueError::new_err(err.to_string()), - other => PyRuntimeError::new_err(other.to_string()), +fn auth_is_value_error(error: &litellm_auth::Error) -> bool { + !matches!(error, litellm_auth::Error::MissingApiKey { .. }) +} + +pub(crate) fn messages_error_to_pyerr(error: messages::Error) -> PyErr { + core_error_to_pyerr(error.into()) +} + +pub(crate) fn audio_transcription_error_to_pyerr(error: audio_transcription::Error) -> PyErr { + core_error_to_pyerr(error.into()) +} + +pub(crate) fn responses_error_to_pyerr(error: responses::Error) -> PyErr { + core_error_to_pyerr(error.into()) +} + +pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr { + let value_error = match &error { + Error::Ocr(error) => matches!( + error, + ocr::Error::Auth(_) + | ocr::Error::InvalidProvider(_) + | ocr::Error::InvalidRequest(_) + | ocr::Error::InvalidType { .. } + | ocr::Error::MissingField(_) + | ocr::Error::MissingDocumentUrl + ), + Error::Messages(error) => match error { + messages::Error::Auth(source) => auth_is_value_error(source), + messages::Error::InvalidProvider(_) + | messages::Error::InvalidRequest(_) + | messages::Error::Headers(_) => true, + _ => false, + }, + Error::AudioTranscription(error) => match error { + audio_transcription::Error::Auth(source) => auth_is_value_error(source), + audio_transcription::Error::InvalidProvider(_) + | audio_transcription::Error::InvalidRequest(_) + | audio_transcription::Error::Headers(_) + | audio_transcription::Error::InvalidType { .. } + | audio_transcription::Error::MissingField(_) + | audio_transcription::Error::Aws(_) => true, + _ => false, + }, + Error::ChatCompletions(error) => match error { + chat_completions::Error::Auth(source) => auth_is_value_error(source), + chat_completions::Error::InvalidProvider(_) + | chat_completions::Error::InvalidRequest(_) + | chat_completions::Error::Headers(_) + | chat_completions::Error::InvalidType { .. } + | chat_completions::Error::MissingField(_) + | chat_completions::Error::Aws(_) => true, + _ => false, + }, + Error::Responses(error) => match error { + responses::Error::Auth(source) => auth_is_value_error(source), + responses::Error::InvalidProvider(_) + | responses::Error::InvalidRequest(_) + | responses::Error::Headers(_) => true, + _ => false, + }, + }; + if value_error { + PyValueError::new_err(error.to_string()) + } else { + PyRuntimeError::new_err(error.to_string()) } } -/// Map a core error for a route whose host keeps a Python implementation. +/// Map a route error for a route whose host keeps a Python implementation. /// /// The distinction the host needs is whether the provider was already called. /// Everything raised before the request goes out is safe for the host to retry /// on its own path; anything after it is not, because the provider has already /// done the work and billed for it. -pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { - match err { +pub(crate) fn chat_completions_error_to_pyerr(error: chat_completions::Error) -> PyErr { + use chat_completions::Error; + match error { Error::Unsupported(_) | Error::Auth(_) + | Error::Aws(_) | Error::InvalidProvider(_) | Error::InvalidRequest(_) | Error::InvalidType { .. } | Error::MissingField(_) - | Error::MissingDocumentUrl - | Error::MissingApiKey { .. } - | Error::MissingAzureAiCredentials - | Error::MissingAzureDocumentIntelligenceCredentials - | Error::MissingReductoApiKey - | Error::Routing(_) - // Nothing reached the provider, so serving it on Python cannot double - // bill and is the only way the caller gets an answer at all. - | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - Error::Http { status, body } => RustUpstreamError::new_err((status, body)), - Error::Network(message) | Error::InvalidResponse(message) => { + | Error::Headers(_) + | Error::Transport(TransportError::Connect(_)) => { + RustBridgeDeclined::new_err(error.to_string()) + } + Error::Transport(TransportError::Http { status, body }) => { + RustUpstreamError::new_err((status, body)) + } + Error::Transport(TransportError::Network(message)) | Error::InvalidResponse(message) => { RustUpstreamError::new_err((0u16, message)) } } @@ -63,3 +120,55 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("RustBridgeDeclined", py.get_type::())?; module.add("RustUpstreamError", py.get_type::()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn transport_status_and_dispatch_certainty_survive_python_mapping() { + Python::initialize(); + Python::attach(|py| { + let connect = chat_completions_error_to_pyerr( + TransportError::Connect("unreachable".into()).into(), + ); + assert!(connect.is_instance_of::(py)); + let network = + chat_completions_error_to_pyerr(TransportError::Network("timed out".into()).into()); + assert!(network.is_instance_of::(py)); + let upstream = chat_completions_error_to_pyerr( + TransportError::Http { + status: 429, + body: "slow down".into(), + } + .into(), + ); + assert_eq!( + upstream + .value(py) + .getattr("args") + .unwrap() + .extract::<(u16, String)>() + .unwrap(), + (429, "slow down".into()) + ); + }); + } + + #[test] + fn missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors() { + Python::initialize(); + Python::attach(|py| { + let missing = messages_error_to_pyerr(messages::Error::Auth( + litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + }, + )); + assert!(missing.is_instance_of::(py)); + let invalid = + messages_error_to_pyerr(messages::Error::Auth(litellm_auth::Error::InvalidHeader)); + assert!(invalid.is_instance_of::(py)); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index d8dda10068d..ffc4c186980 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -165,7 +165,7 @@ mod tests { use std::thread; use std::time::Instant; - use litellm_core::error::Error; + use litellm_core::messages::Error; use pyo3::panic::PanicException; use pyo3::types::{PyDict, PyModule}; use rstest::{fixture, rstest}; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 27ea2c1913e..0306990fd4d 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,7 +13,7 @@ use pyo3::prelude::*; use pyo3::types::PyAny; use serde_json::Value; -use crate::errors::core_error_to_pyerr; +use crate::errors::responses_error_to_pyerr; use crate::marshal::{marshal_headers, optional_timeout}; #[pyclass] @@ -37,7 +37,7 @@ impl ResponsesWebSocketConnection { pyo3_async_runtimes::tokio::future_into_py(py, async move { let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) .await - .map_err(core_error_to_pyerr)?; + .map_err(responses_error_to_pyerr)?; Ok(ResponsesWebSocketConnection { inner }) }) } @@ -45,21 +45,24 @@ impl ResponsesWebSocketConnection { fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult> { let inner = self.inner.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - inner.send_text(text).await.map_err(core_error_to_pyerr) + inner + .send_text(text) + .await + .map_err(responses_error_to_pyerr) }) } fn recv_text<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - inner.recv_text().await.map_err(core_error_to_pyerr) + inner.recv_text().await.map_err(responses_error_to_pyerr) }) } fn close<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - inner.close().await.map_err(core_error_to_pyerr) + inner.close().await.map_err(responses_error_to_pyerr) }) } } diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index 014564ae89d..cf9f31c3c13 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -35,7 +35,8 @@ pub(crate) trait PythonRoute: Send + Sync { fn state_mut(&mut self) -> &mut PythonCallState; fn classify(operation: &::Operation) -> OperationClass; fn lifecycle_result() -> ::Result; - fn map_error(error: litellm_core::Error) -> PyErr; + fn map_error(error: ::Error) -> PyErr; + fn host_error(message: String) -> ::Error; fn invoke( &mut self, py: Python<'_>, @@ -46,8 +47,10 @@ pub(crate) trait PythonRoute: Send + Sync { } type NativeStep = NativeCallStep<::Operation, ::Complete>; -type NativeResult = Result, litellm_core::Error>; +type NativeResult = Result, ::Error>; type HostResumeStep = HostStep::Call>, Py>; +type NativeResume = + Option::Result, HostFailure<::Error>>>; struct NativeCallState { call: C, @@ -102,7 +105,7 @@ impl PythonLifecycle { fn resume_core( &mut self, py: Python<'_>, - result: Option::Result, HostFailure>>, + result: NativeResume, ) -> PyResult> { let call = Arc::clone(self.call.as_ref().ok_or_else(missing_state)?); let future = async move { @@ -154,8 +157,8 @@ impl PythonLifecycle { py: Python<'_>, error: PyErr, phase: Option, - ) -> HostFailure { - let native = litellm_core::Error::InvalidRequest(error.to_string()); + ) -> HostFailure<::Error> { + let native = R::host_error(error.to_string()); let cancelled = !error.is_instance_of::(py); let failure = if !cancelled { HostFailure::Error(native) @@ -667,6 +670,7 @@ mod tests { struct SyntheticCall(bool); impl NativeCall for SyntheticCall { + type Error = litellm_core::messages::Error; type Operation = (); type Result = (); type Complete = (); @@ -674,7 +678,7 @@ mod tests { fn resume( &mut self, result: Option, - ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + ) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error> { Box::pin(async move { match (self.0, result) { (false, None) => { @@ -682,7 +686,7 @@ mod tests { Ok(NativeCallStep::Host(())) } (true, Some(())) => Ok(NativeCallStep::Complete(())), - _ => Err(litellm_core::Error::InvalidRequest( + _ => Err(litellm_core::messages::Error::InvalidRequest( "invalid synthetic lifecycle state".into(), )), } @@ -691,8 +695,8 @@ mod tests { fn interrupt( &mut self, - _: HostFailure, - ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + _: HostFailure, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error> { Box::pin(async { Ok(NativeCallStep::Complete(())) }) } } @@ -716,8 +720,12 @@ mod tests { fn lifecycle_result() {} - fn map_error(error: litellm_core::Error) -> PyErr { - crate::errors::core_error_to_pyerr(error) + fn map_error(error: litellm_core::messages::Error) -> PyErr { + crate::errors::messages_error_to_pyerr(error) + } + + fn host_error(message: String) -> litellm_core::messages::Error { + litellm_core::messages::Error::InvalidRequest(message) } fn invoke(&mut self, py: Python<'_>, _: ()) -> PyResult<()> { diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs b/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs index ba4a8bb3739..e95f642e6ea 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs @@ -1,4 +1,4 @@ -use litellm_core::auth::{credential_default_fields, credential_index}; +use litellm_auth::{credential_default_fields, credential_index}; use pyo3::prelude::*; use pyo3::types::{PyDict, PyList}; diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 5f7633a64a0..9038eb971b3 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -6,7 +6,7 @@ use pyo3::prelude::*; use pyo3::types::PyDict; use serde_json::{Map, Value}; -use litellm_core::auth::InputSource; +use litellm_auth::InputSource; use litellm_python_interop::from_py_preserving_errors as from_py; pub(crate) struct RouteOptions { diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription/value.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription/value.rs index af60515b0e2..5ecca63fcb6 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription/value.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription/value.rs @@ -1,4 +1,4 @@ -use litellm_core::Error; +use litellm_core::audio_transcription::Error; use std::future::Future; use litellm_core::audio_transcription::{ @@ -7,7 +7,7 @@ use litellm_core::audio_transcription::{ use pyo3::prelude::*; use serde_json::Value; -use crate::errors::core_error_to_pyerr; +use crate::errors::audio_transcription_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; fn prepare_transcription( @@ -67,5 +67,5 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_transcription, - errors = core_error_to_pyerr, + errors = audio_transcription_error_to_pyerr, } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs index e67bfa89cc7..09f2ada51a5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs @@ -1,4 +1,4 @@ -use litellm_core::Error; +use litellm_core::chat_completions::Error; use std::future::Future; use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index d7d868fd1f4..4c8d98ebe62 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -79,7 +79,7 @@ mod tests { use std::ffi::CString; use std::sync::atomic::{AtomicBool, Ordering}; - use litellm_core::error::Error; + use litellm_core::messages::Error; use pyo3::exceptions::PyLookupError; use pyo3::types::{PyDict, PyList}; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/value.rs b/litellm-rust/crates/python-bridge/src/routes/messages/value.rs index b741e54f0ca..f5eb80d765c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/value.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/value.rs @@ -1,11 +1,11 @@ -use litellm_core::Error; +use litellm_core::messages::Error; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; use pyo3::prelude::*; use serde_json::Value; use std::future::Future; -use crate::errors::core_error_to_pyerr; +use crate::errors::messages_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, required_object}; fn prepare_messages( @@ -61,5 +61,5 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_messages, - errors = core_error_to_pyerr, + errors = messages_error_to_pyerr, } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index 66bdfb7583e..e4ce813d297 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -1,4 +1,4 @@ -use litellm_core::error::Error; +use litellm_core::ocr::Error; use pyo3::prelude::*; use crate::errors::{RustUpstreamError, core_error_to_pyerr}; @@ -7,7 +7,7 @@ pub(super) fn to_pyerr(error: Error) -> PyErr { let status = error.http_status_code(); let mapped = match error { Error::Http { status, body } => RustUpstreamError::new_err((status, body)), - other => core_error_to_pyerr(other), + other => core_error_to_pyerr(other.into()), }; attach_status(mapped, status) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index 12d902a3544..32794936899 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -1,7 +1,7 @@ use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple}; -use litellm_core::auth::ResolvedCredential; +use litellm_auth::ResolvedCredential; use litellm_core::ocr::hooks::{OcrDuringCallRequest, OcrPostCallRequest, OcrPreCallRequest}; use litellm_core::ocr::{OcrAdmission, OcrCall, OcrClient, OcrHostOperation, OcrHostResult}; use litellm_python_interop::{ @@ -179,10 +179,14 @@ impl PythonRoute for PythonOcrHost { OcrHostResult::Lifecycle(Ok(())) } - fn map_error(error: litellm_core::Error) -> PyErr { + fn map_error(error: litellm_core::ocr::Error) -> PyErr { ocr_error_to_pyerr(error) } + fn host_error(message: String) -> litellm_core::ocr::Error { + litellm_core::ocr::Error::InvalidRequest(message) + } + fn invoke(&mut self, py: Python<'_>, operation: OcrHostOperation) -> PyResult { Ok(match operation { OcrHostOperation::ProjectRequest => { diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 8b6a1b02e19..8d8d5f8c518 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -177,7 +177,7 @@ pub(super) fn admitted_call(outcome: NativeOutcome) -> PyResult