scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

merge litellm_ocr_rust_default into litellm_shared_reqwest_error_mapper

Browse source 
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-07-17 04:40:10 +00:00
commit 0d93689c6d
18 changed files with 1959 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

779
litellm-rust/Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

5
litellm-rust/Cargo.toml
View file

@ -4,12 +4,13 @@ members = [
"crates/ai-gateway",
"crates/python-bridge",
]
resolver = "2"
resolver = "3"
[workspace.package]
edition = "2021"
license = "MIT"
repository = "https://github.com/BerriAI/litellm"
rust-version = "1.86"
[workspace.dependencies]
litellm-core = { path = "crates/core" }
@ -21,6 +22,8 @@ rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
google-cloud-auth = { version = "=1.9.0", default-features = false }
rustls = { version = "=0.23.41", default-features = false, features = ["ring"] }
sha2 = "0.10"
subtle = "2"
thiserror = "2.0"

9
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -4,6 +4,7 @@ version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
rust-version.workspace = true
[lib]
name = "litellm_ai_gateway"
@ -24,17 +25,17 @@ tokio-tungstenite.workspace = true
futures-util.workspace = true
serde_json.workspace = true
base64.workspace = true
google-cloud-auth.workspace = true
rustls.workspace = true
sha2.workspace = true
axum = { workspace = true, features = ["ws"], optional = true }
serde = { workspace = true, optional = true }
subtle = { workspace = true, optional = true }
# sha2 hashes the master key into user_api_key_hash (matches the proxy's
# SHA-256 hash_token) so the plaintext credential never enters a log payload.
sha2 = { workspace = true, optional = true }
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
[features]
default = []
server = ["dep:axum", "dep:subtle", "dep:serde", "dep:sha2"]
server = ["dep:axum", "dep:subtle", "dep:serde"]
# Build the gateway's config from the proxy YAML via an embedded Python
# interpreter (links libpython; requires `litellm` importable at runtime).
python-config = ["dep:pyo3"]

6
litellm-rust/crates/ai-gateway/src/constants.rs
View file

@ -34,4 +34,10 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500;
#[cfg(feature = "server")]
pub(crate) const DEFAULT_PROVIDER: &str = "openai";
pub(crate) const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
pub(crate) const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS";
pub(crate) const VERTEX_CREDENTIALS_CACHE_CAPACITY: usize = 64;
pub(crate) const ENV_REFERENCE_PREFIX: &str = "os.environ/";

1
litellm-rust/crates/ai-gateway/src/io/mod.rs
View file

@ -1,3 +1,4 @@
pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod vertex_ai;

38
litellm-rust/crates/ai-gateway/src/io/ocr.rs
View file

@ -9,19 +9,23 @@ use std::time::Duration;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::{
OcrAuthStrategy, OcrDocumentPreparation, OcrResponseHandling,
OcrAuth, OcrAuthStrategy, OcrDocumentPreparation, OcrResponseHandling,
};
use litellm_core::providers::vertex_ai::ocr::transformation::{
classify_vertex_bearer, validate_vertex_authorization_headers, VertexTokenSource,
};
use litellm_core::CoreResult;
use serde_json::{Map, Value};
use crate::config::resolve_env_reference;
use crate::io::vertex_ai::VertexAiBase;
mod common_utils;
use crate::errors::map_reqwest_error;
use common_utils::{
convert_document_url_to_data_uri, has_header, ocr_provider_config, poll_document_intelligence,
string_headers, truncate_error_body, upload_reducto_document,
convert_document_url_to_data_uri, has_header, header_values, ocr_provider_config,
poll_document_intelligence, string_headers, truncate_error_body, upload_reducto_document,
};
/// OCR over large documents can take a while; bound it generously rather than
@ -96,9 +100,31 @@ async fn ocr_with_env(
let headers = string_headers(request.extra_headers)?;
let auth_strategy = config.auth_strategy();
let api_key = (!has_header(&headers, auth_strategy.header_name()))
.then(|| config.resolve_api_key(api_key.as_deref(), env_lookup))
.transpose()?;
let api_key = if has_header(&headers, auth_strategy.header_name()) {
if config.ocr_auth() == OcrAuth::VertexOauth {
validate_vertex_authorization_headers(&header_values(
&headers,
auth_strategy.header_name(),
))?;
}
None
} else {
Some(match config.ocr_auth() {
OcrAuth::ProviderKey => config.resolve_api_key(api_key.as_deref(), env_lookup)?,
OcrAuth::VertexOauth => match classify_vertex_bearer(api_key.as_deref(), env_lookup)? {
VertexTokenSource::Explicit(token) => token,
VertexTokenSource::Mint => {
let credentials = VertexAiBase::resolve_credential_source(
&request.optional_params,
env_lookup,
);
VertexAiBase::shared()
.get_access_token(credentials.as_deref())
.await?
}
},
})
};
let url = config.complete_url(
api_base.as_deref(),
model,

8
litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs
View file

@ -88,6 +88,14 @@ pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
.any(|(key, _)| key.eq_ignore_ascii_case(name))
}
pub(super) fn header_values<'a>(headers: &'a [(String, String)], name: &str) -> Vec<&'a str> {
headers
.iter()
.filter(|(key, _)| key.eq_ignore_ascii_case(name))
.map(|(_, value)| value.as_str())
.collect()
}
fn document_url_field(document: &Value) -> CoreResult<Option<(&str, &str)>> {
let Some(object) = document.as_object() else {
return Ok(None);

213
litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs
View file

@ -571,3 +571,216 @@ async fn ocr_maps_unregistered_provider_to_invalid_provider() {
assert_eq!(err.public_status_code(), Some(400));
assert_eq!(err.public_message(), "Invalid OCR request");
}
#[tokio::test]
async fn vertex_ocr_sends_explicit_oauth_bearer_to_raw_predict_endpoint() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_headers(&mut socket).await;
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-maas","usage_info":{"pages_processed":1}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
let mut optional_params = Map::new();
optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into()));
optional_params.insert(
"vertex_location".to_string(),
Value::String("global".into()),
);
let response = ocr(OcrRequest {
model: "mistral-ocr-maas",
document: json!({
"type": "image_url",
"image_url": "data:image/png;base64,abc"
}),
api_key: Some("ya29.explicit-oauth-token"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: "vertex_ai",
extra_headers: None,
optional_params,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("vertex ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let request = server.await.expect("server task completes");
let request_line = request.lines().next().unwrap_or_default();
assert!(
request_line
.contains("/v1/projects/proj-1/locations/global/publishers/mistralai/models/mistral-ocr-maas:rawPredict"),
"{request}"
);
let authorization_count = request
.lines()
.filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
.count();
assert_eq!(authorization_count, 1, "{request}");
assert!(
request.contains("authorization: Bearer ya29.explicit-oauth-token")
|| request.contains("Authorization: Bearer ya29.explicit-oauth-token"),
"{request}"
);
}
#[tokio::test]
async fn vertex_ocr_rejects_google_api_key_shaped_token_before_calling_upstream() {
let mut optional_params = Map::new();
optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into()));
optional_params.insert(
"vertex_location".to_string(),
Value::String("global".into()),
);
let err = ocr(OcrRequest {
model: "mistral-ocr-maas",
document: json!({
"type": "image_url",
"image_url": "data:image/png;base64,abc"
}),
api_key: Some("AIzaSyExampleApiKeyValue000000000000000"),
api_base: Some("http://192.0.2.1:9"),
custom_llm_provider: "vertex_ai",
extra_headers: None,
optional_params,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect_err("google api key is rejected");
match err {
CoreError::Auth(message) => assert!(message.contains("OAuth"), "{message}"),
other => panic!("expected auth error, got {other:?}"),
}
}
#[tokio::test]
async fn vertex_ocr_rejects_google_api_key_in_caller_authorization_header_before_upstream() {
let mut optional_params = Map::new();
optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into()));
optional_params.insert(
"vertex_location".to_string(),
Value::String("global".into()),
);
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer AIzaSyExampleApiKeyValue000000000000000".to_string()),
);
let err = ocr(OcrRequest {
model: "mistral-ocr-maas",
document: json!({
"type": "image_url",
"image_url": "data:image/png;base64,abc"
}),
api_key: None,
api_base: Some("http://192.0.2.1:9"),
custom_llm_provider: "vertex_ai",
extra_headers: Some(headers),
optional_params,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect_err("google api key in authorization header is rejected");
match err {
CoreError::Auth(message) => assert!(message.contains("OAuth"), "{message}"),
other => panic!("expected auth error, got {other:?}"),
}
}
#[tokio::test]
async fn vertex_ocr_rejects_malformed_authorization_scheme_before_upstream() {
let mut optional_params = Map::new();
optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into()));
optional_params.insert(
"vertex_location".to_string(),
Value::String("global".into()),
);
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Basic dXNlcjpwYXNzd29yZA==".to_string()),
);
let err = ocr(OcrRequest {
model: "mistral-ocr-maas",
document: json!({
"type": "image_url",
"image_url": "data:image/png;base64,abc"
}),
api_key: None,
api_base: Some("http://192.0.2.1:9"),
custom_llm_provider: "vertex_ai",
extra_headers: Some(headers),
optional_params,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect_err("malformed authorization scheme is rejected");
match err {
CoreError::Auth(message) => assert!(message.contains("Bearer"), "{message}"),
other => panic!("expected auth error, got {other:?}"),
}
}
#[tokio::test]
async fn vertex_ocr_rejects_duplicate_authorization_headers_before_upstream() {
let mut optional_params = Map::new();
optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into()));
optional_params.insert(
"vertex_location".to_string(),
Value::String("global".into()),
);
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer ya29.first-token".to_string()),
);
headers.insert(
"authorization".to_string(),
Value::String("Bearer ya29.second-token".to_string()),
);
let err = ocr(OcrRequest {
model: "mistral-ocr-maas",
document: json!({
"type": "image_url",
"image_url": "data:image/png;base64,abc"
}),
api_key: None,
api_base: Some("http://192.0.2.1:9"),
custom_llm_provider: "vertex_ai",
extra_headers: Some(headers),
optional_params,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect_err("duplicate authorization headers are rejected");
match err {
CoreError::Auth(message) => assert!(message.contains("exactly one"), "{message}"),
other => panic!("expected auth error, got {other:?}"),
}
}

3
litellm-rust/crates/ai-gateway/src/io/vertex_ai/mod.rs Normal file
View file

@ -0,0 +1,3 @@
mod vertex_ai_base;
pub use vertex_ai_base::VertexAiBase;

556
litellm-rust/crates/ai-gateway/src/io/vertex_ai/vertex_ai_base.rs Normal file
View file

@ -0,0 +1,556 @@
//! Shared Vertex AI host-layer authentication.
//!
//! Rust counterpart of Python's `litellm/llms/vertex_ai/vertex_llm_base.py`
//! (`VertexBase`) scoped to what current routes need: resolve a credential
//! source (inline JSON, file path, `VERTEXAI_CREDENTIALS`, or ADC), build
//! Google credentials with the cloud-platform scope, cache them by content,
//! and mint OAuth access tokens. Any Vertex route (OCR today, others later)
//! should use this instead of owning Google auth.
use std::sync::{Once, OnceLock};
use google_cloud_auth::credentials::service_account::AccessSpecifier;
use google_cloud_auth::credentials::{
external_account, impersonated, service_account, user_account, AccessTokenCredentials,
Builder as AdcBuilder,
};
use litellm_core::cache::in_memory::InMemoryCache;
use litellm_core::error::CoreError;
use litellm_core::CoreResult;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use tokio::sync::Mutex;
use crate::config::resolve_env_reference;
use crate::constants::{
CLOUD_PLATFORM_SCOPE, VERTEXAI_CREDENTIALS_ENV, VERTEX_CREDENTIALS_CACHE_CAPACITY,
};
type CacheKey = [u8; 32];
fn ensure_crypto_provider() {
static INSTALL: Once = Once::new();
INSTALL.call_once(|| {
let _ = rustls::crypto::ring::default_provider().install_default();
});
}
fn adc_cache_key() -> CacheKey {
let mut hasher = Sha256::new();
hasher.update(b"adc");
hasher.finalize().into()
}
fn content_cache_key(contents: &str) -> CacheKey {
let mut hasher = Sha256::new();
hasher.update(b"inline:");
hasher.update(contents.trim().as_bytes());
hasher.finalize().into()
}
pub struct VertexAiBase {
cache: Mutex<InMemoryCache<CacheKey, AccessTokenCredentials>>,
}
impl VertexAiBase {
pub fn new() -> Self {
Self {
cache: Mutex::new(InMemoryCache::new(VERTEX_CREDENTIALS_CACHE_CAPACITY)),
}
}
pub fn shared() -> &'static VertexAiBase {
static SHARED: OnceLock<VertexAiBase> = OnceLock::new();
SHARED.get_or_init(VertexAiBase::new)
}
/// Mirrors Python `VertexBase.safe_get_vertex_ai_credentials`: request
/// params (`vertex_credentials`, then `vertex_ai_credentials`) take
/// precedence over the `VERTEXAI_CREDENTIALS` environment variable.
pub fn resolve_credential_source(
optional_params: &Map<String, Value>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Option<String> {
Self::credential_source_param(optional_params)
.and_then(|source| resolve_env_reference(Some(&source), env_lookup))
.or_else(|| {
env_lookup(VERTEXAI_CREDENTIALS_ENV)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
})
}
fn credential_source_param(optional_params: &Map<String, Value>) -> Option<String> {
["vertex_credentials", "vertex_ai_credentials"]
.iter()
.find_map(|key| optional_params.get(*key))
.and_then(|value| match value {
Value::String(raw) => {
let trimmed = raw.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
}
Value::Object(_) => Some(value.to_string()),
_ => None,
})
}
/// Mirrors Python `VertexBase.get_access_token`: load credentials from the
/// given source (or ADC when absent), cache them by content, and return an
/// OAuth access token. The Google auth crate refreshes expired tokens
/// internally, matching Python's cached-credential refresh behavior.
pub async fn get_access_token(&self, credentials: Option<&str>) -> CoreResult<String> {
ensure_crypto_provider();
let token = self
.resolve_credentials(credentials)
.await?
.access_token()
.await
.map_err(|_| CoreError::Auth("Failed to obtain Vertex access token".to_string()))?;
Ok(token.token)
}
async fn resolve_credentials(
&self,
credentials: Option<&str>,
) -> CoreResult<AccessTokenCredentials> {
let (key, built) = match credentials {
None => (adc_cache_key(), None),
Some(raw) => {
let contents = load_credentials_contents(raw).await?;
(content_cache_key(&contents), Some(contents))
}
};
if let Some(existing) = self.cache.lock().await.get(&key) {
return Ok(existing);
}
let credentials = match built {
None => build_adc_credentials()?,
Some(contents) => build_from_json(parse_credentials_json(&contents)?)?,
};
Ok(self.cache.lock().await.get_or_insert(key, credentials))
}
#[cfg(test)]
async fn cached_credential_count(&self) -> usize {
self.cache.lock().await.len()
}
}
impl Default for VertexAiBase {
fn default() -> Self {
Self::new()
}
}
fn build_adc_credentials() -> CoreResult<AccessTokenCredentials> {
AdcBuilder::default()
.with_scopes([CLOUD_PLATFORM_SCOPE])
.build_access_token_credentials()
.map_err(|_| CoreError::Auth("Failed to load Vertex ADC credentials".to_string()))
}
async fn load_credentials_contents(raw: &str) -> CoreResult<String> {
let trimmed = raw.trim();
if trimmed.starts_with('{') {
Ok(trimmed.to_string())
} else {
tokio::fs::read_to_string(trimmed)
.await
.map_err(|_| CoreError::Auth("Failed to read Vertex credentials file".to_string()))
}
}
fn parse_credentials_json(contents: &str) -> CoreResult<Value> {
serde_json::from_str(contents.trim())
.map_err(|_| CoreError::Auth("Vertex credentials are not valid JSON".to_string()))
}
fn build_from_json(json: Value) -> CoreResult<AccessTokenCredentials> {
let scopes = [CLOUD_PLATFORM_SCOPE];
let credentials = match json.get("type").and_then(Value::as_str) {
Some("service_account") => service_account::Builder::new(json)
.with_access_specifier(AccessSpecifier::from_scopes(scopes))
.build_access_token_credentials(),
Some("authorized_user") => user_account::Builder::new(json)
.with_scopes(scopes)
.build_access_token_credentials(),
Some("external_account") => external_account::Builder::new(json)
.with_scopes(scopes)
.build_access_token_credentials(),
Some("impersonated_service_account") => impersonated::Builder::new(json)
.with_scopes(scopes)
.build_access_token_credentials(),
Some(_) => {
return Err(CoreError::Auth(
"Unsupported Vertex credential type".to_string(),
))
}
None => {
return Err(CoreError::Auth(
"Vertex credentials JSON is missing the required `type` field".to_string(),
))
}
};
credentials.map_err(|_| CoreError::Auth("Failed to load Vertex credentials".to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::sync::Arc;
fn inline_service_account_json() -> String {
json!({
"type": "service_account",
"project_id": "proj-1",
"private_key_id": "key-id",
"private_key": "-----BEGIN PRIVATE KEY-----\nnot-a-real-key\n-----END PRIVATE KEY-----\n",
"client_email": "sa@proj-1.iam.gserviceaccount.com"
})
.to_string()
}
#[test]
fn content_cache_key_distinguishes_adc_from_inline_and_matches_on_repeat() {
let adc = adc_cache_key();
let inline = content_cache_key("{\"type\":\"service_account\"}");
assert_ne!(adc, inline);
assert_eq!(
inline,
content_cache_key(" {\"type\":\"service_account\"} ")
);
}
#[test]
fn content_cache_key_differs_for_distinct_credential_sources() {
let adc = adc_cache_key();
let first = content_cache_key("{\"type\":\"service_account\",\"client_email\":\"a\"}");
let second = content_cache_key("{\"type\":\"service_account\",\"client_email\":\"b\"}");
assert_ne!(first, second);
assert_ne!(adc, first);
assert_ne!(adc, second);
}
#[tokio::test]
async fn file_backed_credentials_key_tracks_content_not_path() {
let path = std::env::temp_dir().join(format!(
"vertex-cred-{}-{:?}.json",
std::process::id(),
std::thread::current().id()
));
tokio::fs::write(
&path,
b"{\"type\":\"service_account\",\"client_email\":\"old\"}",
)
.await
.expect("writes first credential file");
let first = load_credentials_contents(path.to_str().expect("utf-8 path"))
.await
.expect("reads first credential file");
tokio::fs::write(
&path,
b"{\"type\":\"service_account\",\"client_email\":\"new\"}",
)
.await
.expect("rotates credential file");
let second = load_credentials_contents(path.to_str().expect("utf-8 path"))
.await
.expect("reads rotated credential file");
tokio::fs::remove_file(&path).await.ok();
assert_ne!(
content_cache_key(&first),
content_cache_key(&second),
"rotating a credential file at the same path must not reuse the old cache entry"
);
}
#[tokio::test]
async fn resolve_credentials_caches_inline_service_account_per_instance() {
let base = VertexAiBase::new();
let inline = inline_service_account_json();
base.resolve_credentials(Some(&inline))
.await
.expect("first resolve builds credentials");
base.resolve_credentials(Some(&inline))
.await
.expect("second resolve reuses cache");
assert_eq!(base.cached_credential_count().await, 1);
}
#[tokio::test]
async fn resolve_credentials_reads_service_account_from_file() {
let path = std::env::temp_dir().join(format!(
"vertex-file-cred-{}-{:?}.json",
std::process::id(),
std::thread::current().id()
));
tokio::fs::write(&path, inline_service_account_json())
.await
.expect("writes credential file");
let base = VertexAiBase::new();
base.resolve_credentials(Some(path.to_str().expect("utf-8 path")))
.await
.expect("file-backed service account resolves");
tokio::fs::remove_file(&path).await.ok();
assert_eq!(base.cached_credential_count().await, 1);
}
#[tokio::test]
async fn resolve_credentials_single_flight_under_concurrency() {
let base = Arc::new(VertexAiBase::new());
let inline = Arc::new(inline_service_account_json());
let tasks: Vec<_> = (0..16)
.map(|_| {
let base = Arc::clone(&base);
let inline = Arc::clone(&inline);
tokio::spawn(async move {
base.resolve_credentials(Some(inline.as_str()))
.await
.expect("concurrent resolve succeeds");
})
})
.collect();
for task in tasks {
task.await.expect("task completes");
}
assert_eq!(base.cached_credential_count().await, 1);
}
#[tokio::test]
async fn build_from_json_builds_service_account_credentials() {
build_from_json(json!({
"type": "service_account",
"project_id": "proj-1",
"private_key_id": "key-id",
"private_key": "-----BEGIN PRIVATE KEY-----\nnot-a-real-key\n-----END PRIVATE KEY-----\n",
"client_email": "sa@proj-1.iam.gserviceaccount.com"
}))
.expect("service account dispatches and builds");
}
#[tokio::test]
async fn build_from_json_builds_authorized_user_credentials() {
build_from_json(json!({
"type": "authorized_user",
"client_id": "client-id.apps.googleusercontent.com",
"client_secret": "client-secret",
"refresh_token": "refresh-token"
}))
.expect("authorized user dispatches and builds");
}
#[tokio::test]
async fn build_from_json_builds_impersonated_service_account_credentials() {
build_from_json(json!({
"type": "impersonated_service_account",
"service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/target@proj-1.iam.gserviceaccount.com:generateAccessToken",
"source_credentials": {
"type": "authorized_user",
"client_id": "client-id.apps.googleusercontent.com",
"client_secret": "client-secret",
"refresh_token": "refresh-token"
}
}))
.expect("impersonated service account dispatches and builds");
}
#[tokio::test]
async fn build_from_json_builds_external_account_for_all_standard_source_mechanisms() {
let base = |source: Value| {
json!({
"type": "external_account",
"audience": "//iam.googleapis.com/projects/1/locations/global/workloadIdentityPools/p/providers/pr",
"subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
"token_url": "https://sts.googleapis.com/v1/token",
"credential_source": source
})
};
build_from_json(base(json!({"file": "/var/run/secrets/token"})))
.expect("file-sourced external account builds");
build_from_json(base(json!({
"url": "https://169.254.169.254/token",
"headers": {"Metadata": "true"},
"format": {"type": "json", "subject_token_field_name": "access_token"}
})))
.expect("url-sourced external account builds");
build_from_json(base(json!({
"executable": {"command": "/usr/bin/token-helper", "timeout_millis": 5000}
})))
.expect("executable-sourced external account builds");
build_from_json(base(json!({
"environment_id": "aws1",
"region_url": "http://169.254.169.254/latest/meta-data/placement/availability-zone",
"regional_cred_verification_url": "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15"
})))
.expect("aws-sourced external account builds");
}
#[test]
fn build_from_json_rejects_unknown_credential_type() {
let err =
build_from_json(json!({"type": "totally_made_up"})).expect_err("unknown type rejected");
assert!(matches!(err, CoreError::Auth(_)), "{err:?}");
}
#[test]
fn build_from_json_requires_type_field() {
let err = build_from_json(json!({"client_email": "x"})).expect_err("missing type rejected");
assert!(matches!(err, CoreError::Auth(_)), "{err:?}");
}
#[test]
fn build_from_json_unknown_type_error_omits_attacker_controlled_type() {
let err = build_from_json(json!({"type": "attacker-controlled-type-string"}))
.expect_err("unknown type rejected");
match err {
CoreError::Auth(message) => {
assert!(
!message.contains("attacker-controlled-type-string"),
"{message}"
);
}
other => panic!("expected auth error, got {other:?}"),
}
}
#[test]
fn parse_credentials_json_reports_invalid_json_without_echoing_contents() {
let err = parse_credentials_json("{not-valid-json-secret-value")
.expect_err("invalid json rejected");
match err {
CoreError::Auth(message) => {
assert!(
!message.contains("not-valid-json-secret-value"),
"{message}"
);
}
other => panic!("expected auth error, got {other:?}"),
}
}
#[tokio::test]
async fn load_credentials_contents_missing_file_error_omits_attacker_controlled_path() {
let err = load_credentials_contents("/attacker/controlled/secret-credentials-path.json")
.await
.expect_err("missing file rejected");
match err {
CoreError::Auth(message) => {
assert!(!message.contains("attacker"), "{message}");
assert!(!message.contains("secret-credentials-path"), "{message}");
}
other => panic!("expected auth error, got {other:?}"),
}
}
#[test]
fn resolve_credential_source_reads_string_object_and_treats_blank_as_absent() {
let mut inline = Map::new();
inline.insert(
"vertex_credentials".to_string(),
Value::String(" /path/to/sa.json ".into()),
);
assert_eq!(
VertexAiBase::resolve_credential_source(&inline, &|_| None).as_deref(),
Some("/path/to/sa.json")
);
let mut object = Map::new();
object.insert(
"vertex_credentials".to_string(),
json!({"type": "service_account"}),
);
assert_eq!(
VertexAiBase::resolve_credential_source(&object, &|_| None).as_deref(),
Some("{\"type\":\"service_account\"}")
);
let mut blank = Map::new();
blank.insert(
"vertex_credentials".to_string(),
Value::String(" ".into()),
);
assert_eq!(
VertexAiBase::resolve_credential_source(&blank, &|_| None),
None
);
assert_eq!(
VertexAiBase::resolve_credential_source(&Map::new(), &|_| None),
None
);
}
#[test]
fn resolve_credential_source_prefers_optional_param_then_env() {
let mut params = Map::new();
params.insert(
"vertex_credentials".to_string(),
Value::String("/from/param.json".into()),
);
let env =
|key: &str| (key == VERTEXAI_CREDENTIALS_ENV).then(|| "/from/env.json".to_string());
assert_eq!(
VertexAiBase::resolve_credential_source(&params, &env).as_deref(),
Some("/from/param.json")
);
assert_eq!(
VertexAiBase::resolve_credential_source(&Map::new(), &env).as_deref(),
Some("/from/env.json")
);
let blank_env = |key: &str| (key == VERTEXAI_CREDENTIALS_ENV).then(|| " ".to_string());
assert_eq!(
VertexAiBase::resolve_credential_source(&Map::new(), &blank_env),
None
);
assert_eq!(
VertexAiBase::resolve_credential_source(&Map::new(), &|_| None),
None
);
}
#[test]
fn resolve_credential_source_falls_back_through_alias_param() {
let mut params = Map::new();
params.insert(
"vertex_ai_credentials".to_string(),
Value::String("/from/alias.json".into()),
);
assert_eq!(
VertexAiBase::resolve_credential_source(&params, &|_| None).as_deref(),
Some("/from/alias.json")
);
}
#[test]
fn resolve_credential_source_resolves_exact_environment_reference() {
let params = Map::from_iter([(
"vertex_credentials".to_string(),
Value::String("os.environ/CUSTOM_VERTEX_CREDENTIALS".into()),
)]);
let env = |key: &str| {
(key == "CUSTOM_VERTEX_CREDENTIALS").then(|| "{\"type\":\"authorized_user\"}".into())
};
assert_eq!(
VertexAiBase::resolve_credential_source(&params, &env).as_deref(),
Some("{\"type\":\"authorized_user\"}")
);
}
#[test]
fn resolve_credential_source_treats_unresolved_reference_as_absent() {
let params = Map::from_iter([(
"vertex_credentials".to_string(),
Value::String("os.environ/MISSING_VERTEX_CREDENTIALS".into()),
)]);
assert_eq!(
VertexAiBase::resolve_credential_source(&params, &|_| None),
None
);
}
}

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

@ -4,6 +4,7 @@ version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
rust-version.workspace = true
[dependencies]
rand.workspace = true

89
litellm-rust/crates/core/src/cache/in_memory.rs vendored Normal file
View file

@ -0,0 +1,89 @@
use std::collections::{HashMap, VecDeque};
use std::hash::Hash;
pub struct InMemoryCache<K, V> {
capacity: usize,
entries: HashMap<K, V>,
order: VecDeque<K>,
}
impl<K, V> InMemoryCache<K, V>
where
K: Clone + Eq + Hash,
V: Clone,
{
pub fn new(capacity: usize) -> Self {
Self {
capacity: capacity.max(1),
entries: HashMap::new(),
order: VecDeque::new(),
}
}
pub fn get(&self, key: &K) -> Option<V> {
self.entries.get(key).cloned()
}
pub fn get_or_insert(&mut self, key: K, value: V) -> V {
if let Some(existing) = self.entries.get(&key) {
return existing.clone();
}
while self.entries.len() >= self.capacity {
let Some(evicted) = self.order.pop_front() else {
break;
};
self.entries.remove(&evicted);
}
self.order.push_back(key.clone());
self.entries.insert(key, value.clone());
value
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::InMemoryCache;
#[test]
fn evicts_oldest_entry_at_capacity() {
let mut cache = InMemoryCache::new(2);
cache.get_or_insert(1, "one");
cache.get_or_insert(2, "two");
cache.get_or_insert(3, "three");
assert_eq!(cache.get(&1), None);
assert_eq!(cache.get(&2), Some("two"));
assert_eq!(cache.get(&3), Some("three"));
assert_eq!(cache.len(), 2);
}
#[test]
fn reuses_existing_entry_without_replacement() {
let mut cache = InMemoryCache::new(2);
assert!(cache.is_empty());
assert_eq!(cache.get_or_insert(1, "first"), "first");
assert_eq!(cache.get_or_insert(1, "replacement"), "first");
assert_eq!(cache.get(&1), Some("first"));
assert_eq!(cache.len(), 1);
}
#[test]
fn zero_capacity_still_retains_one_entry() {
let mut cache = InMemoryCache::new(0);
cache.get_or_insert(1, "one");
cache.get_or_insert(2, "two");
assert_eq!(cache.get(&1), None);
assert_eq!(cache.get(&2), Some("two"));
assert_eq!(cache.len(), 1);
}
}

1
litellm-rust/crates/core/src/cache/mod.rs vendored Normal file
View file

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

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

@ -0,0 +1,4 @@
pub(crate) const VERTEX_GLOBAL_LOCATION: &str = "global";
pub(crate) const VERTEX_GLOBAL_API_BASE: &str = "https://aiplatform.googleapis.com";
pub(crate) const GOOGLE_API_KEY_PREFIX: &str = "AIza";
pub(crate) const BEARER_SCHEME: &str = "Bearer";

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

@ -1,3 +1,5 @@
pub mod cache;
pub(crate) mod constants;
pub mod error;
pub mod ocr;
pub mod providers;

20
litellm-rust/crates/core/src/ocr/transformation.rs
View file

@ -27,6 +27,12 @@ pub enum OcrResponseHandling {
AzureDocumentIntelligencePoll,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrAuth {
ProviderKey,
VertexOauth,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrDocumentPreparation {
None,
@ -72,9 +78,17 @@ pub trait OcrProviderConfig: Sync {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String>;
_api_key: Option<&str>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
Err(crate::error::CoreError::Auth(
"provider does not use direct api-key auth".to_string(),
))
}
fn ocr_auth(&self) -> OcrAuth {
OcrAuth::ProviderKey
}
fn auth_strategy(&self) -> OcrAuthStrategy {
OcrAuthStrategy::Bearer

274
litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs
View file

@ -1,12 +1,15 @@
use crate::constants::{
BEARER_SCHEME, GOOGLE_API_KEY_PREFIX, VERTEX_GLOBAL_API_BASE, VERTEX_GLOBAL_LOCATION,
};
use crate::error::{json_type_name, CoreError, CoreResult};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::transformation::{OcrAuth, OcrProviderConfig};
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use serde_json::{json, Map, Value};
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_DEFAULT_DEEPSEEK_API_BASE: &str = VERTEX_GLOBAL_API_BASE;
const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY";
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
@ -40,22 +43,76 @@ pub fn is_deepseek_model(model: &str) -> bool {
model.to_ascii_lowercase().contains("deepseek")
}
pub fn resolve_vertex_api_key(
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum VertexTokenSource {
Explicit(String),
Mint,
}
fn is_google_api_key(token: &str) -> bool {
token.starts_with(GOOGLE_API_KEY_PREFIX)
}
fn google_api_key_not_oauth_error() -> CoreError {
CoreError::Auth(
"Received a Google API key (AIza...) for Vertex AI, which is not an OAuth access token. \
Provide service-account credentials/ADC or an OAuth access token instead"
.to_string(),
)
}
pub fn classify_vertex_bearer(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
api_key
) -> CoreResult<VertexTokenSource> {
let token = api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.ok_or_else(|| {
CoreError::Auth(
"Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers"
.to_string(),
)
})
.or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty()));
match token {
Some(token) if is_google_api_key(token.trim()) => Err(google_api_key_not_oauth_error()),
Some(token) => Ok(VertexTokenSource::Explicit(token.trim().to_string())),
None => Ok(VertexTokenSource::Mint),
}
}
fn malformed_vertex_authorization_error() -> CoreError {
CoreError::Auth(
"Vertex AI requires exactly one `Authorization: Bearer <OAuth access token>` header. \
Provide a valid OAuth Bearer token, or omit the header to mint one from credentials/ADC"
.to_string(),
)
}
pub fn validate_vertex_authorization_headers(values: &[&str]) -> CoreResult<()> {
match values {
[] => Ok(()),
[single] => validate_vertex_authorization_value(single),
_ => Err(malformed_vertex_authorization_error()),
}
}
fn validate_vertex_authorization_value(header_value: &str) -> CoreResult<()> {
let mut parts = header_value.split_whitespace();
let scheme = parts
.next()
.ok_or_else(malformed_vertex_authorization_error)?;
let token = parts
.next()
.ok_or_else(malformed_vertex_authorization_error)?;
if parts.next().is_some() {
return Err(malformed_vertex_authorization_error());
}
if !scheme.eq_ignore_ascii_case(BEARER_SCHEME) {
return Err(malformed_vertex_authorization_error());
}
if is_google_api_key(token) {
return Err(google_api_key_not_oauth_error());
}
Ok(())
}
fn vertex_project(
@ -84,12 +141,22 @@ fn vertex_location(
.unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string())
}
fn vertex_base_url(location: &str) -> String {
match location {
VERTEX_GLOBAL_LOCATION => VERTEX_GLOBAL_API_BASE.to_string(),
location if !location.contains('-') => {
format!("https://aiplatform.{location}.rep.googleapis.com")
}
location => format!("https://{location}-aiplatform.googleapis.com"),
}
}
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"))
.unwrap_or_else(|| vertex_base_url(location))
.trim_end_matches('/')
.to_string()
}
@ -241,12 +308,8 @@ impl OcrProviderConfig for VertexAiOcrConfig {
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>,
) -> CoreResult<String> {
resolve_vertex_api_key(api_key, env_lookup)
fn ocr_auth(&self) -> OcrAuth {
OcrAuth::VertexOauth
}
fn requires_data_uri_document(&self) -> bool {
@ -346,12 +409,8 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
complete_vertex_deepseek_url(api_base, optional_params, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_vertex_api_key(api_key, env_lookup)
fn ocr_auth(&self) -> OcrAuth {
OcrAuth::VertexOauth
}
}
@ -375,6 +434,171 @@ mod tests {
);
}
#[test]
fn vertex_mistral_url_uses_global_host_without_region_prefix() {
let params = Map::from_iter([
("vertex_project".to_string(), json!("proj-1")),
("vertex_location".to_string(), json!("global")),
]);
let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", &params, &|_| None)
.expect("url builds");
assert_eq!(
url,
"https://aiplatform.googleapis.com/v1/projects/proj-1/locations/global/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
}
#[test]
fn vertex_mistral_url_uses_residency_host_for_single_token_location() {
let params = Map::from_iter([
("vertex_project".to_string(), json!("proj-1")),
("vertex_location".to_string(), json!("eu")),
]);
let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", &params, &|_| None)
.expect("url builds");
assert_eq!(
url,
"https://aiplatform.eu.rep.googleapis.com/v1/projects/proj-1/locations/eu/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
}
#[test]
fn vertex_mistral_url_prefers_explicit_api_base() {
let params = Map::from_iter([
("vertex_project".to_string(), json!("proj-1")),
("vertex_location".to_string(), json!("global")),
]);
let url = complete_vertex_mistral_url(
Some("https://custom.example.com/"),
"mistral-ocr-maas",
&params,
&|_| None,
)
.expect("url builds");
assert_eq!(
url,
"https://custom.example.com/v1/projects/proj-1/locations/global/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
}
#[test]
fn classify_vertex_bearer_uses_explicit_oauth_token() {
let source = classify_vertex_bearer(Some(" ya29.oauth-token "), &|_| None)
.expect("token classifies");
assert_eq!(
source,
VertexTokenSource::Explicit("ya29.oauth-token".to_string())
);
}
#[test]
fn classify_vertex_bearer_reads_oauth_token_from_env() {
let source = classify_vertex_bearer(None, &|key| {
(key == VERTEX_AI_API_KEY_ENV).then(|| "ya29.from-env".to_string())
})
.expect("token classifies");
assert_eq!(
source,
VertexTokenSource::Explicit("ya29.from-env".to_string())
);
}
#[test]
fn classify_vertex_bearer_mints_when_no_token_supplied() {
let source = classify_vertex_bearer(None, &|_| None).expect("token classifies");
assert_eq!(source, VertexTokenSource::Mint);
}
#[test]
fn classify_vertex_bearer_rejects_google_api_key() {
let err = classify_vertex_bearer(Some("AIzaSyExampleApiKeyValue"), &|_| None)
.expect_err("google api key is rejected");
assert!(matches!(err, CoreError::Auth(_)), "{err:?}");
}
#[test]
fn classify_vertex_bearer_rejects_google_api_key_from_env() {
let err = classify_vertex_bearer(None, &|key| {
(key == VERTEXAI_API_KEY_ENV).then(|| "AIzaSyExampleApiKeyValue".to_string())
})
.expect_err("google api key from env is rejected");
assert!(matches!(err, CoreError::Auth(_)), "{err:?}");
}
#[test]
fn validate_vertex_authorization_headers_rejects_api_key_bearer() {
for header in [
"Bearer AIzaSyExampleApiKeyValue",
" bearer AIzaSyExampleApiKeyValue ",
"BEARER AIzaSyExampleApiKeyValue",
] {
let err = validate_vertex_authorization_headers(&[header])
.expect_err(&format!("api key bearer rejected: {header:?}"));
assert!(matches!(err, CoreError::Auth(_)), "{header:?} -> {err:?}");
}
}
#[test]
fn validate_vertex_authorization_headers_rejects_malformed_values() {
for header in [
"",
" ",
"AIzaSyExampleApiKeyValue",
"ya29.raw-token-without-scheme",
"Basic dXNlcjpwYXNz",
"Token ya29.some-token",
"Bearer2 ya29.token",
"Bearer",
"Bearer ",
"Bearer ya29.token extra-part",
"Bearer ya29.token AIzaExtra",
] {
let err = validate_vertex_authorization_headers(&[header])
.expect_err(&format!("expected rejection for {header:?}"));
assert!(matches!(err, CoreError::Auth(_)), "{header:?} -> {err:?}");
}
}
#[test]
fn validate_vertex_authorization_headers_rejects_duplicate_headers() {
let err = validate_vertex_authorization_headers(&[
"Bearer ya29.first-token",
"Bearer ya29.second-token",
])
.expect_err("duplicate authorization headers rejected");
assert!(matches!(err, CoreError::Auth(_)), "{err:?}");
}
#[test]
fn validate_vertex_authorization_headers_allows_single_oauth_bearer() {
for header in [
"Bearer ya29.real-oauth-token",
"bearer ya29.real-oauth-token",
"BEARER ya29.real-oauth-token",
" Bearer ya29.real-oauth-token ",
] {
validate_vertex_authorization_headers(&[header])
.unwrap_or_else(|err| panic!("oauth bearer allowed: {header:?} -> {err:?}"));
}
validate_vertex_authorization_headers(&[])
.expect("no authorization header defers to minting");
}
#[test]
fn vertex_configs_use_google_oauth() {
assert_eq!(VERTEX_AI_OCR_CONFIG.ocr_auth(), OcrAuth::VertexOauth);
assert_eq!(
VERTEX_AI_DEEPSEEK_OCR_CONFIG.ocr_auth(),
OcrAuth::VertexOauth
);
}
#[test]
fn vertex_mistral_reuses_mistral_body_transform() {
let body = VERTEX_AI_OCR_CONFIG

1
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -4,6 +4,7 @@ version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
rust-version.workspace = true
[lib]
name = "_native"