From 726dbf0d0d10ef256a7f2c2e096ba90493c3d60d Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 17 Sep 2026 17:48:52 -0700
Subject: [PATCH] fix(ocr): keep a downloaded document inlined when callbacks
 intercept the request

Providers that cannot fetch a public document URL themselves (Azure AI
mistral document AI, Azure cohere parse, Vertex AI) download it and
inline it as a data URI. When a pre-call callback or debug logging
intercepts the request, the Python host hands the caller's original
document back into the body, so the provider request carried the URL
again and Azure's inline-only check rejected it with "invalid OCR
document data URI". The core now keeps the prepared document when a
hook returns the untouched caller document, while a hook that edits or
replaces the document still wins
---
 .../src/llms/azure_ai/ocr/transformation.rs   | 50 +++++++++++++++++++
 litellm-rust/crates/core/src/ocr/prepare.rs   | 20 ++++++--
 2 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs
index 36a07fca8a9..99f0b2af07b 100644
--- a/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs
+++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs
@@ -402,4 +402,54 @@ mod tests {
         let error = perform_ocr(request).await.unwrap_err();
         assert!(error.to_string().contains("data URI"));
     }
+
+    struct EchoCallerDocument(Value);
+
+    impl OcrHooks for EchoCallerDocument {
+        fn intercepts_requests(&self) -> bool {
+            true
+        }
+
+        fn during_call(
+            &self,
+            mut request: OcrDuringCallRequest,
+        ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
+            let document = self.0.clone();
+            Box::pin(async move {
+                request.body["document"] = document;
+                Ok(request)
+            })
+        }
+    }
+
+    #[tokio::test]
+    async fn remote_document_stays_inlined_when_hook_echoes_caller_document() {
+        let (base, seen, server) = mock_server(vec![
+            MockResponse::json(json!("served document")),
+            MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}],"usage_info":{"pages_processed":1}})),
+        ])
+        .await;
+        let document_url = format!("{base}/document.pdf");
+        let mut request = crate::ocr::test_support::with_source(
+            wire_request("azure_ai/model", &base, json!({})),
+            &document_url,
+        );
+        request.hooks = Arc::new(EchoCallerDocument(
+            json!({"type":"document_url","document_url":document_url}),
+        ));
+
+        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(), 2);
+        assert!(requests[0].starts_with("GET /document.pdf "));
+        let body: Value =
+            serde_json::from_str(requests[1].split_once("\r\n\r\n").unwrap().1).unwrap();
+        assert_eq!(
+            body["document"]["document_url"],
+            json!("data:application/json;base64,InNlcnZlZCBkb2N1bWVudCI=")
+        );
+    }
 }
diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs
index 91da5a9613d..111c5f7e97a 100644
--- a/litellm-rust/crates/core/src/ocr/prepare.rs
+++ b/litellm-rust/crates/core/src/ocr/prepare.rs
@@ -34,6 +34,14 @@ where
                 .then(|| "document".to_string()),
         )
         .collect();
+    let original_document =
+        serde_json::to_value(&request.document).map_err(|_| super::Error::RequestField {
+            path: "document".into(),
+        })?;
+    let prepared_document = composed
+        .get("document")
+        .filter(|prepared| **prepared != original_document)
+        .cloned();
     let (body, headers) = if request.hooks.intercepts_requests() {
         let changed = request
             .hooks
@@ -47,13 +55,19 @@ where
                 retained_fields,
             })
             .await?;
-        if !changed.body.is_object() {
+        let Value::Object(mut fields) = changed.body else {
             return Err(super::Error::RequestField {
                 path: "guardrail.body".into(),
             });
+        };
+        if let Some(prepared) =
+            prepared_document.filter(|_| fields.get("document") == Some(&original_document))
+        {
+            fields.insert("document".into(), prepared);
         }
-        validate(&changed.body)?;
-        (changed.body, changed.headers)
+        let body = Value::Object(fields);
+        validate(&body)?;
+        (body, changed.headers)
     } else {
         (composed, headers.to_vec())
     };