mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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
This commit is contained in:
yujonglee
GitHub
parent
0dd5e6e289
commit
b8928170e9
19 changed files with 1063 additions and 112 deletions
138
litellm-rust/Cargo.lock
generated
138
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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")));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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::<Vec<_>>().join("; "))]
|
||||
CredentialChain(Vec<AuthError>),
|
||||
#[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)]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
mod credential;
|
||||
pub mod error;
|
||||
pub(crate) mod vertex;
|
||||
pub use error::AuthError;
|
||||
pub(crate) mod http;
|
||||
mod policy;
|
||||
|
|
|
|||
592
litellm-rust/crates/core/src/auth/vertex.rs
Normal file
592
litellm-rust/crates/core/src/auth/vertex.rs
Normal file
|
|
@ -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<Sourced<SecretValue>>,
|
||||
project_id: Option<String>,
|
||||
location: Option<String>,
|
||||
}
|
||||
|
||||
impl VertexConfig {
|
||||
pub(crate) fn from_sourced_optional_params(
|
||||
params: &Map<String, Value>,
|
||||
sources: &BTreeMap<String, InputSource>,
|
||||
) -> Result<Self, AuthError> {
|
||||
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<String>,
|
||||
) -> Option<String> {
|
||||
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<String>,
|
||||
) -> Option<String> {
|
||||
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<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
|
||||
loader: Arc<dyn VertexProviderLoader>,
|
||||
}
|
||||
|
||||
impl Default for VertexAuth {
|
||||
fn default() -> Self {
|
||||
Self::new(Arc::new(GcpProviderLoader))
|
||||
}
|
||||
}
|
||||
|
||||
impl VertexAuth {
|
||||
fn new(loader: Arc<dyn VertexProviderLoader>) -> 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<String> + Sync),
|
||||
) -> Result<VertexEnvironment, AuthError> {
|
||||
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<String> + Sync),
|
||||
) -> Result<VertexAccessToken, AuthError> {
|
||||
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<String> + Sync),
|
||||
) -> Result<Arc<dyn VertexTokenSource>, 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<dyn VertexTokenSource>>;
|
||||
}
|
||||
|
||||
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, AuthError>> + Send + 'a>>;
|
||||
|
||||
struct GcpTokenSource(Arc<dyn TokenProvider>);
|
||||
|
||||
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<dyn VertexTokenSource>> {
|
||||
Box::pin(async move {
|
||||
let provider: Arc<dyn TokenProvider> = 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<dyn VertexTokenSource>)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
|
||||
let token_uri = serde_json::from_str::<Value>(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<String>,
|
||||
) -> 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<String, Value>,
|
||||
sources: &BTreeMap<String, InputSource>,
|
||||
names: &[&str],
|
||||
) -> Result<Option<Sourced<SecretValue>>, 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<String, InputSource>, name: &str) -> InputSource {
|
||||
sources.get(name).copied().unwrap_or_default()
|
||||
}
|
||||
|
||||
fn optional_string(
|
||||
params: &Map<String, Value>,
|
||||
names: &[&str],
|
||||
) -> Result<Option<String>, 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<String>, name: &str) -> Option<String> {
|
||||
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<AtomicUsize>,
|
||||
}
|
||||
|
||||
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<AtomicUsize>,
|
||||
provider: Arc<dyn VertexTokenSource>,
|
||||
}
|
||||
|
||||
impl VertexProviderLoader for FakeLoader {
|
||||
fn load(
|
||||
&self,
|
||||
_source: CredentialSource,
|
||||
) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
|
||||
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<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
|
||||
let provider: Arc<dyn VertexTokenSource> = 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
|
|
|||
154
litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs
Normal file
154
litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs
Normal file
|
|
@ -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<reqwest::Request, OcrError> {
|
||||
validate_destination(&request.connection)?;
|
||||
let ParsedProviderParams {
|
||||
known: params,
|
||||
extra_params: _extra_params,
|
||||
} = _prepare_ocr_request::<MistralOcrParams>(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<LiteLLMOcrResponse, OcrResponseError> {
|
||||
mistral::transform_ocr_response(&request.model, response)
|
||||
}
|
||||
}
|
||||
|
||||
fn get_complete_url(
|
||||
api_base: Option<&str>,
|
||||
project: &str,
|
||||
location: &str,
|
||||
model: &str,
|
||||
) -> Result<String, OcrError> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
19
litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs
Normal file
19
litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs
Normal file
|
|
@ -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(())
|
||||
}
|
||||
|
|
@ -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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<String, Value>, 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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
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<String, Value>,
|
||||
|
|
@ -207,53 +179,6 @@ fn ocr_data_from_content(content: Value, usage: Option<Value>, 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<String, Value>,
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
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<String>,
|
||||
) -> Result<String, Error> {
|
||||
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
|
||||
|
|
|
|||
110
litellm-rust/crates/core/tests/vertex_ai_ocr.rs
Normal file
110
litellm-rust/crates/core/tests/vertex_ai_ocr.rs
Normal file
|
|
@ -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")
|
||||
);
|
||||
}
|
||||
|
|
@ -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")));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue