feat(ocr): async-native, typed Rust OCR foundation (ocr/aocr bridge) - #31128
feat(ocr): async-native, typed Rust OCR foundation (ocr/aocr bridge)#31128ishaan-berri wants to merge 16 commits into
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR builds an async-native Rust OCR foundation on top of the Mistral beachhead from #31033. The Rust core is converted from blocking to async (
Confidence Score: 3/5The Rust core and type model are well-structured, but two specific regression-guard tests were removed and the async secret-resolution path can block the event loop for remote backends. Two targeted tests that previously guarded distinct behaviors were removed without replacement: the test verifying the explicit API key takes priority over the resolver (security-relevant, since a silent fallback to a secret-manager key would be wrong), and the test verifying that
|
| Filename | Overview |
|---|---|
| litellm/ocr/main.py | Adds async Rust path (_arun_rust_ocr) and provider-aware routing; _rust_ocr_pre_call runs synchronously on the event loop in the async path, which can block for remote secret backends |
| litellm/ocr/rust_bridge.py | Renames bridge abstraction from callable RustOcr to object RustBridge with ocr/aocr methods; adds RUST_SUPPORTED_PROVIDERS frozenset and rust_supports() guard |
| tests/test_litellm/ocr/test_rust_bridge.py | Adds async aocr coverage and RecordingBridge; removes two regression-guard tests (test_run_rust_ocr_prefers_explicit_key_over_resolver, test_load_rust_ocr_uses_compiled_extension) without equivalent replacements |
| litellm-rust/crates/python-bridge/src/lib.rs | Adds aocr PyO3 function using pyo3-async-runtimes; note_offload() increments before future_into_py completes, potentially overcounting GIL-release stats |
| litellm-rust/crates/providers/src/ocr.rs | Converts run_ocr from blocking to async, dispatches on OcrProvider, switches to shared async reqwest::Client; clean provider dispatch with typed errors |
| litellm-rust/crates/core/src/ocr/types.rs | Replaces serde_json::Value-heavy types with full typed OCR domain model; serde derives match wire JSON; only three Value leaves where Python itself is Any |
| litellm-rust/crates/providers/src/mistral/ocr/transformation.rs | Replaces imperative JSON construction with typed MistralRequestBody struct and parse_response; cleaner and testable pure functions |
Comments Outside Diff (1)
-
tests/test_litellm/ocr/test_rust_bridge.py, line 144-145 (link)Two regression-guard tests removed without replacement
test_run_rust_ocr_prefers_explicit_key_over_resolververified that when an explicitapi_keyis supplied,resolve_api_key(i.e.get_secret_str) is never called. The new code preserves this viacall.api_key or (call.resolve_api_key(secret_env) if secret_env else None), but without the test a future change to key-resolution order would go undetected — this is security-relevant, as it prevents accidentally fetching a wrong key from a vault backend.test_load_rust_ocr_uses_compiled_extensionwas the only test covering the non-injection path inload_rust_bridge(). Critically, the behaviour changed: the old code returnedlitellm_python_bridge.ocr(the callable), while the new code returnslitellm_python_bridge(the module). That difference is now untested — a regression to the old return value would silently make every call fail with anAttributeErrorat runtime.Rule Used: What: Flag any modifications to existing tests and... (source)
Reviews (1): Last reviewed commit: "docs(rust-ocr): AGENTS.md explaining the..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| async def _arun_rust_ocr(bridge: RustBridge, call: _RustOcrCall) -> OCRResponse: | ||
| """Async Rust OCR: pre_call, then await the bridge awaitable. | ||
|
|
||
| The HTTP call runs on the Rust Tokio runtime with the GIL released, so no | ||
| event-loop thread is held during the request — unlike the ``run_in_executor`` | ||
| path, this scales to many concurrent OCR calls on a handful of threads. | ||
| """ | ||
| resolved_api_key = _rust_ocr_pre_call(call) | ||
| raw = await bridge.aocr( | ||
| provider=call.provider, | ||
| model=call.model, | ||
| document=call.document, | ||
| api_key=resolved_api_key, | ||
| api_base=call.api_base, | ||
| extra_headers=call.extra_headers, | ||
| timeout_seconds=call.timeout_seconds, | ||
| params=call.optional_params, | ||
| ) | ||
| return OCRResponse.model_validate(raw) |
There was a problem hiding this comment.
_rust_ocr_pre_call is called synchronously inside this coroutine, which runs on the asyncio event-loop thread when aocr awaits it. The call.resolve_api_key(secret_env) call is get_secret_str, which for remote backends (AWS Secrets Manager, GCP Secret Manager, Azure Key Vault) performs a synchronous network request. That call will block the event loop for the duration of the secret lookup.
For the sync _run_rust_ocr path this is fine — it runs inside run_in_executor. The _arun_rust_ocr path was intended to never hold a thread, but it does block the event loop for the secret resolution step when a remote backend is configured. Moving _rust_ocr_pre_call into the executor step (before returning the coroutine) or wrapping it in asyncio.to_thread would preserve the non-blocking guarantee.
| gil::note_offload(); | ||
| pyo3_async_runtimes::tokio::future_into_py(py, async move { | ||
| match core_ocr(request).await { | ||
| Ok(response) => Python::with_gil(|py| response_to_py(py, response)), | ||
| Err(err) => Err(core_error_to_pyerr(err)), | ||
| } | ||
| }) |
There was a problem hiding this comment.
future_into_py succeeds
note_offload() fires unconditionally before pyo3_async_runtimes::tokio::future_into_py(...). If future_into_py returns an Err (e.g. no active Tokio runtime, Python GIL state issue), the error is propagated to the caller but the offload counter has already been bumped. The gil_stats surface would then overstate the number of times the GIL was genuinely released, making the diagnostic less reliable for profiling and load testing. Moving note_offload() inside the async closure — after core_ocr(request).await is scheduled — would keep the count accurate.
2 participants
Relevant issues
Follow-up to #31033 (now merged). Builds the async-native, typed OCR foundation on top of the opt-in Mistral beachhead.
Pre-Submission checklist
tests/test_litellm/ocr/test_rust_bridge.py(syncocr+ asyncaocrrouting, fallback) plus Rust unit tests acrosscore/providers.cargo fmt/clippy -D warnings/cargo test --workspace --lockedpass.Type
🆕 New Feature
🧹 Refactor
Changes
The proxy calls
aocr(), so the Rust OCR path is now async-first instead of blocking-in-a-thread-pool.litellm_providers::ocr(OcrRequest) -> OcrResponse(renamed fromrun_ocr) on a shared asyncreqwestclient.aocr()returns a Python awaitable (pyo3-async-runtimes+ Tokio);ocr()block_ons the same future with the GIL released. Dropsrun_in_executor, so the HTTP wait no longer pins a thread-pool worker per request.coregets real types (OcrRequest,OcrResponse,OcrProvider,OcrDocument,OcrParams, page/usage types) instead ofserde_json::Value. Serde derives produce the exact wire JSON. The onlyValues left are the three leaves where Python itself isAny/Dict[str, Any].match request.provider(Mistral wired; others return a typedUnsupportedProvider). The Python shell gates withrust_supports()and falls back to the Python path otherwise.litellm-rust/AGENTS.md: explains the bridge — crate layout, async model, request lifecycle, error mapping, and the recipe to add a provider.Names mirror Python end to end:
litellm.ocr/aocr→ bridgeocr/aocr→litellm_providers::ocr. Not a UI change — no screenshots.