From ae6a4a2f2aeba45a70b2201c74c9c0ff8367322e Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 11 Sep 2026 13:03:06 -0700 Subject: [PATCH] feat(ocr): add Azure Mistral adapter and document fetching (#40533) * feat(ocr): add Azure Mistral adapter and document fetching * fix(ocr): decline missing Azure credentials * fix(ocr): map Azure credentials in gateway errors * refactor(ocr): preserve Azure Mistral extra params * refactor(ocr): adopt request preparation contract --- litellm-rust/Cargo.lock | 7 + .../src/audio_transcription/hooks.rs | 5 +- .../crates/ai-gateway/src/ocr/hooks.rs | 5 +- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 12 + .../ai-gateway/src/routes/messages/mod.rs | 4 +- .../crates/ai-gateway/tests/ocr_lifecycle.rs | 51 ++ litellm-rust/crates/core/Cargo.toml | 2 +- litellm-rust/crates/core/src/constants.rs | 7 + litellm-rust/crates/core/src/error.rs | 28 + litellm-rust/crates/core/src/lib.rs | 1 + litellm-rust/crates/core/src/media.rs | 528 ++++++++++++++++++ .../core/src/ocr/adapters/azure_mistral.rs | 146 +++++ .../crates/core/src/ocr/adapters/mod.rs | 3 + litellm-rust/crates/core/src/ocr/client.rs | 19 +- litellm-rust/crates/core/src/ocr/document.rs | 214 +++++++ litellm-rust/crates/core/src/ocr/error.rs | 16 + litellm-rust/crates/core/src/ocr/mod.rs | 4 + litellm-rust/crates/core/src/ocr/registry.rs | 12 + litellm-rust/crates/core/src/ocr/types.rs | 24 + litellm-rust/crates/core/src/ocr/wire.rs | 1 + .../providers/azure_ai/ocr/transformation.rs | 7 +- .../crates/core/tests/azure_ai_ocr.rs | 75 +++ litellm-rust/crates/core/tests/ocr/support.rs | 6 +- .../crates/python-bridge/src/errors.rs | 2 + .../crates/python-bridge/src/routes/ocr.rs | 16 + .../rust_bridge/native_route_wheel_test.py | 26 +- 26 files changed, 1205 insertions(+), 16 deletions(-) create mode 100644 litellm-rust/crates/core/src/media.rs create mode 100644 litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs create mode 100644 litellm-rust/crates/core/src/ocr/document.rs create mode 100644 litellm-rust/crates/core/tests/azure_ai_ocr.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 0e8e6e09a21..a7c514270da 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -878,6 +878,12 @@ version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +[[package]] +name = "data-url" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be1e0bca6c3637f992fc1cc7cbc52a78c1ef6db076dbf1059c4323d6a2048376" + [[package]] name = "deranged" version = "0.5.8" @@ -1608,6 +1614,7 @@ dependencies = [ "aws-smithy-runtime-api", "aws-types", "base64 0.22.1", + "data-url", "rand 0.8.7", "reqwest", "rstest", diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 098ce071efc..6f4a573e6ff 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -269,7 +269,10 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error { fn core_error_kind(error: &Error) -> &'static str { match error { - Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError", + Error::Auth(_) + | Error::MissingApiKey { .. } + | Error::MissingAzureAiCredentials + | Error::MissingAzureAiCredentialsOrAdToken => "AuthError", Error::InvalidProvider(_) => "InvalidProvider", Error::InvalidRequest(_) => "InvalidRequest", Error::InvalidType { .. } => "InvalidType", diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index d1dd811f7ab..c7e8344aafc 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -386,7 +386,10 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error { fn core_error_kind(error: &Error) -> &'static str { match error { - Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError", + Error::Auth(_) + | Error::MissingApiKey { .. } + | Error::MissingAzureAiCredentials + | Error::MissingAzureAiCredentialsOrAdToken => "AuthError", Error::InvalidProvider(_) => "InvalidProvider", Error::InvalidRequest(_) => "InvalidRequest", Error::InvalidType { .. } => "InvalidType", diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index 2acdd232c80..9ed93d779f1 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -31,6 +31,18 @@ mod tests { use super::{OcrRequest, ocr}; use crate::integrations::types::RequestMetadata; + use litellm_core::ocr::wire::is_supported_request; + + #[test] + fn core_activation_excludes_unmigrated_azure_document_intelligence() { + assert!(is_supported_request("model", Some("mistral"))); + assert!(is_supported_request("pixtral-12b", Some("azure_ai"))); + assert!(!is_supported_request( + "doc-intelligence/prebuilt-layout", + Some("azure_ai") + )); + assert!(!is_supported_request("parse-v3", Some("reducto"))); + } async fn read_http_request(socket: &mut TcpStream) -> String { let mut request = Vec::new(); diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index c22d05f5726..adfbf2b5910 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -115,7 +115,9 @@ impl IntoResponse for MessagesRouteError { | Error::InvalidResponse(_) | Error::InvalidType { .. } | Error::MissingField(_) - | Error::MissingApiKey { .. } => ( + | Error::MissingApiKey { .. } + | Error::MissingAzureAiCredentials + | Error::MissingAzureAiCredentialsOrAdToken => ( StatusCode::BAD_GATEWAY, "messages provider request failed".to_string(), ), diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index 60e90ed2a7c..5502a467511 100644 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -236,6 +236,57 @@ fn base_ocr_request(model: &str) -> OcrRequest<'_> { } } +#[tokio::test] +async fn azure_mistral_uses_prepared_authorization_through_gateway() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let api_base = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let request = read_http_request(&mut socket).await; + let body = br#"{"pages":[]}"#; + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ) + .as_bytes(), + ) + .await + .unwrap(); + socket.write_all(body).await.unwrap(); + request + }); + let request = OcrRequest { + model: "mistral-ocr-2505", + document: json!({ + "type":"document_url", + "document_url":"data:application/pdf;base64,YWJj" + }), + api_key: None, + api_base: Some(&api_base), + custom_llm_provider: Some("azure_ai"), + extra_headers: Some(Map::from_iter([( + "Authorization".into(), + json!("Bearer python-prepared-token"), + )])), + optional_params: Map::new(), + timeout: None, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + }; + + ocr(request).await.unwrap(); + let sent = server.await.unwrap(); + assert!(sent.starts_with("POST /providers/mistral/azure/ocr ")); + assert!( + sent.to_ascii_lowercase() + .contains("authorization: bearer python-prepared-token\r\n") + ); +} + #[tokio::test] async fn reducto_during_call_guardrail_blocks_before_upload() { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index b4ca88cf16f..6d0a2fa775b 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -12,6 +12,7 @@ path = "tests/workspace_crate_allowlist.rs" [dependencies] base64.workspace = true +data-url = "0.3.2" rand.workspace = true reqwest.workspace = true serde.workspace = true @@ -44,5 +45,4 @@ observability = ["dep:tracing-subscriber"] [dev-dependencies] rstest.workspace = true -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } tracing-subscriber.workspace = true diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index 108d2a48e30..8a12e186197 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -43,6 +43,13 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str = "[System: Empty message content sanitised to satisfy protocol]"; pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; + +pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10; + pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600; pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10; +pub(crate) const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024; +pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024; +pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10; +pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr"; pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1"; diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 0382314057f..eefa7d606d8 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -21,6 +21,12 @@ pub enum Error { "Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params" )] MissingApiKey { provider: &'static str }, + #[error( + "Missing Azure AI credentials - set AZURE_AI_API_KEY or provide an Authorization header" + )] + MissingAzureAiCredentials, + #[error("Missing Azure AI credentials - set AZURE_AI_API_KEY or provide azure_ad_token")] + MissingAzureAiCredentialsOrAdToken, #[error("upstream request failed with status {status}: {body}")] Http { status: u16, body: String }, #[error("upstream network error: {0}")] @@ -40,6 +46,28 @@ pub enum Error { Unsupported(&'static str), } +#[derive(Debug, ThisError)] +pub(crate) enum MediaError { + #[error("media URL rejected by network policy")] + BlockedUrl, + #[error("media download is disabled")] + DownloadDisabled, + #[error("media download exceeds the maximum size")] + DownloadTooLarge, + #[error("too many redirects while fetching media")] + TooManyRedirects, + #[error("media redirect is missing a Location header")] + MissingRedirectLocation, + #[error("invalid media redirect")] + InvalidRedirect, + #[error("media download failed with status {0}")] + Http(u16), + #[error("media download timed out")] + Timeout, + #[error("{0}")] + Transport(#[from] TransportError), +} + #[derive(Clone, Debug, ThisError, PartialEq, Eq)] pub enum TransportError { #[error("upstream request failed with status {status}: {body}")] diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 3a0896a4d5a..7e81b292441 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -5,6 +5,7 @@ pub mod chat_completions; pub mod constants; pub mod error; pub mod http_utils; +mod media; pub mod messages; #[cfg(any(feature = "observability", test))] pub mod observability; diff --git a/litellm-rust/crates/core/src/media.rs b/litellm-rust/crates/core/src/media.rs new file mode 100644 index 00000000000..5f9a43794c2 --- /dev/null +++ b/litellm-rust/crates/core/src/media.rs @@ -0,0 +1,528 @@ +use std::future::Future; +use std::io; +use std::net::{IpAddr, SocketAddr}; +use std::pin::Pin; +use std::sync::Arc; +use std::time::Duration; + +use reqwest::Url; +use reqwest::dns::{Addrs, Name, Resolve, Resolving}; + +use crate::constants::MEDIA_CONNECT_TIMEOUT_SECS; +use crate::error::{MediaError, TransportError}; + +#[derive(Clone)] +pub(crate) struct MediaFetcher { + client: reqwest::Client, + address_resolver: Arc, + allow_private_network: bool, +} + +type AddressResolution<'a> = Pin>> + Send + 'a>>; + +trait AddressResolver: Send + Sync { + fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a>; +} + +#[derive(Clone, Copy)] +pub(crate) struct DownloadPolicy { + pub(crate) timeout: Duration, + pub(crate) max_bytes: u64, + pub(crate) max_redirects: usize, +} + +#[derive(Debug)] +pub(crate) struct DownloadedMedia { + pub(crate) bytes: Vec, + pub(crate) content_type: String, +} + +impl MediaFetcher { + pub(crate) fn new() -> Result { + Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver)) + } + + fn with_resolvers( + transport_resolver: Arc, + address_resolver: Arc, + ) -> Result + where + R: Resolve + 'static, + { + let client = reqwest::Client::builder() + .connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS)) + .redirect(reqwest::redirect::Policy::none()) + .no_proxy() + .dns_resolver(transport_resolver) + .build()?; + Ok(Self { + client, + address_resolver, + allow_private_network: false, + }) + } + + #[cfg(test)] + pub(crate) fn for_test(client: reqwest::Client) -> Self { + Self { + client, + address_resolver: Arc::new(AllowPrivateResolver), + allow_private_network: true, + } + } + + pub(crate) async fn fetch( + &self, + url: Url, + policy: DownloadPolicy, + ) -> Result { + if policy.max_bytes == 0 { + return Err(MediaError::DownloadDisabled); + } + tokio::time::timeout(policy.timeout, self.fetch_before_deadline(url, policy)) + .await + .map_err(|_| MediaError::Timeout)? + } + + async fn fetch_before_deadline( + &self, + mut url: Url, + policy: DownloadPolicy, + ) -> Result { + let mut redirects_followed = 0; + loop { + self.validate_url(&url).await?; + let mut response = self + .client + .get(url.clone()) + .send() + .await + .map_err(TransportError::from)?; + if response.status().is_redirection() { + if redirects_followed == policy.max_redirects { + return Err(MediaError::TooManyRedirects); + } + let location = response + .headers() + .get(reqwest::header::LOCATION) + .and_then(|value| value.to_str().ok()) + .ok_or(MediaError::MissingRedirectLocation)?; + url = url + .join(location) + .map_err(|_| MediaError::InvalidRedirect)?; + redirects_followed += 1; + continue; + } + if !response.status().is_success() { + return Err(MediaError::Http(response.status().as_u16())); + } + enforce_download_size(response.content_length().unwrap_or(0), policy.max_bytes)?; + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split(';').next()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("application/octet-stream") + .to_string(); + let mut bytes = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(TransportError::from)? { + enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?; + bytes.extend_from_slice(&chunk); + } + return Ok(DownloadedMedia { + bytes, + content_type, + }); + } + } + + async fn validate_url(&self, url: &Url) -> Result<(), MediaError> { + if !matches!(url.scheme(), "http" | "https") + || !url.username().is_empty() + || url.password().is_some() + { + return Err(MediaError::BlockedUrl); + } + let host = url.host_str().ok_or(MediaError::BlockedUrl)?; + if self.allow_private_network { + return Ok(()); + } + if let Ok(ip) = host.parse::() { + return (!is_blocked_ip(ip)) + .then_some(()) + .ok_or(MediaError::BlockedUrl); + } + let port = url.port_or_known_default().ok_or(MediaError::BlockedUrl)?; + let addresses = self + .address_resolver + .resolve(host, port) + .await + .map_err(|error| TransportError::Network(error.to_string()))?; + validate_addresses(&addresses) + } +} + +fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), MediaError> { + if length > max_bytes { + return Err(MediaError::DownloadTooLarge); + } + Ok(()) +} + +fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), MediaError> { + if addresses.is_empty() || addresses.iter().any(|address| is_blocked_ip(address.ip())) { + return Err(MediaError::BlockedUrl); + } + Ok(()) +} + +fn is_blocked_ip(ip: IpAddr) -> bool { + match ip { + IpAddr::V4(ip) => { + let [first, second, third, _] = ip.octets(); + first == 0 + || first == 10 + || first == 127 + || (first == 100 && (64..=127).contains(&second)) + || (first == 169 && second == 254) + || (first == 172 && (16..=31).contains(&second)) + || (first == 192 && second == 0 && (third == 0 || third == 2)) + || (first == 192 && second == 168) + || (first == 192 && second == 88 && third == 99) + || (first == 198 && (second == 18 || second == 19)) + || (first == 198 && second == 51 && third == 100) + || (first == 203 && second == 0 && third == 113) + || first >= 224 + } + IpAddr::V6(ip) => { + let segments = ip.segments(); + ip.is_loopback() + || ip.is_unspecified() + || ip.is_multicast() + || (segments[0] & 0xfe00) == 0xfc00 + || (segments[0] & 0xffc0) == 0xfe80 + || (segments[0] & 0xffc0) == 0xfec0 + || (segments[0] == 0x2001 && segments[1] == 0x0db8) + || ip + .to_ipv4_mapped() + .or_else(|| ip.to_ipv4()) + .map(|ipv4| is_blocked_ip(IpAddr::V4(ipv4))) + .unwrap_or(false) + } + } +} + +#[derive(Default)] +struct PublicDnsResolver; + +struct SystemAddressResolver; + +impl AddressResolver for SystemAddressResolver { + fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> { + Box::pin(async move { + Ok(tokio::net::lookup_host((host, port)) + .await? + .collect::>()) + }) + } +} + +#[cfg(test)] +struct AllowPrivateResolver; + +#[cfg(test)] +impl AddressResolver for AllowPrivateResolver { + fn resolve<'a>(&'a self, _host: &'a str, port: u16) -> AddressResolution<'a> { + Box::pin(async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) }) + } +} + +impl Resolve for PublicDnsResolver { + fn resolve(&self, name: Name) -> Resolving { + let host = name.as_str().to_string(); + Box::pin(async move { + let addresses = tokio::net::lookup_host((host.as_str(), 0)) + .await + .map_err(|error| Box::new(error) as Box)? + .collect::>(); + validate_addresses(&addresses).map_err(|_| { + Box::new(io::Error::other("destination rejected by network policy")) + as Box + })?; + Ok(Box::new(addresses.into_iter()) as Addrs) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has address"); + let task = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let mut request = [0_u8; 1024]; + let bytes_read = socket.read(&mut request).await.expect("reads request"); + assert!(bytes_read > 0); + socket.write_all(response).await.expect("writes response"); + }); + ( + Url::parse(&format!("http://{address}/document")).expect("valid test URL"), + task, + ) + } + + async fn serve_named( + host: &str, + responses: Vec<&'static [u8]>, + ) -> (Url, tokio::task::JoinHandle>, SocketAddr) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has address"); + let task = tokio::spawn(async move { + let mut requests = Vec::with_capacity(responses.len()); + for response in responses { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let mut request = [0_u8; 4096]; + let bytes_read = socket.read(&mut request).await.expect("reads request"); + requests.push(String::from_utf8_lossy(&request[..bytes_read]).into_owned()); + socket.write_all(response).await.expect("writes response"); + } + requests + }); + ( + Url::parse(&format!("http://{host}:{}/document", address.port())) + .expect("valid test URL"), + task, + address, + ) + } + + struct LoopbackDnsResolver(SocketAddr); + + impl Resolve for LoopbackDnsResolver { + fn resolve(&self, _name: Name) -> Resolving { + let address = self.0; + Box::pin(async move { Ok(Box::new(vec![address].into_iter()) as Addrs) }) + } + } + + struct TestAddressResolver { + blocked_hosts: HashSet<&'static str>, + } + + impl AddressResolver for TestAddressResolver { + fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> { + let blocked = self.blocked_hosts.contains(host); + Box::pin(async move { + let ip = if blocked { + IpAddr::from([127, 0, 0, 1]) + } else { + IpAddr::from([8, 8, 8, 8]) + }; + Ok(vec![SocketAddr::new(ip, port)]) + }) + } + } + + fn policy_checked_fetcher( + address: SocketAddr, + blocked_hosts: HashSet<&'static str>, + ) -> MediaFetcher { + MediaFetcher::with_resolvers( + Arc::new(LoopbackDnsResolver(address)), + Arc::new(TestAddressResolver { blocked_hosts }), + ) + .expect("test fetcher builds") + } + + fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy { + DownloadPolicy { + timeout: Duration::from_secs(1), + max_bytes, + max_redirects, + } + } + + #[test] + fn blocks_non_public_addresses() { + for address in [ + "0.0.0.1", + "10.0.0.1", + "100.64.0.1", + "127.0.0.1", + "169.254.1.1", + "172.16.0.1", + "192.168.0.1", + "198.18.0.1", + "198.51.100.1", + "203.0.113.1", + "224.0.0.1", + "::1", + "fc00::1", + "fe80::1", + "2001:db8::1", + "::ffff:127.0.0.1", + ] { + assert!(is_blocked_ip(address.parse().expect("valid test address"))); + } + assert!(!is_blocked_ip( + "8.8.8.8".parse().expect("valid public address") + )); + } + + #[tokio::test] + async fn fetches_exact_limit_and_normalizes_content_type() { + let (url, server) = serve( + b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", + ) + .await; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test client builds"); + let media = MediaFetcher::for_test(client) + .fetch(url, policy(3, 0)) + .await + .expect("download succeeds at exact limit"); + server.await.expect("server completes"); + assert_eq!(media.bytes, b"abc"); + assert_eq!(media.content_type, "application/pdf"); + } + + #[tokio::test] + async fn rejects_declared_oversize_body() { + let (url, server) = serve( + b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", + ) + .await; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test client builds"); + let error = MediaFetcher::for_test(client) + .fetch(url, policy(2, 0)) + .await + .expect_err("oversize body is rejected"); + server.await.expect("server completes"); + assert!(matches!(error, MediaError::DownloadTooLarge)); + } + + #[tokio::test] + async fn rejects_streamed_oversize_body() { + let (url, server) = serve( + b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n", + ) + .await; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test client builds"); + let error = MediaFetcher::for_test(client) + .fetch(url, policy(3, 0)) + .await + .expect_err("stream crossing limit is rejected"); + server.await.expect("server completes"); + assert!(matches!(error, MediaError::DownloadTooLarge)); + } + + #[tokio::test] + async fn follows_allowed_redirects_and_revalidates_each_destination() { + let (url, server, address) = serve_named( + "public.test", + vec![ + b"HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\n\r\n", + b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok", + ], + ) + .await; + let media = policy_checked_fetcher(address, HashSet::new()) + .fetch(url, policy(2, 1)) + .await + .expect("redirected fetch succeeds"); + let requests = server.await.expect("server completes"); + assert_eq!(requests.len(), 2); + assert!(requests[1].starts_with("GET /final ")); + assert_eq!(media.bytes, b"ok"); + } + + #[tokio::test] + async fn blocks_redirected_private_destination_before_second_request() { + let (url, server, address) = serve_named( + "public.test", + vec![b"HTTP/1.1 302 Found\r\nLocation: http://blocked.test/document\r\nContent-Length: 0\r\n\r\n"], + ) + .await; + let error = policy_checked_fetcher(address, HashSet::from(["blocked.test"])) + .fetch(url, policy(10, 1)) + .await + .expect_err("private redirect is rejected"); + let requests = server.await.expect("server completes"); + assert_eq!(requests.len(), 1); + assert!(matches!(error, MediaError::BlockedUrl)); + } + + #[tokio::test] + async fn enforces_total_timeout() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has address"); + let server = tokio::spawn(async move { + let (_socket, _) = listener.accept().await.expect("accepts request"); + tokio::time::sleep(Duration::from_millis(100)).await; + }); + let url = Url::parse(&format!("http://public.test:{}/document", address.port())) + .expect("valid test URL"); + let error = policy_checked_fetcher(address, HashSet::new()) + .fetch( + url, + DownloadPolicy { + timeout: Duration::from_millis(20), + max_bytes: 10, + max_redirects: 0, + }, + ) + .await + .expect_err("fetch times out"); + server.await.expect("server completes"); + assert!(matches!(error, MediaError::Timeout)); + } + + #[tokio::test] + async fn document_client_does_not_send_ambient_credentials() { + let (url, server, address) = serve_named( + "public.test", + vec![b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"], + ) + .await; + policy_checked_fetcher(address, HashSet::new()) + .fetch(url, policy(2, 0)) + .await + .expect("fetch succeeds"); + let requests = server.await.expect("server completes"); + assert!(!requests[0].to_ascii_lowercase().contains("authorization:")); + assert!(!requests[0].to_ascii_lowercase().contains("api-key:")); + } + + #[tokio::test] + async fn rejects_url_credentials_before_network_access() { + let fetcher = MediaFetcher::new().expect("media fetcher builds"); + let url = + Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses"); + assert!(matches!( + fetcher.validate_url(&url).await, + Err(MediaError::BlockedUrl) + )); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs new file mode 100644 index 00000000000..468c883a1dd --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs @@ -0,0 +1,146 @@ +use super::OcrAdapter; +use crate::Error; +use crate::constants::AZURE_AI_OCR_PATH; +use crate::ocr::OcrClient; +use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; +use crate::ocr::document::{inline_remote_document, validate_inline_document}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{ + _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, +}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; +use crate::url_utils::ApiUrl; + +const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; +const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; + +#[derive(Clone, Debug)] +pub(crate) struct AzureMistralAdapter; + +impl OcrAdapter for AzureMistralAdapter { + type ProviderResponse = MistralOcrResponse; + const PROVIDER: OcrProvider = OcrProvider::AzureAi; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let ParsedProviderParams { + known: params, + extra_params: _extra_params, + } = _prepare_ocr_request::(request)?; + let headers = authenticate(&request.connection, &credential_env)?; + let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?; + let document = inline_remote_document( + client.document_fetcher(), + request.document.clone(), + &request.connection, + ) + .await?; + let body = mistral::transform_ocr_request(&request.model, document, ¶ms)?; + transform_request_body(client, request, &url, &headers, body, |body| { + validate_inline_document(&body.document) + }) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + mistral::transform_ocr_response(&request.model, response) + } +} + +fn get_complete_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let base = nonblank(api_base.map(str::to_string)) + .or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV))) + .ok_or_else(|| Error::Auth( + "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter".into(), + ))?; + let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect(); + ApiUrl::parse(&base) + .and_then(|url| url.complete_path(&path)) + .map(|url| url.into_string()) + .map_err(|_| { + OcrRequestError::RequestField { + path: "api_base".into(), + } + .into() + }) +} + +fn authenticate( + connection: &OcrConnection, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result, OcrError> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let key = nonblank(connection.api_key.clone()) + .or_else(|| nonblank(env_lookup(AZURE_AI_API_KEY_ENV))) + .ok_or(Error::MissingAzureAiCredentials)?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) +} + +fn nonblank(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completes_azure_path_and_preserves_query() { + assert_eq!( + get_complete_url(Some("https://example.com/?tenant=a"), &|_| None).unwrap(), + "https://example.com/providers/mistral/azure/ocr?tenant=a" + ); + assert_eq!( + get_complete_url( + Some("https://example.com/providers/mistral/azure/ocr"), + &|_| None + ) + .unwrap(), + "https://example.com/providers/mistral/azure/ocr" + ); + } + + #[test] + fn supplied_authorization_precedes_keys() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + extra_headers: vec![("authorization".into(), "Bearer prepared".into())], + ..Default::default() + }; + assert_eq!( + authenticate(&connection, &|_| Some("environment-key".into())).unwrap(), + connection.extra_headers + ); + } + + #[test] + fn request_key_precedes_environment_key() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + ..Default::default() + }; + assert_eq!( + authenticate(&connection, &|_| Some("environment-key".into())).unwrap()[0], + ("Authorization".into(), "Bearer request-key".into()) + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/mod.rs index 7dfb08b4dc2..bbd6feb6c7b 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/mod.rs @@ -8,8 +8,10 @@ use super::registry::OcrProvider; use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat}; use super::wire::DecodedOcrResponse; +mod azure_mistral; mod mistral; +pub(crate) use azure_mistral::AzureMistralAdapter; pub(crate) use mistral::MistralAdapter; /// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response. @@ -62,6 +64,7 @@ macro_rules! for_each_ocr_adapter { ($callback:ident) => { $callback! { Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral; + AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi; } }; } diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 36bfc036bae..699e44412ac 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -10,15 +10,21 @@ use super::wire::{DecodedOcrResponse, decode_response}; use crate::Error; use crate::constants::OCR_CONNECT_TIMEOUT_SECS; use crate::error::TransportError; +use crate::media::MediaFetcher; #[derive(Clone)] pub struct OcrClient { provider_http: reqwest::Client, + document_fetcher: MediaFetcher, } impl OcrClient { pub fn new(provider_http: reqwest::Client) -> Result { - Ok(Self { provider_http }) + let document_fetcher = MediaFetcher::new().map_err(TransportError::from)?; + Ok(Self { + provider_http, + document_fetcher, + }) } #[tracing::instrument( @@ -35,9 +41,16 @@ impl OcrClient { &self.provider_http } + pub(crate) fn document_fetcher(&self) -> &MediaFetcher { + &self.document_fetcher + } + #[cfg(test)] - pub(crate) fn for_test(provider_http: reqwest::Client) -> Self { - Self { provider_http } + pub(crate) fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self { + Self { + provider_http, + document_fetcher: MediaFetcher::for_test(document_http), + } } } diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs new file mode 100644 index 00000000000..11a6e612a3c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -0,0 +1,214 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +#[cfg(test)] +use data_url::mime::Mime; +use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError}; +use reqwest::Url; + +use super::error::{OcrError, OcrRequestError, OcrResponseError}; +use super::types::{OcrConnection, OcrDocument}; +use crate::constants::OCR_MAX_FETCH_REDIRECTS; +use crate::error::{MediaError, TransportError}; +use crate::media::{DownloadPolicy, MediaFetcher}; + +pub(crate) struct InlineDocument<'a>(DataUrl<'a>); + +impl<'a> InlineDocument<'a> { + pub(crate) fn parse(source: &'a str) -> Result, OcrRequestError> { + match DataUrl::process(source) { + Ok(url) => Ok(Some(Self(url))), + Err(DataUrlError::NotADataUrl) => Ok(None), + Err(DataUrlError::NoComma) => Err(OcrRequestError::InvalidDataUri), + } + } + + #[cfg(test)] + pub(crate) fn mime_type(&self) -> &Mime { + self.0.mime_type() + } + + pub(crate) fn decode(&self, max_bytes: usize) -> Result, OcrRequestError> { + let mut body = Vec::new(); + self.0 + .decode(|bytes| { + if bytes.len() > max_bytes.saturating_sub(body.len()) { + return Err(OcrRequestError::InlineDocumentTooLarge); + } + body.extend_from_slice(bytes); + Ok(()) + }) + .map_err(|error| match error { + DecodeError::InvalidBase64(_) => OcrRequestError::InvalidDataUri, + DecodeError::WriteError(error) => error, + })?; + Ok(body) + } +} + +pub(crate) fn validate_inline_document(document: &OcrDocument) -> Result<(), OcrRequestError> { + let inline = + InlineDocument::parse(document.source())?.ok_or(OcrRequestError::InvalidDataUri)?; + inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; + Ok(()) +} + +pub(crate) async fn inline_remote_document( + fetcher: &MediaFetcher, + document: OcrDocument, + connection: &OcrConnection, +) -> Result { + let source = document.source(); + if !source.starts_with("http://") && !source.starts_with("https://") { + validate_inline_document(&document)?; + return Ok(document); + } + let url = Url::parse(source).map_err(|_| OcrRequestError::RequestField { + path: "document URL".into(), + })?; + let downloaded = fetcher + .fetch( + url, + DownloadPolicy { + timeout: connection.timeout, + max_bytes: connection.max_download_bytes, + max_redirects: OCR_MAX_FETCH_REDIRECTS, + }, + ) + .await + .map_err(map_media_error)?; + let result = document.with_source(format!( + "data:{};base64,{}", + downloaded.content_type, + STANDARD.encode(downloaded.bytes) + )); + validate_inline_document(&result)?; + Ok(result) +} + +fn map_media_error(error: MediaError) -> OcrError { + match error { + MediaError::BlockedUrl => OcrRequestError::BlockedDocumentUrl.into(), + MediaError::DownloadDisabled => OcrRequestError::DownloadDisabled.into(), + MediaError::DownloadTooLarge => OcrRequestError::DownloadTooLarge.into(), + MediaError::TooManyRedirects => OcrRequestError::TooManyRedirects.into(), + MediaError::MissingRedirectLocation => OcrResponseError::MissingRedirectLocation.into(), + MediaError::InvalidRedirect => OcrResponseError::InvalidRedirect.into(), + MediaError::Http(status) => TransportError::Http { + status, + body: "OCR document download failed".into(), + } + .into(), + MediaError::Timeout => { + TransportError::Network("OCR document download timed out".into()).into() + } + MediaError::Transport(error) => error.into(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::Map; + + fn document(source: &str) -> OcrDocument { + OcrDocument::DocumentUrl { + document_url: source.into(), + extra_fields: Map::new(), + } + } + + #[test] + fn decodes_data_urls_and_limits_decoded_size() { + for (source, expected) in [ + ("data:application/pdf;base64,YWJj", b"abc".as_slice()), + ("DATA:application/pdf;BASE64,YWI", b"ab".as_slice()), + ("data:,a%20b%00%FF", b"a b\0\xff".as_slice()), + ] { + let inline = InlineDocument::parse(source).unwrap().unwrap(); + assert_eq!(inline.decode(expected.len()).unwrap(), expected); + assert_eq!( + inline.decode(expected.len() - 1), + Err(OcrRequestError::InlineDocumentTooLarge) + ); + } + } + + #[test] + fn preserves_mime_parameters_and_standard_default() { + let inline = InlineDocument::parse("data:application/pdf;version=1.7;base64,YQ==") + .unwrap() + .unwrap(); + assert!(inline.mime_type().matches("application", "pdf")); + assert_eq!(inline.mime_type().get_parameter("version"), Some("1.7")); + let default = InlineDocument::parse("data:,a").unwrap().unwrap(); + assert!(default.mime_type().matches("text", "plain")); + assert_eq!( + default.mime_type().get_parameter("charset"), + Some("US-ASCII") + ); + } + + #[test] + fn rejects_invalid_inline_documents() { + for source in [ + "https://example.com/document.pdf", + "data:application/pdf;base64", + "data:application/pdf;base64,INVALID!", + ] { + assert!(validate_inline_document(&document(source)).is_err()); + } + } + + #[tokio::test] + async fn remote_conversion_preserves_kind_and_isolates_provider_credentials() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0_u8; 2048]; + let count = socket.read(&mut request).await.unwrap(); + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Type: image/png; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc") + .await + .unwrap(); + String::from_utf8_lossy(&request[..count]).into_owned() + }); + let mut provider_headers = reqwest::header::HeaderMap::new(); + provider_headers.insert( + reqwest::header::AUTHORIZATION, + reqwest::header::HeaderValue::from_static("Bearer provider-secret"), + ); + let provider_http = reqwest::Client::builder() + .default_headers(provider_headers) + .build() + .unwrap(); + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .unwrap(); + let client = super::super::OcrClient::for_test(provider_http, document_http); + let converted = inline_remote_document( + client.document_fetcher(), + OcrDocument::ImageUrl { + image_url: format!("http://{address}/image"), + extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]), + }, + &OcrConnection::default(), + ) + .await + .unwrap(); + let request = server.await.unwrap(); + + assert_eq!( + converted, + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,YWJj".into(), + extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]), + } + ); + assert!(!request.to_ascii_lowercase().contains("authorization")); + assert!(!request.contains("provider-secret")); + } +} diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 50793aad566..395bb60000c 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -10,12 +10,28 @@ pub enum OcrRequestError { RequestField { path: String }, #[error("missing required field: {0}")] MissingField(&'static str), + #[error("invalid OCR document data URI")] + InvalidDataUri, + #[error("inline OCR document exceeds the size limit")] + InlineDocumentTooLarge, + #[error("OCR document URL is blocked by network policy")] + BlockedDocumentUrl, + #[error("OCR document downloads are disabled")] + DownloadDisabled, + #[error("OCR document download exceeds the size limit")] + DownloadTooLarge, + #[error("OCR document download exceeded the redirect limit")] + TooManyRedirects, } #[derive(Debug, Clone, PartialEq, Eq, Error)] pub enum OcrResponseError { #[error("invalid OCR response field: {path}")] ResponseField { path: String }, + #[error("OCR document redirect is missing a location")] + MissingRedirectLocation, + #[error("OCR document redirect location is invalid")] + InvalidRedirect, } #[derive(Debug, Error)] diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 78c86aeaa0e..aa11d0ab3cf 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,6 +1,7 @@ mod adapters; pub mod client; mod codecs; +mod document; pub mod error; mod handler; pub mod hooks; @@ -13,6 +14,9 @@ pub mod wire; pub use client::{OcrClient, ocr}; pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; +#[cfg(test)] +#[path = "../../tests/azure_ai_ocr.rs"] +mod azure_ai_tests; #[cfg(test)] #[path = "../../tests/ocr/support.rs"] pub(crate) mod test_support; diff --git a/litellm-rust/crates/core/src/ocr/registry.rs b/litellm-rust/crates/core/src/ocr/registry.rs index 9bae8153ee8..e70bd7314c0 100644 --- a/litellm-rust/crates/core/src/ocr/registry.rs +++ b/litellm-rust/crates/core/src/ocr/registry.rs @@ -24,12 +24,14 @@ super::adapters::for_each_ocr_adapter!(define_adapter_types); #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum OcrProvider { Mistral, + AzureAi, } impl OcrProvider { pub(crate) const fn as_str(self) -> &'static str { match self { Self::Mistral => "mistral", + Self::AzureAi => "azure_ai", } } } @@ -45,9 +47,19 @@ pub(crate) fn resolve_wire_adapter( }); let typed_provider = match provider.custom_llm_provider { "mistral" => OcrProvider::Mistral, + "azure_ai" => OcrProvider::AzureAi, value => return Err(Error::InvalidProvider(value.to_string())), }; match typed_provider { OcrProvider::Mistral => Ok((provider.model.to_string(), OcrAdapterKind::Mistral)), + OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { + Err(Error::InvalidProvider("azure_ai".into())) + } + OcrProvider::AzureAi => Ok((provider.model.to_string(), OcrAdapterKind::AzureMistral)), } } + +fn is_document_intelligence_model(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + model.contains("doc-intelligence") || model.contains("documentintelligence") +} diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 02b5330188c..eeda94738a0 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -32,6 +32,28 @@ pub enum OcrDocument { }, } +impl OcrDocument { + pub(crate) fn source(&self) -> &str { + match self { + Self::DocumentUrl { document_url, .. } => document_url, + Self::ImageUrl { image_url, .. } => image_url, + } + } + + pub(crate) fn with_source(self, source: String) -> Self { + match self { + Self::DocumentUrl { extra_fields, .. } => Self::DocumentUrl { + document_url: source, + extra_fields, + }, + Self::ImageUrl { extra_fields, .. } => Self::ImageUrl { + image_url: source, + extra_fields, + }, + } + } +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum OcrResponseFormat { @@ -46,6 +68,7 @@ pub struct OcrConnection { pub api_base: Option, pub extra_headers: Vec<(String, String)>, pub timeout: Duration, + pub max_download_bytes: u64, } impl Default for OcrConnection { @@ -55,6 +78,7 @@ impl Default for OcrConnection { api_base: None, extra_headers: Vec::new(), timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS), + max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES, } } } diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index f37c06a01ac..db8d91905c4 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -70,6 +70,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result api_base: nonblank(wire.api_base), extra_headers: headers, timeout: timeout.unwrap_or(defaults.timeout), + max_download_bytes: defaults.max_download_bytes, }; Ok(LiteLLMOcrRequest { connection, diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index 2af25a5639a..2dbec8e2187 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -125,12 +125,7 @@ pub fn validate_azure_ai_environment( } 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(), - ) - }) + .ok_or(Error::MissingAzureAiCredentialsOrAdToken) } pub fn validate_document_intelligence_environment( diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs new file mode 100644 index 00000000000..3d7fb2e54a3 --- /dev/null +++ b/litellm-rust/crates/core/tests/azure_ai_ocr.rs @@ -0,0 +1,75 @@ +use std::sync::Arc; + +use serde_json::{Value, json}; + +use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks}; +use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + +#[tokio::test] +async fn facade_executes_azure_mistral_with_prepared_auth() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"include_image_base64":true}), + ); + request.connection.api_key = None; + request.connection.extra_headers = vec![( + "Authorization".into(), + "Bearer python-prepared-token".into(), + )]; + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0]["markdown"], "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer python-prepared-token\r\n") + ); + let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "include_image_base64":true + }) + ); +} + +struct ReplaceBodyDocument; + +impl OcrHooks for ReplaceBodyDocument { + fn has_guardrails(&self) -> bool { + true + } + + fn during_call( + &self, + mut request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + request.body["document"] = json!({ + "type":"document_url", + "document_url":"https://example.com/not-inline.pdf" + }); + Ok(request) + }) + } +} + +#[tokio::test] +async fn rejects_non_inline_body_after_guardrails() { + let mut request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); + request.hooks = Arc::new(ReplaceBodyDocument); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("data URI")); +} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index d9720019419..45047a0e62a 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -8,7 +8,11 @@ use crate::ocr::wire::{OcrWireRequest, decode_request}; use crate::ocr::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient}; pub(crate) fn ocr_client() -> OcrClient { - OcrClient::for_test(reqwest::Client::new()) + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test document client builds"); + OcrClient::for_test(reqwest::Client::new(), document_http) } pub(crate) async fn perform_ocr( diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 407475dedce..0c30eab8112 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -42,6 +42,8 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { | Error::InvalidType { .. } | Error::MissingField(_) | Error::MissingApiKey { .. } + | Error::MissingAzureAiCredentials + | Error::MissingAzureAiCredentialsOrAdToken | Error::Routing(_) // Nothing reached the provider, so serving it on Python cannot double // bill and is the only way the caller gets an answer at all. diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 951caf4eef4..2e6900f784f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -87,3 +87,19 @@ bridge_route! { prepare = prepare_ocr, errors = ocr_error_to_pyerr, } + +#[cfg(test)] +mod tests { + use litellm_core::ocr::wire::is_supported_request; + + #[test] + fn native_activation_excludes_unmigrated_azure_document_intelligence() { + assert!(is_supported_request("model", Some("mistral"))); + assert!(is_supported_request("pixtral-12b", Some("azure_ai"))); + assert!(!is_supported_request( + "documentintelligence/prebuilt-read", + Some("azure_ai") + )); + assert!(!is_supported_request("mistral-ocr", Some("vertex_ai"))); + } +} diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 9807febbff4..5b2af3cf02e 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -73,7 +73,7 @@ def assert_native_request( headers: HTTPMessage, body: object, ) -> None: - if route not in {"ocr", "transcription", "messages", "chat_completions"}: + if route not in {"ocr", "azure_ocr", "transcription", "messages", "chat_completions"}: raise AssertionError(f"unexpected route marker: {route!r}") if outcome not in {"success", "429", "hang"}: raise AssertionError(f"unexpected outcome marker: {outcome!r}") @@ -86,6 +86,12 @@ def assert_native_request( assert body["document"]["document_url"] == "https://example.com/document.pdf" assert body["include_image_base64"] is True return + if route == "azure_ocr": + assert path == "/providers/mistral/azure/ocr" + assert headers.get("authorization") == "Bearer prepared-azure-token" + assert body["model"] == "mistral-ocr-2505" + assert body["document"]["document_url"] == "data:application/pdf;base64,YWJj" + return if route == "transcription": assert path == "/model/mistral.voxtral-mini-3b-2507/converse" assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ") @@ -107,7 +113,7 @@ def assert_native_request( def native_response(status: int, route: str | None) -> bytes: if status == 429: return b'{"error":"native-rate-limit"}' - if route == "ocr": + if route in {"ocr", "azure_ocr"}: return b'{"pages":[{"index":0,"markdown":"native-ocr"}]}' if route == "transcription": return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}' @@ -181,6 +187,20 @@ def assert_success(route: str, response: object) -> None: raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}") +def azure_ocr_kwargs(api_base: str) -> dict[str, object]: + return { + "model": "mistral-ocr-2505", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_base": api_base, + "custom_llm_provider": "azure_ai", + "extra_headers": { + "Authorization": "Bearer prepared-azure-token", + "x-test-outcome": "success", + "x-test-route": "azure_ocr", + }, + } + + def success_value(route: str, response: dict[object, object]) -> object: if route == "ocr": return response["pages"][0]["markdown"] @@ -211,6 +231,7 @@ def exercise_sync(native: object, api_base: str) -> None: assert_rate_limit(native, route, error) else: raise AssertionError(f"{route} accepted a 429 response") + assert_success("ocr", native.ocr(**azure_ocr_kwargs(api_base))) async def exercise_async(native: object, api_base: str) -> None: @@ -223,6 +244,7 @@ async def exercise_async(native: object, api_base: str) -> None: assert_rate_limit(native, route, error) else: raise AssertionError(f"a{route} accepted a 429 response") + assert_success("ocr", await native.aocr(**azure_ocr_kwargs(api_base))) async def exercise_async_concurrency(native: object, api_base: str) -> None: