From b8928170e93b4937857e947e29224ba1dc30880e Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 11 Sep 2026 16:22:56 -0700 Subject: [PATCH] feat(ocr): add Vertex Mistral adapter (#40507) * feat(ocr): add Vertex Mistral adapter * test(ocr): validate Vertex credentials at adapter boundary * refactor(ocr): preserve Vertex Mistral extra params * refactor(ocr): align Vertex authentication lifecycle * refactor(ocr): keep Vertex preparation behind bridge * fix(ocr): protect Vertex credential destinations * fix(auth): restrict request Vertex token endpoints --- litellm-rust/Cargo.lock | 138 ++++ litellm-rust/Cargo.toml | 1 + .../crates/ai-gateway/src/ocr/common_utils.rs | 6 +- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 2 + litellm-rust/crates/core/Cargo.toml | 1 + litellm-rust/crates/core/src/auth/error.rs | 8 + litellm-rust/crates/core/src/auth/mod.rs | 1 + litellm-rust/crates/core/src/auth/vertex.rs | 592 ++++++++++++++++++ .../crates/core/src/ocr/adapters/mod.rs | 3 + .../core/src/ocr/adapters/vertex/mistral.rs | 154 +++++ .../core/src/ocr/adapters/vertex/mod.rs | 19 + litellm-rust/crates/core/src/ocr/client.rs | 8 + litellm-rust/crates/core/src/ocr/mod.rs | 3 + litellm-rust/crates/core/src/ocr/registry.rs | 7 + .../providers/vertex_ai/ocr/transformation.rs | 106 ---- .../crates/core/tests/vertex_ai_ocr.rs | 110 ++++ .../crates/python-bridge/src/routes/ocr.rs | 3 +- litellm/rust_bridge/ocr.py | 11 +- tests/test_litellm/ocr/test_rust_bridge.py | 2 + 19 files changed, 1063 insertions(+), 112 deletions(-) create mode 100644 litellm-rust/crates/core/src/auth/vertex.rs create mode 100644 litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs create mode 100644 litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs create mode 100644 litellm-rust/crates/core/tests/vertex_ai_ocr.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9c0a8cb7fe7..cf000e85d68 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -40,6 +40,15 @@ dependencies = [ "cc", ] +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + [[package]] name = "anes" version = "0.1.6" @@ -693,6 +702,20 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + [[package]] name = "ciborium" version = "0.2.2" @@ -1287,6 +1310,33 @@ dependencies = [ "slab", ] +[[package]] +name = "gcp_auth" +version = "0.12.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d27dbcc645b60b8e7f6e2868a9d7102ece97d1bb49c1288b5321fcc67f7260" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "http 1.4.2", + "http-body-util", + "hyper 1.10.1", + "hyper-rustls 0.27.9", + "hyper-util", + "ring", + "rustls 0.23.42", + "rustls-pki-types", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tracing", + "tracing-futures", + "url", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -1595,6 +1645,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.2.0" @@ -1875,6 +1949,7 @@ dependencies = [ "azure_identity", "base64 0.22.1", "data-url", + "gcp_auth", "moka", "rand 0.8.7", "reqwest 0.12.28", @@ -3593,6 +3668,16 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "pin-project", + "tracing", +] + [[package]] name = "tracing-subscriber" version = "0.3.23" @@ -3984,12 +4069,65 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.52.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 0f30ac5cf7d..5f25e69a1f8 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -41,6 +41,7 @@ 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"] } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index 10e08253e1a..8305cc80a1d 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -13,9 +13,7 @@ use litellm_core::providers::azure_ai::ocr::transformation::{ }; use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai; -use litellm_core::providers::vertex_ai::ocr::transformation::{ - VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, -}; +use litellm_core::providers::vertex_ai::ocr::transformation::VERTEX_AI_DEEPSEEK_OCR_CONFIG; use crate::client::http_client; @@ -44,7 +42,7 @@ pub(super) fn ocr_provider_config( } "azure_ai" => Some(&AZURE_AI_OCR_CONFIG), "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG), - "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG), + "vertex_ai" => None, _ => None, } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index cdf4a15125f..116c0f4a5e9 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -36,5 +36,7 @@ mod tests { Some("azure_ai") )); assert!(is_supported_request("parse-v3", Some("reducto"))); + assert!(is_supported_request("mistral-ocr", Some("vertex_ai"))); + assert!(!is_supported_request("deepseek-ocr", Some("vertex_ai"))); } } diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index dc4a1acea16..a2433435e34 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -15,6 +15,7 @@ base64.workspace = true azure_core.workspace = true azure_identity.workspace = true data-url = "0.3.2" +gcp_auth.workspace = true moka.workspace = true rand.workspace = true reqwest.workspace = true diff --git a/litellm-rust/crates/core/src/auth/error.rs b/litellm-rust/crates/core/src/auth/error.rs index ddd4a6d016e..e7027c0df10 100644 --- a/litellm-rust/crates/core/src/auth/error.rs +++ b/litellm-rust/crates/core/src/auth/error.rs @@ -6,6 +6,8 @@ pub enum AuthError { 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")] @@ -73,6 +75,12 @@ pub enum AuthConfigurationError { 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)] diff --git a/litellm-rust/crates/core/src/auth/mod.rs b/litellm-rust/crates/core/src/auth/mod.rs index 2ca2f3c3016..35d9c676f65 100644 --- a/litellm-rust/crates/core/src/auth/mod.rs +++ b/litellm-rust/crates/core/src/auth/mod.rs @@ -1,5 +1,6 @@ mod credential; pub mod error; +pub(crate) mod vertex; pub use error::AuthError; pub(crate) mod http; mod policy; diff --git a/litellm-rust/crates/core/src/auth/vertex.rs b/litellm-rust/crates/core/src/auth/vertex.rs new file mode 100644 index 00000000000..00a0a7ea7ee --- /dev/null +++ b/litellm-rust/crates/core/src/auth/vertex.rs @@ -0,0 +1,592 @@ +use std::collections::BTreeMap; +use std::future::Future; +use std::path::Path; +use std::pin::Pin; +use std::sync::Arc; + +use gcp_auth::{CustomServiceAccount, TokenProvider}; +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}; + +const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; +const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; +const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS"; +const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; +const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; +const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS"; +const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; +const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; +const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; + +#[derive(Clone, Debug, Default)] +pub(crate) struct VertexConfig { + credentials: Option>, + project_id: Option, + location: Option, +} + +impl VertexConfig { + pub(crate) fn from_sourced_optional_params( + params: &Map, + sources: &BTreeMap, + ) -> Result { + Ok(Self { + credentials: optional_credentials( + params, + sources, + &["vertex_credentials", "vertex_ai_credentials"], + )?, + project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?, + location: optional_string(params, &["vertex_location", "vertex_ai_location"])?, + }) + } + + pub(crate) fn project_id(&self) -> Option<&str> { + self.project_id.as_deref() + } + + pub(crate) fn location(&self) -> Option<&str> { + self.location.as_deref() + } +} + +pub(crate) struct VertexEnvironment { + pub headers: Vec<(String, String)>, + pub project_id: String, +} + +struct VertexAccessToken { + token: String, + project_id: String, +} + +pub(crate) fn get_vertex_ai_project( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + config + .project_id() + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV)) +} + +pub(crate) fn get_vertex_ai_location( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + config + .location() + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, VERTEXAI_LOCATION_ENV)) + .or_else(|| non_empty_env(env_lookup, VERTEX_LOCATION_ENV)) +} + +#[derive(Clone)] +pub(crate) struct VertexAuth { + providers: Cache>, + loader: Arc, +} + +impl Default for VertexAuth { + fn default() -> Self { + Self::new(Arc::new(GcpProviderLoader)) + } +} + +impl VertexAuth { + fn new(loader: Arc) -> Self { + Self { + providers: Cache::builder().max_capacity(64).build(), + loader, + } + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + pub(crate) async fn validate_environment( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + let has_authorization = headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("Authorization")); + let static_token = api_key + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .or_else(|| non_empty_env(env_lookup, VERTEX_AI_API_KEY_ENV)) + .or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV)); + let project_id = get_vertex_ai_project(config, env_lookup); + + if !has_authorization && static_token.is_none() { + let access = self.get_access_token(config, env_lookup).await?; + return Ok(VertexEnvironment { + headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, + project_id: project_id.unwrap_or(access.project_id), + }); + } + + let project_id = match project_id { + Some(project_id) => project_id, + None => { + self.load_provider(config, env_lookup) + .await? + .project_id() + .await? + } + }; + let headers = if has_authorization { + headers + } else { + apply_credential( + headers, + static_token.as_deref().expect("static token was checked"), + CredentialPlacement::Bearer, + )? + }; + Ok(VertexEnvironment { + headers, + project_id, + }) + } + + async fn get_access_token( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> 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 }) + } + + async fn load_provider( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, AuthError> { + let source = credential_source(config, env_lookup); + let key = source.cache_key(); + self.providers + .try_get_with(key, self.loader.load(source)) + .await + .map_err(|error| (*error).clone()) + } +} + +trait VertexTokenSource: Send + Sync { + fn project_id(&self) -> VertexAuthFuture<'_, String>; + fn token(&self) -> VertexAuthFuture<'_, String>; +} + +trait VertexProviderLoader: Send + Sync { + fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; +} + +type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; + +struct GcpTokenSource(Arc); + +impl VertexTokenSource for GcpTokenSource { + fn project_id(&self) -> VertexAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .project_id() + .await + .map(|project| project.to_string()) + .map_err(auth_acquisition_error) + }) + } + + fn token(&self) -> VertexAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .token(&[CLOUD_PLATFORM_SCOPE]) + .await + .map(|token| token.as_str().to_string()) + .map_err(auth_acquisition_error) + }) + } +} + +struct GcpProviderLoader; + +impl VertexProviderLoader for GcpProviderLoader { + fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { + Box::pin(async move { + let provider: Arc = match source { + CredentialSource::Inline(configured) => Arc::new( + CustomServiceAccount::from_json(validate_request_credentials( + configured.expose(), + )?) + .map_err(auth_acquisition_error)?, + ), + CredentialSource::Trusted(configured) => { + let configured = configured.expose(); + let service_account = if Path::new(configured).is_file() { + CustomServiceAccount::from_file(configured) + } else { + CustomServiceAccount::from_json(configured) + } + .map_err(auth_acquisition_error)?; + Arc::new(service_account) + } + CredentialSource::ApplicationCredentials(path) => { + Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?) + } + CredentialSource::Adc => { + gcp_auth::provider().await.map_err(auth_acquisition_error)? + } + }; + Ok(Arc::new(GcpTokenSource(provider)) as Arc) + }) + } +} + +fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> { + let token_uri = serde_json::from_str::(configured) + .ok() + .and_then(|credentials| { + credentials + .get("token_uri") + .and_then(Value::as_str) + .map(str::to_string) + }); + if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { + return Err(AuthConfigurationError::RequestVertexTokenEndpoint.into()); + } + Ok(configured) +} + +#[derive(Clone, Debug)] +enum CredentialSource { + Inline(SecretValue), + Trusted(SecretValue), + ApplicationCredentials(String), + Adc, +} + +impl CredentialSource { + fn cache_key(&self) -> CredentialCacheKey { + match self { + Self::Inline(configured) => { + CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) + } + Self::Trusted(configured) => { + CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) + } + Self::ApplicationCredentials(path) => { + CredentialCacheKey::ApplicationCredentials(path.clone()) + } + Self::Adc => CredentialCacheKey::Adc, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +enum CredentialCacheKey { + Inline([u8; 32]), + Trusted([u8; 32]), + ApplicationCredentials(String), + Adc, +} + +fn credential_source( + config: &VertexConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> CredentialSource { + if let Some(configured) = config.credentials.clone() { + return match configured.source() { + InputSource::Request => CredentialSource::Inline(configured.into_value()), + InputSource::Deployment | InputSource::Environment => { + CredentialSource::Trusted(configured.into_value()) + } + }; + } + if let Some(configured) = non_empty_env(env_lookup, VERTEXAI_CREDENTIALS_ENV) { + return CredentialSource::Trusted(SecretValue::new(configured)); + } + non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) + .map(CredentialSource::ApplicationCredentials) + .unwrap_or(CredentialSource::Adc) +} + +fn optional_credentials( + params: &Map, + sources: &BTreeMap, + names: &[&str], +) -> Result>, AuthError> { + for name in names { + let source = source_for(sources, name); + 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(Sourced::new(SecretValue::new(value), source))); + } + Some(Value::Object(value)) if value.is_empty() => continue, + Some(Value::Object(value)) => { + return serde_json::to_string(value) + .map(SecretValue::new) + .map(|value| Sourced::new(value, source)) + .map(Some) + .map_err(|error| { + AuthError::Configuration(AuthConfigurationError::InvalidFieldType(format!( + "{}: {error}", + names[0] + ))) + }); + } + Some(_) => { + return Err(AuthError::Configuration( + AuthConfigurationError::InvalidFieldType(names[0].to_string()), + )); + } + } + } + Ok(None) +} + +fn source_for(sources: &BTreeMap, name: &str) -> InputSource { + sources.get(name).copied().unwrap_or_default() +} + +fn optional_string( + params: &Map, + names: &[&str], +) -> Result, AuthError> { + 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()), + )); + } + } + } + Ok(None) +} + +fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option, name: &str) -> Option { + env_lookup(name) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +fn auth_acquisition_error(error: gcp_auth::Error) -> AuthError { + AuthError::VertexTokenAcquisition(error.to_string()) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use serde_json::json; + + use super::*; + + struct FakeProvider { + calls: Arc, + } + + impl VertexTokenSource for FakeProvider { + fn project_id(&self) -> VertexAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok("adc-project".into()) }) + } + + fn token(&self) -> VertexAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok("adc-token".into()) }) + } + } + + struct FakeLoader { + loads: Arc, + provider: Arc, + } + + impl VertexProviderLoader for FakeLoader { + fn load( + &self, + _source: CredentialSource, + ) -> VertexAuthFuture<'_, Arc> { + let loads = self.loads.clone(); + let provider = self.provider.clone(); + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok(provider) + }) + } + } + + fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() + } + + fn auth(calls: Arc, loads: Arc) -> VertexAuth { + let provider: Arc = Arc::new(FakeProvider { calls }); + VertexAuth::new(Arc::new(FakeLoader { loads, provider })) + } + + #[test] + fn config_is_typed_and_secrets_are_redacted() { + let config = config(json!({ + "vertex_credentials":{"private_key":"secret-key"}, + "vertex_project":"project-1", + "vertex_location":"europe-west4" + })); + assert_eq!(config.project_id(), Some("project-1")); + assert_eq!(config.location(), Some("europe-west4")); + assert!(!format!("{config:?}").contains("secret-key")); + assert!( + VertexConfig::from_sourced_optional_params( + json!({"vertex_credentials":true}).as_object().unwrap(), + &BTreeMap::new() + ) + .is_err() + ); + } + + #[test] + fn empty_primary_values_fall_back_to_python_aliases() { + let config = config(json!({ + "vertex_credentials": null, + "vertex_ai_credentials": "alias-credentials", + "vertex_project": " ", + "vertex_ai_project": "alias-project", + "vertex_location": null, + "vertex_ai_location": "alias-location" + })); + assert_eq!( + config.credentials.as_ref().unwrap().value().expose(), + "alias-credentials" + ); + assert_eq!(config.project_id(), Some("alias-project")); + assert_eq!(config.location(), Some("alias-location")); + } + + #[test] + fn project_and_location_prefer_input_then_environment() { + let configured = + config(json!({"vertex_project":"input-project","vertex_location":"input-location"})); + let env = |name: &str| Some(format!("env-{name}")); + assert_eq!( + get_vertex_ai_project(&configured, &env).as_deref(), + Some("input-project") + ); + assert_eq!( + get_vertex_ai_location(&configured, &env).as_deref(), + Some("input-location") + ); + let empty = VertexConfig::default(); + assert_eq!( + get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), + Some("env-project") + ); + assert_eq!( + get_vertex_ai_location(&empty, &|name| (name == VERTEX_LOCATION_ENV) + .then(|| "fallback-location".into())) + .as_deref(), + Some("fallback-location") + ); + } + + #[test] + fn credential_discovery_prefers_input_then_environment_then_adc() { + let params = json!({"vertex_credentials":"input-json"}); + let sources = BTreeMap::from([("vertex_credentials".to_string(), InputSource::Request)]); + let configured = + VertexConfig::from_sourced_optional_params(params.as_object().unwrap(), &sources) + .unwrap(); + assert!( + matches!(credential_source(&configured, &|_| Some("environment-value".into())), CredentialSource::Inline(value) if value.expose() == "input-json") + ); + let empty = VertexConfig::default(); + assert!( + matches!(credential_source(&empty, &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())), CredentialSource::Trusted(value) if value.expose() == "environment-json") + ); + assert!( + matches!(credential_source(&empty, &|name| (name == GOOGLE_APPLICATION_CREDENTIALS_ENV).then(|| "adc.json".into())), CredentialSource::ApplicationCredentials(path) if path == "adc.json") + ); + assert!(matches!( + credential_source(&empty, &|_| None), + CredentialSource::Adc + )); + assert_ne!( + CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), + CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() + ); + } + + #[test] + fn request_credentials_require_canonical_token_endpoint() { + assert!( + validate_request_credentials(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#) + .is_ok() + ); + assert!(matches!( + validate_request_credentials(r#"{"token_uri":"http://127.0.0.1/token"}"#), + Err(AuthError::Configuration( + AuthConfigurationError::RequestVertexTokenEndpoint + )) + )); + assert!(matches!( + validate_request_credentials("{}"), + Err(AuthError::Configuration( + AuthConfigurationError::RequestVertexTokenEndpoint + )) + )); + } + + #[tokio::test] + async fn explicit_token_and_header_do_not_acquire_adc() { + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); + let configured = config(json!({"vertex_project":"project-1"})); + let explicit = auth + .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) + .await + .unwrap(); + assert_eq!(explicit.headers[0].1, "Bearer access-token"); + let existing = auth + .validate_environment( + vec![("authorization".into(), "Bearer existing".into())], + None, + &configured, + &|_| None, + ) + .await + .unwrap(); + assert_eq!(existing.headers[0].1, "Bearer existing"); + assert_eq!(loads.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn provider_is_reused_across_authentication_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(calls.clone(), loads.clone()); + for _ in 0..2 { + let environment = auth + .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None) + .await + .unwrap(); + assert_eq!(environment.project_id, "adc-project"); + assert_eq!(environment.headers[0].1, "Bearer adc-token"); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 4); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/mod.rs index c089ca90605..a96a7fcdf38 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/mod.rs @@ -11,10 +11,12 @@ use super::wire::DecodedOcrResponse; mod azure; mod mistral; mod reducto; +mod vertex; pub(crate) use azure::{AzureDocumentIntelligenceAdapter, AzureMistralAdapter}; pub(crate) use mistral::MistralAdapter; pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter}; +pub(crate) use vertex::VertexMistralAdapter; /// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response. pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static { @@ -70,6 +72,7 @@ macro_rules! for_each_ocr_adapter { AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi; ReductoLegacy, $crate::ocr::adapters::ReductoLegacyAdapter, $crate::ocr::adapters::ReductoLegacyAdapter, Reducto; ReductoV3, $crate::ocr::adapters::ReductoV3Adapter, $crate::ocr::adapters::ReductoV3Adapter, Reducto; + VertexMistral, $crate::ocr::adapters::VertexMistralAdapter, $crate::ocr::adapters::VertexMistralAdapter, VertexAi; } }; } diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs new file mode 100644 index 00000000000..f3335bf497c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs @@ -0,0 +1,154 @@ +use super::super::OcrAdapter; +use super::validate_destination; +use crate::Error; +use crate::auth::vertex::{self, VertexConfig}; +use crate::ocr::OcrClient; +use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; +use crate::ocr::document::{inline_remote_document, validate_inline_document}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{ + _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, +}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; +use crate::url_utils::ApiUrl; +const DEFAULT_LOCATION: &str = "us-central1"; + +#[derive(Clone, Debug)] +pub(crate) struct VertexMistralAdapter; + +impl OcrAdapter for VertexMistralAdapter { + type ProviderResponse = MistralOcrResponse; + const PROVIDER: OcrProvider = OcrProvider::VertexAi; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + validate_destination(&request.connection)?; + let ParsedProviderParams { + known: params, + extra_params: _extra_params, + } = _prepare_ocr_request::(request)?; + let config = VertexConfig::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)?; + let authentication = client + .vertex_auth() + .validate_environment( + request.connection.extra_headers.clone(), + request.connection.api_key.as_deref(), + &config, + &credential_env, + ) + .await + .map_err(Error::from)?; + let location = vertex::get_vertex_ai_location(&config, &credential_env) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + let url = get_complete_url( + request.connection.api_base.as_deref(), + &authentication.project_id, + &location, + &request.model, + )?; + let document = inline_remote_document( + client.document_fetcher(), + request.document.clone(), + &request.connection, + ) + .await?; + let body = mistral::transform_ocr_request(&request.model, document, ¶ms)?; + transform_request_body( + client, + request, + &url, + &authentication.headers, + body, + |body| validate_inline_document(&body.document), + ) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + mistral::transform_ocr_response(&request.model, response) + } +} + +fn get_complete_url( + api_base: Option<&str>, + project: &str, + location: &str, + model: &str, +) -> Result { + validate_location(location)?; + let default_base = format!("https://{location}-aiplatform.googleapis.com"); + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(&default_base); + let prediction = format!("{model}:rawPredict"); + ApiUrl::parse(base) + .and_then(|url| { + url.complete_path(&[ + "v1", + "projects", + project, + "locations", + location, + "publishers", + "mistralai", + "models", + &prediction, + ]) + }) + .map(|url| url.into_string()) + .map_err(|_| { + OcrRequestError::RequestField { + path: "api_base".into(), + } + .into() + }) +} + +fn validate_location(location: &str) -> Result<(), OcrError> { + let valid = !location.is_empty() + && location + .bytes() + .all(|value| value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'-') + && location + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphanumeric) + && location + .as_bytes() + .last() + .is_some_and(u8::is_ascii_alphanumeric); + if valid { + return Ok(()); + } + Err(OcrRequestError::RequestField { + path: "vertex_location".into(), + } + .into()) +} + +#[cfg(test)] +mod tests { + use super::get_complete_url; + + #[test] + fn endpoint_uses_location_project_and_model() { + assert_eq!( + get_complete_url(None, "proj-1", "europe-west4", "mistral-ocr-maas").unwrap(), + "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + assert!(get_complete_url(None, "proj-1", "attacker.example/path", "model").is_err()); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs new file mode 100644 index 00000000000..ce6f884b41d --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs @@ -0,0 +1,19 @@ +mod mistral; + +use crate::Error; +use crate::auth::InputSource; +use crate::auth::error::AuthConfigurationError; +use crate::ocr::error::OcrError; +use crate::ocr::types::OcrConnection; + +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()); + } + Ok(()) +} diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 61b7ae8d995..ab2d098d0bb 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -8,6 +8,7 @@ use super::handler::perform_ocr_request; 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; @@ -17,6 +18,7 @@ pub struct OcrClient { provider_http: reqwest::Client, polling_http: reqwest::Client, document_fetcher: MediaFetcher, + vertex_auth: VertexAuth, } impl OcrClient { @@ -26,6 +28,7 @@ impl OcrClient { provider_http, polling_http: no_redirect_http()?, document_fetcher, + vertex_auth: VertexAuth::default(), }) } @@ -51,12 +54,17 @@ impl OcrClient { &self.document_fetcher } + pub(crate) fn vertex_auth(&self) -> &VertexAuth { + &self.vertex_auth + } + #[cfg(test)] pub(crate) fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self { Self { provider_http, polling_http: no_redirect_http().expect("test polling client builds"), document_fetcher: MediaFetcher::for_test(document_http), + vertex_auth: VertexAuth::default(), } } } diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 5f274186ed0..69b2958483a 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -29,3 +29,6 @@ pub(crate) mod test_support; #[cfg(test)] #[path = "../../tests/ocr.rs"] pub(crate) mod tests; +#[cfg(test)] +#[path = "../../tests/vertex_ai_ocr.rs"] +mod vertex_ai_tests; diff --git a/litellm-rust/crates/core/src/ocr/registry.rs b/litellm-rust/crates/core/src/ocr/registry.rs index 6b100795fc8..1097d102a20 100644 --- a/litellm-rust/crates/core/src/ocr/registry.rs +++ b/litellm-rust/crates/core/src/ocr/registry.rs @@ -26,6 +26,7 @@ pub(crate) enum OcrProvider { Mistral, AzureAi, Reducto, + VertexAi, } impl OcrProvider { @@ -34,6 +35,7 @@ impl OcrProvider { Self::Mistral => "mistral", Self::AzureAi => "azure_ai", Self::Reducto => "reducto", + Self::VertexAi => "vertex_ai", } } } @@ -51,6 +53,7 @@ pub(crate) fn resolve_wire_adapter( "mistral" => OcrProvider::Mistral, "azure_ai" => OcrProvider::AzureAi, "reducto" => OcrProvider::Reducto, + "vertex_ai" => OcrProvider::VertexAi, value => return Err(Error::InvalidProvider(value.to_string())), }; let adapter = match typed_provider { @@ -71,6 +74,10 @@ pub(crate) fn resolve_wire_adapter( provider.model ))); } + OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => { + return Err(Error::Unsupported("Vertex DeepSeek OCR")); + } + OcrProvider::VertexAi => OcrAdapterKind::VertexMistral, }; Ok((provider.model.to_string(), adapter)) } diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index b0e5a278d0b..c2d810822ae 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -3,8 +3,6 @@ use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData}; use serde_json::{Map, Value, json}; -use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; - const VERTEX_DEFAULT_LOCATION: &str = "us-central1"; const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com"; const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; @@ -23,10 +21,8 @@ const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[ "stop", ]; -pub struct VertexAiOcrConfig; pub struct VertexAiDeepSeekOcrConfig; -pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig; pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig; fn string_param<'a>(params: &'a Map, keys: &[&str]) -> Option<&'a str> { @@ -84,30 +80,6 @@ fn vertex_location( .unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string()) } -fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String { - api_base - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com")) - .trim_end_matches('/') - .to_string() -} - -pub fn complete_vertex_mistral_url( - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - let project = vertex_project(optional_params, env_lookup)?; - let location = vertex_location(optional_params, env_lookup); - let base = vertex_mistral_api_base(api_base, &location); - Ok(format!( - "{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict" - )) -} - pub fn complete_vertex_deepseek_url( api_base: Option<&str>, optional_params: &Map, @@ -207,53 +179,6 @@ fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> V } } -impl OcrProviderConfig for VertexAiOcrConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { - MISTRAL_OCR_CONFIG.supported_ocr_params() - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn transform_ocr_request( - &self, - model: &str, - document: Value, - optional_params: Map, - ) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) - } - - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - complete_vertex_mistral_url(api_base, model, optional_params, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_vertex_api_key(api_key, env_lookup) - } - - fn requires_data_uri_document(&self) -> bool { - true - } -} - impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_ocr_params(&self) -> &'static [&'static str] { @@ -379,37 +304,6 @@ mod tests { use super::*; use rstest::rstest; - #[test] - fn vertex_mistral_url_uses_project_location_and_model() { - let params = Map::from_iter([ - ("vertex_project".to_string(), json!("proj-1")), - ("vertex_location".to_string(), json!("europe-west4")), - ]); - - let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None) - .expect("url builds"); - - assert_eq!( - url, - "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - } - - #[test] - fn vertex_mistral_reuses_mistral_body_transform() { - let body = VERTEX_AI_OCR_CONFIG - .transform_ocr_request( - "mistral-ocr-maas", - json!({"type": "image_url", "image_url": "data:image/png;base64,abc"}), - Map::new(), - ) - .expect("request transforms") - .data; - - assert_eq!(body["model"], "mistral-ocr-maas"); - assert_eq!(body["document"]["image_url"], "data:image/png;base64,abc"); - } - #[test] fn vertex_deepseek_request_uses_ocr_endpoint_shape() { let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs new file mode 100644 index 00000000000..2358552c742 --- /dev/null +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -0,0 +1,110 @@ +use serde_json::{Value, json}; + +use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; +use crate::auth::InputSource; +use crate::ocr::wire::{OcrWireRequest, decode_request}; + +fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() +} + +#[tokio::test] +async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/mistral-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "extract_footer":true + }), + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0]["markdown"], "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + assert_eq!( + request_body(&requests[0]), + json!({ + "model":"mistral-ocr-maas", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "extract_footer":true + }) + ); +} + +#[tokio::test] +async fn supplied_authorization_is_forwarded_without_a_static_token() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "vertex_ai/model", + &base, + json!({"vertex_project":"project-1"}), + ); + request.connection.api_key = None; + request.connection.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer supplied") + ); +} + +#[tokio::test] +async fn invalid_credentials_fail_before_provider_http() { + let request = wire_request( + "vertex_ai/model", + "http://127.0.0.1:1", + json!({"vertex_credentials": true}), + ); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("vertex_credentials")); +} + +#[tokio::test] +async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let request = decode_request(OcrWireRequest { + model: "vertex_ai/model".into(), + document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + api_key: Some("test-key".into()), + api_base: Some("https://attacker.example".into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: json!({"vertex_project":"project-1"}) + .as_object() + .unwrap() + .clone(), + input_sources: std::collections::BTreeMap::from([( + "api_base".to_string(), + InputSource::Request, + )]), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + let error = perform_ocr(request).await.unwrap_err(); + + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index d1d4a70007d..0aab11e3cfc 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -111,6 +111,7 @@ mod tests { )); assert!(is_supported_request("parse-v3", Some("reducto"))); assert!(is_supported_request("parse-legacy", Some("reducto"))); - assert!(!is_supported_request("mistral-ocr", Some("vertex_ai"))); + assert!(is_supported_request("mistral-ocr", Some("vertex_ai"))); + assert!(!is_supported_request("deepseek-ocr", Some("vertex_ai"))); } } diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 68bbb186f9b..89eab71ccba 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -179,10 +179,19 @@ def _optional_params(request: LiteLLMOcrRequest, resolve_secret: Callable[[str], or resolve_secret("VERTEXAI_LOCATION") or resolve_secret("VERTEX_LOCATION") ) + credentials: Final = ( + request.kwargs.get("vertex_credentials") + or request.kwargs.get("vertex_ai_credentials") + or resolve_secret("VERTEXAI_CREDENTIALS") + ) vertex_params: Final = MappingProxyType( { name: value - for name, value in (("vertex_project", project), ("vertex_location", location)) + for name, value in ( + ("vertex_project", project), + ("vertex_location", location), + ("vertex_credentials", credentials), + ) if value is not None } ) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 9c1cae6a551..dbb4f822d0b 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -576,6 +576,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana return { "VERTEXAI_PROJECT": "project-from-secret", "VERTEXAI_LOCATION": "us-east5", + "VERTEXAI_CREDENTIALS": "credentials-from-secret", }.get(name) ocr_main._run_rust_ocr( @@ -589,6 +590,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" + assert bridge.calls[0]["optional_params"]["vertex_credentials"] == "credentials-from-secret" def test_prepare_rust_ocr_call_defers_azure_environment_resolution_to_rust():