fix(ocr): preserve native preparation contract

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

@@ -7,16 +7,25 @@ pub mod types;
 use serde_json::Value;
 
 use crate::Error;
-use handler::execute_ocr_provider_call;
-use prepare::prepare_ocr_call;
-pub use types::OcrRequest;
+pub use types::{OcrRequest, PreparedOcrRequest};
 
 /// The injected client must disable automatic redirects so document URLs are
 /// validated before each redirect is followed
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn ocr(client: &reqwest::Client, request: OcrRequest<'_>) -> Result<Value, Error> {
-    let provider_request = prepare_ocr_call(request).await?;
-    execute_ocr_provider_call(client, provider_request).await
+    let provider_request = prepare(request).await?;
+    execute(client, provider_request).await
+}
+
+pub async fn prepare(request: OcrRequest<'_>) -> Result<PreparedOcrRequest, Error> {
+    prepare::prepare_ocr_call(request).await
+}
+
+pub async fn execute(
+    client: &reqwest::Client,
+    request: PreparedOcrRequest,
+) -> Result<Value, Error> {
+    handler::execute_ocr_provider_call(client, request).await
 }
 
 #[cfg(test)]

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

@@ -2,12 +2,18 @@ use serde_json::{Map, Value};
 
 use super::common_utils::ocr_provider_config;
 use super::types::{OcrRequest, PreparedOcrRequest};
-use crate::error::Error;
+use crate::error::{Error, json_type_name};
 use crate::ocr::transformation::OcrProviderConfig;
 use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
 
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) async fn prepare_ocr_call(request: OcrRequest<'_>) -> Result<PreparedOcrRequest, Error> {
+    if !request.document.is_object() {
+        return Err(Error::InvalidType {
+            expected: "object",
+            actual: json_type_name(&request.document),
+        });
+    }
     let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
         .or_else(|| {
             request

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

@@ -17,7 +17,7 @@ pub struct OcrRequest<'a> {
     pub max_document_download_bytes: u64,
 }
 
-pub(super) struct PreparedOcrRequest {
+pub struct PreparedOcrRequest {
     pub(super) model: String,
     pub(super) config: &'static dyn OcrProviderConfig,
     pub(super) document: Value,

litellm-rust/crates/python-bridge/src/routes/ocr.rs

@@ -1,11 +1,10 @@
 use std::future::Future;
 
-use litellm_ai_gateway::io::ocr::{
-    OcrRequest as GatewayOcrRequest, PreparedOcrDispatch, execute_ocr_dispatch,
-    prepare_ocr_dispatch,
-};
 use litellm_core::Error;
-use litellm_core::ocr::{OcrRequest, ocr as run_ocr};
+use litellm_core::ocr::{
+    OcrRequest, PreparedOcrRequest, execute as execute_ocr, ocr as run_ocr,
+    prepare as prepare_core_ocr,
+};
 use litellm_core::routing_utils::provider::get_custom_llm_provider;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
@@ -20,13 +19,14 @@ use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
 
 #[pyclass]
 struct PreparedOcr {
-    dispatch: Option<PreparedOcrDispatch>,
+    request: Option<PreparedOcrRequest>,
 }
 
 struct OwnedOcrRequest {
     document: Value,
     options: RouteOptions,
     optional_params: serde_json::Map<String, Value>,
+    max_document_download_bytes: u64,
 }
 
 #[pyfunction]
@@ -76,6 +76,7 @@ fn owned_ocr_request(inputs: OcrInputs) -> PyResult<OwnedOcrRequest> {
         document,
         options,
         optional_params,
+        max_document_download_bytes: inputs.max_document_download_bytes,
     })
 }
 
@@ -88,6 +89,7 @@ fn prepare_ocr(
             document,
             options,
             optional_params,
+            max_document_download_bytes,
         } = request;
         let client = shared_http_client().map_err(Error::Network)?;
         let RouteOptions {
@@ -109,7 +111,7 @@ fn prepare_ocr(
                 extra_headers,
                 optional_params,
                 timeout,
-                max_document_download_bytes: inputs.max_document_download_bytes,
+                max_document_download_bytes,
             },
         )
         .await
@@ -125,6 +127,7 @@ fn prepare_dispatch(
             document,
             options,
             optional_params,
+            max_document_download_bytes,
         } = request;
         let RouteOptions {
             model,
@@ -139,7 +142,7 @@ fn prepare_dispatch(
         {
             return Err(Error::InvalidRequest(reason));
         }
-        let dispatch = prepare_ocr_dispatch(GatewayOcrRequest {
+        let request = prepare_core_ocr(OcrRequest {
             model: &model,
             document,
             api_key: api_key.as_deref(),
@@ -148,25 +151,23 @@ fn prepare_dispatch(
             extra_headers,
             optional_params,
             timeout,
-            callbacks: Vec::new(),
-            guardrails: Vec::new(),
-            request_metadata: Default::default(),
-            litellm_call_id: None,
+            max_document_download_bytes,
         })
         .await?;
         Ok(PreparedOcr {
-            dispatch: Some(dispatch),
+            request: Some(request),
         })
     })
 }
 
 #[pyfunction(name = "ocr_prepare")]
-#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
+#[pyo3(signature = (model, document, max_document_download_bytes, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
 #[allow(clippy::too_many_arguments)]
 fn ocr_prepare(
     py: Python<'_>,
     model: String,
     #[pyo3(from_py_with = litellm_python_interop::from_py)] document: Value,
+    max_document_download_bytes: u64,
     api_key: Option<String>,
     api_base: Option<String>,
     custom_llm_provider: Option<String>,
@@ -183,18 +184,20 @@ fn ocr_prepare(
         extra_headers,
         optional_params,
         timeout_seconds,
+        max_document_download_bytes,
     })?;
     let prepared = execution::run_sync_value(py, future, ocr_prepare_error_to_pyerr)?;
     Py::new(py, prepared).map(Py::into_any)
 }
 
 #[pyfunction(name = "aocr_prepare")]
-#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
+#[pyo3(signature = (model, document, max_document_download_bytes, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
 #[allow(clippy::too_many_arguments)]
 fn aocr_prepare<'py>(
     py: Python<'py>,
     model: String,
     #[pyo3(from_py_with = litellm_python_interop::from_py)] document: Value,
+    max_document_download_bytes: u64,
     api_key: Option<String>,
     api_base: Option<String>,
     custom_llm_provider: Option<String>,
@@ -211,6 +214,7 @@ fn aocr_prepare<'py>(
         extra_headers,
         optional_params,
         timeout_seconds,
+        max_document_download_bytes,
     })?;
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         let prepared = execution::await_value(future, ocr_prepare_error_to_pyerr).await?;
@@ -218,24 +222,29 @@ fn aocr_prepare<'py>(
     })
 }
 
-fn take_dispatch(py: Python<'_>, prepared: Py<PreparedOcr>) -> PyResult<PreparedOcrDispatch> {
+fn take_request(py: Python<'_>, prepared: Py<PreparedOcr>) -> PyResult<PreparedOcrRequest> {
     prepared
         .borrow_mut(py)
-        .dispatch
+        .request
         .take()
         .ok_or_else(|| PyRuntimeError::new_err("prepared OCR request was already executed"))
 }
 
+async fn execute_prepared_ocr(request: PreparedOcrRequest) -> Result<Value, Error> {
+    let client = shared_http_client().map_err(Error::Network)?;
+    execute_ocr(&client, request).await
+}
+
 #[pyfunction]
 fn ocr_execute(py: Python<'_>, prepared: Py<PreparedOcr>) -> PyResult<Py<PyAny>> {
-    let dispatch = take_dispatch(py, prepared)?;
-    execution::run_sync(py, execute_ocr_dispatch(dispatch), ocr_error_to_pyerr)
+    let request = take_request(py, prepared)?;
+    execution::run_sync(py, execute_prepared_ocr(request), ocr_error_to_pyerr)
 }
 
 #[pyfunction]
 fn aocr_execute<'py>(py: Python<'py>, prepared: Py<PreparedOcr>) -> PyResult<Bound<'py, PyAny>> {
-    let dispatch = take_dispatch(py, prepared)?;
-    execution::run_async(py, execute_ocr_dispatch(dispatch), ocr_error_to_pyerr)
+    let request = take_request(py, prepared)?;
+    execution::run_async(py, execute_prepared_ocr(request), ocr_error_to_pyerr)
 }
 
 bridge_route! {
@@ -306,6 +315,7 @@ mod tests {
             extra_headers: None,
             optional_params: Some(json!({})),
             timeout_seconds: None,
+            max_document_download_bytes: 50 * 1024 * 1024,
         })
         .expect("request should marshal");
 

litellm/ocr/main.py

@@ -188,8 +188,17 @@ def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool:
     return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override)
 
 
+def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
+    if prepared_request.custom_llm_provider != "azure_ai":
+        return True
+    return not any(
+        prepared_request.litellm_params.get(name) is not None
+        for name in ("azure_ad_token_provider", "azure_username", "azure_password")
+    )
+
+
 def _should_attempt_rust_ocr(prepared_request: _PreparedOCRRequest) -> bool:
-    return _rust_ocr_enabled(prepared_request)
+    return _rust_ocr_enabled(prepared_request) and _rust_ocr_supported(prepared_request)
 
 
 def _rust_bridge_optional_params(
@@ -312,6 +321,7 @@ def _prepare_rust_ocr_call(
 def _log_rust_ocr_call(
     request: rust_ocr_bridge.RustOCRRequest,
     logging_obj: LiteLLMLoggingObj,
+    optional_params: Mapping[str, object],
 ) -> None:
     logging_obj.pre_call(
         input="OCR document processing",
@@ -319,8 +329,7 @@ def _log_rust_ocr_call(
         additional_args={
             "complete_input_dict": {
                 "model": request.model,
-                "document": request.document,
-                **request.optional_params,
+                **optional_params,
             },
             "api_base": request.logging_api_base,
             "headers": request.extra_headers,
@@ -371,6 +380,7 @@ def _run_rust_ocr(
         on_accepted=lambda request: _log_rust_ocr_call(
             request=request,
             logging_obj=prepared_request.litellm_logging_obj,
+            optional_params=prepared_request.optional_params,
         ),
     )
     return _rust_ocr_response(attempt_result, context)
@@ -394,6 +404,7 @@ async def _run_rust_aocr(
         on_accepted=lambda request: _log_rust_ocr_call(
             request=request,
             logging_obj=prepared_request.litellm_logging_obj,
+            optional_params=prepared_request.optional_params,
         ),
     )
     return _rust_ocr_response(attempt_result, context)

litellm/rust_bridge/ocr.py

@@ -78,6 +78,7 @@ class RustOcrPrepare(Protocol):
         self,
         model: str,
         document: Mapping[str, object],
+        max_document_download_bytes: int,
         api_key: str | None,
         api_base: str | None,
         custom_llm_provider: str | None,
@@ -94,6 +95,7 @@ class RustAocrPrepare(Protocol):
         self,
         model: str,
         document: Mapping[str, object],
+        max_document_download_bytes: int,
         api_key: str | None,
         api_base: str | None,
         custom_llm_provider: str | None,