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

@@ -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",

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",

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",

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();

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(),
             ),

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")

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

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";

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}")]

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;

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<dyn AddressResolver>,
+    allow_private_network: bool,
+}
+
+type AddressResolution<'a> = Pin<Box<dyn Future<Output = io::Result<Vec<SocketAddr>>> + 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<u8>,
+    pub(crate) content_type: String,
+}
+
+impl MediaFetcher {
+    pub(crate) fn new() -> Result<Self, reqwest::Error> {
+        Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver))
+    }
+
+    fn with_resolvers<R>(
+        transport_resolver: Arc<R>,
+        address_resolver: Arc<dyn AddressResolver>,
+    ) -> Result<Self, reqwest::Error>
+    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<DownloadedMedia, MediaError> {
+        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<DownloadedMedia, MediaError> {
+        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::<IpAddr>() {
+            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::<Vec<_>>())
+        })
+    }
+}
+
+#[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<dyn std::error::Error + Send + Sync>)?
+                .collect::<Vec<_>>();
+            validate_addresses(&addresses).map_err(|_| {
+                Box::new(io::Error::other("destination rejected by network policy"))
+                    as Box<dyn std::error::Error + Send + Sync>
+            })?;
+            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<Vec<String>>, 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)
+        ));
+    }
+}

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<reqwest::Request, OcrError> {
+        let ParsedProviderParams {
+            known: params,
+            extra_params: _extra_params,
+        } = _prepare_ocr_request::<MistralOcrParams>(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, &params)?;
+        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<LiteLLMOcrResponse, OcrResponseError> {
+        mistral::transform_ocr_response(&request.model, response)
+    }
+}
+
+fn get_complete_url(
+    api_base: Option<&str>,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> Result<String, OcrError> {
+    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<String>,
+) -> Result<Vec<(String, String)>, 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<String>) -> Option<String> {
+    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())
+        );
+    }
+}

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;
         }
     };
 }

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<Self, TransportError> {
-        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),
+        }
     }
 }
 

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<Option<Self>, 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<Vec<u8>, 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<OcrDocument, OcrError> {
+    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"));
+    }
+}

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)]

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;

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")
+}

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<String>,
     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,
         }
     }
 }

litellm-rust/crates/core/src/ocr/wire.rs

@@ -70,6 +70,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
         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,

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(

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"));
+}

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(

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.

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")));
+    }
+}

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: