From fa5e27d9bbbcb825975f5212664bda1641992070 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 4 Sep 2026 12:01:30 -0700 Subject: [PATCH] fix(ocr): preserve native preparation contract --- litellm-rust/crates/core/src/ocr/mod.rs | 19 +++++-- litellm-rust/crates/core/src/ocr/prepare.rs | 8 ++- litellm-rust/crates/core/src/ocr/types.rs | 2 +- .../crates/python-bridge/src/routes/ocr.rs | 52 +++++++++++-------- litellm/ocr/main.py | 17 ++++-- litellm/rust_bridge/ocr.py | 2 + 6 files changed, 69 insertions(+), 31 deletions(-) diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 764732fbfaa..189d055bf79 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/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 { - 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 { + prepare::prepare_ocr_call(request).await +} + +pub async fn execute( + client: &reqwest::Client, + request: PreparedOcrRequest, +) -> Result { + handler::execute_ocr_provider_call(client, request).await } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 97e8d42e934..9ec8e3c1c2d 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/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 { + 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 diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index bbf48c856cf..3220d462196 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/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, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 641ba36acfe..8be0727568a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/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, + request: Option, } struct OwnedOcrRequest { document: Value, options: RouteOptions, optional_params: serde_json::Map, + max_document_download_bytes: u64, } #[pyfunction] @@ -76,6 +76,7 @@ fn owned_ocr_request(inputs: OcrInputs) -> PyResult { 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, api_base: Option, custom_llm_provider: Option, @@ -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, api_base: Option, custom_llm_provider: Option, @@ -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) -> PyResult { +fn take_request(py: Python<'_>, prepared: Py) -> PyResult { 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 { + let client = shared_http_client().map_err(Error::Network)?; + execute_ocr(&client, request).await +} + #[pyfunction] fn ocr_execute(py: Python<'_>, prepared: Py) -> PyResult> { - 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) -> PyResult> { - 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"); diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 6882842d493..8d26e50ba11 100644 --- a/litellm/ocr/main.py +++ b/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) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 851f3db010e..60108e3e6e9 100644 --- a/litellm/rust_bridge/ocr.py +++ b/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,