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