scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

feat(ocr): move provider auth into Rust

Browse source
This commit is contained in:
Yujong Lee 2026-09-04 10:45:13 -07:00
parent 41cc4b6020
commit 98f05d52d7
29 changed files with 2282 additions and 229 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
basedpyright-code-budget.json
View file

@ -57,7 +57,7 @@
"limit": 5601
},
"reportMissingTypeArgument": {
"limit": 15285
"limit": 15284
},
"reportMissingTypeStubs": {
"limit": 40
@ -108,7 +108,7 @@
"limit": 38309
},
"reportUnknownParameterType": {
"limit": 19622
"limit": 19621
},
"reportUnknownVariableType": {
"limit": 29846

929
litellm-rust/Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

13
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -88,11 +88,14 @@ impl OcrLifecycleHooks {
) -> Result<ProviderOcrRequest, Error> {
let config = request.config?;
let env_lookup = |key: &str| std::env::var(key).ok();
let upstream_headers = config.validate_environment(
string_headers(request.extra_headers)?,
request.api_key.as_deref(),
&env_lookup,
)?;
let (upstream_headers, _) = config
.validate_environment(
string_headers(request.extra_headers)?,
request.api_key.as_deref(),
&request.optional_params,
&env_lookup,
)
.await?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,

6
litellm-rust/crates/core/CLAUDE.md
View file

@ -30,6 +30,12 @@ Not allowed:
- Provider-specific branching that belongs in `providers`.
- Panics for user/provider-controlled input.
Provider authentication is the narrow exception to the filesystem and global
state rules. Like Bedrock credential discovery, ADC and managed identity may
read SDK-defined credential files or contact metadata services. Credential and
token caches may be process-wide, must be secret-safe, and may contain auth
state only.
## Typed Contracts (core rule)
Trait and function boundaries MUST be strongly typed. No stringly-typed JSON

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

@ -8,12 +8,13 @@ repository.workspace = true
[dependencies]
async-trait = "0.1"
base64.workspace = true
google-cloud-auth = { version = "=1.15.0", default-features = false, features = ["default-rustls-provider"] }
rand.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["net", "time"] }
tokio = { workspace = true, features = ["net", "sync", "time"] }
tracing.workspace = true
tokio = { workspace = true, features = ["sync", "rt"] }
tracing-subscriber = { workspace = true, optional = true }

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

@ -34,6 +34,12 @@ pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600;
pub(crate) const OCR_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
pub(crate) const OCR_MAX_SAFE_FETCH_REDIRECTS: usize = 10;
pub(crate) const AZURE_AUTH_DEFAULT_SCOPE: &str = "https://cognitiveservices.azure.com/.default";
pub(crate) const AZURE_AUTH_TOKEN_REFRESH_WINDOW_SECS: u64 = 120;
pub(crate) const VERTEX_AUTH_CLOUD_PLATFORM_SCOPE: &str =
"https://www.googleapis.com/auth/cloud-platform";
pub(crate) const VERTEX_AUTH_SDK_TOKEN_CACHE_TTL_SECS: u64 = 30;
pub(crate) const VERTEX_AUTH_TOKEN_REFRESH_WINDOW_SECS: u64 = 5;
/// `object` field every non-streaming chat completion response carries.
pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";

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

@ -1,8 +1,9 @@
use crate::auth::ResolvedAuth;
use crate::error::Error;
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use super::transformation::AnthropicMessagesProviderConfig;
use super::types::{MessagesRequest, ProviderMessagesRequest};
use serde_json::{Map, Value};
@ -70,13 +71,7 @@ fn validate_environment(
|| (config.accepts_bearer_auth() && has_bearer_auth(&headers));
if !already_authorized {
let api_key = config.resolve_api_key(api_key, env_lookup)?;
let auth_header = match auth_strategy {
MessagesAuthStrategy::Bearer => {
("authorization".to_string(), format!("Bearer {api_key}"))
}
MessagesAuthStrategy::Header(name) => (name.to_string(), api_key),
};
headers.push(auth_header);
headers.push(ResolvedAuth::from_credential(auth_strategy, api_key).credential_header());
}
for (name, value) in config.default_headers() {

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

@ -1,20 +1,6 @@
use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
use crate::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MessagesAuthStrategy {
Bearer,
Header(&'static str),
}
impl MessagesAuthStrategy {
pub fn header_name(self) -> &'static str {
match self {
Self::Bearer => "authorization",
Self::Header(header_name) => header_name,
}
}
}
use crate::auth::AuthHeaderKind;
pub trait AnthropicMessagesProviderConfig: Sync {
fn complete_url(
@ -30,8 +16,8 @@ pub trait AnthropicMessagesProviderConfig: Sync {
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn auth_strategy(&self) -> MessagesAuthStrategy {
MessagesAuthStrategy::Header("x-api-key")
fn auth_strategy(&self) -> AuthHeaderKind {
AuthHeaderKind::Header("x-api-key")
}
fn accepts_bearer_auth(&self) -> bool {

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

@ -6,6 +6,7 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use reqwest::Url;
use serde_json::{Map, Value};
use crate::auth::ResolvedAuth;
use crate::constants::{OCR_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS, OCR_MAX_SAFE_FETCH_REDIRECTS};
use crate::error::Error;
use crate::http_utils::truncate_error_body;
@ -322,7 +323,7 @@ pub(super) async fn poll_document_intelligence(
client: &reqwest::Client,
operation_url: &str,
original_url: &str,
headers: &[(String, String)],
auth: &ResolvedAuth,
timeout: Option<Duration>,
) -> Result<Value, Error> {
if !same_origin(operation_url, original_url) {
@ -343,12 +344,8 @@ pub(super) async fn poll_document_intelligence(
)));
}
let mut request_builder = client.get(operation_url);
for (key, value) in headers {
if key.eq_ignore_ascii_case("ocp-apim-subscription-key") {
request_builder = request_builder.header(key, value);
}
}
let (name, value) = auth.credential_header();
let request_builder = client.get(operation_url).header(name, value);
let response = request_builder
.send()
.await

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

@ -44,11 +44,14 @@ pub(crate) async fn execute_ocr_provider_call(
max_document_download_bytes,
} = request;
let env_lookup = |key: &str| std::env::var(key).ok();
let upstream_headers = config.validate_environment(
string_headers(extra_headers)?,
api_key.as_deref(),
&env_lookup,
)?;
let (upstream_headers, resolved_auth) = config
.validate_environment(
string_headers(extra_headers)?,
api_key.as_deref(),
&url_params,
&env_lookup,
)
.await?;
let url = config.complete_url(api_base.as_deref(), &model, &url_params, &env_lookup)?;
let document = if config.requires_data_uri_document() {
convert_document_url_to_data_uri(client, document, max_document_download_bytes).await?
@ -97,7 +100,7 @@ pub(crate) async fn execute_ocr_provider_call(
)
})?;
let response_json =
poll_document_intelligence(client, &operation_url, &url, &upstream_headers, timeout)
poll_document_intelligence(client, &operation_url, &url, &resolved_auth, timeout)
.await?;
return public_response(config, &model, &optional_params, response_json);
}

134
litellm-rust/crates/core/src/ocr/tests.rs
View file

@ -50,8 +50,21 @@ async fn read_http_request(socket: &mut TcpStream) -> String {
}
async fn write_response(socket: &mut TcpStream, status: &str, body: &str) {
write_response_with_headers(socket, status, &[], body).await;
}
async fn write_response_with_headers(
socket: &mut TcpStream,
status: &str,
headers: &[(&str, &str)],
body: &str,
) {
let extra_headers = headers
.iter()
.map(|(name, value)| format!("{name}: {value}\r\n"))
.collect::<String>();
let response = format!(
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
"HTTP/1.1 {status}\r\n{extra_headers}content-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
@ -61,6 +74,54 @@ async fn write_response(socket: &mut TcpStream, status: &str, body: &str) {
.expect("writes response");
}
async fn document_intelligence_round_trip(
api_key: Option<&str>,
azure_ad_token: Option<&str>,
) -> (String, String) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let operation_url = format!("http://{address}/operations/1");
let server = tokio::spawn(async move {
let (mut post_socket, _) = listener.accept().await.expect("accepts POST");
let post = read_http_request(&mut post_socket).await;
write_response_with_headers(
&mut post_socket,
"202 Accepted",
&[("operation-location", &operation_url), ("retry-after", "0")],
"",
)
.await;
let (mut poll_socket, _) = listener.accept().await.expect("accepts poll");
let poll = read_http_request(&mut poll_socket).await;
write_response(
&mut poll_socket,
"200 OK",
r#"{"status":"succeeded","analyzeResult":{"content":"hello","pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]}}"#,
)
.await;
(post, poll)
});
let api_base = format!("http://{address}");
let mut request = base_request("azure_ai/doc-intelligence/prebuilt-layout", &api_base);
request.api_key = api_key;
request.extra_headers = Some(Map::from_iter([(
"x-request-only".to_string(),
json!("must-not-poll"),
)]));
request.optional_params = azure_ad_token
.map(|token| Map::from_iter([("azure_ad_token".to_string(), json!(token))]))
.unwrap_or_default();
let client = test_client();
ocr(&client, request)
.await
.expect("Document Intelligence OCR succeeds");
server.await.expect("server task completes")
}
fn base_request<'a>(model: &'a str, api_base: &'a str) -> OcrRequest<'a> {
OcrRequest {
model,
@ -168,6 +229,77 @@ async fn vertex_routing_params_are_not_forwarded_in_provider_body() {
assert!(!provider_request.contains("vertex_location"));
}
#[tokio::test]
async fn document_intelligence_subscription_key_is_reused_for_polling_only() {
let (post, poll) = document_intelligence_round_trip(Some("subscription-key"), None).await;
let post = post.to_ascii_lowercase();
let poll = poll.to_ascii_lowercase();
assert!(post.contains("ocp-apim-subscription-key: subscription-key"));
assert!(poll.contains("ocp-apim-subscription-key: subscription-key"));
assert!(!post.contains("authorization:"));
assert!(!poll.contains("authorization:"));
assert!(post.contains("x-request-only: must-not-poll"));
assert!(!poll.contains("x-request-only:"));
}
#[tokio::test]
async fn document_intelligence_entra_token_is_reused_for_polling_only() {
let (post, poll) = document_intelligence_round_trip(None, Some("entra-token")).await;
let post = post.to_ascii_lowercase();
let poll = poll.to_ascii_lowercase();
assert!(post.contains("authorization: bearer entra-token"));
assert!(poll.contains("authorization: bearer entra-token"));
assert!(!post.contains("ocp-apim-subscription-key:"));
assert!(!poll.contains("ocp-apim-subscription-key:"));
assert!(post.contains("x-request-only: must-not-poll"));
assert!(!poll.contains("x-request-only:"));
}
#[tokio::test]
async fn document_intelligence_rejects_cross_origin_polling_before_sending_credentials() {
let provider = TcpListener::bind("127.0.0.1:0")
.await
.expect("provider listener binds");
let provider_address = provider.local_addr().expect("provider address");
let attacker = TcpListener::bind("127.0.0.1:0")
.await
.expect("attacker listener binds");
let attacker_address = attacker.local_addr().expect("attacker address");
let operation_url = format!("http://{attacker_address}/steal");
let server = tokio::spawn(async move {
let (mut socket, _) = provider.accept().await.expect("accepts POST");
let request = read_http_request(&mut socket).await;
write_response_with_headers(
&mut socket,
"202 Accepted",
&[("operation-location", &operation_url)],
"",
)
.await;
request
});
let api_base = format!("http://{provider_address}");
let client = test_client();
let error = ocr(
&client,
base_request("azure_ai/doc-intelligence/prebuilt-layout", &api_base),
)
.await
.expect_err("cross-origin operation URL must fail");
assert!(error.to_string().contains("cross-origin"));
let initial = server.await.expect("server task completes");
assert!(
initial
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: sk-test")
);
assert!(
tokio::time::timeout(std::time::Duration::from_millis(50), attacker.accept())
.await
.is_err()
);
}
#[tokio::test]
async fn provider_http_error_preserves_status_and_bounds_body() {
let listener = TcpListener::bind("127.0.0.1:0")

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

@ -1,29 +1,18 @@
use async_trait::async_trait;
use crate::Error;
use crate::auth::{AuthHeaderKind, Environment, ResolvedAuth};
use serde_json::{Map, Value};
use super::types::{OcrRequestData, OcrResponseData};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrAuthStrategy {
Bearer,
Header(&'static str),
}
impl OcrAuthStrategy {
pub fn header_name(self) -> &'static str {
match self {
Self::Bearer => "authorization",
Self::Header(header_name) => header_name,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrResponseHandling {
Json,
AzureDocumentIntelligencePoll,
}
#[async_trait]
pub trait OcrProviderConfig: Sync {
fn supported_ocr_params(&self) -> &'static [&'static str];
@ -71,30 +60,50 @@ pub trait OcrProviderConfig: Sync {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error>;
fn forwarded_auth(&self, headers: &[(String, String)]) -> Option<ResolvedAuth> {
let kind = self.auth_header_kind();
headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case(kind.header_name()))
.and_then(|(_, value)| ResolvedAuth::from_header(kind, value))
}
async fn resolve_auth(
&self,
api_key: Option<&str>,
_optional_params: &Map<String, Value>,
environment: &dyn Environment,
) -> Result<ResolvedAuth, Error> {
let env_lookup = |name: &str| environment.get(name);
Ok(ResolvedAuth::from_credential(
self.auth_header_kind(),
self.resolve_api_key(api_key, &env_lookup)?,
))
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn validate_environment(
async fn validate_environment(
&self,
headers: Vec<(String, String)>,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Vec<(String, String)>, Error> {
let strategy = self.auth_strategy();
if crate::http_utils::has_header(&headers, strategy.header_name()) {
return Ok(headers);
}
let api_key = self.resolve_api_key(api_key, env_lookup)?;
let auth_header = match strategy {
OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
OcrAuthStrategy::Header(name) => (name.to_string(), api_key),
optional_params: &Map<String, Value>,
environment: &dyn Environment,
) -> Result<(Vec<(String, String)>, ResolvedAuth), Error> {
let auth = match self.forwarded_auth(&headers) {
Some(auth) => auth,
None => {
self.resolve_auth(api_key, optional_params, environment)
.await?
}
};
Ok(std::iter::once(auth_header).chain(headers).collect())
Ok((auth.apply_preserving_existing(headers), auth))
}
fn auth_strategy(&self) -> OcrAuthStrategy {
OcrAuthStrategy::Bearer
fn auth_header_kind(&self) -> AuthHeaderKind {
AuthHeaderKind::Bearer
}
fn requires_data_uri_document(&self) -> bool {

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

@ -1,5 +1,6 @@
use crate::auth::AuthHeaderKind;
use crate::error::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use crate::messages::transformation::AnthropicMessagesProviderConfig;
const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE";
@ -65,8 +66,8 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
resolve_anthropic_api_key(api_key, env_lookup)
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
MessagesAuthStrategy::Header("x-api-key")
fn auth_strategy(&self) -> AuthHeaderKind {
AuthHeaderKind::Header("x-api-key")
}
}

540
litellm-rust/crates/core/src/providers/azure_ai/auth.rs Normal file
View file

@ -0,0 +1,540 @@
use std::path::PathBuf;
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant, SystemTime};
use async_trait::async_trait;
use azure_core::credentials::{Secret, TokenCredential};
use azure_core::http::ClientMethodOptions;
use azure_identity::{
ClientAssertion, ClientAssertionCredential, ClientSecretCredential, DeveloperToolsCredential,
ManagedIdentityCredential, ManagedIdentityCredentialOptions, UserAssignedId,
WorkloadIdentityCredential, WorkloadIdentityCredentialOptions,
};
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::Error;
use crate::auth::{
AuthHeaderKind, Environment, ExpiringToken, ResolvedAuth, SecretString, TokenCache,
TokenProvider,
};
use crate::constants::{AZURE_AUTH_DEFAULT_SCOPE, AZURE_AUTH_TOKEN_REFRESH_WINDOW_SECS};
static TOKENS: OnceLock<TokenCache> = OnceLock::new();
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AzureCredentialKind {
ApiKey,
StaticToken,
ClientSecret,
ClientAssertion,
WorkloadIdentity,
ManagedIdentity,
Default,
}
pub struct AzureAuthInput<'a> {
pub api_key: Option<&'a str>,
pub api_key_env: &'static str,
pub api_key_kind: AuthHeaderKind,
pub params: &'a Map<String, Value>,
pub environment: &'a dyn Environment,
}
enum AzureTokenCredential {
ClientSecret(Arc<ClientSecretCredential>),
ClientAssertion(Arc<ClientAssertionCredential<StaticClientAssertion>>),
WorkloadIdentity(Arc<WorkloadIdentityCredential>),
ManagedIdentity(Arc<ManagedIdentityCredential>),
Default {
workload_identity: Option<Arc<WorkloadIdentityCredential>>,
managed_identity: Arc<ManagedIdentityCredential>,
developer_tools: Option<Arc<DeveloperToolsCredential>>,
},
}
impl AzureTokenCredential {
async fn token(&self, scope: &str) -> azure_core::Result<azure_core::credentials::AccessToken> {
match self {
Self::ClientSecret(value) => value.get_token(&[scope], None).await,
Self::ClientAssertion(value) => value.get_token(&[scope], None).await,
Self::WorkloadIdentity(value) => value.get_token(&[scope], None).await,
Self::ManagedIdentity(value) => value.get_token(&[scope], None).await,
Self::Default {
workload_identity,
managed_identity,
developer_tools,
} => {
if let Some(value) = workload_identity
&& let Ok(token) = value.get_token(&[scope], None).await
{
return Ok(token);
}
match managed_identity.get_token(&[scope], None).await {
Ok(token) => Ok(token),
Err(managed_identity_error) => match developer_tools {
Some(value) => value.get_token(&[scope], None).await,
None => Err(managed_identity_error),
},
}
}
}
}
}
#[derive(Clone)]
struct StaticClientAssertion(SecretString);
impl std::fmt::Debug for StaticClientAssertion {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("[redacted]")
}
}
#[async_trait]
impl ClientAssertion for StaticClientAssertion {
async fn secret(
&self,
_options: Option<ClientMethodOptions<'_>>,
) -> azure_core::Result<String> {
Ok(self.0.expose().to_string())
}
}
struct AzureTokenProvider {
credential: AzureTokenCredential,
scope: String,
}
#[async_trait]
impl TokenProvider for AzureTokenProvider {
async fn token(&self) -> Result<ExpiringToken, Error> {
let token = self
.credential
.token(&self.scope)
.await
.map_err(|_| Error::Auth("Azure credential refresh failed".to_string()))?;
let seconds = token
.expires_on
.unix_timestamp()
.saturating_sub(
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64,
)
.max(0) as u64;
Ok(ExpiringToken {
token: SecretString::new(token.token.secret()),
expires_at: Instant::now() + Duration::from_secs(seconds),
})
}
}
fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
fn param_or_env(
params: &Map<String, Value>,
param: &str,
env: &str,
environment: &dyn Environment,
) -> Option<String> {
params
.get(param)
.and_then(Value::as_str)
.and_then(|value| non_empty(Some(value)))
.map(str::to_string)
.or_else(|| {
environment
.get(env)
.filter(|value| !value.trim().is_empty())
})
}
fn refresh_enabled(params: &Map<String, Value>) -> bool {
params
.get("azure_ad_token_refresh")
.and_then(Value::as_bool)
.unwrap_or(false)
}
pub fn classify_azure_credential(input: &AzureAuthInput<'_>) -> Result<AzureCredentialKind, Error> {
if non_empty(input.api_key).is_some()
|| input
.environment
.get(input.api_key_env)
.is_some_and(|value| !value.trim().is_empty())
{
return Ok(AzureCredentialKind::ApiKey);
}
let azure_ad_token = param_or_env(
input.params,
"azure_ad_token",
"AZURE_AD_TOKEN",
input.environment,
);
if azure_ad_token
.as_deref()
.is_some_and(|token| !token.starts_with("oidc/"))
{
return Ok(AzureCredentialKind::StaticToken);
}
let tenant = param_or_env(
input.params,
"tenant_id",
"AZURE_TENANT_ID",
input.environment,
);
let client = param_or_env(
input.params,
"client_id",
"AZURE_CLIENT_ID",
input.environment,
);
let secret = param_or_env(
input.params,
"client_secret",
"AZURE_CLIENT_SECRET",
input.environment,
);
if tenant.is_some() && client.is_some() && secret.is_some() {
return Ok(AzureCredentialKind::ClientSecret);
}
if azure_ad_token
.as_deref()
.is_some_and(|token| token.starts_with("oidc/"))
{
if tenant.is_some()
&& client.is_some()
&& param_or_env(
input.params,
"azure_ad_client_assertion",
"AZURE_AD_CLIENT_ASSERTION",
input.environment,
)
.is_some()
{
return Ok(AzureCredentialKind::ClientAssertion);
}
return Err(Error::Auth(
"Azure OIDC authentication requires tenant_id, client_id, and a resolved client assertion"
.to_string(),
));
}
if !refresh_enabled(input.params) {
return Err(Error::Auth(
"Missing Azure credentials; set an API key, azure_ad_token, service principal, or enable Azure token refresh"
.to_string(),
));
}
let credential = input.environment.get("AZURE_CREDENTIAL");
if credential.as_deref() == Some("WorkloadIdentityCredential") {
return Ok(AzureCredentialKind::WorkloadIdentity);
}
if credential.as_deref() == Some("ManagedIdentityCredential")
|| input.environment.get("IDENTITY_ENDPOINT").is_some()
|| input.environment.get("MSI_ENDPOINT").is_some()
{
return Ok(AzureCredentialKind::ManagedIdentity);
}
Ok(AzureCredentialKind::Default)
}
fn cache_key(kind: AzureCredentialKind, values: &[&str]) -> String {
let mut hasher = Sha256::new();
hasher.update(format!("{kind:?}"));
values.iter().for_each(|value| {
hasher.update(value.len().to_le_bytes());
hasher.update(value.as_bytes());
});
hasher
.finalize()
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
fn token_credential(
kind: AzureCredentialKind,
params: &Map<String, Value>,
environment: &dyn Environment,
) -> Result<(String, AzureTokenCredential), Error> {
let scope = param_or_env(params, "azure_scope", "AZURE_SCOPE", environment)
.unwrap_or_else(|| AZURE_AUTH_DEFAULT_SCOPE.to_string());
let tenant = param_or_env(params, "tenant_id", "AZURE_TENANT_ID", environment);
let client = param_or_env(params, "client_id", "AZURE_CLIENT_ID", environment);
let token_file = environment
.get("AZURE_FEDERATED_TOKEN_FILE")
.filter(|value| !value.trim().is_empty());
let result = match kind {
AzureCredentialKind::ClientSecret => {
let tenant =
tenant.ok_or_else(|| Error::Auth("Azure tenant ID is missing".to_string()))?;
let client =
client.ok_or_else(|| Error::Auth("Azure client ID is missing".to_string()))?;
let secret = param_or_env(params, "client_secret", "AZURE_CLIENT_SECRET", environment)
.ok_or_else(|| Error::Auth("Azure client secret is missing".to_string()))?;
let key = cache_key(kind, &[&tenant, &client, &secret, &scope]);
let credential =
ClientSecretCredential::new(&tenant, client, Secret::new(secret), None)
.map_err(|_| Error::Auth("Azure service principal is invalid".to_string()))?;
(key, AzureTokenCredential::ClientSecret(credential))
}
AzureCredentialKind::ClientAssertion => {
let tenant =
tenant.ok_or_else(|| Error::Auth("Azure tenant ID is missing".to_string()))?;
let client =
client.ok_or_else(|| Error::Auth("Azure client ID is missing".to_string()))?;
let assertion = param_or_env(
params,
"azure_ad_client_assertion",
"AZURE_AD_CLIENT_ASSERTION",
environment,
)
.ok_or_else(|| Error::Auth("Azure client assertion is missing".to_string()))?;
let key = cache_key(kind, &[&tenant, &client, &assertion, &scope]);
let credential = ClientAssertionCredential::new(
tenant,
client,
StaticClientAssertion(SecretString::new(assertion)),
None,
)
.map_err(|_| Error::Auth("Azure client assertion is invalid".to_string()))?;
(key, AzureTokenCredential::ClientAssertion(credential))
}
AzureCredentialKind::WorkloadIdentity => {
let credential =
WorkloadIdentityCredential::new(Some(WorkloadIdentityCredentialOptions {
tenant_id: tenant.clone(),
client_id: client.clone(),
token_file_path: token_file.clone().map(PathBuf::from),
..Default::default()
}))
.map_err(|_| Error::Auth("Azure workload identity is invalid".to_string()))?;
(
cache_key(
kind,
&[
tenant.as_deref().unwrap_or(""),
client.as_deref().unwrap_or(""),
token_file.as_deref().unwrap_or(""),
&scope,
],
),
AzureTokenCredential::WorkloadIdentity(credential),
)
}
AzureCredentialKind::ManagedIdentity => {
let options = ManagedIdentityCredentialOptions {
user_assigned_id: client.clone().map(UserAssignedId::ClientId),
..Default::default()
};
let credential = ManagedIdentityCredential::new(Some(options))
.map_err(|_| Error::Auth("Azure managed identity is unavailable".to_string()))?;
(
cache_key(kind, &[client.as_deref().unwrap_or("system"), &scope]),
AzureTokenCredential::ManagedIdentity(credential),
)
}
AzureCredentialKind::Default => {
let workload_identity = match (&tenant, &client, &token_file) {
(Some(tenant_id), Some(client_id), Some(token_file_path)) => {
WorkloadIdentityCredential::new(Some(WorkloadIdentityCredentialOptions {
tenant_id: Some(tenant_id.clone()),
client_id: Some(client_id.clone()),
token_file_path: Some(PathBuf::from(token_file_path)),
..Default::default()
}))
.ok()
}
_ => None,
};
let options = ManagedIdentityCredentialOptions {
user_assigned_id: client.clone().map(UserAssignedId::ClientId),
..Default::default()
};
let managed_identity = ManagedIdentityCredential::new(Some(options)).map_err(|_| {
Error::Auth("Azure default credentials are unavailable".to_string())
})?;
let developer_tools = DeveloperToolsCredential::new(None).ok();
(
cache_key(
kind,
&[
tenant.as_deref().unwrap_or(""),
client.as_deref().unwrap_or("system"),
token_file.as_deref().unwrap_or(""),
&scope,
],
),
AzureTokenCredential::Default {
workload_identity,
managed_identity,
developer_tools,
},
)
}
AzureCredentialKind::ApiKey | AzureCredentialKind::StaticToken => {
return Err(Error::Auth(
"Azure credential classification failed".to_string(),
));
}
};
Ok((format!("{}:{scope}", result.0), result.1))
}
pub async fn resolve_azure_auth(input: AzureAuthInput<'_>) -> Result<ResolvedAuth, Error> {
let kind = classify_azure_credential(&input)?;
match kind {
AzureCredentialKind::ApiKey => {
let key = non_empty(input.api_key)
.map(str::to_string)
.or_else(|| input.environment.get(input.api_key_env))
.ok_or_else(|| Error::Auth("Azure API key is missing".to_string()))?;
Ok(ResolvedAuth::from_credential(input.api_key_kind, key))
}
AzureCredentialKind::StaticToken => {
let token = param_or_env(
input.params,
"azure_ad_token",
"AZURE_AD_TOKEN",
input.environment,
)
.ok_or_else(|| Error::Auth("Azure access token is missing".to_string()))?;
Ok(ResolvedAuth::Bearer(SecretString::new(token)))
}
kind => {
let scope = param_or_env(
input.params,
"azure_scope",
"AZURE_SCOPE",
input.environment,
)
.unwrap_or_else(|| AZURE_AUTH_DEFAULT_SCOPE.to_string());
let (key, credential) = token_credential(kind, input.params, input.environment)?;
let provider = AzureTokenProvider { credential, scope };
let token = TOKENS
.get_or_init(|| {
TokenCache::new(Duration::from_secs(AZURE_AUTH_TOKEN_REFRESH_WINDOW_SECS))
})
.token(key, &provider)
.await?;
Ok(ResolvedAuth::Bearer(token))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn empty_environment(_: &str) -> Option<String> {
None
}
fn input<'a>(
api_key: Option<&'a str>,
params: &'a Map<String, Value>,
env: &'a dyn Environment,
) -> AzureAuthInput<'a> {
AzureAuthInput {
api_key,
api_key_env: "AZURE_DOCUMENT_INTELLIGENCE_API_KEY",
api_key_kind: AuthHeaderKind::Header("Ocp-Apim-Subscription-Key"),
params,
environment: env,
}
}
#[test]
fn api_key_precedes_static_token_and_service_principal() {
let params = Map::from_iter([
(
"azure_ad_token".to_string(),
Value::String("token".to_string()),
),
("tenant_id".to_string(), Value::String("tenant".to_string())),
("client_id".to_string(), Value::String("client".to_string())),
(
"client_secret".to_string(),
Value::String("secret".to_string()),
),
]);
assert_eq!(
classify_azure_credential(&input(Some("key"), &params, &empty_environment)).unwrap(),
AzureCredentialKind::ApiKey
);
}
#[test]
fn static_token_precedes_service_principal_and_refresh_requires_opt_in() {
let token_params = Map::from_iter([
(
"azure_ad_token".to_string(),
Value::String("token".to_string()),
),
("tenant_id".to_string(), Value::String("tenant".to_string())),
("client_id".to_string(), Value::String("client".to_string())),
(
"client_secret".to_string(),
Value::String("secret".to_string()),
),
]);
assert_eq!(
classify_azure_credential(&input(None, &token_params, &empty_environment)).unwrap(),
AzureCredentialKind::StaticToken
);
let params = Map::new();
assert!(classify_azure_credential(&input(None, &params, &empty_environment)).is_err());
}
#[test]
fn oidc_reference_requires_a_resolved_assertion() {
let params = Map::from_iter([
(
"azure_ad_token".to_string(),
Value::String("oidc/token-reference".to_string()),
),
("tenant_id".to_string(), Value::String("tenant".to_string())),
("client_id".to_string(), Value::String("client".to_string())),
]);
assert!(classify_azure_credential(&input(None, &params, &empty_environment)).is_err());
let mut resolved = params;
resolved.insert(
"azure_ad_client_assertion".to_string(),
Value::String("assertion".to_string()),
);
assert_eq!(
classify_azure_credential(&input(None, &resolved, &empty_environment)).unwrap(),
AzureCredentialKind::ClientAssertion
);
}
#[test]
fn refresh_uses_managed_identity_signals_or_default_chain() {
let params = Map::from_iter([("azure_ad_token_refresh".to_string(), Value::Bool(true))]);
let managed_environment = |name: &str| {
(name == "IDENTITY_ENDPOINT").then(|| "https://identity.example".to_string())
};
assert_eq!(
classify_azure_credential(&input(None, &params, &managed_environment)).unwrap(),
AzureCredentialKind::ManagedIdentity
);
assert_eq!(
classify_azure_credential(&input(None, &params, &empty_environment)).unwrap(),
AzureCredentialKind::Default
);
}
#[test]
fn cache_keys_preserve_credential_field_boundaries() {
assert_ne!(
cache_key(AzureCredentialKind::ClientSecret, &["ab", "c"]),
cache_key(AzureCredentialKind::ClientSecret, &["a", "bc"])
);
}
}

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

@ -1,5 +1,6 @@
use crate::auth::AuthHeaderKind;
use crate::error::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use crate::messages::transformation::AnthropicMessagesProviderConfig;
use crate::messages::types::{
AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock,
MessageContent, SystemPrompt,
@ -160,7 +161,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
resolve_azure_api_key(api_key, env_lookup)
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
fn auth_strategy(&self) -> AuthHeaderKind {
self.anthropic.auth_strategy()
}

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

@ -1,2 +1,3 @@
pub mod auth;
pub mod messages;
pub mod ocr;

167
litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs
View file

@ -1,8 +1,12 @@
use std::collections::BTreeSet;
use async_trait::async_trait;
use crate::auth::{AuthHeaderKind, Environment, ResolvedAuth};
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
use crate::ocr::transformation::{OcrProviderConfig, OcrResponseHandling};
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use crate::providers::azure_ai::auth::{AzureAuthInput, resolve_azure_auth};
use serde_json::{Map, Value, json};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
@ -99,68 +103,6 @@ pub fn resolve_document_intelligence_endpoint(
)
}
fn prepend_auth_header(
headers: Vec<(String, String)>,
name: &str,
value: String,
) -> Vec<(String, String)> {
std::iter::once((name.to_string(), value))
.chain(headers)
.collect()
}
pub fn validate_azure_ai_environment(
headers: Vec<(String, String)>,
api_key: Option<&str>,
azure_ad_token: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Vec<(String, String)>, Error> {
if crate::http_utils::has_header(&headers, "Authorization")
|| crate::http_utils::has_header(&headers, "Api-Key")
{
return Ok(headers);
}
if let Ok(api_key) = resolve_azure_ai_api_key(api_key, env_lookup) {
return Ok(prepend_auth_header(headers, "Api-Key", api_key));
}
non_empty(azure_ad_token)
.map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}")))
.ok_or_else(|| {
Error::Auth(
"Missing Azure AI credentials - set AZURE_AI_API_KEY or provide azure_ad_token"
.to_string(),
)
})
}
pub fn validate_document_intelligence_environment(
headers: Vec<(String, String)>,
api_key: Option<&str>,
azure_ad_token: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Vec<(String, String)>, Error> {
if crate::http_utils::has_header(&headers, "Authorization")
|| crate::http_utils::has_header(&headers, "Ocp-Apim-Subscription-Key")
{
return Ok(headers);
}
if let Ok(api_key) = resolve_document_intelligence_api_key(api_key, env_lookup) {
return Ok(prepend_auth_header(
headers,
"Ocp-Apim-Subscription-Key",
api_key,
));
}
non_empty(azure_ad_token)
.map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}")))
.ok_or_else(|| {
Error::Auth(
"Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or provide azure_ad_token"
.to_string(),
)
})
}
fn encode_model_id(model: &str) -> Result<String, Error> {
let model_id = model.rsplit('/').next().unwrap_or(model);
if matches!(model_id, "." | "..") {
@ -502,6 +444,7 @@ fn transform_document_intelligence_response(
})
}
#[async_trait]
impl OcrProviderConfig for AzureAiOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
@ -539,16 +482,33 @@ impl OcrProviderConfig for AzureAiOcrConfig {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
resolve_azure_ai_api_key(api_key, env_lookup)
}
async fn resolve_auth(
&self,
api_key: Option<&str>,
optional_params: &Map<String, Value>,
environment: &dyn Environment,
) -> Result<ResolvedAuth, Error> {
resolve_azure_auth(AzureAuthInput {
api_key,
api_key_env: AZURE_AI_API_KEY_ENV,
api_key_kind: AuthHeaderKind::Bearer,
params: optional_params,
environment,
})
.await
}
fn requires_data_uri_document(&self) -> bool {
true
}
}
#[async_trait]
impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn supported_ocr_params(&self) -> &'static [&'static str] {
@ -636,13 +596,43 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
resolve_document_intelligence_api_key(api_key, env_lookup)
}
fn auth_strategy(&self) -> OcrAuthStrategy {
OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key")
fn forwarded_auth(&self, headers: &[(String, String)]) -> Option<ResolvedAuth> {
[
AuthHeaderKind::Header("Ocp-Apim-Subscription-Key"),
AuthHeaderKind::Bearer,
]
.iter()
.find_map(|kind| {
headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case(kind.header_name()))
.and_then(|(_, value)| ResolvedAuth::from_header(*kind, value))
})
}
async fn resolve_auth(
&self,
api_key: Option<&str>,
optional_params: &Map<String, Value>,
environment: &dyn Environment,
) -> Result<ResolvedAuth, Error> {
resolve_azure_auth(AzureAuthInput {
api_key,
api_key_env: AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV,
api_key_kind: AuthHeaderKind::Header("Ocp-Apim-Subscription-Key"),
params: optional_params,
environment,
})
.await
}
fn auth_header_kind(&self) -> AuthHeaderKind {
AuthHeaderKind::Header("Ocp-Apim-Subscription-Key")
}
fn response_handling(&self) -> OcrResponseHandling {
@ -669,6 +659,10 @@ mod tests {
.map(|(_, value)| value.as_str())
}
fn empty_environment(_: &str) -> Option<String> {
None
}
#[fixture]
fn native_operation() -> Value {
json!({
@ -1103,11 +1097,12 @@ mod tests {
assert!(!url.contains("req_format"));
}
#[test]
fn document_intelligence_validate_environment_uses_subscription_key() {
let headers =
validate_document_intelligence_environment(Vec::new(), Some("my-key"), None, &|_| None)
.expect("api key authenticates");
#[tokio::test]
async fn document_intelligence_validate_environment_uses_subscription_key() {
let (headers, _) = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
.validate_environment(Vec::new(), Some("my-key"), &Map::new(), &empty_environment)
.await
.expect("api key authenticates");
assert_eq!(
header_value(&headers, "Ocp-Apim-Subscription-Key"),
@ -1115,15 +1110,13 @@ mod tests {
);
}
#[test]
fn document_intelligence_validate_environment_falls_back_to_entra_token() {
let headers = validate_document_intelligence_environment(
Vec::new(),
None,
Some("entra-token"),
&|_| None,
)
.expect("Entra token authenticates");
#[tokio::test]
async fn document_intelligence_validate_environment_falls_back_to_entra_token() {
let params = Map::from_iter([("azure_ad_token".to_string(), json!("entra-token"))]);
let (headers, _) = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
.validate_environment(Vec::new(), None, &params, &empty_environment)
.await
.expect("Entra token authenticates");
assert_eq!(
header_value(&headers, "Authorization"),
@ -1373,11 +1366,13 @@ mod tests {
assert_eq!(resolved, "https://generic-azure-ai.example.com");
}
#[test]
fn azure_ai_ocr_authenticates_with_entra_token() {
let headers =
validate_azure_ai_environment(Vec::new(), None, Some("entra-token"), &|_| None)
.expect("Entra token authenticates");
#[tokio::test]
async fn azure_ai_ocr_authenticates_with_entra_token() {
let params = Map::from_iter([("azure_ad_token".to_string(), json!("entra-token"))]);
let (headers, _) = AZURE_AI_OCR_CONFIG
.validate_environment(Vec::new(), None, &params, &empty_environment)
.await
.expect("Entra token authenticates");
assert_eq!(
header_value(&headers, "Authorization"),

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

@ -153,7 +153,7 @@ impl OcrProviderConfig for MistralOcrConfig {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}

21
litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs
View file

@ -1,9 +1,13 @@
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use serde_json::{Map, Value, json};
use super::transformation::*;
use crate::ocr::transformation::OcrProviderConfig;
fn reducto_environment(_: &str) -> Option<String> {
Some("env-reducto-key".to_string())
}
#[fixture]
fn parse_response() -> Value {
json!({
@ -185,15 +189,20 @@ fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) {
}
#[rstest]
fn test_parse_v3_uses_programmatic_api_key_over_env() {
#[tokio::test]
async fn test_parse_v3_uses_programmatic_api_key_over_env() {
let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string()))
.expect("explicit key should resolve");
assert_eq!(key, "passed-key");
let headers = REDUCTO_PARSE_V3_CONFIG
.validate_environment(Vec::new(), Some("passed-key"), &|_| {
Some("env-reducto-key".to_string())
})
let (headers, _) = REDUCTO_PARSE_V3_CONFIG
.validate_environment(
Vec::new(),
Some("passed-key"),
&Map::new(),
&reducto_environment,
)
.await
.expect("headers should validate");
assert_eq!(
headers,

4
litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs
View file

@ -358,7 +358,7 @@ impl OcrProviderConfig for ReductoParseV3Config {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
@ -400,7 +400,7 @@ impl OcrProviderConfig for ReductoParseLegacyConfig {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}

290
litellm-rust/crates/core/src/providers/vertex_ai/auth.rs Normal file
View file

@ -0,0 +1,290 @@
use std::collections::HashMap;
use std::sync::OnceLock;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use google_cloud_auth::credentials::impersonated;
use google_cloud_auth::credentials::{
AccessTokenCredentials, external_account, service_account, user_account,
};
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use tokio::sync::Mutex;
use crate::Error;
use crate::auth::{
Environment, ExpiringToken, ResolvedAuth, SecretString, TokenCache, TokenProvider,
};
use crate::constants::{
VERTEX_AUTH_CLOUD_PLATFORM_SCOPE, VERTEX_AUTH_SDK_TOKEN_CACHE_TTL_SECS,
VERTEX_AUTH_TOKEN_REFRESH_WINDOW_SECS,
};
static CREDENTIALS: OnceLock<Mutex<HashMap<String, AccessTokenCredentials>>> = OnceLock::new();
static TOKENS: OnceLock<TokenCache> = OnceLock::new();
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VertexAuth {
pub auth: ResolvedAuth,
pub project_id: String,
}
struct GoogleTokenProvider {
credentials: AccessTokenCredentials,
}
#[async_trait]
impl TokenProvider for GoogleTokenProvider {
async fn token(&self) -> Result<ExpiringToken, Error> {
let token = self
.credentials
.access_token()
.await
.map_err(|_| Error::Auth("Vertex AI credential refresh failed".to_string()))?;
Ok(ExpiringToken {
token: SecretString::new(token.token),
expires_at: Instant::now() + Duration::from_secs(VERTEX_AUTH_SDK_TOKEN_CACHE_TTL_SECS),
})
}
}
fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
fn value_string<'a>(params: &'a Map<String, Value>, names: &[&str]) -> Option<&'a str> {
names
.iter()
.find_map(|name| params.get(*name).and_then(Value::as_str))
.and_then(|value| non_empty(Some(value)))
}
fn fingerprint(value: &str) -> String {
Sha256::digest(value.as_bytes())
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
fn credentials_json(value: &Value) -> Result<Option<Value>, Error> {
match value {
Value::Null => Ok(None),
Value::Object(_) => Ok(Some(value.clone())),
Value::String(raw) if raw.trim_start().starts_with('{') => serde_json::from_str(raw)
.map(Some)
.map_err(|_| Error::Auth("Vertex AI credentials JSON is invalid".to_string())),
Value::String(path) => std::fs::read_to_string(path)
.map_err(|_| Error::Auth("Vertex AI credentials file could not be read".to_string()))
.and_then(|raw| {
serde_json::from_str(&raw).map(Some).map_err(|_| {
Error::Auth("Vertex AI credentials file contains invalid JSON".to_string())
})
}),
_ => Err(Error::Auth(
"Vertex AI credentials must be a JSON object, JSON string, or file path".to_string(),
)),
}
}
fn build_credentials(json: Option<Value>) -> Result<AccessTokenCredentials, Error> {
let built = match json {
None => google_cloud_auth::credentials::Builder::default()
.with_scopes([VERTEX_AUTH_CLOUD_PLATFORM_SCOPE])
.build_access_token_credentials(),
Some(value) => match value.get("type").and_then(Value::as_str) {
Some("service_account") => service_account::Builder::new(value)
.with_access_specifier(service_account::AccessSpecifier::from_scopes([
VERTEX_AUTH_CLOUD_PLATFORM_SCOPE,
]))
.build_access_token_credentials(),
Some("authorized_user") => user_account::Builder::new(value)
.with_scopes([VERTEX_AUTH_CLOUD_PLATFORM_SCOPE])
.build_access_token_credentials(),
Some("external_account") => external_account::Builder::new(value)
.with_scopes([VERTEX_AUTH_CLOUD_PLATFORM_SCOPE])
.build_access_token_credentials(),
Some("impersonated_service_account") => impersonated::Builder::new(value)
.with_scopes([VERTEX_AUTH_CLOUD_PLATFORM_SCOPE])
.build_access_token_credentials(),
_ => {
return Err(Error::Auth(
"Vertex AI credential type is unsupported".to_string(),
));
}
},
};
built.map_err(|_| Error::Auth("Vertex AI credentials could not be loaded".to_string()))
}
fn configured_project(
explicit: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
non_empty(explicit).map(str::to_string).or_else(|| {
["VERTEXAI_PROJECT", "GOOGLE_CLOUD_PROJECT", "GCLOUD_PROJECT"]
.iter()
.find_map(|name| env_lookup(name))
.filter(|value| !value.trim().is_empty())
})
}
fn inferred_project(
explicit: Option<&str>,
json: Option<&Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
configured_project(explicit, env_lookup)
.or_else(|| {
json.and_then(|value| {
value
.get("project_id")
.or_else(|| value.get("quota_project_id"))
.and_then(Value::as_str)
.and_then(|value| non_empty(Some(value)))
.map(str::to_string)
})
})
.ok_or_else(|| {
Error::InvalidRequest(
"Could not resolve Vertex AI project; pass vertex_project or set VERTEXAI_PROJECT"
.to_string(),
)
})
}
pub fn resolve_vertex_project(
params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
let explicit = value_string(params, &["vertex_project", "vertex_ai_project"]);
if let Some(project) = configured_project(explicit, env_lookup) {
return Ok(project);
}
let json = params
.get("vertex_credentials")
.map(credentials_json)
.transpose()?
.flatten();
inferred_project(None, json.as_ref(), env_lookup)
}
async fn cached_credentials(
key: String,
json: Option<Value>,
) -> Result<AccessTokenCredentials, Error> {
let credentials = CREDENTIALS.get_or_init(|| Mutex::new(HashMap::new()));
let mut entries = credentials.lock().await;
if let Some(credentials) = entries.get(&key) {
return Ok(credentials.clone());
}
let loaded = build_credentials(json)?;
entries.insert(key, loaded.clone());
Ok(loaded)
}
pub async fn resolve_vertex_auth(
api_key: Option<&str>,
params: &Map<String, Value>,
environment: &dyn Environment,
) -> Result<VertexAuth, Error> {
let explicit_project = value_string(params, &["vertex_project", "vertex_ai_project"]);
let env_lookup = |name: &str| environment.get(name);
let configured_project = configured_project(explicit_project, &env_lookup);
let explicit_token = non_empty(api_key)
.map(str::to_string)
.or_else(|| {
environment
.get("VERTEX_AI_API_KEY")
.filter(|value| !value.trim().is_empty())
})
.or_else(|| {
environment
.get("VERTEXAI_API_KEY")
.filter(|value| !value.trim().is_empty())
});
let credential_value = params.get("vertex_credentials");
let json = match (configured_project.as_ref(), explicit_token.as_ref()) {
(Some(_), Some(_)) => None,
_ => credential_value
.map(credentials_json)
.transpose()?
.flatten(),
};
let project_id = configured_project
.map(Ok)
.unwrap_or_else(|| inferred_project(None, json.as_ref(), &env_lookup))?;
let token = match explicit_token {
Some(token) => SecretString::new(token),
None => {
let identity = credential_value
.map(Value::to_string)
.unwrap_or_else(|| "application-default".to_string());
let key = fingerprint(&identity);
let provider = GoogleTokenProvider {
credentials: cached_credentials(key.clone(), json).await?,
};
TOKENS
.get_or_init(|| {
TokenCache::new(Duration::from_secs(VERTEX_AUTH_TOKEN_REFRESH_WINDOW_SECS))
})
.token(key, &provider)
.await?
}
};
Ok(VertexAuth {
auth: ResolvedAuth::Bearer(token),
project_id,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn explicit_token_and_project_precede_environment_without_loading_adc() {
let params = Map::from_iter([
(
"vertex_project".to_string(),
Value::String("param-project".to_string()),
),
(
"vertex_credentials".to_string(),
Value::String("missing-credentials.json".to_string()),
),
]);
let env = |name: &str| match name {
"VERTEXAI_PROJECT" => Some("env-project".to_string()),
"VERTEX_AI_API_KEY" => Some("env-token".to_string()),
_ => None,
};
let resolved = resolve_vertex_auth(Some("explicit-token"), &params, &env)
.await
.expect("explicit credentials resolve");
assert_eq!(resolved.project_id, "param-project");
assert_eq!(
resolved.auth.credential_header(),
(
"Authorization".to_string(),
"Bearer explicit-token".to_string()
)
);
}
#[test]
fn credential_json_infers_service_account_and_authorized_user_projects() {
let env = |_: &str| None;
let service =
serde_json::json!({"type": "service_account", "project_id": "service-project"});
let user =
serde_json::json!({"type": "authorized_user", "quota_project_id": "quota-project"});
assert_eq!(
inferred_project(None, Some(&service), &env).unwrap(),
"service-project"
);
assert_eq!(
inferred_project(None, Some(&user), &env).unwrap(),
"quota-project"
);
}
}

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

@ -1 +1,2 @@
pub mod auth;
pub mod ocr;

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

@ -1,15 +1,18 @@
use async_trait::async_trait;
use crate::auth::{Environment, ResolvedAuth};
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use serde_json::{Map, Value, json};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use crate::providers::vertex_ai::auth::{resolve_vertex_auth, resolve_vertex_project};
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";
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
@ -62,15 +65,7 @@ fn vertex_project(
params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
string_param(params, &["vertex_project", "vertex_ai_project"])
.map(str::to_string)
.or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| {
Error::InvalidRequest(
"Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter"
.to_string(),
)
})
resolve_vertex_project(params, env_lookup)
}
fn vertex_location(
@ -207,6 +202,7 @@ fn ocr_data_from_content(content: Value, usage: Option<Value>, model: &str) -> V
}
}
#[async_trait]
impl OcrProviderConfig for VertexAiOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
@ -244,16 +240,28 @@ impl OcrProviderConfig for VertexAiOcrConfig {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
resolve_vertex_api_key(api_key, env_lookup)
}
async fn resolve_auth(
&self,
api_key: Option<&str>,
optional_params: &Map<String, Value>,
environment: &dyn Environment,
) -> Result<ResolvedAuth, Error> {
Ok(resolve_vertex_auth(api_key, optional_params, environment)
.await?
.auth)
}
fn requires_data_uri_document(&self) -> bool {
true
}
}
#[async_trait]
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn supported_ocr_params(&self) -> &'static [&'static str] {
@ -368,10 +376,21 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
resolve_vertex_api_key(api_key, env_lookup)
}
async fn resolve_auth(
&self,
api_key: Option<&str>,
optional_params: &Map<String, Value>,
environment: &dyn Environment,
) -> Result<ResolvedAuth, Error> {
Ok(resolve_vertex_auth(api_key, optional_params, environment)
.await?
.auth)
}
}
#[cfg(test)]

6
litellm/llms/azure_ai/ocr/document_intelligence/transformation.py
View file

@ -618,7 +618,11 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
except SSRFError as ssrf_err:
raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})")
poll_headers = {"Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "")}
poll_headers: Final = { # mutable-ok: httpx requires a concrete headers dict
name: value
for name in ("Ocp-Apim-Subscription-Key", "Authorization")
if (value := raw_response.request.headers.get(name))
}
return operation_url, poll_headers
@staticmethod

66
litellm/ocr/main.py
View file

@ -196,7 +196,7 @@ def _rust_bridge_optional_params(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> dict[str, object]:
optional_params: Final = dict(prepared_request.optional_params)
optional_params: Final = prepared_request.optional_params
if prepared_request.custom_llm_provider == "vertex_ai":
vertex_project: Final = (
prepared_request.litellm_params.get("vertex_project")
@ -211,11 +211,45 @@ def _rust_bridge_optional_params(
or resolve_secret("VERTEXAI_LOCATION")
or resolve_secret("VERTEX_LOCATION")
)
if vertex_project is not None:
optional_params["vertex_project"] = vertex_project
if vertex_location is not None:
optional_params["vertex_location"] = vertex_location
return optional_params
vertex_credentials: Final = prepared_request.litellm_params.get("vertex_credentials")
vertex_auth_params: Final = { # mutable-ok: the Rust bridge requires a concrete dict
name: value
for name, value in (
("vertex_project", vertex_project),
("vertex_location", vertex_location),
("vertex_credentials", vertex_credentials),
)
if value is not None
}
return { # mutable-ok: the Rust bridge requires a concrete dict
**optional_params,
**vertex_auth_params,
}
if prepared_request.custom_llm_provider == "azure_ai":
azure_ad_token: Final = prepared_request.litellm_params.get("azure_ad_token")
azure_ad_client_assertion: Final = (
resolve_secret(azure_ad_token)
if isinstance(azure_ad_token, str) and azure_ad_token.startswith("oidc/")
else None
)
azure_auth_params: Final = { # mutable-ok: the Rust bridge requires a concrete dict
name: value
for name, value in (
("azure_ad_token", azure_ad_token),
("azure_ad_client_assertion", azure_ad_client_assertion),
("tenant_id", prepared_request.litellm_params.get("tenant_id")),
("client_id", prepared_request.litellm_params.get("client_id")),
("client_secret", prepared_request.litellm_params.get("client_secret")),
("azure_scope", prepared_request.litellm_params.get("azure_scope")),
)
if value is not None
}
return { # mutable-ok: the Rust bridge requires a concrete dict
**optional_params,
**azure_auth_params,
"azure_ad_token_refresh": litellm.enable_azure_ad_token_refresh is True,
}
return dict(optional_params) # mutable-ok: the Rust bridge requires a concrete dict
def _rust_bridge_api_base(
@ -237,15 +271,21 @@ def _prepare_rust_ocr_call(
) -> rust_ocr_bridge.RustOCRRequest:
provider_config: Final = prepared_request.provider_config
api_key_env_var: Final = provider_config.get_api_key_env_var()
rust_owns_auth: Final = prepared_request.custom_llm_provider in ("vertex_ai", "azure_ai")
resolved_api_key: Final = prepared_request.api_key or (
resolve_api_key(api_key_env_var) if api_key_env_var is not None else None
resolve_api_key(api_key_env_var) if api_key_env_var is not None and not rust_owns_auth else None
)
resolved_headers: Final = provider_config.validate_environment(
headers=prepared_request.extra_headers or {},
model=prepared_request.model,
api_key=resolved_api_key,
api_base=prepared_request.api_base,
litellm_params=prepared_request.litellm_params,
resolved_headers: Final = (
dict(prepared_request.extra_headers or {}) # mutable-ok: the Rust bridge requires a concrete dict
if rust_owns_auth
else provider_config.validate_environment(
headers=prepared_request.extra_headers
or {}, # mutable-ok: provider validation requires mutable request headers
model=prepared_request.model,
api_key=resolved_api_key,
api_base=prepared_request.api_base,
litellm_params=prepared_request.litellm_params,
)
)
resolved_complete_url: Final = provider_config.get_complete_url(
api_base=prepared_request.api_base,

27
tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py
View file

@ -1,10 +1,10 @@
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
from litellm.exceptions import UnsupportedParamsError
from litellm.llms.azure_ai.ocr.document_intelligence.transformation import (
AzureDocumentIntelligenceOCRConfig,
)
@ -96,6 +96,31 @@ def _completed_response(payload: dict) -> httpx.Response:
)
@pytest.mark.parametrize(
"credential_header",
[
{"Ocp-Apim-Subscription-Key": "subscription-key"},
{"Authorization": "Bearer entra-token"},
],
)
def test_polling_target_reuses_only_same_origin_credential(credential_header: dict[str, str]) -> None:
operation_url: Final = "https://example.cognitiveservices.azure.com/analyzeResults/xyz"
initial_response: Final = httpx.Response(
status_code=202,
headers={"Operation-Location": operation_url},
request=httpx.Request(
"POST",
"https://example.cognitiveservices.azure.com/documentintelligence/documentModels/layout:analyze",
headers={**credential_header, "x-trace-id": "request-only"},
),
)
target, poll_headers = AzureDocumentIntelligenceOCRConfig()._get_polling_target(initial_response)
assert target == operation_url
assert poll_headers == credential_header
def _assert_native_fields_preserved(serialized: dict) -> None:
assert serialized["content"] == "Invoice\nInvoice No: INV-12345\nTotal: $100.00"
assert serialized["tables"] == AZURE_TABLES

29
tests/test_litellm/llms/test_polling_url_origin_match.py
View file

@ -12,10 +12,8 @@ share an origin with the original request URL.
from unittest.mock import MagicMock, patch
import httpx
import pytest
# Azure DALL-E sync + async paths route through ``assert_same_origin``
# the same way as the cases below. The helper itself is unit-tested in
# ``tests/test_litellm/litellm_core_utils/test_url_utils.py``; the
@ -43,17 +41,22 @@ def test_azure_di_sync_rejects_cross_origin_polling():
)
raw_response.request.headers = {"Ocp-Apim-Subscription-Key": "leak-me"}
with pytest.raises(ValueError, match="rejected polling URL"):
config.transform_ocr_response(
model="azure-doc-intel",
raw_response=raw_response,
logging_obj=MagicMock(),
request_data={},
optional_params={},
litellm_params={},
encoding=None,
response={},
)
with patch( # test-quality-ok: verifies URL validation runs before the internal client factory
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as get_client:
with pytest.raises(ValueError, match="rejected polling URL"):
config.transform_ocr_response(
model="azure-doc-intel",
raw_response=raw_response,
logging_obj=MagicMock(),
request_data={},
optional_params={},
litellm_params={},
encoding=None,
response={},
)
get_client.assert_not_called()
# ── Black Forest Labs polling ─────────────────────────────────────────────────

79
tests/test_litellm/ocr/test_rust_bridge.py
View file

@ -295,6 +295,19 @@ def test_explicit_false_overrides_process_enable():
assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False
@pytest.mark.parametrize(
"litellm_params",
[
{"azure_ad_token_provider": lambda: "token"},
{"azure_username": "user", "azure_password": "password"},
],
)
def test_unsupported_azure_credentials_stay_on_python(litellm_params: dict[str, object]) -> None:
prepared: Final = build_prepared_request(custom_llm_provider="azure_ai", litellm_params=litellm_params)
assert ocr_main._rust_ocr_supported(prepared) is False
def test_load_rust_ocr_returns_injected_impl():
bridge = RecordingBridge()
litellm.rust(True)
@ -662,6 +675,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
"vertex_location": "us-central1",
"vertex_credentials": "redacted",
},
extra_headers={"x-trace-id": "trace-1"},
optional_params={"include_image_base64": True},
timeout=None,
),
@ -672,7 +686,70 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
"include_image_base64": True,
"vertex_project": "project-1",
"vertex_location": "us-central1",
"vertex_credentials": "redacted",
}
assert bridge.calls[0]["api_key"] == "sk-test"
assert bridge.calls[0]["extra_headers"] == {"x-trace-id": "trace-1"}
@pytest.mark.parametrize(
"litellm_params,resolved_assertion,expected_auth_params",
[
(
{"azure_ad_token": "entra-token", "azure_scope": "scope/.default"},
None,
{"azure_ad_token": "entra-token", "azure_scope": "scope/.default"},
),
(
{"tenant_id": "tenant", "client_id": "client", "client_secret": "secret"},
None,
{"tenant_id": "tenant", "client_id": "client", "client_secret": "secret"},
),
(
{"azure_ad_token": "oidc/token-reference", "tenant_id": "tenant", "client_id": "client"},
"assertion",
{
"azure_ad_token": "oidc/token-reference",
"azure_ad_client_assertion": "assertion",
"tenant_id": "tenant",
"client_id": "client",
},
),
],
)
def test_prepare_rust_ocr_call_forwards_azure_auth_configuration(
litellm_params: dict[str, object],
resolved_assertion: str | None,
expected_auth_params: dict[str, object],
) -> None:
bridge: Final = RecordingBridge()
logging_obj: Final = RecordingLogging()
litellm.rust(True)
rust_bridge.set_rust_ocr(ocr=bridge)
ocr_main._run_rust_ocr(
prepared_request=build_prepared_request(
logging_obj=logging_obj,
custom_llm_provider="azure_ai",
model="doc-intelligence/prebuilt-layout",
api_key=None,
extra_headers={"x-trace-id": "trace-1"},
litellm_params=litellm_params,
timeout=None,
),
resolve_api_key=lambda name: resolved_assertion if name == "oidc/token-reference" else None,
)
assert bridge.calls[0]["api_key"] is None
assert bridge.calls[0]["extra_headers"] == {"x-trace-id": "trace-1"}
assert bridge.calls[0]["optional_params"] == {
**expected_auth_params,
"azure_ad_token_refresh": False,
}
assert logging_obj.pre_call_kwargs is not None
logged_input: Final = logging_obj.pre_call_kwargs["additional_args"]
assert "document" not in str(logged_input)
assert all(name not in str(logged_input) for name in expected_auth_params)
def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager():
@ -758,7 +835,7 @@ def test_run_rust_ocr_runs_pre_call_logging():
assert logging_obj.pre_call_kwargs["input"] == "OCR document processing"
additional_args = logging_obj.pre_call_kwargs["additional_args"]
complete_input = additional_args["complete_input_dict"]
assert complete_input["document"] == DOCUMENT
assert "document" not in complete_input
assert complete_input["include_image_base64"] is True
assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr"
assert additional_args["headers"] == {

2
type-discipline-budget.json
View file

@ -3,7 +3,7 @@
"limit": 22292
},
"LIT002": {
"limit": 26749
"limit": 26744
},
"LIT003": {
"limit": 261